diff --git a/backend/common/tests/test_notification_sse.py b/backend/common/tests/test_notification_sse.py index 0b91ce116..95df839c0 100644 --- a/backend/common/tests/test_notification_sse.py +++ b/backend/common/tests/test_notification_sse.py @@ -14,6 +14,7 @@ from rest_framework.test import APIClient from common.models import Notification, Org, Profile, User +from common.views import notification_views from common.views.notification_views import ( _format_keepalive, _format_sse, @@ -119,6 +120,34 @@ async def runner(): assert payload["link"] == "/cases/abc" assert payload["data"] == {"comment_excerpt": "hi"} + def test_stream_self_terminates_at_deadline(self): + """The stream must end on its own once MAX_STREAM_SECONDS elapses, so a + long-open browser tab cannot pin a worker thread / DB connection + forever (the connection-exhaustion regression).""" + async def runner(): + pubsub = _StubPubSub() # no messages -> keepalive-only loop + gen = _stream_events( + "notif:o:p", recipient_id=self.profile.id, pubsub=pubsub + ) + first = await asyncio.wait_for(gen.__anext__(), timeout=1) + stopped = False + try: + # Deadline already passed -> next pull should end the stream. + await asyncio.wait_for(gen.__anext__(), timeout=2) + except StopAsyncIteration: + stopped = True + await gen.aclose() + return first, stopped + + original = notification_views.MAX_STREAM_SECONDS + notification_views.MAX_STREAM_SECONDS = 0 + try: + first, stopped = asyncio.run(runner()) + finally: + notification_views.MAX_STREAM_SECONDS = original + assert first == b": keepalive\n\n" + assert stopped, "stream did not terminate at its deadline" + def test_drops_message_for_other_recipient(self): n_other = Notification.objects.create( org=self.org, recipient=self.other_profile, verb="other" diff --git a/backend/common/views/notification_views.py b/backend/common/views/notification_views.py index c0cce8e79..21a91bc21 100644 --- a/backend/common/views/notification_views.py +++ b/backend/common/views/notification_views.py @@ -16,6 +16,7 @@ import asyncio import json import logging +import time from asgiref.sync import sync_to_async from django.http import StreamingHttpResponse @@ -36,6 +37,17 @@ KEEPALIVE_SECONDS = 15 +# Upper bound on a single SSE connection's lifetime. The browser's EventSource +# transparently reconnects when the server closes the stream, so capping this +# is invisible to the user but critical for the server: under ASGI the *entire* +# stream runs inside one request (one `ThreadSensitiveContext`), so an +# unbounded `while True` stream pins a worker thread — and any DB/redis +# connection bound to it — for as long as the tab stays open (hours). Left +# unbounded this is what slowly exhausts a shared Postgres cluster and leaves +# workers undrainable on deploy. 5 minutes bounds the leak window and lets +# workers recycle. +MAX_STREAM_SECONDS = 300 + DEFAULT_LIMIT = 20 MAX_LIMIT = 100 @@ -141,38 +153,67 @@ async def _aget_serialized(notif_id, recipient_id): """Fetch a notification by id, scoped to recipient. Returns dict or None.""" @sync_to_async def fetch(): - notif = Notification.objects.filter( - pk=notif_id, recipient_id=recipient_id - ).first() - if notif is None: - return None - return NotificationSerializer(notif).data + from django.db import connection + + try: + notif = Notification.objects.filter( + pk=notif_id, recipient_id=recipient_id + ).first() + if notif is None: + return None + return NotificationSerializer(notif).data + finally: + # This runs in asgiref's shared thread-sensitive executor, outside + # Django's request/response cycle — so `close_old_connections` + # never fires for it. Without this explicit close the connection + # dangles open (the SSE tests need TransactionTestCase precisely + # because of it). Close it so each fetch reclaims its connection. + connection.close() return await fetch() +async def _aclose_redis(obj): + """Close a redis.asyncio client/pubsub across versions (aclose vs close).""" + if obj is None: + return + closer = getattr(obj, "aclose", None) or getattr(obj, "close", None) + if closer is None: + return + try: + result = closer() + if asyncio.iscoroutine(result): + await result + except Exception: # pragma: no cover - best effort + pass + + async def _open_pubsub(channel: str): """Open a redis.asyncio pubsub subscribed to ``channel``. - Returns ``None`` if Redis is unreachable; the stream then runs in - keepalive-only mode (the frontend's polling-since path provides - backfill). + Returns ``(client, pubsub)``. Both are ``None`` if Redis is unreachable; + the stream then runs in keepalive-only mode (the frontend's polling-since + path provides backfill). The caller owns closing BOTH the pubsub and the + underlying client connection pool — closing only the pubsub leaks the + client's pooled connections, one per dropped SSE stream. """ try: import redis.asyncio as aioredis # type: ignore except ImportError: # pragma: no cover - return None + return None, None from django.conf import settings url = getattr(settings, "CELERY_BROKER_URL", None) or "redis://localhost:6379/0" + client = None try: client = aioredis.from_url(url) pubsub = client.pubsub() await pubsub.subscribe(channel) - return pubsub + return client, pubsub except Exception as exc: logger.warning("SSE redis subscribe failed (%s); keepalive-only mode", exc) - return None + await _aclose_redis(client) + return None, None async def _stream_events(channel: str, recipient_id, *, pubsub=None): @@ -182,12 +223,18 @@ async def _stream_events(channel: str, recipient_id, *, pubsub=None): opens its own redis.asyncio pubsub via :func:`_open_pubsub`. """ owns_pubsub = pubsub is None + client = None if owns_pubsub: - pubsub = await _open_pubsub(channel) + client, pubsub = await _open_pubsub(channel) + deadline = time.monotonic() + MAX_STREAM_SECONDS try: # Initial comment so the client confirms the stream opened. yield _format_keepalive() while True: + # Bounded lifetime: end the stream so the worker thread and its + # DB/redis connections are released. EventSource reconnects. + if time.monotonic() >= deadline: + return if pubsub is None: await asyncio.sleep(KEEPALIVE_SECONDS) yield _format_keepalive() @@ -211,12 +258,15 @@ async def _stream_events(channel: str, recipient_id, *, pubsub=None): continue yield _format_sse("notification", payload) finally: - if owns_pubsub and pubsub is not None: - try: - await pubsub.unsubscribe(channel) - await pubsub.close() - except Exception: # pragma: no cover - best effort - pass + if owns_pubsub: + if pubsub is not None: + try: + await pubsub.unsubscribe(channel) + except Exception: # pragma: no cover - best effort + pass + await _aclose_redis(pubsub) + # Close the underlying client pool too — not just the pubsub. + await _aclose_redis(client) def _drive_async_gen(agen): @@ -271,6 +321,18 @@ def get(self, request, *args, **kwargs): profile_id = request.profile.id channel = notif_mod.channel_for(org_id, profile_id) + # Release THIS request thread's DB connection before streaming begins. + # Under ASGI the entire stream runs inside one request's + # ThreadSensitiveContext, so the connection opened by auth/middleware + # would otherwise stay checked out for the full (possibly hours-long) + # stream — one leaked Postgres connection per open browser tab, which + # is what exhausts the shared cluster. The stream's own reads go + # through a separate executor (see `_aget_serialized`) and close + # themselves, so dropping this one is safe. + from django.db import connection + + connection.close() + response = StreamingHttpResponse( _drive_async_gen(_stream_events(channel, profile_id)), content_type="text/event-stream",