Add Redis readiness verification (#3555)

Author: ManelCoutinhoSensei

redis/asyncio/client.py

@@ -230,6 +230,7 @@ def __init__(
         encoding: str = "utf-8",
         encoding_errors: str = "strict",
         decode_responses: bool = False,
+        check_server_ready: bool = False,
         retry_on_timeout: bool = False,
         retry: Retry = Retry(
             backoff=ExponentialWithJitterBackoff(base=1, cap=10), retries=3
@@ -279,6 +280,10 @@ def __init__(
 
         When 'connection_pool' is provided - the retry configuration of the
         provided pool will be used.
+
+        Args:
+            check_server_ready: if `True`, an extra handshake is performed by sending a PING command, since
+                connect and send operations work even when Redis server is not ready.
         """
         kwargs: Dict[str, Any]
         if event_dispatcher is None:
@@ -313,6 +318,7 @@ def __init__(
                 "encoding": encoding,
                 "encoding_errors": encoding_errors,
                 "decode_responses": decode_responses,
+                "check_server_ready": check_server_ready,
                 "retry_on_error": retry_on_error,
                 "retry": copy.deepcopy(retry),
                 "max_connections": max_connections,

redis/asyncio/cluster.py

@@ -290,6 +290,7 @@ def __init__(
         encoding_errors: str = "strict",
         decode_responses: bool = False,
         # Connection related kwargs
+        check_server_ready: bool = False,
         health_check_interval: float = 0,
         socket_connect_timeout: Optional[float] = None,
         socket_keepalive: bool = False,
@@ -345,6 +346,7 @@ def __init__(
             "encoding_errors": encoding_errors,
             "decode_responses": decode_responses,
             # Connection related kwargs
+            "check_server_ready": check_server_ready,
             "health_check_interval": health_check_interval,
             "socket_connect_timeout": socket_connect_timeout,
             "socket_keepalive": socket_keepalive,

redis/asyncio/connection.py

@@ -149,6 +149,7 @@ def __init__(
         encoding_errors: str = "strict",
         decode_responses: bool = False,
         parser_class: Type[BaseParser] = DefaultParser,
+        check_server_ready: bool = False,
         socket_read_size: int = 65536,
         health_check_interval: float = 0,
         client_name: Optional[str] = None,
@@ -205,6 +206,7 @@ def __init__(
         self.health_check_interval = health_check_interval
         self.next_health_check: float = -1
         self.encoder = encoder_class(encoding, encoding_errors, decode_responses)
+        self.check_server_ready = check_server_ready
         self.redis_connect_func = redis_connect_func
         self._reader: Optional[asyncio.StreamReader] = None
         self._writer: Optional[asyncio.StreamWriter] = None
@@ -304,11 +306,13 @@ async def connect_check_health(
         try:
             if retry_socket_connect:
                 await self.retry.call_with_retry(
-                    lambda: self._connect(), lambda error: self.disconnect()
+                    lambda: self._connect_check_server_ready(),
+                    lambda error: self.disconnect(),
                 )
             else:
-                await self._connect()
+                await self._connect_check_server_ready()
         except asyncio.CancelledError:
+            self._close()
             raise  # in 3.7 and earlier, this is an Exception, not BaseException
         except (socket.timeout, asyncio.TimeoutError):
             raise TimeoutError("Timeout connecting to server")
@@ -343,6 +347,33 @@ async def connect_check_health(
             if task and inspect.isawaitable(task):
                 await task
 
+    async def _connect_check_server_ready(self):
+        await self._connect()
+
+        # Doing handshake since connect and send operations work even when Redis is not ready
+        if self.check_server_ready:
+            try:
+                await self.send_command("PING", check_health=False)
+
+                if self.socket_timeout is not None:
+                    async with async_timeout(self.socket_timeout):
+                        response = str_if_bytes(await self._reader.read(1024))
+                else:
+                    response = str_if_bytes(await self._reader.read(1024))
+
+                if not (response.startswith("+PONG") or response.startswith("-NOAUTH")):
+                    raise ResponseError(f"Invalid PING response: {response}")
+            except (
+                socket.timeout,
+                asyncio.TimeoutError,
+                ResponseError,
+                ConnectionResetError,
+            ) as e:
+                # `socket_keepalive_options` might contain invalid options
+                # causing an error. Do not leave the connection open.
+                self._close()
+                raise ConnectionError(self._error_message(e))
+
     @abstractmethod
     async def _connect(self):
         pass
@@ -532,8 +563,7 @@ async def send_packed_command(
                     self._send_packed_command(command), self.socket_timeout
                 )
             else:
-                self._writer.writelines(command)
-                await self._writer.drain()
+                await self._send_packed_command(command)
         except asyncio.TimeoutError:
             await self.disconnect(nowait=True)
             raise TimeoutError("Timeout writing to socket") from None
@@ -776,7 +806,7 @@ async def _connect(self):
             except (OSError, TypeError):
                 # `socket_keepalive_options` might contain invalid options
                 # causing an error. Do not leave the connection open.
-                writer.close()
+                self._close()
                 raise
 
     def _host_error(self) -> str:
@@ -961,7 +991,6 @@ async def _connect(self):
             reader, writer = await asyncio.open_unix_connection(path=self.path)
         self._reader = reader
         self._writer = writer
-        await self.on_connect()
 
     def _host_error(self) -> str:
         return self.path

redis/client.py

@@ -274,10 +274,9 @@ def __init__(
         provided pool will be used.
 
         Args:
-
-        single_connection_client:
-            if `True`, connection pool is not used. In that case `Redis`
-            instance use is not thread safe.
+            single_connection_client:
+                if `True`, connection pool is not used. In that case `Redis`
+                instance use is not thread safe.
         """
         if event_dispatcher is None:
             self._event_dispatcher = EventDispatcher()

redis/connection.py

@@ -575,17 +575,17 @@ def connect_check_health(
             return
         try:
             if retry_socket_connect:
-                sock = self.retry.call_with_retry(
-                    lambda: self._connect(), lambda error: self.disconnect(error)
+                self.retry.call_with_retry(
+                    lambda: self._connect(),
+                    lambda error: self.disconnect(error),
                 )
             else:
-                sock = self._connect()
+                self._connect()
         except socket.timeout:
             raise TimeoutError("Timeout connecting to server")
         except OSError as e:
             raise ConnectionError(self._error_message(e))
 
-        self._sock = sock
         try:
             if self.redis_connect_func is None:
                 # Use the default on_connect function
@@ -608,7 +608,7 @@ def connect_check_health(
                 callback(self)
 
     @abstractmethod
-    def _connect(self):
+    def _connect(self) -> None:
         pass
 
     @abstractmethod
@@ -626,6 +626,12 @@ def on_connect_check_health(self, check_health: bool = True):
         self._parser.on_connect(self)
         parser = self._parser
 
+        if check_health:
+            self.retry.call_with_retry(
+                lambda: self._send_ping(),
+                lambda error: self.disconnect(error),
+            )
+
         auth_args = None
         # if credential provider or username and/or password are set, authenticate
         if self.credential_provider or (self.username or self.password):
@@ -680,7 +686,7 @@ def on_connect_check_health(self, check_health: bool = True):
                 # update cluster exception classes
                 self._parser.EXCEPTION_CLASSES = parser.EXCEPTION_CLASSES
                 self._parser.on_connect(self)
-            self.send_command("HELLO", self.protocol, check_health=check_health)
+            self.send_command("HELLO", self.protocol, check_health=False)
             self.handshake_metadata = self.read_response()
             if (
                 self.handshake_metadata.get(b"proto") != self.protocol
@@ -711,7 +717,7 @@ def on_connect_check_health(self, check_health: bool = True):
                     "ON",
                     "moving-endpoint-type",
                     endpoint_type.value,
-                    check_health=check_health,
+                    check_health=False,
                 )
                 response = self.read_response()
                 if str_if_bytes(response) != "OK":
@@ -737,7 +743,7 @@ def on_connect_check_health(self, check_health: bool = True):
                 "CLIENT",
                 "SETNAME",
                 self.client_name,
-                check_health=check_health,
+                check_health=False,
             )
             if str_if_bytes(self.read_response()) != "OK":
                 raise ConnectionError("Error setting client name")
@@ -750,7 +756,7 @@ def on_connect_check_health(self, check_health: bool = True):
                     "SETINFO",
                     "LIB-NAME",
                     self.lib_name,
-                    check_health=check_health,
+                    check_health=False,
                 )
                 self.read_response()
         except ResponseError:
@@ -763,15 +769,15 @@ def on_connect_check_health(self, check_health: bool = True):
                     "SETINFO",
                     "LIB-VER",
                     self.lib_version,
-                    check_health=check_health,
+                    check_health=False,
                 )
                 self.read_response()
         except ResponseError:
             pass
 
         # if a database is specified, switch to it
         if self.db:
-            self.send_command("SELECT", self.db, check_health=check_health)
+            self.send_command("SELECT", self.db, check_health=False)
             if str_if_bytes(self.read_response()) != "OK":
                 raise ConnectionError("Invalid Database")
 
@@ -800,8 +806,16 @@ def disconnect(self, *args):
     def _send_ping(self):
         """Send PING, expect PONG in return"""
         self.send_command("PING", check_health=False)
-        if str_if_bytes(self.read_response()) != "PONG":
-            raise ConnectionError("Bad response from PING health check")
+        try:
+            # Do not disconnect on error here, since we want to keep the connection in case of AuthenticationError
+            # since we are raising ConnectionError in all other cases and ping_failed already disconnects,
+            # connection reload is already handled
+            if str_if_bytes(self.read_response(disconnect_on_error=False)) != "PONG":
+                raise ConnectionError("Bad response from PING health check")
+        except AuthenticationError:
+            # if we get an authentication error, the server is healthy
+            self._parser.on_disconnect()
+            self._parser.on_connect(self)
 
     def _ping_failed(self, error):
         """Function to call when PING fails"""
@@ -1097,7 +1111,7 @@ def repr_pieces(self):
             pieces.append(("client_name", self.client_name))
         return pieces
 
-    def _connect(self):
+    def _connect(self) -> None:
         "Create a TCP socket connection"
         # we want to mimic what socket.create_connection does to support
         # ipv4/ipv6, but we want to set options prior to calling
@@ -1128,7 +1142,8 @@ def _connect(self):
 
                 # set the socket_timeout now that we're connected
                 sock.settimeout(self.socket_timeout)
-                return sock
+                self._sock = sock
+                return
 
             except OSError as _:
                 err = _
@@ -1448,15 +1463,15 @@ def __init__(
         self.ssl_ciphers = ssl_ciphers
         super().__init__(**kwargs)
 
-    def _connect(self):
+    def _connect(self) -> None:
         """
         Wrap the socket with SSL support, handling potential errors.
         """
-        sock = super()._connect()
+        super()._connect()
         try:
-            return self._wrap_socket_with_ssl(sock)
+            self._sock = self._wrap_socket_with_ssl(self._sock)
         except (OSError, RedisError):
-            sock.close()
+            self._sock.close()
             raise
 
     def _wrap_socket_with_ssl(self, sock):
@@ -1559,7 +1574,7 @@ def repr_pieces(self):
             pieces.append(("client_name", self.client_name))
         return pieces
 
-    def _connect(self):
+    def _connect(self) -> None:
         "Create a Unix domain socket connection"
         sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
         sock.settimeout(self.socket_connect_timeout)
@@ -1574,7 +1589,7 @@ def _connect(self):
             sock.close()
             raise
         sock.settimeout(self.socket_timeout)
-        return sock
+        self._sock = sock
 
     def _host_error(self):
         return self.path

tests/test_asyncio/test_cluster.py

@@ -729,7 +729,7 @@ async def test_reading_with_load_balancing_strategies(
             Connection,
             send_command=mock.DEFAULT,
             read_response=mock.DEFAULT,
-            _connect=mock.DEFAULT,
+            _connect_check_server_ready=mock.DEFAULT,
             can_read_destructive=mock.DEFAULT,
             on_connect=mock.DEFAULT,
         ) as mocks:
@@ -761,7 +761,7 @@ def execute_command_mock_third(self, *args, **options):
                 execute_command.side_effect = execute_command_mock_first
                 mocks["send_command"].return_value = True
                 mocks["read_response"].return_value = "OK"
-                mocks["_connect"].return_value = True
+                mocks["_connect_check_server_ready"].return_value = True
                 mocks["can_read_destructive"].return_value = False
                 mocks["on_connect"].return_value = True
 
@@ -3117,13 +3117,19 @@ async def execute_command(self, *args, **kwargs):
 
         return _create_client
 
+    @pytest.mark.parametrize("check_server_ready", [True, False])
     async def test_ssl_connection_without_ssl(
-        self, create_client: Callable[..., Awaitable[RedisCluster]]
+        self, create_client: Callable[..., Awaitable[RedisCluster]], check_server_ready
     ) -> None:
         with pytest.raises(RedisClusterException) as e:
-            await create_client(mocked=False, ssl=False)
+            await create_client(
+                mocked=False, ssl=False, check_server_ready=check_server_ready
+            )
         e = e.value.__cause__
-        assert "Connection closed by server" in str(e)
+        if check_server_ready:
+            assert "Invalid PING response" in str(e)
+        else:
+            assert "Connection closed by server" in str(e)
 
     async def test_ssl_with_invalid_cert(
         self, create_client: Callable[..., Awaitable[RedisCluster]]