Repository: redis/redis-py
Pull request #2771 (2023-05-21 12:29:33)

Fix dead weakref in sentinel connection causing ReferenceError (#2767)
### Pull Request check-list

- [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?
- [x] Is a documentation update included (if this change modifies existing APIs, or introduces new ones)?
- [x] Is there an example added to the examples folder (if applicable)?
- [x] Was the change added to CHANGES file?

### Description of change

Instead of passing the sentinel connection a `weakref.proxy` to the sentinel connection pool, I pass it a wrapper object (of type `SentinelConnectionPoolProxy`) that holds a `weakref.ref` to the sentinel connection pool (and checks its liveness before invocation).
Please see related issue #2767

codecov-commenter: ## [Codecov](https://app.codecov.io/gh/redis/redis-py/pull/2771?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) Report
Patch coverage: **`97.91`**% and project coverage change: **`-14.34`** :warning:
> Comparison is base [(`2d9b5ac`)](https://app.codecov.io/gh/redis/redis-py/commit/2d9b5ac6fe03fdc572b8ca47f7134082bae2a5e2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) 92.30% compared to head [(`7efd4a7`)](https://app.codecov.io/gh/redis/redis-py/pull/2771?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) 77.96%.

:exclamation: Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please [install the Github App Integration](https://github.com/apps/codecov) for your organization. [Read more](https://about.codecov.io/blog/codecov-is-updating-its-github-integration/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis).

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           master    #2771       +/-   ##
===========================================
- Coverage   92.30%   77.96%   -14.34%     
===========================================
  Files         116      116               
  Lines       30214    30215        +1     
===========================================
- Hits        27888    23557     -4331     
- Misses       2326     6658     +4332     
```


| [Impacted Files](https://app.codecov.io/gh/redis/redis-py/pull/2771?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) | Coverage Δ | |
|---|---|---|
| [redis/sentinel.py](https://app.codecov.io/gh/redis/redis-py/pull/2771?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-cmVkaXMvc2VudGluZWwucHk=) | `86.95% <97.56%> (+0.89%)` | :arrow_up: |
| [tests/test\_sentinel.py](https://app.codecov.io/gh/redis/redis-py/pull/2771?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-dGVzdHMvdGVzdF9zZW50aW5lbC5weQ==) | `99.34% <100.00%> (+0.03%)` | :arrow_up: |

... and [46 files with indirect coverage changes](https://app.codecov.io/gh/redis/redis-py/pull/2771/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/redis/redis-py/pull/2771?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis).

chayim: @shahar-lev kindly resolve the conflict, and we can include it in the next release
shahar-lev: @chayim pulled from master.
shahar-lev: @chayim do you need anything else from me?
chayim: Nope. We're good. It's just not yet merged. Will happen shortly. 

Patch:
diff --git a/CHANGES b/CHANGES
index 76b3393a..040e4452 100644
--- a/CHANGES
+++ b/CHANGES
@@ -1,3 +1,4 @@
+    * Fix dead weakref in sentinel connection causing ReferenceError (#2767)
     * Fix #2768, Fix KeyError: 'first-entry' in parse_xinfo_stream.
     * Fix #2749, remove unnecessary __del__ logic to close connections.
     * Fix #2754, adding a missing argument to SentinelManagedConnection
diff --git a/redis/sentinel.py b/redis/sentinel.py
index f9f8f1c3..0ba179b9 100644
--- a/redis/sentinel.py
+++ b/redis/sentinel.py
@@ -78,6 +78,54 @@ class SentinelManagedSSLConnection(SentinelManagedConnection, SSLConnection):
     pass
 
 
+class SentinelConnectionPoolProxy:
+    def __init__(
+        self,
+        connection_pool,
+        is_master,
+        check_connection,
+        service_name,
+        sentinel_manager,
+    ):
+        self.connection_pool_ref = weakref.ref(connection_pool)
+        self.is_master = is_master
+        self.check_connection = check_connection
+        self.service_name = service_name
+        self.sentinel_manager = sentinel_manager
+        self.reset()
+
+    def reset(self):
+        self.master_address = None
+        self.slave_rr_counter = None
+
+    def get_master_address(self):
+        master_address = self.sentinel_manager.discover_master(self.service_name)
+        if self.is_master and self.master_address != master_address:
+            self.master_address = master_address
+            # disconnect any idle connections so that they reconnect
+            # to the new master the next time that they are used.
+            connection_pool = self.connection_pool_ref()
+            if connection_pool is not None:
+                connection_pool.disconnect(inuse_connections=False)
+        return master_address
+
+    def rotate_slaves(self):
+        slaves = self.sentinel_manager.discover_slaves(self.service_name)
+        if slaves:
+            if self.slave_rr_counter is None:
+                self.slave_rr_counter = random.randint(0, len(slaves) - 1)
+            for _ in range(len(slaves)):
+                self.slave_rr_counter = (self.slave_rr_counter + 1) % len(slaves)
+                slave = slaves[self.slave_rr_counter]
+                yield slave
+        # Fallback to the master connection
+        try:
+            yield self.get_master_address()
+        except MasterNotFoundError:
+            pass
+        raise SlaveNotFoundError(f"No slave found for {self.service_name!r}")
+
+
 class SentinelConnectionPool(ConnectionPool):
     """
     Sentinel backed connection pool.
@@ -95,8 +143,15 @@ class SentinelConnectionPool(ConnectionPool):
         )
         self.is_master = kwargs.pop("is_master", True)
         self.check_connection = kwargs.pop("check_connection", False)
+        self.proxy = SentinelConnectionPoolProxy(
+            connection_pool=self,
+            is_master=self.is_master,
+            check_connection=self.check_connection,
+            service_name=service_name,
+            sentinel_manager=sentinel_manager,
+        )
         super().__init__(**kwargs)
-        self.connection_kwargs["connection_pool"] = weakref.proxy(self)
+        self.connection_kwargs["connection_pool"] = self.proxy
         self.service_name = service_name
         self.sentinel_manager = sentinel_manager
 
@@ -106,8 +161,11 @@ class SentinelConnectionPool(ConnectionPool):
 
     def reset(self):
         super().reset()
-        self.master_address = None
-        self.slave_rr_counter = None
+        self.proxy.reset()
+
+    @property
+    def master_address(self):
+        return self.proxy.master_address
 
     def owns_connection(self, connection):
         check = not self.is_master or (
@@ -117,31 +175,11 @@ class SentinelConnectionPool(ConnectionPool):
         return check and parent.owns_connection(connection)
 
     def get_master_address(self):
-        master_address = self.sentinel_manager.discover_master(self.service_name)
-        if self.is_master:
-            if self.master_address != master_address:
-                self.master_address = master_address
-                # disconnect any idle connections so that they reconnect
-                # to the new master the next time that they are used.
-                self.disconnect(inuse_connections=False)
-        return master_address
+        return self.proxy.get_master_address()
 
     def rotate_slaves(self):
         "Round-robin slave balancer"
-        slaves = self.sentinel_manager.discover_slaves(self.service_name)
-        if slaves:
-            if self.slave_rr_counter is None:
-                self.slave_rr_counter = random.randint(0, len(slaves) - 1)
-            for _ in range(len(slaves)):
-                self.slave_rr_counter = (self.slave_rr_counter + 1) % len(slaves)
-                slave = slaves[self.slave_rr_counter]
-                yield slave
-        # Fallback to the master connection
-        try:
-            yield self.get_master_address()
-        except MasterNotFoundError:
-            pass
-        raise SlaveNotFoundError(f"No slave found for {self.service_name!r}")
+        return self.proxy.rotate_slaves()
 
 
 class Sentinel(SentinelCommands):


Test patch:
diff --git a/tests/ssl_utils.py b/tests/ssl_utils.py
new file mode 100644
index 00000000..50937638
--- /dev/null
+++ b/tests/ssl_utils.py
@@ -0,0 +1,14 @@
+import os
+
+
+def get_ssl_filename(name):
+    root = os.path.join(os.path.dirname(__file__), "..")
+    cert_dir = os.path.abspath(os.path.join(root, "docker", "stunnel", "keys"))
+    if not os.path.isdir(cert_dir):  # github actions package validation case
+        cert_dir = os.path.abspath(
+            os.path.join(root, "..", "docker", "stunnel", "keys")
+        )
+        if not os.path.isdir(cert_dir):
+            raise IOError(f"No SSL certificates found. They should be in {cert_dir}")
+
+    return os.path.join(cert_dir, name)
diff --git a/tests/test_asyncio/test_cluster.py b/tests/test_asyncio/test_cluster.py
index 17aa879b..c41d4a21 100644
--- a/tests/test_asyncio/test_cluster.py
+++ b/tests/test_asyncio/test_cluster.py
@@ -1,7 +1,6 @@
 import asyncio
 import binascii
 import datetime
-import os
 import warnings
 from typing import Any, Awaitable, Callable, Dict, List, Optional, Type, Union
 from urllib.parse import urlparse
@@ -36,6 +35,7 @@ from tests.conftest import (
     skip_unless_arch_bits,
 )
 
+from ..ssl_utils import get_ssl_filename
 from .compat import mock
 
 pytestmark = pytest.mark.onlycluster
@@ -2744,17 +2744,8 @@ class TestSSL:
     appropriate port.
     """
 
-    ROOT = os.path.join(os.path.dirname(__file__), "../..")
-    CERT_DIR = os.path.abspath(os.path.join(ROOT, "docker", "stunnel", "keys"))
-    if not os.path.isdir(CERT_DIR):  # github actions package validation case
-        CERT_DIR = os.path.abspath(
-            os.path.join(ROOT, "..", "docker", "stunnel", "keys")
-        )
-        if not os.path.isdir(CERT_DIR):
-            raise IOError(f"No SSL certificates found. They should be in {CERT_DIR}")
-
-    SERVER_CERT = os.path.join(CERT_DIR, "server-cert.pem")
-    SERVER_KEY = os.path.join(CERT_DIR, "server-key.pem")
+    SERVER_CERT = get_ssl_filename("server-cert.pem")
+    SERVER_KEY = get_ssl_filename("server-key.pem")
 
     @pytest_asyncio.fixture()
     def create_client(self, request: FixtureRequest) -> Callable[..., RedisCluster]:
diff --git a/tests/test_asyncio/test_connect.py b/tests/test_asyncio/test_connect.py
new file mode 100644
index 00000000..8e3209fd
--- /dev/null
+++ b/tests/test_asyncio/test_connect.py
@@ -0,0 +1,145 @@
+import asyncio
+import logging
+import re
+import socket
+import ssl
+
+import pytest
+
+from redis.asyncio.connection import (
+    Connection,
+    SSLConnection,
+    UnixDomainSocketConnection,
+)
+
+from ..ssl_utils import get_ssl_filename
+
+_logger = logging.getLogger(__name__)
+
+
+_CLIENT_NAME = "test-suite-client"
+_CMD_SEP = b"\r\n"
+_SUCCESS_RESP = b"+OK" + _CMD_SEP
+_ERROR_RESP = b"-ERR" + _CMD_SEP
+_SUPPORTED_CMDS = {f"CLIENT SETNAME {_CLIENT_NAME}": _SUCCESS_RESP}
+
+
+@pytest.fixture
+def tcp_address():
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()
+
+
+@pytest.fixture
+def uds_address(tmpdir):
+    return tmpdir / "uds.sock"
+
+
+async def test_tcp_connect(tcp_address):
+    host, port = tcp_address
+    conn = Connection(host=host, port=port, client_name=_CLIENT_NAME, socket_timeout=10)
+    await _assert_connect(conn, tcp_address)
+
+
+async def test_uds_connect(uds_address):
+    path = str(uds_address)
+    conn = UnixDomainSocketConnection(
+        path=path, client_name=_CLIENT_NAME, socket_timeout=10
+    )
+    await _assert_connect(conn, path)
+
+
+@pytest.mark.ssl
+async def test_tcp_ssl_connect(tcp_address):
+    host, port = tcp_address
+    certfile = get_ssl_filename("server-cert.pem")
+    keyfile = get_ssl_filename("server-key.pem")
+    conn = SSLConnection(
+        host=host,
+        port=port,
+        client_name=_CLIENT_NAME,
+        ssl_ca_certs=certfile,
+        socket_timeout=10,
+    )
+    await _assert_connect(conn, tcp_address, certfile=certfile, keyfile=keyfile)
+
+
+async def _assert_connect(conn, server_address, certfile=None, keyfile=None):
+    stop_event = asyncio.Event()
+    finished = asyncio.Event()
+
+    async def _handler(reader, writer):
+        try:
+            return await _redis_request_handler(reader, writer, stop_event)
+        finally:
+            finished.set()
+
+    if isinstance(server_address, str):
+        server = await asyncio.start_unix_server(_handler, path=server_address)
+    elif certfile:
+        host, port = server_address
+        context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+        context.minimum_version = ssl.TLSVersion.TLSv1_2
+        context.load_cert_chain(certfile=certfile, keyfile=keyfile)
+        server = await asyncio.start_server(_handler, host=host, port=port, ssl=context)
+    else:
+        host, port = server_address
+        server = await asyncio.start_server(_handler, host=host, port=port)
+
+    async with server as aserver:
+        await aserver.start_serving()
+        try:
+            await conn.connect()
+            await conn.disconnect()
+        finally:
+            stop_event.set()
+            aserver.close()
+            await aserver.wait_closed()
+            await finished.wait()
+
+
+async def _redis_request_handler(reader, writer, stop_event):
+    buffer = b""
+    command = None
+    command_ptr = None
+    fragment_length = None
+    while not stop_event.is_set() or buffer:
+        _logger.info(str(stop_event.is_set()))
+        try:
+            buffer += await asyncio.wait_for(reader.read(1024), timeout=0.5)
+        except TimeoutError:
+            continue
+        if not buffer:
+            continue
+        parts = re.split(_CMD_SEP, buffer)
+        buffer = parts[-1]
+        for fragment in parts[:-1]:
+            fragment = fragment.decode()
+            _logger.info("Command fragment: %s", fragment)
+
+            if fragment.startswith("*") and command is None:
+                command = [None for _ in range(int(fragment[1:]))]
+                command_ptr = 0
+                fragment_length = None
+                continue
+
+            if fragment.startswith("$") and command[command_ptr] is None:
+                fragment_length = int(fragment[1:])
+                continue
+
+            assert len(fragment) == fragment_length
+            command[command_ptr] = fragment
+            command_ptr += 1
+
+            if command_ptr < len(command):
+                continue
+
+            command = " ".join(command)
+            _logger.info("Command %s", command)
+            resp = _SUPPORTED_CMDS.get(command, _ERROR_RESP)
+            _logger.info("Response from %s", resp)
+            writer.write(resp)
+            await writer.drain()
+            command = None
+    _logger.info("Exit handler")
diff --git a/tests/test_connect.py b/tests/test_connect.py
new file mode 100644
index 00000000..b4ec7020
--- /dev/null
+++ b/tests/test_connect.py
@@ -0,0 +1,185 @@
+import logging
+import re
+import socket
+import socketserver
+import ssl
+import threading
+
+import pytest
+
+from redis.connection import Connection, SSLConnection, UnixDomainSocketConnection
+
+from .ssl_utils import get_ssl_filename
+
+_logger = logging.getLogger(__name__)
+
+
+_CLIENT_NAME = "test-suite-client"
+_CMD_SEP = b"\r\n"
+_SUCCESS_RESP = b"+OK" + _CMD_SEP
+_ERROR_RESP = b"-ERR" + _CMD_SEP
+_SUPPORTED_CMDS = {f"CLIENT SETNAME {_CLIENT_NAME}": _SUCCESS_RESP}
+
+
+@pytest.fixture
+def tcp_address():
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()
+
+
+@pytest.fixture
+def uds_address(tmpdir):
+    return tmpdir / "uds.sock"
+
+
+def test_tcp_connect(tcp_address):
+    host, port = tcp_address
+    conn = Connection(host=host, port=port, client_name=_CLIENT_NAME, socket_timeout=10)
+    _assert_connect(conn, tcp_address)
+
+
+def test_uds_connect(uds_address):
+    path = str(uds_address)
+    conn = UnixDomainSocketConnection(path, client_name=_CLIENT_NAME, socket_timeout=10)
+    _assert_connect(conn, path)
+
+
+@pytest.mark.ssl
+def test_tcp_ssl_connect(tcp_address):
+    host, port = tcp_address
+    certfile = get_ssl_filename("server-cert.pem")
+    keyfile = get_ssl_filename("server-key.pem")
+    conn = SSLConnection(
+        host=host,
+        port=port,
+        client_name=_CLIENT_NAME,
+        ssl_ca_certs=certfile,
+        socket_timeout=10,
+    )
+    _assert_connect(conn, tcp_address, certfile=certfile, keyfile=keyfile)
+
+
+def _assert_connect(conn, server_address, certfile=None, keyfile=None):
+    if isinstance(server_address, str):
+        server = _RedisUDSServer(server_address, _RedisRequestHandler)
+    else:
+        server = _RedisTCPServer(
+            server_address, _RedisRequestHandler, certfile=certfile, keyfile=keyfile
+        )
+    with server as aserver:
+        t = threading.Thread(target=aserver.serve_forever)
+        t.start()
+        try:
+            aserver.wait_online()
+            conn.connect()
+            conn.disconnect()
+        finally:
+            aserver.stop()
+            t.join(timeout=5)
+
+
+class _RedisTCPServer(socketserver.TCPServer):
+    def __init__(self, *args, certfile=None, keyfile=None, **kw) -> None:
+        self._ready_event = threading.Event()
+        self._stop_requested = False
+        self._certfile = certfile
+        self._keyfile = keyfile
+        super().__init__(*args, **kw)
+
+    def service_actions(self):
+        self._ready_event.set()
+
+    def wait_online(self):
+        self._ready_event.wait()
+
+    def stop(self):
+        self._stop_requested = True
+        self.shutdown()
+
+    def is_serving(self):
+        return not self._stop_requested
+
+    def get_request(self):
+        if self._certfile is None:
+            return super().get_request()
+        newsocket, fromaddr = self.socket.accept()
+        connstream = ssl.wrap_socket(
+            newsocket,
+            server_side=True,
+            certfile=self._certfile,
+            keyfile=self._keyfile,
+            ssl_version=ssl.PROTOCOL_TLSv1_2,
+        )
+        return connstream, fromaddr
+
+
+class _RedisUDSServer(socketserver.UnixStreamServer):
+    def __init__(self, *args, **kw) -> None:
+        self._ready_event = threading.Event()
+        self._stop_requested = False
+        super().__init__(*args, **kw)
+
+    def service_actions(self):
+        self._ready_event.set()
+
+    def wait_online(self):
+        self._ready_event.wait()
+
+    def stop(self):
+        self._stop_requested = True
+        self.shutdown()
+
+    def is_serving(self):
+        return not self._stop_requested
+
+
+class _RedisRequestHandler(socketserver.StreamRequestHandler):
+    def setup(self):
+        _logger.info("%s connected", self.client_address)
+
+    def finish(self):
+        _logger.info("%s disconnected", self.client_address)
+
+    def handle(self):
+        buffer = b""
+        command = None
+        command_ptr = None
+        fragment_length = None
+        while self.server.is_serving() or buffer:
+            try:
+                buffer += self.request.recv(1024)
+            except socket.timeout:
+                continue
+            if not buffer:
+                continue
+            parts = re.split(_CMD_SEP, buffer)
+            buffer = parts[-1]
+            for fragment in parts[:-1]:
+                fragment = fragment.decode()
+                _logger.info("Command fragment: %s", fragment)
+
+                if fragment.startswith("*") and command is None:
+                    command = [None for _ in range(int(fragment[1:]))]
+                    command_ptr = 0
+                    fragment_length = None
+                    continue
+
+                if fragment.startswith("$") and command[command_ptr] is None:
+                    fragment_length = int(fragment[1:])
+                    continue
+
+                assert len(fragment) == fragment_length
+                command[command_ptr] = fragment
+                command_ptr += 1
+
+                if command_ptr < len(command):
+                    continue
+
+                command = " ".join(command)
+                _logger.info("Command %s", command)
+                resp = _SUPPORTED_CMDS.get(command, _ERROR_RESP)
+                _logger.info("Response %s", resp)
+                self.request.sendall(resp)
+                command = None
+        _logger.info("Exit handler")
diff --git a/tests/test_sentinel.py b/tests/test_sentinel.py
index 8542a0bf..e5e3d26f 100644
--- a/tests/test_sentinel.py
+++ b/tests/test_sentinel.py
@@ -98,6 +98,15 @@ def test_discover_master_error(sentinel):
         sentinel.discover_master("xxx")
 
 
+@pytest.mark.onlynoncluster
+def test_dead_pool(sentinel):
+    master = sentinel.master_for("mymaster", db=9)
+    conn = master.connection_pool.get_connection("_")
+    conn.disconnect()
+    del master
+    conn.connect()
+
+
 @pytest.mark.onlynoncluster
 def test_discover_master_sentinel_down(cluster, sentinel, master_ip):
     # Put first sentinel 'foo' down
diff --git a/tests/test_ssl.py b/tests/test_ssl.py
index ed38a316..c1a981d3 100644
--- a/tests/test_ssl.py
+++ b/tests/test_ssl.py
@@ -1,4 +1,3 @@
-import os
 import socket
 import ssl
 from urllib.parse import urlparse
@@ -9,6 +8,7 @@ import redis
 from redis.exceptions import ConnectionError, RedisError
 
 from .conftest import skip_if_cryptography, skip_if_nocryptography
+from .ssl_utils import get_ssl_filename
 
 
 @pytest.mark.ssl
@@ -19,17 +19,8 @@ class TestSSL:
     and connecting to the appropriate port.
     """
 
-    ROOT = os.path.join(os.path.dirname(__file__), "..")
-    CERT_DIR = os.path.abspath(os.path.join(ROOT, "docker", "stunnel", "keys"))
-    if not os.path.isdir(CERT_DIR):  # github actions package validation case
-        CERT_DIR = os.path.abspath(
-            os.path.join(ROOT, "..", "docker", "stunnel", "keys")
-        )
-        if not os.path.isdir(CERT_DIR):
-            raise IOError(f"No SSL certificates found. They should be in {CERT_DIR}")
-
-    SERVER_CERT = os.path.join(CERT_DIR, "server-cert.pem")
-    SERVER_KEY = os.path.join(CERT_DIR, "server-key.pem")
+    SERVER_CERT = get_ssl_filename("server-cert.pem")
+    SERVER_KEY = get_ssl_filename("server-key.pem")
 
     def test_ssl_with_invalid_cert(self, request):
         ssl_url = request.config.option.redis_ssl_url