[Service Bus] Drain receive link and release messages on receiver close (#42917)

[pranz1996]

Summary

Fixes #42917. A PEEK_LOCK Service Bus receiver, on close, dropped messages
that had been prefetched into the client buffer or were still in flight on
the wire (leftover link credit), leaving them locked at the broker until lock
expiry — delaying redelivery and inflating delivery_count (risking premature
dead-lettering under repeated receiver churn).

Repro

The reporter's scenario, reproduced against a live namespace:

• async receiver, prefetch_count=0, receiver recreated each loop;
• a sender trickles messages (~2 msg/s) so receive_messages returns partial
  batches, leaving outstanding link credit;
• each received message is completed, then the receiver is closed — repeat.

Bug metric: messages redelivered with delivery_count > 0 — i.e. left
locked at close and redelivered only after the lock expired.

Why the test queue uses LockDuration = 10s: this is a test-environment
choice, not part of the fix. The default lock duration is up to 5 minutes, so a
lock-expiry redelivery would take minutes to observe. A short 10s lock makes the
bug surface within seconds. The fix itself uses no sleep and no
lock-duration-dependent timing — the reporter's asyncio.sleep(10) workaround
is removed entirely.

Fix (transport layer — vendored _pyamqp untouched)

The _pyamqp package is vendored and shared with azure-eventhub, so the fix
lives in the Service Bus transport layer instead of the vendored client:

• PyamqpTransport.drain_and_release_messages (and the async twin) sends a
  drain flow (flow(link_credit=0, drain=True)), pumps the connection until
  quiescent (bounded by RECEIVE_LINK_DRAIN_TIMEOUT = 5s, which exits early on
  quiescence — not a sleep), then releases every buffered/in-flight message with
  the released disposition. No-op for the deprecated uamqp transport.
• ServiceBusReceiver._close_handler calls it, gated to non-session
  PEEK_LOCK receivers (RECEIVE_AND_DELETE / Event Hubs use First mode where
  the drain is a no-op). This is a superset of .NET's drain-on-close: like
  .NET it drains in-flight transfers, and additionally releases the already
  buffered prefetch so nothing is left locked.

released (not modified/abandon) means the broker redelivers immediately
without incrementing delivery_count.

Tested after the fix

Unit (tests/unittests/test_receiver_drain_on_close.py, 19 tests):
transport drain/release mechanics (drain flow sent, in-flight pulled in,
released; skipped on no-credit / closed / null link; drain-failure still
releases; timeout-bounded) + receiver gating (PEEK_LOCK-non-session drains;
RECEIVE_AND_DELETE and session skip) — sync and async. Verified all behavior
tests fail when the fix is disabled. Full offline unit-test suite passes;
pylint/mypy clean on changed files.

Live A/B on a real namespace (queue LockDuration 10s), reporter's scenario,
no sleep workaround, ~60–90s per run — lock-expiry redeliveries
(delivery_count > 0) per run:

Build	Redeliveries per run
Unpatched (7.14.2)	5, 9
Release-only (interim)	4, 4
Drain + release (this PR)	0, 0, 0 (+ 0 under a POST_DELAY amplifier)

Skipped paths: RECEIVE_AND_DELETE and session receivers receive all messages
and close cleanly (drain correctly skipped).

Are there any breaking changes?

No. Behavior only improves (locks released sooner); no public API change.

[Copilot]

Pull request overview

Fixes delayed redelivery of buffered PEEK_LOCK messages when Service Bus receivers close.

Changes:

• Releases buffered messages during synchronous and asynchronous close.
• Adds seven regression tests.
• Documents the fix in the changelog.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
CHANGELOG.md	Documents the receiver-close fix.
_pyamqp/client.py	Drains buffered messages during synchronous close.
_pyamqp/aio/_client_async.py	Implements asynchronous buffer draining.
test_receiver_drain_on_close.py	Tests settlement modes, empty buffers, and failures.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 4 comments.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.

[pranz1996]

@microsoft-github-policy-service agree company="Microsoft"

[EldertGrootenboer]

Not ready to merge yet: CI is red on the pylint guidelines checker.

Rename drain_receive_link_and_release_messages_async in the three async transport files; it's over the 40-char limit (C4751), e.g. drain_and_release_messages_async.

The fix itself looks solid.

A few smaller things, none blocking:

The except Exception: pass swallows the error with no log; a debug log there would help when triaging a close that didn't release.

Drain and release share one try block, so if the drain flow throws, the already-buffered messages don't get released either.

The "parity with .NET" wording is slightly off: .NET gates on non-session only and just drains, while this also releases the prefetch buffer, so it's a superset.

[pranz1996]

Thanks Eldert — all addressed in f301213833:

1. Name length (C4751) — renamed to drain_and_release_messages / _async (32 chars).
2. Split try — drain and release are now separate try/except blocks, so a drain failure still releases the buffer. Tests updated to assert this.
3. Debug log — added on both failure paths.
4. ".NET parity" — reworded in the PR description to superset of .NET's drain-on-close (drains in-flight and releases buffered prefetch).

Also re-verified all 19 tests fail with the fix disabled. Ready for another look.

[EldertGrootenboer]

The rename, the split try, and the debug logs all look good.

Two things I'd still like before this merges, then a couple of smaller ones.

The PEEK_LOCK and not session gate is described as ".NET parity", but that isn't quite what .NET does.
AmqpReceiver.CloseAsync gates on !_isSessionReceiver && LinkCredit > 0 only, with no receive-mode check, so .NET drains non-session RECEIVE_AND_DELETE too.
The non-session half is genuine parity; the PEEK_LOCK restriction is a Python-specific refinement (you can't released-settle a pre-settled RECEIVE_AND_DELETE delivery), which is the right call but a Python addition, not something .NET does.
Can you reword the code comment and the test docstring so the two gates aren't both attributed to .NET?

The release loop wraps the whole while in one try, so if a single settle_messages throws (a stale delivery tag, or the broker detaching mid-drain) every message after it is skipped and left locked until expiry, which is the symptom this PR fixes.
Moving the try inside the loop keeps one bad tag from stranding the rest, and folds in the empty()/get_nowait() race at the same time.
Same in the async twin.

Smaller:

The async drain tests don't cover the null-link, closed-link, or empty-buffer cases that the sync class does, even though the async impl has the same guards.
Adding the three async parity tests keeps a broken async guard from shipping green.
Related: three of the 19 tests pass against a no-op body (the skip-when-closed / null / empty ones), so "all 19 fail with the fix disabled" is really 16.

_pyamqp_transport_async.py imports _LOGGER from the sync module, so the async drain and release debug logs get attributed to the sync namespace.
A local _LOGGER = logging.getLogger(__name__) matches the rest of the package.

One optional thing: the tests assert the released disposition, not the broker outcome.
If there's a live run behind the 0,0,0 table, linking it on the PR would be good to have in the record.
Not a blocker.

[pranz1996]

All addressed in 6db78e4909:

1. .NET gate — reworded so only the non-session gate is attributed to .NET; the PEEK_LOCK gate is called out as a Python-specific refinement.
2. Release loop — try now inside the loop (one bad tag won't strand the rest) + queue.Empty race handled. Same async.
3. Async tests — added null-link / closed-link / empty-buffer parity tests.
4. Async _LOGGER — now local to the async module.

[EldertGrootenboer]

CI is green now, and the rename plus the try-split, per-message release, and logging all landed cleanly.

Thanks for the quick turnaround. Approving.

[SwayGom]

Approving — this is a strong, well-scoped fix with excellent test coverage (16 transport + 6 gating tests, sync+async, verified to fail when the fix is disabled) and a convincing live A/B (0 lock-expiry redeliveries). The transport-layer approach (leaving vendored _pyamqp untouched), the use of released (so delivery_count isn't bumped), the separate try-blocks so a drain failure still releases the buffer, and using listen() instead of do_work() (which would re-issue credit and undo the drain) are all correct and nicely commented. Approving with a couple of clarifying questions and small nits below — none blocking.

Questions

Q1 — Session receivers are excluded; please confirm they don't have the same leak. The gate is PEEK_LOCK and not self._session, and the comment justifies the non-session part as ".NET parity." I believe sessions are safe because closing a session receiver releases the session lock, making its unsettled messages immediately available — but that only holds if close releases the session lock immediately. If it instead relies on session-lock expiry, session receivers have the same delayed-redelivery / inflated-delivery_count problem and should also drain+release. Could you confirm which it is, and capture the session-lock-release rationale in the comment (it currently reads as parity, not mechanism)?

Q2 — Worst-case close() latency when _socket_timeout is None. When _socket_timeout is None, wait = remaining (up to the full RECEIVE_LINK_DRAIN_TIMEOUT = 5s). test_bounded_by_deadline_even_with_blocking_listen correctly proves it stays bounded — but is a ~5s worst-case acceptable for close()? In practice, is _socket_timeout effectively always non-None (keeping idle cycles sub-second)? If it can be None, consider capping wait with a small idle-poll interval so normal closes stay snappy while still bounded.

Minor / nits (non-blocking)

• 2-idle-cycle quiescence heuristic: since pyamqp drops the drain echo, quiescence is inferred from 2 consecutive empty listen()s. A slow in-flight transfer arriving after 2 idle cycles (but before the deadline) would be missed and left in-flight → still locked. Bounded and empirically 0 in your A/B, so fine — worth a one-line note on the tradeoff (and whether the threshold should scale with outstanding on high-latency links).
• task_done() skipped on settle failure (both sync and async release loops): it runs only after a successful settle_messages, so the queue accounting drifts on the except path. Harmless on close, but a finally keeps it consistent:

try:
    handler.settle_messages(frame[1], frame[2], "released")
except Exception:  # pylint: disable=broad-except
    _LOGGER.debug("Releasing a buffered message on close failed.", exc_info=True)
finally:
    handler._received_messages.task_done()

(and the await/settle_messages_async equivalent in the async transport.)

• Sync/async duplication: the two drain_and_release_messages bodies are ~55 near-identical lines. Understandable given the sync/async split and it reads clearly — just flagging, not blocking.

Nice work — this cleanly resolves #42917.

[j7nw4r]

This 2 is the one drain knob that didn't get promoted to a constant, unlike RECEIVE_LINK_DRAIN_TIMEOUT above it. Can we pull it out as RECEIVE_LINK_DRAIN_IDLE_CYCLES = 2 next to the timeout, so both knobs sit together and the "why 2" comment attaches to a name? Same literal lives in the async twin at _pyamqp_transport_async.py:344.