diff --git a/bluesky_httpserver/app.py b/bluesky_httpserver/app.py index f09acb3..0b0ba0b 100644 --- a/bluesky_httpserver/app.py +++ b/bluesky_httpserver/app.py @@ -16,7 +16,7 @@ from fastapi.openapi.utils import get_openapi from .authentication import Mode -from .console_output import CollectPublishedConsoleOutput, ConsoleOutputStream, SystemInfoStream +from .console_output import CollectPublishedConsoleOutput, ConsoleOutputStream, ProgressStream, SystemInfoStream from .core import PatchedStreamingResponse from .database.core import purge_expired from .resources import SERVER_RESOURCES as SR @@ -351,6 +351,8 @@ async def purge_expired_sessions_and_api_keys(): SR.console_output_loader.subscribe(SR.console_output_stream.add_message) SR.set_system_info_stream(SystemInfoStream(rm_ref=RM)) SR.system_info_stream.start() + SR.set_progress_stream(ProgressStream(rm_ref=RM)) + SR.progress_stream.start() # Import module with custom code module_names_str = os.getenv("QSERVER_CUSTOM_MODULES", None) @@ -394,6 +396,7 @@ async def shutdown_event(): await SR.console_output_loader.stop() await SR.console_output_stream.stop() await SR.system_info_stream.stop() + await SR.progress_stream.stop() @lru_cache(1) def override_get_authenticators(): diff --git a/bluesky_httpserver/console_output.py b/bluesky_httpserver/console_output.py index 24b142d..4335cf9 100644 --- a/bluesky_httpserver/console_output.py +++ b/bluesky_httpserver/console_output.py @@ -330,3 +330,72 @@ def start(self): async def stop(self): await self._stop_background_task() await self._RM.system_info_monitor.disable_wait() + + +class ProgressStream: + def __init__(self, *, rm_ref): + self._RM = rm_ref + self._queues = {} + self._background_task = None + self._background_task_running = False + self._background_task_stopped = asyncio.Event() + self._background_task_stopped.set() + self._queue_max_size = 1000 + + @property + def background_task_running(self): + return self._background_task_running + + @property + def queues(self): + return self._queues + + def add_queue(self, key): + """ + Add a new queue to the dictionary of queues. The key is a reference to the socket for + for connection with the client. + """ + queue = asyncio.Queue(maxsize=self._queue_max_size) + self._queues[key] = queue + return queue + + def remove_queue(self, key): + """ + Remove the queue identified by the key from the dictionary of queues. + """ + if key in self._queues: + del self._queues[key] + + def _start_background_task(self): + if not self._background_task_running: + self._background_task = asyncio.create_task(self._load_msgs_task()) + + async def _stop_background_task(self): + self._background_task_running = False + await self._background_task_stopped.wait() + + async def _load_msgs_task(self): + self._background_task_stopped.clear() + self._background_task_running = True + while self._background_task_running: + try: + msg = await self._RM.progress_monitor.next_msg(timeout=0.5) + + if isinstance(msg, dict) and "msg" in msg: + msg_json = json.dumps(msg) + for q in self._queues.values(): + # Protect from overflow. It's ok to discard old messages. + if q.full(): + q.get_nowait() + await q.put(msg_json) + except self._RM.RequestTimeoutError: + pass + self._background_task_stopped.set() + + def start(self): + self._RM.progress_monitor.enable() + self._start_background_task() + + async def stop(self): + await self._stop_background_task() + await self._RM.progress_monitor.disable_wait() diff --git a/bluesky_httpserver/resources.py b/bluesky_httpserver/resources.py index 1dca2ca..108146b 100644 --- a/bluesky_httpserver/resources.py +++ b/bluesky_httpserver/resources.py @@ -60,5 +60,16 @@ def system_info_stream(self): def system_info_stream(self, _): raise RuntimeError("Attempting to set read-only property 'system_info_stream'") + def set_progress_stream(self, progress_stream): + self._progress_stream = progress_stream + + @property + def progress_stream(self): + return self._progress_stream + + @progress_stream.setter + def progress_stream(self, _): + raise RuntimeError("Attempting to set read-only property 'progress_stream'") + SERVER_RESOURCES = _ServerResources() diff --git a/bluesky_httpserver/routers/core_api.py b/bluesky_httpserver/routers/core_api.py index 397972b..a7c1f0b 100644 --- a/bluesky_httpserver/routers/core_api.py +++ b/bluesky_httpserver/routers/core_api.py @@ -1215,3 +1215,29 @@ async def info_ws(websocket: WebSocket, scopes=["read:monitor"]): pass finally: SR.system_info_stream.remove_queue_info(websocket) + + +@router.websocket("/progress/ws") +async def progress_ws(websocket: WebSocket, scopes=["read:monitor"]): + principal = get_current_principal_websocket(websocket=websocket, scopes=scopes) + if not principal: + await websocket.close(code=4001, reason="Invalid token") + return + + await websocket.accept() + q = SR.progress_stream.add_queue(websocket) + wsmon = WebSocketMonitor(websocket) + wsmon.start() + try: + while wsmon.is_alive: + try: + msg = await asyncio.wait_for(q.get(), timeout=1) + await websocket.send_text(msg) + except asyncio.TimeoutError: + pass + except RuntimeError: # 'send' after the client is disconnected + pass + except WebSocketDisconnect: + pass + finally: + SR.progress_stream.remove_queue(websocket) diff --git a/bluesky_httpserver/tests/test_progress_socket.py b/bluesky_httpserver/tests/test_progress_socket.py new file mode 100644 index 0000000..6abbd58 --- /dev/null +++ b/bluesky_httpserver/tests/test_progress_socket.py @@ -0,0 +1,211 @@ +import json +import pprint +import threading +import time as ttime + +import pytest +from bluesky_queueserver.manager.tests.common import ( # noqa F401 + append_code_to_last_startup_file, + copy_default_profile_collection, + re_manager_cmd, +) +from websockets.sync.client import connect + +from bluesky_httpserver.tests.conftest import ( # noqa F401 + API_KEY_FOR_TESTS, + SERVER_ADDRESS, + SERVER_PORT, + fastapi_server_fs, + request_to_json, + set_qserver_zmq_encoding, + wait_for_environment_to_be_closed, + wait_for_environment_to_be_created, + wait_for_queue_execution_to_complete, +) + +# Startup code that defines a slow (delayed) motor and a plan that moves it. Moving a delayed +# motor makes the RE 'waiting_hook' emit incremental watcher updates (name/current/fraction/...), +# as opposed to the instantaneous demo motors which only produce completion messages. +_progress_test_startup_code = """ +from ophyd.sim import SynAxis as _SynAxisProgressTest +from bluesky.plan_stubs import mv as _mv_progress_test + +motor_slow = _SynAxisProgressTest(name="motor_slow", delay=0.6) + + +def progress_test_plan(npts: int = 5): + for _i in range(npts): + yield from _mv_progress_test(motor_slow, _i + 1) +""" + + +class _ReceiveProgressSocket(threading.Thread): + """ + Catch streaming progress updates by connecting to the ``/progress/ws`` socket and + save messages to the buffer. + """ + + def __init__(self, *, endpoint="/progress/ws", api_key=API_KEY_FOR_TESTS, **kwargs): + super().__init__(**kwargs) + self.received_data_buffer = [] + self._exit = False + self._api_key = api_key + self._endpoint = endpoint + + def run(self): + websocket_uri = f"ws://{SERVER_ADDRESS}:{SERVER_PORT}/api{self._endpoint}" + additional_headers = {"Authorization": f"ApiKey {self._api_key}"} + try: + with connect(websocket_uri, additional_headers=additional_headers) as websocket: + while not self._exit: + try: + msg_json = websocket.recv(timeout=0.1, decode=False) + try: + msg = json.loads(msg_json) + self.received_data_buffer.append(msg) + except json.JSONDecodeError: + pass + except TimeoutError: + pass + except Exception as ex: + print(f"Failed to connect to server: {ex}") + + def stop(self): + self._exit = True + + def __del__(self): + self.stop() + + +@pytest.mark.parametrize("zmq_port", (None, 60619)) +def test_http_server_progress_socket_1( + monkeypatch, re_manager_cmd, fastapi_server_fs, zmq_port # noqa F811 +): + """ + Test for the ``/progress/ws`` websocket. Runs a plan that moves a motor so that the + RE Manager publishes ``waiting_hook`` (watcher) progress updates over 0MQ, and verifies + that the updates are streamed to the connected websocket client. + """ + # Start HTTP Server + if zmq_port is not None: + monkeypatch.setenv("QSERVER_ZMQ_INFO_ADDRESS", f"tcp://localhost:{zmq_port}") + fastapi_server_fs() + + # Start RE Manager with progress publishing enabled + params = ["--zmq-publish-progress", "ON"] + if zmq_port is not None: + params.extend(["--zmq-info-addr", f"tcp://*:{zmq_port}"]) + re_manager_cmd(params) + + rps = _ReceiveProgressSocket() + rps.start() + ttime.sleep(1) # Wait until the client connects to the socket + + resp1 = request_to_json("post", "/environment/open") + assert resp1["success"] is True, pprint.pformat(resp1) + assert wait_for_environment_to_be_created(timeout=10) + + # A 'scan' moves 'motor', which makes the RE wait on status objects and publish progress updates + plan = {"name": "scan", "args": [["det"], "motor", -1, 1, 10], "item_type": "plan"} + resp2 = request_to_json("post", "/queue/item/add", json={"item": plan}) + assert resp2["success"] is True, pprint.pformat(resp2) + + resp3 = request_to_json("post", "/queue/start") + assert resp3["success"] is True, pprint.pformat(resp3) + + assert wait_for_queue_execution_to_complete(timeout=30) + + resp4 = request_to_json("post", "/environment/close") + assert resp4["success"] is True, pprint.pformat(resp4) + assert wait_for_environment_to_be_closed(timeout=10) + + # Wait until capture is complete + ttime.sleep(2) + rps.stop() + rps.join() + + buffer = rps.received_data_buffer + assert len(buffer) > 0, "No progress updates were received" + for msg in buffer: + assert "time" in msg, msg + assert isinstance(msg["time"], float), msg + assert "msg" in msg, msg + assert isinstance(msg["msg"], dict), msg + + # The RE sends a completion message each time it finishes waiting on status objects. + completion_msgs = [_ for _ in buffer if _["msg"].get("completed") is True] + assert len(completion_msgs) > 0, pprint.pformat(buffer) + + +def test_http_server_progress_socket_2( + tmp_path, monkeypatch, re_manager_cmd, fastapi_server_fs # noqa F811 +): + """ + Test that incremental (watcher) progress updates are streamed over ``/progress/ws``. + Runs a plan that moves a delayed motor, which makes the RE Manager publish per-status + progress updates (with ``name``/``current``/``fraction``) and not just completion messages. + """ + # Prepare a startup profile with a delayed motor and a plan that moves it + pc_path = copy_default_profile_collection(tmp_path) + append_code_to_last_startup_file(pc_path, _progress_test_startup_code) + + fastapi_server_fs() + + # 'ENVIRONMENT_OPEN' regenerates the list of existing plans/devices so 'progress_test_plan' is known + params = [ + "--zmq-publish-progress", + "ON", + "--startup-dir", + pc_path, + "--update-existing-plans-devices", + "ENVIRONMENT_OPEN", + ] + re_manager_cmd(params) + + rps = _ReceiveProgressSocket() + rps.start() + ttime.sleep(1) # Wait until the client connects to the socket + + resp1 = request_to_json("post", "/environment/open") + assert resp1["success"] is True, pprint.pformat(resp1) + assert wait_for_environment_to_be_created(timeout=10) + + plan = {"name": "progress_test_plan", "kwargs": {"npts": 5}, "item_type": "plan"} + resp2 = request_to_json("post", "/queue/item/add", json={"item": plan}) + assert resp2["success"] is True, pprint.pformat(resp2) + + resp3 = request_to_json("post", "/queue/start") + assert resp3["success"] is True, pprint.pformat(resp3) + + assert wait_for_queue_execution_to_complete(timeout=30) + + resp4 = request_to_json("post", "/environment/close") + assert resp4["success"] is True, pprint.pformat(resp4) + assert wait_for_environment_to_be_closed(timeout=10) + + # Wait until capture is complete + ttime.sleep(2) + rps.stop() + rps.join() + + buffer = rps.received_data_buffer + assert len(buffer) > 0, "No progress updates were received" + for msg in buffer: + assert "time" in msg, msg + assert isinstance(msg["time"], float), msg + assert "msg" in msg, msg + assert isinstance(msg["msg"], dict), msg + + # Incremental watcher updates carry the moved device name (not just {"completed": True}) + incremental_msgs = [_ for _ in buffer if _["msg"].get("name") == "motor_slow"] + assert len(incremental_msgs) > 0, pprint.pformat(buffer) + for msg in incremental_msgs: + m = msg["msg"] + assert "current" in m, msg + assert "fraction" in m, msg + assert "done" in m, msg + + # Completion messages are still sent when each wait finishes + completion_msgs = [_ for _ in buffer if _["msg"].get("completed") is True] + assert len(completion_msgs) > 0, pprint.pformat(buffer) +