diff --git a/CHANGELOG.md b/CHANGELOG.md index cce5c66..32623b7 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] +## [2.10.1] - Unreleased + +### Fixed + +- Keep worker deregistration inside the remaining shutdown deadline. When the + runtime advertises generation-fenced receipts, retry matching transient or + lost responses with the original registration token. Without advertised + fencing, use one legacy attempt. Exhausted or uncertain shutdown keeps + registration removal unconfirmed instead of reporting success. + ## [2.10.0] - 2026-10-23 ### Added diff --git a/docs/reference/worker.md b/docs/reference/worker.md index 32b7a5f..9e78421 100644 --- a/docs/reference/worker.md +++ b/docs/reference/worker.md @@ -65,8 +65,27 @@ worker-stop behavior. Neither uses this generic poll-failure backoff. before polling. After registration succeeds, shutdown stops every poller and waits for accepted work within the shared `shutdown_timeout` deadline before removing that worker-plane registration. Work still running at the deadline is -cancelled. Calling `stop()` more than once is safe and sends at most one -deregistration request. No deregistration is sent if registration failed. +cancelled. Calling `stop()` more than once shares the same shutdown operation +and result. No deregistration is sent if registration failed. + +**Unreleased source behavior, prepared for SDK 2.10.1:** the final HTTP call and any recovery use the +original remaining shutdown budget. If Server advertises +`worker_deregistration_fencing`, Worker retains its registration token and uses +the fenced receipt endpoint. Matching transient responses and transport losses +can retry that same generation. A completed receipt can replay safely even if +a replacement now owns the worker ID. Authentication, malformed or mismatched +receipts, lost authority and expired receipts fail without retrying. Without +advertised fencing, Worker keeps one attempt through the legacy worker-plane +DELETE and does not replay an uncertain result. That legacy route has no +generation fence. Keep worker IDs unique to a process, and do not reuse an ID +until the prior worker has fully stopped. + +If drain consumes the deadline, no late deregistration request is sent. A final +HTTP timeout or exhausted recovery raises `RuntimeError` and leaves removal +unconfirmed. The Server might have committed a request whose response was lost; +an exception is not proof that the registration still exists. Inspect the +runtime before taking an operator recovery action. These changes are not in a +published package yet. Cancelling an await cannot stop a running synchronous activity thread. Its physical slot remains occupied until the callback returns, and a late return diff --git a/pyproject.toml b/pyproject.toml index cd610fa..e118a7f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "durable-workflow" -version = "2.10.0" +version = "2.10.1" description = "Python client and worker SDK for Durable Workflow Cloud and self-hosted Server" readme = "README.md" requires-python = ">=3.10" @@ -76,8 +76,8 @@ durable-workflow-replay-conformance = "durable_workflow.replay_conformance:main" durable-workflow-workflow-updates-conformance = "durable_workflow.workflow_updates_conformance:main" [tool.durable-workflow] -product-train = "2.10.0" -registry-version = "2.10.0" +product-train = "2.10.1" +registry-version = "2.10.1" supported-server-versions = "2.8.0" worker-protocol-version = "1.20" control-plane-version = "2" diff --git a/src/durable_workflow/_metric_routes.py b/src/durable_workflow/_metric_routes.py index 9eb35f7..ad2e1d5 100644 --- a/src/durable_workflow/_metric_routes.py +++ b/src/durable_workflow/_metric_routes.py @@ -12,6 +12,7 @@ for suffix in ("", "/pass")), "/workers", "/workers/{worker_id}", "/worker/register", "/worker/heartbeat", "/worker/registrations/{worker_id}", + "/worker/registrations/{worker_id}/deregister", "/worker/sessions", "/worker/sessions/{session_id}", "/worker/sessions/{session_id}/heartbeat", *(f"/worker/{kind}-tasks/poll" for kind in ("workflow", "activity", "query", "update-validation")), *(f"/worker/{kind}-tasks/{{task_id}}/{action}" for kind, actions in ( diff --git a/src/durable_workflow/client.py b/src/durable_workflow/client.py index 3c9040e..36bb07b 100644 --- a/src/durable_workflow/client.py +++ b/src/durable_workflow/client.py @@ -1853,6 +1853,7 @@ async def _request( json: Any = None, timeout: float | None = None, context: str = "", + retry: bool = True, ) -> Any: if worker and discovery: raise ValueError("a request cannot be both worker-plane and cluster discovery") @@ -1920,7 +1921,7 @@ async def _do_request() -> httpx.Response: try: try: - resp = await self.retry_policy.execute(_do_request) + resp = await self.retry_policy.execute(_do_request) if retry else await _do_request() except httpx.HTTPStatusError as exc: status_code = str(exc.response.status_code) outcome = "http_error" @@ -5245,39 +5246,69 @@ async def close_worker_session( json={"worker_id": worker_id, "reason": reason}, )) - async def deregister_worker_registration(self, worker_id: str) -> dict[str, Any]: + async def deregister_worker_registration( + self, worker_id: str, *, registration_token: str | None = None, timeout: float | None = None, + expected_namespace: str | None = None, + ) -> dict[str, Any]: """Remove this runtime's successful worker-plane registration. Called by :class:`~durable_workflow.Worker` after its pollers and in-flight tasks have drained. This is distinct from the operator-only :meth:`deregister_worker` control-plane management operation. Returns the worker deregistration envelope, including recovered workflow-task - count. + count. With a registration token, use the runtime's generation-fenced + receipt endpoint. Each call makes one attempt: Worker owns safe retries + and the remaining shutdown budget. An unfenced DELETE cannot safely be + repeated after an uncertain response because the worker ID may be reused. + ``expected_namespace`` is a local registration-scope precondition, not + a namespace override. """ + if expected_namespace is not None and self.namespace != expected_namespace: + raise RuntimeError("worker registration namespace changed; registration removal is unconfirmed") + if registration_token is not None and ( + not isinstance(registration_token, str) or re.fullmatch(r"[a-f0-9]{32}", registration_token) is None + ): + raise ValueError("registration_token must contain exactly 32 lowercase hexadecimal characters") + fenced = registration_token is not None try: data = await self._request( - "DELETE", - f"/worker/registrations/{quote(worker_id, safe='')}", + "POST" if fenced else "DELETE", + f"/worker/registrations/{quote(worker_id, safe='')}" + ("/deregister" if fenced else ""), worker=True, context=worker_id, + json={"registration_token": registration_token} if fenced else None, + timeout=timeout, + retry=False, ) except ServerError as error: - if error.status == 404 and error.reason() == "worker_not_found": + if not fenced and error.status == 404 and error.reason() == "worker_not_found": return { "worker_id": worker_id, "outcome": "already_deregistered", "recovered_workflow_task_count": 0, } raise - if not isinstance(data, dict): + if not self._valid_worker_deregistration_response(data, worker_id, registration_token): raise ServerError( 200, { "reason": "invalid_worker_deregistration_response", - "message": f"expected JSON object, got {type(data).__name__}", + "message": "expected a deregistration receipt for this worker registration", }, ) - return data + return cast(dict[str, Any], data) + + @staticmethod + def _valid_worker_deregistration_response(data: Any, worker_id: str, token: str | None) -> bool: + if not isinstance(data, dict): + return False + count = data.get("recovered_workflow_task_count") + return ( + data.get("worker_id") == worker_id + and data.get("outcome") == "deregistered" + and type(count) is int and count >= 0 + and (token is None or data.get("registration_token") == token) + ) async def heartbeat_worker( self, diff --git a/src/durable_workflow/errors.py b/src/durable_workflow/errors.py index e18f958..9550403 100644 --- a/src/durable_workflow/errors.py +++ b/src/durable_workflow/errors.py @@ -219,6 +219,21 @@ def reason(self) -> str | None: return self.body.get("reason") return None + def worker_deregistration_retry_delay(self, worker_id: str, registration_token: str) -> int | None: + """Return a delay only for an identity-preserving fenced transient.""" + body = self.body + if ( + self.status != 503 or not isinstance(body, dict) + or body.get("reason") not in ("backend_lock_pressure", "backend_unavailable") + or body.get("operation") != "deregister_worker" or body.get("outcome") != "unknown" + or body.get("worker_id") != worker_id or body.get("registration_token") != registration_token + or body.get("retryable") is not True + or type(body.get("retry_after_seconds")) is not int + or not 0 < body["retry_after_seconds"] <= 60 + ): + return None + return int(body["retry_after_seconds"]) + def poll_capacity_backpressure_delay(self, task_kind: str, task_queue: str) -> int | None: """Return the advertised wait only for a complete poll-capacity refusal.""" body = self.body diff --git a/src/durable_workflow/worker.py b/src/durable_workflow/worker.py index 4fd9d15..83f440c 100644 --- a/src/durable_workflow/worker.py +++ b/src/durable_workflow/worker.py @@ -37,6 +37,8 @@ from types import FunctionType from typing import Annotated, Any, Concatenate, Literal, ParamSpec, TypeVar, Union, get_args, get_origin, get_type_hints +import httpx + from . import serializer from ._activity_lease import activity_lease_renewal_supported, current_activity_authority from ._activity_process import CallbackFailure, CallbackInvocation, CallbackProcessLost, SupervisedCallback @@ -1169,6 +1171,10 @@ def __init__( self._in_flight: set[asyncio.Task[Any]] = set() self._poller_tasks: set[asyncio.Task[Any]] = set() self._registered = False + self._deregistration_fenced = False + self._registration_token: str | None = None + self._registration_worker_id: str | None = None + self._registration_namespace: object = None self._run_started = False self._registration_done = asyncio.Event() self._shutdown_task: asyncio.Task[None] | None = None @@ -1295,6 +1301,18 @@ async def _register(self) -> None: _validate_server_compatibility(info) protocol = info.get("worker_protocol") server_capabilities = protocol.get("server_capabilities") if isinstance(protocol, Mapping) else None + deregistration = ( + server_capabilities.get("worker_deregistration_fencing") + if isinstance(server_capabilities, Mapping) else None + ) + self._deregistration_fenced = isinstance(deregistration, Mapping) and deregistration.get("supported") is True + if isinstance(deregistration, Mapping) and self._deregistration_fenced and ( + deregistration.get("schema") != "durable-workflow.v2.worker-deregistration.v1" + or deregistration.get("endpoint") != "/worker/registrations/{workerId}/deregister" + or type(deregistration.get("receipt_retention_seconds")) is not int + or deregistration["receipt_retention_seconds"] <= 0 + ): + raise RuntimeError("Server compatibility error: invalid worker_deregistration_fencing contract") sticky_support = ( server_capabilities.get("sticky_execution") if isinstance(server_capabilities, Mapping) else None ) @@ -1449,6 +1467,20 @@ async def _register(self) -> None: task_slots=self._current_task_slots(), process_metrics=self._current_process_metrics(), ) + if self._deregistration_fenced: + token = ack.get("registration_token") if isinstance(ack, dict) else None + if ( + not isinstance(token, str) or len(token) != 32 + or any(character not in "0123456789abcdef" for character in token) + or ack.get("worker_id") != self.worker_id or ack.get("registered") is not True + ): + raise RuntimeError( + "Server returned an invalid registration_token acknowledgement; " + "registration may remain active and unfenced cleanup is unsafe" + ) + self._registration_token = token + self._registration_worker_id = self.worker_id + self._registration_namespace = getattr(self.client, "namespace", None) self._registered = True # Adapt to the server-advertised cadence when present so a cluster # can pin the worker fleet's heartbeat beat without each worker @@ -4567,8 +4599,9 @@ def _track_worker_session_from_task(self, task: dict[str, Any]) -> None: async def stop(self) -> None: """Stop polling, drain tasks, and remove the worker registration. - Repeated calls share one shutdown result. A failed deregistration is - therefore propagated to every caller without sending another DELETE. + Repeated calls share one shutdown result. Generation-fenced recovery + stays within the original shutdown deadline. Unfenced deregistration + makes one attempt and does not retry an uncertain outcome. """ self._stop.set() if self._shutdown_task is None: @@ -4664,8 +4697,9 @@ async def _shutdown(self) -> None: log.warning("worker session %s close failed: %s", session.options.session_id, error) if self._registered: - await self.client.deregister_worker_registration(self.worker_id) + await self._deregister_before_deadline(deadline) self._registered = False + self._registration_token = None log.info("worker %s deregistered", self.worker_id) self._sticky_cache.clear() payload_observers = getattr(self.client, "_workflow_payload_observers", None) @@ -4676,3 +4710,62 @@ async def _shutdown(self) -> None: @staticmethod def _remaining_shutdown_time(deadline: float) -> float: return max(0.0, deadline - asyncio.get_running_loop().time()) + + def _check_deregistration_identity(self) -> None: + if self._deregistration_fenced and self._registration_token is None: + raise RuntimeError("worker registration token is missing; registration removal is unconfirmed") + if self._registration_worker_id is not None and ( + self.worker_id != self._registration_worker_id + or getattr(self.client, "namespace", None) != self._registration_namespace + ): + raise RuntimeError("worker registration identity changed; registration removal is unconfirmed") + + async def _deregistration_attempt(self, remaining: float) -> tuple[float | None, Exception] | None: + kwargs: dict[str, Any] = {} + if self._deregistration_fenced: + kwargs = {"registration_token": self._registration_token, "timeout": remaining} + if isinstance(self._registration_namespace, str): + kwargs["expected_namespace"] = self._registration_namespace + try: + await asyncio.wait_for( + self.client.deregister_worker_registration(self.worker_id, **kwargs), + timeout=remaining, + ) + except asyncio.TimeoutError as error: + raise RuntimeError( + "worker shutdown timed out during deregistration; registration removal is unconfirmed" + ) from error + except ServerError as error: + retry_delay = ( + error.worker_deregistration_retry_delay(self.worker_id, self._registration_token) + if self._deregistration_fenced and self._registration_token is not None else None + ) + if retry_delay is None: + raise + return float(retry_delay), error + except (httpx.NetworkError, httpx.TimeoutException, httpx.RemoteProtocolError) as error: + if not self._deregistration_fenced: + raise + return None, error + return None + + async def _deregister_before_deadline(self, deadline: float) -> None: + attempt = 0 + while True: + self._check_deregistration_identity() + remaining = self._remaining_shutdown_time(deadline) + if remaining <= 0: + raise RuntimeError("worker shutdown deadline exhausted; registration removal is unconfirmed") + retry = await self._deregistration_attempt(remaining) + if retry is None: + return + retry_delay, failure = retry + delay = min(1.0, 0.1 * 2 ** min(attempt, 4)) if retry_delay is None else retry_delay + if delay >= self._remaining_shutdown_time(deadline): + raise RuntimeError( + "worker shutdown cannot retry within its deadline; registration removal is unconfirmed" + ) from failure + if attempt == 0: + log.warning("worker deregistration is unconfirmed; retrying the fenced request within its deadline") + attempt += 1 + await asyncio.sleep(delay) diff --git a/tests/deregistration_recovery.py b/tests/deregistration_recovery.py new file mode 100644 index 0000000..9d977aa --- /dev/null +++ b/tests/deregistration_recovery.py @@ -0,0 +1,51 @@ +"""Shared workflow and committed-response fault for deregistration controls.""" +from __future__ import annotations + +from typing import Any + +import httpx + +from durable_workflow import Worker, activity, workflow + + +@activity.defn(name="tests.deregistration.echo") +def echo(value: str) -> str: + return value + + +@workflow.defn(name="tests.deregistration.cold-recovery") +class RecoveredWorkflow: + def run(self, ctx: Any, value: str) -> Any: + returned = yield ctx.schedule_activity("tests.deregistration.echo", [value], retry_policy={"max_attempts": 1}) + return {"value": returned, "run": ctx.run_id} + + +class DiscardCompletedResponse(httpx.AsyncBaseTransport): + def __init__(self, replacement: Worker) -> None: + self.inner = httpx.AsyncHTTPTransport() + self.replacement = replacement + self.discarded = False + self.deregistration_bodies: list[bytes] = [] + self.recovered_workflow_task_count: int | None = None + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + response = await self.inner.handle_async_request(request) + if request.method != "POST" or not request.url.path.endswith("/deregister"): + return response + self.deregistration_bodies.append(request.content) + if self.discarded: + return response + # The real Server commits before only the caller's response is lost. + await response.aread() + assert response.status_code == 200 + receipt = response.json() + assert receipt["outcome"] == "deregistered" + self.recovered_workflow_task_count = receipt["recovered_workflow_task_count"] + await response.aclose() + await self.replacement._register() + assert self.replacement._registration_token != receipt["registration_token"] + self.discarded = True + raise httpx.RemoteProtocolError("test discarded a committed response", request=request) + + async def aclose(self) -> None: + await self.inner.aclose() diff --git a/tests/fixtures/replay_regressions/fenced-deregistration-cold-recovery.json b/tests/fixtures/replay_regressions/fenced-deregistration-cold-recovery.json new file mode 100644 index 0000000..956f271 --- /dev/null +++ b/tests/fixtures/replay_regressions/fenced-deregistration-cold-recovery.json @@ -0,0 +1,64 @@ +{ + "$schema": "https://raw.githubusercontent.com/durable-workflow/.github/main/regression-corpus/evidence-schema.json", + "fixture_schema": "durable-workflow.replay-regression/v1", + "id": "fenced-deregistration-cold-recovery", + "protocol_version": "1.20", + "bindings": [ + "python" + ], + "workflow": { + "type": "tests.deregistration.cold-recovery", + "input": [ + "handoff" + ], + "run_id": "01m5mpg2jp77pe59qqcjrvg51x", + "payload_codec": "avro" + }, + "history": [ + { + "event_type": "ActivityScheduled", + "payload": { + "activity_execution_id": "01m5mpg4cw11vhx6w2b9mdssmg", + "activity_type": "tests.deregistration.echo", + "sequence": 1 + } + }, + { + "event_type": "ActivityCompleted", + "payload": { + "activity_execution_id": "01m5mpg4cw11vhx6w2b9mdssmg", + "activity_attempt_id": "01M5MPG4MRZJB84ND7MZ1VRPRX", + "activity_type": "tests.deregistration.echo", + "sequence": 1, + "result": { + "codec": "avro", + "blob": "wwHioz3/VYAiNwoOaGFuZG9mZg==" + }, + "payload_codec": "avro" + } + } + ], + "expected": { + "command_type": "CompleteWorkflow", + "result": { + "value": "handoff", + "run": "01m5mpg2jp77pe59qqcjrvg51x" + } + }, + "provenance": { + "qualification": "fresh published Server with candidate Python source; not registry publication", + "server_version": "2.11.1", + "server_source": "5336a98dd256e046edfdb67c81764cde4babd360", + "core_version": "2.10.2", + "core_source": "aca7928484083635e68a5a384b2dcecef842108e", + "image_digest": "sha256:bf8bed13f5f809ee3425233180a581353eac4683adbee5c4ecbf63bb18bbd102", + "workflow_id": "sdk175-cold-af7d38fbd6fa", + "observed_at": "2026-10-23T15:36:31.757789+00:00", + "observation_sha256": "90a75f09517c1ad852d48df2ad43d31b68837d2cc3f4895eaabfae2b095765eb", + "recovered_claim_count": 1, + "replacement_preserved": true, + "old_attempt": 1, + "replacement_attempt": 2, + "fixture_form": "minimal replay fields derived from the retained fresh observation; not a raw export" + } +} diff --git a/tests/integration/test_worker_deregistration.py b/tests/integration/test_worker_deregistration.py new file mode 100644 index 0000000..cf364a0 --- /dev/null +++ b/tests/integration/test_worker_deregistration.py @@ -0,0 +1,189 @@ +"""Qualify generation-fenced receipts against the connected real Server.""" + +from __future__ import annotations + +import asyncio +import os +import uuid +from typing import Any + +import httpx +import pytest + +from durable_workflow import Client, ServerError, Worker +from durable_workflow.workflow import replay +from tests.deregistration_recovery import DiscardCompletedResponse, RecoveredWorkflow, echo + + +def make_worker(client: Client, worker_id: str) -> Worker: + return Worker(client, task_queue="python-deregistration", worker_id=worker_id) + + +@pytest.mark.asyncio +async def test_completed_receipt_replays_without_deleting_replacement( + server_url: str, server_token: str, +) -> None: + worker_id = "py-dereg-receipt-" + uuid.uuid4().hex + async with Client(server_url, token=server_token) as client: + first = make_worker(client, worker_id) + await first._register() + token = first._registration_token + # This integration cohort must expose the contract being qualified; + # SDK unit controls separately preserve capability-absent operation. + assert first._deregistration_fenced + assert token + result = await client.deregister_worker_registration(worker_id, registration_token=token) + assert result["outcome"] == "deregistered" + replacement = make_worker(client, worker_id) + await replacement._register() + try: + assert replacement._registration_token != token + await first.stop() + assert not first._registered + assert (await client.heartbeat_worker(worker_id=worker_id))["acknowledged"] is True + finally: + await replacement.stop() + + +@pytest.mark.asyncio +async def test_superseded_generation_is_refused_without_touching_replacement( + server_url: str, server_token: str, +) -> None: + worker_id = "py-dereg-replaced-" + uuid.uuid4().hex + async with Client(server_url, token=server_token) as client: + first = make_worker(client, worker_id) + await first._register() + assert first._deregistration_fenced + assert first._registration_token + replacement = make_worker(client, worker_id) + await replacement._register() + try: + with pytest.raises(ServerError) as failure: + await first.stop() + assert failure.value.status == 409 + assert failure.value.reason() == "worker_registration_lost_authority" + assert first._registered + assert (await client.heartbeat_worker(worker_id=worker_id))["acknowledged"] is True + finally: + await replacement.stop() + + +@pytest.mark.asyncio +async def test_unknown_registration_token_is_not_already_deregistered( + server_url: str, server_token: str, +) -> None: + worker_id = "py-dereg-unknown-" + uuid.uuid4().hex + async with Client(server_url, token=server_token) as client: + worker = make_worker(client, worker_id) + await worker._register() + assert worker._deregistration_fenced + assert worker._registration_token + try: + with pytest.raises(ServerError) as failure: + await client.deregister_worker_registration(worker_id, registration_token=uuid.uuid4().hex) + assert failure.value.status == 404 + assert failure.value.reason() == "worker_registration_token_not_found" + assert (await client.heartbeat_worker(worker_id=worker_id))["acknowledged"] is True + finally: + await worker.stop() + + +@pytest.mark.asyncio +async def test_lost_completed_response_retries_receipt_after_replacement_registration( + server_url: str, server_token: str, +) -> None: + worker_id = "py-dereg-lost-" + uuid.uuid4().hex + async with ( + Client(server_url, token=server_token) as client, + Client(server_url, token=server_token) as replacement_client, + ): + replacement = make_worker(replacement_client, worker_id) + + fault = DiscardCompletedResponse(replacement) + await client._http.aclose() + client._http = httpx.AsyncClient(base_url=client.base_url, transport=fault) + first = make_worker(client, worker_id) + await first._register() + assert first._deregistration_fenced + assert first._registration_token + try: + await first.stop() + assert fault.discarded + assert len(fault.deregistration_bodies) == 2 + assert fault.deregistration_bodies[0] == fault.deregistration_bodies[1] + assert not first._registered + assert (await replacement_client.heartbeat_worker(worker_id=worker_id))["acknowledged"] is True + finally: + await replacement.stop() + + +async def exercise_recovered_claim(client: Client, replacement_client: Client, identity: str) -> dict[str, Any]: + first = Worker(client, task_queue=identity, worker_id=identity, workflows=[RecoveredWorkflow], activities=[echo]) + replacement = Worker( + replacement_client, task_queue=identity, worker_id=identity, workflows=[RecoveredWorkflow], activities=[echo], + ) + fault = DiscardCompletedResponse(replacement) + await client._http.aclose() + client._http = httpx.AsyncClient(base_url=client.base_url, transport=fault) + await first._register() + handle = await client.start_workflow( + workflow_type="tests.deregistration.cold-recovery", task_queue=identity, + workflow_id=identity, input=["handoff"], + ) + old = await client.poll_workflow_task(worker_id=identity, task_queue=identity, timeout=2) + assert isinstance(old, dict) + try: + await first.stop() + assert not first._registered + assert fault.discarded + assert fault.recovered_workflow_task_count == 1 + assert len(fault.deregistration_bodies) == 2 + assert fault.deregistration_bodies[0] == fault.deregistration_bodies[1] + claimed = await replacement_client.poll_workflow_task(worker_id=identity, task_queue=identity, timeout=2) + assert isinstance(claimed, dict) + assert claimed["task_id"] == old["task_id"] + assert claimed["workflow_task_attempt"] > old["workflow_task_attempt"] + outcome = replay(RecoveredWorkflow, claimed["history_events"], ["handoff"], run_id=handle.run_id) + assert len(outcome.commands) == 1 + await replacement_client.complete_workflow_task( + task_id=claimed["task_id"], lease_owner=identity, workflow_task_attempt=claimed["workflow_task_attempt"], + commands=[command.to_server_command(identity) for command in outcome.commands], + ) + task = await replacement_client.poll_activity_task(worker_id=identity, task_queue=identity, timeout=2) + assert isinstance(task, dict) + await replacement_client.complete_activity_task( + task_id=task["task_id"], activity_attempt_id=task["activity_attempt_id"], + lease_owner=identity, result="handoff", + ) + resumed = await replacement_client.poll_workflow_task(worker_id=identity, task_queue=identity, timeout=2) + assert isinstance(resumed, dict) + cold = replay(RecoveredWorkflow, resumed["history_events"], ["handoff"], run_id=handle.run_id) + assert len(cold.commands) == 1 + assert cold.commands[0].to_server_command(identity)["type"] == "complete_workflow" + await replacement_client.complete_workflow_task( + task_id=resumed["task_id"], lease_owner=identity, workflow_task_attempt=resumed["workflow_task_attempt"], + commands=[command.to_server_command(identity) for command in cold.commands], + ) + result = await handle.result(timeout=3, poll_interval=.05) + assert result == {"value": "handoff", "run": handle.run_id} + return {"history": resumed["history_events"], "run_id": handle.run_id, "workflow_id": handle.workflow_id, + "result": result, "old_attempt": old["workflow_task_attempt"], + "replacement_attempt": claimed["workflow_task_attempt"], "recovered_claim_count": 1, + "deregistration_attempts": 2, "same_request_bytes": True, + "replacement_preserved": True, "cold_replay_did_not_reschedule_activity": True} + finally: + await replacement.stop() + + +@pytest.mark.asyncio +async def test_recovered_claim_cold_replays_after_lost_receipt_and_replacement( + server_url: str, server_token: str, +) -> None: + async with ( + Client(server_url, worker_token=server_token, + control_token=os.environ.get("DURABLE_WORKFLOW_CONTROL_TOKEN", server_token)) as client, + Client(server_url, worker_token=server_token) as replacement_client, + ): + await asyncio.wait_for( + exercise_recovered_claim(client, replacement_client, "py-dereg-cold-" + uuid.uuid4().hex[:12]), timeout=30, + ) diff --git a/tests/integration/test_worker_task_families.py b/tests/integration/test_worker_task_families.py index 1682d04..9e6a82c 100644 --- a/tests/integration/test_worker_task_families.py +++ b/tests/integration/test_worker_task_families.py @@ -5,6 +5,7 @@ import asyncio import importlib.metadata import json +import os import uuid from collections import Counter from contextvars import ContextVar @@ -69,9 +70,14 @@ async def test_workers_execute_without_polling_empty_task_families( suffix = uuid.uuid4().hex[:12] queue = f"task-family-{suffix}" async with ( - Client(server_url, token=server_token, namespace="default") as control, - Client(server_url, token=server_token, namespace="default") as workflow_client, - Client(server_url, token=server_token, namespace="default") as activity_client, + Client(server_url, control_token=os.environ.get("DURABLE_WORKFLOW_CONTROL_TOKEN", server_token), + namespace="default") as control, + Client(server_url, worker_token=server_token, + control_token=os.environ.get("DURABLE_WORKFLOW_CONTROL_TOKEN", server_token), + namespace="default") as workflow_client, + Client(server_url, worker_token=server_token, + control_token=os.environ.get("DURABLE_WORKFLOW_CONTROL_TOKEN", server_token), + namespace="default") as activity_client, ): observation.clients.update({ id(control._http): "control", id(workflow_client._http): "workflow", @@ -143,7 +149,12 @@ async def test_workers_execute_without_polling_empty_task_families( assert workflow_polls["workflow"] > 0, observation.report() assert (workflow_polls["activity"] > 0) is (mode == "mixed"), observation.report() assert observation.successful["POST registration"] >= len(workers), observation.report() - assert observation.successful["DELETE registration"] >= len(workers), observation.report() + deregistrations = (observation.successful["POST deregistration"] + + observation.successful["DELETE deregistration"]) + assert deregistrations >= len(workers), observation.report() + assert all(not selected._registered for selected in workers), observation.report() + assert not workflow_client._http.is_closed + assert not activity_client._http.is_closed assert observation.active_count == 0, observation.report() if mode == "split": assert activity_polls["activity"] > 0, observation.report() diff --git a/tests/test_cooperative_cancellation_worker.py b/tests/test_cooperative_cancellation_worker.py index 0a3fc30..d732c60 100644 --- a/tests/test_cooperative_cancellation_worker.py +++ b/tests/test_cooperative_cancellation_worker.py @@ -607,8 +607,12 @@ def synchronous_compensation(request_id: str) -> object: await asyncio.wait_for(entered.wait(), timeout=1) original_history = deepcopy(server.history) heartbeats_before_stop = server.client.heartbeat_workflow_task.await_count - await asyncio.wait_for(worker.stop(), timeout=1) - release_thread.set() + stopping = worker.stop() + try: + with pytest.raises(RuntimeError, match="deadline exhausted.*registration.*unconfirmed"): + await asyncio.wait_for(stopping, timeout=1) + finally: + release_thread.set() await asyncio.wait_for(discarded.wait(), timeout=1) assert fenced_heartbeats == [True] assert server.client.heartbeat_workflow_task.await_count == heartbeats_before_stop @@ -616,7 +620,8 @@ def synchronous_compensation(request_id: str) -> object: assert server.history == original_history server.client.complete_workflow_task.assert_not_awaited() server.client.fail_workflow_task.assert_not_awaited() - server.client.deregister_worker_registration.assert_awaited_once_with("cooperative-worker") + assert worker._registered + server.client.deregister_worker_registration.assert_not_awaited() replacement = ClaimServer(history=original_history, lease_owner="replacement-worker", workflow_task_attempt=5) cleanup_ids: list[str] = [] diff --git a/tests/test_cooperative_remote_worker.py b/tests/test_cooperative_remote_worker.py index 6d6610d..82f6e72 100644 --- a/tests/test_cooperative_remote_worker.py +++ b/tests/test_cooperative_remote_worker.py @@ -92,8 +92,7 @@ async def application_failure(kind: str) -> None: raise ActivityCancelled() -@pytest_asyncio.fixture -async def owner(monkeypatch: pytest.MonkeyPatch): # type: ignore[no-untyped-def] +async def make_owner(monkeypatch: pytest.MonkeyPatch) -> tuple[Worker, AsyncMock]: monkeypatch.delenv("DURABLE_WORKFLOW_WORKER_PROTOCOL_VERSION", raising=False) client = AsyncMock(spec=Client) client.get_cluster_info.return_value = compatible_cluster_info(worker_protocol={ @@ -106,6 +105,12 @@ async def owner(monkeypatch: pytest.MonkeyPatch): # type: ignore[no-untyped-def worker = Worker(client, task_queue="queue", worker_id="owner", capabilities=["cooperative_cancellation"], max_concurrent_activity_tasks=1) await worker._register() + return worker, client + + +@pytest_asyncio.fixture +async def owner(monkeypatch: pytest.MonkeyPatch): # type: ignore[no-untyped-def] + worker, client = await make_owner(monkeypatch) try: yield worker, client finally: @@ -152,18 +157,27 @@ async def acknowledge(**options: Any) -> dict[str, Any]: client.fail_activity_task.assert_not_awaited() -async def test_shutdown_timeout_joins_the_remote_callback_before_returning(owner: Any, tmp_path: Path) -> None: - worker, client = owner +async def test_shutdown_timeout_joins_the_remote_callback_before_returning( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +) -> None: + worker, client = await make_owner(monkeypatch) worker._shutdown_timeout = 0.01 marker = tmp_path / "callback" worker.activities["remote"] = blocked_async execution = worker._track(worker._run_activity_task(task(str(marker)))) await wait_for_file(marker) pid = int(marker.read_text()) - await asyncio.wait_for(worker.stop(), timeout=4.0) + stopping = worker.stop() + with pytest.raises(RuntimeError, match="deadline exhausted.*registration.*unconfirmed") as failed: + await asyncio.wait_for(stopping, timeout=4.0) await wait_for_exit(pid) assert execution.cancelled() or execution.result() == "claim_aborted" assert not worker._remote_activity_processes + assert worker._registered + client.deregister_worker_registration.assert_not_awaited() + with pytest.raises(RuntimeError) as repeated: + await worker.stop() + assert repeated.value is failed.value client.complete_activity_task.assert_not_awaited() client.fail_activity_task.assert_not_awaited() diff --git a/tests/test_replay_regression_corpus.py b/tests/test_replay_regression_corpus.py index 2619639..3a7c201 100644 --- a/tests/test_replay_regression_corpus.py +++ b/tests/test_replay_regression_corpus.py @@ -25,6 +25,7 @@ commands_to_server_commands, query_state, ) +from tests.deregistration_recovery import RecoveredWorkflow as FencedRecoveryWorkflow from tests.test_condition_replay_drift import ConditionIdentityWorkflow from tests.test_golden_history_replay import ( GoldenSagaCompensationWorkflow, @@ -371,6 +372,7 @@ def run(self, ctx: WorkflowContext, cleanup_seconds: int = 0): # type: ignore[n RepeatedPatchWorkflow, PendingOperation, RecoveredWorkflow, + FencedRecoveryWorkflow, PagedRecordedSideEffectsWorkflow, PostConditionReceiversWorkflow, RecordedSideEffectWorkflow, diff --git a/tests/test_storage_admission.py b/tests/test_storage_admission.py index 63c4b82..a8cc675 100644 --- a/tests/test_storage_admission.py +++ b/tests/test_storage_admission.py @@ -8,6 +8,7 @@ from pathlib import Path from types import SimpleNamespace from typing import Any +from urllib.parse import unquote import httpx import pytest @@ -544,7 +545,12 @@ def handler(request: httpx.Request) -> httpx.Response: info = compatible_cluster_info() info["worker_protocol"]["server_capabilities"]["query_tasks"] = False return httpx.Response(200, json=info) - if path.endswith("/register") or request.method == "DELETE": + if request.method == "DELETE": + return httpx.Response(200, json={ + "worker_id": unquote(path.rsplit("/", 1)[-1]), + "outcome": "deregistered", "recovered_workflow_task_count": 0, + }) + if path.endswith("/register"): return httpx.Response(200, json={"registered": True}) if path.endswith("/activity-tasks/poll"): claims += 1 @@ -656,7 +662,10 @@ def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(200, json=info) if request.method == "DELETE": deregistered += 1 - return httpx.Response(200, json={}) + return httpx.Response(200, json={ + "worker_id": unquote(path.rsplit("/", 1)[-1]), + "outcome": "deregistered", "recovered_workflow_task_count": 0, + }) if ( (phase == "registration" and path.endswith("/register")) or (phase == "poll" and path.endswith("/poll")) @@ -813,7 +822,12 @@ def handler(request: httpx.Request) -> httpx.Response: info = compatible_cluster_info() info["worker_protocol"]["server_capabilities"]["query_tasks"] = False return httpx.Response(200, json=info) - if path.endswith("/register") or request.method == "DELETE": + if request.method == "DELETE": + return httpx.Response(200, json={ + "worker_id": unquote(path.rsplit("/", 1)[-1]), + "outcome": "deregistered", "recovered_workflow_task_count": 0, + }) + if path.endswith("/register"): return httpx.Response(200, json={}) if path.endswith("/workflows/cold-replay"): return httpx.Response(200, json={"status": "completed" if completed else "running"}) diff --git a/tests/test_worker.py b/tests/test_worker.py index e06405d..7320d63 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -3226,7 +3226,7 @@ async def complete_validation(**_: object) -> dict[str, str]: assert worker._current_task_slots()["workflow_available"] == 1 @pytest.mark.asyncio - async def test_shutdown_cancels_validation_and_deregisters_after_releasing_capacity( + async def test_shutdown_releases_validation_capacity_without_deregistering_after_deadline( self, mock_client: AsyncMock ) -> None: validation_started = asyncio.Event() @@ -3279,15 +3279,16 @@ async def blocked_validation(_: dict[str, object]) -> str: await asyncio.wait_for(validation_started.wait(), timeout=1.0) assert worker._current_task_slots()["workflow_available"] == 0 - await worker.stop() + with pytest.raises(RuntimeError, match="deadline exhausted.*registration.*unconfirmed"): + await worker.stop() assert validation_cancelled.is_set() assert worker._workflow_inflight == 0 assert worker._workflow_reserved == 0 assert worker._wf_semaphore._value == 1 assert worker._current_task_slots()["workflow_available"] == 1 - assert worker._registered is False - mock_client.deregister_worker_registration.assert_awaited_once_with(worker.worker_id) + assert worker._registered is True + mock_client.deregister_worker_registration.assert_not_awaited() @pytest.mark.asyncio async def test_cancellation_before_dispatch_starts_releases_admitted_capacity(self, mock_client: AsyncMock) -> None: @@ -4516,7 +4517,7 @@ async def deregister_worker(_: str) -> dict[str, object]: mock_client.deregister_worker_registration.assert_awaited_once_with(worker.worker_id) @pytest.mark.asyncio - async def test_shutdown_deadline_cancels_work_before_deregistering( + async def test_shutdown_deadline_cancels_work_without_sending_late_deregistration( self, mock_client: AsyncMock ) -> None: events: list[str] = [] @@ -4539,11 +4540,13 @@ async def deregister_worker(_: str) -> dict[str, object]: tracked_task = worker._track(accepted_work()) await work_started.wait() - await worker.stop() + with pytest.raises(RuntimeError, match="deadline exhausted.*registration.*unconfirmed"): + await worker.stop() assert tracked_task.cancelled() - assert events == ["work_cancelled", "deregister"] - mock_client.deregister_worker_registration.assert_awaited_once_with(worker.worker_id) + assert events == ["work_cancelled"] + assert worker._registered is True + mock_client.deregister_worker_registration.assert_not_awaited() @pytest.mark.asyncio async def test_failed_registration_does_not_deregister(self, mock_client: AsyncMock) -> None: diff --git a/tests/test_worker_deregistration.py b/tests/test_worker_deregistration.py new file mode 100644 index 0000000..cea081b --- /dev/null +++ b/tests/test_worker_deregistration.py @@ -0,0 +1,361 @@ +from __future__ import annotations + +import asyncio +from typing import Any +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from durable_workflow import Client, ServerError, Worker +from durable_workflow.retry_policy import TransportRetryPolicy +from tests.test_worker import compatible_cluster_info + +TOKEN = "a" * 32 +FENCING = { + "schema": "durable-workflow.v2.worker-deregistration.v1", + "supported": True, + "receipt_retention_seconds": 600, + "endpoint": "/worker/registrations/{workerId}/deregister", +} + + +def completed(worker_id: str = "worker / one", **changes: Any) -> dict[str, Any]: + return { + "worker_id": worker_id, + "registration_token": TOKEN, + "outcome": "deregistered", + "recovered_workflow_task_count": 2, + **changes, + } + + +def refusal(**changes: Any) -> ServerError: + return ServerError(503, { + "reason": "backend_lock_pressure", + "operation": "deregister_worker", + "worker_id": "worker / one", + "registration_token": TOKEN, + "outcome": "unknown", + "retryable": True, + "retry_after_seconds": 1, + **changes, + }) + + +def worker_client(*, fenced: bool = True, token: Any = TOKEN) -> AsyncMock: + client = AsyncMock(spec=Client) + info = compatible_cluster_info() + if fenced: + info["worker_protocol"]["server_capabilities"]["worker_deregistration_fencing"] = dict(FENCING) # type: ignore[index] + client.get_cluster_info.return_value = info + client.register_worker.return_value = { + "registered": True, "worker_id": "worker / one", "registration_token": token, + } + client.deregister_worker_registration.return_value = completed() + return client + + +@pytest.mark.asyncio +async def test_fenced_client_uses_worker_credentials_and_exact_generation() -> None: + client = Client("http://runtime", namespace="orders", worker_token="worker", control_token="operator") + response = httpx.Response(200, request=httpx.Request("POST", "http://runtime"), json=completed()) + try: + with patch.object(client._http, "request", new_callable=AsyncMock, return_value=response) as request: + result = await client.deregister_worker_registration("worker / one", registration_token=TOKEN) + assert result == completed() + request.assert_awaited_once() + assert request.call_args.args[:2] == ("POST", "/api/worker/registrations/worker%20%2F%20one/deregister") + assert request.call_args.kwargs["json"] == {"registration_token": TOKEN} + headers = request.call_args.kwargs["headers"] + assert headers["Authorization"] == "Bearer worker" + assert headers["X-Namespace"] == "orders" + assert "X-Durable-Workflow-Control-Plane-Version" not in headers + finally: + await client.aclose() + + +@pytest.mark.parametrize("token", ["", "a" * 31, "a" * 33, "A" * 32, "z" * 32, 1, True]) +@pytest.mark.asyncio +async def test_fenced_client_rejects_invalid_token_before_http(token: Any) -> None: + async with Client("http://runtime") as client: + with ( + patch.object(client._http, "request", new_callable=AsyncMock) as request, + pytest.raises(ValueError, match="registration_token"), + ): + await client.deregister_worker_registration("worker / one", registration_token=token) + request.assert_not_awaited() + + +@pytest.mark.parametrize("changes", [ + {"worker_id": "replacement"}, {"registration_token": "b" * 32}, + {"outcome": "already_deregistered"}, {"outcome": "unknown"}, + {"recovered_workflow_task_count": -1}, {"recovered_workflow_task_count": True}, + {"recovered_workflow_task_count": 1.0}, +]) +@pytest.mark.asyncio +async def test_fenced_client_refuses_mismatched_or_malformed_receipt(changes: dict[str, Any]) -> None: + async with Client("http://runtime") as client: + response = httpx.Response(200, request=httpx.Request("POST", "http://runtime"), json=completed(**changes)) + with ( + patch.object(client._http, "request", new_callable=AsyncMock, return_value=response), + pytest.raises(ServerError) as error, + ): + await client.deregister_worker_registration("worker / one", registration_token=TOKEN) + assert error.value.reason() == "invalid_worker_deregistration_response" + + +@pytest.mark.parametrize("changes", [ + {"worker_id": "replacement"}, {"worker_id": None}, {"outcome": "unknown"}, {"outcome": None}, + {"recovered_workflow_task_count": -1}, {"recovered_workflow_task_count": True}, + {"recovered_workflow_task_count": 1.0}, {"recovered_workflow_task_count": None}, +]) +@pytest.mark.asyncio +async def test_legacy_http_success_still_requires_the_existing_receipt_identity(changes: dict[str, Any]) -> None: + async with Client("http://runtime") as client: + response = httpx.Response(200, request=httpx.Request("DELETE", "http://runtime"), json=completed(**changes)) + with ( + patch.object(client._http, "request", new_callable=AsyncMock, return_value=response) as request, + pytest.raises(ServerError) as error, + ): + await client.deregister_worker_registration("worker / one") + assert error.value.reason() == "invalid_worker_deregistration_response" + request.assert_awaited_once() + + +@pytest.mark.parametrize("fenced", [False, True]) +@pytest.mark.asyncio +async def test_deregistration_does_not_inherit_generic_503_retries(fenced: bool) -> None: + async with Client("http://runtime", retry_policy=TransportRetryPolicy(max_attempts=3)) as client: + response = httpx.Response(503, request=httpx.Request("POST", "http://runtime"), json={"reason": "unrelated"}) + kwargs = {"registration_token": TOKEN} if fenced else {} + with ( + patch.object(client._http, "request", new_callable=AsyncMock, return_value=response) as request, + pytest.raises(ServerError), + ): + await client.deregister_worker_registration("worker / one", **kwargs) + request.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_worker_keeps_registration_token_and_coalesces_fenced_shutdown() -> None: + client = worker_client() + worker = Worker(client, task_queue="orders", worker_id="worker / one") + await worker._register() + await asyncio.gather(worker.stop(), worker.stop()) + await worker.stop() + assert not worker._registered + client.deregister_worker_registration.assert_awaited_once() + call = client.deregister_worker_registration.call_args + assert call.args == ("worker / one",) + assert call.kwargs["registration_token"] == TOKEN + assert 0 < call.kwargs["timeout"] <= worker._shutdown_timeout + + +@pytest.mark.parametrize("token", [None, "", "bad-token", "A" * 32]) +@pytest.mark.asyncio +async def test_advertised_fencing_never_falls_back_after_invalid_registration_ack(token: Any) -> None: + client = worker_client(token=token) + worker = Worker(client, task_queue="orders", worker_id="worker / one") + with pytest.raises(RuntimeError, match="registration_token"): + await worker._register() + client.deregister_worker_registration.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_worker_retries_matching_transient_with_same_token_and_one_deadline() -> None: + client = worker_client() + client.deregister_worker_registration.side_effect = [refusal(), completed()] + worker = Worker(client, task_queue="orders", worker_id="worker / one", shutdown_timeout=3) + await worker._register() + await worker.stop() + calls = client.deregister_worker_registration.call_args_list + assert len(calls) == 2 + assert all(call.args == (worker.worker_id,) and call.kwargs["registration_token"] == TOKEN for call in calls) + assert 0 < calls[1].kwargs["timeout"] < calls[0].kwargs["timeout"] <= 3 + assert not worker._registered + + +@pytest.mark.parametrize("changes", [ + {"reason": "unrelated"}, {"operation": "destroy"}, {"worker_id": "replacement"}, + {"registration_token": "b" * 32}, {"outcome": "deferred"}, {"retryable": False}, + {"retry_after_seconds": True}, {"retry_after_seconds": 0}, {"retry_after_seconds": 1.0}, + {"retry_after_seconds": 61}, {"reason": {"unexpected": "object"}}, {"worker_id": None}, +]) +@pytest.mark.asyncio +async def test_worker_never_retries_incomplete_or_mismatched_transient(changes: dict[str, Any]) -> None: + client = worker_client() + failure = refusal(**changes) + client.deregister_worker_registration.side_effect = failure + worker = Worker(client, task_queue="orders", worker_id="worker / one") + await worker._register() + with pytest.raises(ServerError) as error: + await worker.stop() + assert error.value is failure + assert worker._registered + client.deregister_worker_registration.assert_awaited_once() + + +@pytest.mark.parametrize("changes", [ + {"schema": "unknown"}, {"endpoint": "https://other-runtime/deregister"}, + {"receipt_retention_seconds": 0}, {"receipt_retention_seconds": True}, + {"receipt_retention_seconds": 1.0}, +]) +@pytest.mark.asyncio +async def test_worker_rejects_malformed_advertised_fencing_before_registration(changes: dict[str, Any]) -> None: + client = worker_client() + client.get_cluster_info.return_value["worker_protocol"]["server_capabilities"][ + "worker_deregistration_fencing" + ].update(changes) + worker = Worker(client, task_queue="orders", worker_id="worker / one") + with pytest.raises(RuntimeError, match="worker_deregistration_fencing"): + await worker._register() + client.register_worker.assert_not_awaited() + client.deregister_worker_registration.assert_not_awaited() + + +@pytest.mark.parametrize("changed", ["worker_id", "namespace", "missing_token"]) +@pytest.mark.asyncio +async def test_worker_does_not_deregister_after_registration_identity_changes(changed: str) -> None: + client = worker_client() + client.namespace = "orders" + worker = Worker(client, task_queue="orders", worker_id="worker / one") + await worker._register() + if changed == "worker_id": + worker.worker_id = "replacement" + elif changed == "namespace": + client.namespace = "another-namespace" + else: + worker._registration_token = None + with pytest.raises(RuntimeError, match="registration.*unconfirmed"): + await worker.stop() + assert worker._registered + client.deregister_worker_registration.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_original_worker_error_remains_primary_during_fenced_cleanup_failure() -> None: + client = worker_client() + cleanup = ServerError(403, {"reason": "worker_registration_forbidden"}) + client.deregister_worker_registration.side_effect = cleanup + worker = Worker(client, task_queue="orders", worker_id="worker / one") + original = RuntimeError("original worker failure") + worker._heartbeat_loop = AsyncMock(side_effect=original) # type: ignore[method-assign] + with pytest.raises(RuntimeError) as error: + await worker.run() + assert error.value is original + assert error.value.__cause__ is cleanup + client.deregister_worker_registration.assert_awaited_once() + + +@pytest.mark.parametrize("fenced", [False, True]) +@pytest.mark.asyncio +async def test_client_checks_registration_namespace_before_constructing_http_headers(fenced: bool) -> None: + async with Client("http://runtime", namespace="replacement-namespace") as client: + kwargs = {"registration_token": TOKEN} if fenced else {} + with ( + patch.object(client._http, "request", new_callable=AsyncMock) as request, + pytest.raises(RuntimeError, match="namespace changed.*unconfirmed"), + ): + await client.deregister_worker_registration("worker / one", expected_namespace="original", **kwargs) + request.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_namespace_cannot_change_between_fenced_retry_attempts() -> None: + client = worker_client() + client.namespace = "original" + worker = Worker(client, task_queue="orders", worker_id="worker / one", shutdown_timeout=3) + await worker._register() + + async def change_namespace(*_: Any, **kwargs: Any) -> None: + assert kwargs["expected_namespace"] == "original" + client.namespace = "replacement-namespace" + raise refusal() + + client.deregister_worker_registration.side_effect = change_namespace + with pytest.raises(RuntimeError, match="identity changed.*unconfirmed"): + await worker.stop() + assert worker._registered + client.deregister_worker_registration.assert_awaited_once() + + +@pytest.mark.parametrize("fenced", [False, True]) +@pytest.mark.asyncio +async def test_shutdown_bounds_final_http_and_keeps_uncertain_registration(fenced: bool) -> None: + client = worker_client(fenced=fenced) + cancelled = asyncio.Event() + + async def stalled(*_: Any, **__: Any) -> None: + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + client.deregister_worker_registration.side_effect = stalled + worker = Worker(client, task_queue="orders", worker_id="worker / one", shutdown_timeout=0.01) + await worker._register() + # The outer bound diagnoses the original unbounded implementation without + # leaving its shielded shutdown task alive after the regression fails. + try: + stopping = worker.stop() + with pytest.raises(RuntimeError, match="registration.*unconfirmed"): + await asyncio.wait_for(stopping, timeout=0.3) + finally: + if worker._shutdown_task is not None and not worker._shutdown_task.done(): + worker._shutdown_task.cancel() + await asyncio.gather(worker._shutdown_task, return_exceptions=True) + assert cancelled.is_set() + assert worker._registered + client.deregister_worker_registration.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_lost_fenced_response_retries_same_generation_without_legacy_delete() -> None: + client = worker_client() + lost = httpx.RemoteProtocolError("lost response", request=httpx.Request("POST", "http://runtime")) + client.deregister_worker_registration.side_effect = [lost, completed()] + worker = Worker(client, task_queue="orders", worker_id="worker / one", shutdown_timeout=1) + await worker._register() + await worker.stop() + calls = client.deregister_worker_registration.call_args_list + assert len(calls) == 2 + assert all(call.kwargs["registration_token"] == TOKEN for call in calls) + + +@pytest.mark.parametrize("status,reason", [ + (403, "worker_registration_forbidden"), (404, "worker_registration_token_not_found"), + (409, "worker_registration_lost_authority"), (410, "worker_registration_receipt_expired"), + (412, "worker_registration_transaction_boundary_unavailable"), +]) +@pytest.mark.asyncio +async def test_worker_surfaces_definitive_refusal_without_retry(status: int, reason: str) -> None: + client = worker_client() + failure = ServerError(status, {"reason": reason}) + client.deregister_worker_registration.side_effect = failure + worker = Worker(client, task_queue="orders", worker_id="worker / one") + await worker._register() + with pytest.raises(ServerError) as error: + await worker.stop() + assert error.value is failure + assert worker._registered + client.deregister_worker_registration.assert_awaited_once() + + +def test_recovered_registration_replays_genuine_completed_activity() -> None: + import json + from pathlib import Path + + from durable_workflow import serializer + from durable_workflow.workflow import replay + from tests.deregistration_recovery import RecoveredWorkflow + + path = Path(__file__).parent / "fixtures/replay_regressions/fenced-deregistration-cold-recovery.json" + fixture = json.loads(path.read_text()) + outcome = replay( + RecoveredWorkflow, fixture["history"], fixture["workflow"]["input"], run_id=fixture["workflow"]["run_id"], + ) + assert len(outcome.commands) == 1 + command = outcome.commands[0].to_server_command("recovered") + assert command["type"] == "complete_workflow" + assert serializer.decode_envelope(command["result"]) == fixture["expected"]["result"] diff --git a/tests/test_worker_observation.py b/tests/test_worker_observation.py index 34b2422..1fc85cb 100644 --- a/tests/test_worker_observation.py +++ b/tests/test_worker_observation.py @@ -297,3 +297,15 @@ async def fail(context): assert report["milestones"]["activity_handler_error"]["count"] == 1 assert report["milestones"]["workflow_task_enter"]["first"]["task_id"] == "original" assert "SECRET" not in json.dumps(report) + + +@pytest.mark.parametrize(("method", "path", "expected"), [ + ("POST", "/api/worker/register", "registration"), + ("POST", "/api/worker/registrations/worker/deregister", "deregistration"), + ("DELETE", "/api/worker/registrations/worker", "deregistration"), + ("GET", "/api/worker/registrations/worker", "registration"), +]) +def test_registration_and_deregistration_are_counted_separately(method: str, path: str, expected: str) -> None: + from tests.worker_observation import route_name + + assert route_name(httpx.Request(method, "http://runtime" + path)) == expected diff --git a/tests/worker_observation.py b/tests/worker_observation.py index 5e090f9..8d38a4e 100644 --- a/tests/worker_observation.py +++ b/tests/worker_observation.py @@ -27,7 +27,11 @@ def route_name(request: httpx.Request) -> str: return f"{family}_poll" if f"/worker/{family}-tasks/" in path: return f"{family}_submission" - if path.endswith("/worker/register") or "/worker/registrations" in path: + if "/worker/registrations/" in path: + if request.method == "DELETE" or request.method == "POST" and path.endswith("/deregister"): + return "deregistration" + return "registration" + if path.endswith("/worker/register"): return "registration" if path.endswith("/cluster/info"): return "discovery"