fix(channels): make subscribe cancellation-safe to prevent subscriber leak

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:

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(