diff --git a/src/durable_workflow/_activity_lease.py b/src/durable_workflow/_activity_lease.py new file mode 100644 index 0000000..6b289a7 --- /dev/null +++ b/src/durable_workflow/_activity_lease.py @@ -0,0 +1,319 @@ +"""Runtime lease maintenance for one original activity delivery.""" + +from __future__ import annotations + +import asyncio +import re +import time +from collections.abc import Awaitable, Callable, Mapping +from contextvars import ContextVar +from dataclasses import dataclass, field +from datetime import datetime, timedelta +from typing import Any, TypeVar +from urllib.parse import quote + +import httpx + +from .errors import ActivityLeaseAuthorityLost, ServerError + +_T = TypeVar("_T") +_UTC6 = re.compile(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{6}Z\Z") + + +def _timestamp(value: Any) -> datetime: + if not isinstance(value, str) or _UTC6.fullmatch(value) is None: + raise ValueError("activity timing must be UTC with six fractional digits") + return datetime.fromisoformat(value.replace("Z", "+00:00")) + + +def activity_lease_renewal_supported(client: Any) -> bool: + """Use the same cached discovery that validates actual renewal requests.""" + info = getattr(client, "_cluster_info", None) + protocol = info.get("worker_protocol") if isinstance(info, Mapping) else None + capabilities = protocol.get("server_capabilities") if isinstance(protocol, Mapping) else None + return isinstance(capabilities, Mapping) and capabilities.get("activity_lease_renewal") is True + + +def _heartbeat_terms(task: Mapping[str, Any], deadlines: Mapping[str, Any], + reference: datetime) -> tuple[datetime | None, int | None]: + heartbeat = _timestamp(deadlines["heartbeat"]) if deadlines.get("heartbeat") is not None else None + policy = task.get("retry_policy") + interval = policy.get("heartbeat_timeout") if isinstance(policy, Mapping) else None + if heartbeat is not None and (type(interval) is not int or interval <= 0): + raise ValueError("missing recorded activity heartbeat interval") + interval = interval if type(interval) is int and interval > 0 else None + if heartbeat is not None and interval is not None and heartbeat - reference > timedelta(seconds=interval): + raise ValueError("activity heartbeat deadline predates its claimed progress") + return heartbeat, interval + + +@dataclass +class ActivityLeaseAuthority: + client: Any + namespace: str + identity: dict[str, Any] + task_identity: dict[str, Any] + captured_at: float + lease_until: float + business_until: float | None + heartbeat_until: float | None + heartbeat_timeout: int | None + server_time: datetime + lease_expiry: datetime + heartbeat_time: datetime | None + heartbeat_deadline: datetime | None + last_renewal: tuple[datetime, datetime, float] | None = None + lost_reason: str | None = field(default=None, init=False) + + @classmethod + async def confirm_delivery(cls, client: Any, namespace: str, task: Mapping[str, Any], *, owner: str, + task_queue: str, server_time: Any) -> ActivityLeaseAuthority: + # A long poll may wait before a task or its lease exists. Establish the + # first execution budget with a separate, original-claim confirmation; + # its entire round trip is charged and it cannot revive an expired lease. + provisional = cls.delivery(client, namespace, task, owner=owner, task_queue=task_queue, + server_time=server_time, request_started=time.monotonic()) + started = time.monotonic() + try: + reply = await client.renew_activity_task_lease( + task_id=provisional.task_id, activity_attempt_id=provisional.identity["activity_attempt_id"], + lease_owner=owner, timeout=min(5.0, provisional.remaining()), + ) + provisional.acknowledge(reply, started) + except ActivityLeaseAuthorityLost: + raise + except Exception as error: + reason = error.reason() if isinstance(error, ServerError) else None + raise ActivityLeaseAuthorityLost( + provisional.task_id, reason or "activity_initial_lease_unavailable", + ) from error + observed = dict(task) + observed["lease_expires_at"] = reply["lease_expires_at"] + if provisional.heartbeat_deadline is not None: + observed["deadlines"] = {**dict(task.get("deadlines", {})), + "heartbeat": provisional.heartbeat_deadline.isoformat( + timespec="microseconds").replace("+00:00", "Z")} + return cls.delivery(client, namespace, observed, owner=owner, task_queue=task_queue, + server_time=reply["server_time"], request_started=started) + + @classmethod + def delivery(cls, client: Any, namespace: str, task: Mapping[str, Any], *, owner: str, + task_queue: str, server_time: Any, request_started: float) -> ActivityLeaseAuthority: + task_id = str(task.get("task_id") or "") + try: + keys = ("task_id", "workflow_id", "run_id", "activity_execution_id", "activity_attempt_id", + "lease_owner", "task_queue") + if (any(not isinstance(task.get(key), str) or not task[key] for key in keys) + or task["lease_owner"] != owner or task["task_queue"] != task_queue + or task.get("namespace", namespace) != namespace + or type(task.get("attempt_number")) is not int or task["attempt_number"] < 1): + raise ValueError("invalid original activity identity") + reference = _timestamp(server_time) + expiry = _timestamp(task.get("lease_expires_at")) + duration = (expiry - reference).total_seconds() + if not 0 < duration <= 86400 or not 0 <= request_started <= time.monotonic(): + raise ValueError("invalid original activity lease") + deadlines = task.get("deadlines", {}) + if not isinstance(deadlines, Mapping): + raise ValueError("invalid activity deadlines") + bounds = [_timestamp(deadlines[key]) for key in ("start_to_close", "schedule_to_close", "run", "execution") + if deadlines.get(key) is not None] + business = request_started + (min(bounds) - reference).total_seconds() if bounds else None + heartbeat, interval = _heartbeat_terms(task, deadlines, reference) + result = cls( + client, namespace, + {"workflow_task_id": task["task_id"], "workflow_instance_id": task["workflow_id"], + "workflow_run_id": task["run_id"], "activity_execution_id": task["activity_execution_id"], + "activity_attempt_id": task["activity_attempt_id"], "attempt_number": task["attempt_number"], + "lease_owner": owner}, + {key: task[key] for key in (*keys, "attempt_number")}, + request_started, request_started + duration, business, + request_started + (heartbeat - reference).total_seconds() if heartbeat is not None else None, + interval, reference, expiry, + heartbeat - timedelta(seconds=interval) if heartbeat is not None and interval is not None else None, + heartbeat, + ) + except (KeyError, TypeError, ValueError, OverflowError) as error: + raise ActivityLeaseAuthorityLost(task_id, "invalid_activity_claim") from error + result.remaining() + return result + + @property + def task_id(self) -> str: + return str(self.identity["workflow_task_id"]) + + def matches(self, client: Any, namespace: str, task: Mapping[str, Any], owner: str, task_queue: str) -> bool: + return (self.client is client and self.namespace == namespace and self.identity["lease_owner"] == owner + and self.task_identity["task_queue"] == task_queue + and type(task.get("attempt_number")) is int + and all(task.get(key) == value for key, value in self.task_identity.items())) + + def governs(self, client: Any, namespace: str, path: str) -> bool: + prefixes = (f"/worker/activity-tasks/{self.task_id}/", + f"/worker/activity-tasks/{quote(self.task_id, safe='')}/") + return (self.client is client and self.namespace == namespace + and (any(path == prefix + action for prefix in prefixes + for action in ("complete", "fail", "heartbeat", "renew-lease")) + or path == "/external-payloads/v1")) + + def acknowledge_progress(self, reply: Any, request_started: float) -> None: + """Use an actual application heartbeat, never a runtime liveness ping. + + The existing heartbeat contract records last_heartbeat_at at the write + and resets its original recorded timeout interval from that instant. + Charging the whole request duration consumes storage/network latency. + """ + self.remaining() + received = time.monotonic() + try: + expected = {"task_id": self.task_id, "activity_attempt_id": self.identity["activity_attempt_id"], + "lease_owner": self.identity["lease_owner"]} + if (not isinstance(reply, Mapping) or reply.get("heartbeat_recorded") is not True + or reply.get("can_continue") is not True or reply.get("cancel_requested") is not False + or reply.get("reason") is not None + or any(reply.get(key) != value for key, value in expected.items()) + or not self.captured_at <= request_started <= received): + raise ValueError("heartbeat did not confirm the original activity") + progress = _timestamp(reply.get("last_heartbeat_at")) + expiry = _timestamp(reply.get("lease_expires_at")) + duration = (expiry - progress).total_seconds() + if not 0 < duration <= 86400: + raise ValueError("invalid heartbeat lease grant") + if self.heartbeat_time is not None and progress <= self.heartbeat_time: + return + if request_started + duration <= received: + raise ValueError("heartbeat lease grant was consumed in transport") + if self.heartbeat_timeout is not None: + until = request_started + self.heartbeat_timeout + if until <= received: + raise ValueError("heartbeat progress window was consumed in transport") + self.heartbeat_until = max(self.heartbeat_until or until, until) + self.heartbeat_deadline = progress + timedelta(seconds=self.heartbeat_timeout) + self.heartbeat_time = progress + self.lease_until = max(self.lease_until, request_started + duration) + self.lease_expiry = max(self.lease_expiry, expiry) + except (TypeError, ValueError, OverflowError) as error: + self.lost_reason = "invalid_activity_heartbeat" + raise ActivityLeaseAuthorityLost(self.task_id, self.lost_reason) from error + + def remaining(self) -> float: + if self.lost_reason is not None: + raise ActivityLeaseAuthorityLost(self.task_id, self.lost_reason) + bounds = [self.lease_until, *(value for value in (self.business_until, self.heartbeat_until) + if value is not None)] + remaining = min(bounds) - time.monotonic() + if remaining <= 0: + self.lost_reason = "activity_execution_budget_expired" + raise ActivityLeaseAuthorityLost(self.task_id, self.lost_reason) + return remaining + + def acknowledge(self, reply: Any, request_started: float) -> None: + self.remaining() + received = time.monotonic() + try: + if (not isinstance(reply, Mapping) or reply.get("lease_renewed") is not True + or reply.get("can_continue") is not True or reply.get("cancel_requested") is not False + or reply.get("heartbeat_recorded") is not False or reply.get("reason") is not None + or type(reply.get("attempt_number")) is not int + or any(reply.get(key) != value for key, value in self.identity.items()) + or not self.captured_at <= request_started <= received): + raise ValueError("renewal did not preserve the original activity claim") + server = _timestamp(reply.get("server_time")) + expiry = _timestamp(reply.get("lease_expires_at")) + duration = (expiry - server).total_seconds() + if not 0 < duration <= 86400 or request_started + duration <= received: + raise ValueError("renewal did not prove a live grant") + if (self.heartbeat_time is not None and server < self.heartbeat_time and expiry <= self.lease_expiry): + # An independently acknowledged application heartbeat can + # overtake an in-flight renewal. Its older reply grants nothing. + return + if expiry < self.lease_expiry or server < self.server_time: + raise ValueError("renewal predates the confirmed grant") + confirmation = (server, expiry, request_started) + if self.last_renewal is not None and server <= self.last_renewal[0]: + if confirmation == self.last_renewal: + return + raise ValueError("replayed activity renewal") + self.lease_until = max(self.lease_until, request_started + duration) + self._observe_progress(reply, server, request_started, received) + self.server_time, self.lease_expiry, self.last_renewal = server, expiry, confirmation + except (TypeError, ValueError, OverflowError) as error: + self.lost_reason = "invalid_activity_renewal" + raise ActivityLeaseAuthorityLost(self.task_id, self.lost_reason) from error + + def _observe_progress(self, reply: Mapping[str, Any], server: datetime, started: float, received: float) -> None: + if self.heartbeat_until is None or self.heartbeat_timeout is None: + return + try: + progress = _timestamp(reply.get("last_heartbeat_at")) + deadline = _timestamp(reply.get("heartbeat_deadline_at")) + except (TypeError, ValueError): + return + until = started + (deadline - server).total_seconds() + if (progress <= server < deadline and until > received + and (self.heartbeat_time is None or progress > self.heartbeat_time) + and (self.heartbeat_deadline is None or deadline > self.heartbeat_deadline) + and deadline - progress == timedelta(seconds=self.heartbeat_timeout)): + self.heartbeat_until = max(self.heartbeat_until, until) + self.heartbeat_time, self.heartbeat_deadline = progress, deadline + + async def renew(self) -> bool: + started = time.monotonic() + try: + reply = await asyncio.wait_for(self.client.renew_activity_task_lease( + task_id=self.task_id, activity_attempt_id=self.identity["activity_attempt_id"], + lease_owner=self.identity["lease_owner"], timeout=min(5.0, self.remaining()), + ), timeout=min(5.0, self.remaining())) + self.acknowledge(reply, started) + return True + except ActivityLeaseAuthorityLost: + raise + except Exception as error: + if (isinstance(error, (httpx.TransportError, asyncio.TimeoutError)) + or isinstance(error, ServerError) and (error.status >= 500 or error.status == 429)): + # A failed exchange grants nothing. The existing lease may + # still authorize work until its original local bound. + self.remaining() + return False + self.lost_reason = "activity_renewal_unavailable" + raise ActivityLeaseAuthorityLost(self.task_id, self.lost_reason) from error + + async def _maintain(self) -> None: + confirmed = True + while True: + delay = min(30.0 if confirmed else 1.0, max(0.0, self.lease_until - time.monotonic()) / 3) + if self.heartbeat_timeout is not None: + delay = min(delay, self.heartbeat_timeout / 3) + await asyncio.sleep(min(delay, self.remaining())) + confirmed = await self.renew() + + async def run(self, operation: Callable[[], Awaitable[_T]]) -> _T: + self.remaining() + execution = asyncio.ensure_future(operation()) + maintenance = asyncio.create_task(self._maintain()) + try: + while True: + done, _ = await asyncio.wait({execution, maintenance}, timeout=self.remaining(), + return_when=asyncio.FIRST_COMPLETED) + if maintenance in done: + await maintenance + raise ActivityLeaseAuthorityLost(self.task_id, "activity_renewal_stopped") + if execution in done: + result = await execution + self.remaining() + return result + except BaseException: + self.lost_reason = self.lost_reason or "activity_execution_abandoned" + raise + finally: + # A Python thread may continue after cancellation. Its existing + # executor keeps the physical slot until it actually exits, while + # this context independently fences every late HTTP publication. + for future in (execution, maintenance): + future.cancel() + await asyncio.gather(execution, maintenance, return_exceptions=True) + + +current_activity_authority: ContextVar[ActivityLeaseAuthority | None] = ContextVar( + "activity_lease_authority", default=None, +) diff --git a/src/durable_workflow/client.py b/src/durable_workflow/client.py index b345cd8..9810e6e 100644 --- a/src/durable_workflow/client.py +++ b/src/durable_workflow/client.py @@ -39,6 +39,7 @@ import httpx from . import serializer +from ._activity_lease import ActivityLeaseAuthority, current_activity_authority from ._cancellation_scope import CancellationScopeOpenReceipt, scope_identity from ._cancellation_scope_history import ( CancellationScopeBudget, @@ -1871,8 +1872,14 @@ async def _request( raise ValueError("cluster discovery credentials can only authorize /cluster/info") query_poll = worker and path == "/worker/query-tasks/poll" + activity_poll = worker and path == "/worker/activity-tasks/poll" if query_poll: current_query_authority.set(None) + if activity_poll: + current_activity_authority.set(None) + activity_authority = current_activity_authority.get() + if worker and activity_authority is not None and activity_authority.governs(self, self.namespace, path): + activity_authority.remaining() request_json = json if ( not discovery @@ -1904,14 +1911,20 @@ async def _request( plane = "worker" if worker or discovery_uses_worker_token else "control" status_code = "none" outcome = "pending" + request_started = time.monotonic() async def _do_request() -> httpx.Response: + nonlocal request_started + request_started = time.monotonic() + request_timeout = self.timeout if timeout is None else timeout + if worker and activity_authority is not None and activity_authority.governs(self, self.namespace, path): + request_timeout = min(request_timeout, activity_authority.remaining()) resp = await self._http.request( method, f"/api{path}", headers=self._discovery_headers() if discovery else self._headers(worker=worker), json=request_json, - timeout=self.timeout if timeout is None else timeout, + timeout=request_timeout, ) # Raise HTTPStatusError for 4xx/5xx so retry policy can catch it resp.raise_for_status() @@ -1936,11 +1949,28 @@ async def _do_request() -> httpx.Response: outcome = "ok" return None result = resp.json() + if (worker and activity_authority is not None + and activity_authority.governs(self, self.namespace, path) + and path.endswith("/heartbeat") and isinstance(result, dict) + and result.get("heartbeat_recorded") is True): + activity_authority.acknowledge_progress(result, request_started) + if activity_poll and isinstance(result, dict) and isinstance(result.get("task"), dict): + capabilities = result.get("server_capabilities") + if isinstance(capabilities, dict) and capabilities.get("activity_lease_renewal") is True: + activity_authority = await ActivityLeaseAuthority.confirm_delivery( + self, self.namespace, result["task"], owner=request_json["worker_id"], + task_queue=request_json["task_queue"], server_time=result.get("server_time"), + ) + current_activity_authority.set(activity_authority) query_authority = None if query_poll and isinstance(result, dict) and isinstance(result.get("task"), dict): query_authority = await self._query_delivery_authority(result["task"], request_json) current_query_authority.set(query_authority) - if not discovery and _contains_runtime_payload_envelope(result): + if activity_authority is not None and activity_poll and _contains_runtime_payload_envelope(result): + result = await activity_authority.run( + lambda: self._resolve_activity_claim_payloads(result, worker=worker), + ) + elif not discovery and _contains_runtime_payload_envelope(result): transport = (await query_authority.run(self._runtime_external_payload_transport) if query_authority is not None else await self._runtime_external_payload_transport()) if transport is None: @@ -1957,6 +1987,8 @@ async def _do_request() -> httpx.Response: result, worker=worker, transport=transport, path=path, request_json=request_json, ) result = self._hydrate_result_envelopes(result) + if activity_poll and activity_authority is not None: + activity_authority.remaining() if query_authority is not None: await query_authority.confirm() outcome = "ok" @@ -1976,6 +2008,12 @@ async def _do_request() -> httpx.Response: self.metrics.increment(CLIENT_REQUESTS, tags=tags) self.metrics.record(CLIENT_REQUEST_DURATION_SECONDS, time.perf_counter() - start, tags=tags) + async def _resolve_activity_claim_payloads(self, value: Any, *, worker: bool) -> Any: + transport = await self._runtime_external_payload_transport() + if transport is None: + raise ExternalPayloadUnsupported("activity payload reference has no authenticated runtime transport") + return await self._resolve_runtime_payloads(value, worker=worker, transport=transport, resolved={}) + async def _runtime_external_payload_transport( self, ) -> _RuntimeExternalPayloadTransport | None: @@ -6160,6 +6198,48 @@ async def fail_activity_task( "POST", f"/worker/activity-tasks/{task_id}/fail", worker=True, json=body ) + async def renew_activity_task_lease( + self, + *, + task_id: str, + activity_attempt_id: str, + lease_owner: str, + timeout: float = 5.0, + ) -> Any: + """Renew one live activity claim without recording application progress. + + Requires affirmative ``activity_lease_renewal`` discovery. The caller + must validate the returned original identity and grant before using it. + This neither revives an expired claim nor extends business timeouts. + Discovery, retries and transport share the supplied bounded window. + """ + for field, value in ( + ("task_id", task_id), ("activity_attempt_id", activity_attempt_id), ("lease_owner", lease_owner), + ): + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{field} must be a non-empty string") + if type(timeout) not in (int, float) or not math.isfinite(timeout) or not 0 < timeout <= 5: + raise ValueError("activity lease renewal timeout must be greater than zero and at most five seconds") + + async def renew() -> Any: + operation = "Client.renew_activity_task_lease" + path = "worker_protocol.server_capabilities.activity_lease_renewal" + info = await self._runtime_discovery(operation=operation, required_path=path) + protocol = info.get("worker_protocol") + capabilities = protocol.get("server_capabilities") if isinstance(protocol, dict) else None + supported = capabilities.get("activity_lease_renewal") if isinstance(capabilities, dict) else None + if supported is False: + raise RuntimeCapabilityUnsupported(operation, path, f"This runtime does not support {path}.") + if supported is not True: + raise RuntimeDiscoveryUnavailable(operation, path, f"Runtime discovery must advertise {path}=true.") + + return await self._request( + "POST", f"/worker/activity-tasks/{quote(task_id, safe='')}/renew-lease", worker=True, + json={"activity_attempt_id": activity_attempt_id, "lease_owner": lease_owner}, timeout=timeout, + ) + + return await asyncio.wait_for(renew(), timeout=timeout) + async def activity_task_status( self, *, diff --git a/src/durable_workflow/errors.py b/src/durable_workflow/errors.py index 7b45aaf..e18f958 100644 --- a/src/durable_workflow/errors.py +++ b/src/durable_workflow/errors.py @@ -36,6 +36,15 @@ def __init__(self, query_task_id: str, reason: str) -> None: self.reason = reason +class ActivityLeaseAuthorityLost(DurableWorkflowError): + """An activity delivery lost its original live execution authority.""" + + def __init__(self, task_id: str, reason: str) -> None: + super().__init__(f"activity task {task_id!r} cannot continue: {reason}") + self.task_id = task_id + self.reason = reason + + class InvalidWorkflowDefinition(DurableWorkflowError, ValueError): """Workflow authoring uses an unsupported entry point or result.""" diff --git a/src/durable_workflow/worker.py b/src/durable_workflow/worker.py index 05a9d13..9d81217 100644 --- a/src/durable_workflow/worker.py +++ b/src/durable_workflow/worker.py @@ -37,6 +37,7 @@ from typing import Annotated, Any, Concatenate, Literal, ParamSpec, TypeVar, Union, get_args, get_origin, get_type_hints from . import serializer +from ._activity_lease import activity_lease_renewal_supported, current_activity_authority from ._activity_process import CallbackFailure, CallbackInvocation, CallbackProcessLost, SupervisedCallback from ._cancellation_scope_history import CancellationScopeBudget, CancellationScopeDeliveryReceipt from ._cooperative_cancellation import CancellationRequest, read_cancellation_history @@ -67,6 +68,7 @@ from .errors import ( ActivityCancelled, ActivityFailed, + ActivityLeaseAuthorityLost, AvroNotInstalledError, DuplicateRegistrationError, DurableWorkflowError, @@ -3053,76 +3055,36 @@ async def observe_ownership() -> None: # stop receipt, result or failure. Owner shutdown closes its channel. async def _run_activity_task(self, task: dict[str, Any]) -> str: + authority = current_activity_authority.get() + if authority is None and not activity_lease_renewal_supported(self.client): + return await self._run_activity_task_body(task) + if authority is None or not authority.matches( + self.client, self.client.namespace, task, self.worker_id, self.task_queue, + ): + log.warning("activity task %s has no original lease authority", task.get("task_id")) + return "claim_aborted" + try: + return await authority.run(lambda: self._run_activity_task_body(task)) + except ActivityLeaseAuthorityLost as error: + log.warning("activity claim abandoned: %s", error) + return "claim_aborted" + + async def _run_activity_task_body(self, task: dict[str, Any]) -> str: self._track_worker_session_from_task(task) task_id: str = task["task_id"] attempt_id: str = task.get("activity_attempt_id") or task.get("attempt_id", "") activity_type: str = task.get("activity_type", "") attempt_number: int = task.get("attempt_number", 1) - raw_args = task.get("arguments") - inbound_codec = task.get("payload_codec") - try: - inbound_codec = _validate_payload_codec(inbound_codec) - args = serializer.decode_envelope( - raw_args, - codec=inbound_codec, - external_storage=self.external_storage, - external_storage_cache=self.external_storage_cache, - ) or [] - except AvroNotInstalledError as e: - log.exception("activity %s arguments Avro decode failed (avro dependency unavailable)", task_id) - try: - await self.client.fail_activity_task( - task_id=task_id, - activity_attempt_id=attempt_id, - lease_owner=self.worker_id, - message=( - f"cannot decode activity arguments with codec 'avro': {e}. " - f"Reinstall durable-workflow with its runtime dependencies." - ), - failure_type=type(e).__name__, - failure_class=_exception_class_name(e), - failure_code=_exception_code(e), - stack_trace=traceback.format_exc(), - non_retryable=True, - ) - except Exception as fe: - log.warning("failed to report Avro-missing activity decode failure: %s", fe) - return "decode_error" - except (ValueError, TypeError) as e: - log.exception("activity %s arguments decode failed (codec=%r)", task_id, inbound_codec) - try: - await self.client.fail_activity_task( - task_id=task_id, - activity_attempt_id=attempt_id, - lease_owner=self.worker_id, - message=( - f"cannot decode activity arguments with codec {inbound_codec!r}: {e}. " - f"Verify the argument bytes match the declared codec and writer schema." - ), - failure_type=type(e).__name__, - failure_class=_exception_class_name(e), - failure_code=_exception_code(e), - stack_trace=traceback.format_exc(), - non_retryable=True, - ) - except Exception as fe: - log.warning("failed to report activity decode failure: %s", fe) + decoded = await self._decode_activity_arguments(task) + if decoded is None: return "decode_error" - if not isinstance(args, list): - args = [args] - result_codec = inbound_codec + result_codec, args = decoded fn = self.activities.get(activity_type) if fn is None: - try: - await self.client.fail_activity_task( - task_id=task_id, - activity_attempt_id=attempt_id, - lease_owner=self.worker_id, - message=f"no activity registered for {activity_type!r}", - ) - except Exception as e: - log.warning("failed to report unknown activity type: %s", e) + await self._report_activity_task_failure( + task_id, attempt_id, "unknown activity type", message=f"no activity registered for {activity_type!r}", + ) return "no_handler" act_ctx = ActivityContext( @@ -3138,12 +3100,7 @@ async def _run_activity_task(self, task: dict[str, Any]) -> str: ) _set_context(act_ctx) try: - if self._cooperative_cancellation_supported: - result = await self._execute_cooperative_remote_callable(task, fn, tuple(args), act_ctx.info) - else: - result = await self._execute_activity_callable( - task, activity_type, tuple(args), fn, run_sync_in_thread=True, remote=True, - ) + result = await self._invoke_remote_activity(task, activity_type, fn, args, act_ctx.info) except _RemoteActivityExecutionAborted as error: log.warning("remote activity %s claim abandoned: %s", task_id, error) return "claim_aborted" @@ -3152,67 +3109,39 @@ async def _run_activity_task(self, task: dict[str, Any]) -> str: return "claim_aborted" except _RemoteActivityApplicationFailure as error: failure = error.failure - try: - await self.client.fail_activity_task( - task_id=task_id, activity_attempt_id=attempt_id, lease_owner=self.worker_id, - message="activity cancelled" if failure.cancelled else failure.message, - failure_type=failure.failure_type, failure_class=failure.failure_class, - failure_code=failure.failure_code, stack_trace=failure.stack_trace, - non_retryable=failure.non_retryable or failure.cancelled, - ) - except Exception as report_error: - log.warning("failed to report remote activity failure: %s", report_error) - return ("cancelled" if failure.cancelled else - "failed_non_retryable" if failure.non_retryable else "failed") + await self._report_activity_task_failure( + task_id, attempt_id, "remote activity failure", + message="activity cancelled" if failure.cancelled else failure.message, + failure_type=failure.failure_type, failure_class=failure.failure_class, + failure_code=failure.failure_code, stack_trace=failure.stack_trace, + non_retryable=failure.non_retryable or failure.cancelled, + ) + return self._remote_activity_failure_status(failure) except ActivityCancelled: log.info("activity %s cancelled via heartbeat", task_id) - try: - await self.client.fail_activity_task( - task_id=task_id, - activity_attempt_id=attempt_id, - lease_owner=self.worker_id, - message="activity cancelled", - failure_type="ActivityCancelled", - failure_class=_exception_class_name(ActivityCancelled), - non_retryable=True, - ) - except Exception as fe: - log.warning("failed to report activity cancellation: %s", fe) + await self._report_activity_task_failure( + task_id, attempt_id, "activity cancellation", message="activity cancelled", + failure_type="ActivityCancelled", failure_class=_exception_class_name(ActivityCancelled), + non_retryable=True, + ) return "cancelled" except NonRetryableError as e: log.exception("activity failed (non-retryable)") - try: - await self.client.fail_activity_task( - task_id=task_id, - activity_attempt_id=attempt_id, - lease_owner=self.worker_id, - message=str(e), - failure_type=type(e).__name__, - failure_class=_exception_class_name(e), - failure_code=_exception_code(e), - stack_trace=traceback.format_exc(), - non_retryable=True, - ) - except Exception as fe: - log.warning("failed to report activity failure: %s", fe) + await self._report_activity_task_failure( + task_id, attempt_id, "activity failure", message=str(e), failure_type=type(e).__name__, + failure_class=_exception_class_name(e), failure_code=_exception_code(e), + stack_trace=traceback.format_exc(), non_retryable=True, + ) return "failed_non_retryable" except Exception as e: if _is_storage_admission_error(e): raise log.exception("activity failed") - try: - await self.client.fail_activity_task( - task_id=task_id, - activity_attempt_id=attempt_id, - lease_owner=self.worker_id, - message=str(e), - failure_type=type(e).__name__, - failure_class=_exception_class_name(e), - failure_code=_exception_code(e), - stack_trace=traceback.format_exc(), - ) - except Exception as fe: - log.warning("failed to report activity failure: %s", fe) + await self._report_activity_task_failure( + task_id, attempt_id, "activity failure", message=str(e), failure_type=type(e).__name__, + failure_class=_exception_class_name(e), failure_code=_exception_code(e), + stack_trace=traceback.format_exc(), + ) return "failed" finally: _set_context(None) @@ -3233,6 +3162,63 @@ async def _run_activity_task(self, task: dict[str, Any]) -> str: return "complete_error" return "completed" + async def _decode_activity_arguments(self, task: dict[str, Any]) -> tuple[str, list[Any]] | None: + task_id: str = task["task_id"] + attempt_id: str = task.get("activity_attempt_id") or task.get("attempt_id", "") + codec = task.get("payload_codec") + try: + codec = _validate_payload_codec(codec) + args = serializer.decode_envelope( + task.get("arguments"), codec=codec, external_storage=self.external_storage, + external_storage_cache=self.external_storage_cache, + ) or [] + except AvroNotInstalledError as error: + log.exception("activity %s arguments Avro decode failed (avro dependency unavailable)", task_id) + await self._report_activity_task_failure( + task_id, attempt_id, "Avro-missing activity decode failure", + message=(f"cannot decode activity arguments with codec 'avro': {error}. " + "Reinstall durable-workflow with its runtime dependencies."), + failure_type=type(error).__name__, failure_class=_exception_class_name(error), + failure_code=_exception_code(error), stack_trace=traceback.format_exc(), non_retryable=True, + ) + return None + except (ValueError, TypeError) as error: + log.exception("activity %s arguments decode failed (codec=%r)", task_id, codec) + await self._report_activity_task_failure( + task_id, attempt_id, "activity decode failure", + message=(f"cannot decode activity arguments with codec {codec!r}: {error}. " + "Verify the argument bytes match the declared codec and writer schema."), + failure_type=type(error).__name__, failure_class=_exception_class_name(error), + failure_code=_exception_code(error), stack_trace=traceback.format_exc(), non_retryable=True, + ) + return None + return codec, args if isinstance(args, list) else [args] + + async def _report_activity_task_failure(self, task_id: str, attempt_id: str, + context: str, **failure: Any) -> None: + try: + await self.client.fail_activity_task( + task_id=task_id, activity_attempt_id=attempt_id, lease_owner=self.worker_id, **failure, + ) + except Exception as error: + log.warning("failed to report %s: %s", context, error) + + async def _invoke_remote_activity(self, task: dict[str, Any], activity_type: str, + fn: Callable[..., Any], args: list[Any], info: ActivityInfo) -> Any: + if self._cooperative_cancellation_supported: + return await self._execute_cooperative_remote_callable(task, fn, tuple(args), info) + return await self._execute_activity_callable( + task, activity_type, tuple(args), fn, run_sync_in_thread=True, remote=True, + ) + + @staticmethod + def _remote_activity_failure_status(failure: CallbackFailure) -> str: + if failure.cancelled: + return "cancelled" + if failure.non_retryable: + return "failed_non_retryable" + return "failed" + async def _execute_activity_callable( self, task: dict[str, Any], diff --git a/tests/test_activity_lease_authority.py b/tests/test_activity_lease_authority.py new file mode 100644 index 0000000..43b0416 --- /dev/null +++ b/tests/test_activity_lease_authority.py @@ -0,0 +1,176 @@ +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from typing import Any +from unittest.mock import AsyncMock + +import pytest + +from durable_workflow._activity_lease import ActivityLeaseAuthority +from durable_workflow.errors import ActivityLeaseAuthorityLost, ServerError + +ORIGINAL = datetime(2026, 10, 21, tzinfo=timezone.utc) + + +def timestamp(seconds: float) -> str: + return (ORIGINAL + timedelta(seconds=seconds)).isoformat(timespec="microseconds").replace("+00:00", "Z") + + +@pytest.fixture +def clock(monkeypatch: pytest.MonkeyPatch) -> list[float]: + value = [100.0] + monkeypatch.setattr("durable_workflow._activity_lease.time.monotonic", lambda: value[0]) + return value + + +def task(**changes: Any) -> dict[str, Any]: + return { + "task_id": "task-1", "workflow_id": "workflow-1", "run_id": "run-1", + "activity_execution_id": "execution-1", "activity_attempt_id": "attempt-1", "attempt_number": 1, + "lease_owner": "worker-1", "task_queue": "queue-1", "lease_expires_at": timestamp(5), + "deadlines": {"start_to_close": timestamp(4)}, "retry_policy": {}, **changes, + } + + +def delivery(value: dict[str, Any] | None = None, **changes: Any) -> ActivityLeaseAuthority: + return ActivityLeaseAuthority.delivery( + object(), "17", value if value is not None else task(), owner="worker-1", + task_queue="queue-1", server_time=timestamp(0), request_started=100, **changes, + ) + + +def reply(**changes: Any) -> dict[str, Any]: + return { + "task_id": "task-1", "workflow_task_id": "task-1", "workflow_instance_id": "workflow-1", + "workflow_run_id": "run-1", "activity_execution_id": "execution-1", "activity_attempt_id": "attempt-1", + "attempt_number": 1, "lease_owner": "worker-1", "lease_renewed": True, "can_continue": True, + "cancel_requested": False, "heartbeat_recorded": False, "reason": None, + "server_time": timestamp(1), "lease_expires_at": timestamp(11), **changes, + } + + +@pytest.mark.parametrize(("field", "value"), [ + ("task_id", ""), ("activity_attempt_id", None), ("attempt_number", True), ("attempt_number", 0), + ("lease_owner", "other"), ("task_queue", "other"), ("namespace", "017"), + ("lease_expires_at", "2026-02-30T00:00:10.000000Z"), + ("lease_expires_at", "2026-10-21T00:00:10Z"), ("deadlines", []), + ("deadlines", {"heartbeat": timestamp(3)}), +]) +def test_invalid_original_delivery_never_authorizes_execution(clock: list[float], field: str, value: Any) -> None: + original = task(**{field: value}) + with pytest.raises(ActivityLeaseAuthorityLost): + delivery(original) + + +@pytest.mark.parametrize(("field", "value"), [ + ("workflow_task_id", "other"), ("workflow_instance_id", "other"), ("workflow_run_id", "other"), + ("activity_execution_id", "other"), ("activity_attempt_id", "other"), ("attempt_number", True), + ("lease_owner", "other"), ("lease_renewed", "true"), ("can_continue", 1), + ("cancel_requested", True), ("heartbeat_recorded", True), ("reason", "lease_expired"), + ("server_time", "2026-02-30T00:00:00.000000Z"), ("server_time", None), + ("lease_expires_at", timestamp(0)), ("lease_expires_at", timestamp(90000)), + ("lease_expires_at", timestamp(4)), +]) +def test_wrong_or_unbounded_renewal_cannot_extend_the_grant( + clock: list[float], field: str, value: Any, +) -> None: + authority = delivery() + original = authority.lease_until + clock[0] = 101.1 + renewal = reply(**{field: value}) + with pytest.raises(ActivityLeaseAuthorityLost): + authority.acknowledge(renewal, 101) + assert authority.lease_until == original + + +def test_confirmed_renewal_preserves_every_original_business_deadline(clock: list[float]) -> None: + authority = delivery(task(deadlines={"start_to_close": timestamp(4), "run": timestamp(3), + "execution": timestamp(3.5), "schedule_to_close": timestamp(3.75)})) + clock[0] = 101.1 + authority.acknowledge(reply(), 101) + assert authority.lease_until == 111 + assert authority.business_until == 103 + assert authority.remaining() == pytest.approx(1.9) + + +def test_acknowledgement_arriving_after_expiry_cannot_revive_the_claim(clock: list[float]) -> None: + authority = delivery(task(deadlines={})) + clock[0] = 105.1 + renewal = reply(server_time=timestamp(4), lease_expires_at=timestamp(20)) + with pytest.raises(ActivityLeaseAuthorityLost, match="budget_expired"): + authority.acknowledge(renewal, 104) + assert authority.lease_until == 105 + + +def test_replaying_a_successful_reply_at_a_later_request_cannot_reset_latency(clock: list[float]) -> None: + authority = delivery(task(deadlines={})) + clock[0] = 101.1 + authority.acknowledge(reply(), 101) + authority.acknowledge(reply(), 101) + assert authority.lease_until == 111 + clock[0] = 101.2 + replayed = reply() + with pytest.raises(ActivityLeaseAuthorityLost, match="invalid_activity_renewal"): + authority.acknowledge(replayed, 101.15) + assert authority.lease_until == 111 + + +def test_only_confirmed_progress_with_the_original_interval_moves_heartbeat_budget(clock: list[float]) -> None: + authority = delivery(task(deadlines={"heartbeat": timestamp(2), "run": timestamp(3)}, + retry_policy={"heartbeat_timeout": 2})) + clock[0] = 101.1 + authority.acknowledge(reply(last_heartbeat_at=timestamp(0.8), heartbeat_deadline_at=timestamp(2.8)), 101) + assert authority.heartbeat_until == pytest.approx(102.8) + assert authority.business_until == 103 + + +@pytest.mark.parametrize(("progress", "deadline"), [(1.1, 3.1), (0.8, 3.8), (0, 2), (None, None)]) +def test_runtime_renewal_cannot_fabricate_application_progress( + clock: list[float], progress: float | None, deadline: float | None, +) -> None: + authority = delivery(task(deadlines={"heartbeat": timestamp(2)}, retry_policy={"heartbeat_timeout": 2})) + clock[0] = 101.1 + authority.acknowledge(reply(last_heartbeat_at=timestamp(progress) if progress is not None else None, + heartbeat_deadline_at=timestamp(deadline) if deadline is not None else None), 101) + assert authority.heartbeat_until == 102 + + +def test_real_heartbeat_receipt_moves_only_its_recorded_interval_and_not_the_business_bound(clock: list[float]) -> None: + authority = delivery(task(deadlines={"heartbeat": timestamp(2), "run": timestamp(4)}, + retry_policy={"heartbeat_timeout": 2})) + clock[0] = 100.9 + authority.acknowledge_progress(reply(heartbeat_recorded=True, last_heartbeat_at=timestamp(0.8), + lease_expires_at=timestamp(5.8)), 100.8) + assert authority.heartbeat_until == pytest.approx(102.8) + assert authority.business_until == 104 + assert authority.lease_until == pytest.approx(105.8) + clock[0] = 101 + authority.acknowledge(reply(server_time=timestamp(0.7), lease_expires_at=timestamp(5.7)), 100.7) + assert authority.heartbeat_until == pytest.approx(102.8) + assert authority.lease_until == pytest.approx(105.8) + + +def test_lost_execution_authority_does_not_block_readonly_status_or_callback_stop_receipts(clock: list[float]) -> None: + authority = delivery() + for action in ("complete", "fail", "heartbeat", "renew-lease"): + assert authority.governs(authority.client, "17", f"/worker/activity-tasks/task-1/{action}") + for action in ("status", "cancellation/acknowledge"): + assert not authority.governs(authority.client, "17", f"/worker/activity-tasks/task-1/{action}") + + +async def test_temporary_renewal_failure_grants_nothing_but_does_not_revoke_a_live_original() -> None: + import time + + now = datetime.now(timezone.utc) + expires = (now + timedelta(seconds=10)).isoformat(timespec="microseconds").replace("+00:00", "Z") + value = task(lease_expires_at=expires, deadlines={}) + client = AsyncMock() + client.renew_activity_task_lease.side_effect = ServerError(503, {"reason": "backend_unavailable"}) + authority = ActivityLeaseAuthority.delivery( + client, "17", value, owner="worker-1", task_queue="queue-1", + server_time=now.isoformat(timespec="microseconds").replace("+00:00", "Z"), request_started=time.monotonic(), + ) + original = authority.lease_until + assert await authority.renew() is False + assert authority.lease_until == original + assert authority.remaining() > 0 diff --git a/tests/test_activity_lease_renewal.py b/tests/test_activity_lease_renewal.py new file mode 100644 index 0000000..9ac0614 --- /dev/null +++ b/tests/test_activity_lease_renewal.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock + +import httpx +import pytest + +from durable_workflow.client import Client +from durable_workflow.errors import RuntimeCapabilityUnsupported, RuntimeDiscoveryUnavailable + + +def renewal_client(supported: object = True) -> Client: + client = Client("http://unused.invalid", worker_token="worker-secret", namespace="17") + client._cluster_info = { + "worker_protocol": {"server_capabilities": {"activity_lease_renewal": supported}}, + } + return client + + +@pytest.mark.asyncio +async def test_renewal_uses_worker_credentials_and_exact_original_identity() -> None: + client = renewal_client() + reply = {"lease_renewed": True, "heartbeat_recorded": False, "server_time": "original-clock"} + client._request = AsyncMock(return_value=reply) + try: + assert await client.renew_activity_task_lease( + task_id="task/17", activity_attempt_id="attempt0", lease_owner="worker017", + ) is reply + client._request.assert_awaited_once_with( + "POST", "/worker/activity-tasks/task%2F17/renew-lease", worker=True, + json={"activity_attempt_id": "attempt0", "lease_owner": "worker017"}, timeout=5.0, + ) + finally: + await client.aclose() + + +@pytest.mark.asyncio +async def test_real_client_transport_preserves_worker_plane_and_namespace() -> None: + client = renewal_client() + client._http.request = AsyncMock(return_value=httpx.Response( + 200, json={"lease_renewed": True, "heartbeat_recorded": False}, + request=httpx.Request("POST", "http://unused.invalid/api/worker/activity-tasks/task/renew-lease"), + )) + try: + reply = await client.renew_activity_task_lease( + task_id="task", activity_attempt_id="attempt", lease_owner="worker", + ) + assert reply == {"lease_renewed": True, "heartbeat_recorded": False} + sent = client._http.request.await_args + assert sent.args == ("POST", "/api/worker/activity-tasks/task/renew-lease") + assert sent.kwargs["headers"]["Authorization"] == "Bearer worker-secret" + assert sent.kwargs["headers"]["X-Namespace"] == "17" + assert sent.kwargs["json"] == {"activity_attempt_id": "attempt", "lease_owner": "worker"} + finally: + await client.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("supported", [False, None, "true", 1]) +async def test_renewal_requires_affirmative_capability_before_a_worker_request(supported: object) -> None: + client = renewal_client(supported) + client._request = AsyncMock() + try: + expected = RuntimeCapabilityUnsupported if supported is False else RuntimeDiscoveryUnavailable + with pytest.raises(expected, match="activity_lease_renewal"): + await client.renew_activity_task_lease( + task_id="task", activity_attempt_id="attempt", lease_owner="worker", + ) + client._request.assert_not_awaited() + finally: + await client.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ["task_id", "activity_attempt_id", "lease_owner"]) +async def test_empty_claim_identity_is_refused_before_discovery(field: str) -> None: + client = renewal_client() + client._runtime_discovery = AsyncMock() + identity = {"task_id": "task", "activity_attempt_id": "attempt", "lease_owner": "worker", field: ""} + try: + with pytest.raises(ValueError, match=field): + await client.renew_activity_task_lease(**identity) + client._runtime_discovery.assert_not_awaited() + finally: + await client.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("timeout", [0, -1, 6, True, float("inf"), float("nan")]) +async def test_invalid_window_cannot_disable_or_expand_the_bound(timeout: float) -> None: + client = renewal_client() + client._runtime_discovery = AsyncMock() + try: + with pytest.raises(ValueError, match="timeout"): + await client.renew_activity_task_lease( + task_id="task", activity_attempt_id="attempt", lease_owner="worker", timeout=timeout, + ) + client._runtime_discovery.assert_not_awaited() + finally: + await client.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["discovery", "renewal"]) +async def test_discovery_and_transport_share_one_bounded_request_window(phase: str) -> None: + client = renewal_client() + cancelled = asyncio.Event() + + async def held(*args: object, **kwargs: object) -> None: + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + if phase == "discovery": + client._runtime_discovery = AsyncMock(side_effect=held) + client._request = AsyncMock() + else: + client._request = AsyncMock(side_effect=held) + try: + renewal = client.renew_activity_task_lease( + task_id="task", activity_attempt_id="attempt", lease_owner="worker", timeout=0.02, + ) + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(renewal, timeout=0.5) + assert cancelled.is_set() + if phase == "discovery": + client._request.assert_not_awaited() + finally: + await client.aclose() diff --git a/tests/test_activity_lease_worker.py b/tests/test_activity_lease_worker.py new file mode 100644 index 0000000..b021bd9 --- /dev/null +++ b/tests/test_activity_lease_worker.py @@ -0,0 +1,266 @@ +from __future__ import annotations + +import asyncio +import threading +import time +from datetime import datetime, timedelta, timezone +from typing import Any +from unittest.mock import AsyncMock + +import httpx +import pytest + +from durable_workflow import Client, Worker, activity, serializer +from durable_workflow._activity_lease import current_activity_authority + + +@pytest.mark.parametrize("capability", [True, False, None, "true"]) +async def test_worker_uses_client_discovery_for_original_authority(capability: Any) -> None: + client = Client("http://runtime.invalid") + client._cluster_info = {"worker_protocol": {"server_capabilities": {"activity_lease_renewal": capability}}} + worker = Worker(client, task_queue="queue-1", worker_id="worker-1") + worker._run_activity_task_body = AsyncMock(return_value="completed") + context = current_activity_authority.set(None) + try: + outcome = await worker._run_activity_task({"task_id": "original-task"}) + if capability is True: + assert outcome == "claim_aborted" + worker._run_activity_task_body.assert_not_awaited() + else: + assert outcome == "completed" + worker._run_activity_task_body.assert_awaited_once() + finally: + current_activity_authority.reset(context) + await client.aclose() + + +def timestamp(value: datetime) -> str: + return value.isoformat(timespec="microseconds").replace("+00:00", "Z") + + +class LeaseServer: + def __init__(self, *, lease: float = 0.18, deadline: float | None = None) -> None: + self.lease = lease + self.started = datetime.now(timezone.utc) + self.monotonic_started = time.monotonic() + self.task = { + "task_id": "task/17", "activity_attempt_id": "attempt-017", "attempt_number": 1, + "activity_execution_id": "execution-1", "workflow_id": "workflow-1", "run_id": "run-1", + "task_queue": "queue-1", "activity_type": "maintained", "lease_owner": "worker-1", + "lease_expires_at": timestamp(self.started + timedelta(seconds=lease)), + "retry_policy": {"heartbeat_timeout": None}, "deadlines": {}, + "arguments": serializer.envelope([], codec="avro"), "payload_codec": "avro", + } + if deadline is not None: + self.task["deadlines"] = {"start_to_close": timestamp(self.started + timedelta(seconds=deadline))} + self.paths: list[str] = [] + self.renewals = 0 + self.deny = False + self.delay_poll = 0.0 + self.delay_before_claim = 0.0 + self.delay_renewal = 0.0 + self.fail_completion = False + self.publication_times: list[float] = [] + + async def respond(self, request: httpx.Request) -> httpx.Response: + self.paths.append(request.url.path) + if request.url.path.endswith("/poll"): + if self.delay_before_claim: + await asyncio.sleep(self.delay_before_claim) + claimed_at = datetime.now(timezone.utc) + # Start-to-close and heartbeat begin when this synthetic server + # claims the attempt, after client bootstrap and any empty wait. + offset = claimed_at - self.started + for key in ("start_to_close", "heartbeat"): + if key in self.task["deadlines"]: + deadline = datetime.fromisoformat(self.task["deadlines"][key].replace("Z", "+00:00")) + self.task["deadlines"][key] = timestamp(deadline + offset) + self.started = claimed_at + self.monotonic_started = time.monotonic() + self.task["lease_expires_at"] = timestamp(claimed_at + timedelta(seconds=self.lease)) + response_clock = datetime.now(timezone.utc) + await asyncio.sleep(self.delay_poll) + return httpx.Response(200, json={ + "task": self.task, "server_time": timestamp(response_clock), + "server_capabilities": {"activity_lease_renewal": True}, + }) + if request.url.path.endswith("/renew-lease"): + self.renewals += 1 + await asyncio.sleep(self.delay_renewal if self.renewals > 1 else 0) + now = datetime.now(timezone.utc) + denied = (self.deny and self.renewals > 1 + or now >= datetime.fromisoformat(self.task["lease_expires_at"].replace("Z", "+00:00"))) + if not denied: + self.task["lease_expires_at"] = timestamp(now + timedelta(seconds=self.lease)) + return httpx.Response(409 if denied else 200, json={ + "task_id": self.task["task_id"], "workflow_task_id": self.task["task_id"], + "workflow_instance_id": self.task["workflow_id"], "workflow_run_id": self.task["run_id"], + "activity_execution_id": self.task["activity_execution_id"], + "activity_attempt_id": self.task["activity_attempt_id"], "attempt_number": 1, + "lease_owner": self.task["lease_owner"], "lease_renewed": not denied, + "can_continue": not denied, "cancel_requested": False, "heartbeat_recorded": False, + "server_time": timestamp(now), "lease_expires_at": self.task["lease_expires_at"], + "reason": "lease_expired" if denied else None, + }) + if request.url.path.endswith("/complete"): + self.publication_times.append(time.monotonic() - self.monotonic_started) + if self.fail_completion: + return httpx.Response(503, json={"reason": "temporary_failure"}) + return httpx.Response(200, json={"outcome": "completed", "recorded": True}) + if request.url.path.endswith("/heartbeat"): + now = datetime.now(timezone.utc) + return httpx.Response(200, json={ + "task_id": self.task["task_id"], "activity_attempt_id": self.task["activity_attempt_id"], + "lease_owner": self.task["lease_owner"], "heartbeat_recorded": True, + "can_continue": True, "cancel_requested": False, "reason": None, + "last_heartbeat_at": timestamp(now), + "lease_expires_at": timestamp(now + timedelta(seconds=self.lease)), + }) + if request.url.path.endswith("/fail"): + return httpx.Response(200, json={"outcome": "failed", "recorded": True}) + raise AssertionError(f"Unexpected request {request.method} {request.url.path}") + + def publications(self) -> list[str]: + return [path for path in self.paths if path.endswith(("/complete", "/fail"))] + + +async def execute(server: LeaseServer, handler: Any) -> str: + async with Client("http://lease.invalid", namespace="17", worker_token="synthetic-worker-token") as client: + await client._http.aclose() + client._http = httpx.AsyncClient(base_url="http://lease.invalid", transport=httpx.MockTransport(server.respond)) + client._cluster_info = {"worker_protocol": {"server_capabilities": {"activity_lease_renewal": True}}} + worker = Worker(client, worker_id="worker-1", task_queue="queue-1", workflows=[], activities=[handler]) + task = await client.poll_activity_task(worker_id="worker-1", task_queue="queue-1") + return await worker._run_activity_task(task) + + +async def test_healthy_callback_runs_beyond_original_lease_without_application_heartbeats() -> None: + server = LeaseServer() + + @activity.defn(name="maintained") + async def handler() -> str: + await asyncio.sleep(0.55) + return "completed once" + + assert await execute(server, handler) == "completed" + assert server.renewals >= 2 + assert len(server.publications()) == 1 + assert not any(path.endswith("/heartbeat") for path in server.paths) + + +@pytest.mark.parametrize("loss", ["denied", "slow-renewal", "business-deadline"]) +async def test_lost_claim_cannot_publish_a_late_result_or_failure(loss: str) -> None: + server = LeaseServer(deadline=0.13 if loss == "business-deadline" else None) + server.deny = loss == "denied" + server.delay_renewal = 0.3 if loss == "slow-renewal" else 0.0 + called = False + cancelled = False + + @activity.defn(name="maintained") + async def handler() -> str: + nonlocal called, cancelled + called = True + try: + await asyncio.sleep(0.45) + except asyncio.CancelledError: + # A callback can swallow cancellation. Publication still needs its + # original live authority, independently of callback cooperation. + cancelled = True + await asyncio.sleep(0.02) + return "late result" + + assert await execute(server, handler) == "claim_aborted" + assert called is True + assert cancelled is True + assert server.publications() == [] + + +async def test_original_deadline_fences_transport_retries_during_terminal_publication() -> None: + server = LeaseServer(lease=1, deadline=0.12) + server.fail_completion = True + + @activity.defn(name="maintained") + async def handler() -> str: + return "ready before the deadline" + + assert await execute(server, handler) == "claim_aborted" + assert server.publications() + assert all(elapsed < 0.12 for elapsed in server.publication_times) + + +async def test_cancelled_thread_keeps_its_real_slot_until_the_callback_exits() -> None: + server = LeaseServer() + server.deny = True + started, release = threading.Event(), threading.Event() + + @activity.defn(name="maintained") + def handler() -> str: + started.set() + release.wait(timeout=2) + return "late thread result" + + async with Client("http://lease.invalid", namespace="17", worker_token="synthetic-worker-token") as client: + await client._http.aclose() + client._http = httpx.AsyncClient(base_url="http://lease.invalid", transport=httpx.MockTransport(server.respond)) + client._cluster_info = {"worker_protocol": {"server_capabilities": {"activity_lease_renewal": True}}} + worker = Worker(client, worker_id="worker-1", task_queue="queue-1", workflows=[], activities=[handler], + max_concurrent_activity_tasks=1) + try: + task = await client.poll_activity_task(worker_id="worker-1", task_queue="queue-1") + assert await worker._run_activity_task(task) == "claim_aborted" + assert started.is_set() + assert len(worker._remote_activity_threads) == 1 + assert worker._remote_activity_thread_slots.locked() + assert server.publications() == [] + finally: + release.set() + if worker._remote_activity_executor is not None: + worker._remote_activity_executor.shutdown(wait=True) + await asyncio.sleep(0) + assert worker._remote_activity_threads == set() + + +async def test_actual_application_heartbeat_acknowledges_recorded_progress_before_original_expiry() -> None: + server = LeaseServer(lease=3) + server.task["retry_policy"] = {"heartbeat_timeout": 1} + server.task["deadlines"] = {"heartbeat": timestamp(server.started + timedelta(seconds=0.15))} + + @activity.defn(name="maintained") + async def handler() -> str: + await asyncio.sleep(0.04) + await activity.context().heartbeat({"step": 1}) + await asyncio.sleep(0.28) + return "progress acknowledged" + + assert await execute(server, handler) == "completed" + assert len(server.publications()) == 1 + assert len([path for path in server.paths if path.endswith("/heartbeat")]) == 1 + + +async def test_poll_round_trip_consumes_the_original_delivery_budget() -> None: + server = LeaseServer(lease=0.06) + server.delay_poll = 0.1 + called = False + + @activity.defn(name="maintained") + async def handler() -> str: + nonlocal called + called = True + return "must not run" + + with pytest.raises(Exception, match="activity.*(expired|budget)"): + await execute(server, handler) + assert called is False + assert server.publications() == [] + + +async def test_waiting_for_a_new_task_does_not_consume_its_not_yet_created_lease() -> None: + server = LeaseServer() + server.delay_before_claim = 0.25 + + @activity.defn(name="maintained") + async def handler() -> str: + return "newly claimed after long poll" + + assert await execute(server, handler) == "completed" + assert len(server.publications()) == 1