Refactored health check logic for MultiDBClient

docs/geographic_failover.rst

@@ -193,6 +193,7 @@ MultiDbConfig
         health_check_interval: float = DEFAULT_HEALTH_CHECK_INTERVAL # seconds
         health_check_probes: int = DEFAULT_HEALTH_CHECK_PROBES
         health_check_delay: float = DEFAULT_HEALTH_CHECK_DELAY # seconds
+        health_check_timeout: float = DEFAULT_HEALTH_CHECK_TIMEOUT # seconds
         health_check_policy: HealthCheckPolicies = DEFAULT_HEALTH_CHECK_POLICY,
 
         # Failure detector
@@ -336,26 +337,31 @@ reverse proxy behind an actual REST API endpoint.
 
 **Custom Health Checks**
 ~~~~~~~~~~~~~~~~~~~~~
-You can add custom health checks for specific requirements:
+You can add custom health checks for specific requirements. Please notice that all health checks are executed within
+asyncio event loop, so please ensure that `check_health` method is async:
 
 .. code-block:: python
 
-    from redis.multidb.healthcheck import AbstractHealthCheck
-    from redis.retry import Retry
-    from redis.utils import dummy_fail
+    from redis.asyncio.multidb.healthcheck import AbstractHealthCheck
+    from redis.asyncio import Connection
+
     class PingHealthCheck(AbstractHealthCheck):
-        def __init__(self, retry: Retry):
-            super().__init__(retry=retry)
-        def check_health(self, database) -> bool:
-            return self._retry.call_with_retry(
-                lambda: self._returns_pong(database),
-                lambda _: dummy_fail()
+        def __init__(
+            self,
+            health_check_probes: int = 10,
+            health_check_delay: float = 0.1,
+            health_check_timeout: float = 2,
+        ):
+            super().__init__(
+                health_check_probes=health_check_probes,
+                health_check_delay=health_check_delay,
+                health_check_timeout=health_check_timeout,
             )
-        def _returns_pong(self, database) -> bool:
-            expected_message = ["PONG", b"PONG"]
-            actual_message = database.client.execute_command("PING")
-            return actual_message in expected_message
 
+        async def check_health(self, database, connection: Connection) -> bool:
+            await connection.send_command("PING")
+            response = await connection.read_response()
+            return response in (b"PONG", "PONG")
 
 Failure Detection (Reactive Monitoring)
 -----------------

pyproject.toml

@@ -86,6 +86,7 @@ markers = [
     "replica: replica tests",
     "experimental: run only experimental tests",
     "cp_integration: credential provider integration tests",
+    "no_mock_connections: skip the autouse mock_health_check_connections fixture",
 ]
 asyncio_default_fixture_loop_scope = "function"
 asyncio_mode = "auto"

redis/asyncio/multidb/client.py

@@ -45,10 +45,9 @@ def __init__(self, config: MultiDbConfig):
             if not config.health_checks
             else config.health_checks
         )
-
         self._health_check_interval = config.health_check_interval
-        self._health_check_policy: HealthCheckPolicy = config.health_check_policy.value(
-            config.health_check_probes, config.health_check_delay
+        self._health_check_policy: HealthCheckPolicy = (
+            config.health_check_policy.value()
         )
         self._failure_detectors = (
             config.default_failure_detectors()
@@ -89,14 +88,25 @@ async def __aenter__(self: "MultiDBClient") -> "MultiDBClient":
             await self.initialize()
         return self
 
-    async def __aexit__(self, exc_type, exc_value, traceback):
+    async def aclose(self):
+        # Cancel background tasks
         if self._recurring_hc_task:
             self._recurring_hc_task.cancel()
         if self._half_open_state_task:
             self._half_open_state_task.cancel()
         for hc_task in self._hc_tasks:
             hc_task.cancel()
 
+        # Close health check connection pools
+        await self._health_check_policy.close()
+
+        # Close database client
+        if self.command_executor.active_database:
+            await self.command_executor.active_database.client.aclose()
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self.aclose()
+
     async def initialize(self):
         """
         Perform initialization of databases to define their initial state.
@@ -326,23 +336,15 @@ async def _check_databases_health(self) -> dict[Database, bool]:
         Runs health checks as a recurring task.
         Runs health checks against all databases.
         """
-        try:
-            task_to_db: dict[asyncio.Task, Database] = {}
+        task_to_db: dict[asyncio.Task, Database] = {}
 
-            self._hc_tasks = []
-            for database, _ in self._databases:
-                task = asyncio.create_task(self._check_db_health(database))
-                task_to_db[task] = database
-                self._hc_tasks.append(task)
+        self._hc_tasks = []
+        for database, _ in self._databases:
+            task = asyncio.create_task(self._check_db_health(database))
+            task_to_db[task] = database
+            self._hc_tasks.append(task)
 
-            results = await asyncio.wait_for(
-                asyncio.gather(*self._hc_tasks, return_exceptions=True),
-                timeout=self._health_check_interval,
-            )
-        except asyncio.TimeoutError:
-            raise asyncio.TimeoutError(
-                "Health check execution exceeds health_check_interval"
-            )
+        results = await asyncio.gather(*self._hc_tasks, return_exceptions=True)
 
         # Map end results to databases
         db_results = {
@@ -360,8 +362,6 @@ async def _check_databases_health(self) -> dict[Database, bool]:
                 )
 
                 db_results[unhealthy_db] = False
-            elif isinstance(result, Exception):
-                db_results[database] = False
 
         return db_results
 
@@ -427,10 +427,6 @@ def _on_circuit_state_change_callback(
         if old_state != CBState.CLOSED and new_state == CBState.CLOSED:
             logger.info(f"Database {circuit.database} is reachable again.")
 
-    async def aclose(self):
-        if self.command_executor.active_database:
-            await self.command_executor.active_database.client.aclose()
-
 
 def _half_open_circuit(circuit: CircuitBreaker):
     circuit.state = CBState.HALF_OPEN

redis/asyncio/multidb/config.py

@@ -21,6 +21,7 @@
     DEFAULT_HEALTH_CHECK_INTERVAL,
     DEFAULT_HEALTH_CHECK_POLICY,
     DEFAULT_HEALTH_CHECK_PROBES,
+    DEFAULT_HEALTH_CHECK_TIMEOUT,
     HealthCheck,
     HealthCheckPolicies,
     PingHealthCheck,
@@ -110,6 +111,8 @@ class MultiDbConfig:
         health_check_interval: Time interval for executing health checks.
         health_check_probes: Number of attempts to evaluate the health of a database.
         health_check_delay: Delay between health check attempts.
+        health_check_timeout: Timeout for the full health check operation (including all probes).
+        health_check_policy: Policy for determining database health based on health checks.
         failover_strategy: Optional strategy for handling database failover scenarios.
         failover_attempts: Number of retries allowed for failover operations.
         failover_delay: Delay between failover attempts.
@@ -149,6 +152,7 @@ class MultiDbConfig:
     health_check_interval: float = DEFAULT_HEALTH_CHECK_INTERVAL
     health_check_probes: int = DEFAULT_HEALTH_CHECK_PROBES
     health_check_delay: float = DEFAULT_HEALTH_CHECK_DELAY
+    health_check_timeout: float = DEFAULT_HEALTH_CHECK_TIMEOUT
     health_check_policy: HealthCheckPolicies = DEFAULT_HEALTH_CHECK_POLICY
     failover_strategy: Optional[AsyncFailoverStrategy] = None
     failover_attempts: int = DEFAULT_FAILOVER_ATTEMPTS
@@ -213,7 +217,11 @@ def default_failure_detectors(self) -> List[AsyncFailureDetector]:
 
     def default_health_checks(self) -> List[HealthCheck]:
         return [
-            PingHealthCheck(),
+            PingHealthCheck(
+                health_check_probes=self.health_check_probes,
+                health_check_delay=self.health_check_delay,
+                health_check_timeout=self.health_check_timeout,
+            ),
         ]
 
     def default_failover_strategy(self) -> AsyncFailoverStrategy: