fix races

Author: luke-beamcloud

fastloop/monitor.py

@@ -100,9 +100,10 @@ async def _check_orphaned_loops(self) -> None:
         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.state_manager.try_claim_loop_recovery(loop.loop_id):
+                continue
+
+            logger.info("Orphaned loop recovered", 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
@@ -115,12 +116,14 @@ async def _check_orphaned_workflows(self) -> None:
         for workflow in running_workflows:
             if await self.state_manager.workflow_has_claim(workflow.workflow_run_id):
                 continue
+            if not await self.state_manager.try_claim_workflow_recovery(
+                workflow.workflow_run_id
+            ):
+                continue
+
             logger.info(
-                "Workflow has no claim, restarting",
-                extra={
-                    "workflow_run_id": workflow.workflow_run_id,
-                    "block_index": workflow.current_block_index,
-                },
+                "Orphaned workflow recovered",
+                extra={"workflow_run_id": workflow.workflow_run_id},
             )
             if not await self.fastloop_instance.restart_workflow(
                 workflow.workflow_run_id
@@ -136,16 +139,17 @@ async def _check_orphaned_tasks(self) -> None:
         for task in running_tasks:
             if await self.state_manager.task_has_claim(task.task_id):
                 continue
+            if not await self.state_manager.try_claim_task_recovery(task.task_id):
+                continue
+
+            await self.state_manager.update_task_status(task.task_id, TaskStatus.FAILED)
 
             metadata = self.fastloop_instance._task_metadata.get(task.task_name)
             if not metadata:
-                await self.state_manager.update_task_status(
-                    task.task_id, TaskStatus.FAILED
-                )
                 continue
 
             logger.info(
-                "Task has no claim, restarting",
+                "Orphaned task recovered",
                 extra={"task_id": task.task_id, "task_name": task.task_name},
             )
             await self.fastloop_instance.task_manager.submit(
@@ -157,12 +161,6 @@ async def _check_orphaned_tasks(self) -> None:
             )
 
     async def _check_scheduled_workflows(self) -> None:
-        """Check for IDLE workflows with past-due scheduled wake times.
-
-        This is a backup mechanism that catches workflows that may have been
-        removed from the ZSET but not yet processed (e.g., if the wake queue
-        consumer failed or the wake monitoring thread died).
-        """
         now = time.time()
         idle_workflows = await self.state_manager.get_all_workflows(
             status=LoopStatus.IDLE
@@ -174,17 +172,14 @@ async def _check_scheduled_workflows(self) -> None:
                 continue
             if await self.state_manager.workflow_has_claim(workflow.workflow_run_id):
                 continue
-            claimed_from_zset = await self.state_manager.try_claim_workflow_wake(
+            if not await self.state_manager.try_claim_workflow_wake(
                 workflow.workflow_run_id
-            )
+            ):
+                continue
+
             logger.info(
-                "IDLE workflow has past-due wake time, restarting",
-                extra={
-                    "workflow_run_id": workflow.workflow_run_id,
-                    "scheduled_wake_time": workflow.scheduled_wake_time,
-                    "block_index": workflow.current_block_index,
-                    "claimed_from_zset": claimed_from_zset,
-                },
+                "IDLE workflow past-due, restarting",
+                extra={"workflow_run_id": workflow.workflow_run_id},
             )
             if await self.fastloop_instance.restart_workflow(workflow.workflow_run_id):
                 await self.state_manager.clear_workflow_wake_time(
@@ -200,8 +195,12 @@ async def _check_scheduled_workflows(self) -> None:
 
     async def _check_scheduled_tasks(self) -> None:
         for schedule_id, schedule in await self.state_manager.get_due_schedules():
+            if not await self.state_manager.try_claim_schedule(schedule_id):
+                continue
+
             metadata = self.fastloop_instance._task_metadata.get(schedule.task_name)
             if not metadata:
+                await self.state_manager.advance_schedule(schedule_id, schedule)
                 continue
 
             try:
@@ -212,12 +211,13 @@ async def _check_scheduled_tasks(self) -> None:
                     retry_policy=metadata.get("retry"),
                     executor_type=metadata.get("executor"),
                 )
-                await self.state_manager.advance_schedule(schedule_id, schedule)
             except Exception as e:
                 logger.error(
                     "Scheduled task failed",
                     extra={"schedule_id": schedule_id, "error": str(e)},
                 )
+            finally:
+                await self.state_manager.advance_schedule(schedule_id, schedule)
 
     async def _check_disconnect_stops(self) -> None:
         active_ids = await self.loop_manager.active_loop_ids()

fastloop/state/state_redis.py

@@ -449,6 +449,12 @@ async def has_claim(self, loop_id: str) -> bool:
         )
         return result is not None
 
+    async def try_claim_loop_recovery(self, loop_id: str) -> bool:
+        """Atomically claim right to recover an orphaned loop. Returns True if won."""
+        claim_key = f"{RedisKeys.LOOP_CLAIM.format(app_name=self.app_name, loop_id=loop_id)}:recovery"
+        acquired = await self.rdb.set(claim_key, "1", nx=True, ex=60)
+        return acquired is not None
+
     async def get_all_loop_ids(self) -> set[str]:
         members = await self.rdb.smembers(
             RedisKeys.LOOP_INDEX.format(app_name=self.app_name)
@@ -864,6 +870,12 @@ async def workflow_has_claim(self, workflow_run_id: str) -> bool:
         )
         return result is not None
 
+    async def try_claim_workflow_recovery(self, workflow_run_id: str) -> bool:
+        """Atomically claim right to recover an orphaned workflow. Returns True if won."""
+        claim_key = f"{RedisKeys.WORKFLOW_CLAIM.format(app_name=self.app_name, workflow_run_id=workflow_run_id)}:recovery"
+        acquired = await self.rdb.set(claim_key, "1", nx=True, ex=60)
+        return acquired is not None
+
     @asynccontextmanager
     async def with_workflow_claim(
         self, workflow_run_id: str
@@ -1072,6 +1084,12 @@ async def task_has_claim(self, task_id: str) -> bool:
         result = await self.rdb.get(claim_key)
         return result is not None
 
+    async def try_claim_task_recovery(self, task_id: str) -> bool:
+        """Atomically claim right to recover an orphaned task. Returns True if won."""
+        claim_key = f"{RedisKeys.TASK_CLAIM.format(app_name=self.app_name, task_id=task_id)}:recovery"
+        acquired = await self.rdb.set(claim_key, "1", nx=True, ex=60)
+        return acquired is not None
+
     @asynccontextmanager
     async def with_task_claim(self, task_id: str) -> AsyncGenerator[None, None]:
         lease_key = RedisKeys.TASK_CLAIM.format(app_name=self.app_name, task_id=task_id)
@@ -1170,6 +1188,12 @@ async def get_due_schedules(self) -> list[tuple[str, Schedule]]:
 
         return results
 
+    async def try_claim_schedule(self, schedule_id: str) -> bool:
+        """Atomically claim a due schedule. Returns True if this replica won."""
+        queue_key = RedisKeys.SCHEDULE_QUEUE.format(app_name=self.app_name)
+        removed = await self.rdb.zrem(queue_key, schedule_id)
+        return removed > 0
+
     async def advance_schedule(self, schedule_id: str, schedule: Schedule) -> None:
         schedule.next_run = schedule.compute_next_run()
         await self.save_schedule(schedule_id, schedule)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "fastloop"
-version = "0.1.99"
+version = "0.1.100"
 description = "A Python package for deploying stateful loops"
 readme = "README.md"
 requires-python = ">=3.12"