Refactor loop monitor and state management for efficiency

Co-authored-by: luke <luke@smartshare.io>

Author: cursoragent

fastloop/context.py

@@ -21,9 +21,19 @@
     from .integrations import Integration
 
 logger = setup_logger(__name__)
-
 T = TypeVar("T", bound="LoopContext")
 
+_DURATION_RE = re.compile(
+    r"^(\d+(?:\.\d+)?)\s*(seconds?|secs?|minutes?|mins?|hours?|hrs?|days?)$"
+)
+_UNIT_MULTIPLIERS = {
+    "sec": 1,
+    "min": 60,
+    "hour": 3600,
+    "hr": 3600,
+    "day": 86400,
+}
+
 
 class LoopContext:
     def __init__(
@@ -231,26 +241,12 @@ def should_pause(self) -> bool:
         return self._pause_requested
 
     def _parse_duration(self, duration_str: str) -> float:
-        duration_str = duration_str.lower().strip()
-
-        match = re.match(
-            r"^(\d+(?:\.\d+)?)\s*(seconds?|secs?|minutes?|mins?|hours?|hrs?|days?)$",
-            duration_str,
-        )
+        match = _DURATION_RE.match(duration_str.lower().strip())
         if not match:
             raise ValueError(f"Invalid duration format: {duration_str}")
 
-        value = float(match.group(1))
-        unit = match.group(2)
-
-        # Convert to seconds
-        if unit.startswith("sec"):
-            return value
-        elif unit.startswith("min"):
-            return value * 60
-        elif unit.startswith("hour") or unit.startswith("hr"):
-            return value * 3600
-        elif unit.startswith("day"):
-            return value * 86400
-        else:
-            raise ValueError(f"Unknown time unit: {unit}")
+        value, unit = float(match.group(1)), match.group(2)
+        for prefix, mult in _UNIT_MULTIPLIERS.items():
+            if unit.startswith(prefix):
+                return value * mult
+        return value

fastloop/fastloop.py

@@ -22,7 +22,7 @@
 from .integrations import Integration
 from .logging import configure_logging, setup_logger
 from .loop import LoopEvent, LoopManager
-from .state.state import LoopState, StateManager, create_state_manager
+from .state.state import StateManager, create_state_manager
 from .types import BaseConfig, LoopStatus
 from .utils import get_func_import_path, import_func_from_path, infer_application_path
 
@@ -499,120 +499,76 @@ def __init__(
         wake_queue: Queue[str],
         fastloop_instance: FastLoop,
     ):
-        self.state_manager: StateManager = state_manager
-        self.loop_manager: LoopManager = loop_manager
-        self.restart_callback: Callable[[str], Coroutine[Any, Any, bool]] = (
-            restart_callback
-        )
-        self.wake_queue: Queue[str] = wake_queue
-        self.fastloop_instance: FastLoop = fastloop_instance
-        self._stop_event: asyncio.Event = asyncio.Event()
+        self.state_manager = state_manager
+        self.loop_manager = loop_manager
+        self.restart_callback = restart_callback
+        self.wake_queue = wake_queue
+        self.fastloop_instance = fastloop_instance
+        self._stop_event = asyncio.Event()
 
     def stop(self) -> None:
         self._stop_event.set()
 
-    async def run(self):
-        while not self._stop_event.is_set():
-            try:
-                if not self.wake_queue.empty():
-                    loop_id = self.wake_queue.get()
-                    if await self.state_manager.has_claim(loop_id):
-                        continue
-
-                    logger.info(
-                        "Loop woke up, restarting",
-                        extra={"loop_id": loop_id},
-                    )
-                    if not await self.restart_callback(loop_id):
-                        await self.state_manager.update_loop_status(
-                            loop_id, LoopStatus.STOPPED
-                        )
-                        await self.loop_manager.stop(loop_id)
-
-                    continue
-
-                loop_ids: set[str] = await self.state_manager.get_all_loop_ids()
-                active_loop_ids: set[str] = await self.loop_manager.active_loop_ids()
-                loops_running: set[str] = active_loop_ids.intersection(loop_ids)
-
-                for loop_id in loops_running:
-                    loop = await self.state_manager.get_loop(loop_id)
-
-                    if (
-                        loop.status in LoopStatus.IDLE
-                        or loop.status == LoopStatus.STOPPED
-                    ):
-                        if await self.state_manager.has_claim(loop_id):
-                            continue
-
-                        logger.info(
-                            "Loop is idle or stopped, stopping",
-                            extra={"loop_id": loop_id},
-                        )
+    async def _process_wake(self, loop_id: str) -> None:
+        if await self.state_manager.has_claim(loop_id):
+            return
+        logger.info("Loop woke up, restarting", extra={"loop_id": loop_id})
+        if not await self.restart_callback(loop_id):
+            await self.state_manager.update_loop_status(loop_id, LoopStatus.STOPPED)
 
-                        await self.loop_manager.stop(loop_id)
-                        continue
+    async def _check_orphaned_loops(self) -> None:
+        running_loops = await self.state_manager.get_all_loops(
+            status=LoopStatus.RUNNING
+        )
+        for loop in running_loops:
+            if await self.state_manager.has_claim(loop.loop_id):
+                continue
+            logger.info(
+                "Loop has no claim, restarting", extra={"loop_id": loop.loop_id}
+            )
+            if not await self.restart_callback(loop.loop_id):
+                await self.state_manager.update_loop_status(
+                    loop.loop_id, LoopStatus.STOPPED
+                )
 
-                loops: list[LoopState] = await self.state_manager.get_all_loops(
-                    status=LoopStatus.RUNNING
+    async def _check_disconnect_stops(self) -> None:
+        active_ids = await self.loop_manager.active_loop_ids()
+        for loop_id in active_ids:
+            try:
+                loop = await self.state_manager.get_loop(loop_id)
+            except LoopNotFoundError:
+                continue
+            if not loop.loop_name:
+                continue
+            metadata = self.fastloop_instance._loop_metadata.get(loop.loop_name)
+            if not metadata or not metadata.get("stop_on_disconnect"):
+                continue
+            if not await self.fastloop_instance.has_active_clients(loop_id):
+                logger.info(
+                    "Loop has no clients, stopping",
+                    extra={"loop_id": loop_id, "loop_name": loop.loop_name},
                 )
-                for loop in loops:
-                    # Restart loop if it has no claim and is not idle (maybe the task crashed or was interrupted)
-                    if not await self.state_manager.has_claim(loop.loop_id):
-                        logger.info(
-                            "Loop has no claim, restarting",
-                            extra={
-                                "loop_id": loop.loop_id,
-                            },
-                        )
+                await self.state_manager.update_loop_status(loop_id, LoopStatus.STOPPED)
+                await self.loop_manager.stop(loop_id)
+
+    async def run(self):
+        while not self._stop_event.is_set():
+            try:
+                while not self.wake_queue.empty():
+                    await self._process_wake(self.wake_queue.get_nowait())
 
-                        if not await self.restart_callback(loop.loop_id):
-                            await self.state_manager.update_loop_status(
-                                loop.loop_id, LoopStatus.STOPPED
-                            )
-                            await self.loop_manager.stop(loop.loop_id)
-
-                        continue
-
-                # Check for loops with stop_on_disconnect=true that have no active clients
-                for loop in loops:
-                    if (
-                        loop.loop_name
-                        and loop.loop_name in self.fastloop_instance._loop_metadata
-                    ):
-                        metadata = self.fastloop_instance._loop_metadata[loop.loop_name]
-                        if metadata.get("stop_on_disconnect", False):
-                            has_clients = (
-                                await self.fastloop_instance.has_active_clients(
-                                    loop.loop_id
-                                )
-                            )
-                            if not has_clients:
-                                logger.info(
-                                    "Loop has stop_on_disconnect=true and no active clients, stopping",
-                                    extra={
-                                        "loop_id": loop.loop_id,
-                                        "loop_name": loop.loop_name,
-                                    },
-                                )
-                                await self.state_manager.update_loop_status(
-                                    loop.loop_id, LoopStatus.STOPPED
-                                )
-                                await self.loop_manager.stop(loop.loop_id)
+                await self._check_orphaned_loops()
+                await self._check_disconnect_stops()
 
                 try:
                     await asyncio.wait_for(
                         self._stop_event.wait(), timeout=WATCHDOG_INTERVAL_S
                     )
                     break
                 except TimeoutError:
-                    continue
-
+                    pass
             except asyncio.CancelledError:
                 break
-            except BaseException as e:
-                logger.error(
-                    "Error in loop monitor",
-                    extra={"error": str(e)},
-                )
+            except Exception as e:
+                logger.error("Error in loop monitor", extra={"error": str(e)})
                 await asyncio.sleep(WATCHDOG_INTERVAL_S)

fastloop/state/state.py

@@ -36,16 +36,15 @@ def from_json(cls, json_str: str) -> "LoopState":
 
 class StateManager(ABC):
     @abstractmethod
-    async def get_all_loop_ids(
-        self,
-    ) -> set[str]:
+    async def get_all_loop_ids(self) -> set[str]:
         pass
 
     @abstractmethod
-    async def get_all_loops(
-        self,
-        status: LoopStatus | None = None,
-    ) -> list[LoopState]:
+    async def get_running_loop_ids(self) -> set[str]:
+        pass
+
+    @abstractmethod
+    async def get_all_loops(self, status: LoopStatus | None = None) -> list[LoopState]:
         pass
 
     @abstractmethod