Preserve worker poll identity across backend outages (#73)

* Preserve Python worker poll identity during backend outages

* Fail closed on ambiguous backend refusal contracts

---------

Co-authored-by: Durable Workflow <support@durable-workflow.com>

Author: rmcdaniel

src/durable_workflow/retry_policy.py

@@ -65,6 +65,61 @@ def _storage_refusal(exc: Exception) -> tuple[ServerError, str | None] | None:
     return error, poll_id
 
 
+def _backend_unavailable_refusal(exc: Exception) -> tuple[bool, int | None]:
+    if not isinstance(exc, httpx.HTTPStatusError) or exc.response.status_code != 503:
+        return False, None
+    try:
+        body = exc.response.json()
+    except ValueError:
+        return False, None
+    if not isinstance(body, dict) or body.get("reason") != "backend_unavailable":
+        return False, None
+    request = exc.request
+    if request.method != "POST" or "X-Durable-Workflow-Protocol-Version" not in request.headers:
+        return True, None
+    operations = {
+        "/api/worker/workflow-tasks/poll": "poll_workflow_task",
+        "/api/worker/activity-tasks/poll": "poll_activity_task",
+        "/api/worker/query-tasks/poll": "poll_query_task",
+        "/api/worker/register": "register_worker",
+        "/api/worker/heartbeat": "heartbeat_worker",
+    }
+    operation = next((name for path, name in operations.items() if request.url.path.endswith(path)), None)
+    if operation is None:
+        return True, None
+    try:
+        submitted = json.loads(request.content)
+    except ValueError:
+        return True, None
+    if not isinstance(submitted, dict):
+        return True, None
+    worker_id = submitted.get("worker_id")
+    queue = submitted.get("task_queue")
+    delay = body.get("retry_after_seconds")
+    if (
+        not isinstance(worker_id, str) or not worker_id
+        or body.get("operation") != operation
+        or body.get("outcome") != "unknown"
+        or body.get("worker_id") != worker_id
+        or body.get("task_queue") != queue
+        or body.get("retryable") is not True
+        or type(delay) is not int or delay <= 0
+        or (operation != "heartbeat_worker" and (not isinstance(queue, str) or not queue))
+    ):
+        return True, None
+    if operation.startswith("poll_"):
+        poll_id = submitted.get("poll_request_id")
+        if (
+            not isinstance(poll_id, str) or not poll_id
+            or "task" not in body or body["task"] is not None
+            or body.get("poll_status") != "backend_unavailable"
+            or body.get("poll_request_id") != poll_id
+            or body.get("retry_same_poll_request_id") is not True
+        ):
+            return True, None
+    return True, delay
+
+
 @dataclass
 class TransportRetryPolicy:
     """
@@ -123,7 +178,7 @@ async def execute(self, fn: Callable[[], Awaitable[T]]) -> T:
         Raises the last exception if all retries are exhausted.
         """
         attempt = 0
-        storage_attempt = 0
+        worker_pause_attempt = 0
         last_exc: Exception | None = None
 
         while attempt < self.max_attempts:
@@ -134,20 +189,30 @@ async def execute(self, fn: Callable[[], Awaitable[T]]) -> T:
                 last_exc = exc
                 stop = _worker_storage_admission_stop.get()
                 refusal = _storage_refusal(exc) if stop is not None else None
-                if refusal is not None and stop is not None:
+                backend_refusal, backend_delay = _backend_unavailable_refusal(exc)
+                if backend_refusal and backend_delay is None:
+                    raise
+                pause: tuple[str, int] | None = None
+                if refusal is not None:
                     error, poll_id = refusal
-                    if not error.is_storage_admission_failure(poll_id) or stop():
+                    if not error.is_storage_admission_failure(poll_id):
                         raise
-                    storage_attempt += 1
                     assert isinstance(error.body, dict)
+                    pause = ("storage admission paused", error.body["retry_after_seconds"])
+                elif backend_delay is not None:
+                    pause = ("worker backend unavailable", backend_delay)
+                if pause is not None and stop is not None:
+                    if stop():
+                        raise
+                    worker_pause_attempt += 1
                     delay = min(
                         5.0,
                         max(
-                            self.backoff_seconds(min(storage_attempt - 1, 6)),
-                            error.body["retry_after_seconds"],
+                            self.backoff_seconds(min(worker_pause_attempt - 1, 6)),
+                            pause[1],
                         ),
                     )
-                    log.warning("storage admission paused; retrying the same worker request in %.2fs", delay)
+                    log.warning("%s; retrying the same worker request in %.2fs", pause[0], delay)
                     # Do not consume the finite transport budget or repeat serialization/uploads.
                     while delay > 0:
                         if stop():

tests/test_storage_admission.py

@@ -39,6 +39,30 @@ def pressure(poll_id: str | None = None, *, reason: str = "storage_pressure", en
     return body
 
 
+def backend_unavailable(request: httpx.Request) -> dict[str, Any]:
+    submitted = json.loads(request.content)
+    path = request.url.path
+    operation = {
+        "/api/worker/workflow-tasks/poll": "poll_workflow_task",
+        "/api/worker/activity-tasks/poll": "poll_activity_task",
+        "/api/worker/query-tasks/poll": "poll_query_task",
+        "/api/worker/register": "register_worker",
+        "/api/worker/heartbeat": "heartbeat_worker",
+    }[path]
+    body: dict[str, Any] = {
+        "reason": "backend_unavailable", "operation": operation, "outcome": "unknown",
+        "worker_id": submitted["worker_id"], "task_queue": submitted.get("task_queue"),
+        "retryable": True, "retry_after_seconds": 1,
+    }
+    if path.endswith("/poll"):
+        body.update({
+            "task": None, "poll_status": "backend_unavailable",
+            "poll_request_id": submitted["poll_request_id"],
+            "retry_same_poll_request_id": True,
+        })
+    return body
+
+
 @contextmanager
 def worker_scope(stop: Callable[[], bool] = lambda: False) -> Iterator[None]:
     token = _worker_storage_admission_stop.set(stop)
@@ -104,6 +128,147 @@ def handler(request: httpx.Request) -> httpx.Response:
     assert sum(retry_sleeps) == pytest.approx(5)
 
 
+@pytest.mark.parametrize("kind", ["workflow", "activity", "query", "multiplexed"])
+async def test_backend_outage_preserves_logical_poll_after_transport_budget(
+    kind: str, retry_sleeps: list[float],
+) -> None:
+    requests: list[bytes] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        requests.append(request.content)
+        if len(requests) == 1:
+            raise httpx.ReadTimeout("Response lost after a possible claim.", request=request)
+        if len(requests) <= 5:
+            return httpx.Response(503, json=backend_unavailable(request))
+        return httpx.Response(200, json={"task": {"task_id": "reconciled-claim"}})
+
+    async with client_for(handler) as client:
+        with worker_scope():
+            if kind == "multiplexed":
+                task = await client.poll_workflow_task(
+                    worker_id="backend-worker", task_queue="orders",
+                    task_kinds=("workflow", "update_validation"),
+                )
+            else:
+                task = await getattr(client, f"poll_{kind}_task")(
+                    worker_id="backend-worker", task_queue="orders",
+                )
+    assert task == {"task_id": "reconciled-claim"}
+    assert len(requests) == 6
+    assert len(set(requests)) == 1
+    assert sum(retry_sleeps) == pytest.approx(4)
+
+
+@pytest.mark.parametrize("method,kwargs", [
+    ("register_worker", {
+        "worker_id": "backend-worker", "task_queue": "orders",
+        "capability_manifest": PORTABLE_WORKER_AFFINITY_CAPABILITY_MANIFEST,
+    }),
+    ("heartbeat_worker", {"worker_id": "backend-worker"}),
+])
+async def test_backend_outage_retries_worker_registration_and_heartbeat(
+    method: str, kwargs: dict[str, Any], retry_sleeps: list[float],
+) -> None:
+    requests: list[bytes] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        requests.append(request.content)
+        if len(requests) <= 4:
+            return httpx.Response(503, json=backend_unavailable(request))
+        return httpx.Response(200, json={"recorded": True})
+
+    async with client_for(handler) as client:
+        with worker_scope():
+            result = await getattr(client, method)(**kwargs)
+    assert result == {"recorded": True}
+    assert len(requests) == 5
+    assert len(set(requests)) == 1
+    assert sum(retry_sleeps) == pytest.approx(4)
+
+
+@pytest.mark.parametrize("override", [
+    {"operation": "poll_activity_task"},
+    {"outcome": "failed"}, {"worker_id": "other-worker"},
+    {"task_queue": "other-queue"}, {"retryable": False},
+    {"retry_after_seconds": 0}, {"retry_after_seconds": True},
+    {"retry_after_seconds": "1"}, {"task": {"task_id": "claimed"}},
+    {"poll_status": "empty"}, {"poll_request_id": "another-poll"},
+    {"retry_same_poll_request_id": False},
+])
+async def test_invalid_backend_outage_contract_remains_bounded(
+    override: dict[str, Any], retry_sleeps: list[float],
+) -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(503, json={**backend_unavailable(request), **override})
+
+    async with client_for(handler) as client:
+        with worker_scope(), pytest.raises(ServerError):
+            await client.poll_workflow_task(worker_id="backend-worker", task_queue="orders")
+    assert calls == 1
+    assert not retry_sleeps
+
+
+async def test_legacy_validation_backend_refusal_does_not_repeat_ambiguous_poll(
+    retry_sleeps: list[float],
+) -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(503, json={
+            "reason": "backend_unavailable", "operation": "poll_update_validation_task",
+            "outcome": "unknown", "worker_id": "backend-worker", "task_queue": "orders",
+            "retryable": False, "retry_after_seconds": 1, "task": None,
+            "poll_status": "backend_unavailable", "poll_request_id": None,
+            "retry_same_poll_request_id": True,
+        })
+
+    async with client_for(handler) as client:
+        with worker_scope(), pytest.raises(ServerError):
+            await client.poll_update_validation_task(worker_id="backend-worker", task_queue="orders")
+    assert calls == 1
+    assert not retry_sleeps
+
+
+async def test_direct_backend_outage_poll_remains_bounded(retry_sleeps: list[float]) -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(503, json=backend_unavailable(request))
+
+    async with client_for(handler) as client:
+        with pytest.raises(ServerError):
+            await client.poll_workflow_task(worker_id="backend-worker", task_queue="orders")
+    assert calls == 2
+
+
+async def test_shutdown_interrupts_backend_outage_without_new_poll(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls = 0
+    stopped = False
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(503, json=backend_unavailable(request))
+
+    async def stop_on_sleep(delay: float) -> None:
+        nonlocal stopped
+        stopped = True
+
+    monkeypatch.setattr(retry_module, "asyncio", SimpleNamespace(sleep=stop_on_sleep))
+    async with client_for(handler) as client:
+        with worker_scope(lambda: stopped), pytest.raises(ServerError):
+            await client.poll_workflow_task(worker_id="backend-worker", task_queue="orders")
+    assert calls == 1
+
+
 @pytest.mark.parametrize("method,kwargs", [
     ("register_worker", {
         "worker_id": "storage-worker", "task_queue": "orders",