Adding retries for the overall connect - socket connect + handshake. Fix for pubsub reconnect issues. (#3863)

* Adding retries for the overall connect - socket connect + handshake. Fix for pubsub reconnect issues.

* Update tests/test_connection.py

Co-authored-by: Copilot <[email protected]>

* Update tests/test_asyncio/test_connection.py

Co-authored-by: Copilot <[email protected]>

* Apply suggestions from code review

Co-authored-by: Copilot <[email protected]>

* Fix linters

---------

Co-authored-by: Copilot <[email protected]>

Author: petyaslavova

redis/asyncio/connection.py

@@ -296,7 +296,14 @@ def set_parser(self, parser_class: Type[BaseParser]) -> None:
 
     async def connect(self):
         """Connects to the Redis server if not already connected"""
-        await self.connect_check_health(check_health=True)
+        # try once the socket connect with the handshake, retry the whole
+        # connect/handshake flow based on retry policy
+        await self.retry.call_with_retry(
+            lambda: self.connect_check_health(
+                check_health=True, retry_socket_connect=False
+            ),
+            lambda error: self.disconnect(),
+        )
 
     async def connect_check_health(
         self, check_health: bool = True, retry_socket_connect: bool = True

redis/connection.py

@@ -843,7 +843,14 @@ def _get_parser(self) -> Union[_HiredisParser, _RESP3Parser, _RESP2Parser]:
 
     def connect(self):
         "Connects to the Redis server if not already connected"
-        self.connect_check_health(check_health=True)
+        # try once the socket connect with the handshake, retry the whole
+        # connect/handshake flow based on retry policy
+        self.retry.call_with_retry(
+            lambda: self.connect_check_health(
+                check_health=True, retry_socket_connect=False
+            ),
+            lambda error: self.disconnect(error),
+        )
 
     def connect_check_health(
         self, check_health: bool = True, retry_socket_connect: bool = True

tests/test_asyncio/test_connection.py

@@ -155,16 +155,33 @@ async def mock_connect():
     await conn.disconnect()
 
 
-async def test_connect_without_retry_on_os_error():
-    """Test that the _connect function is not being retried in case of a OSError"""
+async def test_connect_without_retry_on_non_retryable_error():
+    """
+    Test that the _connect function is not being retried in case of a CancelledError -
+    error that is not in the list of retry-able errors"""
     with patch.object(Connection, "_connect") as _connect:
-        _connect.side_effect = OSError("")
+        _connect.side_effect = asyncio.CancelledError("")
         conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 2))
-        with pytest.raises(ConnectionError):
+        with pytest.raises(asyncio.CancelledError):
             await conn.connect()
         assert _connect.call_count == 1
 
 
+async def test_connect_with_retries():
+    """
+    Test that retries occur for the entire connect+handshake flow when OSError happens during the handshake phase.
+    """
+    with patch.object(asyncio.StreamWriter, "writelines") as writelines:
+        writelines.side_effect = OSError(ECONNREFUSED)
+        conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 2))
+        with pytest.raises(ConnectionError):
+            await conn.connect()
+        # the handshake commands are the failing ones
+        # validate that we don't execute too many commands on each retry
+        # 3 retries --> 3 commands
+        assert writelines.call_count == 3
+
+
 async def test_connect_timeout_error_without_retry():
     """Test that the _connect function is not being retried if retry_on_timeout is
     set to False"""

tests/test_connection.py

@@ -124,16 +124,31 @@ def mock_connect():
         assert conn._connect.call_count == 3
         self.clear(conn)
 
-    def test_connect_without_retry_on_os_error(self):
-        """Test that the _connect function is not being retried in case of a OSError"""
+    def test_connect_without_retry_on_non_retryable_error(self):
+        """Test that the _connect function is not being retried in case of a non-retryable error"""
         with patch.object(Connection, "_connect") as _connect:
-            _connect.side_effect = OSError("")
+            _connect.side_effect = RedisError("")
             conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 2))
-            with pytest.raises(ConnectionError):
+            with pytest.raises(RedisError):
                 conn.connect()
             assert _connect.call_count == 1
             self.clear(conn)
 
+    def test_connect_with_retries(self):
+        """
+        Validate that retries occur for the entire connect+handshake flow when OSError
+        happens during the handshake phase.
+        """
+        with patch.object(socket.socket, "sendall") as sendall:
+            sendall.side_effect = OSError(ECONNREFUSED)
+            conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 2))
+            with pytest.raises(ConnectionError):
+                conn.connect()
+            # the handshake commands are the failing ones
+            # validate that we don't execute too many commands on each retry
+            # 3 retries --> 3 commands
+            assert sendall.call_count == 3
+
     def test_connect_timeout_error_without_retry(self):
         """Test that the _connect function is not being retried if retry_on_timeout is
         set to False"""