Refactored bg scheduler

Author: vladvildanov

redis/asyncio/multidb/client.py

@@ -2,7 +2,7 @@
 import logging
 from typing import Any, Awaitable, Callable, Coroutine, List, Optional, Union
 
-from redis.asyncio.client import PSWorkerThreadExcHandlerT, PubSubHandler
+from redis.asyncio.client import PubSubHandler
 from redis.asyncio.multidb.command_executor import DefaultCommandExecutor
 from redis.asyncio.multidb.config import DEFAULT_GRACE_PERIOD, MultiDbConfig
 from redis.asyncio.multidb.database import AsyncDatabase, Databases
@@ -507,7 +507,7 @@ async def get_message(
     async def run(
         self,
         *,
-        exception_handler: Optional["PSWorkerThreadExcHandlerT"] = None,
+        exception_handler=None,
         poll_timeout: float = 1.0,
     ) -> None:
         """Process pub/sub messages using registered callbacks.
@@ -524,5 +524,5 @@ async def run(
             >>> await task
         """
         return await self._client.command_executor.execute_pubsub_run(
-            exception_handler=exception_handler, sleep_time=poll_timeout, pubsub=self
+            sleep_time=poll_timeout, exception_handler=exception_handler, pubsub=self
         )

redis/asyncio/multidb/command_executor.py

@@ -266,9 +266,15 @@ async def callback():
 
         return await self._execute_with_failure_detection(callback, *args)
 
-    async def execute_pubsub_run(self, sleep_time: float, **kwargs) -> Any:
+    async def execute_pubsub_run(
+        self, sleep_time: float, exception_handler=None, pubsub=None
+    ) -> Any:
         async def callback():
-            return await self._active_pubsub.run(poll_timeout=sleep_time, **kwargs)
+            return await self._active_pubsub.run(
+                poll_timeout=sleep_time,
+                exception_handler=exception_handler,
+                pubsub=pubsub,
+            )
 
         return await self._execute_with_failure_detection(callback)
 

redis/background.py

@@ -10,17 +10,47 @@ class BackgroundScheduler:
 
     def __init__(self):
         self._next_timer = None
+        self._event_loops = []
+        self._lock = threading.Lock()
+        self._stopped = False
 
     def __del__(self):
-        if self._next_timer:
-            self._next_timer.cancel()
+        self.stop()
+
+    def stop(self):
+        """
+        Stop all scheduled tasks and clean up resources.
+        """
+        with self._lock:
+            if self._stopped:
+                return
+            self._stopped = True
+
+            if self._next_timer:
+                self._next_timer.cancel()
+                self._next_timer = None
+
+            # Stop all event loops
+            for loop in self._event_loops:
+                if loop.is_running():
+                    loop.call_soon_threadsafe(loop.stop)
+
+            self._event_loops.clear()
 
     def run_once(self, delay: float, callback: Callable, *args):
         """
         Runs callable task once after certain delay in seconds.
         """
+        with self._lock:
+            if self._stopped:
+                return
+
         # Run loop in a separate thread to unblock main thread.
         loop = asyncio.new_event_loop()
+
+        with self._lock:
+            self._event_loops.append(loop)
+
         thread = threading.Thread(
             target=_start_event_loop_in_thread,
             args=(loop, self._call_later, delay, callback, *args),
@@ -32,9 +62,16 @@ def run_recurring(self, interval: float, callback: Callable, *args):
         """
         Runs recurring callable task with given interval in seconds.
         """
+        with self._lock:
+            if self._stopped:
+                return
+
         # Run loop in a separate thread to unblock main thread.
         loop = asyncio.new_event_loop()
 
+        with self._lock:
+            self._event_loops.append(loop)
+
         thread = threading.Thread(
             target=_start_event_loop_in_thread,
             args=(loop, self._call_later_recurring, interval, callback, *args),
@@ -49,10 +86,17 @@ async def run_recurring_async(
         Runs recurring coroutine with given interval in seconds in the current event loop.
         To be used only from an async context. No additional threads are created.
         """
+        with self._lock:
+            if self._stopped:
+                return
+
         loop = asyncio.get_running_loop()
         wrapped = _async_to_sync_wrapper(loop, coro, *args)
 
         def tick():
+            with self._lock:
+                if self._stopped:
+                    return
             # Schedule the coroutine
             wrapped()
             # Schedule next tick
@@ -64,6 +108,9 @@ def tick():
     def _call_later(
         self, loop: asyncio.AbstractEventLoop, delay: float, callback: Callable, *args
     ):
+        with self._lock:
+            if self._stopped:
+                return
         self._next_timer = loop.call_later(delay, callback, *args)
 
     def _call_later_recurring(
@@ -73,6 +120,9 @@ def _call_later_recurring(
         callback: Callable,
         *args,
     ):
+        with self._lock:
+            if self._stopped:
+                return
         self._call_later(
             loop, interval, self._execute_recurring, loop, interval, callback, *args
         )
@@ -87,7 +137,19 @@ def _execute_recurring(
         """
         Executes recurring callable task with given interval in seconds.
         """
-        callback(*args)
+        with self._lock:
+            if self._stopped:
+                return
+
+        try:
+            callback(*args)
+        except Exception:
+            # Silently ignore exceptions during shutdown
+            pass
+
+        with self._lock:
+            if self._stopped:
+                return
 
         self._call_later(
             loop, interval, self._execute_recurring, loop, interval, callback, *args
@@ -106,7 +168,22 @@ def _start_event_loop_in_thread(
     """
     asyncio.set_event_loop(event_loop)
     event_loop.call_soon(call_soon_cb, event_loop, *args)
-    event_loop.run_forever()
+    try:
+        event_loop.run_forever()
+    finally:
+        try:
+            # Clean up pending tasks
+            pending = asyncio.all_tasks(event_loop)
+            for task in pending:
+                task.cancel()
+            # Run loop once more to process cancellations
+            event_loop.run_until_complete(
+                asyncio.gather(*pending, return_exceptions=True)
+            )
+        except Exception:
+            pass
+        finally:
+            event_loop.close()
 
 
 def _async_to_sync_wrapper(loop, coro_func, *args, **kwargs):

redis/event.py

@@ -117,24 +117,22 @@ async def dispatch_async(self, event: object):
 
     def register_listeners(
         self,
-        event_listeners: Dict[
+        mappings: Dict[
             Type[object],
             List[Union[EventListenerInterface, AsyncEventListenerInterface]],
         ],
     ):
         with self._lock:
-            for event_type in event_listeners:
+            for event_type in mappings:
                 if event_type in self._event_listeners_mapping:
                     self._event_listeners_mapping[event_type] = list(
                         set(
                             self._event_listeners_mapping[event_type]
-                            + event_listeners[event_type]
+                            + mappings[event_type]
                         )
                     )
                 else:
-                    self._event_listeners_mapping[event_type] = event_listeners[
-                        event_type
-                    ]
+                    self._event_listeners_mapping[event_type] = mappings[event_type]
 
 
 class AfterConnectionReleasedEvent:

tests/test_event.py

@@ -31,7 +31,7 @@ def callback(event):
         mock_another_event_listener = Mock(spec=EventListenerInterface)
         mock_another_event_listener.listen = callback
         dispatcher.register_listeners(
-            event_listeners={type(mock_event): [mock_another_event_listener]}
+            mappings={type(mock_event): [mock_another_event_listener]}
         )
         dispatcher.dispatch(mock_event)
 
@@ -60,7 +60,7 @@ async def callback(event):
         mock_another_event_listener = Mock(spec=AsyncEventListenerInterface)
         mock_another_event_listener.listen = callback
         dispatcher.register_listeners(
-            event_listeners={type(mock_event): [mock_another_event_listener]}
+            mappings={type(mock_event): [mock_another_event_listener]}
         )
         await dispatcher.dispatch_async(mock_event)