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

When receive_messages grants AMQP link credit, the broker may transfer more
messages than the application consumes before the call returns. Previously,
closing the receiver cleared the client buffer without settling it and never
drained the link, leaving those messages locked at the broker until lock
expiry (delaying redelivery and inflating the delivery count).

This adds a transport-layer drain-and-release on close (the vendored _pyamqp
package is left untouched):

- PyamqpTransport.drain_and_release_messages (and the async twin) drains the
  link via flow(link_credit=0, drain=True), pumps the connection until
  quiescent (bounded by a timeout), then releases every buffered message with
  the released disposition.
- ServiceBusReceiver._close_handler calls it, gated to non-session PEEK_LOCK
  receivers. The non-session gate matches the .NET SDK; the PEEK_LOCK gate is a
  Python refinement (a pre-settled RECEIVE_AND_DELETE delivery cannot be
  released-settled).

Adds unit tests covering the transport drain/release paths and the receiver
close gating (sync + async).

Fixes #42917

Author: pranz1996

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -12,6 +12,7 @@
 - Read `com.microsoft:max-message-batch-size` vendor property from the AMQP sender link to correctly limit batch size on Premium large-message entities, where `max-message-size` can be up to 100 MB but the batch limit is 1 MB.
 - Fixed a bug where sending a batched or multi-message payload with `uamqp_transport=True` raised `TypeError: 'BatchMessage' object is not subscriptable` (and a masked `AttributeError` on the list path) when the first message carried a `message_id`, `session_id`, or `partition_key`. The batch envelope properties are now set through the transport-appropriate code path. (regression from [#42598](https://github.com/Azure/azure-sdk-for-python/pull/42598))
 - Fixed a bug where the async receiver factory methods on `azure.servicebus.aio.ServiceBusClient` (`get_queue_receiver`, `get_subscription_receiver`) and the async `ServiceBusReceiver` annotated the `auto_lock_renewer` keyword with the synchronous `AutoLockRenewer`, causing static type checkers to reject the documented `azure.servicebus.aio.AutoLockRenewer` usage. The annotation now references the async `AutoLockRenewer`, matching the docstrings and runtime behavior. ([#47948](https://github.com/Azure/azure-sdk-for-python/issues/47948))
+- Fixed a bug where closing a `PEEK_LOCK` receiver did not release messages that had been prefetched into the client buffer or were still in flight, so they remained locked at the broker until lock expiry — delaying their redelivery and inflating their delivery count. On close, a non-session `PEEK_LOCK` receiver now drains the link (stopping the broker and flushing in-flight transfers) and releases the buffered messages (`released` disposition), so the broker can redeliver them immediately without incrementing the delivery count. ([#42917](https://github.com/Azure/azure-sdk-for-python/issues/42917))
 
 ## 7.14.3 (2025-11-11)
 

sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py

@@ -538,6 +538,17 @@ def _renew_locks(self, *lock_tokens: str, **kwargs: Any) -> Any:
 
     def _close_handler(self):
         self._message_iter = None
+        if (
+            self._handler
+            and self._receive_mode == ServiceBusReceiveMode.PEEK_LOCK
+            and not self._session
+        ):
+            # Drain the link and release buffered/in-flight messages so they are not
+            # left locked at the broker until lock expiry (delaying redelivery,
+            # inflating delivery count). Non-session gate matches .NET; the PEEK_LOCK
+            # gate is Python-specific (a pre-settled RECEIVE_AND_DELETE delivery
+            # cannot be released-settled).
+            self._amqp_transport.drain_and_release_messages(self._handler)
         super(ServiceBusReceiver, self)._close_handler()
 
     @property

sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py

@@ -302,6 +302,14 @@ def reset_link_credit(handler, link_credit):
         :param int link_credit: The link credit.
         """
 
+    @staticmethod
+    @abstractmethod
+    def drain_and_release_messages(handler):
+        """
+        Drain the receive link and release buffered/in-flight messages on close.
+        :param ~uamqp.ReceiveClient or ~pyamqp.ReceiveClient handler: The handler.
+        """
+
     @staticmethod
     @abstractmethod
     def settle_message_via_receiver_link(

sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py

@@ -11,6 +11,7 @@
 from datetime import timezone
 from typing import Optional, Tuple, cast, List, TYPE_CHECKING, Any, Callable, Dict, Union, Iterator, Type
 import logging
+import queue
 
 from .._pyamqp import (
     utils,
@@ -157,6 +158,10 @@ def is_retryable(self, error):
     ERROR_CODE_TIMEOUT: OperationTimeoutError,
 }
 
+# Safety cap (seconds) for draining a receive link on close; the drain loop exits
+# early on quiescence, so a normal close returns well before this.
+RECEIVE_LINK_DRAIN_TIMEOUT = 5
+
 
 class PyamqpTransport(AmqpTransport):  # pylint: disable=too-many-public-methods
     """
@@ -799,6 +804,60 @@ def reset_link_credit(handler: "ReceiveClient", link_credit: int) -> None:
         """
         handler._link.flow(link_credit=link_credit)  # pylint: disable=protected-access
 
+    @staticmethod
+    def drain_and_release_messages(handler: "ReceiveClient") -> None:
+        """
+        Drain the receive link and release buffered/in-flight messages on close, so
+        they are not left locked at the broker until lock expiry. Intended for
+        non-session PEEK_LOCK receivers only (gated by the caller).
+        :param ReceiveClient handler: Client whose receive link to drain.
+        :rtype: None
+        """
+        # pylint: disable=protected-access
+        link = handler._link
+        if link is None or link._is_closed:
+            return  # cannot drain or settle through a closed/absent link
+        if link.current_link_credit > 0:
+            # Drain in-flight transfers into the buffer. Own try so a drain
+            # failure still falls through to the release below.
+            try:
+                outstanding = link.current_link_credit
+                link.flow(link_credit=0, drain=True)
+                deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT
+                idle_cycles = 0
+                while True:
+                    remaining = deadline - time.time()
+                    if remaining <= 0:
+                        break
+                    before = handler._received_messages.qsize()
+                    # listen() directly, not do_work() (which re-issues credit at 0 and
+                    # undoes the drain). batch=outstanding assembles multi-frame transfers
+                    # in one cycle; cap each read by the remaining budget so close stays
+                    # bounded. pyamqp drops the drain echo, so stop after 2 idle cycles.
+                    wait = remaining if handler._socket_timeout is None else min(handler._socket_timeout, remaining)
+                    handler._connection.listen(wait=wait, batch=outstanding)
+                    if handler._received_messages.qsize() == before:
+                        idle_cycles += 1
+                        if idle_cycles >= 2:
+                            break
+                    else:
+                        idle_cycles = 0
+            except Exception:  # pylint: disable=broad-except
+                _LOGGER.debug("Draining the receive link on close failed.", exc_info=True)
+        # Release buffered deliveries (incl. credit==0 prefetch) so they are not
+        # left locked until lock expiry. Per-message try so one bad tag doesn't
+        # strand the rest; get_nowait handles the empty()/get race.
+        while True:
+            try:
+                frame, _ = handler._received_messages.get_nowait()
+            except queue.Empty:
+                break
+            try:
+                handler.settle_messages(frame[1], frame[2], "released")
+                handler._received_messages.task_done()
+            except Exception:  # pylint: disable=broad-except
+                _LOGGER.debug("Releasing a buffered message on close failed.", exc_info=True)
+
     @staticmethod
     def settle_message_via_receiver_link(
         handler: ReceiveClient,

sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py

@@ -848,6 +848,15 @@ def reset_link_credit(handler: "ReceiveClient", link_credit: int) -> None:
             """
             handler.message_handler.reset_link_credit(link_credit)
 
+        @staticmethod
+        def drain_and_release_messages(handler: "ReceiveClient") -> None:
+            """
+            No-op for uamqp: drain-on-close is only implemented for the pyamqp
+            transport (the default). uamqp is deprecated.
+            :param ~uamqp.ReceiveClient handler: The handler.
+            :rtype: None
+            """
+
         # Executes message settlement, implementation is in settle_message_via_receiver_link_impl
         # May be able to remove and just call methods in private method.
         @staticmethod

sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py

@@ -527,6 +527,17 @@ async def _renew_locks(self, *lock_tokens: str, timeout: Optional[float] = None)
 
     async def _close_handler(self):
         self._message_iter = None
+        if (
+            self._handler
+            and self._receive_mode == ServiceBusReceiveMode.PEEK_LOCK
+            and not self._session
+        ):
+            # Drain the link and release buffered/in-flight messages so they are not
+            # left locked at the broker until lock expiry (delaying redelivery,
+            # inflating delivery count). Non-session gate matches .NET; the PEEK_LOCK
+            # gate is Python-specific (a pre-settled RECEIVE_AND_DELETE delivery
+            # cannot be released-settled).
+            await self._amqp_transport.drain_and_release_messages_async(self._handler)
         await super(ServiceBusReceiver, self)._close_handler()
 
     @property

sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py

@@ -244,6 +244,16 @@ async def reset_link_credit_async(handler, link_credit):
         :rtype: None
         """
 
+    @staticmethod
+    @abstractmethod
+    async def drain_and_release_messages_async(handler):
+        """
+        Drain the receive link and release buffered/in-flight messages on close.
+        :param ~uamqp.ReceiveClientAsync
+         or ~pyamqp.aio.ReceiveClientAsync handler: Client whose receive link to drain.
+        :rtype: None
+        """
+
     @staticmethod
     @abstractmethod
     async def settle_message_via_receiver_link_async(

sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py

@@ -9,6 +9,8 @@
 import time
 import math
 import random
+import logging
+import queue
 
 from ..._pyamqp import constants
 from ..._pyamqp.message import BatchMessage
@@ -42,7 +44,7 @@
     OPERATION_TIMEOUT,
     NEXT_AVAILABLE_SESSION,
 )
-from ..._transport._pyamqp_transport import PyamqpTransport
+from ..._transport._pyamqp_transport import PyamqpTransport, RECEIVE_LINK_DRAIN_TIMEOUT
 from ...exceptions import (
     OperationTimeoutError,
     ServiceBusConnectionError,
@@ -59,6 +61,9 @@
     from ..._pyamqp.aio._client_async import AMQPClientAsync
 
 
+_LOGGER = logging.getLogger(__name__)
+
+
 class PyamqpTransportAsync(PyamqpTransport, AmqpTransportAsync):
     """
     Class which defines pyamqp-based methods used by the sender and receiver.
@@ -301,6 +306,61 @@ async def reset_link_credit_async(handler: "ReceiveClientAsync", link_credit: in
         """
         await handler._link.flow(link_credit=link_credit)  # pylint: disable=protected-access
 
+    @staticmethod
+    async def drain_and_release_messages_async(handler: "ReceiveClientAsync") -> None:
+        """
+        Drain the receive link and release buffered/in-flight messages on close, so
+        they are not left locked at the broker until lock expiry. Intended for
+        non-session PEEK_LOCK receivers only (gated by the caller).
+        :param ReceiveClientAsync handler: Client whose receive link to drain.
+        :rtype: None
+        """
+        # pylint: disable=protected-access
+        link = handler._link
+        if link is None or link._is_closed:
+            return  # cannot drain or settle through a closed/absent link
+        if link.current_link_credit > 0:
+            # Drain in-flight transfers into the buffer. Own try so a drain
+            # failure still falls through to the release below.
+            try:
+                outstanding = link.current_link_credit
+                await link.flow(link_credit=0, drain=True)
+                deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT
+                idle_cycles = 0
+                while True:
+                    remaining = deadline - time.time()
+                    if remaining <= 0:
+                        break
+                    before = handler._received_messages.qsize()
+                    # listen() directly, not do_work_async() (which re-issues credit at
+                    # 0 and undoes the drain). batch=outstanding assembles multi-frame
+                    # transfers in one cycle; cap each read by the remaining budget so
+                    # close stays bounded. pyamqp drops the drain echo, so stop after 2
+                    # idle cycles.
+                    wait = remaining if handler._socket_timeout is None else min(handler._socket_timeout, remaining)
+                    await handler._connection.listen(wait=wait, batch=outstanding)
+                    if handler._received_messages.qsize() == before:
+                        idle_cycles += 1
+                        if idle_cycles >= 2:
+                            break
+                    else:
+                        idle_cycles = 0
+            except Exception:  # pylint: disable=broad-except
+                _LOGGER.debug("Draining the receive link on close failed.", exc_info=True)
+        # Release buffered deliveries (incl. credit==0 prefetch) so they are not
+        # left locked until lock expiry. Per-message try so one bad tag doesn't
+        # strand the rest; get_nowait handles the empty()/get race.
+        while True:
+            try:
+                frame, _ = handler._received_messages.get_nowait()
+            except queue.Empty:
+                break
+            try:
+                await handler.settle_messages_async(frame[1], frame[2], "released")
+                handler._received_messages.task_done()
+            except Exception:  # pylint: disable=broad-except
+                _LOGGER.debug("Releasing a buffered message on close failed.", exc_info=True)
+
     @staticmethod
     async def settle_message_via_receiver_link_async(
         handler: "ReceiveClientAsync",

sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py

@@ -248,6 +248,15 @@ async def reset_link_credit_async(handler: "ReceiveClientAsync", link_credit: in
             """
             await handler.message_handler.reset_link_credit_async(link_credit)
 
+        @staticmethod
+        async def drain_and_release_messages_async(handler: "ReceiveClientAsync") -> None:
+            """
+            No-op for uamqp: drain-on-close is only implemented for the pyamqp
+            transport (the default). uamqp is deprecated.
+            :param ReceiveClientAsync handler: The handler.
+            :rtype: None
+            """
+
         @staticmethod
         async def settle_message_via_receiver_link_async(
             handler: "ReceiveClientAsync",