Added graceful connection closing, added graceful hc tasks termination

Author: vladvildanov

redis/asyncio/multidb/client.py

@@ -59,7 +59,8 @@ def __init__(self, config: MultiDbConfig):
         self._hc_lock = asyncio.Lock()
         self._bg_scheduler = BackgroundScheduler()
         self._config = config
-        self._hc_task = None
+        self._recurring_hc_task = None
+        self._hc_tasks = []
         self._half_open_state_task = None
 
     async def __aenter__(self: "MultiDBClient") -> "MultiDBClient":
@@ -68,10 +69,12 @@ async def __aenter__(self: "MultiDBClient") -> "MultiDBClient":
         return self
 
     async def __aexit__(self, exc_type, exc_value, traceback):
-        if self._hc_task:
-            self._hc_task.cancel()
+        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()
 
     async def initialize(self):
         """
@@ -84,7 +87,7 @@ async def raise_exception_on_failed_hc(error):
         await self._check_databases_health(on_error=raise_exception_on_failed_hc)
 
         # Starts recurring health checks on the background.
-        self._hc_task = asyncio.create_task(self._bg_scheduler.run_recurring_async(
+        self._recurring_hc_task = asyncio.create_task(self._bg_scheduler.run_recurring_async(
             self._health_check_interval,
             self._check_databases_health,
         ))
@@ -251,12 +254,10 @@ async def _check_databases_health(
         Runs health checks against all databases.
         """
         try:
+            self._hc_tasks = [asyncio.create_task(self._check_db_health(database)) for database, _ in self._databases]
             results = await asyncio.wait_for(
                 asyncio.gather(
-                    *(
-                        asyncio.create_task(self._check_db_health(database))
-                        for database, _ in self._databases
-                    ),
+                    *self._hc_tasks,
                     return_exceptions=True,
                 ),
                 timeout=self._health_check_interval,

redis/asyncio/multidb/command_executor.py

@@ -7,7 +7,7 @@
 from redis.asyncio.client import PubSub, Pipeline
 from redis.asyncio.multidb.database import Databases, AsyncDatabase, Database
 from redis.asyncio.multidb.event import AsyncActiveDatabaseChanged, RegisterCommandFailure, \
-    ResubscribeOnActiveDatabaseChanged
+    ResubscribeOnActiveDatabaseChanged, CloseConnectionOnActiveDatabaseChanged
 from redis.asyncio.multidb.failover import AsyncFailoverStrategy, FailoverStrategyExecutor, DefaultFailoverStrategyExecutor, \
     DEFAULT_FAILOVER_ATTEMPTS, DEFAULT_FAILOVER_DELAY
 from redis.asyncio.multidb.failure_detector import AsyncFailureDetector
@@ -286,7 +286,8 @@ def _setup_event_dispatcher(self):
         """
         failure_listener = RegisterCommandFailure(self._failure_detectors)
         resubscribe_listener = ResubscribeOnActiveDatabaseChanged()
+        close_connection_listener = CloseConnectionOnActiveDatabaseChanged()
         self._event_dispatcher.register_listeners({
             AsyncOnCommandsFailEvent: [failure_listener],
-            AsyncActiveDatabaseChanged: [resubscribe_listener],
+            AsyncActiveDatabaseChanged: [close_connection_listener, resubscribe_listener],
         })

redis/asyncio/multidb/event.py

@@ -1,5 +1,6 @@
 from typing import List
 
+from redis.asyncio import Redis
 from redis.asyncio.multidb.database import AsyncDatabase
 from redis.asyncio.multidb.failure_detector import AsyncFailureDetector
 from redis.event import AsyncEventListenerInterface, AsyncOnCommandsFailEvent
@@ -53,6 +54,16 @@ async def listen(self, event: AsyncActiveDatabaseChanged):
             event.command_executor.active_pubsub = new_pubsub
             await old_pubsub.aclose()
 
+class CloseConnectionOnActiveDatabaseChanged(AsyncEventListenerInterface):
+    """
+    Close connection to the old active database.
+    """
+    async def listen(self, event: AsyncActiveDatabaseChanged):
+        await event.old_database.client.aclose()
+
+        if isinstance(event.old_database.client, Redis):
+            await event.old_database.client.connection_pool.disconnect()
+
 class RegisterCommandFailure(AsyncEventListenerInterface):
     """
     Event listener that registers command failures and passing it to the failure detectors.

redis/multidb/command_executor.py

@@ -7,7 +7,8 @@
 from redis.multidb.config import DEFAULT_AUTO_FALLBACK_INTERVAL
 from redis.multidb.database import Database, Databases, SyncDatabase
 from redis.multidb.circuit import State as CBState
-from redis.multidb.event import RegisterCommandFailure, ActiveDatabaseChanged, ResubscribeOnActiveDatabaseChanged
+from redis.multidb.event import RegisterCommandFailure, ActiveDatabaseChanged, ResubscribeOnActiveDatabaseChanged, \
+    CloseConnectionOnActiveDatabaseChanged
 from redis.multidb.failover import FailoverStrategy, FailoverStrategyExecutor, DEFAULT_FAILOVER_ATTEMPTS, \
     DEFAULT_FAILOVER_DELAY, DefaultFailoverStrategyExecutor
 from redis.multidb.failure_detector import FailureDetector
@@ -303,7 +304,8 @@ def _setup_event_dispatcher(self):
         """
         failure_listener = RegisterCommandFailure(self._failure_detectors)
         resubscribe_listener = ResubscribeOnActiveDatabaseChanged()
+        close_connection_listener = CloseConnectionOnActiveDatabaseChanged()
         self._event_dispatcher.register_listeners({
             OnCommandsFailEvent: [failure_listener],
-            ActiveDatabaseChanged: [resubscribe_listener],
+            ActiveDatabaseChanged: [close_connection_listener, resubscribe_listener],
         })

redis/multidb/event.py

@@ -1,5 +1,8 @@
 from typing import List
 
+from redis.client import Redis
+from sphinx.events import EventListener
+
 from redis.event import EventListenerInterface, OnCommandsFailEvent
 from redis.multidb.database import SyncDatabase
 from redis.multidb.failure_detector import FailureDetector
@@ -53,6 +56,16 @@ def listen(self, event: ActiveDatabaseChanged):
             event.command_executor.active_pubsub = new_pubsub
             old_pubsub.close()
 
+class CloseConnectionOnActiveDatabaseChanged(EventListenerInterface):
+    """
+    Close connection to the old active database.
+    """
+    def listen(self, event: ActiveDatabaseChanged):
+        event.old_database.client.close()
+
+        if isinstance(event.old_database.client, Redis):
+            event.old_database.client.connection_pool.disconnect()
+
 class RegisterCommandFailure(EventListenerInterface):
     """
     Event listener that registers command failures and passing it to the failure detectors.