148 lines
3.6 KiB
Python
148 lines
3.6 KiB
Python
"""
|
|
Async bridge for synchronous OCI SSE streams.
|
|
|
|
The bridge runs blocking SSE iteration in a worker thread and
|
|
delivers events to an asynchronous consumer with backpressure.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import threading
|
|
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def stream_sse_events(
|
|
stream: Any,
|
|
*,
|
|
max_queue_size: int = 16,
|
|
shutdown_timeout: float = 5.0,
|
|
) -> AsyncIterator[str]:
|
|
"""
|
|
Convert a synchronous SSE stream into asynchronous text events.
|
|
|
|
The stream must provide:
|
|
events() -> Iterator
|
|
close() -> None
|
|
|
|
Cancellation stops delivery, requests stream closure and
|
|
waits briefly for the worker to finish.
|
|
|
|
Worker termination is not guaranteed if the underlying HTTP
|
|
read ignores stream closure.
|
|
"""
|
|
if max_queue_size < 1:
|
|
raise ValueError("max_queue_size must be positive.")
|
|
|
|
if shutdown_timeout <= 0:
|
|
raise ValueError("shutdown_timeout must be positive.")
|
|
|
|
loop = asyncio.get_running_loop()
|
|
|
|
queue: asyncio.Queue[tuple[str, Any]] = asyncio.Queue(
|
|
maxsize=max_queue_size,
|
|
)
|
|
|
|
stop_requested = threading.Event()
|
|
|
|
def publish(kind: str, value: Any = None) -> bool:
|
|
if stop_requested.is_set():
|
|
return False
|
|
|
|
future = asyncio.run_coroutine_threadsafe(
|
|
queue.put((kind, value)),
|
|
loop,
|
|
)
|
|
|
|
while not stop_requested.is_set():
|
|
try:
|
|
future.result(timeout=0.1)
|
|
return True
|
|
except TimeoutError:
|
|
continue
|
|
except Exception:
|
|
future.cancel()
|
|
raise
|
|
|
|
future.cancel()
|
|
return False
|
|
|
|
def worker() -> None:
|
|
try:
|
|
for event in stream.events():
|
|
if stop_requested.is_set():
|
|
break
|
|
|
|
data = getattr(event, "data", None)
|
|
|
|
if not isinstance(data, str):
|
|
continue
|
|
|
|
if not publish("data", data):
|
|
break
|
|
|
|
except Exception as exc:
|
|
if not stop_requested.is_set():
|
|
publish("error", exc)
|
|
|
|
finally:
|
|
if not stop_requested.is_set():
|
|
publish("done")
|
|
|
|
worker_task = asyncio.create_task(
|
|
asyncio.to_thread(worker)
|
|
)
|
|
|
|
try:
|
|
while True:
|
|
kind, value = await queue.get()
|
|
|
|
if kind == "done":
|
|
break
|
|
|
|
if kind == "error":
|
|
raise value
|
|
|
|
yield value
|
|
|
|
finally:
|
|
stop_requested.set()
|
|
|
|
# Request transport closure outside the event loop.
|
|
try:
|
|
await asyncio.wait_for(
|
|
asyncio.to_thread(stream.close),
|
|
timeout=shutdown_timeout,
|
|
)
|
|
except TimeoutError:
|
|
logger.error(
|
|
"OCI SSE stream closure timed out."
|
|
)
|
|
except Exception:
|
|
logger.exception(
|
|
"Failed to close OCI SSE stream."
|
|
)
|
|
|
|
# Do not assume that closing the stream immediately
|
|
# terminates the synchronous reader.
|
|
try:
|
|
await asyncio.wait_for(
|
|
asyncio.shield(worker_task),
|
|
timeout=shutdown_timeout,
|
|
)
|
|
except TimeoutError:
|
|
logger.error(
|
|
"OCI SSE worker did not terminate after closure."
|
|
)
|
|
except Exception:
|
|
logger.exception(
|
|
"OCI SSE worker failed during shutdown."
|
|
)
|
|
|