Fixing pubsub's listen method to be blocking.

Author: petyaslavova

redis/asyncio/client.py

@@ -1,6 +1,7 @@
 import asyncio
 import copy
 import inspect
+import math
 import re
 import time
 import warnings
@@ -1276,7 +1277,21 @@ async def parse_response(self, block: bool = True, timeout: float = 0):
         if not conn.is_connected:
             await conn.connect()
 
-        read_timeout = None if block else timeout
+        # Block=True: signal "no timeout" to conn.read_response via
+        # math.inf. The connection treats math.inf as the per-read
+        # opt-in for blocking indefinitely without falling back to
+        # self.socket_timeout. Reconnect/AUTH/HELLO/resubscribe
+        # operations performed by the retry layer continue to honor
+        # self.socket_timeout because they do not pass math.inf.
+        #
+        # TODO(next-major): when the async Connection.read_response
+        # default for ``timeout`` is changed to SENTINEL, passing
+        # ``timeout=None`` from this method will become the natural
+        # "no timeout" signal and the math.inf hand-off can be
+        # removed. That swap is a breaking change to the
+        # Connection.read_response signature so it must wait for a
+        # major release.
+        read_timeout = math.inf if block else timeout
         response = await self._execute(
             conn,
             conn.read_response,

redis/asyncio/connection.py

@@ -1,6 +1,7 @@
 import asyncio
 import copy
 import inspect
+import math
 import socket
 import sys
 import time
@@ -746,8 +747,35 @@ async def read_response(
         disconnect_on_error: bool = True,
         push_request: Optional[bool] = False,
     ):
-        """Read the response from a previously sent command"""
-        read_timeout = timeout if timeout is not None else self.socket_timeout
+        """Read the response from a previously sent command.
+
+        ``timeout`` semantics:
+        - ``None`` (default): fall back to ``self.socket_timeout``.
+        - ``math.inf``: block indefinitely with no timeout. Used by PubSub
+          blocking reads (``listen()`` / ``get_message(timeout=None)`` /
+          ``parse_response(block=True)``) where the configured
+          ``socket_timeout`` must not abort the read.
+        - ``float``: apply that timeout in seconds for this single read.
+
+        TODO(next-major): replace the ``math.inf`` opt-in with a SENTINEL
+        default for ``timeout``. After that change, ``timeout=None`` will
+        mean "no timeout, block until a response arrives" (matching the
+        long-standing PubSub docstring contract) and the SENTINEL default
+        will be the value that falls back to ``self.socket_timeout``.
+        That swap is a breaking change, so it must wait for a major
+        release. Until then, callers that need an indefinitely blocking
+        read pass ``math.inf`` explicitly.
+        """
+        # TODO(next-major): drop the math.inf branch. Use SENTINEL as the
+        # default for ``timeout`` and treat ``timeout is None`` as the
+        # "no timeout" signal (matching the PubSub docstring contract).
+        # Match only positive infinity here. ``-math.inf`` is not a valid
+        # "block forever" signal and historically behaved as an already-
+        # expired timeout; preserve that.
+        if timeout == math.inf:
+            read_timeout = None
+        else:
+            read_timeout = timeout if timeout is not None else self.socket_timeout
         host_error = self._host_error()
         try:
             if read_timeout is not None and self.protocol in ["3", 3]:

redis/client.py

@@ -1255,7 +1255,11 @@ def try_read():
                 read_timeout = timeout
             else:
                 conn.connect()
-                read_timeout = SENTINEL  # Use default socket timeout for blocking
+                # Block indefinitely waiting for a pubsub message. timeout=None
+                # makes the socket layer call sock.settimeout(None) for this read
+                # (and restore the original socket_timeout afterwards), so the
+                # configured socket_timeout does not abort the read.
+                read_timeout = None
             return conn.read_response(
                 disconnect_on_error=False, push_request=True, timeout=read_timeout
             )

tests/test_asyncio/test_pubsub.py

@@ -1618,6 +1618,193 @@ async def test_timeout_with_no_subscription(self, r):
         await p.aclose()
 
 
+@pytest.mark.asyncio
+@pytest.mark.onlynoncluster
+class TestAsyncPubSubBlockingListen:
+    """
+    Tests that PubSub.listen() and parse_response(block=True) block
+    indefinitely waiting for a message, ignoring the connection's
+    configured socket_timeout. Regression tests for the issue where
+    listen() would raise TimeoutError once socket_timeout elapsed.
+
+    Retries are disabled on the subscriber so a TimeoutError from the
+    blocking read surfaces immediately instead of being silently retried
+    by the client's default retry policy (which would mask the bug).
+    """
+
+    SOCKET_TIMEOUT = 0.3
+    PUBLISH_DELAY = SOCKET_TIMEOUT * 3
+
+    async def _make_subscriber(self, create_redis):
+        from redis.asyncio.retry import Retry
+        from redis.backoff import NoBackoff
+
+        return await create_redis(
+            socket_timeout=self.SOCKET_TIMEOUT,
+            retry=Retry(NoBackoff(), 0),
+            retry_on_timeout=False,
+        )
+
+    async def test_listen_blocks_longer_than_socket_timeout(self, create_redis):
+        """
+        listen() must keep blocking past socket_timeout. Configure a very
+        short socket_timeout, then publish after a delay that exceeds it,
+        and assert listen() yields the message instead of raising
+        TimeoutError.
+        """
+        r = await self._make_subscriber(create_redis)
+        publisher = await create_redis()
+        p = r.pubsub()
+        await p.subscribe("foo")
+        msg = await wait_for_message(p, timeout=1.0)
+        assert msg is not None and msg["type"] == "subscribe"
+
+        async def publish_after_delay():
+            await asyncio.sleep(self.PUBLISH_DELAY)
+            await publisher.publish("foo", "hello")
+
+        task = asyncio.create_task(publish_after_delay())
+
+        start = asyncio.get_running_loop().time()
+        msg = await p.listen().__anext__()
+        elapsed = asyncio.get_running_loop().time() - start
+        assert msg["type"] == "message"
+        assert msg["data"] == b"hello"
+        # Should have blocked past socket_timeout to receive the message.
+        assert elapsed > self.SOCKET_TIMEOUT
+        await task
+        await p.aclose()
+
+    async def test_get_message_block_indefinitely_past_socket_timeout(
+        self, create_redis
+    ):
+        """
+        get_message(timeout=None) must block past socket_timeout.
+        """
+        r = await self._make_subscriber(create_redis)
+        publisher = await create_redis()
+        p = r.pubsub()
+        await p.subscribe("foo")
+        msg = await wait_for_message(p, timeout=1.0)
+        assert msg is not None
+
+        async def publish_after_delay():
+            await asyncio.sleep(self.PUBLISH_DELAY)
+            await publisher.publish("foo", "delayed")
+
+        task = asyncio.create_task(publish_after_delay())
+
+        start = asyncio.get_running_loop().time()
+        msg = await p.get_message(timeout=None)
+        elapsed = asyncio.get_running_loop().time() - start
+        assert msg is not None
+        assert msg["type"] == "message"
+        assert msg["data"] == b"delayed"
+        assert elapsed > self.SOCKET_TIMEOUT
+        await task
+        await p.aclose()
+
+    async def test_parse_response_block_true_blocks_past_socket_timeout(
+        self, create_redis
+    ):
+        """
+        parse_response(block=True) must block past socket_timeout.
+        """
+        r = await self._make_subscriber(create_redis)
+        publisher = await create_redis()
+        p = r.pubsub(ignore_subscribe_messages=True)
+        await p.subscribe("foo")
+        # Drain the subscribe ack before calling parse_response directly.
+        await wait_for_message(p, timeout=1.0)
+
+        async def publish_after_delay():
+            await asyncio.sleep(self.PUBLISH_DELAY)
+            await publisher.publish("foo", "raw")
+
+        task = asyncio.create_task(publish_after_delay())
+
+        start = asyncio.get_running_loop().time()
+        response = await p.parse_response(block=True)
+        elapsed = asyncio.get_running_loop().time() - start
+        # Response is the raw pubsub frame [b"message", b"foo", b"raw"].
+        assert response is not None
+        assert response[0] == b"message"
+        assert response[-1] == b"raw"
+        assert elapsed > self.SOCKET_TIMEOUT
+        await task
+        await p.aclose()
+
+    async def test_blocking_listen_does_not_mutate_socket_timeout(
+        self, create_redis
+    ):
+        """
+        A blocking read must not mutate the connection's configured
+        socket_timeout. Otherwise reconnect/AUTH/HELLO/resubscribe paths
+        triggered by retries inside _execute would inherit no timeout and
+        could hang indefinitely. The fix uses a per-read sentinel rather
+        than clearing socket_timeout for the duration of the call.
+        """
+        r = await self._make_subscriber(create_redis)
+        publisher = await create_redis()
+        p = r.pubsub()
+        await p.subscribe("foo")
+        await wait_for_message(p, timeout=1.0)
+
+        orig_timeout = p.connection.socket_timeout
+        assert orig_timeout == self.SOCKET_TIMEOUT
+
+        # Sample socket_timeout while the blocking read is in progress, to
+        # catch any mutation that happens only for the duration of the call.
+        async def sample_timeout_during_read():
+            # Yield a few times so the get_message read is in flight.
+            for _ in range(10):
+                await asyncio.sleep(0)
+            assert p.connection.socket_timeout == orig_timeout
+
+        async def publish_after_delay():
+            await asyncio.sleep(self.PUBLISH_DELAY)
+            await publisher.publish("foo", "msg")
+
+        sample_task = asyncio.create_task(sample_timeout_during_read())
+        publish_task = asyncio.create_task(publish_after_delay())
+
+        msg = await p.get_message(timeout=None)
+        await sample_task
+        assert msg is not None and msg["data"] == b"msg"
+        assert p.connection.socket_timeout == orig_timeout
+        await publish_task
+        await p.aclose()
+
+    async def test_block_indefinitely_does_not_alter_subsequent_reads(
+        self, create_redis
+    ):
+        """
+        After a blocking read returns, a follow-up non-blocking read must
+        still honor the configured ``socket_timeout`` (i.e. the per-read
+        ``block_indefinitely`` does not bleed into later operations and
+        does not require mutating any connection-wide state).
+        """
+        r = await self._make_subscriber(create_redis)
+        publisher = await create_redis()
+        p = r.pubsub()
+        await p.subscribe("foo")
+        await wait_for_message(p, timeout=1.0)
+
+        # First a blocking read.
+        await publisher.publish("foo", "first")
+        msg = await p.get_message(timeout=None)
+        assert msg is not None and msg["data"] == b"first"
+
+        # Then a non-blocking read with a short timeout and no message.
+        # If block_indefinitely leaked, this would block past the timeout.
+        start = asyncio.get_running_loop().time()
+        msg = await p.get_message(timeout=0.1)
+        elapsed = asyncio.get_running_loop().time() - start
+        assert msg is None
+        assert elapsed < 0.5
+        await p.aclose()
+
+
 @pytest.mark.asyncio
 class TestPubSubHandleMessageMetrics:
     """Tests for handle_message recording pubsub metrics."""