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..ea1bd9e0f2 100644 --- a/tests/unit/test_channels/test_plugin.py +++ b/tests/unit/test_channels/test_plugin.py @@ -279,6 +279,63 @@ 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 + 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 + + +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 + 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 + + @pytest.mark.parametrize("unsubscribe_all", [False, True]) @pytest.mark.parametrize("channels", ["foo", ["foo", "bar"]]) async def test_unsubscribe(