worker: make drain deadlines explicit

proto/CONTRACT.md

@@ -503,7 +503,7 @@ Producers: autoscaler retirement, deploy rollover, operator action.
 
 | field | producer | consumer | semantics |
 |---|---|---|---|
-| `deadline_ms` | O policy | W drain loop | relative grace budget from receipt |
+| `deadline_ms` | O policy | W drain loop | `0` waits indefinitely; `>0` is the total grace budget from receipt |
 
 W behavior: stop admitting (`RunJob` after Drain ⇒ `JobResult{RETRYABLE,
 safe_message:"worker draining"}` without JobAccepted), finish in-flight jobs,

proto/worker_scheduler.proto

@@ -449,7 +449,7 @@ message FnDegraded {
 // Orchestrator -> worker: stop accepting work, finish in-flight jobs, ship
 // results, close the stream, exit. No reply message; stream close is the ack.
 message Drain {
-  int64 deadline_ms = 1; // relative grace budget; worker aborts remaining work at expiry
+  int64 deadline_ms = 1; // 0 waits indefinitely; >0 is the total grace budget from receipt
 }
 
 // ---------------------------------------------------------------------------

src/gen_worker/lifecycle.py

@@ -95,6 +95,7 @@ def __init__(self, settings: Settings, executor: Executor) -> None:
         # mere membership) so a runtime ladder demotion (gw#463) re-emits.
         self._emitted_degraded: dict[str, str] = {}
         self._drain_task: Optional[asyncio.Task] = None
+        self._drain_deadline_at: Optional[float] = None
 
     # ---- snapshots -----------------------------------------------------------
 
@@ -346,29 +347,44 @@ async def startup(self) -> None:
     # ---- drain -------------------------------------------------------------------
 
     def start_drain(self, deadline_ms: int) -> None:
-        if self._drain_task is None or self._drain_task.done():
-            self._drain_task = asyncio.create_task(
-                self.drain(deadline_ms), name="drain"
-            )
+        if self.draining:
+            return
+        self._begin_drain(deadline_ms)
+        self._drain_task = asyncio.create_task(self._finish_drain(), name="drain")
 
     async def drain(self, deadline_ms: int = 0) -> None:
         """stop admitting -> finish in-flight -> ship buffered results ->
-        close the stream -> signal exit 0."""
+        close the stream -> signal exit 0. Zero waits without a cutoff."""
         if self.draining:
             return
+        self._begin_drain(deadline_ms)
+        await self._finish_drain()
+
+    def _begin_drain(self, deadline_ms: int) -> None:
+        """Synchronously stop admission and anchor the deadline at receipt."""
         self.draining = True
         self.executor.draining = True
         logger.info("drain started (deadline_ms=%d)", deadline_ms)
+        deadline_s = (deadline_ms / 1000.0) if deadline_ms > 0 else None
+        loop = asyncio.get_running_loop()
+        self._drain_deadline_at = loop.time() + deadline_s if deadline_s is not None else None
+        if self._drain_deadline_at is None and self.transport is not None:
+            self.transport.begin_non_expiring_drain()
+
+    async def _finish_drain(self) -> None:
+        deadline_at = self._drain_deadline_at
+        loop = asyncio.get_running_loop()
         await self.maybe_send_state_delta()
 
-        deadline_s = (deadline_ms / 1000.0) if deadline_ms > 0 else None
-        finished = await self.executor.wait_idle(timeout=deadline_s)
+        wait_timeout = None if deadline_at is None else max(0.0, deadline_at - loop.time())
+        finished = await self.executor.wait_idle(timeout=wait_timeout)
         if not finished:
             logger.warning("drain deadline expired; aborting remaining jobs as RETRYABLE")
             await self.executor.abort_all(safe_message="worker draining")
 
         await self.executor.shutdown_instances()
         if self.transport is not None:
-            await self.transport.close_after_flush(timeout=30.0)
+            flush_timeout = None if deadline_at is None else max(0.0, deadline_at - loop.time())
+            await self.transport.close_after_flush(timeout=flush_timeout)
         self.drained.set()
         logger.info("drain complete")

src/gen_worker/transport.py

@@ -188,6 +188,7 @@ def __init__(
         self._stopping = asyncio.Event()
         self._connected = asyncio.Event()
         self._clean_close = False
+        self._non_expiring_drain = False
         self.reconnect_delays: List[float] = []  # observability + tests
         self._consecutive_auth_failures = 0
         self._first_auth_failure_at: Optional[float] = None
@@ -221,12 +222,21 @@ async def wait_connected(self, timeout: Optional[float] = None) -> bool:
 
     # ---- drain / shutdown --------------------------------------------------
 
-    async def close_after_flush(self, timeout: float = 30.0) -> None:
-        """Drain path: ship everything queued, then end the stream and stop."""
-        await self.queue.wait_empty(timeout=timeout)
-        self._clean_close = True
+    def begin_non_expiring_drain(self) -> None:
+        """Keep reconnecting until a zero-deadline drain ships its results."""
+        self._non_expiring_drain = True
+
+    async def close_after_flush(self, timeout: Optional[float] = None) -> bool:
+        """Ship the queue, then stop; ``None`` waits until every result ships."""
+        flushed = await self.queue.wait_empty(timeout=timeout)
+        self._clean_close = flushed
+        if not flushed:
+            logger.warning(
+                "drain flush deadline expired; closing abruptly for hub reconciliation"
+            )
         self._stopping.set()
         await self.queue.notify()
+        return flushed
 
     def stop(self) -> None:
         self._stopping.set()
@@ -340,7 +350,11 @@ async def run(self) -> None:
                 if now - self._connected_at >= _BACKOFF_RESET_AFTER_S:
                     attempt = 0
             timeout_s = float(self._settings.worker_disconnected_timeout_s or 0)
-            if timeout_s > 0 and now - last_connected > timeout_s:
+            if (
+                timeout_s > 0
+                and now - last_connected > timeout_s
+                and not self._non_expiring_drain
+            ):
                 raise FatalTransportError(
                     f"no successful connection for {timeout_s:.0f}s; exiting for reap"
                 )

tests/test_drain_flush.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import asyncio
+
+from gen_worker.config import Settings
+from gen_worker.executor import Executor
+from gen_worker.lifecycle import Lifecycle
+from gen_worker.pb import worker_scheduler_pb2 as pb
+from gen_worker.transport import FatalTransportError, Transport
+
+
+class _IdleExecutor(Executor):
+    def __init__(self, idle_delay: float = 0.0) -> None:
+        self.draining = False
+        self.idle_delay = idle_delay
+        self.wait_timeout: float | None = None
+        self.shutdown_started = asyncio.Event()
+
+    async def wait_idle(self, timeout: float | None = None) -> bool:
+        self.wait_timeout = timeout
+        await asyncio.sleep(self.idle_delay)
+        return True
+
+    async def abort_all(self, safe_message: str = "worker draining") -> None:
+        raise AssertionError("idle executor must not be aborted")
+
+    async def shutdown_instances(self) -> None:
+        self.shutdown_started.set()
+
+
+class _DrainLifecycle(Lifecycle):
+    def __init__(self, executor: _IdleExecutor, transport: Transport) -> None:
+        self.executor = executor
+        self.transport = transport
+        self.draining = False
+        self.drained = asyncio.Event()
+
+    async def maybe_send_state_delta(self) -> None:
+        pass
+
+
+class _RecordingTransport(Transport):
+    def __init__(self) -> None:
+        super().__init__(Settings(orchestrator_public_addr="localhost:1"), object())
+        self.flush_timeout: float | None = None
+
+    async def close_after_flush(self, timeout: float | None = None) -> bool:
+        self.flush_timeout = timeout
+        return await super().close_after_flush(timeout)
+
+
+class _DisconnectHandlers:
+    async def on_disconnect(self) -> None:
+        pass
+
+
+class _ReconnectTransport(Transport):
+    def __init__(self) -> None:
+        super().__init__(
+            Settings(
+                orchestrator_public_addr="localhost:1",
+                worker_disconnected_timeout_s=1,
+            ),
+            _DisconnectHandlers(),
+            backoff_base_s=0.0,
+            backoff_cap_s=0.0,
+        )
+        self.connect_attempts = 0
+
+    async def _connect_once(self, target: str, use_tls: bool) -> None:
+        self.connect_attempts += 1
+        if self.connect_attempts == 1:
+            await asyncio.sleep(1.01)
+            raise ConnectionError("orchestrator unavailable")
+        kind, result = await self.queue.get()
+        assert kind == "result"
+        await self.queue.mark_result_shipped(result)
+        await asyncio.sleep(0)
+
+
+def _result() -> pb.WorkerMessage:
+    return pb.WorkerMessage(job_result=pb.JobResult(
+        request_id="request-1", attempt=1, status=pb.JOB_STATUS_OK,
+    ))
+
+
+def test_non_expiring_drain_waits_for_blocked_result_queue() -> None:
+    async def _run() -> None:
+        transport = Transport(Settings(orchestrator_public_addr="localhost:1"), object())
+        executor = _IdleExecutor()
+        lifecycle = _DrainLifecycle(executor, transport)
+        await transport.send(_result())
+
+        drain = asyncio.create_task(lifecycle.drain(0))
+        await executor.shutdown_started.wait()
+        await asyncio.sleep(0)
+
+        assert not drain.done()
+        assert not transport._stopping.is_set()
+        kind, result = await transport.queue.get()
+        assert kind == "result"
+        await transport.queue.mark_result_shipped(result)
+
+        await asyncio.wait_for(drain, 1.0)
+        assert lifecycle.drained.is_set()
+        assert transport._stopping.is_set()
+        assert transport._clean_close
+
+    asyncio.run(_run())
+
+
+def test_explicit_drain_deadline_still_bounds_blocked_result_queue() -> None:
+    async def _run() -> None:
+        transport = _RecordingTransport()
+        executor = _IdleExecutor(idle_delay=0.02)
+        lifecycle = _DrainLifecycle(executor, transport)
+        await transport.send(_result())
+
+        await asyncio.wait_for(lifecycle.drain(100), 1.0)
+
+        assert lifecycle.drained.is_set()
+        assert transport._stopping.is_set()
+        assert not transport._clean_close
+        assert not transport._non_expiring_drain
+        assert transport.queue.pending_result_keys == [("request-1", 1)]
+        assert executor.wait_timeout is not None
+        assert transport.flush_timeout is not None
+        assert transport.flush_timeout < executor.wait_timeout - 0.01
+
+    asyncio.run(_run())
+
+
+def test_start_drain_stops_admission_and_anchors_deadline_synchronously() -> None:
+    async def _run() -> None:
+        transport = _RecordingTransport()
+        executor = _IdleExecutor()
+        lifecycle = _DrainLifecycle(executor, transport)
+        loop = asyncio.get_running_loop()
+
+        before = loop.time()
+        lifecycle.start_drain(100)
+
+        assert lifecycle.draining
+        assert executor.draining
+        assert lifecycle._drain_deadline_at is not None
+        received_at = lifecycle._drain_deadline_at - 0.1
+        assert before <= received_at <= loop.time()
+        assert lifecycle._drain_task is not None
+        await asyncio.wait_for(lifecycle._drain_task, 1.0)
+
+    asyncio.run(_run())
+
+
+def test_non_expiring_drain_reconnects_past_disconnected_timeout() -> None:
+    async def _run() -> None:
+        transport = _ReconnectTransport()
+        executor = _IdleExecutor()
+        lifecycle = _DrainLifecycle(executor, transport)
+        await transport.send(_result())
+
+        drain = asyncio.create_task(lifecycle.drain(0))
+        await executor.shutdown_started.wait()
+        run = asyncio.create_task(transport.run())
+
+        await asyncio.wait_for(asyncio.gather(drain, run), 3.0)
+        assert transport.connect_attempts == 2
+        assert lifecycle.drained.is_set()
+        assert transport._clean_close
+        assert transport.queue.pending_result_keys == []
+
+    asyncio.run(_run())
+
+
+def test_disconnected_timeout_still_applies_without_non_expiring_drain() -> None:
+    async def _run() -> None:
+        transport = _ReconnectTransport()
+        try:
+            await transport.run()
+        except FatalTransportError as exc:
+            assert "no successful connection" in str(exc)
+        else:
+            raise AssertionError("ordinary disconnect unexpectedly retried past its lifetime")
+        assert transport.connect_attempts == 1
+
+    asyncio.run(_run())

tests/test_worker_grpc_e2e.py

@@ -323,7 +323,6 @@ def test_full_contract_round_trip(scheduler_and_worker) -> None:
         input_payload=_msgpack("marco")))
     conn3.wait_for(_is_accept_for("r-last"))
     conn3.send(drain=pb.Drain(deadline_ms=5000))
-    time.sleep(0.05)
     conn3.send(run_job=pb.RunJob(
         request_id="r-after-drain", attempt=1, function_name="echo",
         input_payload=_msgpack("marco")))