From ada637a64752974153e420f726eea1996940fad7 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Sat, 20 Jun 2026 11:38:28 +0300 Subject: [PATCH 1/2] fix(channels): make subscribe cancellation-safe to prevent subscriber leak ChannelsPlugin.subscribe registered the new subscriber into self._channels synchronously, before awaiting backend.subscribe and put_subscriber_history. The history fetch performs real I/O (e.g. redis xrevrange), so a cancellation there raised CancelledError without returning the subscriber: it stayed registered in _channels, unreachable by the caller and never unsubscribed, and the backend stayed subscribed. Under bursty connect/disconnect (e.g. SSE clients that disconnect mid-subscribe) _channels grew without bound until restart. Fetch history before registering the subscriber so a cancellation discards an unregistered subscriber instead of leaking it. start_subscription is fixed by the same change, since subscribe now leaves nothing behind on cancel. Closes #4871 Co-Authored-By: Claude Opus 4.8 (1M context) --- litestar/channels/plugin.py | 11 +++-- tests/unit/test_channels/test_plugin.py | 55 +++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 3 deletions(-) diff --git a/litestar/channels/plugin.py b/litestar/channels/plugin.py index 7110ff7739..e6041f2599 100644 --- a/litestar/channels/plugin.py +++ b/litestar/channels/plugin.py @@ -186,6 +186,14 @@ async def subscribe(self, channels: str | Iterable[str], history: int | None = N max_backlog=self._max_backlog, backlog_strategy=self._backlog_strategy, ) + + # Fetch history before registering the subscriber. ``put_subscriber_history`` performs the only real + # I/O here (the cancellation window); doing it first means a cancellation discards an unregistered + # subscriber instead of leaking it into ``_channels``, where the caller could never reach it to + # unsubscribe. See https://github.com/litestar-org/litestar/issues/4871. + if history: + await self.put_subscriber_history(subscriber=subscriber, limit=history, channels=channels) + channels_to_subscribe = set() for channel in channels: @@ -205,9 +213,6 @@ async def subscribe(self, channels: str | Iterable[str], history: int | None = N if channels_to_subscribe: await self._backend.subscribe(channels_to_subscribe) - if history: - await self.put_subscriber_history(subscriber=subscriber, limit=history, channels=channels) - return subscriber async def unsubscribe(self, subscriber: Subscriber, channels: str | Iterable[str] | None = None) -> None: diff --git a/tests/unit/test_channels/test_plugin.py b/tests/unit/test_channels/test_plugin.py index 51aadf01f0..85317ce8bb 100644 --- a/tests/unit/test_channels/test_plugin.py +++ b/tests/unit/test_channels/test_plugin.py @@ -279,6 +279,61 @@ async def test_subscribe_non_existent_channel_raises(memory_backend: MemoryChann await plugin.subscribe("bar") +class _SlowHistoryBackend(MemoryChannelsBackend): + """Stand-in for any backend whose ``get_history`` suspends (e.g. redis ``xrevrange``).""" + + def __init__(self, history: int = 1) -> None: + super().__init__(history=history) + self.in_history = asyncio.Event() + self.release = asyncio.Event() + + async def get_history(self, channel: str, limit: int | None = None) -> list[bytes]: + self.in_history.set() + await self.release.wait() # suspend, as a real backend would on I/O + return await super().get_history(channel, limit) + + +async def test_subscribe_cancellation_does_not_leak() -> None: + """A ``subscribe`` cancelled while fetching history must leave no registered subscriber. + + Regression test for https://github.com/litestar-org/litestar/issues/4871: the subscriber was + registered into ``_channels`` before the history ``await``, so a cancellation there leaked it + permanently (the caller never received it and could never ``unsubscribe``). + """ + backend = _SlowHistoryBackend(history=1) + plugin = ChannelsPlugin(backend=backend, arbitrary_channels_allowed=True) + channel = "user:1" + + task = asyncio.create_task(plugin.subscribe([channel], history=1)) + await backend.in_history.wait() # subscribe is now suspended in get_history + + task.cancel() # client disconnects mid-subscribe + await asyncio.gather(task, return_exceptions=True) + + assert not plugin._channels.get(channel) # no subscriber left behind + assert channel not in backend._channels # backend not left subscribed + + +async def test_start_subscription_cancellation_does_not_leak() -> None: + """``start_subscription`` cancelled while fetching history must leave no registered subscriber.""" + backend = _SlowHistoryBackend(history=1) + plugin = ChannelsPlugin(backend=backend, arbitrary_channels_allowed=True) + channel = "user:1" + + async def consume() -> None: + async with plugin.start_subscription([channel], history=1): + pass + + task = asyncio.create_task(consume()) + await backend.in_history.wait() # entering the context manager, suspended in get_history + + task.cancel() # client disconnects mid-subscribe + await asyncio.gather(task, return_exceptions=True) + + assert not plugin._channels.get(channel) # no subscriber left behind + assert channel not in backend._channels # backend not left subscribed + + @pytest.mark.parametrize("unsubscribe_all", [False, True]) @pytest.mark.parametrize("channels", ["foo", ["foo", "bar"]]) async def test_unsubscribe( From b16381d88704f63e0ca077615ccd8e4fdbc2ff31 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Sat, 27 Jun 2026 21:41:35 +0300 Subject: [PATCH 2/2] review fix --- tests/unit/test_channels/test_plugin.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_channels/test_plugin.py b/tests/unit/test_channels/test_plugin.py index 85317ce8bb..ea1bd9e0f2 100644 --- a/tests/unit/test_channels/test_plugin.py +++ b/tests/unit/test_channels/test_plugin.py @@ -308,7 +308,8 @@ async def test_subscribe_cancellation_does_not_leak() -> None: await backend.in_history.wait() # subscribe is now suspended in get_history task.cancel() # client disconnects mid-subscribe - await asyncio.gather(task, return_exceptions=True) + results = await asyncio.gather(task, return_exceptions=True) + assert isinstance(results[0], asyncio.CancelledError) # the cancellation propagated assert not plugin._channels.get(channel) # no subscriber left behind assert channel not in backend._channels # backend not left subscribed @@ -328,7 +329,8 @@ async def consume() -> None: await backend.in_history.wait() # entering the context manager, suspended in get_history task.cancel() # client disconnects mid-subscribe - await asyncio.gather(task, return_exceptions=True) + results = await asyncio.gather(task, return_exceptions=True) + assert isinstance(results[0], asyncio.CancelledError) # the cancellation propagated assert not plugin._channels.get(channel) # no subscriber left behind assert channel not in backend._channels # backend not left subscribed