Allow MJPEGStream to stop cleanly

It's now possible to tell MJPEGStream that the stream has stopped. This terminates the streaming
response, and gets around the lack of streaming
support in starlette's test client.

With a bit of thought, this could potentially fix the longstanding issue that MJPEG Streams
prevent a labthings server from restarting.

Author: rwb27

src/labthings_fastapi/outputs/mjpeg_stream.py

@@ -107,10 +107,11 @@ def reset(self, ringbuffer_size: Optional[int] = None):
             ]
             self.last_frame_i = -1
 
-    def stop(self):
+    def stop(self, portal: BlockingPortal):
         """Stop the stream"""
         with self._lock:
             self._streaming = False
+            portal.start_task_soon(self.notify_stream_stopped)
 
     async def ringbuffer_entry(self, i: int) -> RingbufferEntry:
         """Return the ith frame acquired by the camera
@@ -139,9 +140,13 @@ async def buffer_for_reading(self, i: int) -> AsyncIterator[bytes]:
         yield entry.frame
 
     async def next_frame(self) -> int:
-        """Wait for the next frame, and return its index"""
+        """Wait for the next frame, and return its index
+
+        :raises 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
 
     async def grab_frame(self) -> bytes:
@@ -170,6 +175,8 @@ async def frame_async_generator(self) -> AsyncGenerator[bytes, None]:
                 i = await self.next_frame()
                 async with self.buffer_for_reading(i) as frame:
                     yield frame
+            except StopAsyncIteration:
+                break
             except Exception as e:
                 logging.error(f"Error in stream: {e}, stream stopped")
                 return
@@ -178,7 +185,7 @@ async def mjpeg_stream_response(self) -> MJPEGStreamResponse:
         """Return a StreamingResponse that streams an MJPEG stream"""
         return MJPEGStreamResponse(self.frame_async_generator())
 
-    def add_frame(self, frame: bytes, portal: BlockingPortal):
+    def add_frame(self, frame: bytes, portal: BlockingPortal) -> None:
         """Return the next buffer in the ringbuffer to write to
 
         :param frame: The frame to add
@@ -196,12 +203,18 @@ def add_frame(self, frame: bytes, portal: BlockingPortal):
             entry.index = self.last_frame_i + 1
             portal.start_task_soon(self.notify_new_frame, entry.index)
 
-    async def notify_new_frame(self, i):
+    async def notify_new_frame(self, i: int) -> None:
         """Notify any waiting tasks that a new frame is available"""
         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."""
+        assert self._streaming is False
+        async with self.condition:
+            self.condition.notify_all()
+
 
 class MJPEGStreamDescriptor:
     """A descriptor that returns a MJPEGStream object when accessed

tests/test_mjpeg_stream.py

@@ -32,34 +32,36 @@ def _make_images(self):
             image.save(dest, "jpeg")
             jpegs.append(dest.getvalue())
 
-        i = -1
-        start_time = time.time()
-        while self._streaming:
-            i = (i + 1) % len(jpegs)
-            print(f"sending frame {i}")
+        i = 0
+        while self._streaming and i < len(jpegs):
             self.stream.add_frame(jpegs[i], self._labthings_blocking_portal)
             time.sleep(1 / self.framerate)
-
-            if time.time() - start_time > 10:
-                break
-        print("stopped sending frames")
+            i = i + 1
+        self.stream.stop(self._labthings_blocking_portal)
         self._streaming = False
 
 
 def test_mjpeg_stream():
+    """Verify the MJPEG stream contains at least one frame marker.
+
+    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.
+    """
     server = lt.ThingServer()
     telly = Telly()
     server.add_thing(telly, "telly")
     with TestClient(server.app) as client:
-        with client.stream("GET", "/telly/stream", timeout=0.1) as stream:
+        with client.stream("GET", "/telly/stream") as stream:
             stream.raise_for_status()
             received = 0
             for b in stream.iter_bytes():
                 received += 1
-                print(f"Got packet {received}")
                 assert b.startswith(b"--frame")
-                if received > 5:
-                    break
 
 
 if __name__ == "__main__":