Refactor event pipe handling into context manager and iterable

Author: tpoliaw

src/blueapi/service/interface.py

@@ -1,5 +1,6 @@
 import logging
 from collections.abc import Mapping
+from dataclasses import dataclass
 from functools import cache
 from multiprocessing.connection import Connection
 from typing import Any
@@ -283,27 +284,35 @@ def get_python_env(
     return get_python_environment(config=scratch, name=name, source=source)
 
 
-SubHandle = tuple[int, int, int]
+@dataclass
+class SubHandles:
+    worker: int
+    progress: int
+    data: int
 
 
-def pipe_events(tx: Connection) -> SubHandle:
+def pipe_events(tx: Connection) -> SubHandles:
+    tw = worker()
 
     def handler(
         worker_event: WorkerEvent | DataEvent | ProgressEvent,
         _cor_id: str | None,
     ) -> None:
-        tx.send(worker_event)
 
-    task_worker = worker()
-    w_id = task_worker.worker_events.subscribe(handler)
-    d_id = task_worker.data_events.subscribe(handler)
-    p_id = task_worker.progress_events.subscribe(handler)
-    return (w_id, d_id, p_id)
+        try:
+            tx.send(worker_event)
+        except BrokenPipeError:
+            LOGGER.warning("Sending event to broken pipe")
+            pass
 
+    w = tw.worker_events.subscribe(handler)
+    d = tw.data_events.subscribe(handler)
+    p = tw.progress_events.subscribe(handler)
+    return SubHandles(worker=w, data=d, progress=p)
 
-def unpipe_events(hnd: SubHandle) -> None:
-    task_worker = worker()
-    w, d, p = hnd
-    task_worker.worker_events.unsubscribe(w)
-    task_worker.data_events.unsubscribe(d)
-    task_worker.progress_events.unsubscribe(p)
+
+def unpipe_events(hnd: SubHandles):
+    tw = worker()
+    tw.worker_events.unsubscribe(hnd.worker)
+    tw.data_events.unsubscribe(hnd.data)
+    tw.progress_events.unsubscribe(hnd.progress)

src/blueapi/service/main.py

@@ -2,7 +2,6 @@
 import urllib.parse
 from collections.abc import Awaitable, Callable
 from contextlib import asynccontextmanager
-from multiprocessing import Pipe
 from typing import Annotated, Any
 
 import jwt
@@ -16,9 +15,9 @@
     Request,
     Response,
     WebSocket,
+    WebSocketDisconnect,
     status,
 )
-from fastapi.concurrency import run_in_threadpool
 from fastapi.datastructures import Address
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import RedirectResponse, StreamingResponse
@@ -555,44 +554,45 @@ async def run_plan(
 ):
     user = "alice"
 
-    # ack ws
+    LOGGER.info("Starting WS plan")
     await ws.accept()
-    # accept task request through socket
     rq = await ws.receive_json()
-    # submit task to runner
+    LOGGER.info("Raw request: %s", rq)
     try:
         task_request: TaskRequest = TaskRequest.model_validate(rq)
+        LOGGER.info("Plan request: %s", task_request)
         task_id: str = runner.run(interface.submit_task, task_request, {"user": user})
+        LOGGER.info("Task ID: %s", task_id)
     except ValidationError:
+        LOGGER.error("Args not valid", exc_info=True)
         await ws.close(code=1003, reason="invalid args")
         return
     except KeyError:
+        LOGGER.error("Plan not found", exc_info=True)
         await ws.close(code=1003, reason="unknown plan")
         return
 
-    # add listener to runner
-    tx, rx = Pipe()
-    h = runner.run(interface.pipe_events, tx=tx)
-    # start task
     try:
-        task = WorkerTask(task_id=task_id)
-        runner.run(
-            interface.begin_task,
-            task=task,
-        )
+        with runner.event_pipe() as events:
+            LOGGER.info("Created event pipe")
+            runner.run(interface.begin_task, task=WorkerTask(task_id=task_id))
+            async for evt in events:
+                LOGGER.debug("Event: %s", evt)
+                await ws.send_json(evt.model_dump(mode="json"))
+                if isinstance(evt, WorkerEvent) and evt.is_complete():
+                    LOGGER.info("End of stream")
+                    break
     except WorkerBusyError:
+        LOGGER.error("Worker was busy")
         await ws.close(code=1013, reason="Worker busy")
-        return
-    # pipe events to ws
-    try:
-        while True:
-            event: AnyEvent = await run_in_threadpool(rx.recv)
-            await ws.send_json(event.model_dump(mode="json"))
-            if isinstance(event, WorkerEvent) and event.is_complete():
-                break
-    finally:
+    except WebSocketDisconnect:
+        LOGGER.info("Client disconnected")
+        runner.run(
+            interface.cancel_active_task, failure=True, reason="Client disconnected"
+        )
+    else:
+        LOGGER.info("Plan complete")
         await ws.close()
-        runner.run(interface.unpipe_events, hnd=h)
 
 
 @start_as_current_span(TRACER, "config")

src/blueapi/service/runner.py

@@ -1,10 +1,12 @@
+import asyncio
 import inspect
 import logging
 import signal
 import uuid
-from collections.abc import Callable
+from collections.abc import AsyncIterator, Callable
 from importlib import import_module
 from multiprocessing import Pool, set_start_method
+from multiprocessing.connection import Connection, Pipe
 from multiprocessing.pool import Pool as PoolClass
 from typing import Any, ParamSpec, TypeVar
 
@@ -15,11 +17,13 @@
 )
 from opentelemetry.context import attach
 from opentelemetry.propagate import get_global_textmap
-from pydantic import TypeAdapter
 
 from blueapi.config import ApplicationConfig
-from blueapi.service.interface import setup, teardown
+from blueapi.core.bluesky_types import DataEvent
+from blueapi.service import interface
+from blueapi.service.interface import SubHandles, setup, teardown
 from blueapi.service.model import EnvironmentResponse
+from blueapi.worker.event import ProgressEvent, WorkerEvent
 
 # The default multiprocessing start method is fork
 set_start_method("spawn", force=True)
@@ -145,11 +149,57 @@ def run(
             kwargs,
         )
 
+    def event_pipe(self):
+        return EventPipe(self)
+
     @property
     def state(self) -> EnvironmentResponse:
         return self._state
 
 
+class EventStream:
+    def __init__(self, rx: Connection):
+        self._rx = rx
+
+    def __aiter__(self) -> AsyncIterator[WorkerEvent | DataEvent | ProgressEvent]:
+        return self
+
+    async def __anext__(self) -> WorkerEvent | DataEvent | ProgressEvent:
+        data_available = asyncio.Event()
+        asyncio.get_event_loop().add_reader(self._rx.fileno(), data_available.set)
+        try:
+            while not self._rx.poll():
+                await data_available.wait()
+                data_available.clear()
+            return self._rx.recv()
+        except BrokenPipeError:
+            raise StopAsyncIteration() from None
+        finally:
+            asyncio.get_event_loop().remove_reader(self._rx.fileno())
+
+
+class EventPipe:
+    runner: WorkerDispatcher
+    handles: list[tuple[SubHandles, Connection]]
+
+    def __init__(self, runner: WorkerDispatcher):
+        self.runner = runner
+        self.handles = []
+
+    def __enter__(self) -> EventStream:
+        tx, rx = Pipe()
+        hnd = self.runner.run(interface.pipe_events, tx)
+        LOGGER.debug("Subscribing new event pipe: %s", hnd)
+        self.handles.append((hnd, tx))
+        return EventStream(rx)
+
+    def __exit__(self, *exc):
+        hnd, conn = self.handles.pop()
+        LOGGER.debug("Unsubscribing event pipe: %s", hnd)
+        conn.close()
+        self.runner.run(interface.unpipe_events, hnd)
+
+
 class InvalidRunnerStateError(Exception):
     def __init__(self, message):
         super().__init__(message)
@@ -173,15 +223,7 @@ def import_and_run_function(
     func: Callable[..., T] = _validate_function(
         mod.__dict__.get(function_name, None), function_name
     )
-    value = func(*args, **kwargs)
-    return _valid_return(value, expected_type)
-
-
-def _valid_return(value: Any, expected_type: type[T] | None = None) -> T:
-    if expected_type is None:
-        return value
-    else:
-        return TypeAdapter(expected_type).validate_python(value)
+    return func(*args, **kwargs)
 
 
 def _validate_function(func: Any, function_name: str) -> Callable: