diff --git a/api/config.py b/api/config.py index f02262b1..fc1034bb 100644 --- a/api/config.py +++ b/api/config.py @@ -27,6 +27,16 @@ class PubSubSettings(BaseSettings): cloud_events_source: str = "https://api.kernelci.org/" redis_host: str = "redis" redis_db_number: int = 1 + # Every pub/sub subscription pins one dedicated Redis connection, so the + # pool must be large enough for all concurrent subscribers plus regular + # commands. redis-py >= 8 defaults to only 100. + redis_max_connections: int = 1000 + # Maximum number of concurrent subscriptions a single user may hold. + # Before the limit is enforced, that user's subscriptions which were + # never polled and are older than abandoned_subscription_seconds are + # evicted, so a crash-looping client cannot pile them up. + max_subscriptions_per_user: int = 20 + abandoned_subscription_seconds: int = 60 keep_alive_period: int = 45 # MongoDB durable pub/sub settings event_ttl_days: int = 7 # Auto-delete events after N days diff --git a/api/main.py b/api/main.py index a3782066..be3a1adc 100644 --- a/api/main.py +++ b/api/main.py @@ -57,6 +57,7 @@ parse_node_obj, ) from pydantic import BaseModel +from redis import exceptions as redis_exceptions from .auth import Authentication from .config import AuthSettings @@ -80,7 +81,7 @@ UserUpdate, UserUpdateRequest, ) -from .pubsub_mongo import PubSub +from .pubsub_mongo import PubSub, SubscriptionLimitError from .user_manager import create_user_manager, get_user_manager SUBSCRIPTION_CLEANUP_INTERVAL_MINUTES = 15 # How often to run cleanup task @@ -2146,7 +2147,13 @@ async def subscribe( options["promiscuous"] = promisc if subscriber_id: options["subscriber_id"] = subscriber_id - return await pubsub.subscribe(channel, user.username, options) + try: + return await pubsub.subscribe(channel, user.username, options) + except SubscriptionLimitError as error: + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail=str(error), + ) from error @app.post("/unsubscribe/{sub_id}") @@ -2385,10 +2392,27 @@ def traceback_exception_handler(request: Request, exc: Exception): ) +def redis_unavailable_exception_handler(request: Request, exc: Exception): + """Report Redis connection problems as a retryable 503, not a 500""" + print(f"Redis unavailable: {exc}") + return JSONResponse( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + headers={"Retry-After": "5"}, + content={"message": "Redis unavailable, please retry"}, + ) + + versioned_app.add_exception_handler(ValueError, value_error_exception_handler) versioned_app.add_exception_handler( errors.InvalidId, invalid_id_exception_handler ) +versioned_app.add_exception_handler( + redis_exceptions.MaxConnectionsError, + redis_unavailable_exception_handler, +) +versioned_app.add_exception_handler( + redis_exceptions.ConnectionError, redis_unavailable_exception_handler +) versioned_app.add_exception_handler(Exception, traceback_exception_handler) diff --git a/api/pubsub_mongo.py b/api/pubsub_mongo.py index ffe2f299..f534bbe8 100644 --- a/api/pubsub_mongo.py +++ b/api/pubsub_mongo.py @@ -33,6 +33,10 @@ logger = logging.getLogger(__name__) +class SubscriptionLimitError(Exception): + """Raised when a user already holds too many active subscriptions""" + + class PubSub: """Hybrid Pub/Sub implementation with MongoDB durability @@ -80,7 +84,9 @@ def __init__( db_number = self._settings.redis_db_number self._redis = aioredis.from_url( - "redis://" + host + "/" + str(db_number), health_check_interval=30 + "redis://" + host + "/" + str(db_number), + health_check_interval=30, + max_connections=self._settings.redis_max_connections, ) # MongoDB setup @@ -226,9 +232,26 @@ async def _keep_alive(self): break for channel in channels: # Use _publish_keepalive to avoid storing BEEP in MongoDB - await self._publish_keepalive(channel, "BEEP") + try: + await self._publish_keepalive(channel, "BEEP") + except aioredis.RedisError as exc: + # A transient Redis error (e.g. exhausted connection + # pool) must not kill the keep-alive loop for good + logger.warning("Keep-alive publish failed: %s", exc) await asyncio.sleep(self._settings.keep_alive_period) + @staticmethod + async def _close_redis_sub(redis_sub): + """Best-effort close of a Redis PubSub, releasing its connection""" + try: + await redis_sub.unsubscribe() + except aioredis.RedisError as exc: + logger.debug("PubSub unsubscribe failed: %s", exc) + try: + await redis_sub.aclose() + except aioredis.RedisError as exc: + logger.debug("PubSub close failed: %s", exc) + async def _publish_keepalive(self, channel: str, data: str): """Publish keep-alive message (Redis only, no MongoDB storage)""" attributes = { @@ -372,16 +395,23 @@ async def subscribe( Returns: Subscription object with id, channel, user, promiscuous fields """ - sub_id = await self._redis.incr(self.ID_KEY) subscriber_id = options.get("subscriber_id") if options else None promiscuous = options.get("promiscuous", False) if options else False + await self._enforce_subscription_limit(user) + sub_id = await self._redis.incr(self.ID_KEY) + async with self._lock: redis_sub = self._redis.pubsub() sub = Subscription( id=sub_id, channel=channel, user=user, promiscuous=promiscuous ) - await redis_sub.subscribe(channel) + try: + await redis_sub.subscribe(channel) + except aioredis.RedisError: + # Never leave a half-initialised PubSub holding a connection + await self._close_redis_sub(redis_sub) + raise self._subscriptions[sub_id] = { "redis_sub": redis_sub, @@ -407,6 +437,54 @@ async def subscribe( return sub + async def _enforce_subscription_limit(self, user: str): + """Evict the user's abandoned subscriptions, then enforce the cap + + A subscription that has never been polled and is older than + abandoned_subscription_seconds belongs to a client that died (or + crash-loops) right after subscribing. Those are evicted first so + they do not count against the user. If the user still holds + max_subscriptions_per_user active subscriptions afterwards, a + SubscriptionLimitError is raised. + """ + limit = self._settings.max_subscriptions_per_user + cutoff = datetime.utcnow() - timedelta( + seconds=self._settings.abandoned_subscription_seconds + ) + + async with self._lock: + user_sub_ids = [] + abandoned_ids = [] + for sub_id, sub_data in self._subscriptions.items(): + if sub_data["sub"].user != user: + continue + user_sub_ids.append(sub_id) + if ( + sub_data.get("last_poll") is None + and sub_data["created"] < cutoff + ): + abandoned_ids.append(sub_id) + + # unsubscribe() takes the lock itself, so call it outside of it + for sub_id in abandoned_ids: + try: + await self.unsubscribe(sub_id) + except KeyError: + pass + if abandoned_ids: + logger.warning( + "Evicted %d abandoned (never polled) subscriptions of %s", + len(abandoned_ids), + user, + ) + + active = len(user_sub_ids) - len(abandoned_ids) + if limit and active >= limit: + raise SubscriptionLimitError( + f"User {user} already holds {active} active subscriptions " + f"(limit {limit}); unsubscribe or poll them first" + ) + async def _setup_durable_subscription( self, sub_id: int, @@ -479,8 +557,7 @@ async def unsubscribe(self, sub_id: int, user: Optional[str] = None): self._subscriptions.pop(sub_id) self._update_channels() - await sub["redis_sub"].unsubscribe() - await sub["redis_sub"].close() + await self._close_redis_sub(sub["redis_sub"]) async def _get_listen_subscription( self, sub_id: int, user: Optional[str] = None @@ -532,10 +609,14 @@ async def _rebuild_redis_subscription( ): async with self._lock: channel = sub.channel + old_redis_sub = sub_data.get("redis_sub") new_redis_sub = self._redis.pubsub() await new_redis_sub.subscribe(channel) self._subscriptions[sub_id]["redis_sub"] = new_redis_sub sub_data["redis_sub"] = new_redis_sub + # Release the broken PubSub, otherwise its pooled connection leaks + if old_redis_sub is not None: + await self._close_redis_sub(old_redis_sub) def _maybe_update_delivery_offset( self, subscriber_id: Optional[str], sub_data: dict, msg_data: Any @@ -695,6 +776,10 @@ async def cleanup_stale_subscriptions( ) -> int: """Remove subscriptions not polled recently + Subscriptions that were never polled at all are aged from their + creation time, otherwise a client that subscribes and dies before + its first listen() leaks a Redis connection forever. + For durable subscriptions, only the in-memory state is cleaned up. The MongoDB subscriber state is preserved for reconnection. """ @@ -703,8 +788,10 @@ async def cleanup_stale_subscriptions( async with self._lock: for sub_id, sub_data in self._subscriptions.items(): - last_poll = sub_data.get("last_poll") - if last_poll and last_poll < cutoff: + last_activity = sub_data.get("last_poll") or sub_data.get( + "created" + ) + if last_activity and last_activity < cutoff: stale_ids.append(sub_id) for sub_id in stale_ids: diff --git a/tests/unit_tests/test_pubsub_mongo.py b/tests/unit_tests/test_pubsub_mongo.py new file mode 100644 index 00000000..ce79d729 --- /dev/null +++ b/tests/unit_tests/test_pubsub_mongo.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: LGPL-2.1-or-later +# +# Copyright (C) 2026 Collabora Limited +# Author: Denys Fedoryshchenko + +"""Unit tests for the MongoDB-backed Pub/Sub Redis connection handling""" + +from datetime import datetime, timedelta +from unittest.mock import AsyncMock, MagicMock + +import fakeredis.aioredis +import pytest +from redis import exceptions as redis_exceptions + +from api.pubsub_mongo import PubSub, SubscriptionLimitError + + +@pytest.fixture +def mongo_pubsub(mocker): + """PubSub instance backed by fakeredis and a dummy Mongo client""" + pubsub = PubSub(mongo_client=MagicMock()) + mocker.patch.object(pubsub, "_redis", fakeredis.aioredis.FakeRedis()) + # Keep-alive timer is irrelevant here and needs a running loop thread + mocker.patch.object(pubsub, "_start_keep_alive_timer") + return pubsub + + +@pytest.mark.asyncio +async def test_cleanup_never_polled_subscription(mongo_pubsub): + """ + Test Case: A subscription that was created but never polled must be + removed once it is older than the stale cutoff. + + Expected Result: + cleanup_stale_subscriptions() returns 1 and the subscription is + gone from PubSub._subscriptions, releasing its Redis connection. + """ + sub = await mongo_pubsub.subscribe("node", "test") + assert mongo_pubsub._subscriptions[sub.id]["last_poll"] is None + mongo_pubsub._subscriptions[sub.id]["created"] = ( + datetime.utcnow() - timedelta(minutes=20) + ) + + cleaned = await mongo_pubsub.cleanup_stale_subscriptions(15) + + assert cleaned == 1 + assert not mongo_pubsub._subscriptions + + +@pytest.mark.asyncio +async def test_cleanup_keeps_fresh_subscriptions(mongo_pubsub): + """ + Test Case: Subscriptions that are either recently created or recently + polled must survive the stale cleanup. + + Expected Result: + cleanup_stale_subscriptions() returns 0 and both entries remain. + """ + fresh = await mongo_pubsub.subscribe("node", "test") + polled = await mongo_pubsub.subscribe("node", "test") + mongo_pubsub._subscriptions[polled.id]["created"] = ( + datetime.utcnow() - timedelta(hours=5) + ) + mongo_pubsub._subscriptions[polled.id]["last_poll"] = datetime.utcnow() + + cleaned = await mongo_pubsub.cleanup_stale_subscriptions(15) + + assert cleaned == 0 + assert set(mongo_pubsub._subscriptions) == {fresh.id, polled.id} + + +@pytest.mark.asyncio +async def test_rebuild_closes_old_redis_subscription(mongo_pubsub): + """ + Test Case: Rebuilding the Redis subscription after a connection error + must close the previous PubSub object. + + Expected Result: + The old PubSub is unsubscribed and closed, and the subscription + entry points at a new PubSub object. + """ + sub = await mongo_pubsub.subscribe("node", "test") + sub_data = mongo_pubsub._subscriptions[sub.id] + old_redis_sub = AsyncMock() + sub_data["redis_sub"] = old_redis_sub + + await mongo_pubsub._rebuild_redis_subscription(sub.id, sub, sub_data) + + old_redis_sub.unsubscribe.assert_awaited_once() + old_redis_sub.aclose.assert_awaited_once() + assert sub_data["redis_sub"] is not old_redis_sub + assert ( + mongo_pubsub._subscriptions[sub.id]["redis_sub"] + is sub_data["redis_sub"] + ) + + +@pytest.mark.asyncio +async def test_subscribe_failure_releases_connection(mongo_pubsub, mocker): + """ + Test Case: When the Redis SUBSCRIBE command fails, the half-initialised + PubSub must be closed and nothing recorded in PubSub._subscriptions. + + Expected Result: + The Redis error propagates, the PubSub object is closed and no + subscription entry is left behind. + """ + failing_sub = AsyncMock() + failing_sub.subscribe.side_effect = redis_exceptions.MaxConnectionsError( + "Too many connections" + ) + mocker.patch.object( + mongo_pubsub._redis, "pubsub", MagicMock(return_value=failing_sub) + ) + + with pytest.raises(redis_exceptions.MaxConnectionsError): + await mongo_pubsub.subscribe("node", "test") + + failing_sub.aclose.assert_awaited_once() + assert not mongo_pubsub._subscriptions + + +@pytest.mark.asyncio +async def test_keep_alive_survives_redis_error(mongo_pubsub, mocker): + """ + Test Case: A Redis error while publishing the keep-alive BEEP must not + terminate the keep-alive loop. + + Expected Result: + _keep_alive() keeps iterating after the failed publish and only + stops once there are no channels left to keep alive. + """ + await mongo_pubsub.subscribe("node", "test") + publish = mocker.patch.object( + mongo_pubsub, + "_publish_keepalive", + AsyncMock(side_effect=redis_exceptions.MaxConnectionsError("Too many")), + ) + + async def drop_channels(_period): + # Simulate the last subscriber going away during the sleep + mongo_pubsub._channels = set() + + mocker.patch("api.pubsub_mongo.asyncio.sleep", side_effect=drop_channels) + + await mongo_pubsub._keep_alive() + + publish.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_subscription_limit_per_user(mongo_pubsub): + """ + Test Case: A user may hold at most max_subscriptions_per_user active + subscriptions; another user is not affected by that count. + + Expected Result: + The 21st subscribe() for the same user raises + SubscriptionLimitError while a different user can still subscribe. + """ + limit = mongo_pubsub._settings.max_subscriptions_per_user + assert limit == 20 + for _ in range(limit): + sub = await mongo_pubsub.subscribe("node", "greedy") + # Recently polled subscriptions are legitimately active + mongo_pubsub._subscriptions[sub.id]["last_poll"] = datetime.utcnow() + + with pytest.raises(SubscriptionLimitError): + await mongo_pubsub.subscribe("node", "greedy") + + other = await mongo_pubsub.subscribe("node", "other") + assert other.user == "other" + assert len(mongo_pubsub._subscriptions) == limit + 1 + + +@pytest.mark.asyncio +async def test_subscribe_evicts_abandoned_subscriptions(mongo_pubsub): + """ + Test Case: Never-polled subscriptions older than one minute belonging + to the subscribing user are evicted before the limit is enforced, but + fresh never-polled ones (client may still be about to listen) are kept. + + Expected Result: + The old abandoned subscriptions disappear, the fresh one survives, + and the new subscribe() succeeds despite the user having been at + the limit. + """ + limit = mongo_pubsub._settings.max_subscriptions_per_user + old = datetime.utcnow() - timedelta(minutes=2) + abandoned_ids = [] + for _ in range(limit - 1): + sub = await mongo_pubsub.subscribe("node", "crashy") + mongo_pubsub._subscriptions[sub.id]["created"] = old + abandoned_ids.append(sub.id) + fresh = await mongo_pubsub.subscribe("node", "crashy") + + new = await mongo_pubsub.subscribe("node", "crashy") + + remaining = set(mongo_pubsub._subscriptions) + assert not remaining & set(abandoned_ids) + assert remaining == {fresh.id, new.id}