From d9fd54696a7eb7bd824a17557ef6eb96eba6a5e1 Mon Sep 17 00:00:00 2001
From: Julian Stirling <julian@julianstirling.co.uk>
Date: Tue, 8 Jul 2025 10:45:45 +0100
Subject: [PATCH] Monkey patch Uvicorn's handle_exit function to stop streams

Before the uvicorn handle_exit function is called kill the camera
streams.
---
 .../server/__init__.py                        | 44 ++++++++++++++++++-
 .../things/camera/__init__.py                 | 15 +++++++
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/src/openflexure_microscope_server/server/__init__.py b/src/openflexure_microscope_server/server/__init__.py
index 485ea0c7..5a8412c6 100644
--- a/src/openflexure_microscope_server/server/__init__.py
+++ b/src/openflexure_microscope_server/server/__init__.py
@@ -1,15 +1,42 @@
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, Callable
 from copy import copy
 
 import labthings_fastapi as lt
 import uvicorn
+import logging
+from uvicorn.main import Server
 from .serve_static_files import add_static_files
 from .legacy_api import add_v2_endpoints
 from ..logging import configure_logging, retrieve_log, retrieve_log_from_file
 
 
+def set_shutdown_function(shutdown_function: Callable[[], None]):
+    """
+    Ensure a function is called before the shutdown
+
+    This monkey patches the Uvicorn Server's handle_exit. This is needed because
+    the uvicorn `lifecycle` events and FastAPI `shutdown` events only fire once
+    background tasks have completed.
+
+    Without this the system exits cleanly only if no client is receiving a
+    StreamingResponse. This patch is used to stop the async generators that
+    send streaming responses.
+
+    :param shutdown_function: A callable with no arguments or outputs. This
+    should stop any async generators that may be sending to streaming responses.
+    """
+
+    original_handler = Server.handle_exit
+
+    def handle_exit(*args, **kwargs):
+        shutdown_function()
+        original_handler(*args, **kwargs)
+
+    Server.handle_exit = handle_exit
+
+
 def customise_server(
     server: lt.ThingServer, log_folder: str, scans_folder: Optional[str]
 ):
@@ -57,6 +84,21 @@ def serve_from_cli(argv: Optional[list[str]] = None):
         scans_folder = _get_scans_dir(config)
         server = lt.cli.server_from_config(config)
         customise_server(server, log_folder, scans_folder)
+
+        def shutdown_call():
+            try:
+                # Kill any mjpeg streams so that StreamingResponses close.
+                server.things["/camera/"].kill_mjpeg_streams()
+            except BaseException as e:
+                # Catch anything and log as it is essential that this
+                # function cannot raise an unhandled exception or Uvicorn
+                # will never get a shutdown signal.
+                logging.error(e, exc_info=True)
+
+        # Monkey patch uvicorn's exit handling to stop the MJPEG Streams
+        # before waiting for background tasks to complete.
+        set_shutdown_function(shutdown_call)
+
         uvicorn.run(
             server.app,
             host=args.host,
diff --git a/src/openflexure_microscope_server/things/camera/__init__.py b/src/openflexure_microscope_server/things/camera/__init__.py
index 74531053..3ac1a3b1 100644
--- a/src/openflexure_microscope_server/things/camera/__init__.py
+++ b/src/openflexure_microscope_server/things/camera/__init__.py
@@ -169,6 +169,21 @@ class BaseCamera(lt.Thing):
             "CameraThings must define their own start_streaming method"
         )
 
+    def kill_mjpeg_streams(self):
+        """
+        Kill the streams now as the server is shutting down.
+
+        This is called when uvicorn gets the a shutdown signal. As this is called from
+        the event loop it cannot interact with the our ThingProperties or run
+        `self.mjpeg_stream.stop()` as the portal cannot be called from this loop.
+
+        Instead we just set the `_streaming` value to False. This stops the async frame
+        generator when the next frame notifies.
+        """
+        if self.stream_active:
+            self.mjpeg_stream._streaming = False
+            self.lores_mjpeg_stream._streaming = False
+
     @lt.thing_property
     def stream_active(self) -> bool:
         "Whether the MJPEG stream is active"