Repository: bluesky/bluesky
Pull request #1914 (2025-05-28 16:38:11)

Threaded Buffer to Allow Non-Blocking Callbacks
Add buffering capability to Bluesky callbacks

## Description
An implementation of a simple wrapper for general bluesky callbacks that allows the downstream processing to happen in a separate thread without blocking the `RunEngine`. This could be useful for slow I/O operations (e.g. `TiledWriter`).

Intended usage:
```
tw = TiltedWriter(client)
buff_tw = BufferingWrapper(tw)
RE.subscribe(buff_tw)
```

Note that the `BufferingWrapper` must be used with great caution: the wrapped callback should be thread-safe and must not be subscribed to the RE directly. The errors encountered in the wrapped callback will not be propagated to the RE, and hence there is a possibility for data loss.

## Motivation and Context
This is mainly a development feature to allow subscribing (potentially many) experimental callbacks to a `RunEngine` without the risk of blocking the experiment. The use in production is not encouraged; for critical workflows, it is generally recommended to have a backup (blocking) writing callback.

## How Has This Been Tested?
* `src/bluesky/tests/test_buffering.py`: Added tests for `BufferingWrapper` covering:
  - Delegation of calls to the wrapped callback.
  - Graceful shutdown, ensuring no new items are processed post-shutdown.
  - Proper handling of signal-triggered shutdowns.
  - Validation of shutdown behavior with and without waiting.
  - Thread safety and robustness against double shutdowns



Patch:
diff --git a/src/bluesky/callbacks/buffer.py b/src/bluesky/callbacks/buffer.py
new file mode 100644
index 00000000..3820aa5c
--- /dev/null
+++ b/src/bluesky/callbacks/buffer.py
@@ -0,0 +1,97 @@
+import atexit
+import logging
+import threading
+from queue import Empty, Full, Queue
+from typing import Callable
+
+logger = logging.getLogger(__name__)
+
+
+class BufferingWrapper:
+    """A wrapper for callbacks that processes documents in a separate thread.
+
+    This class allows a callback to be executed in a background thread, processing
+    documents as they are received. This prevent the blocking of RunEngine on any
+    slow I/O operations by the callback. It handles graceful shutdown on exit or signal
+    termination, ensuring that no new documents are accepted after shutdown has been
+    initiated.
+
+    The wrapped callback should be thread-safe and not subscribed to the RE directly.
+    If it maintains shared mutable state, it must protect it using internal locking.
+
+    This is mainly a development feature to allow subscribing (potentially many)
+    experimental callbacks to a `RunEngine` without the risk of blocking the experiment.
+    The use in production is currently not encouraged (at least not without a proper
+    testing and risk assessment).
+
+    Parameters
+    ----------
+        target : callable
+            The instance of a callback that will be called with the documents.
+            It should accept two parameters: `name` and `doc`.
+        queue_size : int, optional
+            The maximum size of the internal queue. Default is 1,000,000.
+
+    Usage
+    -----
+        tw = TiltedWriter(client)
+        buff_tw = BufferingWrapper(tw)
+        RE.subscribe(buff_tw)
+    """
+
+    def __init__(self, target: Callable, queue_size: int = 1_000_000):
+        self._wrapped_callback = target
+        self._queue: Queue = Queue(maxsize=queue_size)
+        self._stop_event = threading.Event()
+        self._shutdown_lock = threading.Lock()
+
+        self._thread = threading.Thread(target=self._process_queue, daemon=True)
+        self._thread.start()
+
+        atexit.register(self.shutdown)
+
+    def __call__(self, name, doc):
+        if self._stop_event.is_set():
+            raise RuntimeError("Cannot accept new data after shutdown.")
+            # TODO: This can be refactored using the upstream functionality (in Python >= 3.13)
+            # https://docs.python.org/3/library/queue.html#queue.Queue.shutdown
+        try:
+            self._queue.put((name, doc))
+        except Full as e:
+            logger.exception(
+                f"The buffer is full. The {self._wrapped_callback.__class__.__name__} can not keep up with the incoming data. "  # noqa
+                f"Consider increasing the queue size or optimizing the callback processing: {e}"
+            )
+            raise RuntimeError(
+                f"The buffer is full. The {self._wrapped_callback.__class__.__name__} can not keep up with the incoming data. "  # noqa
+                "Consider increasing the queue size or optimizing the callback processing."
+            ) from e
+        except Exception as e:
+            logger.exception(f"Failed to put document {name} in queue: {e}")
+            raise RuntimeError(f"Failed to put document {name} in queue: {e}") from e
+
+    def _process_queue(self):
+        while True:
+            try:
+                if item := self._queue.get(timeout=1):
+                    self._wrapped_callback(*item)  # Delegate to wrapped callback
+                else:
+                    break  # Received sentinel value to stop processing
+            except Empty:
+                if self._stop_event.is_set():
+                    break
+            except Exception as e:
+                logger.exception(f"Exception in {self._wrapped_callback.__class__.__name__}: {e}")
+
+    def shutdown(self, wait: bool = True):
+        if self._stop_event.is_set():
+            return
+        self._stop_event.set()
+        self._queue.put(None)
+
+        atexit.unregister(self.shutdown)
+
+        if wait:
+            self._thread.join()
+
+        logger.info(f"{self._wrapped_callback.__class__.__name__} shut down gracefully.")
diff --git a/src/bluesky/run_engine.py b/src/bluesky/run_engine.py
index ad864f91..f7bd3c54 100644
--- a/src/bluesky/run_engine.py
+++ b/src/bluesky/run_engine.py
@@ -2902,8 +2902,8 @@ def call_in_bluesky_event_loop(coro: typing.Awaitable[T], timeout: typing.Option
         if iscoroutine(coro):
             coro.close()
         raise RuntimeError("Bluesky event loop not running")
-    fut = asyncio.run_coroutine_threadsafe(
-        coro,
+    fut: concurrent.futures.Future = asyncio.run_coroutine_threadsafe(
+        coro,  # type: ignore
         loop=_bluesky_event_loop,
     )
     return fut.result(timeout=timeout)


Test patch:
diff --git a/src/bluesky/tests/test_buffering.py b/src/bluesky/tests/test_buffering.py
new file mode 100644
index 00000000..e73eac1c
--- /dev/null
+++ b/src/bluesky/tests/test_buffering.py
@@ -0,0 +1,204 @@
+import time
+from contextlib import contextmanager
+
+import pytest
+
+from bluesky.callbacks import CallbackBase
+from bluesky.callbacks.buffer import BufferingWrapper
+
+
+class SlowDummyCallback(CallbackBase):
+    """Simulates a slow (blocking) callback for testing threaded wrappers."""
+
+    def __init__(self, delay=0.1):
+        self.delay = delay
+        self.called = []
+
+    def __call__(self, name, doc):
+        time.sleep(self.delay)  # Simulate slow work
+        self.called.append((name, doc))
+
+
+class CallbackWithException(CallbackBase):
+    """A callback that raises an exception when called."""
+
+    def __call__(self, name, doc):
+        raise RuntimeError("This callback always raises an exception.")
+
+
+@pytest.fixture
+def fast_cb():
+    yield SlowDummyCallback(delay=0)
+
+
+@pytest.fixture
+def slow_cb():
+    yield SlowDummyCallback(delay=0.1)
+
+
+@contextmanager
+def wait_for_condition(condition, timeout=3, interval=0.01):
+    """Wait for a condition to become True within a timeout period.
+
+    Parameters
+    ----------
+        condition : callable
+            A function that returns True when the condition is met.
+        timeout : float
+            Maximum time to wait for the condition to become True.
+        interval : float, optional
+            Time to wait between checks of the condition (default is 0.01 seconds).
+
+    Usage
+    -----
+        with wait_for_condition(1, lambda: x > y):
+            pass
+    """
+    start = time.time()
+    while time.time() - start < timeout:
+        if condition():
+            yield
+            return
+        time.sleep(interval)
+    raise TimeoutError("Condition not met within timeout")
+
+
+@pytest.mark.parametrize("cb", ["fast_cb", "slow_cb"])
+def test_calls_are_delegated(cb, request):
+    cb = request.getfixturevalue(cb)
+    buff_cb = BufferingWrapper(cb)
+
+    buff_cb("start", {"x": 1})
+    with wait_for_condition(lambda: ("start", {"x": 1}) in cb.called):
+        assert len(cb.called) == 1
+
+    buff_cb("stop", {"x": 2})
+    with wait_for_condition(lambda: ("stop", {"x": 2}) in cb.called):
+        assert len(cb.called) == 2
+
+
+@pytest.mark.parametrize("cb", ["fast_cb", "slow_cb"])
+def test_calls_are_delegated_and_finished(cb, request):
+    cb = request.getfixturevalue(cb)
+    buff_cb = BufferingWrapper(cb)
+
+    assert buff_cb._thread.is_alive()
+    assert len(cb.called) == 0
+
+    buff_cb("start", {"x": 1})
+    buff_cb("stop", {"x": 2})
+
+    assert buff_cb._thread.is_alive()
+    buff_cb.shutdown()
+    assert not buff_cb._thread.is_alive()
+
+    assert ("start", {"x": 1}) in cb.called
+    assert ("stop", {"x": 2}) in cb.called
+
+
+@pytest.mark.parametrize("cb", ["fast_cb", "slow_cb"])
+def test_graceful_shutdown_blocks_queue(cb, request):
+    cb = request.getfixturevalue(cb)
+    buff_cb = BufferingWrapper(cb)
+
+    buff_cb("event", {"data": 42})
+    buff_cb.shutdown()
+
+    with pytest.raises(RuntimeError):
+        buff_cb("post-shutdown", {"fail": True})
+
+    assert ("event", {"data": 42}) in cb.called
+    assert ("post-shutdown", {"fail": True}) not in cb.called
+
+
+@pytest.mark.parametrize("cb", ["fast_cb", "slow_cb"])
+def test_double_shutdown_does_not_fail(cb, request):
+    cb = request.getfixturevalue(cb)
+    buff_cb = BufferingWrapper(cb)
+
+    buff_cb("one", {})
+    buff_cb.shutdown()
+    buff_cb.shutdown()  # Second shutdown should be a no-op
+
+    assert ("one", {}) in cb.called
+
+
+@pytest.mark.parametrize("cb, expected_min_duration", [("fast_cb", 0.0), ("slow_cb", 0.5)])
+def test_shutdown_waits_for_processing(cb, expected_min_duration, request):
+    cb = request.getfixturevalue(cb)
+    buff_cb = BufferingWrapper(cb)
+
+    for i in range(5):
+        buff_cb("event", {"val": i})
+
+    t0 = time.time()
+    time.sleep(0.02)  # Let the last document enter the queue before shutdown
+    buff_cb.shutdown()
+    duration = time.time() - t0
+    assert duration >= expected_min_duration  # Ensure shutdown waited for processing
+
+    # After shutdown, all should be processed
+    assert len(cb.called) == 5
+    for i in range(5):
+        assert ("event", {"val": i}) in cb.called
+
+
+@pytest.mark.parametrize("cb, expected_max_duration", [("fast_cb", 0.1), ("slow_cb", 0.1)])
+def test_shutdown_without_wait(cb, expected_max_duration, request):
+    cb = request.getfixturevalue(cb)
+    buff_cb = BufferingWrapper(cb)
+
+    for i in range(5):
+        buff_cb("event", {"val": i})
+
+    # Shutdown without waiting — thread may still be running briefly
+    t0 = time.time()
+    buff_cb.shutdown(wait=False)
+    duration = time.time() - t0
+    assert duration < expected_max_duration
+    assert len(cb.called) <= 5
+
+
+@pytest.mark.parametrize("cb", ["fast_cb", "slow_cb"])
+def test_shutdown_stops_processing_new_items(cb, request):
+    cb = request.getfixturevalue(cb)
+    buff_cb = BufferingWrapper(cb)
+
+    buff_cb("one", {})
+    buff_cb.shutdown()
+
+    with pytest.raises(RuntimeError):
+        buff_cb("two", {})
+
+    assert ("one", {}) in cb.called
+    assert ("two", {}) not in cb.called
+
+
+def test_execption_in_callback():
+    """Test that exceptions in the callback are handled gracefully."""
+    cb = CallbackWithException()
+    buff_cb = BufferingWrapper(cb)
+
+    with pytest.raises(RuntimeError, match="This callback always raises an exception."):
+        cb("test", {"data": 123})
+
+    # Ensure that the exception does not crash the buffering wrapper
+    buff_cb("test", {"data": 123})
+    assert buff_cb._thread.is_alive()  # The thread should still be running
+    buff_cb("test", {"data": 123})
+
+
+def test_callback_logging_exceptions(monkeypatch):
+    from types import SimpleNamespace
+    from unittest.mock import MagicMock
+
+    logger = SimpleNamespace(exception=MagicMock())
+    monkeypatch.setattr("bluesky.callbacks.buffer.logger", logger)
+
+    cb = CallbackWithException()
+    buff_cb = BufferingWrapper(cb)
+
+    assert logger.exception.call_count == 0
+    buff_cb("test", {"data": 123})
+    with wait_for_condition(lambda: logger.exception.call_count == 1):
+        assert True