Repository: redis/redis-py
Pull request #2281 (2022-07-17 10:26:21)

automatically reconnect pubsub when reading messages in blocking mode
### Pull Request check-list

_Please make sure to review and check all of these items:_

- [x] Does `$ tox` pass with this change (including linting)?
- [x] Do the CI tests pass with this change (enable it first in your forked repo and wait for the github action build to finish)?
- [x] Is the new or changed code fully tested?
- [ ] Is a documentation update included (if this change modifies existing APIs, or introduces new ones)?
- [ ] Is there an example added to the examples folder (if applicable)?
- [ ] Was the change added to CHANGES file?

_NOTE: these things are not required to open a PR and can be done
afterwards / while the PR is open._

### Description of change

Issue #2089 demonstrates how a pub-sub connection doesn't automatically reconnect.  Further investigation showed that this was due to the fact that the `connection` only performed a just-in-time connect in the `can_read()` method, which is only invoked if the operation to read a message has a _timeout_.  In particular, using `listen()` to get the message, instead of the `get_message()` with a timeout, the `can_read()` method isn't called, and so a `connect()` call is not made.

This issue is present in both non-async and async versions of the library as demonstrated with the unit-test code.

The PR fixes this by adding a just-in-time connect to the `read_response` call of both the `Connection` object.

codecov-commenter: # [Codecov](https://codecov.io/gh/redis/redis-py/pull/2281?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) Report
> Merging [#2281](https://codecov.io/gh/redis/redis-py/pull/2281?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) (8664478) into [master](https://codecov.io/gh/redis/redis-py/commit/4b0543d567aef36ac467ce495d831a24575d8d5b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) (4b0543d) will **decrease** coverage by `14.57%`.
> The diff coverage is `95.00%`.

```diff
@@             Coverage Diff             @@
##           master    #2281       +/-   ##
===========================================
- Coverage   91.87%   77.30%   -14.58%     
===========================================
  Files         108      108               
  Lines       27783    27082      -701     
===========================================
- Hits        25527    20935     -4592     
- Misses       2256     6147     +3891     
```


| [Impacted Files](https://codecov.io/gh/redis/redis-py/pull/2281?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) | Coverage Δ | |
|---|---|---|
| [tests/conftest.py](https://codecov.io/gh/redis/redis-py/pull/2281/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-dGVzdHMvY29uZnRlc3QucHk=) | `74.89% <58.33%> (-12.41%)` | :arrow_down: |
| [tests/test\_asyncio/compat.py](https://codecov.io/gh/redis/redis-py/pull/2281/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-dGVzdHMvdGVzdF9hc3luY2lvL2NvbXBhdC5weQ==) | `90.90% <83.33%> (-9.10%)` | :arrow_down: |
| [tests/test\_asyncio/test\_pubsub.py](https://codecov.io/gh/redis/redis-py/pull/2281/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-dGVzdHMvdGVzdF9hc3luY2lvL3Rlc3RfcHVic3ViLnB5) | `98.76% <97.46%> (-0.64%)` | :arrow_down: |
| [tests/test\_pubsub.py](https://codecov.io/gh/redis/redis-py/pull/2281/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-dGVzdHMvdGVzdF9wdWJzdWIucHk=) | `99.57% <98.73%> (-0.18%)` | :arrow_down: |
| [redis/asyncio/connection.py](https://codecov.io/gh/redis/redis-py/pull/2281/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-cmVkaXMvYXN5bmNpby9jb25uZWN0aW9uLnB5) | `76.39% <100.00%> (-7.68%)` | :arrow_down: |
| [redis/connection.py](https://codecov.io/gh/redis/redis-py/pull/2281/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-cmVkaXMvY29ubmVjdGlvbi5weQ==) | `77.73% <100.00%> (-8.70%)` | :arrow_down: |
| [tests/test\_cluster.py](https://codecov.io/gh/redis/redis-py/pull/2281/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-dGVzdHMvdGVzdF9jbHVzdGVyLnB5) | `14.66% <0.00%> (-82.43%)` | :arrow_down: |
| [tests/test\_asyncio/test\_cluster.py](https://codecov.io/gh/redis/redis-py/pull/2281/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-dGVzdHMvdGVzdF9hc3luY2lvL3Rlc3RfY2x1c3Rlci5weQ==) | `18.62% <0.00%> (-78.96%)` | :arrow_down: |
| [redis/asyncio/cluster.py](https://codecov.io/gh/redis/redis-py/pull/2281/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-cmVkaXMvYXN5bmNpby9jbHVzdGVyLnB5) | `16.83% <0.00%> (-73.13%)` | :arrow_down: |
| [redis/cluster.py](https://codecov.io/gh/redis/redis-py/pull/2281/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-cmVkaXMvY2x1c3Rlci5weQ==) | `18.71% <0.00%> (-71.71%)` | :arrow_down: |
| ... and [66 more](https://codecov.io/gh/redis/redis-py/pull/2281/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/redis/redis-py/pull/2281?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/redis/redis-py/pull/2281?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis). Last update [4b0543d...8664478](https://codecov.io/gh/redis/redis-py/pull/2281?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis).

dvora-h: @kristjanvalur Thanks, great PR!
dvora-h: @kristjanvalur Can you resolve the conflicts?

Patch:
diff --git a/redis/asyncio/client.py b/redis/asyncio/client.py
index 3d59016b..9c8caae0 100644
--- a/redis/asyncio/client.py
+++ b/redis/asyncio/client.py
@@ -754,9 +754,15 @@ class PubSub:
 
         await self.check_health()
 
-        if not block and not await self._execute(conn, conn.can_read, timeout=timeout):
-            return None
-        response = await self._execute(conn, conn.read_response)
+        async def try_read():
+            if not block:
+                if not await conn.can_read(timeout=timeout):
+                    return None
+            else:
+                await conn.connect()
+            return await conn.read_response()
+
+        response = await self._execute(conn, try_read)
 
         if conn.health_check_interval and response == self.health_check_response:
             # ignore the health check message as user might not expect it
diff --git a/redis/client.py b/redis/client.py
index 035edb16..39fbcdea 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -1497,9 +1497,15 @@ class PubSub:
 
         self.check_health()
 
-        if not block and not self._execute(conn, conn.can_read, timeout=timeout):
-            return None
-        response = self._execute(conn, conn.read_response)
+        def try_read():
+            if not block:
+                if not conn.can_read(timeout=timeout):
+                    return None
+            else:
+                conn.connect()
+            return conn.read_response()
+
+        response = self._execute(conn, try_read)
 
         if self.is_health_check_response(response):
             # ignore the health check message as user might not expect it


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 8a907cc3..3d403750 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -130,15 +130,25 @@ def _get_info(redis_url):
 
 
 def pytest_sessionstart(session):
+    # during test discovery, e.g. with VS Code, we may not
+    # have a server running.
     redis_url = session.config.getoption("--redis-url")
-    info = _get_info(redis_url)
-    version = info["redis_version"]
-    arch_bits = info["arch_bits"]
-    cluster_enabled = info["cluster_enabled"]
+    try:
+        info = _get_info(redis_url)
+        version = info["redis_version"]
+        arch_bits = info["arch_bits"]
+        cluster_enabled = info["cluster_enabled"]
+        enterprise = info["enterprise"]
+    except redis.ConnectionError:
+        # provide optimistic defaults
+        version = "10.0.0"
+        arch_bits = 64
+        cluster_enabled = False
+        enterprise = False
     REDIS_INFO["version"] = version
     REDIS_INFO["arch_bits"] = arch_bits
     REDIS_INFO["cluster_enabled"] = cluster_enabled
-    REDIS_INFO["enterprise"] = info["enterprise"]
+    REDIS_INFO["enterprise"] = enterprise
     # store REDIS_INFO in config so that it is available from "condition strings"
     session.config.REDIS_INFO = REDIS_INFO
 
diff --git a/tests/test_asyncio/compat.py b/tests/test_asyncio/compat.py
index ced49741..ed5c6804 100644
--- a/tests/test_asyncio/compat.py
+++ b/tests/test_asyncio/compat.py
@@ -1,6 +1,15 @@
+import asyncio
+import sys
 from unittest import mock
 
 try:
     mock.AsyncMock
 except AttributeError:
     import mock
+
+
+def create_task(coroutine):
+    if sys.version_info[:2] >= (3, 7):
+        return asyncio.create_task(coroutine)
+    else:
+        return asyncio.ensure_future(coroutine)
diff --git a/tests/test_asyncio/test_pubsub.py b/tests/test_asyncio/test_pubsub.py
index 0c602671..2ff0c726 100644
--- a/tests/test_asyncio/test_pubsub.py
+++ b/tests/test_asyncio/test_pubsub.py
@@ -1,5 +1,6 @@
 import asyncio
 import functools
+import socket
 from typing import Optional
 
 import async_timeout
@@ -11,7 +12,7 @@ from redis.exceptions import ConnectionError
 from redis.typing import EncodableT
 from tests.conftest import skip_if_server_version_lt
 
-from .compat import mock
+from .compat import create_task, mock
 
 
 def with_timeout(t):
@@ -786,3 +787,130 @@ class TestPubSubRun:
             "pattern": None,
             "type": "message",
         }
+
+
+# @pytest.mark.xfail
+@pytest.mark.parametrize("method", ["get_message", "listen"])
+@pytest.mark.onlynoncluster
+class TestPubSubAutoReconnect:
+    timeout = 2
+
+    async def mysetup(self, r, method):
+        self.messages = asyncio.Queue()
+        self.pubsub = r.pubsub()
+        # State: 0 = initial state , 1 = after disconnect, 2 = ConnectionError is seen,
+        # 3=successfully reconnected 4 = exit
+        self.state = 0
+        self.cond = asyncio.Condition()
+        if method == "get_message":
+            self.get_message = self.loop_step_get_message
+        else:
+            self.get_message = self.loop_step_listen
+
+        self.task = create_task(self.loop())
+        # get the initial connect message
+        message = await self.messages.get()
+        assert message == {
+            "channel": b"foo",
+            "data": 1,
+            "pattern": None,
+            "type": "subscribe",
+        }
+
+    async def mycleanup(self):
+        message = await self.messages.get()
+        assert message == {
+            "channel": b"foo",
+            "data": 1,
+            "pattern": None,
+            "type": "subscribe",
+        }
+        # kill thread
+        async with self.cond:
+            self.state = 4  # quit
+        await self.task
+
+    async def test_reconnect_socket_error(self, r: redis.Redis, method):
+        """
+        Test that a socket error will cause reconnect
+        """
+        async with async_timeout.timeout(self.timeout):
+            await self.mysetup(r, method)
+            # now, disconnect the connection, and wait for it to be re-established
+            async with self.cond:
+                assert self.state == 0
+                self.state = 1
+                with mock.patch.object(self.pubsub.connection, "_parser") as mockobj:
+                    mockobj.read_response.side_effect = socket.error
+                    mockobj.can_read.side_effect = socket.error
+                    # wait until task noticies the disconnect until we undo the patch
+                    await self.cond.wait_for(lambda: self.state >= 2)
+                    assert not self.pubsub.connection.is_connected
+                    # it is in a disconnecte state
+                # wait for reconnect
+                await self.cond.wait_for(lambda: self.pubsub.connection.is_connected)
+                assert self.state == 3
+
+            await self.mycleanup()
+
+    async def test_reconnect_disconnect(self, r: redis.Redis, method):
+        """
+        Test that a manual disconnect() will cause reconnect
+        """
+        async with async_timeout.timeout(self.timeout):
+            await self.mysetup(r, method)
+            # now, disconnect the connection, and wait for it to be re-established
+            async with self.cond:
+                self.state = 1
+                await self.pubsub.connection.disconnect()
+                assert not self.pubsub.connection.is_connected
+                # wait for reconnect
+                await self.cond.wait_for(lambda: self.pubsub.connection.is_connected)
+                assert self.state == 3
+
+            await self.mycleanup()
+
+    async def loop(self):
+        # reader loop, performing state transitions as it
+        # discovers disconnects and reconnects
+        await self.pubsub.subscribe("foo")
+        while True:
+            await asyncio.sleep(0.01)  # give main thread chance to get lock
+            async with self.cond:
+                old_state = self.state
+                try:
+                    if self.state == 4:
+                        break
+                    # print("state a ", self.state)
+                    got_msg = await self.get_message()
+                    assert got_msg
+                    if self.state in (1, 2):
+                        self.state = 3  # successful reconnect
+                except redis.ConnectionError:
+                    assert self.state in (1, 2)
+                    self.state = 2  # signal that we noticed the disconnect
+                finally:
+                    self.cond.notify()
+                # make sure that we did notice the connection error
+                # or reconnected without any error
+                if old_state == 1:
+                    assert self.state in (2, 3)
+
+    async def loop_step_get_message(self):
+        # get a single message via get_message
+        message = await self.pubsub.get_message(timeout=0.1)
+        # print(message)
+        if message is not None:
+            await self.messages.put(message)
+            return True
+        return False
+
+    async def loop_step_listen(self):
+        # get a single message via listen()
+        try:
+            async with async_timeout.timeout(0.1):
+                async for message in self.pubsub.listen():
+                    await self.messages.put(message)
+                    return True
+        except asyncio.TimeoutError:
+            return False
diff --git a/tests/test_pubsub.py b/tests/test_pubsub.py
index 6456370c..26df598b 100644
--- a/tests/test_pubsub.py
+++ b/tests/test_pubsub.py
@@ -1,4 +1,6 @@
 import platform
+import queue
+import socket
 import threading
 import time
 from unittest import mock
@@ -608,3 +610,128 @@ class TestPubSubDeadlock:
             p = r.pubsub()
             p.subscribe("my-channel-1", "my-channel-2")
             pool.reset()
+
+
+@pytest.mark.timeout(5, method="thread")
+@pytest.mark.parametrize("method", ["get_message", "listen"])
+@pytest.mark.onlynoncluster
+class TestPubSubAutoReconnect:
+    def mysetup(self, r, method):
+        self.messages = queue.Queue()
+        self.pubsub = r.pubsub()
+        self.state = 0
+        self.cond = threading.Condition()
+        if method == "get_message":
+            self.get_message = self.loop_step_get_message
+        else:
+            self.get_message = self.loop_step_listen
+
+        self.thread = threading.Thread(target=self.loop)
+        self.thread.daemon = True
+        self.thread.start()
+        # get the initial connect message
+        message = self.messages.get(timeout=1)
+        assert message == {
+            "channel": b"foo",
+            "data": 1,
+            "pattern": None,
+            "type": "subscribe",
+        }
+
+    def wait_for_reconnect(self):
+        self.cond.wait_for(lambda: self.pubsub.connection._sock is not None, timeout=2)
+        assert self.pubsub.connection._sock is not None  # we didn't time out
+        assert self.state == 3
+
+        message = self.messages.get(timeout=1)
+        assert message == {
+            "channel": b"foo",
+            "data": 1,
+            "pattern": None,
+            "type": "subscribe",
+        }
+
+    def mycleanup(self):
+        # kill thread
+        with self.cond:
+            self.state = 4  # quit
+            self.cond.notify()
+        self.thread.join()
+
+    def test_reconnect_socket_error(self, r: redis.Redis, method):
+        """
+        Test that a socket error will cause reconnect
+        """
+        self.mysetup(r, method)
+        try:
+            # now, disconnect the connection, and wait for it to be re-established
+            with self.cond:
+                self.state = 1
+                with mock.patch.object(self.pubsub.connection, "_parser") as mockobj:
+                    mockobj.read_response.side_effect = socket.error
+                    mockobj.can_read.side_effect = socket.error
+                    # wait until thread notices the disconnect until we undo the patch
+                    self.cond.wait_for(lambda: self.state >= 2)
+                    assert (
+                        self.pubsub.connection._sock is None
+                    )  # it is in a disconnected state
+                self.wait_for_reconnect()
+
+        finally:
+            self.mycleanup()
+
+    def test_reconnect_disconnect(self, r: redis.Redis, method):
+        """
+        Test that a manual disconnect() will cause reconnect
+        """
+        self.mysetup(r, method)
+        try:
+            # now, disconnect the connection, and wait for it to be re-established
+            with self.cond:
+                self.state = 1
+                self.pubsub.connection.disconnect()
+                assert self.pubsub.connection._sock is None
+                # wait for reconnect
+                self.wait_for_reconnect()
+        finally:
+            self.mycleanup()
+
+    def loop(self):
+        # reader loop, performing state transitions as it
+        # discovers disconnects and reconnects
+        self.pubsub.subscribe("foo")
+        while True:
+            time.sleep(0.01)  # give main thread chance to get lock
+            with self.cond:
+                old_state = self.state
+                try:
+                    if self.state == 4:
+                        break
+                    # print ('state, %s, sock %s' % (state, pubsub.connection._sock))
+                    got_msg = self.get_message()
+                    assert got_msg
+                    if self.state in (1, 2):
+                        self.state = 3  # successful reconnect
+                except redis.ConnectionError:
+                    assert self.state in (1, 2)
+                    self.state = 2
+                finally:
+                    self.cond.notify()
+                # assert that we noticed a connect error, or automatically
+                # reconnected without error
+                if old_state == 1:
+                    assert self.state in (2, 3)
+
+    def loop_step_get_message(self):
+        # get a single message via listen()
+        message = self.pubsub.get_message(timeout=0.1)
+        if message is not None:
+            self.messages.put(message)
+            return True
+        return False
+
+    def loop_step_listen(self):
+        # get a single message via listen()
+        for message in self.pubsub.listen():
+            self.messages.put(message)
+            return True