fix(SSE): better handle reconnections and keep alive

requirements.txt

@@ -70,4 +70,3 @@ werkzeug==3.1.8
     #   flask
     #   flask-cors
 
-opengeodeweb-microservice==1.*,>=1.3.2

src/opengeodeweb_back/app.py

@@ -6,6 +6,7 @@
 import json
 import logging
 import queue
+import threading
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -25,6 +26,9 @@
 
 logger = logging.getLogger(__name__)
 
+# Below nginx's default 60 s proxy_read_timeout, which closes idle streams
+EVENTS_KEEP_ALIVE_SECONDS = 15
+
 
 def _register_request_hooks(app: flask.Flask) -> None:
     @app.before_request
@@ -41,25 +45,60 @@ def teardown_request(exception: BaseException | None) -> None:
         utils_functions.teardown_request(flask.current_app, exception)
 
 
+type _Event = tuple[str, dict[str, Any]]
+
+
+# One queue per connected stream: with a single shared queue, a stream left over from a
+# Dropped connection would take events meant for the others.
+class _EventListeners:
+    def __init__(self) -> None:
+        self._queues: set[queue.Queue[_Event]] = set()
+        self._lock = threading.Lock()
+
+    def publish(self, event: str, data: dict[str, Any]) -> None:
+        with self._lock:
+            for listener in self._queues:
+                listener.put((event, data))
+
+    def stream(self, keep_alive_seconds: float) -> flask.Response:
+        listener: queue.Queue[_Event] = queue.Queue()
+        with self._lock:
+            self._queues.add(listener)
+
+        def unregister() -> None:
+            with self._lock:
+                self._queues.discard(listener)
+
+        response = flask.Response(
+            _stream_events(listener, keep_alive_seconds), mimetype="text/event-stream"
+        )
+        response.call_on_close(unregister)
+        return response
+
+
+def _stream_events(listener: queue.Queue[_Event], keep_alive_seconds: float) -> Generator[str]:
+    while True:
+        try:
+            event, data = listener.get(timeout=keep_alive_seconds)
+        except queue.Empty:
+            # Comment line, ignored by EventSource: keeps proxies from closing the idle
+            # Stream and lets a dropped connection be detected
+            yield ": keep-alive\n\n"
+            continue
+        yield f"event: {event}\ndata: {json.dumps(data)}\n\n"
+
+
 def _register_event_stream(app: flask.Flask) -> None:
     def wants_event_stream() -> bool:
         accept = flask.request.headers.get("Accept", "")
         return "text/event-stream" in accept
 
-    _event_queue: queue.Queue[tuple[str, dict[str, Any]]] = queue.Queue()
-
-    def publish_event(event: str, data: dict[str, Any]) -> None:
-        _event_queue.put((event, data))
-
-    def stream_events() -> Generator[str]:
-        while True:
-            event, data = _event_queue.get()
-            yield f"event: {event}\ndata: {json.dumps(data)}\n\n"
+    listeners = _EventListeners()
 
     @app.after_request
     def after_request(response: flask.Response) -> flask.Response:
         endpoint = flask.request.endpoint.replace(".", "/") if flask.request.endpoint else None
-        if endpoint in {"events", None}:
+        if endpoint is None or endpoint == "events":
             return response
 
         if wants_event_stream():
@@ -68,12 +107,14 @@ def after_request(response: flask.Response) -> flask.Response:
                 payload = response.get_json()
             except ValueError:
                 payload = {"status": response.status_code}
-            publish_event(endpoint, payload)
+            listeners.publish(endpoint, payload)
         return response
 
     @app.route("/events")
     def events() -> flask.Response:
-        return flask.Response(stream_events(), mimetype="text/event-stream")
+        return listeners.stream(
+            flask.current_app.config.get("EVENTS_KEEP_ALIVE_SECONDS", EVENTS_KEEP_ALIVE_SECONDS)
+        )
 
 
 def _register_error_handlers(app: flask.Flask) -> None:

tests/test_events.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+# Standard library imports
+import queue
+import threading
+import time
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from flask.testing import FlaskClient
+
+CHUNK_TIMEOUT_SECONDS = 2.0
+LISTENERS_CONNECT_SECONDS = 0.5
+
+
+# Opens a new event stream and reads its first chunk in a daemon thread: the test client
+# Already reads that chunk inside client.get, which blocks until something is published.
+def _read_next_chunk(client: FlaskClient) -> queue.Queue[bytes]:
+    result: queue.Queue[bytes] = queue.Queue()
+
+    def read() -> None:
+        result.put(next(client.get("/events").iter_encoded()))
+
+    threading.Thread(target=read, daemon=True).start()
+    return result
+
+
+def _chunk(result: queue.Queue[bytes]) -> str | None:
+    try:
+        return result.get(timeout=CHUNK_TIMEOUT_SECONDS).decode()
+    except queue.Empty:
+        return None
+
+
+def test_every_listener_receives_each_event(client: FlaskClient) -> None:
+    # A listener left over from a dropped connection must not take events meant for the others
+    first = _read_next_chunk(client)
+    second = _read_next_chunk(client)
+    time.sleep(LISTENERS_CONNECT_SECONDS)
+    client.get("/health", headers={"Accept": "text/event-stream"})
+    assert _chunk(first) == "event: health\ndata: {}\n\n"
+    assert _chunk(second) == "event: health\ndata: {}\n\n"
+
+
+def test_idle_stream_sends_keep_alive(client: FlaskClient) -> None:
+    app = client.application
+    app.config["EVENTS_KEEP_ALIVE_SECONDS"] = 0.1
+    try:
+        assert _chunk(_read_next_chunk(client)) == ": keep-alive\n\n"
+    finally:
+        app.config.pop("EVENTS_KEEP_ALIVE_SECONDS")