Skip to content

Commit d6da755

Browse files
authored
Stop activity publication after heartbeat authority refusal (#187)
Stop default activities when a heartbeat refuses continued execution
2 parents fe40c64 + 0c215c9 commit d6da755

7 files changed

Lines changed: 237 additions & 14 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
77
## [Unreleased]
88

99
### Fixed
10+
11+
- Stop default async and synchronous activities when a heartbeat explicitly
12+
refuses continued execution. Preserve the Server reason separately from
13+
cancellation and fence completion/failure even when user code catches the
14+
refusal. The worker frees its slot after the callback exits.
15+
1016
- Reject canonical version markers missing their change ID or supported-range
1117
bounds. Consistent historical duplicate ranges may still differ when each
1218
contains the frozen version.

‎README.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,17 @@ independently of application heartbeats. Existing terminal cancellation remains
142142
available. See the [cancellation guide](docs/cooperative-cancellation-design.md)
143143
for immutable context, operation policies, shielded cleanup and recovery.
144144

145+
Both `await activity.context().heartbeat()` and synchronous
146+
`activity.context().heartbeat_sync()` stop at an explicit Server refusal.
147+
They raise `ActivityLeaseAuthorityLost` with its original `.reason` for an
148+
elapsed deadline, closed attempt or closed run. Cancellation continues to raise
149+
`ActivityCancelled`. The worker submits no result or failure for an attempt
150+
whose heartbeat refused continuation, even if application code catches that
151+
exception. A heartbeat check acts only when called. Code that never heartbeats
152+
can continue until it returns, subject to independent lease/cancellation
153+
supervision; a running native thread cannot be forcibly stopped or its effects
154+
undone by the SDK.
155+
145156
## Development
146157

147158
The installed `durable-workflow-replay-conformance` command runs an offline

‎docs/sdk-reference.md‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -966,6 +966,23 @@ compatibility manifest. A client configured with only the opposite role's token
966966
still fails before transport for actual worker or control operations; use
967967
`token` when one shared credential intentionally authorizes both planes.
968968

969+
## Activity heartbeat refusals
970+
971+
Use `await activity.context().heartbeat(details)` in an async activity or
972+
`activity.context().heartbeat_sync(details)` in a synchronous activity thread.
973+
Both raise `ActivityCancelled` when cancellation is requested. An explicit
974+
`can_continue=false` instead raises `ActivityLeaseAuthorityLost` with the
975+
original `.reason`, such as `activity_deadline_elapsed`, `attempt_closed` or
976+
`run_closed`. This is lost execution authority, not a new activity failure.
977+
The worker releases the slot after the callback exits and submits neither a
978+
result nor a failure for that refused attempt.
979+
980+
Checks happen at the next heartbeat. Code that never heartbeats may continue
981+
until it returns, subject to separately negotiated lease/cancellation
982+
supervision. The SDK cannot forcibly stop a running native thread or undo its
983+
external effects. Keep side effects idempotent and heartbeat at useful progress
984+
boundaries.
985+
969986
## Metrics
970987

971988
Pass a recorder to `Client(metrics=...)` or `Worker(metrics=...)` to collect request, poll, and task metrics. The SDK ships a no-op default, an `InMemoryMetrics` recorder for tests or custom exporter loops, and `PrometheusMetrics` for deployments that install the optional extra:

‎src/durable_workflow/activity.py‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
from dataclasses import dataclass
88
from typing import TYPE_CHECKING, Any
99

10-
from .errors import ActivityCancelled
10+
from ._activity_lease import current_activity_authority
11+
from .errors import ActivityCancelled, ActivityLeaseAuthorityLost
1112

1213
if TYPE_CHECKING:
1314
from .client import Client
@@ -45,6 +46,7 @@ def __init__(
4546
self._client = client
4647
self._heartbeat_callback = heartbeat_callback
4748
self._cancel_requested = False
49+
self._cannot_continue_reason: str | None = None
4850
try:
4951
self._event_loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop()
5052
except RuntimeError:
@@ -61,7 +63,7 @@ def is_cancelled(self) -> bool:
6163
return self._cancel_requested
6264

6365
def heartbeat_sync(self, details: dict[str, Any] | None = None) -> None:
64-
"""Heartbeat from a synchronous activity thread and observe cancellation.
66+
"""Heartbeat from a synchronous activity thread and check attempt authority.
6567
6668
The request runs on the worker's event loop using the original attempt
6769
and client. Async handlers must use ``await heartbeat()`` instead.
@@ -84,7 +86,7 @@ def heartbeat_sync(self, details: dict[str, Any] | None = None) -> None:
8486
pending.result()
8587

8688
async def heartbeat(self, details: dict[str, Any] | None = None) -> None:
87-
"""Report liveness to the server and check for a cancellation request.
89+
"""Report progress and check cancellation and continued attempt authority.
8890
8991
Long-running activities should call ``heartbeat()`` periodically so the
9092
server can distinguish a slow-but-alive attempt from a dead worker.
@@ -94,7 +96,12 @@ async def heartbeat(self, details: dict[str, Any] | None = None) -> None:
9496
Raises :class:`~durable_workflow.errors.ActivityCancelled` when the
9597
owning workflow has requested cancellation, so the activity can exit
9698
cleanly at its next natural break point.
99+
100+
Raises :class:`~durable_workflow.errors.ActivityLeaseAuthorityLost`
101+
with the original Server reason when the attempt cannot continue,
102+
including an elapsed deadline, closed attempt or closed run.
97103
"""
104+
self._raise_if_cannot_continue()
98105
if self._heartbeat_callback is not None:
99106
try:
100107
await self._heartbeat_callback(details)
@@ -111,6 +118,21 @@ async def heartbeat(self, details: dict[str, Any] | None = None) -> None:
111118
if isinstance(resp, dict) and resp.get("cancel_requested"):
112119
self._cancel_requested = True
113120
raise ActivityCancelled()
121+
if isinstance(resp, dict) and resp.get("can_continue") is False:
122+
reason = resp.get("reason")
123+
self._cannot_continue_reason = reason if isinstance(reason, str) and reason else "activity_cannot_continue"
124+
self._raise_if_cannot_continue()
125+
126+
def _raise_if_cannot_continue(self) -> None:
127+
if self._cannot_continue_reason is not None:
128+
authority = current_activity_authority.get()
129+
if (authority is not None and authority.client is self._client
130+
and authority.task_id == self._info.task_id
131+
and authority.identity["activity_attempt_id"] == self._info.activity_attempt_id
132+
and authority.identity["attempt_number"] == self._info.attempt_number
133+
and authority.identity["lease_owner"] == self._info.worker_id):
134+
authority.lost_reason = authority.lost_reason or self._cannot_continue_reason
135+
raise ActivityLeaseAuthorityLost(self._info.task_id, self._cannot_continue_reason)
114136

115137

116138
def context() -> ActivityContext:

‎src/durable_workflow/client.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6385,7 +6385,8 @@ async def heartbeat_activity_task(
63856385
Worker-plane endpoint. Most code calls
63866386
:meth:`~durable_workflow.ActivityContext.heartbeat` instead, which
63876387
additionally raises :class:`~durable_workflow.errors.ActivityCancelled`
6388-
when the server reports the activity should stop.
6388+
for cancellation or :class:`~durable_workflow.errors.ActivityLeaseAuthorityLost`
6389+
with the original reason when execution cannot continue.
63896390
"""
63906391
body: dict[str, Any] = {
63916392
"activity_attempt_id": activity_attempt_id,

‎src/durable_workflow/worker.py‎

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3100,7 +3100,15 @@ async def _run_activity_task_body(self, task: dict[str, Any]) -> str:
31003100
)
31013101
_set_context(act_ctx)
31023102
try:
3103-
result = await self._invoke_remote_activity(task, activity_type, fn, args, act_ctx.info)
3103+
try:
3104+
result = await self._invoke_remote_activity(task, activity_type, fn, args, act_ctx.info)
3105+
finally:
3106+
# A caught heartbeat refusal cannot restore outcome authority,
3107+
# including when user code subsequently raises another error.
3108+
act_ctx._raise_if_cannot_continue()
3109+
except ActivityLeaseAuthorityLost as error:
3110+
log.warning("activity %s cannot continue after heartbeat: %s", task_id, error.reason)
3111+
return "claim_aborted"
31043112
except _RemoteActivityExecutionAborted as error:
31053113
log.warning("remote activity %s claim abandoned: %s", task_id, error)
31063114
return "claim_aborted"
@@ -3134,15 +3142,7 @@ async def _run_activity_task_body(self, task: dict[str, Any]) -> str:
31343142
)
31353143
return "failed_non_retryable"
31363144
except Exception as e:
3137-
if _is_storage_admission_error(e):
3138-
raise
3139-
log.exception("activity failed")
3140-
await self._report_activity_task_failure(
3141-
task_id, attempt_id, "activity failure", message=str(e), failure_type=type(e).__name__,
3142-
failure_class=_exception_class_name(e), failure_code=_exception_code(e),
3143-
stack_trace=traceback.format_exc(),
3144-
)
3145-
return "failed"
3145+
return await self._handle_activity_exception(task_id, attempt_id, e)
31463146
finally:
31473147
_set_context(None)
31483148

@@ -3162,6 +3162,17 @@ async def _run_activity_task_body(self, task: dict[str, Any]) -> str:
31623162
return "complete_error"
31633163
return "completed"
31643164

3165+
async def _handle_activity_exception(self, task_id: str, attempt_id: str, error: Exception) -> str:
3166+
if _is_storage_admission_error(error):
3167+
raise error
3168+
log.error("activity failed", exc_info=(type(error), error, error.__traceback__))
3169+
await self._report_activity_task_failure(
3170+
task_id, attempt_id, "activity failure", message=str(error), failure_type=type(error).__name__,
3171+
failure_class=_exception_class_name(error), failure_code=_exception_code(error),
3172+
stack_trace="".join(traceback.format_exception(type(error), error, error.__traceback__)),
3173+
)
3174+
return "failed"
3175+
31653176
async def _decode_activity_arguments(self, task: dict[str, Any]) -> tuple[str, list[Any]] | None:
31663177
task_id: str = task["task_id"]
31673178
attempt_id: str = task.get("activity_attempt_id") or task.get("attempt_id", "")
Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
from __future__ import annotations
2+
3+
import asyncio
4+
import time
5+
from contextlib import suppress
6+
from datetime import datetime, timedelta, timezone
7+
from unittest.mock import AsyncMock
8+
9+
import httpx
10+
import pytest
11+
12+
from durable_workflow import activity, serializer
13+
from durable_workflow._activity_lease import ActivityLeaseAuthority, current_activity_authority
14+
from durable_workflow.activity import ActivityContext, ActivityInfo
15+
from durable_workflow.client import Client
16+
from durable_workflow.errors import ActivityLeaseAuthorityLost, NonRetryableError
17+
from durable_workflow.worker import Worker
18+
19+
REASONS = ("activity_deadline_elapsed", "attempt_closed", "run_closed")
20+
21+
22+
def refusing_client(reason: str) -> AsyncMock:
23+
client = AsyncMock(spec=Client)
24+
client.heartbeat_activity_task.return_value = {
25+
"can_continue": False, "cancel_requested": False, "reason": reason,
26+
}
27+
return client
28+
29+
30+
@pytest.mark.asyncio
31+
@pytest.mark.parametrize("reason", REASONS)
32+
@pytest.mark.parametrize("synchronous", (False, True))
33+
async def test_heartbeat_refusal_carries_original_reason(reason: str, synchronous: bool) -> None:
34+
ctx = ActivityContext(info=ActivityInfo(
35+
task_id="original-task", activity_type="heartbeat-refusal", activity_attempt_id="original-attempt",
36+
attempt_number=1, task_queue="original-queue", worker_id="original-worker",
37+
), client=refusing_client(reason))
38+
39+
invocation = (asyncio.to_thread(ctx.heartbeat_sync, {"progress": 1}) if synchronous
40+
else ctx.heartbeat({"progress": 1}))
41+
with pytest.raises(ActivityLeaseAuthorityLost) as refusal:
42+
await invocation
43+
44+
assert refusal.value.task_id == "original-task"
45+
assert refusal.value.reason == reason
46+
assert not ctx.is_cancelled
47+
48+
49+
@pytest.mark.asyncio
50+
@pytest.mark.parametrize("reason", REASONS)
51+
@pytest.mark.parametrize("synchronous", (False, True))
52+
async def test_default_worker_does_not_publish_after_heartbeat_refusal(
53+
reason: str, synchronous: bool,
54+
) -> None:
55+
client = refusing_client(reason)
56+
continued: list[str] = []
57+
58+
def sync_handler() -> str:
59+
activity.context().heartbeat_sync()
60+
continued.append("after refusal")
61+
return "late result"
62+
63+
async def async_handler() -> str:
64+
await activity.context().heartbeat()
65+
continued.append("after refusal")
66+
return "late result"
67+
68+
handler = activity.defn(name="heartbeat-refusal")(sync_handler if synchronous else async_handler)
69+
worker = Worker(client, task_queue="original-queue", workflows=[], activities=[handler])
70+
task = {
71+
"task_id": "original-task", "activity_attempt_id": "original-attempt",
72+
"activity_type": "heartbeat-refusal", "attempt_number": 1,
73+
"arguments": serializer.encode([], codec="avro"), "payload_codec": "avro",
74+
}
75+
76+
assert await worker._run_activity_task(task) == "claim_aborted"
77+
assert continued == []
78+
client.complete_activity_task.assert_not_called()
79+
client.fail_activity_task.assert_not_called()
80+
with pytest.raises(RuntimeError, match="outside of an activity"):
81+
activity.context()
82+
83+
84+
@pytest.mark.asyncio
85+
@pytest.mark.parametrize("same_original", (True, False))
86+
async def test_refusal_fences_only_the_matching_original_client_authority(same_original: bool) -> None:
87+
paths: list[str] = []
88+
89+
async def respond(request: httpx.Request) -> httpx.Response:
90+
paths.append(request.url.path)
91+
assert request.url.path.endswith("/heartbeat")
92+
return httpx.Response(200, json={
93+
"heartbeat_recorded": False, "can_continue": False,
94+
"cancel_requested": False, "reason": "attempt_closed",
95+
})
96+
97+
async with Client("http://heartbeat.invalid", worker_token="synthetic-worker-token") as client:
98+
await client._http.aclose()
99+
client._http = httpx.AsyncClient(base_url="http://heartbeat.invalid", transport=httpx.MockTransport(respond))
100+
reference = datetime.now(timezone.utc)
101+
def timestamp(value: datetime) -> str:
102+
return value.isoformat(timespec="microseconds").replace("+00:00", "Z")
103+
104+
authority = ActivityLeaseAuthority.delivery(client, client.namespace, {
105+
"task_id": "original-task" if same_original else "different-task",
106+
"workflow_id": "workflow-1", "run_id": "run-1", "activity_execution_id": "execution-1",
107+
"activity_attempt_id": "original-attempt", "lease_owner": "original-worker",
108+
"task_queue": "original-queue", "attempt_number": 1,
109+
"lease_expires_at": timestamp(reference + timedelta(seconds=30)), "deadlines": {},
110+
}, owner="original-worker", task_queue="original-queue", server_time=timestamp(reference),
111+
request_started=time.monotonic())
112+
token = current_activity_authority.set(authority)
113+
try:
114+
ctx = ActivityContext(info=ActivityInfo(
115+
task_id="original-task", activity_type="heartbeat-refusal", activity_attempt_id="original-attempt",
116+
attempt_number=1, task_queue="original-queue", worker_id="original-worker",
117+
), client=client)
118+
with pytest.raises(ActivityLeaseAuthorityLost, match="attempt_closed"):
119+
await ctx.heartbeat()
120+
if same_original:
121+
with pytest.raises(ActivityLeaseAuthorityLost, match="attempt_closed"):
122+
await client.complete_activity_task(
123+
task_id="original-task", activity_attempt_id="original-attempt",
124+
lease_owner="original-worker", result="late result",
125+
)
126+
assert paths == ["/api/worker/activity-tasks/original-task/heartbeat"]
127+
else:
128+
assert authority.remaining() > 0
129+
finally:
130+
current_activity_authority.reset(token)
131+
132+
133+
@pytest.mark.asyncio
134+
@pytest.mark.parametrize("late_error", (None, ValueError, NonRetryableError))
135+
async def test_swallowed_heartbeat_refusal_still_fences_worker_publication(late_error: type[Exception] | None) -> None:
136+
client = refusing_client("attempt_closed")
137+
138+
@activity.defn(name="heartbeat-refusal-caught")
139+
async def handler() -> str:
140+
with suppress(Exception):
141+
await activity.context().heartbeat()
142+
if late_error is not None:
143+
raise late_error("application error after closed attempt")
144+
return "late result"
145+
146+
worker = Worker(client, task_queue="original-queue", workflows=[], activities=[handler])
147+
task = {
148+
"task_id": "original-task", "activity_attempt_id": "original-attempt",
149+
"activity_type": "heartbeat-refusal-caught", "attempt_number": 1,
150+
"arguments": serializer.encode([], codec="avro"), "payload_codec": "avro",
151+
}
152+
153+
assert await worker._run_activity_task(task) == "claim_aborted"
154+
client.complete_activity_task.assert_not_called()
155+
client.fail_activity_task.assert_not_called()

0 commit comments

Comments
 (0)