From 5f574327ceb2efc24c4f6bce6efc6669e8b02eac Mon Sep 17 00:00:00 2001 From: Richard Bowman Date: Fri, 9 Oct 2026 00:09:32 +0100 Subject: [PATCH 01/10] Extend MessageBroker I've added "stream" as a message type, and provided "next message" and "close streams by affordance" methods. This should prepare the way for using MessageBroker for MJPEG streams. --- src/labthings_fastapi/message_broker.py | 34 ++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/src/labthings_fastapi/message_broker.py b/src/labthings_fastapi/message_broker.py index d7ce98c7..1621170a 100644 --- a/src/labthings_fastapi/message_broker.py +++ b/src/labthings_fastapi/message_broker.py @@ -36,7 +36,7 @@ class Message: thing: str affordance: str - message_type: Literal["property", "action"] + message_type: Literal["property", "action", "stream"] payload: Any @@ -83,6 +83,19 @@ async def subscribe( streams = affordances.setdefault(affordance, WeakSet()) streams.add(stream) + async def next_message(self, thing: str, affordance: str) -> Message: + """Get the next message from a particular affordance. + + Note that there's no timeout: standard `anyio` commands may be used to do that. + + :param thing: The name of the `.Thing` being subscribed to. + :param affordance: The name of the affordance being subscribed to. + :return: the next message from the specified affordance. + """ + send, recv = anyio.create_memory_object_stream[Message](max_buffer_size=1) + await self.subscribe(thing, affordance, send) + return await recv.receive() + async def unsubscribe( self, thing: str, affordance: str, stream: MemoryObjectSendStream[Message] ) -> None: @@ -148,3 +161,22 @@ async def close_streams(self) -> None: for subs in thing_subs.values(): for stream in subs: tg.start_soon(stream.aclose) + + async def close_streams_for_affordance(self, thing: str, affordance: str) -> int: + """Close all streams subscribed to a particular affordance. + + This may be called to signal that a stream has stopped. + + :param thing: the Thing name. + :param affordance: the affordance name. + :return: the number of streams that were closed. + """ + try: + streams = self._subscriptions[thing][affordance] + except KeyError: + return 0 # If there are no streams, ignore it. + async with anyio.create_task_group() as tg: + n = len(streams) + for stream in streams: + tg.start_soon(stream.aclose) + return n From 80d4e59ac49aad4b2dd9745218a20607abc1c7c0 Mon Sep 17 00:00:00 2001 From: Richard Bowman Date: Fri, 9 Oct 2026 00:10:20 +0100 Subject: [PATCH 02/10] Expose the message broker. The thing server interface now exposes the message broker. --- src/labthings_fastapi/thing_server_interface.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/labthings_fastapi/thing_server_interface.py b/src/labthings_fastapi/thing_server_interface.py index 3acbf011..a9d0bb8c 100644 --- a/src/labthings_fastapi/thing_server_interface.py +++ b/src/labthings_fastapi/thing_server_interface.py @@ -20,7 +20,7 @@ from labthings_fastapi.exceptions import FeatureNotEnabledError, ServerNotRunningError from labthings_fastapi.global_lock import GlobalLock -from labthings_fastapi.message_broker import Message +from labthings_fastapi.message_broker import Message, MessageBroker if TYPE_CHECKING: from labthings_fastapi.actions import ActionManager @@ -156,6 +156,11 @@ def publish(self, message: Message) -> None: except ServerNotRunningError: pass # If the server isn't running yet, we can't publish events. + @property + def message_broker(self) -> MessageBroker: + """The message broker associated with the current server.""" + return self._get_server().message_broker + @property def settings_folder(self) -> str: """The path to a folder where persistent files may be saved.""" From ca8460dcd742a4d021b15ad81105091fc22b3e26 Mon Sep 17 00:00:00 2001 From: Richard Bowman Date: Fri, 9 Oct 2026 00:18:29 +0100 Subject: [PATCH 03/10] Refactor MJPEG stream This strips out the pub/sub logic from MJPEGStream in favour of using MessageBroker, as is already done for property notifications. The resulting class should be much simpler. I've also inherited from BaseDescriptor to save duplication in MJPEGStreamDescriptor. --- src/labthings_fastapi/outputs/mjpeg_stream.py | 259 +++++------------- tests/test_mjpeg_stream.py | 1 - 2 files changed, 76 insertions(+), 184 deletions(-) diff --git a/src/labthings_fastapi/outputs/mjpeg_stream.py b/src/labthings_fastapi/outputs/mjpeg_stream.py index 39f4d7e2..5bc60bd9 100644 --- a/src/labthings_fastapi/outputs/mjpeg_stream.py +++ b/src/labthings_fastapi/outputs/mjpeg_stream.py @@ -6,33 +6,28 @@ from __future__ import annotations -import logging import threading from dataclasses import dataclass from datetime import datetime from typing import ( - TYPE_CHECKING, Any, AsyncGenerator, - Literal, Optional, - Union, - overload, ) import anyio +from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from fastapi import FastAPI from fastapi.responses import HTMLResponse, StreamingResponse -from typing_extensions import Self -if TYPE_CHECKING: - from labthings_fastapi.thing import Thing - from labthings_fastapi.thing_server_interface import ThingServerInterface +from labthings_fastapi.base_descriptor import BaseDescriptor +from labthings_fastapi.message_broker import Message +from labthings_fastapi.thing import Thing @dataclass -class RingbufferEntry: - """A single entry in a ringbuffer. +class Frame: + """A single frame in the stream. This structure comprises one frame as a JPEG, plus a timestamp and a buffer index. Each time a frame is added to the stream, it is @@ -51,6 +46,24 @@ class RingbufferEntry: """The index of the frame within the stream.""" +def frame_payload(message: Message) -> Frame: + """Extract a Frame object from a Message. + + This checks that the message is of "stream" type, and checks the + type of its payload. + + :param message: the message containing the Frame. + :return: the Frame object. + :raises TypeError: if the message isn't a stream message containing + a `Frame` object. + """ + if message.message_type != "stream": + raise TypeError("The message {message} was not part of a stream.") + if not isinstance(message.payload, Frame): + raise TypeError("The message {message} didn't contain a Frame.") + return message.payload + + class MJPEGStreamResponse(StreamingResponse): """A StreamingResponse that streams an MJPEG stream. @@ -64,7 +77,10 @@ class MJPEGStreamResponse(StreamingResponse): """The media_type used to describe the endpoint in FastAPI.""" def __init__( - self, gen: AsyncGenerator[bytes, None], status_code: int = 200 + self, + send: MemoryObjectSendStream[Message], + recv: MemoryObjectReceiveStream[Message], + status_code: int = 200, ) -> None: """Set up StreamingResponse that streams an MJPEG stream. @@ -73,18 +89,20 @@ def __init__( types that mark it as an MJPEG stream. This is sufficient to enable it to work in an `img` tag, with the `src` set to the MJPEG stream's endpoint. - It expects an async generator that supplies individual JPEGs to be streamed, - such as the one provided by `.MJPEGStream`. + It expects to get a stream that receives `Message` objects with `Frame` + payloads. Both the send and receive streams are retained, because the send + stream is only weakly referenced by the message broker. NB the ``status_code`` argument is used by FastAPI to set the status code of the response in OpenAPI. - :param gen: an async generator, yielding `bytes` objects each of which is - one image, in JPEG format. + :param send: the send stream subscribed to the MJPEG stream. + :param recv: the receive stream subscribed to the MJPEG stream. :param status_code: The status code associated with the response, by default a 200 code is returned. """ - self.frame_async_generator = gen + self._send_stream = send + self._receive_stream = recv StreamingResponse.__init__( self, self.mjpeg_async_generator(), @@ -101,9 +119,10 @@ async def mjpeg_async_generator(self) -> AsyncGenerator[bytes, None]: :yield: JPEG frames, each with a ``--frame`` marker prepended. """ - async for frame in self.frame_async_generator: + async for message in self._receive_stream: + frame = frame_payload(message) yield b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" - yield frame + yield frame.frame yield b"\r\n" @@ -128,26 +147,19 @@ class MJPEGStream: of new frames, and then retrieving the frame (shortly) afterwards. """ - def __init__( - self, thing_server_interface: ThingServerInterface, ringbuffer_size: int = 10 - ) -> None: + def __init__(self, thing: Thing, name: str) -> None: """Initialise an MJPEG stream. See the class docstring for `.MJPEGStream`. Note that it will often be initialised by `.MJPEGStreamDescriptor`. - :param thing_server_interface: the `~lt.ThingServerInterface` of the - `~lt.Thing` associated with this stream. It's used to run the async - code that relays frames to open connections. - :param ringbuffer_size: The number of frames to retain in - memory, to allow retrieval after the frame has been sent. + :param thing: the `~lt.Thing` on which this stream is defined. + :param name: The attribute name of this stream. """ self._lock = threading.Lock() - self.condition = anyio.Condition() - self._streaming = False - self._ringbuffer: list[RingbufferEntry] = [] - self._thing_server_interface = thing_server_interface - self.reset(ringbuffer_size=ringbuffer_size) + self._thing = thing + self._name = name + self.reset() def reset(self, ringbuffer_size: Optional[int] = None) -> None: """Reset the stream and optionally change the ringbuffer size. @@ -157,16 +169,6 @@ def reset(self, ringbuffer_size: Optional[int] = None) -> None: :param ringbuffer_size: the number of frames to keep in memory. """ with self._lock: - self._streaming = True - n = ringbuffer_size or len(self._ringbuffer) - self._ringbuffer = [ - RingbufferEntry( - frame=b"", - index=-1, - timestamp=datetime.min, - ) - for i in range(n) - ] self.last_frame_i = -1 def stop(self) -> None: @@ -174,53 +176,12 @@ def stop(self) -> None: Stop the stream and cause all clients to disconnect. """ - with self._lock: - self._streaming = False - self._thing_server_interface.start_async_task_soon( - self.notify_stream_stopped - ) - - async def ringbuffer_entry(self, i: int) -> RingbufferEntry: - """Return the ith frame acquired by the camera. - - The ringbuffer means we can retrieve frames even if they are not - the latest frame. Specifying ``i`` also makes it simple to ensure - that every frame in a stream is acquired. - - :param i: The index of the frame to read. - - :return: the frame, together with a timestamp and its index. - - :raise ValueError: if the frame is not available. - """ - if i < 0: - raise ValueError("i must be >= 0") - if i < self.last_frame_i - len(self._ringbuffer) + 2: - raise ValueError("the ith frame has been overwritten") - if i > self.last_frame_i: - # TODO: await the ith frame - raise ValueError("the ith frame has not yet been acquired") - entry = self._ringbuffer[i % len(self._ringbuffer)] - if entry.index != i: - raise ValueError("the ith frame has been overwritten") - return entry - - async def next_frame(self) -> int: - """Wait for the next frame, and return its index. - - This async function will yield until a new frame arrives, then return - its index. The index may then be used to retrieve the new frame - with `.MJPEGStream.ringbuffer_entry`. - - :return: the index of the next frame to arrive. - - :raise StopAsyncIteration: if the stream has stopped. - """ - async with self.condition: - await self.condition.wait() - if not self._streaming: - raise StopAsyncIteration() - return self.last_frame_i + tsi = self._thing._thing_server_interface + tsi.call_async_task( + tsi.message_broker.close_streams_for_affordance, + self._thing.name, + self._name, + ) async def grab_frame(self) -> bytes: """Wait for the next frame, and return it. @@ -230,9 +191,11 @@ async def grab_frame(self) -> bytes: :return: The next JPEG frame, as a `bytes` object. """ - i = await self.next_frame() - entry = await self.ringbuffer_entry(i) - return entry.frame + message = await self._thing._thing_server_interface.message_broker.next_message( + self._thing.name, self._name + ) + payload = frame_payload(message) + return payload.frame async def next_frame_size(self) -> int: """Wait for the next frame and return its size. @@ -241,34 +204,7 @@ async def next_frame_size(self) -> int: :return: The size of the next JPEG frame, in bytes. """ - i = await self.next_frame() - entry = await self.ringbuffer_entry(i) - return len(entry.frame) - - async def frame_async_generator(self) -> AsyncGenerator[bytes, None]: - """Yield frames as bytes objects. - - This generator will return frames from the MJPEG stream. - - Note that this will wait for a new frame each time. There is no - guarantee that we won't skip frames. - - :yield: the frames in sequence, as a `bytes` object containing - JPEG data. - """ - while self._streaming: - try: - i = await self.next_frame() - entry = await self.ringbuffer_entry(i) - yield entry.frame - except StopAsyncIteration: - break - except Exception as e: # noqa: BLE001 - # It's important that errors in the stream don't crash the server. - # This may be something we can remove in the future, now streams stop - # more elegantly. However, it will require careful testing.f - logging.exception(f"Error in stream: {e}, stream stopped") - return + return len(await self.grab_frame()) async def mjpeg_stream_response(self) -> MJPEGStreamResponse: """Return a StreamingResponse that streams an MJPEG stream. @@ -281,7 +217,11 @@ async def mjpeg_stream_response(self) -> MJPEGStreamResponse: :return: a streaming response in MJPEG format. """ - return MJPEGStreamResponse(self.frame_async_generator()) + send, recv = anyio.create_memory_object_stream[Message](max_buffer_size=1) + await self._thing._thing_server_interface.message_broker.subscribe( + self._thing.name, self._name, send + ) + return MJPEGStreamResponse(send, recv) def add_frame(self, frame: bytes, timestamp: Optional[datetime] = None) -> None: """Add a JPEG to the MJPEG stream. @@ -307,40 +247,18 @@ def add_frame(self, frame: bytes, timestamp: Optional[datetime] = None) -> None: ): raise ValueError("Invalid JPEG") with self._lock: - entry = self._ringbuffer[(self.last_frame_i + 1) % len(self._ringbuffer)] - entry.timestamp = timestamp if timestamp is not None else datetime.now() - entry.frame = frame - entry.index = self.last_frame_i + 1 - self._thing_server_interface.start_async_task_soon( - self.notify_new_frame, entry.index - ) - - async def notify_new_frame(self, i: int) -> None: - """Notify any waiting tasks that a new frame is available. - - :param i: The number of the frame (which counts up since the server starts) - """ - async with self.condition: - self.last_frame_i = i - self.condition.notify_all() - - async def notify_stream_stopped(self) -> None: - """Raise an exception in any waiting tasks to signal the stream has stopped. - - This should be run only when streaming has stopped, i.e. ``self._streaming`` - is ``False`` and an error will be raised if this isn't the case. - - :raises RuntimeError: if the stream is still streaming. - """ - if self._streaming is True: - raise RuntimeError( - "This function should only be called when the stream is stopped." + self.last_frame_i += 1 + payload = Frame( + frame=frame, + timestamp=timestamp if timestamp is not None else datetime.now(), + index=self.last_frame_i, ) - async with self.condition: - self.condition.notify_all() + self._thing._thing_server_interface.publish( + Message(self._thing.name, self._name, "stream", payload) + ) -class MJPEGStreamDescriptor: +class MJPEGStreamDescriptor(BaseDescriptor[Thing, MJPEGStream]): """A descriptor that returns a MJPEGStream object when accessed. If this descriptor is added to a `~lt.Thing`, it will create an `.MJPEGStream` @@ -357,48 +275,23 @@ def __init__(self, **kwargs: Any) -> None: :param \**kwargs: keyword arguments are passed to the initialiser of `.MJPEGStream`. """ + super().__init__() self._kwargs: Any = kwargs - def __set_name__(self, _owner: Thing, name: str) -> None: - """Remember the name to which we are assigned. - - The name is important, as it will set the URL of the HTTP endpoint used - to access the stream. - - :param _owner: the `~lt.Thing` to which we are attached. - :param name: the name to which this descriptor is assigned. - """ - self.name = name - - @overload - def __get__(self, obj: Literal[None], type: type | None = None) -> Self: ... # noqa: D105 - - @overload - def __get__(self, obj: Thing, type: type | None = None) -> MJPEGStream: ... # noqa: D105 - - def __get__( - self, obj: Optional[Thing], type: type[Thing] | None = None - ) -> Union[MJPEGStream, Self]: - """Return the MJPEG Stream, or the descriptor object. - - When accessed on the class, this ``__get__`` method will return the descriptor - object. This allows LabThings to add it to the HTTP API. - - When accessed on the object, an `.MJPEGStream` is returned. + def instance_get(self, obj: Thing) -> MJPEGStream: + """Return the MJPEG Stream. - :param obj: the host `~lt.Thing`, or ``None`` if accessed on the class. - :param type: the class on which we are defined. + :param obj: the host `~lt.Thing`. :return: an `.MJPEGStream`, or this descriptor. """ - if obj is None: - return self try: return obj.__dict__[self.name] except KeyError: obj.__dict__[self.name] = MJPEGStream( + thing=obj, + name=self.name, **self._kwargs, - thing_server_interface=obj._thing_server_interface, ) return obj.__dict__[self.name] diff --git a/tests/test_mjpeg_stream.py b/tests/test_mjpeg_stream.py index 1aefa2b8..f8b30cb8 100644 --- a/tests/test_mjpeg_stream.py +++ b/tests/test_mjpeg_stream.py @@ -41,7 +41,6 @@ def _make_images(self): time.sleep(1 / self.framerate) i = i + 1 self.stream.stop() - self._streaming = False @pytest.fixture From a989150162b7d3898103321d2503b0288ee96422 Mon Sep 17 00:00:00 2001 From: Richard Bowman Date: Fri, 9 Oct 2026 00:32:42 +0100 Subject: [PATCH 04/10] Pull over improved tests from #372 This uses an event to speed up tests by eliminating `time.sleep` and adds some more checks on the data received. --- tests/test_mjpeg_stream.py | 171 ++++++++++++++++++++++++++++++++----- 1 file changed, 148 insertions(+), 23 deletions(-) diff --git a/tests/test_mjpeg_stream.py b/tests/test_mjpeg_stream.py index f8b30cb8..76198abc 100644 --- a/tests/test_mjpeg_stream.py +++ b/tests/test_mjpeg_stream.py @@ -12,7 +12,9 @@ class Telly(lt.Thing): _stream_thread: threading.Thread _streaming: bool = False framerate: float = 1000 - frame_limit: int = 3 + frame_limit: int = 999 + frame_event: threading.Event | None = None + initial_delay: float = 0 stream = lt.outputs.MJPEGStreamDescriptor() @@ -23,6 +25,10 @@ def __enter__(self): def __exit__(self, exc_t, exc_v, exc_tb): self._streaming = False + if self.frame_event: + # Trigger an iteration of the loop, so that it + # will terminate rather than hang forever + self.frame_event.set() self._stream_thread.join() def _make_images(self): @@ -35,47 +41,166 @@ def _make_images(self): image.save(dest, "jpeg") jpegs.append(dest.getvalue()) + if self.initial_delay > 0: + time.sleep(self.initial_delay) + i = 0 while self._streaming and (i < self.frame_limit or self.frame_limit < 0): self.stream.add_frame(jpegs[i % len(jpegs)]) - time.sleep(1 / self.framerate) i = i + 1 + if self.frame_event: + self.frame_event.wait() + self.frame_event.clear() + else: + time.sleep(1 / self.framerate) self.stream.stop() + self._streaming = False + + +def assert_magic_bytes(frame: bytes) -> None: + """Check that a `bytes` object starts and ends with the JPEG markers.""" + assert frame[0] == 0xFF + assert frame[1] == 0xD8 + assert frame[-2] == 0xFF + assert frame[-1] == 0xD9 @pytest.fixture -def client(): - """Yield a test client connected to a ThingServer""" - server = lt.ThingServer.from_things({"telly": Telly}) - with server.test_client() as client: - yield client +def server(): + """Yield a ThingServer with the `Telly` thing.""" + return lt.ThingServer.from_things({"telly": Telly}) + + +@pytest.fixture +def telly(server): + """Yield a Telly thing from the server.""" + telly = server.things["telly"] + assert isinstance(telly, Telly) + return telly + +def test_grab_and_shutdown(server: lt.ThingServer, telly: Telly): + """Check we can grab frames, and shut down cleanly. -def test_mjpeg_stream(client): - """Verify the MJPEG stream contains at least one frame marker. + This test uses an Event to synchronise new frames with the various + methods we're calling to retrieve them. This is intended to make the + test suite faster and less reliant on `time.sleep`. + + We check that `grab_frame` and `next_frame_size` both work when the + camera emits frames, and also that they raise an error if they're called + once the camera has stopped. + + We verify that the stream can be stopped, though we don't shut down any + long-running listeners, because limitations in TestClient mean this + isn't possible. It would be possible if we spun up an actual HTTP server, + but that's quite high-effort. + + The async grab functions need to run in the event loop, which happens in + a background thread during the `with server.test_client()` block. We use + the thing server interface to run tasks in the event loop for convenience. + """ + telly.frame_event = threading.Event() # Make timings more deterministic. + with server.test_client(): + # this `with` block starts an event loop and runs the server. + + # Grab some frames and check we get a JPEG back + for _ in range(3): + assert telly._stream_thread.is_alive() # Catch premature termination + # Start the coroutine to grab a frame + future = telly._thing_server_interface.start_async_task_soon( + telly.stream.grab_frame + ) + # then use the event to trigger the next frame. + telly.frame_event.set() + # then wait for the coroutine to finish + frame = future.result() + # this should return a valid JPEG, which we can (sort of) check below + assert_magic_bytes(frame) + + # Repeat the process for grabbing a frame to test `next_frame_size` + future = telly._thing_server_interface.start_async_task_soon( + telly.stream.next_frame_size + ) + telly.frame_event.set() + size = future.result() # wait for the frame to be grabbed + assert isinstance(size, int) + assert size > 0 + + # Close all streams + telly._thing_server_interface.call_async_task(telly.stream.close_streams) + + # We shouldn't be able to get any more frames now + # This means that the stream won't generate any new stream pairs + with pytest.raises(StopAsyncIteration): + telly.stream.connected_stream() + # The grab functions depend on the function above, so they should also fail. + with pytest.raises(StopAsyncIteration): + telly._thing_server_interface.call_async_task(telly.stream.grab_frame) + with pytest.raises(StopAsyncIteration): + telly._thing_server_interface.call_async_task(telly.stream.next_frame_size) + # The background thread gets shut down by `Telly.__exit__`. + + +def test_mjpeg_stream_http(server: lt.ThingServer, telly: Telly): + """Verify the MJPEG stream works, and is shut down cleanly. A limitation of the TestClient is that it can't actually stream. This means that all of the frames sent by our test Thing will arrive in a single packet. - For now, we just check it starts with the frame separator, - but it might be possible in the future to check there are three - images there. + For now, we download all the data and then chop it up afterwards. + + The `Telly` will send exactly 3 JPEGs, with a short delay at + the start to make sure the `StreamingResponse` is created and + doesn't miss any frames. + + This test also verifies the stream is shut down by the server + when it's stopped - if it wasn't terminated by the server, the + `client.stream()` call would hang indefinitely. """ - with client.stream("GET", "/telly/stream") as stream: - stream.raise_for_status() - received = 0 - for b in stream.iter_bytes(): - received += 1 - assert b.startswith(b"--frame") + telly.frame_limit = 3 # stream 3 frames and then stop + telly.initial_delay = 0.05 # give the client time to start listening + with server.test_client() as client: + with client.stream("GET", "/telly/stream") as response: + # Note: we don't actually enter this `with` block until after + # the camera background thread has finished and the stream + # has been closed. + response.raise_for_status() + parts = 0 + mjpeg_data = b"" + for b in response.iter_bytes(): + parts += 1 + mjpeg_data += b + # Due to a quirk in TestClient, we get all the data in a single + # chunk - this limits our ability to test the stream. + # If that's fixed in the future, the assertion below will fail, + # which should prompt us to improve these tests. + assert parts == 1 + + # Check the received data contained the expected number of frames + # We split chunks using the known header, and remove extra white + # space before checking for JPEG start/end bytes. + chunks = mjpeg_data.split(b"--frame\r\nContent-Type: image/jpeg") + n = 0 + for chunk in chunks: + # Check each chunk is a JPEG + stripped = chunk.strip(b"\r\n") + if len(stripped) > 10: # If the chunk doesn't look empty + assert_magic_bytes(stripped) + n += 1 + assert telly.stream.last_frame_i == 2 + assert n == telly.frame_limit if __name__ == "__main__": - import uvicorn - - server = lt.ThingServer.from_things({"telly": Telly}) - telly = server.things["telly"] + """This block allows you to connect manually with a web browser. + + That's helpful, because the tests above don't actually stream anything + as noted in `test_mjpeg_stream_http`. + """ + thing_server = lt.ThingServer.from_things({"telly": Telly}) + telly = thing_server.things["telly"] assert isinstance(telly, Telly) telly.framerate = 6 telly.frame_limit = -1 - uvicorn.run(server.app, port=5000, ws="websockets-sansio") + thing_server.serve() From d15ac8073714e57e0041397dc20dfba819f5c49b Mon Sep 17 00:00:00 2001 From: Richard Bowman Date: Fri, 9 Oct 2026 00:45:13 +0100 Subject: [PATCH 05/10] Remove unused ringbuffer argument --- src/labthings_fastapi/outputs/mjpeg_stream.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/labthings_fastapi/outputs/mjpeg_stream.py b/src/labthings_fastapi/outputs/mjpeg_stream.py index 5bc60bd9..4318007a 100644 --- a/src/labthings_fastapi/outputs/mjpeg_stream.py +++ b/src/labthings_fastapi/outputs/mjpeg_stream.py @@ -161,13 +161,8 @@ def __init__(self, thing: Thing, name: str) -> None: self._name = name self.reset() - def reset(self, ringbuffer_size: Optional[int] = None) -> None: - """Reset the stream and optionally change the ringbuffer size. - - Discard all frames from the ringbuffer and reset the frame index. - - :param ringbuffer_size: the number of frames to keep in memory. - """ + def reset(self) -> None: + """Reset the frame index.""" with self._lock: self.last_frame_i = -1 From 352e6419bb09651faadb152a29d58155a7ae41b1 Mon Sep 17 00:00:00 2001 From: Richard Bowman Date: Fri, 9 Oct 2026 00:46:21 +0100 Subject: [PATCH 06/10] Test the type checking function and remove broken tests. The stream no longer has the notion of being "active" so we can't check for errors there. This might be a feature we want to revive - I'm not totally sure. --- tests/test_mjpeg_stream.py | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/tests/test_mjpeg_stream.py b/tests/test_mjpeg_stream.py index 76198abc..5028e8da 100644 --- a/tests/test_mjpeg_stream.py +++ b/tests/test_mjpeg_stream.py @@ -1,11 +1,14 @@ import io import threading import time +from datetime import datetime import pytest from PIL import Image import labthings_fastapi as lt +from labthings_fastapi.message_broker import Message +from labthings_fastapi.outputs.mjpeg_stream import Frame, frame_payload class Telly(lt.Thing): @@ -79,6 +82,21 @@ def telly(server): return telly +def test_frame_payload(): + """Check that a frame is returned, or we get an error.""" + frame = Frame(b"payload", datetime.now(), 1) + good_message = Message("thing", "affordance", "stream", frame) + assert frame_payload(good_message) is frame + + bad_message = Message("thing", "affordance", "action", frame) + with pytest.raises(TypeError): + frame_payload(bad_message) + + bad_message2 = Message("thing", "affordance", "stream", b"payload") + with pytest.raises(TypeError): + frame_payload(bad_message2) + + def test_grab_and_shutdown(server: lt.ThingServer, telly: Telly): """Check we can grab frames, and shut down cleanly. @@ -127,17 +145,8 @@ def test_grab_and_shutdown(server: lt.ThingServer, telly: Telly): assert size > 0 # Close all streams - telly._thing_server_interface.call_async_task(telly.stream.close_streams) - - # We shouldn't be able to get any more frames now - # This means that the stream won't generate any new stream pairs - with pytest.raises(StopAsyncIteration): - telly.stream.connected_stream() - # The grab functions depend on the function above, so they should also fail. - with pytest.raises(StopAsyncIteration): - telly._thing_server_interface.call_async_task(telly.stream.grab_frame) - with pytest.raises(StopAsyncIteration): - telly._thing_server_interface.call_async_task(telly.stream.next_frame_size) + telly.stream.stop() + # The background thread gets shut down by `Telly.__exit__`. From 4401bef5dd55e4cd3c792fd67c986db3eb7c76a5 Mon Sep 17 00:00:00 2001 From: Richard Bowman Date: Fri, 9 Oct 2026 01:28:31 +0100 Subject: [PATCH 07/10] Add a function to get the next frame + metadata. This is tested by a better-isolated test that mocks the server and uses a real MessageBroker. --- src/labthings_fastapi/outputs/mjpeg_stream.py | 16 +++-- tests/test_mjpeg_stream.py | 70 ++++++++++++++++--- 2 files changed, 73 insertions(+), 13 deletions(-) diff --git a/src/labthings_fastapi/outputs/mjpeg_stream.py b/src/labthings_fastapi/outputs/mjpeg_stream.py index 4318007a..7752499f 100644 --- a/src/labthings_fastapi/outputs/mjpeg_stream.py +++ b/src/labthings_fastapi/outputs/mjpeg_stream.py @@ -178,18 +178,24 @@ def stop(self) -> None: self._name, ) - async def grab_frame(self) -> bytes: + async def grab_frame_with_metadata(self) -> Frame: """Wait for the next frame, and return it. - This copies the frame for safety, so there is no need to release - or return the buffer. + This includes metadata such as the timestamp. - :return: The next JPEG frame, as a `bytes` object. + :return: The next JPEG frame, in a dataclass with metadata. """ message = await self._thing._thing_server_interface.message_broker.next_message( self._thing.name, self._name ) - payload = frame_payload(message) + return frame_payload(message) + + async def grab_frame(self) -> bytes: + """Wait for the next frame, and return it. + + :return: The next JPEG frame, as a `bytes` object. + """ + payload = await self.grab_frame_with_metadata() return payload.frame async def next_frame_size(self) -> int: diff --git a/tests/test_mjpeg_stream.py b/tests/test_mjpeg_stream.py index 5028e8da..94b31762 100644 --- a/tests/test_mjpeg_stream.py +++ b/tests/test_mjpeg_stream.py @@ -3,12 +3,22 @@ import time from datetime import datetime +import anyio import pytest +from anyio import from_thread, to_thread from PIL import Image import labthings_fastapi as lt -from labthings_fastapi.message_broker import Message -from labthings_fastapi.outputs.mjpeg_stream import Frame, frame_payload +from labthings_fastapi.message_broker import Message, MessageBroker +from labthings_fastapi.outputs.mjpeg_stream import Frame, MJPEGStream, frame_payload + + +def make_jpeg(colour: str) -> bytes: + """Return a 10x10 JPEG of the given colour.""" + image = Image.new("RGB", (10, 10), colour) + dest = io.BytesIO() + image.save(dest, "jpeg") + return dest.getvalue() class Telly(lt.Thing): @@ -37,12 +47,7 @@ def __exit__(self, exc_t, exc_v, exc_tb): def _make_images(self): """Stream a series of solid colours""" colours = ["#F00", "#0F0", "#00F"] - jpegs = [] - for c in colours: - image = Image.new("RGB", (10, 10), c) - dest = io.BytesIO() - image.save(dest, "jpeg") - jpegs.append(dest.getvalue()) + jpegs = [make_jpeg(c) for c in colours] if self.initial_delay > 0: time.sleep(self.initial_delay) @@ -150,6 +155,55 @@ def test_grab_and_shutdown(server: lt.ThingServer, telly: Telly): # The background thread gets shut down by `Telly.__exit__`. +@pytest.fixture +def mjpeg_stream(mocker): + """An MJPEGStream object with a mocked Thing.""" + thing = mocker.Mock() + broker = MessageBroker() + tsi = thing._thing_server_interface + tsi.message_broker = broker + tsi.start_async_task_soon.side_effect = from_thread.run + + def publish(message: Message) -> None: + from_thread.run(broker.publish, message) + + tsi.publish.side_effect = publish + thing.name = "thing" + yield MJPEGStream(thing, "stream") + + +@pytest.fixture +def jpeg_bytes(): + return make_jpeg("#F0F") + + +async def test_grab_frame(mjpeg_stream, jpeg_bytes): + """Test we can add a frame in a thread and grab its bytes.""" + # Listen for the next frame (grab_frame) then add a frame. + async with anyio.create_task_group() as tg: + handle = tg.start_soon(mjpeg_stream.grab_frame) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + assert handle.return_value is jpeg_bytes + + +async def test_grab_frame_with_metadata(mjpeg_stream, jpeg_bytes): + """Test we can add a frame in a thread and grab it.""" + # Listen for the next frame (grab_frame) then add a frame. + async with anyio.create_task_group() as tg: + handle = tg.start_soon(mjpeg_stream.grab_frame_with_metadata) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + assert handle.return_value.frame is jpeg_bytes + + +async def test_next_frame_size(mjpeg_stream, jpeg_bytes): + """Test we can add a frame in a thread and grab its size.""" + # Listen for the next frame (grab_frame) then add a frame. + async with anyio.create_task_group() as tg: + handle = tg.start_soon(mjpeg_stream.next_frame_size) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + assert handle.return_value == len(jpeg_bytes) + + def test_mjpeg_stream_http(server: lt.ThingServer, telly: Telly): """Verify the MJPEG stream works, and is shut down cleanly. From bf011e14ab949ace19ae1bd65b4aa245fd9fd5b8 Mon Sep 17 00:00:00 2001 From: Richard Bowman Date: Mon, 12 Oct 2026 22:17:16 +0100 Subject: [PATCH 08/10] Reinstate the async generator. The async generator actually works quite well with the new stream based structure, so I've resinstated it. I've taken the opportunity to rename it, and return Frame objects, not just bytes. The old name remains for backwards compatibility. This shrinks the code changes slightly, I think, and should also make it neater for client code as there's less need to create/use/type check streams directly. I've added explicit tests for the new functions in message broker. --- src/labthings_fastapi/message_broker.py | 27 +++- src/labthings_fastapi/outputs/mjpeg_stream.py | 131 ++++++++++-------- tests/test_message_broker.py | 102 ++++++++++++++ tests/test_mjpeg_stream.py | 98 ++++++++++--- 4 files changed, 284 insertions(+), 74 deletions(-) diff --git a/src/labthings_fastapi/message_broker.py b/src/labthings_fastapi/message_broker.py index 1621170a..d683e291 100644 --- a/src/labthings_fastapi/message_broker.py +++ b/src/labthings_fastapi/message_broker.py @@ -6,6 +6,7 @@ import logging import warnings +from collections.abc import AsyncGenerator from typing import Any, Literal from weakref import WeakSet @@ -83,6 +84,24 @@ async def subscribe( streams = affordances.setdefault(affordance, WeakSet()) streams.add(stream) + async def messages_for_affordance( + self, thing: str, affordance: str + ) -> AsyncGenerator[Message, None]: + """Yield messages from one affordance as an async generator. + + Note that there's no timeout: standard `anyio` commands may be used to do that. + + :param thing: The name of the `.Thing` being subscribed to. + :param affordance: The name of the affordance being subscribed to. + :yield: messages from the specified affordance. + """ + send, recv = anyio.create_memory_object_stream[Message](max_buffer_size=1) + await self.subscribe(thing, affordance, send) + async for message in recv: + yield message + # The stream isn't explicitly closed because it will now be discarded: the + # broker doesn't hold a strong reference. + async def next_message(self, thing: str, affordance: str) -> Message: """Get the next message from a particular affordance. @@ -91,10 +110,12 @@ async def next_message(self, thing: str, affordance: str) -> Message: :param thing: The name of the `.Thing` being subscribed to. :param affordance: The name of the affordance being subscribed to. :return: the next message from the specified affordance. + :raises GeneratorExit: if the stream closes before a message is received. """ - send, recv = anyio.create_memory_object_stream[Message](max_buffer_size=1) - await self.subscribe(thing, affordance, send) - return await recv.receive() + # Note: the for loop below will do at most one iteration. + async for message in self.messages_for_affordance(thing, affordance): + return message + raise GeneratorExit("The message stream closed before anything was received.") async def unsubscribe( self, thing: str, affordance: str, stream: MemoryObjectSendStream[Message] diff --git a/src/labthings_fastapi/outputs/mjpeg_stream.py b/src/labthings_fastapi/outputs/mjpeg_stream.py index 7752499f..78ede767 100644 --- a/src/labthings_fastapi/outputs/mjpeg_stream.py +++ b/src/labthings_fastapi/outputs/mjpeg_stream.py @@ -15,13 +15,11 @@ Optional, ) -import anyio -from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from fastapi import FastAPI from fastapi.responses import HTMLResponse, StreamingResponse from labthings_fastapi.base_descriptor import BaseDescriptor -from labthings_fastapi.message_broker import Message +from labthings_fastapi.message_broker import Message, MessageBroker from labthings_fastapi.thing import Thing @@ -37,6 +35,7 @@ class Frame: frame: bytes """The frame as a `bytes` object, which is a JPEG image for an MJPEG stream.""" + timestamp: datetime """The time the frame was captured. If the capture time is unknown then the timestamp is @@ -46,22 +45,18 @@ class Frame: """The index of the frame within the stream.""" -def frame_payload(message: Message) -> Frame: - """Extract a Frame object from a Message. - - This checks that the message is of "stream" type, and checks the - type of its payload. +def check_jpeg_bytes(frame: bytes) -> bool: + """Check a JPEG file starts and ends with the right magic bytes. - :param message: the message containing the Frame. - :return: the Frame object. - :raises TypeError: if the message isn't a stream message containing - a `Frame` object. + :param frame: the data, which should be in JPEG format. + :return: `True` if the data is a valid JPEG (i.e. has valid start/end bytes). """ - if message.message_type != "stream": - raise TypeError("The message {message} was not part of a stream.") - if not isinstance(message.payload, Frame): - raise TypeError("The message {message} didn't contain a Frame.") - return message.payload + return ( + frame[0] == 0xFF + and frame[1] == 0xD8 + and frame[-2] == 0xFF + and frame[-1] == 0xD9 + ) class MJPEGStreamResponse(StreamingResponse): @@ -77,10 +72,7 @@ class MJPEGStreamResponse(StreamingResponse): """The media_type used to describe the endpoint in FastAPI.""" def __init__( - self, - send: MemoryObjectSendStream[Message], - recv: MemoryObjectReceiveStream[Message], - status_code: int = 200, + self, gen: AsyncGenerator[bytes, None], status_code: int = 200 ) -> None: """Set up StreamingResponse that streams an MJPEG stream. @@ -89,20 +81,18 @@ def __init__( types that mark it as an MJPEG stream. This is sufficient to enable it to work in an `img` tag, with the `src` set to the MJPEG stream's endpoint. - It expects to get a stream that receives `Message` objects with `Frame` - payloads. Both the send and receive streams are retained, because the send - stream is only weakly referenced by the message broker. + It expects an async generator that supplies individual JPEGs to be streamed, + such as the one provided by `.MJPEGStream`. NB the ``status_code`` argument is used by FastAPI to set the status code of the response in OpenAPI. - :param send: the send stream subscribed to the MJPEG stream. - :param recv: the receive stream subscribed to the MJPEG stream. + :param gen: an async generator, yielding `bytes` objects each of which is + one image, in JPEG format. :param status_code: The status code associated with the response, by default a 200 code is returned. """ - self._send_stream = send - self._receive_stream = recv + self._frame_async_generator = gen StreamingResponse.__init__( self, self.mjpeg_async_generator(), @@ -119,10 +109,9 @@ async def mjpeg_async_generator(self) -> AsyncGenerator[bytes, None]: :yield: JPEG frames, each with a ``--frame`` marker prepended. """ - async for message in self._receive_stream: - frame = frame_payload(message) + async for frame in self._frame_async_generator: yield b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" - yield frame.frame + yield frame yield b"\r\n" @@ -161,6 +150,11 @@ def __init__(self, thing: Thing, name: str) -> None: self._name = name self.reset() + @property + def _broker(self) -> MessageBroker: + """The message broker associated with our host Thing.""" + return self._thing._thing_server_interface.message_broker + def reset(self) -> None: """Reset the frame index.""" with self._lock: @@ -170,25 +164,62 @@ def stop(self) -> None: """Stop the stream. Stop the stream and cause all clients to disconnect. + This uses the thing server interface to allow it to be called from + a thread, for convenience. + + To close all streams from the event loop, use the message broker directly. """ - tsi = self._thing._thing_server_interface - tsi.call_async_task( - tsi.message_broker.close_streams_for_affordance, + self._thing._thing_server_interface.call_async_task( + self._broker.close_streams_for_affordance, self._thing.name, self._name, ) + async def frames(self) -> AsyncGenerator[Frame, None]: + """Yield Frame objects (data and metadata). + + This generator will return frames from the MJPEG stream along with timestamps, + using the `Frame` dataclass. + + Note that this will wait for a new frame each time. There is no + guarantee that we won't skip frames. There's no timeout, if the MJPEG stream + doesn't get new frames this may block indefinitely. + + :yield: the frames in sequence. + """ + async for message in self._broker.messages_for_affordance( + self._thing.name, self._name + ): + if isinstance(message.payload, Frame): + # We simply ignore messages that don't contain Frame objects. + # We don't expect to get any payloads that are not Frame objects. + yield message.payload + + async def frame_async_generator(self) -> AsyncGenerator[bytes, None]: + """Yield frames as bytes objects. + + This generator will return frames from the MJPEG stream. + + See `frames` for how this works. This function wraps that one, returning + only the `bytes` object rather than the full `Frame` dataclass which also + includes metadata. + + :yield: the frames in sequence, as a `bytes` object containing + JPEG data. + """ + async for frame in self.frames(): + yield frame.frame + async def grab_frame_with_metadata(self) -> Frame: """Wait for the next frame, and return it. This includes metadata such as the timestamp. + This method calls `frames` and returns the first frame that's yielded. + :return: The next JPEG frame, in a dataclass with metadata. """ - message = await self._thing._thing_server_interface.message_broker.next_message( - self._thing.name, self._name - ) - return frame_payload(message) + return await anext(self.frames()) async def grab_frame(self) -> bytes: """Wait for the next frame, and return it. @@ -218,11 +249,7 @@ async def mjpeg_stream_response(self) -> MJPEGStreamResponse: :return: a streaming response in MJPEG format. """ - send, recv = anyio.create_memory_object_stream[Message](max_buffer_size=1) - await self._thing._thing_server_interface.message_broker.subscribe( - self._thing.name, self._name, send - ) - return MJPEGStreamResponse(send, recv) + return MJPEGStreamResponse(self.frame_async_generator()) def add_frame(self, frame: bytes, timestamp: Optional[datetime] = None) -> None: """Add a JPEG to the MJPEG stream. @@ -240,20 +267,16 @@ def add_frame(self, frame: bytes, timestamp: Optional[datetime] = None) -> None: :raise ValueError: if the supplied frame does not start with the JPEG start bytes and end with the end bytes. """ - if not ( - frame[0] == 0xFF - and frame[1] == 0xD8 - and frame[-2] == 0xFF - and frame[-1] == 0xD9 - ): + if not check_jpeg_bytes(frame): raise ValueError("Invalid JPEG") with self._lock: self.last_frame_i += 1 - payload = Frame( - frame=frame, - timestamp=timestamp if timestamp is not None else datetime.now(), - index=self.last_frame_i, - ) + frame_index = self.last_frame_i + payload = Frame( + frame=frame, + timestamp=timestamp if timestamp is not None else datetime.now(), + index=frame_index, + ) self._thing._thing_server_interface.publish( Message(self._thing.name, self._name, "stream", payload) ) diff --git a/tests/test_message_broker.py b/tests/test_message_broker.py index dbfea33d..ea29abfe 100644 --- a/tests/test_message_broker.py +++ b/tests/test_message_broker.py @@ -23,6 +23,7 @@ class Unjsonable: ("test_thing", "prop", "property", Unjsonable()), ("test_thing", "do_it", "action", None), ("test_thing", "notify", "action", {"key": "value"}), + ("test_thing", "mjpeg", "stream", b"whatever"), ], ) def test_message_valid(message): @@ -91,6 +92,23 @@ async def test_subscribe_unsubscribe(): await broker.unsubscribe("whatever", 42, send_stream) # type: ignore +async def test_unsubscribe_on_deletion(): + """Check that streams auto-unsubscribe when they are finalised.""" + broker = MessageBroker() + assert broker._subscriptions == {} + + # subscribe to a stream + send_stream, receive_stream = anyio.create_memory_object_stream[Message]() + await broker.subscribe("thing_name", "prop", send_stream) + assert send_stream in broker._subscriptions["thing_name"]["prop"] + + # delete the streams + del send_stream + del receive_stream + + assert len(broker._subscriptions["thing_name"]["prop"]) == 0 + + async def append_messages( stream: ObjectReceiveStream[Message], dest: list[Message], @@ -260,3 +278,87 @@ async def test_weakset_garbage_collection(): assert len(iterated) == 1 # We only iterated once, but there wasn't an error. # If we complete the test, it confirms we don't need to worry about streams being # finalised during iteration. + + +async def collect_messages(generator, limit: int | None = None) -> list[Message]: + """Collect messages from a generator into a list.""" + received = [] + async for message in generator: + received.append(message) + if limit and len(received) >= limit: + break + return received + + +async def test_messages_for_affordance(): + """Test the shorthand to iterate over messages for an affordance works. + + This also tests close_streams_for_affordance(). + + I don't love that this test needs `await anyio.sleep(0.001)` in + order to work, but because the message broker is deliberately + lossy (to prevent back-pressure from a slow consumer) this is + the least horrible way I've found to stop messages being dropped. + """ + broker = MessageBroker() + message = Message("thing", "prop", "property", None) + + async with anyio.create_task_group() as tg: + handle = tg.start_soon( + collect_messages, + broker.messages_for_affordance("thing", "prop"), + 3, + ) + for _ in range(3): + await anyio.sleep(0.001) # allow other task to run + await broker.publish(message) + assert handle.return_value == [message] * 3 + # As the listener stopped, its stream should no longer be subscribed + assert len(broker._subscriptions["thing"]["prop"]) == 0 + + # Do the same again, but this time stop the stream from the broker + async with anyio.create_task_group() as tg: + handle = tg.start_soon( + collect_messages, + broker.messages_for_affordance("thing", "prop"), + # Note: no "limit" value + ) + for _ in range(3): + await anyio.sleep(0.001) # allow other task to run + await broker.publish(message) + # The background task will run until its stream is closed. + # the next line should close the stream: if the test doesn't + # hang, it's worked. + await broker.close_streams_for_affordance("thing", "prop") + # Calling this twice shouldn't error. + await broker.close_streams_for_affordance("thing", "prop") + assert handle.return_value == [message] * 3 + # As the listener stopped, its stream should no longer be subscribed + assert len(broker._subscriptions["thing"]["prop"]) == 0 + + +async def test_next_message(): + """Test the shorthand to get the next message for an affordance. + + This also tests close_streams_for_affordance(). + """ + broker = MessageBroker() + message = Message("thing", "prop", "property", None) + + # Listen for three messages (and then stop listening) + async with anyio.create_task_group() as tg: + handle = tg.start_soon(broker.next_message, "thing", "prop") + await anyio.sleep(0.001) # allow other task to run + await broker.publish(message) + assert handle.return_value == message + # The stream should no longer be subscribed + assert len(broker._subscriptions["thing"]["prop"]) == 0 + + # Test for error if the stream closes without a message + with pytest.RaisesGroup(GeneratorExit, flatten_subgroups=True): + async with anyio.create_task_group() as tg: + handle = tg.start_soon(broker.next_message, "thing", "prop") + await anyio.sleep(0.001) # allow other task to run + await broker.close_streams_for_affordance("thing", "prop") + # The stream should no longer be subscribed + assert len(broker._subscriptions["thing"]["prop"]) == 0 diff --git a/tests/test_mjpeg_stream.py b/tests/test_mjpeg_stream.py index 94b31762..8fe94b8f 100644 --- a/tests/test_mjpeg_stream.py +++ b/tests/test_mjpeg_stream.py @@ -1,7 +1,6 @@ import io import threading import time -from datetime import datetime import anyio import pytest @@ -10,7 +9,11 @@ import labthings_fastapi as lt from labthings_fastapi.message_broker import Message, MessageBroker -from labthings_fastapi.outputs.mjpeg_stream import Frame, MJPEGStream, frame_payload +from labthings_fastapi.outputs.mjpeg_stream import ( + Frame, + MJPEGStream, + check_jpeg_bytes, +) def make_jpeg(colour: str) -> bytes: @@ -87,19 +90,10 @@ def telly(server): return telly -def test_frame_payload(): - """Check that a frame is returned, or we get an error.""" - frame = Frame(b"payload", datetime.now(), 1) - good_message = Message("thing", "affordance", "stream", frame) - assert frame_payload(good_message) is frame - - bad_message = Message("thing", "affordance", "action", frame) - with pytest.raises(TypeError): - frame_payload(bad_message) - - bad_message2 = Message("thing", "affordance", "stream", b"payload") - with pytest.raises(TypeError): - frame_payload(bad_message2) +def test_check_jpeg_bytes(jpeg_bytes): + """Check the check_jpeg_bytes function works correctly""" + assert check_jpeg_bytes(jpeg_bytes) is True + assert check_jpeg_bytes(b"00000000") is False def test_grab_and_shutdown(server: lt.ThingServer, telly: Telly): @@ -156,10 +150,15 @@ def test_grab_and_shutdown(server: lt.ThingServer, telly: Telly): @pytest.fixture -def mjpeg_stream(mocker): +def broker() -> MessageBroker: + """A MessageBroker object.""" + return MessageBroker() + + +@pytest.fixture +def mjpeg_stream(mocker, broker): """An MJPEGStream object with a mocked Thing.""" thing = mocker.Mock() - broker = MessageBroker() tsi = thing._thing_server_interface tsi.message_broker = broker tsi.start_async_task_soon.side_effect = from_thread.run @@ -204,6 +203,71 @@ async def test_next_frame_size(mjpeg_stream, jpeg_bytes): assert handle.return_value == len(jpeg_bytes) +async def grab_three_frames_with_generator(stream: MJPEGStream) -> list[Frame]: + """Grab three frames and return them in a list.""" + frames = [] + async for frame in stream.frames(): + frames.append(frame) + if len(frames) >= 3: + break + return frames + + +async def test_frames(mjpeg_stream, jpeg_bytes, broker): + """Check we can get a few frames from the generator. + + This test also makes sure we don't end up with dangling + streams that are no longer wanted. + """ + async with anyio.create_task_group() as tg: + handle = tg.start_soon(grab_three_frames_with_generator, mjpeg_stream) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + frames = handle.return_value + assert len(frames) == 3 + for i, f in enumerate(frames): + assert isinstance(f, Frame) + assert f.frame is jpeg_bytes + assert f.index == i + # Emit another frame, to ensure it deletes the stream + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + # Check there's no streams left connected + assert len(broker._subscriptions["thing"]["stream"]) == 0 + + +async def grab_three_frames_from_old_generator(stream: MJPEGStream) -> list[bytes]: + """Grab three frames from the old-style generator""" + frames = [] + async for frame in stream.frame_async_generator(): + frames.append(frame) + if len(frames) >= 3: + break + return frames + + +async def test_frame_async_generator(mjpeg_stream, jpeg_bytes, broker): + """Check we can get a few frames from the old generator. + + This test also makes sure we don't end up with dangling + streams that are no longer wanted. + """ + async with anyio.create_task_group() as tg: + handle = tg.start_soon(grab_three_frames_from_old_generator, mjpeg_stream) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + frames = handle.return_value + assert len(frames) == 3 + for f in frames: + assert isinstance(f, bytes) + assert f is jpeg_bytes + # Emit another frame, to ensure it deletes the stream + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + # Check there's no streams left connected + assert len(broker._subscriptions["thing"]["stream"]) == 0 + + def test_mjpeg_stream_http(server: lt.ThingServer, telly: Telly): """Verify the MJPEG stream works, and is shut down cleanly. From 86c5d5ace45e224fb8ec4b55272dc29ee7eba6b5 Mon Sep 17 00:00:00 2001 From: Richard Bowman Date: Tue, 13 Oct 2026 22:53:11 +0100 Subject: [PATCH 09/10] Don't clutter the logs when a client stalls. The old MJPEG stream silently dropped frames. We were logging every dropped frame at WARNING. I've downgraded this to DEBUG for streams. We may choose to turn it off entirely to replicate the old behaviour. I have also altered the `repr` of `Frame` objects, to avoid dumping JPEG data to console/log accidentally. --- src/labthings_fastapi/message_broker.py | 9 +++- src/labthings_fastapi/outputs/mjpeg_stream.py | 10 +++++ tests/test_message_broker.py | 13 ++++++ tests/test_mjpeg_stream.py | 42 +++++++++++++++++++ 4 files changed, 72 insertions(+), 2 deletions(-) diff --git a/src/labthings_fastapi/message_broker.py b/src/labthings_fastapi/message_broker.py index d683e291..fe944049 100644 --- a/src/labthings_fastapi/message_broker.py +++ b/src/labthings_fastapi/message_broker.py @@ -163,8 +163,13 @@ async def publish(self, message: Message) -> None: subscriptions_to_remove.add(stream) except anyio.WouldBlock: msg = f"Could not pass notification to {stream} as it was full." - LOGGER.warning(msg) - warnings.warn(MessageDroppedWarning(msg), stacklevel=1) + if message.message_type == "stream": + # It's expected that streams might drop messages, so we + # only log this at debug level. + LOGGER.debug(msg) + else: + LOGGER.warning(msg) + warnings.warn(MessageDroppedWarning(msg), stacklevel=1) for stream in subscriptions_to_remove: # discard rather than remove, so that if the stream has been finalised # since it was closed, we don't get an error. diff --git a/src/labthings_fastapi/outputs/mjpeg_stream.py b/src/labthings_fastapi/outputs/mjpeg_stream.py index 78ede767..ebd8cf08 100644 --- a/src/labthings_fastapi/outputs/mjpeg_stream.py +++ b/src/labthings_fastapi/outputs/mjpeg_stream.py @@ -44,6 +44,16 @@ class Frame: index: int """The index of the frame within the stream.""" + def __repr__(self) -> str: + """Represent the object as a string. + + :return: a representation of the object as a string. + """ + return ( + f"Frame(index={self.index}, timestamp={self.timestamp}, " + f"frame=<{len(self.frame)} bytes>)" + ) + def check_jpeg_bytes(frame: bytes) -> bool: """Check a JPEG file starts and ends with the right magic bytes. diff --git a/tests/test_message_broker.py b/tests/test_message_broker.py index ea29abfe..68a4d404 100644 --- a/tests/test_message_broker.py +++ b/tests/test_message_broker.py @@ -246,6 +246,19 @@ async def test_sending_to_full_stream(caplog): with pytest.warns(MessageDroppedWarning): await broker.publish(message) assert len(caplog.records) == 1 + assert caplog.records[0].levelno == logging.WARNING + msg = caplog.records[0].getMessage() + assert msg.startswith("Could not pass notification to") + assert msg.endswith("as it was full.") + caplog.clear() + + # Send a fourth message marked as being part of a stream. + # This should only log at DEBUG level. + stream_msg = Message("thing", "prop", "stream", None) + caplog.set_level(logging.DEBUG) + await broker.publish(stream_msg) + assert len(caplog.records) == 1 + assert caplog.records[0].levelno == logging.DEBUG msg = caplog.records[0].getMessage() assert msg.startswith("Could not pass notification to") assert msg.endswith("as it was full.") diff --git a/tests/test_mjpeg_stream.py b/tests/test_mjpeg_stream.py index 8fe94b8f..1fc93039 100644 --- a/tests/test_mjpeg_stream.py +++ b/tests/test_mjpeg_stream.py @@ -1,6 +1,8 @@ import io +import logging import threading import time +from datetime import datetime import anyio import pytest @@ -96,6 +98,20 @@ def test_check_jpeg_bytes(jpeg_bytes): assert check_jpeg_bytes(b"00000000") is False +def test_frame_repr(jpeg_bytes): + """Make sure Frame objects don't print JPEG data.""" + frame = Frame(jpeg_bytes, datetime(2005, 7, 14, 12, 30), 1) + # I'll be explicit here: I want to make sure there's no dump + # of the bytes. + assert 'b"' not in repr(frame) + # I can also just test the known repr: + assert repr(frame) == ( + "Frame(index=1, timestamp=2005-07-14 12:30:00, " + f"frame=<{len(jpeg_bytes)} bytes>)" + ) + assert str(frame) == repr(frame) + + def test_grab_and_shutdown(server: lt.ThingServer, telly: Telly): """Check we can grab frames, and shut down cleanly. @@ -268,6 +284,32 @@ async def test_frame_async_generator(mjpeg_stream, jpeg_bytes, broker): assert len(broker._subscriptions["thing"]["stream"]) == 0 +async def test_dropped_frames(mjpeg_stream, jpeg_bytes, broker, caplog): + """Test what happens when a subscribed stream is full. + + Note that the MessageBroker functionality is also tested in + test_message_broker.py::test_sending_to_full_stream. + """ + send, recv = anyio.create_memory_object_stream[Message](1) + await broker.subscribe("thing", "stream", send) + caplog.set_level(logging.DEBUG) + caplog.clear() + # The first message sits in the buffer, the second gets dropped + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + print(caplog.records) + assert len(caplog.records) == 1 + assert caplog.records[0].levelno == logging.DEBUG + msg = caplog.records[0].getMessage() + assert msg.startswith("Could not pass notification to") + assert msg.endswith("as it was full.") + # We should summarise the length of the JPEG, not include + # the JPEG bytes. + assert f"<{len(jpeg_bytes)} bytes>" in msg + # We want to make sure we're not dumping JPEGs to the log + assert 'b"' not in msg + + def test_mjpeg_stream_http(server: lt.ThingServer, telly: Telly): """Verify the MJPEG stream works, and is shut down cleanly. From e4959f37339d89477a95a087cc647d2520b0d505 Mon Sep 17 00:00:00 2001 From: Richard Bowman Date: Mon, 26 Oct 2026 23:20:41 +0000 Subject: [PATCH 10/10] Disable buffering for mjpegstream There are now no buffers used when streaming MJPEG frames, exactly replicating the previous behaviour. This is enabled by a better understanding (by me) of the anyio event loop. Specifically, not all `await` statements yield to the event loop, and so it's necessary to manually yield to the event loop during the test. Now that I've added this, I should be able to write better async tests for a few things. --- src/labthings_fastapi/message_broker.py | 11 ++- src/labthings_fastapi/outputs/mjpeg_stream.py | 6 +- tests/test_message_broker.py | 84 +++++++++++-------- 3 files changed, 61 insertions(+), 40 deletions(-) diff --git a/src/labthings_fastapi/message_broker.py b/src/labthings_fastapi/message_broker.py index fe944049..17eba015 100644 --- a/src/labthings_fastapi/message_broker.py +++ b/src/labthings_fastapi/message_broker.py @@ -85,7 +85,10 @@ async def subscribe( streams.add(stream) async def messages_for_affordance( - self, thing: str, affordance: str + self, + thing: str, + affordance: str, + max_buffer_size: int = 1, ) -> AsyncGenerator[Message, None]: """Yield messages from one affordance as an async generator. @@ -93,9 +96,13 @@ async def messages_for_affordance( :param thing: The name of the `.Thing` being subscribed to. :param affordance: The name of the affordance being subscribed to. + :param max_buffer_size: The buffer size of the underlying streams. + The default will buffer up to one message. :yield: messages from the specified affordance. """ - send, recv = anyio.create_memory_object_stream[Message](max_buffer_size=1) + send, recv = anyio.create_memory_object_stream[Message]( + max_buffer_size=max_buffer_size, + ) await self.subscribe(thing, affordance, send) async for message in recv: yield message diff --git a/src/labthings_fastapi/outputs/mjpeg_stream.py b/src/labthings_fastapi/outputs/mjpeg_stream.py index ebd8cf08..3a0f64b1 100644 --- a/src/labthings_fastapi/outputs/mjpeg_stream.py +++ b/src/labthings_fastapi/outputs/mjpeg_stream.py @@ -193,12 +193,14 @@ async def frames(self) -> AsyncGenerator[Frame, None]: Note that this will wait for a new frame each time. There is no guarantee that we won't skip frames. There's no timeout, if the MJPEG stream - doesn't get new frames this may block indefinitely. + doesn't get new frames this may block indefinitely. Frames are never buffered. :yield: the frames in sequence. """ async for message in self._broker.messages_for_affordance( - self._thing.name, self._name + self._thing.name, + self._name, + max_buffer_size=0, ): if isinstance(message.payload, Frame): # We simply ignore messages that don't contain Frame objects. diff --git a/tests/test_message_broker.py b/tests/test_message_broker.py index 68a4d404..e3de8925 100644 --- a/tests/test_message_broker.py +++ b/tests/test_message_broker.py @@ -4,6 +4,7 @@ from weakref import WeakSet import anyio +import anyio.lowlevel import pytest from anyio.abc import ObjectReceiveStream from pydantic import ValidationError @@ -204,7 +205,6 @@ async def test_sending_to_closed_streams(caplog, action): assert send_stream in broker._subscriptions["thing"]["prop"] message = Message("thing", "prop", "property", None) - print("sending first message") # Check we can send and receive a message await broker.publish(message) received = await receive_stream.receive() @@ -218,7 +218,6 @@ async def test_sending_to_closed_streams(caplog, action): else: del send_stream # streams are unsubscribed when they are finalised assert len(caplog.records) == 0 - print("sending second message") await broker.publish(message) assert len(caplog.records) == 0 # Shouldn't be any warnings in the log # Check we've been unsubscribed @@ -303,49 +302,62 @@ async def collect_messages(generator, limit: int | None = None) -> list[Message] return received -async def test_messages_for_affordance(): +@pytest.mark.parametrize("buffer_size", (0, 1, 3)) +@pytest.mark.parametrize("close_stream", ("subscriber", "broker")) +async def test_messages_for_affordance(buffer_size, close_stream): """Test the shorthand to iterate over messages for an affordance works. This also tests close_streams_for_affordance(). - I don't love that this test needs `await anyio.sleep(0.001)` in - order to work, but because the message broker is deliberately - lossy (to prevent back-pressure from a slow consumer) this is - the least horrible way I've found to stop messages being dropped. + Note that there are various uses of `await anyio.lowlevel.checkpoint()` in + this test. It's necessary to add these in order to yield to the event + loop so that the task that listens is able to run. + + For buffer sizes from 0 to 3, we send (b+1) messages at which point the + buffer is full. We then send another message that is dropped. + + After sending the messages, we then await (b+2) checkpoints, to allow + the task to clear the buffer, and try again. """ broker = MessageBroker() - message = Message("thing", "prop", "property", None) - - async with anyio.create_task_group() as tg: - handle = tg.start_soon( - collect_messages, - broker.messages_for_affordance("thing", "prop"), - 3, - ) - for _ in range(3): - await anyio.sleep(0.001) # allow other task to run - await broker.publish(message) - assert handle.return_value == [message] * 3 - # As the listener stopped, its stream should no longer be subscribed - assert len(broker._subscriptions["thing"]["prop"]) == 0 + messages = [Message("thing", "prop", "property", i) for i in range(buffer_size + 1)] + dropped_message = Message("thing", "prop", "property", "dropped") - # Do the same again, but this time stop the stream from the broker + # We subscribe to the broker in a background task, and send messages. + # With no buffer, we should process one message each time we yield to the + # event loop. With a finite buffer, we process (b+1) messages each time. + # We make b+2 messages, and check we get the first (b+1) of them. async with anyio.create_task_group() as tg: + # Start a background task to listen for messages, using a pair of streams + # with the specified buffer size. handle = tg.start_soon( collect_messages, - broker.messages_for_affordance("thing", "prop"), - # Note: no "limit" value + broker.messages_for_affordance("thing", "prop", buffer_size), + 2 * (buffer_size + 1) if close_stream == "subscriber" else None, ) - for _ in range(3): - await anyio.sleep(0.001) # allow other task to run - await broker.publish(message) - # The background task will run until its stream is closed. - # the next line should close the stream: if the test doesn't - # hang, it's worked. - await broker.close_streams_for_affordance("thing", "prop") - # Calling this twice shouldn't error. - await broker.close_streams_for_affordance("thing", "prop") - assert handle.return_value == [message] * 3 + # Allow the background task to run (subscribe and start listening) + await anyio.lowlevel.checkpoint() + await anyio.lowlevel.checkpoint() + # We loop around twice, sending all the messages + for _ in range(2): + for message in messages: + await broker.publish(message) + # After sending these, the buffer should be full + with pytest.warns(MessageDroppedWarning): + await broker.publish(dropped_message) + # Yield to the event loop, so the task can clear its buffer. + for _ in range(buffer_size + 2): + await anyio.lowlevel.checkpoint() + if close_stream == "broker": + # The background task will run until its stream is closed. + # the next line should close the stream: if the test doesn't + # hang, it's worked. + await broker.close_streams_for_affordance("thing", "prop") + # Calling this twice shouldn't error. + await broker.close_streams_for_affordance("thing", "prop") + # Each time, we should have dropped the last message (because we send two + # more than the buffer size). + assert handle.return_value == messages * 2 # As the listener stopped, its stream should no longer be subscribed assert len(broker._subscriptions["thing"]["prop"]) == 0 @@ -361,7 +373,7 @@ async def test_next_message(): # Listen for three messages (and then stop listening) async with anyio.create_task_group() as tg: handle = tg.start_soon(broker.next_message, "thing", "prop") - await anyio.sleep(0.001) # allow other task to run + await anyio.lowlevel.checkpoint() await broker.publish(message) assert handle.return_value == message # The stream should no longer be subscribed @@ -371,7 +383,7 @@ async def test_next_message(): with pytest.RaisesGroup(GeneratorExit, flatten_subgroups=True): async with anyio.create_task_group() as tg: handle = tg.start_soon(broker.next_message, "thing", "prop") - await anyio.sleep(0.001) # allow other task to run + await anyio.lowlevel.checkpoint() await broker.close_streams_for_affordance("thing", "prop") # The stream should no longer be subscribed assert len(broker._subscriptions["thing"]["prop"]) == 0