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
@@ -344,19 +345,23 @@ You can add custom health checks for specific requirements:
     from redis.retry import Retry
     from redis.utils import dummy_fail
     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:
+
+        def check_health(self, database) -> bool:
             expected_message = ["PONG", b"PONG"]
             actual_message = database.client.execute_command("PING")
             return actual_message in expected_message
 
-
 Failure Detection (Reactive Monitoring)
 -----------------
 

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()
@@ -326,23 +325,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 = {

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: