diff --git a/CHANGELOG.md b/CHANGELOG.md index 0047a6de..51d1d12a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- `Worker.status()` returns a read-only `WorkerStatus` snapshot for health + endpoints and orchestrator probes. Readiness requires this worker's own + acknowledged registration, running pollers and a successful latest Server + exchange per role; liveness reports event-loop progress, exited pollers and + stalled requests. Snapshots perform no network I/O and can be read from + another thread. Adds `WorkerState` and `WorkerRoleStatus`, a runnable probe + example and a Kubernetes guide. Existing worker behavior is unchanged. + ## [2.12.4] - 2026-11-02 ### Added diff --git a/docs/guides/worker-health.md b/docs/guides/worker-health.md new file mode 100644 index 00000000..9bd1c1a7 --- /dev/null +++ b/docs/guides/worker-health.md @@ -0,0 +1,128 @@ +# Worker readiness and liveness + +`Worker.status()` returns a read-only `WorkerStatus` snapshot for an +application health endpoint or an orchestrator probe. It reports this worker's +own registration and local progress. A healthy Server, or another worker on the +same queue, does not make this worker ready. + +This API is in unreleased SDK source. Published packages need the release +that includes it. + +## Ready versus live + +| | Ready | Live | +|---|---|---| +| Question | Should this worker count as serving its queue? | Is this worker's local machinery progressing? | +| True when | Registration is acknowledged for its namespace and queue, every configured poller is running, the latest Server exchange of every role succeeded, and no request is overdue | The event loop advanced recently, no configured poller has exited, and no Server request has outlived its stall bound | +| Transient Server or network outage | `False` (`server_unavailable:`) | `True` | +| Use it for | Readiness probes, load-balancer membership, rollout gates | Liveness probes that restart a stuck process | + +The roles are `heartbeat`, plus `workflow_poll`, `activity_poll` and +`query_poll` when the worker registers workflows, activities or Server query +tasks. Readiness returns without a restart on the next successful exchange of +each failed role. A failed heartbeat stays unready until a later heartbeat +succeeds, so recovery can take up to one heartbeat interval. + +An idle long poll is healthy. Two bounds apply to an outstanding poll or +heartbeat: + +- `request_overdue_seconds` (readiness) defaults to one HTTP attempt deadline, + the larger of the client request timeout and the poll's absolute HTTP + deadline. A request unanswered for longer reports `server_unavailable` even + while the client is still retrying it, because the Server did not answer in + time. +- `request_stall_seconds` (liveness) defaults to twice the client's complete + transport retry budget for one request. Only a request stuck beyond any + legitimate retry window marks the worker not live, so a Server outage alone + does not trigger a liveness restart. + +Long-running activity handlers never affect either signal. +Poll-capacity backpressure is a Server answer for an admitted worker and keeps +the worker ready. + +## States + +| `state` | `ready` | `live` | Meaning | +|---|---|---|---| +| `starting` | `False` (`not_started` or `registering`) | `False` before `run()`, then `True` | Constructed, or waiting for the registration acknowledgement | +| `registered` | per the rules above | per the rules above | Polling and heartbeating | +| `draining` | `False` (`draining`) | `True` while the loop progresses | `stop()` was requested; in-flight work is draining | +| `stopped` | `False` | `False` | Shutdown finished and the registration was removed | +| `failed` | `False` | `False` (`failed:registration`, `failed:runtime` or `failed:shutdown`) | Registration was refused, a poller failed, or shutdown could not confirm cleanup | + +A dead query-poller thread, a stalled event loop, a registration refusal and a +completed shutdown never report ready. + +## Checking a stalled event loop + +`status()` performs no network I/O, claims no tasks, renews no leases and does +not wait on the event loop. While the worker runs, a small local task advances a +progress tick once per second. Serve probes from a separate thread, as in the +example below, and the probe can answer even when the event loop is blocked: +`live` turns `False` with `event_loop_stalled` once the tick is older than +`event_loop_stall_seconds` (default 30). A probe handler that runs on the +worker's own event loop cannot report its own stall. + +```python +status = worker.status() +status.ready, status.ready_reason # (True, "ready") +status.live, status.live_reason # (True, "live") +status.as_dict() # JSON-serializable; identity and progress only +``` + +The snapshot contains the worker ID, namespace, task queue, state, reasons, +ages in seconds, and per-role progress. It never contains payloads, task +arguments, exception messages (only exception class names) or credentials. + +## Kubernetes example + +`examples/worker_health_probe.py` runs a worker with a stdlib HTTP server on its +own thread, serving `/readyz`, `/livez` and `/status`. Each answers 200 or 503 +with the snapshot as JSON. Bind it to the pod network only. + +```yaml +containers: + - name: worker + image: example/python-worker:1.0.0 + env: + - name: DW_HEALTH_HOST + value: "0.0.0.0" + ports: + - name: health + containerPort: 8081 + startupProbe: + httpGet: {path: /readyz, port: health} + periodSeconds: 2 + failureThreshold: 60 + readinessProbe: + httpGet: {path: /readyz, port: health} + periodSeconds: 5 + failureThreshold: 2 + livenessProbe: + httpGet: {path: /livez, port: health} + periodSeconds: 10 + failureThreshold: 6 +terminationGracePeriodSeconds: 45 +``` + +Give the liveness probe a failure window longer than ordinary event-loop pauses, +such as a large replay. Restarting a worker that is only unready because the +Server is unreachable does not help, so don't probe liveness with `/readyz`. + +## Shutdown + +Three different timers apply. Keep them distinct: + +- `Worker(shutdown_timeout=...)` is how long `stop()` drains in-flight work + before cancelling it and removing the registration. +- The orchestrator's stop grace (`terminationGracePeriodSeconds`) is how long + the process may run after SIGTERM before it is killed. Make it longer than + `shutdown_timeout`, plus time for deregistration. +- Handler termination is separate again. Cancelling an await cannot stop a + running synchronous activity thread; if one is still running at the + deadline, `stop()` raises and the status becomes `failed:shutdown`. + +On SIGTERM, call `worker.stop()` and await the run task. While draining, the +worker reports `draining` (not ready, still live), so traffic moves away +without a liveness restart interrupting the drain. See +[Lifecycle and shutdown](../reference/worker.md#lifecycle-and-shutdown). diff --git a/docs/reference/worker.md b/docs/reference/worker.md index c05043ce..adf579d0 100644 --- a/docs/reference/worker.md +++ b/docs/reference/worker.md @@ -197,4 +197,15 @@ before its reply reaches the SDK. A worker that never receives that reply has not begun polling and becomes stale under the Server's heartbeat policy. Use the existing workflow handle and queue when bringing up a replacement. +## Readiness and liveness + +`Worker.status()` returns a read-only snapshot that separates readiness +(registered, every configured poller running, latest Server exchange of each +role succeeded) from liveness (event loop progressing, no poller exited, no +request outliving its bound). It performs no network I/O and can be called from +another thread. See [Worker readiness and liveness](../guides/worker-health.md) +for states, probe configuration and shutdown guidance. + ::: durable_workflow.worker + +::: durable_workflow.worker_status diff --git a/examples/worker_health_probe.py b/examples/worker_health_probe.py new file mode 100644 index 00000000..ff26f432 --- /dev/null +++ b/examples/worker_health_probe.py @@ -0,0 +1,114 @@ +"""Serve worker readiness and liveness probes next to a long-running worker. + +Run against a Server: + + DURABLE_WORKFLOW_SERVER_URL=http://localhost:8080 \\ + DURABLE_WORKFLOW_AUTH_TOKEN=dev-token-123 \\ + python examples/worker_health_probe.py + +Then probe it: + + curl -i http://127.0.0.1:8081/readyz + curl -i http://127.0.0.1:8081/livez + +The probe server runs on its own thread. ``Worker.status()`` performs no +network I/O and does not need the worker's event loop, so ``/livez`` still +answers, with 503, when that loop is stalled. Bind the probe port only to the +pod or host network your orchestrator uses; the response carries worker +identity and progress, not payloads or credentials. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import signal +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +from durable_workflow import Client, Worker, activity, workflow + + +@activity.defn(name="greet") +def greet(name: str) -> str: + return f"hello, {name}" + + +@workflow.defn(name="greeter") +class GreeterWorkflow: + def run(self, ctx, name): # type: ignore[no-untyped-def] + result = yield ctx.schedule_activity("greet", [name]) + return result + + +def start_probe_server(worker: Worker, host: str, port: int) -> ThreadingHTTPServer: + class ProbeHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: # noqa: N802 - http.server API + status = worker.status() + if self.path == "/readyz": + healthy = status.ready + elif self.path == "/livez": + healthy = status.live + elif self.path == "/status": + healthy = True + else: + self.send_error(404) + return + body = json.dumps(status.as_dict()).encode() + self.send_response(200 if healthy else 503) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: # noqa: A002 - http.server API + return + + server = ThreadingHTTPServer((host, port), ProbeHandler) + threading.Thread(target=server.serve_forever, name="worker-health-probe", daemon=True).start() + return server + + +async def main() -> None: + async with Client( + os.environ.get("DURABLE_WORKFLOW_SERVER_URL", "http://localhost:8080"), + token=os.environ.get("DURABLE_WORKFLOW_AUTH_TOKEN"), + namespace=os.environ.get("DURABLE_WORKFLOW_NAMESPACE", "default"), + ) as client: + worker = Worker( + client, + task_queue=os.environ.get("DURABLE_WORKFLOW_TASK_QUEUE", "python-workers"), + workflows=[GreeterWorkflow], + activities=[greet], + # Keep this below the orchestrator's termination grace period. + shutdown_timeout=25.0, + ) + probe = start_probe_server( + worker, + os.environ.get("DW_HEALTH_HOST", "127.0.0.1"), + int(os.environ.get("DW_HEALTH_PORT", "8081")), + ) + shutdown_requested = asyncio.Event() + loop = asyncio.get_running_loop() + for signum in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(signum, shutdown_requested.set) + + run_task = asyncio.create_task(worker.run()) + try: + await asyncio.wait( + {run_task, asyncio.create_task(shutdown_requested.wait())}, + return_when=asyncio.FIRST_COMPLETED, + ) + finally: + try: + await worker.stop() + finally: + try: + await run_task + finally: + probe.shutdown() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/mkdocs.yml b/mkdocs.yml index 52eb995f..e4c03548 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -79,6 +79,7 @@ nav: - Execution telemetry: telemetry.md - Sticky execution: guides/sticky-execution.md - Separate worker processes: guides/separate-workers.md + - Worker readiness and liveness: guides/worker-health.md - Query deadlines: query-task-authority.md - SDK guide: https://durable-workflow.com/docs/2.0/polyglot/python/ - PyPI: https://pypi.org/project/durable-workflow/ diff --git a/scripts/smoke-built-package.py b/scripts/smoke-built-package.py index 0c739aaa..1ad3b696 100644 --- a/scripts/smoke-built-package.py +++ b/scripts/smoke-built-package.py @@ -112,6 +112,7 @@ def completed(result): "durable_workflow.sync", "durable_workflow.testing", "durable_workflow.worker", + "durable_workflow.worker_status", "durable_workflow.workflow", "durable_workflow.workflow_updates_conformance", ] diff --git a/src/durable_workflow/__init__.py b/src/durable_workflow/__init__.py index dde93989..0dfd811c 100644 --- a/src/durable_workflow/__init__.py +++ b/src/durable_workflow/__init__.py @@ -225,6 +225,7 @@ ) from .worker import Worker from .worker_session import WorkerSession, WorkerSessionOptions +from .worker_status import WorkerRoleStatus, WorkerState, WorkerStatus from .workflow import ( ActivityRetryPolicy, ChildWorkflowRetryPolicy, @@ -314,6 +315,9 @@ "Worker", "WorkerSession", "WorkerSessionOptions", + "WorkerRoleStatus", + "WorkerState", + "WorkerStatus", "WorkerInterceptor", "WorkflowCommandResult", "WorkflowExecution", diff --git a/src/durable_workflow/worker.py b/src/durable_workflow/worker.py index fd13a74c..dc4a3615 100644 --- a/src/durable_workflow/worker.py +++ b/src/durable_workflow/worker.py @@ -22,6 +22,7 @@ import inspect import json import logging +import math import pickle import random import sys @@ -68,6 +69,7 @@ WorkflowExecution, _protocol_version_from_env, _supports_cooperative_cancellation_protocol, + _worker_poll_http_timeout, ) from .errors import ( ActivityCancelled, @@ -108,6 +110,12 @@ from .telemetry import ExecutionTelemetry, current_execution_telemetry from .telemetry_carrier import identity as observability_identity from .worker_session import WorkerSession, WorkerSessionOptions +from .worker_status import ( + DEFAULT_EVENT_LOOP_STALL_SECONDS, + EVENT_LOOP_TICK_SECONDS, + WorkerStatus, + _WorkerStatusTracker, +) from .workflow import ( Command, FailWorkflow, @@ -233,6 +241,15 @@ def _is_storage_admission_error(error: BaseException) -> bool: ) +def _positive_seconds(value: object) -> bool: + return ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and math.isfinite(value) + and value > 0 + ) + + def _poll_error_summary(error: BaseException) -> str: error_class = type(error).__name__ if isinstance(error, Unauthorized) or isinstance(error, ServerError) and error.status == 401: @@ -1209,6 +1226,8 @@ def __init__( self._query_thread: threading.Thread | None = None self._query_thread_loop: asyncio.AbstractEventLoop | None = None self._query_thread_task: asyncio.Task[None] | None = None + self._status = _WorkerStatusTracker() + self._status_tick_task: asyncio.Task[None] | None = None # In-flight slot accounting feeds the periodic heartbeat so operators # see free-slot counts without the worker having to re-derive them at # shutdown. A workflow reservation covers the multiplexed long poll, @@ -1514,6 +1533,7 @@ async def _register(self) -> None: advertised = ack.get("heartbeat_interval_seconds") if isinstance(advertised, int) and advertised > 0: self._heartbeat_interval = float(advertised) + self._status.registered() log.info("worker %s registered on %s", self.worker_id, self.task_queue) async def _run_workflow_task(self, task: dict[str, Any]) -> list[dict[str, Any]] | None: @@ -3793,6 +3813,7 @@ async def _poll_workflow_tasks(self) -> None: return try: poll_start = time.perf_counter() + self._status.request_started("workflow_poll") task = await self.client.poll_workflow_task( worker_id=self.worker_id, task_queue=self.task_queue, @@ -3802,15 +3823,18 @@ async def _poll_workflow_tasks(self) -> None: history_page_size=WORKFLOW_HISTORY_PAGE_SIZE, ) except asyncio.CancelledError: + self._status.request_abandoned("workflow_poll") self._release_workflow_capacity() raise except Exception as e: self._release_workflow_capacity() if self._stop.is_set(): + self._status.request_abandoned("workflow_poll") return + delay = _poll_capacity_delay(e, "workflow_task", self.task_queue) + self._record_poll_status("workflow_poll", e, answered=delay is not None) if _is_storage_admission_error(e): raise - delay = _poll_capacity_delay(e, "workflow_task", self.task_queue) if delay is not None: self._record_poll_metrics("workflow", "backpressure", time.perf_counter() - poll_start) log.debug("workflow poll capacity backpressure; retrying in %d s", delay) @@ -3820,6 +3844,7 @@ async def _poll_workflow_tasks(self) -> None: await asyncio.sleep(failures.retry_delay(e, "workflow")) continue failures.reset() + self._status.request_succeeded("workflow_poll") if task is None: self._release_workflow_capacity() self._record_poll_metrics("workflow", "empty", time.perf_counter() - poll_start) @@ -3905,6 +3930,7 @@ async def _poll_activity_tasks(self) -> None: return try: poll_start = time.perf_counter() + self._status.request_started("activity_poll") task = await self.client.poll_activity_task( worker_id=self.worker_id, task_queue=self.task_queue, @@ -3912,15 +3938,18 @@ async def _poll_activity_tasks(self) -> None: build_id=self.build_id, ) except asyncio.CancelledError: + self._status.request_abandoned("activity_poll") self._act_semaphore.release() raise except Exception as e: self._act_semaphore.release() if self._stop.is_set(): + self._status.request_abandoned("activity_poll") return + delay = _poll_capacity_delay(e, "activity_task", self.task_queue) + self._record_poll_status("activity_poll", e, answered=delay is not None) if _is_storage_admission_error(e): raise - delay = _poll_capacity_delay(e, "activity_task", self.task_queue) if delay is not None: self._record_poll_metrics("activity", "backpressure", time.perf_counter() - poll_start) log.debug("activity poll capacity backpressure; retrying in %d s", delay) @@ -3930,6 +3959,7 @@ async def _poll_activity_tasks(self) -> None: await asyncio.sleep(failures.retry_delay(e, "activity")) continue failures.reset() + self._status.request_succeeded("activity_poll") if task is None: self._act_semaphore.release() self._record_poll_metrics("activity", "empty", time.perf_counter() - poll_start) @@ -3965,20 +3995,26 @@ async def _poll_query_tasks(self, *, client: Client | None = None, track_tasks: ): try: poll_start = time.perf_counter() + self._status.request_started("query_poll") task = await client.poll_query_task( worker_id=self.worker_id, task_queue=self.task_queue, timeout=self._poll_timeout, build_id=self.build_id, ) + except asyncio.CancelledError: + self._status.request_abandoned("query_poll") + raise except Exception as e: if self._stop.is_set() or ( query_thread_stop is not None and query_thread_stop.is_set() ): + self._status.request_abandoned("query_poll") return + delay = _poll_capacity_delay(e, "query_task", self.task_queue) + self._record_poll_status("query_poll", e, answered=delay is not None) if _is_storage_admission_error(e): raise - delay = _poll_capacity_delay(e, "query_task", self.task_queue) if delay is not None: self._record_poll_metrics("query", "backpressure", time.perf_counter() - poll_start) log.debug("query poll capacity backpressure; retrying in %d s", delay) @@ -3988,6 +4024,7 @@ async def _poll_query_tasks(self, *, client: Client | None = None, track_tasks: await asyncio.sleep(failures.retry_delay(e, "query")) continue failures.reset() + self._status.request_succeeded("query_poll") if task is None: self._record_poll_metrics("query", "empty", time.perf_counter() - poll_start) if self._stop.is_set() or ( @@ -4231,6 +4268,8 @@ def _run_query_task_thread(self) -> None: log.exception("query task poller thread was cancelled unexpectedly") except Exception: log.exception("query task poller thread stopped unexpectedly") + finally: + self._status.role_exited("query_poll") @_with_storage_admission_retries async def _query_task_thread_main(self) -> None: @@ -4274,24 +4313,12 @@ async def _stop_query_task_thread(self, *, deadline: float) -> None: async def run(self) -> None: """Register the worker and poll until `stop()` is called or the task is cancelled.""" self._begin_run() - try: - await self._register() - finally: - self._registration_done.set() + await self._register_for_run() try: if self._stop.is_set(): return - loops = [asyncio.create_task(self._heartbeat_loop())] - if self.workflows: - loops.append(asyncio.create_task(self._poll_workflow_tasks())) - if self.activities: - loops.append(asyncio.create_task(self._poll_activity_tasks())) - if self.workflows and self._query_tasks_supported: - if self._can_clone_client_for_query_tasks(): - self._start_query_task_thread() - else: - loops.append(asyncio.create_task(self._poll_query_tasks())) + loops = self._start_run_pollers() self._poller_tasks.update(loops) try: @@ -4302,6 +4329,8 @@ async def run(self) -> None: finally: primary_error = sys.exc_info()[1] primary_traceback = primary_error.__traceback__ if primary_error is not None else None + if isinstance(primary_error, Exception): + self._status.failed("runtime") try: await self.stop() except BaseException as shutdown_error: @@ -4309,6 +4338,35 @@ async def run(self) -> None: raise primary_error.with_traceback(primary_traceback) from shutdown_error raise + async def _register_for_run(self) -> None: + try: + await self._register() + except BaseException as error: + if isinstance(error, Exception): + self._status.failed("registration") + if self._shutdown_task is None: + # run() leaves no shutdown behind a failed registration. + self._cancel_status_tick() + raise + finally: + self._registration_done.set() + + def _start_run_pollers(self) -> list[asyncio.Task[None]]: + roles: list[tuple[str, Coroutine[Any, Any, None] | None]] = [("heartbeat", self._heartbeat_loop())] + if self.workflows: + roles.append(("workflow_poll", self._poll_workflow_tasks())) + if self.activities: + roles.append(("activity_poll", self._poll_activity_tasks())) + query_thread = False + if self.workflows and self._query_tasks_supported: + query_thread = self._can_clone_client_for_query_tasks() + roles.append(("query_poll", None if query_thread else self._poll_query_tasks())) + self._status.roles_started(role for role, _ in roles) + loops = [self._start_status_role(role, loop) for role, loop in roles if loop is not None] + if query_thread: + self._start_query_task_thread() + return loops + async def _heartbeat_loop(self) -> None: """Periodically refresh the server-side worker registration. @@ -4325,22 +4383,36 @@ async def _heartbeat_loop(self) -> None: await asyncio.wait_for(self._stop.wait(), timeout=self._heartbeat_interval) if self._stop.is_set(): return - try: - ack = await self.client.heartbeat_worker( - worker_id=self.worker_id, - task_slots=self._current_task_slots(), - process_metrics=self._current_process_metrics(), - ) - except Exception as e: - if _is_storage_admission_error(e) and not self._stop.is_set(): - raise - log.warning("worker heartbeat failed: %s", e) - continue + ack = await self._send_heartbeat() if isinstance(ack, dict): advertised = ack.get("heartbeat_interval_seconds") if isinstance(advertised, int) and advertised > 0: self._heartbeat_interval = float(advertised) + async def _send_heartbeat(self) -> Any: + """Send one heartbeat; return its acknowledgement, or ``None`` after a logged failure.""" + self._status.request_started("heartbeat") + try: + ack = await self.client.heartbeat_worker( + worker_id=self.worker_id, + task_slots=self._current_task_slots(), + process_metrics=self._current_process_metrics(), + ) + except asyncio.CancelledError: + self._status.request_abandoned("heartbeat") + raise + except Exception as e: + if self._stop.is_set(): + self._status.request_abandoned("heartbeat") + else: + self._status.request_failed("heartbeat", e) + if _is_storage_admission_error(e) and not self._stop.is_set(): + raise + log.warning("worker heartbeat failed: %s", e) + return None + self._status.request_succeeded("heartbeat") + return ack + def _current_task_slots(self) -> dict[str, int]: return { "workflow_available": max( @@ -4442,20 +4514,12 @@ async def run_until( self._begin_run() try: - try: - await asyncio.wait_for(self._register(), timeout=timeout) - finally: - self._registration_done.set() + await self._register_within(timeout) startup_completed = True if self._stop.is_set(): raise asyncio.CancelledError - background_tasks.append(asyncio.create_task(self._heartbeat_loop())) - if self.workflows and self._query_tasks_supported: - if self._can_clone_client_for_query_tasks(): - self._start_query_task_thread() - else: - background_tasks.append(asyncio.create_task(self._poll_query_tasks())) + background_tasks.extend(self._start_run_until_background()) self._poller_tasks.update(background_tasks) run_until_loop = asyncio.create_task( self._run_until_loop( @@ -4485,6 +4549,26 @@ async def run_until( raise primary_error.with_traceback(primary_traceback) from shutdown_error raise + async def _register_within(self, timeout: float) -> None: + try: + await asyncio.wait_for(self._register(), timeout=timeout) + except Exception: + self._status.failed("registration") + raise + finally: + self._registration_done.set() + + def _start_run_until_background(self) -> list[asyncio.Task[None]]: + query_tasks = bool(self.workflows and self._query_tasks_supported) + query_thread = query_tasks and self._can_clone_client_for_query_tasks() + self._status.roles_started(("heartbeat", "query_poll") if query_tasks else ("heartbeat",)) + tasks = [self._start_status_role("heartbeat", self._heartbeat_loop())] + if query_thread: + self._start_query_task_thread() + elif query_tasks: + tasks.append(self._start_status_role("query_poll", self._poll_query_tasks())) + return tasks + async def _run_until_loop( self, *, @@ -4644,10 +4728,113 @@ def _begin_run(self) -> None: if self._stop.is_set(): raise RuntimeError("worker cannot run after stop() has been called") self._run_started = True + self._status.run_started() + self._status_tick_task = asyncio.get_running_loop().create_task(self._status_tick_loop()) payload_observers = getattr(self.client, "_workflow_payload_observers", None) if isinstance(payload_observers, dict): payload_observers[self.worker_id] = self._renew_workflow_payload_lease + def status( + self, + *, + event_loop_stall_seconds: float = DEFAULT_EVENT_LOOP_STALL_SECONDS, + request_overdue_seconds: float | None = None, + request_stall_seconds: float | None = None, + ) -> WorkerStatus: + """Return a read-only readiness and liveness snapshot of this worker. + + The snapshot performs no network I/O, claims no tasks, renews no leases + and does not wait on the worker's event loop, so a health endpoint may + call it from any thread, including while the event loop is stalled. + + Readiness follows this worker's own registration: it becomes ready + once the Server acknowledges the registration, while every configured + poller is running, the latest Server exchange of every role succeeded + and no outstanding request is older than ``request_overdue_seconds``. + A transient Server or network failure makes the worker unready; the + next successful exchange of that role restores readiness. Liveness + reports local progress: the event loop advanced within + ``event_loop_stall_seconds``, no configured poller has exited, and no + outstanding Server request is older than ``request_stall_seconds``. + Idle long polls, Server outages within the transport retry budget and + long-running activity handlers keep the worker live. Draining workers + stay live but are not ready; stopped and failed workers are neither. + + Args: + event_loop_stall_seconds: Maximum age of the event loop's progress + tick before the worker is reported not live. + request_overdue_seconds: Maximum age of one outstanding poll or + heartbeat before the worker is reported not ready. Defaults to + one HTTP attempt deadline: the larger of the client request + timeout and the poll's absolute HTTP deadline. + request_stall_seconds: Maximum age of one outstanding poll or + heartbeat before the worker is reported not live. Defaults to + twice the client's complete transport retry budget for one + such request. + + Returns: + A :class:`~durable_workflow.WorkerStatus` snapshot. + """ + if not _positive_seconds(event_loop_stall_seconds): + raise ValueError("event_loop_stall_seconds must be a positive finite number") + if request_overdue_seconds is None: + request_overdue_seconds = self._request_attempt_deadline_seconds() + elif not _positive_seconds(request_overdue_seconds): + raise ValueError("request_overdue_seconds must be a positive finite number") + if request_stall_seconds is None: + request_stall_seconds = self._default_request_stall_seconds() + elif not _positive_seconds(request_stall_seconds): + raise ValueError("request_stall_seconds must be a positive finite number") + namespace = self._registration_namespace if self._registration_worker_id is not None else getattr( + self.client, "namespace", None + ) + return self._status.snapshot( + worker_id=self.worker_id, + namespace=namespace if isinstance(namespace, str) else None, + task_queue=self.task_queue, + event_loop_stall_seconds=float(event_loop_stall_seconds), + request_overdue_seconds=float(request_overdue_seconds), + request_stall_seconds=float(request_stall_seconds), + ) + + def _request_attempt_deadline_seconds(self) -> float: + client_timeout = getattr(self.client, "timeout", None) + request_timeout = ( + float(client_timeout) + if isinstance(client_timeout, (int, float)) and _positive_seconds(client_timeout) + else 60.0 + ) + poll_deadline = _worker_poll_http_timeout(self._poll_timeout) or 0.0 + return max(request_timeout, poll_deadline) + + def _default_request_stall_seconds(self) -> float: + # One logical request may retry transport failures; only a request + # outliving twice that complete budget indicates a local stall. + policy = getattr(self.client, "retry_policy", None) + attempts = getattr(policy, "max_attempts", 3) + backoff = getattr(policy, "max_backoff_seconds", 5.0) + attempts = attempts if isinstance(attempts, int) and attempts >= 1 else 3 + backoff = float(backoff) if isinstance(backoff, (int, float)) and _positive_seconds(backoff) else 5.0 + budget = attempts * self._request_attempt_deadline_seconds() + (attempts - 1) * backoff + return 2.0 * budget + + async def _status_tick_loop(self) -> None: + while True: + self._status.loop_tick() + await asyncio.sleep(EVENT_LOOP_TICK_SECONDS) + + def _start_status_role(self, role: str, loop: Coroutine[Any, Any, None]) -> asyncio.Task[None]: + task = asyncio.create_task(loop) + task.add_done_callback(lambda _: self._status.role_exited(role)) + return task + + def _record_poll_status(self, role: str, error: Exception, *, answered: bool) -> None: + # Capacity backpressure is a Server answer for an admitted worker. + if answered: + self._status.request_succeeded(role) + else: + self._status.request_failed(role, error) + def worker_session(self, options: WorkerSessionOptions) -> WorkerSession: """Return a typed lease handle held by this worker until shutdown.""" existing = self._worker_sessions.get(options.session_id) @@ -4679,11 +4866,28 @@ async def stop(self) -> None: stays within the original shutdown deadline. Unfenced deregistration makes one attempt and does not retry an uncertain outcome. """ + self._status.draining() self._stop.set() if self._shutdown_task is None: - self._shutdown_task = asyncio.create_task(self._shutdown()) + self._shutdown_task = asyncio.create_task(self._shutdown_and_record_status()) await asyncio.shield(self._shutdown_task) + async def _shutdown_and_record_status(self) -> None: + try: + await self._shutdown() + except BaseException: + self._status.failed("shutdown") + raise + else: + self._status.stopped() + finally: + self._cancel_status_tick() + + def _cancel_status_tick(self) -> None: + tick = self._status_tick_task + if tick is not None and not tick.done(): + tick.cancel() + async def _shutdown(self) -> None: if self._run_started and not self._registration_done.is_set(): await self._registration_done.wait() diff --git a/src/durable_workflow/worker_status.py b/src/durable_workflow/worker_status.py new file mode 100644 index 00000000..453f32c9 --- /dev/null +++ b/src/durable_workflow/worker_status.py @@ -0,0 +1,385 @@ +"""Read-only worker readiness and liveness status. + +:meth:`durable_workflow.Worker.status` returns a :class:`WorkerStatus` +snapshot that an application health endpoint or orchestrator probe can serve. +Taking a snapshot never performs network I/O, claims tasks, renews leases or +waits on the worker's event loop, so it is safe to call from another thread, +including when the worker's event loop itself is stalled. + +Readiness and liveness answer different questions: + +- **Ready** means this worker holds a current registration for its own + namespace and task queue, every configured poller is running, the most + recent Server exchange of every role succeeded, and no outstanding request + is overdue (unanswered beyond one HTTP attempt deadline). A transient Server + or network outage makes the worker unready without making it unlive, and + the next successful exchange of each affected role restores readiness + without a restart. +- **Live** means the worker's local machinery is progressing: its event loop + is advancing, no configured poller has exited, and no outstanding Server + request has outlived its stall bound, which exceeds the transport retry + budget. An idle long poll, a Server outage within that budget, or a + long-running activity handler is not a liveness failure. + +Missing observations are reported as unavailable (``None``), never as +success. +""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Iterable +from dataclasses import dataclass +from enum import Enum +from typing import Any + +__all__ = ["WorkerRoleStatus", "WorkerState", "WorkerStatus"] + +# Event-loop progress is sampled at this cadence by a local task that performs +# no I/O. Probes compare the age of the latest sample with their stall bound. +EVENT_LOOP_TICK_SECONDS = 1.0 +DEFAULT_EVENT_LOOP_STALL_SECONDS = 30.0 + + +class WorkerState(str, Enum): + """Lifecycle state of a :class:`~durable_workflow.Worker`.""" + + STARTING = "starting" + """Constructed or registering; not yet admitted by the Server.""" + + REGISTERED = "registered" + """Registered with the Server; pollers and heartbeats are running.""" + + DRAINING = "draining" + """``stop()`` was requested; pollers stop and in-flight tasks drain.""" + + STOPPED = "stopped" + """Shutdown completed and the registration was removed.""" + + FAILED = "failed" + """Registration was refused, a poller failed, or shutdown could not confirm cleanup.""" + + +@dataclass(frozen=True) +class WorkerRoleStatus: + """Progress of one configured worker role. + + Ages are seconds on the local monotonic clock, measured when the snapshot + was taken. ``None`` means the observation has not happened yet. + """ + + role: str + """``workflow_poll``, ``activity_poll``, ``query_poll`` or ``heartbeat``.""" + + running: bool + """Whether the role's loop or thread is still running.""" + + in_flight_seconds: float | None + """Age of the role's outstanding Server request, or ``None`` when idle.""" + + last_success_age_seconds: float | None + """Age of the role's most recent successful Server exchange.""" + + last_failure_age_seconds: float | None + """Age of the role's most recent failed Server exchange.""" + + consecutive_failures: int + """Failed exchanges since the role's last success.""" + + last_failure_type: str | None + """Exception class name of the most recent failure; never its message.""" + + overdue: bool + """Whether the outstanding request is unanswered beyond one HTTP attempt deadline.""" + + stalled: bool + """Whether the outstanding request has outlived the snapshot's stall bound.""" + + def as_dict(self) -> dict[str, Any]: + """Return a JSON-serializable mapping of this role status.""" + + return { + "role": self.role, + "running": self.running, + "in_flight_seconds": self.in_flight_seconds, + "last_success_age_seconds": self.last_success_age_seconds, + "last_failure_age_seconds": self.last_failure_age_seconds, + "consecutive_failures": self.consecutive_failures, + "last_failure_type": self.last_failure_type, + "overdue": self.overdue, + "stalled": self.stalled, + } + + +@dataclass(frozen=True) +class WorkerStatus: + """Point-in-time readiness and liveness of a worker. + + The snapshot carries identity and progress only: no payloads, task + arguments, exception messages or credentials. + """ + + worker_id: str + namespace: str | None + task_queue: str + state: WorkerState + ready: bool + """Whether an orchestrator should route work to, or count, this worker.""" + + ready_reason: str + """``ready`` or the first reason the worker is not ready.""" + + live: bool + """Whether the worker's local machinery is progressing.""" + + live_reason: str + """``live`` or the first reason the worker is not live.""" + + registered_age_seconds: float | None + """Seconds since the Server acknowledged this worker's registration.""" + + event_loop_progress_age_seconds: float | None + """Seconds since the worker's event loop last advanced its progress tick.""" + + roles: tuple[WorkerRoleStatus, ...] + + def as_dict(self) -> dict[str, Any]: + """Return a JSON-serializable mapping suitable for a health endpoint.""" + + return { + "worker_id": self.worker_id, + "namespace": self.namespace, + "task_queue": self.task_queue, + "state": self.state.value, + "ready": self.ready, + "ready_reason": self.ready_reason, + "live": self.live, + "live_reason": self.live_reason, + "registered_age_seconds": self.registered_age_seconds, + "event_loop_progress_age_seconds": self.event_loop_progress_age_seconds, + "roles": [role.as_dict() for role in self.roles], + } + + +class _RoleProgress: + __slots__ = ( + "consecutive_failures", + "in_flight_since", + "last_failure_at", + "last_failure_type", + "last_success_at", + "running", + ) + + def __init__(self) -> None: + self.running = True + self.in_flight_since: float | None = None + self.last_success_at: float | None = None + self.last_failure_at: float | None = None + self.last_failure_type: str | None = None + self.consecutive_failures = 0 + + +class _WorkerStatusTracker: + """Thread-safe progress recorder owned by one worker. + + Pollers record from the worker's event loop and from the query thread; + snapshots are read from any thread. Every operation holds one lock briefly + and never blocks on I/O. + """ + + def __init__(self, *, clock: Any = time.monotonic) -> None: + self._clock = clock + self._lock = threading.Lock() + self._state = WorkerState.STARTING + self._started = False + self._registered_at: float | None = None + self._loop_tick_at: float | None = None + self._failure_reason: str | None = None + self._roles: dict[str, _RoleProgress] = {} + + def now(self) -> float: + return float(self._clock()) + + def run_started(self) -> None: + with self._lock: + self._started = True + self._loop_tick_at = self.now() + + def loop_tick(self) -> None: + with self._lock: + self._loop_tick_at = self.now() + + def registered(self) -> None: + with self._lock: + if self._state is WorkerState.STARTING: + self._state = WorkerState.REGISTERED + self._registered_at = self.now() + + def roles_started(self, roles: Iterable[str]) -> None: + with self._lock: + for role in roles: + self._roles[role] = _RoleProgress() + + def request_started(self, role: str) -> None: + with self._lock: + progress = self._roles.get(role) + if progress is not None: + progress.in_flight_since = self.now() + + def request_succeeded(self, role: str) -> None: + with self._lock: + progress = self._roles.get(role) + if progress is not None: + progress.in_flight_since = None + progress.last_success_at = self.now() + progress.consecutive_failures = 0 + + def request_failed(self, role: str, error: BaseException) -> None: + with self._lock: + progress = self._roles.get(role) + if progress is not None: + progress.in_flight_since = None + progress.last_failure_at = self.now() + progress.last_failure_type = type(error).__name__ + progress.consecutive_failures += 1 + + def request_abandoned(self, role: str) -> None: + with self._lock: + progress = self._roles.get(role) + if progress is not None: + progress.in_flight_since = None + + def role_exited(self, role: str) -> None: + with self._lock: + progress = self._roles.get(role) + if progress is not None: + progress.running = False + progress.in_flight_since = None + + def draining(self) -> None: + with self._lock: + if self._state in (WorkerState.STARTING, WorkerState.REGISTERED): + self._state = WorkerState.DRAINING + + def failed(self, reason: str) -> None: + with self._lock: + if self._state is not WorkerState.FAILED: + self._state = WorkerState.FAILED + self._failure_reason = reason + + def stopped(self) -> None: + with self._lock: + if self._state is not WorkerState.FAILED: + self._state = WorkerState.STOPPED + + def snapshot( + self, + *, + worker_id: str, + namespace: str | None, + task_queue: str, + event_loop_stall_seconds: float, + request_overdue_seconds: float, + request_stall_seconds: float, + ) -> WorkerStatus: + with self._lock: + now = self.now() + state = self._state + started = self._started + failure_reason = self._failure_reason + registered_age = _age(now, self._registered_at) + loop_age = _age(now, self._loop_tick_at) + roles = tuple( + WorkerRoleStatus( + role=role, + running=progress.running, + in_flight_seconds=_age(now, progress.in_flight_since), + last_success_age_seconds=_age(now, progress.last_success_at), + last_failure_age_seconds=_age(now, progress.last_failure_at), + consecutive_failures=progress.consecutive_failures, + last_failure_type=progress.last_failure_type, + overdue=( + progress.in_flight_since is not None + and now - progress.in_flight_since > request_overdue_seconds + ), + stalled=( + progress.in_flight_since is not None + and now - progress.in_flight_since > request_stall_seconds + ), + ) + for role, progress in self._roles.items() + ) + + live, live_reason = _liveness(state, started, failure_reason, loop_age, event_loop_stall_seconds, roles) + ready, ready_reason = _readiness(state, started, live, live_reason, roles) + return WorkerStatus( + worker_id=worker_id, + namespace=namespace, + task_queue=task_queue, + state=state, + ready=ready, + ready_reason=ready_reason, + live=live, + live_reason=live_reason, + registered_age_seconds=registered_age if state is WorkerState.REGISTERED else None, + event_loop_progress_age_seconds=loop_age, + roles=roles, + ) + + +def _age(now: float, at: float | None) -> float | None: + return None if at is None else max(0.0, now - at) + + +def _liveness( + state: WorkerState, + started: bool, + failure_reason: str | None, + loop_age: float | None, + event_loop_stall_seconds: float, + roles: tuple[WorkerRoleStatus, ...], +) -> tuple[bool, str]: + if not started: + return False, "not_started" + if state is WorkerState.FAILED: + return False, f"failed:{failure_reason}" if failure_reason else "failed" + if state is WorkerState.STOPPED: + return False, "stopped" + if loop_age is None or loop_age > event_loop_stall_seconds: + return False, "event_loop_stalled" + if state is WorkerState.REGISTERED: + return _role_liveness(roles) + return True, "live" + + +def _role_liveness(roles: tuple[WorkerRoleStatus, ...]) -> tuple[bool, str]: + stopped = next((role.role for role in roles if not role.running), None) + if stopped is not None: + return False, f"role_stopped:{stopped}" + stalled = next((role.role for role in roles if role.stalled), None) + if stalled is not None: + return False, f"role_stalled:{stalled}" + return True, "live" + + +def _readiness( + state: WorkerState, + started: bool, + live: bool, + live_reason: str, + roles: tuple[WorkerRoleStatus, ...], +) -> tuple[bool, str]: + if not started: + return False, "not_started" + if state is WorkerState.STARTING: + return False, "registering" + if state is WorkerState.DRAINING: + return False, "draining" + if state is not WorkerState.REGISTERED or not live: + return False, live_reason + for role in roles: + if role.consecutive_failures or role.overdue: + return False, f"server_unavailable:{role.role}" + return True, "ready" diff --git a/tests/fixtures/replay_regressions/worker-status-outage-recovery-cold-history.json b/tests/fixtures/replay_regressions/worker-status-outage-recovery-cold-history.json new file mode 100644 index 00000000..5260fca0 --- /dev/null +++ b/tests/fixtures/replay_regressions/worker-status-outage-recovery-cold-history.json @@ -0,0 +1,48 @@ +{ + "$schema": "https://raw.githubusercontent.com/durable-workflow/.github/main/regression-corpus/evidence-schema.json", + "fixture_schema": "durable-workflow.replay-regression/v1", + "id": "worker-status-outage-recovery-cold-history", + "protocol_version": "1.20", + "bindings": [ + "python" + ], + "workflow": { + "type": "tests.worker-status.recovery", + "input": [ + "042" + ], + "run_id": "status-run", + "payload_codec": "avro" + }, + "history": [ + { + "event_type": "ActivityScheduled", + "payload": { + "sequence": 1, + "activity_type": "reserve" + } + }, + { + "event_type": "ActivityCompleted", + "payload": { + "sequence": 1, + "activity_type": "reserve", + "result": { + "codec": "avro", + "blob": "wwHioz3/VYAiNw4ECHNsb3QEDgpzdGF0ZQoIaGVsZAA=" + } + } + } + ], + "expected": { + "command_type": "CompleteWorkflow", + "result": { + "order": "042", + "reserved": { + "slot": 7, + "state": "held" + }, + "run": "status-run" + } + } +} diff --git a/tests/integration/test_worker_status.py b/tests/integration/test_worker_status.py new file mode 100644 index 00000000..0b45668a --- /dev/null +++ b/tests/integration/test_worker_status.py @@ -0,0 +1,110 @@ +"""Worker readiness and liveness against a running Server (#168).""" + +from __future__ import annotations + +import asyncio +import threading +import time +import uuid +from typing import Any + +import pytest + +from durable_workflow import Client, Worker, WorkerState, activity, workflow + +release_activity = threading.Event() +activity_started = threading.Event() + + +@activity.defn(name="tests.worker-status.slow") +def slow(value: str) -> str: + activity_started.set() + release_activity.wait(30) + return f"slow:{value}" + + +@workflow.defn(name="tests.worker-status.workflow") +class StatusWorkflow: + def __init__(self) -> None: + self.stage = "started" + + @workflow.query("stage") + def current_stage(self) -> str: + return self.stage + + def run(self, ctx: Any, value: str) -> Any: + self.stage = "running_activity" + result = yield ctx.schedule_activity("tests.worker-status.slow", [value]) + self.stage = "done" + return result + + +async def eventually(predicate: Any, timeout: float = 30.0) -> None: + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() > deadline: + raise AssertionError("condition was not reached") + await asyncio.sleep(0.05) + + +@pytest.mark.asyncio +async def test_worker_status_follows_a_real_worker_lifecycle(server_url: str, server_token: str) -> None: + release_activity.clear() + activity_started.clear() + suffix = uuid.uuid4().hex[:8] + queue = f"worker-status-{suffix}" + async with Client(server_url, token=server_token, namespace="default") as client: + worker = Worker( + client, + task_queue=queue, + workflows=[StatusWorkflow], + activities=[slow], + poll_timeout=2.0, + shutdown_timeout=20.0, + ) + assert worker.status().ready_reason == "not_started" + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: worker.status().ready) + status = worker.status() + assert status.state is WorkerState.REGISTERED + assert status.namespace == "default" + assert status.task_queue == queue + assert {"heartbeat", "workflow_poll", "activity_poll"} <= {role.role for role in status.roles} + + # Idle long polls keep the worker ready across several poll windows. + await asyncio.sleep(5) + assert worker.status().ready, worker.status().as_dict() + + handle = await client.start_workflow( + workflow_type="tests.worker-status.workflow", + workflow_id=f"worker-status-{suffix}", + task_queue=queue, + input=["x"], + ) + await eventually(activity_started.is_set) + # A blocked synchronous activity does not affect readiness or liveness. + await asyncio.sleep(1) + status = worker.status() + assert status.ready, status.as_dict() + assert status.live, status.as_dict() + if "query_poll" in {role.role for role in status.roles}: + response = await client.query_workflow(handle.workflow_id, "stage") + assert response["result"] == "running_activity" + assert worker.status().ready + + stop_task = asyncio.create_task(worker.stop()) + await eventually(lambda: worker.status().state is WorkerState.DRAINING) + draining = worker.status() + assert (draining.ready, draining.ready_reason, draining.live) == (False, "draining", True) + release_activity.set() + await stop_task + await run_task + stopped = worker.status() + assert stopped.state is WorkerState.STOPPED + assert (stopped.ready, stopped.live) == (False, False) + finally: + release_activity.set() + if not run_task.done(): + await worker.stop() + await run_task diff --git a/tests/test_replay_regression_corpus.py b/tests/test_replay_regression_corpus.py index 37857750..3592c5f2 100644 --- a/tests/test_replay_regression_corpus.py +++ b/tests/test_replay_regression_corpus.py @@ -50,6 +50,7 @@ ) from tests.test_wait_condition import SignalCounterUntilFinished from tests.test_worker import compatible_cluster_info +from tests.test_worker_status import RecoveryWorkflow as StatusRecoveryWorkflow from tests.test_workflow_identity import IdentityAcrossReplay, IdentityContinueAsNew FIXTURE_SCHEMA = "durable-workflow.replay-regression/v1" @@ -394,6 +395,7 @@ def run(self, ctx: WorkflowContext, message: str) -> None: WorkflowStreamAuthorWorkflow, WorkflowMemoAuthorWorkflow, YieldedContinueAfterMetadataWorkflow, + StatusRecoveryWorkflow, ] WORKFLOW_TYPES = {str(getattr(workflow, "__workflow_name__", workflow.__name__)): workflow for workflow in WORKFLOWS} WORKFLOW_TYPES["tests.replay.condition-identity"] = ConditionIdentityWorkflow diff --git a/tests/test_worker_status.py b/tests/test_worker_status.py new file mode 100644 index 00000000..1468a767 --- /dev/null +++ b/tests/test_worker_status.py @@ -0,0 +1,574 @@ +"""Worker readiness and liveness status (#168).""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import threading +import time +from collections.abc import Callable +from pathlib import Path +from unittest.mock import AsyncMock + +import httpx +import pytest + +import durable_workflow.worker as worker_module +from durable_workflow import Worker, WorkerState, WorkerStatus, activity, serializer, workflow +from durable_workflow.client import Client +from durable_workflow.errors import ServerError +from tests.test_worker import TestWorkflow, compatible_cluster_info, echo_activity + +pytestmark = pytest.mark.asyncio + + +async def eventually(predicate: Callable[[], bool], *, timeout: float = 10.0) -> None: + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() > deadline: + raise AssertionError("condition was not reached") + await asyncio.sleep(0.005) + + +def roles(status: WorkerStatus) -> dict[str, object]: + return {role.role: role for role in status.roles} + + +async def empty_poll(**_: object) -> None: + await asyncio.sleep(0.002) + return None + + +@pytest.fixture +def client() -> AsyncMock: + client = AsyncMock(spec=Client) + client.get_cluster_info = AsyncMock(return_value=compatible_cluster_info()) + client.register_worker = AsyncMock(return_value={"worker_id": "w1", "registered": True}) + client.deregister_worker_registration = AsyncMock( + return_value={"worker_id": "w1", "outcome": "deregistered", "recovered_workflow_task_count": 0} + ) + client.heartbeat_worker = AsyncMock(return_value={"worker_id": "w1", "acknowledged": True}) + # Empty polls yield briefly so mocked pollers do not starve other threads. + client.poll_workflow_task = AsyncMock(side_effect=empty_poll) + client.poll_activity_task = AsyncMock(side_effect=empty_poll) + client.poll_query_task = AsyncMock(side_effect=empty_poll) + client.complete_activity_task = AsyncMock(return_value={"outcome": "completed"}) + client.fail_activity_task = AsyncMock(return_value={"outcome": "failed"}) + return client + + +@pytest.fixture(autouse=True) +def fast_status(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(worker_module, "EVENT_LOOP_TICK_SECONDS", 0.01) + monkeypatch.setattr(worker_module._PollFailureBackoff, "retry_delay", lambda *_args, **_kwargs: 0.01) + + +def make_worker(client: AsyncMock, **kwargs: object) -> Worker: + options: dict[str, object] = { + "task_queue": "q1", + "workflows": [TestWorkflow], + "activities": [echo_activity], + "poll_timeout": 0.01, + "shutdown_timeout": 1.0, + } + options.update(kwargs) + return Worker(client, **options) # type: ignore[arg-type] + + +async def stop(worker: Worker, run_task: asyncio.Task[object]) -> None: + await worker.stop() + with contextlib.suppress(asyncio.CancelledError): + await run_task + + +async def test_unstarted_worker_is_neither_ready_nor_live(client: AsyncMock) -> None: + worker = make_worker(client, worker_id="w-status") + + status = worker.status() + + assert status.state is WorkerState.STARTING + assert (status.ready, status.ready_reason) == (False, "not_started") + assert (status.live, status.live_reason) == (False, "not_started") + assert status.worker_id == "w-status" + assert status.task_queue == "q1" + assert status.roles == () + assert status.event_loop_progress_age_seconds is None + assert status.registered_age_seconds is None + + +async def test_registering_worker_is_live_but_not_ready(client: AsyncMock) -> None: + release = asyncio.Event() + + async def slow_registration(**_: object) -> dict[str, object]: + await release.wait() + return {"worker_id": "w1", "registered": True} + + client.register_worker = AsyncMock(side_effect=slow_registration) + worker = make_worker(client) + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: client.register_worker.await_count == 1) + status = worker.status() + assert status.state is WorkerState.STARTING + assert (status.ready, status.ready_reason) == (False, "registering") + assert (status.live, status.live_reason) == (True, "live") + + release.set() + await eventually(lambda: worker.status().ready) + status = worker.status() + assert status.state is WorkerState.REGISTERED + assert status.registered_age_seconds is not None + finally: + release.set() + await stop(worker, run_task) + + +async def test_registered_worker_reports_every_configured_role_ready(client: AsyncMock) -> None: + worker = make_worker(client) + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: client.poll_activity_task.await_count >= 3) + status = worker.status() + + assert status.state is WorkerState.REGISTERED + assert (status.ready, status.ready_reason) == (True, "ready") + assert (status.live, status.live_reason) == (True, "live") + assert set(roles(status)) == {"heartbeat", "workflow_poll", "activity_poll", "query_poll"} + for role in status.roles: + assert role.running is True + assert role.consecutive_failures == 0 + assert role.stalled is False + assert roles(status)["workflow_poll"].last_success_age_seconds is not None + assert roles(status)["heartbeat"].last_success_age_seconds is None + finally: + await stop(worker, run_task) + + +async def test_only_configured_roles_are_reported(client: AsyncMock) -> None: + worker = make_worker(client, workflows=[]) + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: worker.status().ready) + assert set(roles(worker.status())) == {"heartbeat", "activity_poll"} + finally: + await stop(worker, run_task) + + +async def test_empty_and_in_flight_long_polls_stay_ready(client: AsyncMock) -> None: + hold = asyncio.Event() + + async def long_poll(**_: object) -> None: + await hold.wait() + return None + + client.poll_workflow_task = AsyncMock(side_effect=long_poll) + worker = make_worker(client) + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: client.poll_workflow_task.await_count == 1) + await asyncio.sleep(0.05) + status = worker.status() + workflow_poll = roles(status)["workflow_poll"] + + assert status.ready + assert status.live + assert workflow_poll.in_flight_seconds is not None + assert workflow_poll.in_flight_seconds >= 0.04 + assert workflow_poll.overdue is False + assert workflow_poll.stalled is False + + overdue = worker.status(request_overdue_seconds=0.01) + assert (overdue.ready, overdue.ready_reason) == (False, "server_unavailable:workflow_poll") + assert (overdue.live, overdue.live_reason) == (True, "live") + assert roles(overdue)["workflow_poll"].overdue is True + + stalled = worker.status(request_stall_seconds=0.01) + assert (stalled.live, stalled.live_reason) == (False, "role_stalled:workflow_poll") + assert (stalled.ready, stalled.ready_reason) == (False, "role_stalled:workflow_poll") + assert roles(stalled)["workflow_poll"].stalled is True + finally: + hold.set() + await stop(worker, run_task) + + +async def test_transient_server_loss_is_unready_but_live_and_recovers_without_restart( + client: AsyncMock, +) -> None: + outage = threading.Event() + outage.set() + + async def poll(**_: object) -> None: + await asyncio.sleep(0.001) + if outage.is_set(): + raise httpx.ConnectError("connection refused") + return None + + client.poll_activity_task = AsyncMock(side_effect=poll) + worker = make_worker(client, workflows=[]) + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: client.poll_activity_task.await_count >= 2) + status = worker.status() + activity_poll = roles(status)["activity_poll"] + assert status.state is WorkerState.REGISTERED + assert (status.ready, status.ready_reason) == (False, "server_unavailable:activity_poll") + assert (status.live, status.live_reason) == (True, "live") + assert activity_poll.consecutive_failures >= 1 + assert activity_poll.last_failure_type == "ConnectError" + assert "connection refused" not in json.dumps(status.as_dict()) + + outage.clear() + await eventually(lambda: worker.status().ready) + recovered = roles(worker.status())["activity_poll"] + assert recovered.consecutive_failures == 0 + assert recovered.last_failure_age_seconds is not None + assert client.register_worker.await_count == 1 + finally: + await stop(worker, run_task) + + +async def test_poll_capacity_backpressure_keeps_the_worker_ready(client: AsyncMock) -> None: + refusal = ServerError( + 429, + { + "reason": "long_poll_capacity_exhausted", + "poll_status": "long_poll_capacity_exhausted", + "task": None, + "task_kind": "activity_task", + "task_queue": "q1", + "retryable": True, + "retry_after_seconds": 1, + }, + ) + client.poll_activity_task = AsyncMock(side_effect=refusal) + worker = make_worker(client, workflows=[]) + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: client.poll_activity_task.await_count >= 1) + status = worker.status() + assert status.ready, status.ready_reason + assert roles(status)["activity_poll"].consecutive_failures == 0 + finally: + await stop(worker, run_task) + + +async def test_heartbeat_failure_is_unready_until_the_next_heartbeat_succeeds(client: AsyncMock) -> None: + failing = threading.Event() + failing.set() + + async def heartbeat(**_: object) -> dict[str, object]: + if failing.is_set(): + raise httpx.ReadTimeout("timed out") + return {"worker_id": "w1", "acknowledged": True} + + client.heartbeat_worker = AsyncMock(side_effect=heartbeat) + worker = make_worker(client, heartbeat_interval=0.02) + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: client.heartbeat_worker.await_count >= 1) + await eventually(lambda: worker.status().ready_reason == "server_unavailable:heartbeat") + assert worker.status().live + + failing.clear() + await eventually(lambda: worker.status().ready) + assert roles(worker.status())["heartbeat"].last_success_age_seconds is not None + finally: + await stop(worker, run_task) + + +async def test_registration_refusal_reports_failed(client: AsyncMock) -> None: + client.register_worker = AsyncMock(side_effect=ServerError(403, {"reason": "forbidden"})) + worker = make_worker(client) + + with pytest.raises(ServerError): + await worker.run() + + status = worker.status() + assert status.state is WorkerState.FAILED + assert (status.ready, status.ready_reason) == (False, "failed:registration") + assert (status.live, status.live_reason) == (False, "failed:registration") + assert worker._status_tick_task is not None + await asyncio.sleep(0) + assert worker._status_tick_task.done() + + +async def test_poller_failure_reports_failed(client: AsyncMock, monkeypatch: pytest.MonkeyPatch) -> None: + async def broken_poller() -> None: + raise RuntimeError("poller bug") + + worker = make_worker(client) + monkeypatch.setattr(worker, "_poll_activity_tasks", broken_poller) + + with pytest.raises(RuntimeError, match="poller bug"): + await worker.run() + + status = worker.status() + assert status.state is WorkerState.FAILED + assert (status.ready, status.live, status.live_reason) == (False, False, "failed:runtime") + + +async def test_dead_query_thread_is_neither_live_nor_ready(client: AsyncMock) -> None: + class BrokenQueryClient: + async def __aenter__(self) -> BrokenQueryClient: + raise RuntimeError("query client unavailable") + + async def __aexit__(self, *_: object) -> None: + return None + + worker = make_worker(client) + worker._can_clone_client_for_query_tasks = lambda: True # type: ignore[method-assign] + worker._clone_client_for_query_tasks = lambda: BrokenQueryClient() # type: ignore[method-assign,return-value] + run_task = asyncio.create_task(worker.run()) + try: + def query_role_exited() -> bool: + query_poll = roles(worker.status()).get("query_poll") + return query_poll is not None and not query_poll.running # type: ignore[attr-defined] + + await eventually(query_role_exited) + status = worker.status() + assert status.state is WorkerState.REGISTERED + assert (status.live, status.live_reason) == (False, "role_stopped:query_poll") + assert (status.ready, status.ready_reason) == (False, "role_stopped:query_poll") + finally: + await stop(worker, run_task) + + +async def test_stalled_event_loop_is_reported_from_another_thread(client: AsyncMock) -> None: + worker = make_worker(client) + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: worker.status().ready) + observed: list[WorkerStatus] = [] + + def probe() -> None: + time.sleep(0.3) + observed.append(worker.status(event_loop_stall_seconds=0.15)) + + prober = threading.Thread(target=probe) + prober.start() + time.sleep(0.5) # Block the worker's event loop. + prober.join() + + assert (observed[0].live, observed[0].live_reason) == (False, "event_loop_stalled") + assert observed[0].ready is False + await eventually(lambda: worker.status(event_loop_stall_seconds=0.15).live) + finally: + await stop(worker, run_task) + + +async def test_long_running_activity_and_drain_states(client: AsyncMock) -> None: + started = threading.Event() + release = threading.Event() + + @activity.defn(name="status-slow-activity") + def slow_activity() -> str: + started.set() + release.wait(5) + return "done" + + delivered = False + + async def poll_activity(**_: object) -> dict[str, object] | None: + nonlocal delivered + await asyncio.sleep(0.001) + if delivered: + return None + delivered = True + return { + "task_id": "at-1", + "activity_attempt_id": "attempt-1", + "activity_type": "status-slow-activity", + "attempt_number": 1, + "lease_owner": "w1", + "payload_codec": "avro", + "arguments": None, + } + + client.poll_activity_task = AsyncMock(side_effect=poll_activity) + worker = make_worker(client, workflows=[], activities=[slow_activity], shutdown_timeout=5.0) + run_task = asyncio.create_task(worker.run()) + stop_task: asyncio.Task[None] | None = None + try: + await eventually(started.is_set) + await asyncio.sleep(0.05) + status = worker.status(request_stall_seconds=0.02) + assert status.ready, status.ready_reason + assert status.live, status.live_reason + + stop_task = asyncio.create_task(worker.stop()) + await eventually(lambda: worker.status().state is WorkerState.DRAINING) + draining = worker.status() + assert (draining.ready, draining.ready_reason) == (False, "draining") + assert (draining.live, draining.live_reason) == (True, "live") + + release.set() + await stop_task + await run_task + stopped = worker.status() + assert stopped.state is WorkerState.STOPPED + assert (stopped.ready, stopped.ready_reason) == (False, "stopped") + assert (stopped.live, stopped.live_reason) == (False, "stopped") + assert worker._status_tick_task is not None + assert worker._status_tick_task.done() + finally: + release.set() + if stop_task is not None: + await stop_task + if not run_task.done(): + await stop(worker, run_task) + + +async def test_shutdown_that_cannot_confirm_cleanup_reports_failed( + client: AsyncMock, monkeypatch: pytest.MonkeyPatch, +) -> None: + worker = make_worker(client) + run_task = asyncio.create_task(worker.run()) + await eventually(lambda: worker.status().ready) + + async def unconfirmed_shutdown() -> None: + raise RuntimeError("the worker registration remains active") + + monkeypatch.setattr(worker, "_shutdown", unconfirmed_shutdown) + with pytest.raises(RuntimeError): + await worker.stop() + with contextlib.suppress(RuntimeError, asyncio.CancelledError): + await run_task + + status = worker.status() + assert status.state is WorkerState.FAILED + assert (status.ready, status.live, status.live_reason) == (False, False, "failed:shutdown") + + +async def test_status_snapshot_performs_no_server_io(client: AsyncMock) -> None: + worker = make_worker(client, workflows=[], activities=[]) + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: worker.status().ready) + calls = [(name, getattr(client, name).await_count) for name in ( + "register_worker", "heartbeat_worker", "poll_workflow_task", "poll_activity_task", "poll_query_task", + "get_cluster_info", + )] + for _ in range(50): + worker.status() + await asyncio.to_thread(worker.status) + assert calls == [(name, getattr(client, name).await_count) for name, _ in calls] + finally: + await stop(worker, run_task) + + +async def test_status_is_json_serializable(client: AsyncMock) -> None: + worker = make_worker(client) + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: worker.status().ready) + payload = json.loads(json.dumps(worker.status().as_dict())) + assert payload["state"] == "registered" + assert payload["ready"] is True + assert payload["live"] is True + assert {role["role"] for role in payload["roles"]} == { + "heartbeat", "workflow_poll", "activity_poll", "query_poll", + } + finally: + await stop(worker, run_task) + + +@pytest.mark.parametrize("value", [0, -1, float("inf"), float("nan"), True, "5"]) +async def test_status_rejects_invalid_bounds(client: AsyncMock, value: object) -> None: + worker = make_worker(client) + with pytest.raises(ValueError, match="event_loop_stall_seconds"): + worker.status(event_loop_stall_seconds=value) # type: ignore[arg-type] + with pytest.raises(ValueError, match="request_overdue_seconds"): + worker.status(request_overdue_seconds=value) # type: ignore[arg-type] + with pytest.raises(ValueError, match="request_stall_seconds"): + worker.status(request_stall_seconds=value) # type: ignore[arg-type] + + +async def test_default_request_bounds_follow_the_attempt_deadline_and_retry_budget(client: AsyncMock) -> None: + from durable_workflow.retry_policy import TransportRetryPolicy + + client.timeout = 60.0 + client.retry_policy = TransportRetryPolicy() + worker = make_worker(client, poll_timeout=35.0) + assert worker._request_attempt_deadline_seconds() == 60.0 + assert worker._default_request_stall_seconds() == 2 * (3 * 60.0 + 2 * 5.0) + client.timeout = 10.0 + assert worker._request_attempt_deadline_seconds() == 40.0 + client.retry_policy = TransportRetryPolicy(max_attempts=1) + assert worker._default_request_stall_seconds() == 2 * 40.0 + + +async def test_run_until_reports_registration_and_stop(client: AsyncMock) -> None: + from durable_workflow.client import WorkflowExecution + + client.describe_workflow = AsyncMock(return_value=WorkflowExecution( + workflow_id="wf-1", run_id="run-1", workflow_type="test-wf", status="completed", + )) + worker = make_worker(client) + await worker.run_until(workflow_id="wf-1", timeout=2.0) + + status = worker.status() + assert status.state is WorkerState.STOPPED + assert set(roles(status)) == {"heartbeat", "query_poll"} + assert (status.ready, status.live) == (False, False) + + +RECOVERY_FIXTURE = Path(__file__).parent / "fixtures/replay_regressions/worker-status-outage-recovery-cold-history.json" + + +@workflow.defn(name="tests.worker-status.recovery") +class RecoveryWorkflow: + def run(self, ctx, order: str): # type: ignore[no-untyped-def] + reserved = yield ctx.schedule_activity("reserve", [order]) + return {"order": order, "reserved": reserved, "run": ctx.run_id} + + +async def test_outage_recovery_replays_cold_history_and_restores_readiness(client: AsyncMock) -> None: + fixture = json.loads(RECOVERY_FIXTURE.read_text()) + outage = threading.Event() + outage.set() + delivered = False + completed = asyncio.Event() + task = { + "task_id": "status-recovered-task", "workflow_task_attempt": 2, + "workflow_type": fixture["workflow"]["type"], "workflow_id": "order-042", + "run_id": fixture["workflow"]["run_id"], "payload_codec": "avro", + "arguments": serializer.encode(fixture["workflow"]["input"]), "history_events": fixture["history"], + } + + async def poll(**_: object) -> dict[str, object] | None: + nonlocal delivered + await asyncio.sleep(0.002) + if outage.is_set(): + raise httpx.ConnectError("runtime unavailable") + if not delivered: + delivered = True + return task + return None + + async def complete(**_: object) -> dict[str, str]: + completed.set() + return {"outcome": "completed"} + + client.poll_workflow_task = AsyncMock(side_effect=poll) + client.complete_workflow_task = AsyncMock(side_effect=complete) + worker = make_worker(client, workflows=[RecoveryWorkflow], activities=[]) + run_task = asyncio.create_task(worker.run()) + try: + await eventually(lambda: worker.status().ready_reason == "server_unavailable:workflow_poll") + assert worker.status().live + + outage.clear() + await eventually(completed.is_set) + await eventually(lambda: worker.status().ready) + finally: + await stop(worker, run_task) + + client.complete_workflow_task.assert_awaited_once() + publication = client.complete_workflow_task.call_args.kwargs + assert publication["task_id"] == "status-recovered-task" + [command] = publication["commands"] + assert command["type"] == "complete_workflow" + assert serializer.decode_envelope(command["result"], codec="avro") == fixture["expected"]["result"] + client.fail_workflow_task.assert_not_awaited() + assert client.register_worker.await_count == 1