Add Redis readiness verification (redis#3555)

Author: ManelCoutinhoSensei

Diff for: redis/asyncio/client.py

@@ -224,6 +224,7 @@ def __init__(
         encoding: str = "utf-8",
         encoding_errors: str = "strict",
         decode_responses: bool = False,
+        check_ready: bool = False,
         retry_on_timeout: bool = False,
         retry_on_error: Optional[list] = None,
         ssl: bool = False,
@@ -291,6 +292,7 @@ def __init__(
                 "encoding": encoding,
                 "encoding_errors": encoding_errors,
                 "decode_responses": decode_responses,
+                "check_ready": check_ready,
                 "retry_on_timeout": retry_on_timeout,
                 "retry_on_error": retry_on_error,
                 "retry": copy.deepcopy(retry),

Diff for: redis/asyncio/cluster.py

@@ -258,6 +258,7 @@ def __init__(
         encoding_errors: str = "strict",
         decode_responses: bool = False,
         # Connection related kwargs
+        check_ready: bool = False,
         health_check_interval: float = 0,
         socket_connect_timeout: Optional[float] = None,
         socket_keepalive: bool = False,
@@ -313,6 +314,7 @@ def __init__(
             "encoding_errors": encoding_errors,
             "decode_responses": decode_responses,
             # Connection related kwargs
+            "check_ready": check_ready,
             "health_check_interval": health_check_interval,
             "socket_connect_timeout": socket_connect_timeout,
             "socket_keepalive": socket_keepalive,

Diff for: redis/asyncio/connection.py

@@ -148,6 +148,7 @@ def __init__(
         encoding_errors: str = "strict",
         decode_responses: bool = False,
         parser_class: Type[BaseParser] = DefaultParser,
+        check_ready: bool = False,
         socket_read_size: int = 65536,
         health_check_interval: float = 0,
         client_name: Optional[str] = None,
@@ -204,6 +205,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_ready = check_ready
         self.redis_connect_func = redis_connect_func
         self._reader: Optional[asyncio.StreamReader] = None
         self._writer: Optional[asyncio.StreamWriter] = None
@@ -295,14 +297,48 @@ async def connect(self):
         """Connects to the Redis server if not already connected"""
         await self.connect_check_health(check_health=True)
 
+    async def _connect_check_ready(self):
+        await self._connect()
+
+        # Doing handshake since connect and send operations work even when Redis is not ready
+        if self.check_ready:
+            try:
+                ping_cmd = self.pack_command("PING")
+                if self.socket_timeout:
+                    await asyncio.wait_for(
+                        self._send_packed_command(ping_cmd), self.socket_timeout
+                    )
+                else:
+                    await self._send_packed_command(ping_cmd)
+
+                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))
+
     async def connect_check_health(self, check_health: bool = True):
         if self.is_connected:
             return
         try:
             await self.retry.call_with_retry(
-                lambda: self._connect(), lambda error: self.disconnect()
+                lambda: self._connect_check_ready(), lambda error: self.disconnect()
             )
         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")
@@ -526,8 +562,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
@@ -774,7 +809,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:
@@ -933,7 +968,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

Diff for: redis/client.py

@@ -206,6 +206,7 @@ def __init__(
         charset: Optional[str] = None,
         errors: Optional[str] = None,
         decode_responses: bool = False,
+        check_ready: bool = False,
         retry_on_timeout: bool = False,
         retry_on_error: Optional[List[Type[Exception]]] = None,
         ssl: bool = False,
@@ -282,6 +283,7 @@ def __init__(
                 "encoding": encoding,
                 "encoding_errors": encoding_errors,
                 "decode_responses": decode_responses,
+                "check_ready": check_ready,
                 "retry_on_error": retry_on_error,
                 "retry": copy.deepcopy(retry),
                 "max_connections": max_connections,

Diff for: redis/connection.py

@@ -236,6 +236,7 @@ def __init__(
         encoding: str = "utf-8",
         encoding_errors: str = "strict",
         decode_responses: bool = False,
+        check_ready: bool = False,
         parser_class=DefaultParser,
         socket_read_size: int = 65536,
         health_check_interval: int = 0,
@@ -302,6 +303,7 @@ def __init__(
         self.redis_connect_func = redis_connect_func
         self.encoder = Encoder(encoding, encoding_errors, decode_responses)
         self.handshake_metadata = None
+        self.check_ready = check_ready
         self._sock = None
         self._socket_read_size = socket_read_size
         self.set_parser(parser_class)
@@ -378,12 +380,35 @@ def connect(self):
         "Connects to the Redis server if not already connected"
         self.connect_check_health(check_health=True)
 
+    def _connect_check_ready(self):
+        sock = self._connect()
+
+        # Doing handshake since connect and send operations work even when Redis is not ready
+        if self.check_ready:
+            try:
+                ping_parts = self._command_packer.pack("PING")
+                for part in ping_parts:
+                    sock.sendall(part)
+
+                response = str_if_bytes(sock.recv(1024))
+                if not (response.startswith("+PONG") or response.startswith("-NOAUTH")):
+                    raise ResponseError(f"Invalid PING response: {response}")
+            except (ConnectionResetError, ResponseError) as err:
+                try:
+                    sock.shutdown(socket.SHUT_RDWR)  # ensure a clean close
+                except OSError:
+                    pass
+                sock.close()
+                raise ConnectionError(self._error_message(err))
+        return sock
+
     def connect_check_health(self, check_health: bool = True):
         if self._sock:
             return
         try:
             sock = self.retry.call_with_retry(
-                lambda: self._connect(), lambda error: self.disconnect(error)
+                lambda: self._connect_check_ready(),
+                lambda error: self.disconnect(error),
             )
         except socket.timeout:
             raise TimeoutError("Timeout connecting to server")

Diff for: tests/test_asyncio/test_cluster.py

@@ -716,7 +716,7 @@ async def test_reading_with_load_balancing_strategies(
             Connection,
             send_command=mock.DEFAULT,
             read_response=mock.DEFAULT,
-            _connect=mock.DEFAULT,
+            _connect_check_ready=mock.DEFAULT,
             can_read_destructive=mock.DEFAULT,
             on_connect=mock.DEFAULT,
         ) as mocks:
@@ -748,7 +748,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_ready"].return_value = True
                 mocks["can_read_destructive"].return_value = False
                 mocks["on_connect"].return_value = True
 
@@ -3090,13 +3090,17 @@ async def execute_command(self, *args, **kwargs):
 
         return _create_client
 
+    @pytest.mark.parametrize("check_ready", [True, False])
     async def test_ssl_connection_without_ssl(
-        self, create_client: Callable[..., Awaitable[RedisCluster]]
+        self, create_client: Callable[..., Awaitable[RedisCluster]], check_ready
     ) -> None:
         with pytest.raises(RedisClusterException) as e:
-            await create_client(mocked=False, ssl=False)
+            await create_client(mocked=False, ssl=False, check_ready=check_ready)
         e = e.value.__cause__
-        assert "Connection closed by server" in str(e)
+        if check_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]]