From eb110dce55ad08b111f1f1e8129647f6e102a208 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Tue, 14 Jul 2026 02:03:37 -0600 Subject: [PATCH] worker: make drain deadlines explicit --- proto/CONTRACT.md | 2 +- proto/worker_scheduler.proto | 2 +- src/gen_worker/lifecycle.py | 32 ++++-- src/gen_worker/transport.py | 24 ++++- tests/test_drain_flush.py | 185 ++++++++++++++++++++++++++++++++++ tests/test_worker_grpc_e2e.py | 1 - 6 files changed, 230 insertions(+), 16 deletions(-) create mode 100644 tests/test_drain_flush.py diff --git a/proto/CONTRACT.md b/proto/CONTRACT.md index 67c56af..582a992 100644 --- a/proto/CONTRACT.md +++ b/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, diff --git a/proto/worker_scheduler.proto b/proto/worker_scheduler.proto index 395049b..26ee666 100644 --- a/proto/worker_scheduler.proto +++ b/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 } // --------------------------------------------------------------------------- diff --git a/src/gen_worker/lifecycle.py b/src/gen_worker/lifecycle.py index 5c34b6a..d10c33b 100644 --- a/src/gen_worker/lifecycle.py +++ b/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") diff --git a/src/gen_worker/transport.py b/src/gen_worker/transport.py index 9a542ea..e63b071 100644 --- a/src/gen_worker/transport.py +++ b/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" ) diff --git a/tests/test_drain_flush.py b/tests/test_drain_flush.py new file mode 100644 index 0000000..77ca104 --- /dev/null +++ b/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()) diff --git a/tests/test_worker_grpc_e2e.py b/tests/test_worker_grpc_e2e.py index 7752c54..3b69d6a 100644 --- a/tests/test_worker_grpc_e2e.py +++ b/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")))