Stop activity lease maintenance after settlement when cancellation is consumed

.github/workflows/ci.yml

@@ -135,7 +135,27 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - run: pip install -e '.[dev]'
-      - run: pytest tests/ -m "not integration" -q
+      - name: Run unit tests with stalled-thread diagnostics
+        id: unit_tests
+        shell: bash
+        run: |
+          set -o pipefail
+          # Leave room within the 15-minute job for bounded failure diagnostics.
+          timeout --signal=INT --kill-after=10s 13m \
+            pytest tests/ -m "not integration" -vv \
+            -o faulthandler_timeout=120 --durations=20 \
+            2>&1 | tee "$RUNNER_TEMP/sdk-unit-tests.log"
+      - name: Attach bounded failure diagnostics to the check
+        if: ${{ always() && steps.unit_tests.outcome == 'failure' }}
+        shell: bash
+        run: |
+          if [[ -f "$RUNNER_TEMP/sdk-unit-tests.log" ]]; then
+            phase_unit_tail="$(tail -c 12000 "$RUNNER_TEMP/sdk-unit-tests.log")"
+            phase_unit_tail="${phase_unit_tail//'%'/'%25'}"
+            phase_unit_tail="${phase_unit_tail//$'\r'/'%0D'}"
+            phase_unit_tail="${phase_unit_tail//$'\n'/'%0A'}"
+            printf '::error title=Unit-test failure diagnostics::%s\n' "$phase_unit_tail"
+          fi
 
   package:
     runs-on: ubuntu-latest

src/durable_workflow/_activity_lease.py

@@ -278,19 +278,22 @@ async def renew(self) -> bool:
             self.lost_reason = "activity_renewal_unavailable"
             raise ActivityLeaseAuthorityLost(self.task_id, self.lost_reason) from error
 
-    async def _maintain(self) -> None:
+    async def _maintain(self, stopped: asyncio.Event) -> None:
         confirmed = True
-        while True:
+        while not stopped.is_set():
             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()))
+            if stopped.is_set():
+                return
             confirmed = await self.renew()
 
     async def run(self, operation: Callable[[], Awaitable[_T]]) -> _T:
         self.remaining()
+        stopped = asyncio.Event()
         execution = asyncio.ensure_future(operation())
-        maintenance = asyncio.create_task(self._maintain())
+        maintenance = asyncio.create_task(self._maintain(stopped))
         try:
             while True:
                 done, _ = await asyncio.wait({execution, maintenance}, timeout=self.remaining(),
@@ -306,6 +309,9 @@ async def run(self, operation: Callable[[], Awaitable[_T]]) -> _T:
             self.lost_reason = self.lost_reason or "activity_execution_abandoned"
             raise
         finally:
+            # An I/O wrapper may complete concurrently with cancel() and consume
+            # that cancellation. The settled scope must still stop renewing.
+            stopped.set()
             # 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.

tests/test_activity_lease_worker.py

@@ -11,7 +11,7 @@
 import pytest
 
 from durable_workflow import Client, Worker, activity, serializer
-from durable_workflow._activity_lease import current_activity_authority
+from durable_workflow._activity_lease import ActivityLeaseAuthority, current_activity_authority
 
 
 @pytest.mark.parametrize("capability", [True, False, None, "true"])
@@ -148,6 +148,63 @@ async def handler() -> str:
     assert not any(path.endswith("/heartbeat") for path in server.paths)
 
 
+async def test_settled_callback_stops_maintenance_when_inflight_renewal_swallows_cancellation(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    server = LeaseServer(lease=1)
+    authority = ActivityLeaseAuthority.delivery(object(), "17", server.task, owner="worker-1",
+        task_queue="queue-1", server_time=timestamp(server.started), request_started=server.monotonic_started)
+    renewal_started = asyncio.Event()
+    calls = 0
+    swallowed = False
+
+    async def renewal(_self: ActivityLeaseAuthority) -> bool:
+        nonlocal calls, swallowed
+        calls += 1
+        renewal_started.set()
+        try:
+            await asyncio.get_running_loop().create_future()
+        except asyncio.CancelledError:
+            if not swallowed:
+                # Model an I/O wrapper completing concurrently with cancel().
+                swallowed = True
+                return True
+            raise
+        raise AssertionError("The controlled renewal should remain pending until cancellation")
+
+    async def operation() -> str:
+        await renewal_started.wait()
+        return "completed once"
+
+    monkeypatch.setattr(ActivityLeaseAuthority, "renew", renewal)
+    running = asyncio.create_task(authority.run(operation))
+    try:
+        done, _ = await asyncio.wait({running}, timeout=0.8)
+        assert running in done, "Settled callbacks must not await an indefinitely renewing helper"
+        assert await running == "completed once"
+        assert swallowed is True
+        assert calls == 1, "A settled operation must not start a subsequent renewal"
+    finally:
+        running.cancel()
+        await asyncio.gather(running, return_exceptions=True)
+
+
+async def test_loading_supervision_can_settle_before_the_same_claim_runs_its_callback() -> None:
+    server = LeaseServer(lease=1)
+    authority = ActivityLeaseAuthority.delivery(object(), "17", server.task, owner="worker-1",
+        task_queue="queue-1", server_time=timestamp(server.started), request_started=server.monotonic_started)
+
+    async def loading() -> str:
+        await asyncio.sleep(0)
+        return "loaded"
+
+    assert await authority.run(loading) == "loaded"
+    original_deadline = authority.lease_until
+    assert await authority.run(loading) == "loaded"
+    assert authority.lease_until == original_deadline
+    assert authority.remaining() > 0
+
+
 @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)