fix(concurrency): release leases granted to callers that were cancelled mid-acquire

[devin-ai-integration]

Fixes the flake in tests/concurrency/test_context.py::test_concurrency_context_releases_slots_async seen on main (https://github.com/PrefectHQ/prefect/actions/runs/30122656936): the limit was left with active_slots == 1 after the context exited.

Root cause

aacquire_concurrency_slots_with_lease awaits asyncio.wrap_future(future) on a concurrent.futures.Future that the acquisition service fulfills from the global event loop thread. When the caller's cancel scope fires while the increment request is in flight, wrap_future cancels that future. The request then succeeds server-side, but:

future.set_result(response)   # InvalidStateError: CANCELLED

so the service crashed handling the item and the lease id was never returned to anyone — the caller couldn't record it on ConcurrencyContext.cleanup_lease_ids, so the slots stayed occupied until lease expiry. Exactly this sequence is in the CI log (increment-with-lease 200 at 39.998, InvalidStateError at 40.000, read of active_slots == 1 at 40.001).

Fix

Two orderings of "caller goes away while the acquisition is in flight" are now both handled:

1. Cancelled before the result is delivered. FutureQueueService._handle transitions the future to running before delivering the result, which closes the cancel/set race atomically and lets the service see that the caller gave up:

   if future.set_running_or_notify_cancel():
       future.set_result(response)
   else:
       await self._discard(response)   # no-op by default

   ConcurrencySlotAcquisitionWithLeaseService._discard releases the orphaned lease.

2. Cancelled after the future is running, where wrap_future cancels only its wrapper and the delivered result is dropped: aacquire_concurrency_slots_with_lease catches CancelledError and attaches a done callback that releases the lease if one ends up being granted (off the cancelled caller, on a short-lived daemon thread).

Regression tests added for both paths; each fails on main and passes here.

Checklist

• This pull request references any related issue by including "closes <link to issue>"
• If this is a complex change, a maintainer has confirmed the proposed approach on the linked issue.
• If this pull request adds or changes functionality, it includes tests or explains why tests are not needed.
• If this pull request changes user-facing behavior, it updates documentation or explains why documentation is not needed.
• If this pull request removes docs files, it includes redirect settings in mint.json.
• If this pull request adds functions or classes, it includes helpful docstrings.

Link to Devin session: https://app.devin.ai/sessions/d6c085dd33e84b6bb8ad36d739e8414e

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8384724c3c

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[codspeed-hq]

Merging this PR will not alter performance

✅ 8 untouched benchmarks
⏩ 1 skipped benchmark¹

---

_{Comparing devin/1784924513-fix-cancelled-lease-acquire (cb3353b) with main (d630146)}

Footnotes

1. 1 benchmark was skipped, so the baseline result was used instead. If it was deleted from the codebase, click here and archive it to remove it from the performance reports. ↩

[desertaxle]

The acquisition service owns the client that granted the lease, but cleanup creates another client from the caller’s context. Can we put release behind the acquisition service so acquisition and cleanup always use the same API?

[devin-ai-integration]

Agreed — moved it in a256dad. ConcurrencySlotAcquisitionWithLeaseService now exposes release_orphaned_lease(response), which schedules self._discard(response) on the service's loop so cleanup reuses the client (and therefore the API) that granted the lease. The caller-side helper is just service.release_orphaned_lease(completed.result()) in the future's done callback — no extra thread, no second client, and contextvars/get_client(sync_client=True) are gone from _asyncio.py. If the service's loop is already gone we log a warning instead of silently leaking.

[desertaxle]

future.set_result() only starts the release thread, so patching get_client still leaves this test sensitive to scheduling. Can the test inject the release operation instead and verify that the orphaned lease is released?

[devin-ai-integration]

Reworked in a256dad. The caller-side test no longer patches get_client or waits on a thread: it injects a MagicMock(spec=ConcurrencySlotAcquisitionWithLeaseService) as the service instance and asserts service.release_orphaned_lease.assert_called_once_with(response) — the done callback runs synchronously inside future.set_result(), so there's no scheduling involved. The actual release is verified separately at the service level in test_release_orphaned_lease_uses_the_acquiring_client, which awaits the future returned by release_orphaned_lease and asserts the acquiring client got release_concurrency_slots_with_lease(lease_id=...).

[desertaxle]

This suppresses JSON parsing failures from a successful acquisition response, so a malformed response can leave an active lease without any diagnostic. Can you handle cancelled or failed futures separately and log invalid successful responses?

[devin-ai-integration]

Fixed in a256dad. The done callback now checks completed.cancelled() / completed.exception() first and returns for those (nothing was granted, nothing to log), and parsing moved into ConcurrencySlotAcquisitionWithLeaseService._lease_id_from_response, which logs a warning with exc_info when a successful response can't be read: "Unable to read the lease id from the acquisition response for ...; if a lease was granted its slots remain occupied until it expires." Covered by test_release_orphaned_lease_logs_unreadable_response and a parametrized test_does_not_release_when_no_lease_was_granted.

[desertaxle]

drain() can finish and close _client while this fire-and-forget coroutine is still pending, so the orphaned lease remains occupied until expiry. Can you make this cleanup service-owned and joined, with a cancellation-plus-immediate-drain regression test?

[devin-ai-integration]

Good catch — fixed in cb3353b. release_orphaned_lease now registers the scheduled release under the same lock _stop() takes, so a release can only be scheduled while the service is still running, and the service's _lifespan awaits any pending releases before the client closes. If the service is already stopped we log and return an immediately-completed future instead of scheduling work that cannot finish.

Added test_drain_waits_for_orphaned_lease_release: it schedules a release that yields repeatedly and then drains immediately. It fails without the tracking (drain returns before the release completes) and passes with it.