fix wake condition

Author: luke-beamcloud

fastloop/fastloop.py

@@ -801,13 +801,24 @@ async def _process_wake(self, wake_id: str) -> None:
             if await self.state_manager.workflow_has_claim(workflow_id):
                 return
             logger.info(
-                "Workflow woke up, restarting", extra={"workflow_id": workflow_id}
+                "Workflow woke up, attempting restart",
+                extra={"workflow_id": workflow_id},
             )
-            await self.state_manager.clear_workflow_wake_time(workflow_id)
-            if not await self.fastloop_instance.restart_workflow(workflow_id):
+            if await self.fastloop_instance.restart_workflow(workflow_id):
+                await self.state_manager.clear_workflow_wake_time(workflow_id)
+                logger.info(
+                    "Workflow restarted successfully",
+                    extra={"workflow_id": workflow_id},
+                )
+            else:
+                await self.state_manager.clear_workflow_wake_time(workflow_id)
                 await self.state_manager.update_workflow_status(
                     workflow_id, LoopStatus.STOPPED
                 )
+                logger.warning(
+                    "Workflow restart failed, marked as stopped",
+                    extra={"workflow_id": workflow_id},
+                )
         else:
             loop_id = wake_id
             if await self.state_manager.has_claim(loop_id):
@@ -850,6 +861,38 @@ async def _check_orphaned_workflows(self) -> None:
                     workflow.workflow_id, LoopStatus.STOPPED
                 )
 
+    async def _check_scheduled_workflows(self) -> None:
+        """Check for IDLE workflows with past-due scheduled wake times."""
+        import time
+
+        now = time.time()
+        idle_workflows = await self.state_manager.get_all_workflows(
+            status=LoopStatus.IDLE
+        )
+        for workflow in idle_workflows:
+            if not workflow.scheduled_wake_time:
+                continue
+            if workflow.scheduled_wake_time > now:
+                continue
+            if await self.state_manager.workflow_has_claim(workflow.workflow_id):
+                continue
+            if not await self.state_manager.try_claim_workflow_wake(
+                workflow.workflow_id
+            ):
+                continue
+            logger.info(
+                "IDLE workflow has past-due wake time, restarting",
+                extra={
+                    "workflow_id": workflow.workflow_id,
+                    "scheduled_wake_time": workflow.scheduled_wake_time,
+                    "block_index": workflow.current_block_index,
+                },
+            )
+            if not await self.fastloop_instance.restart_workflow(workflow.workflow_id):
+                await self.state_manager.update_workflow_status(
+                    workflow.workflow_id, LoopStatus.STOPPED
+                )
+
     async def _check_disconnect_stops(self) -> None:
         active_ids = await self.loop_manager.active_loop_ids()
         for loop_id in active_ids:
@@ -928,6 +971,7 @@ async def run(self):
 
                 await self._check_orphaned_loops()
                 await self._check_orphaned_workflows()
+                await self._check_scheduled_workflows()
                 await self._check_disconnect_stops()
 
                 try:

fastloop/state/state.py

@@ -247,6 +247,11 @@ async def set_workflow_wake_time(self, workflow_id: str, timestamp: float) -> No
     async def clear_workflow_wake_time(self, workflow_id: str) -> None:
         pass
 
+    @abstractmethod
+    async def try_claim_workflow_wake(self, workflow_id: str) -> bool:
+        """Atomically try to claim a workflow wake. Returns True if this caller won the race."""
+        pass
+
     @abstractmethod
     async def set_workflow_block_output(self, workflow_id: str, output: Any) -> None:
         pass

fastloop/state/state_redis.py

@@ -128,7 +128,13 @@ def _run_wake_monitoring(self):
             with suppress(sync_redis.exceptions.ResponseError):
                 rdb.config_set("notify-keyspace-events", "Ex")
 
-            self._process_due_wakes(rdb)
+            logger.info("Wake monitoring thread started, processing due wakes")
+            due_count = self._process_due_wakes(rdb)
+            if due_count > 0:
+                logger.info(
+                    "Processed due wakes on startup",
+                    extra={"count": due_count},
+                )
 
             pubsub = rdb.pubsub()
             pubsub.psubscribe("__keyevent@*__:expired")
@@ -142,9 +148,17 @@ def _run_wake_monitoring(self):
                         key = message["data"].decode("utf-8")
                         if f":{self.app_name}:wake:" in key:
                             loop_id = key.split(":")[-1]
+                            logger.info(
+                                "Loop wake key expired",
+                                extra={"loop_id": loop_id},
+                            )
                             self._queue_wake(rdb, loop_id)
                         elif f":{self.app_name}:workflow_wake:" in key:
                             workflow_id = key.split(":")[-1]
+                            logger.info(
+                                "Workflow wake key expired",
+                                extra={"workflow_id": workflow_id},
+                            )
                             self._queue_wake(rdb, workflow_id)
                     except Exception as e:
                         logger.error(f"Error processing wake notification: {e}")
@@ -174,6 +188,10 @@ def _process_due_wakes(self, rdb) -> int:
         for loop_id_bytes in due_loop_wakes:
             loop_id = loop_id_bytes.decode("utf-8")
             if rdb.zrem(loop_schedule_key, loop_id):
+                logger.info(
+                    "Due loop wake found, queuing",
+                    extra={"loop_id": loop_id},
+                )
                 self.wake_queue.put(loop_id)
                 processed += 1
 
@@ -186,6 +204,10 @@ def _process_due_wakes(self, rdb) -> int:
         for workflow_id_bytes in due_workflow_wakes:
             workflow_id = workflow_id_bytes.decode("utf-8")
             if rdb.zrem(workflow_schedule_key, workflow_id):
+                logger.info(
+                    "Due workflow wake found, queuing",
+                    extra={"workflow_id": workflow_id},
+                )
                 self.wake_queue.put(f"workflow:{workflow_id}")
                 processed += 1
 
@@ -858,6 +880,15 @@ async def set_workflow_wake_time(self, workflow_id: str, timestamp: float) -> No
             )
             await pipe.execute()
 
+        logger.info(
+            "Workflow wake scheduled",
+            extra={
+                "workflow_id": workflow_id,
+                "wake_timestamp": timestamp,
+                "ttl_ms": ttl_ms,
+            },
+        )
+
     async def clear_workflow_wake_time(self, workflow_id: str) -> None:
         """Clear any scheduled workflow wake time."""
         schedule_key = RedisKeys.WORKFLOW_WAKE_SCHEDULE.format(app_name=self.app_name)
@@ -870,6 +901,28 @@ async def clear_workflow_wake_time(self, workflow_id: str) -> None:
             pipe.delete(wake_key)
             await pipe.execute()
 
+        logger.info(
+            "Workflow wake cleared",
+            extra={"workflow_id": workflow_id},
+        )
+
+    async def try_claim_workflow_wake(self, workflow_id: str) -> bool:
+        """Atomically try to claim a workflow wake. Returns True if this caller won the race."""
+        schedule_key = RedisKeys.WORKFLOW_WAKE_SCHEDULE.format(app_name=self.app_name)
+        wake_key = RedisKeys.WORKFLOW_WAKE_KEY.format(
+            app_name=self.app_name, workflow_id=workflow_id
+        )
+
+        removed = await self.rdb.zrem(schedule_key, workflow_id)
+        if removed:
+            await self.rdb.delete(wake_key)
+            logger.info(
+                "Workflow wake claimed",
+                extra={"workflow_id": workflow_id},
+            )
+            return True
+        return False
+
     async def set_workflow_block_output(self, workflow_id: str, output: Any) -> None:
         """Store the block output for a workflow using cloudpickle."""
         output_key = RedisKeys.WORKFLOW_BLOCK_OUTPUT.format(

fastloop/state/state_s3.py

@@ -701,3 +701,41 @@ async def get_all_workflows(
                 continue
 
         return workflows
+
+    async def set_workflow_wake_time(self, workflow_id: str, timestamp: float) -> None:
+        workflow = await self.get_workflow(workflow_id)
+        workflow.scheduled_wake_time = timestamp
+        workflow.status = LoopStatus.IDLE
+        await self.update_workflow(workflow_id, workflow)
+
+    async def clear_workflow_wake_time(self, workflow_id: str) -> None:
+        try:
+            workflow = await self.get_workflow(workflow_id)
+            workflow.scheduled_wake_time = None
+            await self.update_workflow(workflow_id, workflow)
+        except WorkflowNotFoundError:
+            pass
+
+    async def try_claim_workflow_wake(self, workflow_id: str) -> bool:
+        try:
+            workflow = await self.get_workflow(workflow_id)
+            if workflow.scheduled_wake_time is None:
+                return False
+            workflow.scheduled_wake_time = None
+            await self.update_workflow(workflow_id, workflow)
+            return True
+        except WorkflowNotFoundError:
+            return False
+
+    async def set_workflow_block_output(self, workflow_id: str, output: Any) -> None:
+        key = f"{self.prefix}/{self.app_name}/workflow_output/{workflow_id}.pkl"
+        output_bytes: bytes = cloudpickle.dumps(output)
+        self.s3.put_object(Bucket=self.bucket, Key=key, Body=output_bytes)
+
+    async def get_workflow_block_output(self, workflow_id: str) -> Any:
+        key = f"{self.prefix}/{self.app_name}/workflow_output/{workflow_id}.pkl"
+        try:
+            response = self.s3.get_object(Bucket=self.bucket, Key=key)
+            return cloudpickle.loads(response["Body"].read())
+        except ClientError:
+            return None

pyproject.toml

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