Efficiency 2

Author: codeSamuraii

offwork/worker/backends/ws.py

@@ -186,6 +186,7 @@ async def _request(
         backoff = _RECONNECT_BACKOFF_MIN
         last_exc: Exception | None = None
         for attempt in range(3):
+            req_id: str | None = None
             try:
                 ws = await self._ensure_connected()
                 req_id = uuid.uuid4().hex
@@ -203,6 +204,10 @@ async def _request(
                     return await future
                 return await asyncio.wait_for(future, timeout=timeout)
             except (ConnectionError, OSError) as exc:
+                # Drop any orphaned pending entry before reconnecting so a
+                # late response can't resolve a future no one awaits.
+                if req_id is not None:
+                    self._pending.pop(req_id, None)
                 last_exc = exc
                 if self._closed:
                     raise
@@ -212,7 +217,8 @@ async def _request(
             except asyncio.TimeoutError:
                 # Drop the pending entry — the response, if it ever
                 # arrives, has no waiter.
-                self._pending.pop(req_id, None)
+                if req_id is not None:
+                    self._pending.pop(req_id, None)
                 raise
         assert last_exc is not None
         raise last_exc

offwork/worker/remote.py

@@ -540,6 +540,7 @@ async def _handle_task(
     known_clients: KnownClients | None = None,
     nonce_lru: NonceLRU | None = None,
     clock_skew: float = DEFAULT_CLOCK_SKEW,
+    sem: asyncio.Semaphore | None = None,
 ) -> None:
     """Process a single task: deserialize, execute with policy, send result.
 
@@ -573,13 +574,32 @@ async def _handle_task(
 
     logger.debug("Received task %s: %s", task.task_id, task.function_name)
 
-    # Wait for scheduled time
+    # Wait for scheduled time *before* acquiring a concurrency slot, so a
+    # task scheduled far in the future does not occupy one of the worker's
+    # limited execution slots while it sleeps.
     if task.scheduled_at is not None:
         delay = task.scheduled_at - time.time()
         if delay > 0:
             logger.debug("Task %s scheduled in %.1fs", task.task_id, delay)
             await asyncio.sleep(delay)
 
+    # Acquire the concurrency slot only for the actual execution phase.
+    slot = sem if sem is not None else contextlib.nullcontext()
+    async with slot:
+        await _execute_task(
+            worker, backend, task,
+            root_token=root_token,
+        )
+
+
+async def _execute_task(
+    worker: Worker,
+    backend: Backend,
+    task: Task,
+    *,
+    root_token: bytes | None = None,
+) -> None:
+    """Run a ready task (scheduled wait already elapsed) and send its result."""
     # Any failure in the backend checks below must still surface to the
     # client, otherwise it would hang forever polling for a result.
     try:
@@ -793,14 +813,17 @@ def _on_shutdown_signal() -> None:
             pass
 
     async def bounded_handle(task_json: str) -> None:
-        async with sem:
-            await _handle_task(
-                worker, backend, task_json,
-                root_token=root_token,
-                known_clients=known_clients,
-                nonce_lru=nonce_lru,
-                clock_skew=clock_skew,
-            )
+        # The semaphore is acquired inside _handle_task, *after* any
+        # scheduled-time wait, so far-future tasks don't hold a slot while
+        # sleeping.
+        await _handle_task(
+            worker, backend, task_json,
+            root_token=root_token,
+            known_clients=known_clients,
+            nonce_lru=nonce_lru,
+            clock_skew=clock_skew,
+            sem=sem,
+        )
 
     async def _listen() -> None:
         async for task_json in backend.listen():

offwork/worker/result.py

@@ -324,18 +324,24 @@ async def _do_cancel(self, resolved: float | None) -> bool:
 
         deadline = None if resolved is None else time.monotonic() + resolved
         while True:
-            raw = await self._backend.try_get_result(self._task_id)
-            if raw is not None:
-                self._envelope = ResultEnvelope.from_json(raw)
-                return True
-            if deadline is not None and time.monotonic() >= deadline:
+            remaining = None if deadline is None else deadline - time.monotonic()
+            if remaining is not None and remaining <= 0:
                 # Timed out waiting for worker confirmation — seed a cancelled
                 # envelope so future reads don't hang forever.
                 env = ResultEnvelope.cancelled(self._task_id)
                 await self._backend.send_result(self._task_id, env.to_json())
                 self._envelope = env
                 return False
-            await asyncio.sleep(0.5)
+            # Long-poll for the worker's confirmation envelope instead of
+            # busy-polling — wakes immediately on push-capable backends.
+            wait_for = 5.0 if remaining is None else min(5.0, 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)
+                return True
 
     # -- awaiting the result ---------------------------------------------------
 

tests/fixtures/_ws_broker.py

@@ -149,6 +149,8 @@ async def _dispatch(
         return {"cancelled": True}
     if op == "is_cancelled":
         return {"cancelled": payload["task_id"] in state.cancelled}
+    if op == "send_log_line":
+        return None
     if op == "send_progress":
         state.progress[payload["task_id"]] = payload["progress_json"]
         return {"ok": True}