Performance improvements

Author: codeSamuraii

.gitignore

@@ -14,3 +14,4 @@ build/
 .claudeignore
 node_modules/
 cloud_poc/frontend/dist/
+examples/stress_task.py

offwork/worker/backends/base.py

@@ -41,6 +41,17 @@ async def send_heartbeat(self, task_id: str) -> None:
         heartbeat-based stall detection should override this.
         """
 
+    async def heartbeat_and_check_cancel(self, task_id: str) -> bool:
+        """Send a heartbeat and return whether the task is cancelled.
+
+        Backends with a single round-trip combining both (e.g. an
+        HTTP POST whose response carries the cancel flag) should
+        override this to halve worker→broker chatter while a task is
+        running. The default implementation issues two calls.
+        """
+        await self.send_heartbeat(task_id)
+        return await self.is_cancelled(task_id)
+
     async def get_heartbeat(self, task_id: str) -> float | None:
         """Return the timestamp of the last heartbeat for *task_id*.
 

offwork/worker/backends/http.py

@@ -171,6 +171,12 @@ async def try_get_result(self, task_id: str) -> str | None:
     async def send_heartbeat(self, task_id: str) -> None:
         await self._request("POST", f"/tasks/{task_id}/heartbeat")
 
+    async def heartbeat_and_check_cancel(self, task_id: str) -> bool:
+        _status, body = await self._request(
+            "POST", f"/tasks/{task_id}/heartbeat", allow_not_found=True,
+        )
+        return bool(body and body.get("cancelled"))
+
     async def get_heartbeat(self, task_id: str) -> float | None:
         _status, body = await self._request(
             "GET", f"/tasks/{task_id}/heartbeat", allow_not_found=True,

offwork/worker/remote.py

@@ -347,7 +347,7 @@ def _build_detail_tags(worker: Worker) -> str:
     return ", ".join(parts)
 
 
-_HEARTBEAT_INTERVAL = 1.0
+_HEARTBEAT_INTERVAL = 5.0
 
 
 async def _heartbeat_loop(
@@ -358,23 +358,23 @@ async def _heartbeat_loop(
 ) -> None:
     """Send periodic heartbeats and check for cancellation.
 
-    When *exec_task* is provided and the backend reports the task as
-    cancelled, the execution task is cancelled via
+    Backends override ``heartbeat_and_check_cancel`` to combine both
+    into a single round-trip (e.g. an HTTP POST whose response carries
+    the cancel flag). When *exec_task* is provided and the backend
+    reports the task as cancelled, the execution task is cancelled via
     :meth:`asyncio.Task.cancel`, which raises :class:`CancelledError`
     at the next ``await`` in async user functions.
     """
     while not cancel_event.is_set():
         try:
-            await backend.send_heartbeat(task_id)
-        except Exception:
-            logger.debug("Heartbeat send failed for task %s", task_id, exc_info=True)
-        if exec_task is not None:
-            try:
-                if await backend.is_cancelled(task_id):
+            if exec_task is not None:
+                if await backend.heartbeat_and_check_cancel(task_id):
                     exec_task.cancel()
                     return
-            except Exception:
-                logger.debug("Cancellation check failed for task %s", task_id, exc_info=True)
+            else:
+                await backend.send_heartbeat(task_id)
+        except Exception:
+            logger.debug("Heartbeat send failed for task %s", task_id, exc_info=True)
         try:
             await asyncio.wait_for(cancel_event.wait(), timeout=_HEARTBEAT_INTERVAL)
         except asyncio.TimeoutError:

offwork/worker/result.py

@@ -171,13 +171,29 @@ async def _wait_with_stall_detection(
         timeout: float | None,
         stall_timeout: float,
     ) -> None:
-        """Poll for result with heartbeat-based stall detection."""
+        """Wait for the result with heartbeat-based stall detection.
+
+        Uses the backend's blocking ``get_result`` (long-poll on HTTP
+        backends) in bounded slices, with a heartbeat check between
+        slices. The slice length is derived from ``stall_timeout`` so
+        we still detect a silent worker in time.
+        """
         deadline = None if timeout is None else time.monotonic() + timeout
+        slice_seconds = max(1.0, min(stall_timeout / 2, 30.0))
         last_hb_value: float | None = None
         last_hb_change: float | None = None
 
         while True:
-            raw = await self._backend.try_get_result(self._task_id)
+            remaining = None if deadline is None else deadline - time.monotonic()
+            if remaining is not None and remaining <= 0:
+                raise TimeoutError(
+                    f"Timed out waiting for result of task {self._task_id}"
+                )
+            wait_for = slice_seconds if remaining is None else min(slice_seconds, remaining)
+            try:
+                raw = await self._backend.get_result(self._task_id, timeout=wait_for)
+            except TimeoutError:
+                raw = None
             if raw is not None:
                 self._envelope = ResultEnvelope.from_json(raw)
                 logger.debug(
@@ -186,14 +202,6 @@ async def _wait_with_stall_detection(
                 )
                 return
 
-            if deadline is not None:
-                remaining = deadline - time.monotonic()
-                if remaining <= 0:
-                    raise TimeoutError(
-                        f"Timed out waiting for result of task {self._task_id}"
-                    )
-
-            logger.debug("Polling heartbeat for task %s", self._task_id[:8])
             hb = await self._backend.get_heartbeat(self._task_id)
             now = time.monotonic()
             if hb is not None and hb != last_hb_value:
@@ -206,8 +214,6 @@ async def _wait_with_stall_detection(
                     f"{elapsed:.1f}s (threshold: {stall_timeout}s)"
                 )
 
-            await asyncio.sleep(1.0)
-
     def __await__(self) -> Generator[Any, None, Any]:
         """Allow ``await result`` as shorthand for ``await result.result()``."""
         return self.result().__await__()