generalize "forever" tasks to centralize exception handling/retry

masenf · masenf · commit 6a4db937ca44 · 2025-11-03T16:18:04.000-08:00
diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py
@@ -29,6 +29,7 @@
     LockExpiredError,
     StateSchemaMismatchError,
 )
+from reflex.utils.tasks import ensure_task
 
 
 def _default_lock_expiration() -> int:
@@ -155,7 +156,7 @@ def __post_init__(self):
             raise InvalidLockWarningThresholdError(msg)
         with contextlib.suppress(RuntimeError):
             asyncio.get_running_loop()  # Check if we're in an event loop.
-            self._lock_task = asyncio.create_task(self._lock_updates_forever())
+            self._ensure_lock_task()
 
     def _get_required_state_classes(
         self,
@@ -586,7 +587,7 @@ async def _create_lease_break_task(
         Returns:
             The lease break task, or None when there is contention.
         """
-        await self._ensure_lock_task()
+        self._ensure_lock_task()
 
         client_token, _ = _split_substate_key(token)
 
@@ -765,6 +766,9 @@ async def _subscribe_lock_updates(self, redis_db: int = 0):
         Args:
             redis_db: The logical database number to subscribe to.
         """
+        await self._enable_keyspace_notifications()
+        redis_db = self.redis.get_connection_kwargs().get("db", 0)
+
         lock_key_pattern = f"__keyspace@{redis_db}__:*_lock"
         lock_waiter_key_pattern = f"__keyspace@{redis_db}__:*_lock_waiters"
         handlers = {
@@ -776,27 +780,14 @@ async def _subscribe_lock_updates(self, redis_db: int = 0):
             async for _ in pubsub.listen():
                 pass
 
-    async def _lock_updates_forever(self) -> None:
-        """Background task to monitor Redis keyspace notifications for lock updates."""
-        await self._enable_keyspace_notifications()
-        redis_db = self.redis.get_connection_kwargs().get("db", 0)
-        while True:
-            try:
-                await self._subscribe_lock_updates(redis_db)
-            except asyncio.CancelledError:  # noqa: PERF203
-                raise
-            except Exception as e:
-                if isinstance(e, RuntimeError) and str(e) == "no running event loop":
-                    # Happens when shutting down, break out of the loop.
-                    raise
-                console.error(f"StateManagerRedis lock update task error: {e}")
-
-    async def _ensure_lock_task(self) -> None:
+    def _ensure_lock_task(self) -> None:
         """Ensure the lock updates subscriber task is running."""
-        if self._lock_task is None or self._lock_task.done():
-            async with self._state_manager_lock:
-                if self._lock_task is None or self._lock_task.done():
-                    self._lock_task = asyncio.create_task(self._lock_updates_forever())
+        ensure_task(
+            owner=self,
+            task_attribute="_lock_task",
+            coro_function=self._subscribe_lock_updates,
+            suppress_exceptions=[Exception],
+        )
 
     async def _enable_keyspace_notifications(self):
         """Enable keyspace notifications for the redis server.
@@ -954,7 +945,7 @@ async def _wait_lock(self, lock_key: bytes, lock_id: bytes) -> None:
                 )
             return
         # Make sure lock waiter task is running.
-        await self._ensure_lock_task()
+        self._ensure_lock_task()
         async with (
             self._lock_waiter(lock_key) as lock_released_event,
             self._request_lock_release(lock_key, lock_id),
diff --git a/reflex/utils/tasks.py b/reflex/utils/tasks.py
@@ -0,0 +1,109 @@
+"""Helpers for managing asyncio tasks."""
+
+import asyncio
+import time
+from collections.abc import Callable, Coroutine
+from typing import Any
+
+from reflex.utils import console
+
+
+async def _run_forever(
+    coro_function: Callable[..., Coroutine],
+    *args: Any,
+    suppress_exceptions: list[type[BaseException]],
+    exception_delay: float,
+    exception_limit: int,
+    exception_limit_window: float,
+    **kwargs: Any,
+):
+    """Wrapper to continuously run a coroutine function, suppressing certain exceptions.
+
+    Args:
+        coro_function: The coroutine function to run.
+        *args: The arguments to pass to the coroutine function.
+        suppress_exceptions: The exceptions to suppress.
+        exception_delay: The delay between retries when an exception is suppressed.
+        exception_limit: The maximum number of suppressed exceptions within the limit window before raising.
+        exception_limit_window: The time window in seconds for counting suppressed exceptions.
+        **kwargs: The keyword arguments to pass to the coroutine function.
+    """
+    last_regular_loop_start = 0
+    exception_count = 0
+
+    while True:
+        # Reset the exception count when the limit window has elapsed since the last non-exception loop started.
+        if last_regular_loop_start + exception_limit_window < time.monotonic():
+            exception_count = 0
+        if not exception_count:
+            last_regular_loop_start = time.monotonic()
+        try:
+            await coro_function(*args, **kwargs)
+        except (asyncio.CancelledError, RuntimeError):
+            raise
+        except Exception as e:
+            if any(isinstance(e, ex) for ex in suppress_exceptions):
+                exception_count += 1
+                if exception_count >= exception_limit:
+                    console.error(
+                        f"{coro_function.__name__}: task exceeded exception limit {exception_limit} within {exception_limit_window}s: {e}"
+                    )
+                    raise
+                console.error(f"{coro_function.__name__}: task error suppressed: {e}")
+                await asyncio.sleep(exception_delay)
+                continue
+            raise
+
+
+def ensure_task(
+    owner: Any,
+    task_attribute: str,
+    coro_function: Callable[..., Coroutine],
+    *args: Any,
+    suppress_exceptions: list[type[BaseException]] | None = None,
+    exception_delay: float = 1.0,
+    exception_limit: int = 5,
+    exception_limit_window: float = 60.0,
+    **kwargs: Any,
+) -> asyncio.Task:
+    """Ensure that a task is running for the given coroutine function.
+
+    Note: if the task is already running, args and kwargs are ignored.
+
+    Args:
+        owner: The owner of the task.
+        task_attribute: The attribute name to store/retrieve the task from the owner object.
+        coro_function: The coroutine function to run as a task.
+        suppress_exceptions: The exceptions to log and continue when running the coroutine.
+        exception_delay: The delay between retries when an exception is suppressed.
+        exception_limit: The maximum number of suppressed exceptions within the limit window before raising.
+        exception_limit_window: The time window in seconds for counting suppressed exceptions.
+        *args: The arguments to pass to the coroutine function.
+        **kwargs: The keyword arguments to pass to the coroutine function.
+
+    Returns:
+        The asyncio task running the coroutine function.
+    """
+    if suppress_exceptions is None:
+        suppress_exceptions = []
+    if RuntimeError in suppress_exceptions:
+        msg = "Cannot suppress RuntimeError exceptions which may be raised by asyncio machinery."
+        raise RuntimeError(msg)
+
+    task = getattr(owner, task_attribute, None)
+    if task is None or task.done():
+        asyncio.get_running_loop()  # Ensure we're in an event loop.
+        task = asyncio.create_task(
+            _run_forever(
+                coro_function,
+                *args,
+                suppress_exceptions=suppress_exceptions,
+                exception_delay=exception_delay,
+                exception_limit=exception_limit,
+                exception_limit_window=exception_limit_window,
+                **kwargs,
+            ),
+            name=f"reflex_ensure_task|{type(owner).__name__}.{task_attribute}={coro_function.__name__}|{time.time()}",
+        )
+        setattr(owner, task_attribute, task)
+    return task
diff --git a/reflex/utils/token_manager.py b/reflex/utils/token_manager.py
@@ -14,6 +14,7 @@
 from reflex.istate.manager.redis import StateManagerRedis
 from reflex.state import BaseState, StateUpdate
 from reflex.utils import console, prerequisites
+from reflex.utils.tasks import ensure_task
 
 if TYPE_CHECKING:
     from redis.asyncio import Redis
@@ -241,6 +242,11 @@ def _handle_socket_record_del(self, token: str) -> None:
 
     async def _subscribe_socket_record_updates(self, redis_db: int) -> None:
         """Subscribe to Redis keyspace notifications for socket record updates."""
+        await StateManagerRedis(
+            state=BaseState, redis=self.redis
+        )._enable_keyspace_notifications()
+        redis_db = self.redis.get_connection_kwargs().get("db", 0)
+
         async with self.redis.pubsub() as pubsub:
             await pubsub.psubscribe(
                 f"__keyspace@{redis_db}__:{self._get_redis_key('*')}"
@@ -260,29 +266,14 @@ async def _subscribe_socket_record_updates(self, redis_db: int) -> None:
                     elif event == "set":
                         await self._get_token_owner(token, refresh=True)
 
-    async def _socket_record_updates_forever(self) -> None:
-        """Background task to monitor Redis keyspace notifications for socket record updates."""
-        await StateManagerRedis(
-            state=BaseState, redis=self.redis
-        )._enable_keyspace_notifications()
-        redis_db = self.redis.get_connection_kwargs().get("db", 0)
-        while True:
-            try:
-                await self._subscribe_socket_record_updates(redis_db)
-            except asyncio.CancelledError:  # noqa: PERF203
-                raise
-            except Exception as e:
-                if isinstance(e, RuntimeError) and str(e) == "no running event loop":
-                    # Happens when shutting down, break out of the loop.
-                    raise
-                console.error(f"RedisTokenManager socket record update task error: {e}")
-
     def _ensure_socket_record_task(self) -> None:
         """Ensure the socket record updates subscriber task is running."""
-        if self._socket_record_task is None or self._socket_record_task.done():
-            self._socket_record_task = asyncio.create_task(
-                self._socket_record_updates_forever()
-            )
+        ensure_task(
+            owner=self,
+            task_attribute="_socket_record_task",
+            coro_function=self._subscribe_socket_record_updates,
+            suppress_exceptions=[Exception],
+        )
 
     async def link_token_to_sid(self, token: str, sid: str) -> str | None:
         """Link a token to a session ID with Redis-based duplicate detection.
@@ -389,26 +380,6 @@ async def _subscribe_lost_and_found_updates(
                     record = LostAndFoundRecord(**json.loads(message["data"].decode()))
                     await emit_update(StateUpdate(**record.update), record.token)
 
-    async def _lost_and_found_updates_forever(
-        self,
-        emit_update: Callable[[StateUpdate, str], Coroutine[None, None, None]],
-    ):
-        """Background task to monitor Redis lost and found deltas.
-
-        Args:
-            emit_update: The function to emit state updates.
-        """
-        while True:
-            try:
-                await self._subscribe_lost_and_found_updates(emit_update)
-            except asyncio.CancelledError:  # noqa: PERF203
-                raise
-            except Exception as e:
-                if isinstance(e, RuntimeError) and str(e) == "no running event loop":
-                    # Happens when shutting down, break out of the loop.
-                    raise
-                console.error(f"RedisTokenManager lost and found task error: {e}")
-
     def ensure_lost_and_found_task(
         self,
         emit_update: Callable[[StateUpdate, str], Coroutine[None, None, None]],
@@ -418,10 +389,13 @@ def ensure_lost_and_found_task(
         Args:
             emit_update: The function to emit state updates.
         """
-        if self._lost_and_found_task is None or self._lost_and_found_task.done():
-            self._lost_and_found_task = asyncio.create_task(
-                self._lost_and_found_updates_forever(emit_update)
-            )
+        ensure_task(
+            owner=self,
+            task_attribute="_lost_and_found_task",
+            coro_function=self._subscribe_lost_and_found_updates,
+            suppress_exceptions=[Exception],
+            emit_update=emit_update,
+        )
 
     async def _get_token_owner(self, token: str, refresh: bool = False) -> str | None:
         """Get the instance ID of the owner of a token.
diff --git a/tests/units/utils/test_tasks.py b/tests/units/utils/test_tasks.py
