From b396c8738fc7502b75a4a1eff0628a335edef666 Mon Sep 17 00:00:00 2001 From: Pranjal Patel Date: Fri, 10 Jul 2026 10:37:13 -0700 Subject: [PATCH 1/4] fix(servicebus): release buffered messages on receiver close On close, a PEEK_LOCK receiver dropped messages that had been prefetched into the client buffer but not yet returned to the application, leaving them locked at the broker until lock expiry (delaying redelivery and inflating delivery_count). Both the sync (ReceiveClient.close) and async (ReceiveClientAsync.close_async) pyamqp paths now release buffered-but- unconsumed messages with the `released` disposition before clearing the buffer. Gated on PEEK_LOCK (ReceiverSettleMode.Second) and guarded so a faulted link cannot block close. Adds deterministic regression tests. Fixes #42917 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bef8092b-2a9f-4862-aef1-680eeb2e97f0 --- sdk/servicebus/azure-servicebus/CHANGELOG.md | 1 + .../servicebus/_pyamqp/aio/_client_async.py | 17 ++ .../azure/servicebus/_pyamqp/client.py | 16 ++ .../unittests/test_receiver_drain_on_close.py | 171 ++++++++++++++++++ 4 files changed, 205 insertions(+) create mode 100644 sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py diff --git a/sdk/servicebus/azure-servicebus/CHANGELOG.md b/sdk/servicebus/azure-servicebus/CHANGELOG.md index afe2fa018c4d..292ef452f084 100644 --- a/sdk/servicebus/azure-servicebus/CHANGELOG.md +++ b/sdk/servicebus/azure-servicebus/CHANGELOG.md @@ -11,6 +11,7 @@ - Fixed a bug where messages returned by `receive_deferred_messages` had a `lock_token` of `None`, which prevented settling (completing, abandoning, dead-lettering, deferring) or renewing the lock on a deferred message in `PEEK_LOCK` mode. The lock token is now read from the `lock-token` field of the management-link response for deferred messages. ([#42454](https://github.com/Azure/azure-sdk-for-python/issues/42454)) - 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 closing a `PEEK_LOCK` receiver did not release messages that had been prefetched into the client buffer but not yet returned to the application. These messages remained locked at the broker until lock expiry, which delayed their redelivery and inflated their delivery count. On close, buffered-but-unconsumed messages are now released (`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) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py index 087a0c48b243..ffc4d260e7d7 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py @@ -25,6 +25,7 @@ from ..message import _MessageDelivery, Message from ..constants import ( MessageDeliveryState, + ReceiverSettleMode, SEND_DISPOSITION_ACCEPT, SEND_DISPOSITION_REJECT, LinkDeliverySettleReason, @@ -817,6 +818,22 @@ async def _receive_message_batch_impl_async( return batch async def close_async(self): + # Release any buffered-but-unconsumed messages before closing so they are + # not left locked at the broker until lock expiry (which also inflates the + # delivery count and delays redelivery). Only applies to PEEK_LOCK + # (ReceiverSettleMode.Second); in RECEIVE_AND_DELETE (First) the transfers + # are already settled, so there is nothing to release. + if self._receive_settle_mode != ReceiverSettleMode.First: + while not self._received_messages.empty(): + try: + frame, _ = self._received_messages.get_nowait() + await self.settle_messages_async(frame[1], frame[2], "released") + self._received_messages.task_done() + except queue.Empty: + break + except Exception: # pylint:disable=broad-except + # A faulted or already-detached link must never block close. + break self._received_messages = queue.Queue() await super(ReceiveClientAsync, self).close_async() diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py index 657c9bd8a879..4bb4378f19b1 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py @@ -931,6 +931,22 @@ def _receive_message_batch_impl( return batch def close(self): + # Release any buffered-but-unconsumed messages before closing so they are + # not left locked at the broker until lock expiry (which also inflates the + # delivery count and delays redelivery). Only applies to PEEK_LOCK + # (ReceiverSettleMode.Second); in RECEIVE_AND_DELETE (First) the transfers + # are already settled, so there is nothing to release. + if self._receive_settle_mode != ReceiverSettleMode.First: + while not self._received_messages.empty(): + try: + frame, _ = self._received_messages.get_nowait() + self.settle_messages(frame[1], frame[2], "released") + self._received_messages.task_done() + except queue.Empty: + break + except Exception: # pylint:disable=broad-except + # A faulted or already-detached link must never block close. + break self._received_messages = queue.Queue() super(ReceiveClient, self).close() diff --git a/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py b/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py new file mode 100644 index 000000000000..74b96345fe74 --- /dev/null +++ b/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py @@ -0,0 +1,171 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +# pylint: disable=protected-access + +""" +Unit tests for releasing buffered-but-unconsumed messages when a receive client +is closed (regression guard for azure-sdk-for-python#42917). + +When ``receive_messages`` grants AMQP link credit, the broker may transfer more +messages than the application consumes before the call returns. Those surplus +messages sit in the client-side buffer ``_received_messages``. Prior to the fix, +``close``/``close_async`` cleared that buffer without settling it, leaving the +messages locked at the broker until lock expiry (delaying redelivery and +inflating the delivery count). + +These tests assert the drain-on-close behavior directly against the pyamqp +``ReceiveClient`` / ``ReceiveClientAsync`` close paths, with no network: + +* PEEK_LOCK (``ReceiverSettleMode.Second``): every buffered message is released + (``released`` disposition) before the buffer is cleared. +* RECEIVE_AND_DELETE (``ReceiverSettleMode.First``): messages are already settled + on transfer, so nothing is released (asymmetry — the SAME buffered input yields + releases in one mode and none in the other). +* A failure while releasing must never block close. +""" + +import queue + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from azure.servicebus._pyamqp.client import ReceiveClient, AMQPClient +from azure.servicebus._pyamqp.aio._client_async import ReceiveClientAsync, AMQPClientAsync +from azure.servicebus._pyamqp.constants import ReceiverSettleMode +from azure.servicebus._pyamqp.performatives import TransferFrame + + +def _buffer_with(*deliveries): + """Build a _received_messages queue holding (TransferFrame, message) tuples. + + Mirrors what the receive callback puts on the buffer: the transfer frame + (which carries delivery_id at index 1 and delivery_tag at index 2) paired + with the decoded message. + """ + buf = queue.Queue() + for delivery_id, delivery_tag in deliveries: + frame = TransferFrame(handle=0, delivery_id=delivery_id, delivery_tag=delivery_tag) + buf.put((frame, MagicMock(name=f"message-{delivery_id}"))) + return buf + + +def _make_sync_client(settle_mode, buffer): + """Create an uninitialized ReceiveClient with only the attributes close() touches.""" + client = ReceiveClient.__new__(ReceiveClient) + client._received_messages = buffer + client._receive_settle_mode = settle_mode + client.settle_messages = MagicMock(name="settle_messages") + return client + + +def _make_async_client(settle_mode, buffer): + """Create an uninitialized ReceiveClientAsync with only the attributes close_async() touches.""" + client = ReceiveClientAsync.__new__(ReceiveClientAsync) + client._received_messages = buffer + client._receive_settle_mode = settle_mode + client.settle_messages_async = AsyncMock(name="settle_messages_async") + return client + + +class TestSyncReceiverDrainOnClose: + def test_peek_lock_releases_all_buffered_messages(self): + buffer = _buffer_with((10, b"tag-10"), (11, b"tag-11"), (12, b"tag-12")) + client = _make_sync_client(ReceiverSettleMode.Second, buffer) + + with patch.object(AMQPClient, "close") as super_close: + client.close() + + # Every buffered delivery is released with the 'released' disposition, + # keyed by (delivery_id, delivery_tag) from the transfer frame. + assert client.settle_messages.call_count == 3 + released = [ + (args[0], args[1], args[2]) for args, _ in client.settle_messages.call_args_list + ] + assert released == [ + (10, b"tag-10", "released"), + (11, b"tag-11", "released"), + (12, b"tag-12", "released"), + ] + # Buffer is drained and the parent close still runs. + assert client._received_messages.empty() + super_close.assert_called_once() + + def test_receive_and_delete_does_not_release(self): + # SAME buffered input as the PEEK_LOCK case — asymmetry is driven only by mode. + buffer = _buffer_with((10, b"tag-10"), (11, b"tag-11"), (12, b"tag-12")) + client = _make_sync_client(ReceiverSettleMode.First, buffer) + + with patch.object(AMQPClient, "close") as super_close: + client.close() + + client.settle_messages.assert_not_called() + super_close.assert_called_once() + + def test_empty_buffer_releases_nothing(self): + client = _make_sync_client(ReceiverSettleMode.Second, queue.Queue()) + + with patch.object(AMQPClient, "close") as super_close: + client.close() + + client.settle_messages.assert_not_called() + super_close.assert_called_once() + + def test_release_failure_does_not_block_close(self): + buffer = _buffer_with((10, b"tag-10"), (11, b"tag-11")) + client = _make_sync_client(ReceiverSettleMode.Second, buffer) + # Simulate a faulted / already-detached link raising on disposition. + client.settle_messages.side_effect = ValueError("link detached") + + with patch.object(AMQPClient, "close") as super_close: + client.close() # must not raise + + # We stop releasing on the first failure but still tear down cleanly. + assert client.settle_messages.call_count == 1 + assert client._received_messages.empty() + super_close.assert_called_once() + + +class TestAsyncReceiverDrainOnClose: + @pytest.mark.asyncio + async def test_peek_lock_releases_all_buffered_messages(self): + buffer = _buffer_with((20, b"tag-20"), (21, b"tag-21")) + client = _make_async_client(ReceiverSettleMode.Second, buffer) + + with patch.object(AMQPClientAsync, "close_async", new=AsyncMock()) as super_close: + await client.close_async() + + assert client.settle_messages_async.call_count == 2 + released = [ + (args[0], args[1], args[2]) for args, _ in client.settle_messages_async.call_args_list + ] + assert released == [ + (20, b"tag-20", "released"), + (21, b"tag-21", "released"), + ] + assert client._received_messages.empty() + super_close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_receive_and_delete_does_not_release(self): + buffer = _buffer_with((20, b"tag-20"), (21, b"tag-21")) + client = _make_async_client(ReceiverSettleMode.First, buffer) + + with patch.object(AMQPClientAsync, "close_async", new=AsyncMock()) as super_close: + await client.close_async() + + client.settle_messages_async.assert_not_called() + super_close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_release_failure_does_not_block_close(self): + buffer = _buffer_with((20, b"tag-20"), (21, b"tag-21")) + client = _make_async_client(ReceiverSettleMode.Second, buffer) + client.settle_messages_async.side_effect = ValueError("link detached") + + with patch.object(AMQPClientAsync, "close_async", new=AsyncMock()) as super_close: + await client.close_async() # must not raise + + assert client.settle_messages_async.call_count == 1 + assert client._received_messages.empty() + super_close.assert_awaited_once() From 383e856204717bd408c37e40cb3b4eebd5cb69ea Mon Sep 17 00:00:00 2001 From: Pranjal Patel Date: Fri, 10 Jul 2026 14:43:47 -0700 Subject: [PATCH 2/4] fix(servicebus): drain receive link on close in the transport layer Reworks the #42917 fix to keep the vendored _pyamqp package untouched (it is shared with azure-eventhub and maintained in lockstep). The earlier edit to _pyamqp/client.py and _client_async.py is reverted; drain-on-close now lives in the Service Bus transport layer: - PyamqpTransport.drain_receive_link_and_release_messages (+ async twin) sends a drain flow (flow link_credit=0, drain=True), pumps the connection until quiescent (bounded by RECEIVE_LINK_DRAIN_TIMEOUT), then releases buffered/in-flight messages with the `released` disposition. No-op for the deprecated uamqp transport. - ServiceBusReceiver._close_handler calls it, gated to non-session PEEK_LOCK receivers (parity with the .NET SDK). Draining in-flight (not just already-buffered) messages fully eliminates the delayed redelivery and inflated delivery count in the reporter's scenario. Fixes #42917 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bef8092b-2a9f-4862-aef1-680eeb2e97f0 --- sdk/servicebus/azure-servicebus/CHANGELOG.md | 2 +- .../servicebus/_pyamqp/aio/_client_async.py | 17 - .../azure/servicebus/_pyamqp/client.py | 16 - .../azure/servicebus/_servicebus_receiver.py | 9 + .../azure/servicebus/_transport/_base.py | 8 + .../_transport/_pyamqp_transport.py | 34 ++ .../servicebus/_transport/_uamqp_transport.py | 9 + .../aio/_servicebus_receiver_async.py | 9 + .../servicebus/aio/_transport/_base_async.py | 10 + .../aio/_transport/_pyamqp_transport_async.py | 32 +- .../aio/_transport/_uamqp_transport_async.py | 9 + .../unittests/test_receiver_drain_on_close.py | 372 ++++++++++++------ 12 files changed, 361 insertions(+), 166 deletions(-) diff --git a/sdk/servicebus/azure-servicebus/CHANGELOG.md b/sdk/servicebus/azure-servicebus/CHANGELOG.md index 292ef452f084..d3476e03bfff 100644 --- a/sdk/servicebus/azure-servicebus/CHANGELOG.md +++ b/sdk/servicebus/azure-servicebus/CHANGELOG.md @@ -11,7 +11,7 @@ - Fixed a bug where messages returned by `receive_deferred_messages` had a `lock_token` of `None`, which prevented settling (completing, abandoning, dead-lettering, deferring) or renewing the lock on a deferred message in `PEEK_LOCK` mode. The lock token is now read from the `lock-token` field of the management-link response for deferred messages. ([#42454](https://github.com/Azure/azure-sdk-for-python/issues/42454)) - 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 closing a `PEEK_LOCK` receiver did not release messages that had been prefetched into the client buffer but not yet returned to the application. These messages remained locked at the broker until lock expiry, which delayed their redelivery and inflated their delivery count. On close, buffered-but-unconsumed messages are now released (`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)) +- 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) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py index ffc4d260e7d7..087a0c48b243 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py @@ -25,7 +25,6 @@ from ..message import _MessageDelivery, Message from ..constants import ( MessageDeliveryState, - ReceiverSettleMode, SEND_DISPOSITION_ACCEPT, SEND_DISPOSITION_REJECT, LinkDeliverySettleReason, @@ -818,22 +817,6 @@ async def _receive_message_batch_impl_async( return batch async def close_async(self): - # Release any buffered-but-unconsumed messages before closing so they are - # not left locked at the broker until lock expiry (which also inflates the - # delivery count and delays redelivery). Only applies to PEEK_LOCK - # (ReceiverSettleMode.Second); in RECEIVE_AND_DELETE (First) the transfers - # are already settled, so there is nothing to release. - if self._receive_settle_mode != ReceiverSettleMode.First: - while not self._received_messages.empty(): - try: - frame, _ = self._received_messages.get_nowait() - await self.settle_messages_async(frame[1], frame[2], "released") - self._received_messages.task_done() - except queue.Empty: - break - except Exception: # pylint:disable=broad-except - # A faulted or already-detached link must never block close. - break self._received_messages = queue.Queue() await super(ReceiveClientAsync, self).close_async() diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py index 4bb4378f19b1..657c9bd8a879 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py @@ -931,22 +931,6 @@ def _receive_message_batch_impl( return batch def close(self): - # Release any buffered-but-unconsumed messages before closing so they are - # not left locked at the broker until lock expiry (which also inflates the - # delivery count and delays redelivery). Only applies to PEEK_LOCK - # (ReceiverSettleMode.Second); in RECEIVE_AND_DELETE (First) the transfers - # are already settled, so there is nothing to release. - if self._receive_settle_mode != ReceiverSettleMode.First: - while not self._received_messages.empty(): - try: - frame, _ = self._received_messages.get_nowait() - self.settle_messages(frame[1], frame[2], "released") - self._received_messages.task_done() - except queue.Empty: - break - except Exception: # pylint:disable=broad-except - # A faulted or already-detached link must never block close. - break self._received_messages = queue.Queue() super(ReceiveClient, self).close() diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py index 3c8c1dd0c87f..074032215d58 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py @@ -538,6 +538,15 @@ 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). + self._amqp_transport.drain_receive_link_and_release_messages(self._handler) super(ServiceBusReceiver, self)._close_handler() @property diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py index 5ee0409c4f34..d95a6cd89ab6 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_base.py +++ b/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_receive_link_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( diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py index ac5a3e1add7e..12045ea82680 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py @@ -157,6 +157,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 +803,36 @@ def reset_link_credit(handler: "ReceiveClient", link_credit: int) -> None: """ handler._link.flow(link_credit=link_credit) # pylint: disable=protected-access + @staticmethod + def drain_receive_link_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 or link.current_link_credit <= 0: + return + try: + link.flow(link_credit=0, drain=True) + deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT + while time.time() < deadline: + before = handler._received_messages.qsize() + # listen() directly, not do_work(): do_work re-issues credit at 0 + # and would undo the drain. Stop on quiescence (no new transfers). + handler._connection.listen(wait=handler._socket_timeout) + if handler._received_messages.qsize() == before: + break + while not handler._received_messages.empty(): + frame, _ = handler._received_messages.get_nowait() + handler.settle_messages(frame[1], frame[2], "released") + handler._received_messages.task_done() + except Exception: # pylint: disable=broad-except + pass # a faulted/detached link must not block close + @staticmethod def settle_message_via_receiver_link( handler: ReceiveClient, diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py index 1df6cbae992f..bc0c5a1cf90d 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py +++ b/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_receive_link_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 diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py index fa9ba1a0d679..ecf8cee38eb6 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py @@ -527,6 +527,15 @@ 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). + await self._amqp_transport.drain_receive_link_and_release_messages_async(self._handler) await super(ServiceBusReceiver, self)._close_handler() @property diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py index 00697accbd42..3e8c8d049593 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_base_async.py +++ b/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_receive_link_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( diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py index 105eda4572b2..01a02934add1 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py @@ -42,7 +42,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, @@ -301,6 +301,36 @@ 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_receive_link_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 or link.current_link_credit <= 0: + return + try: + await link.flow(link_credit=0, drain=True) + deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT + while time.time() < deadline: + before = handler._received_messages.qsize() + # listen() directly, not do_work_async(): do_work re-issues credit + # at 0 and would undo the drain. Stop on quiescence (no new transfers). + await handler._connection.listen(wait=handler._socket_timeout) + if handler._received_messages.qsize() == before: + break + while not handler._received_messages.empty(): + frame, _ = handler._received_messages.get_nowait() + await handler.settle_messages_async(frame[1], frame[2], "released") + handler._received_messages.task_done() + except Exception: # pylint: disable=broad-except + pass # a faulted/detached link must not block close + @staticmethod async def settle_message_via_receiver_link_async( handler: "ReceiveClientAsync", diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py index f59a33f0f511..c7ac0038eec1 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_uamqp_transport_async.py +++ b/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_receive_link_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", diff --git a/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py b/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py index 74b96345fe74..dc3d94a878b3 100644 --- a/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py +++ b/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py @@ -1,28 +1,32 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. -# pylint: disable=protected-access +# pylint: disable=protected-access,unused-argument """ -Unit tests for releasing buffered-but-unconsumed messages when a receive client -is closed (regression guard for azure-sdk-for-python#42917). +Unit tests for draining and releasing buffered/in-flight messages when a Service +Bus receiver is closed (regression guard for azure-sdk-for-python#42917). When ``receive_messages`` grants AMQP link credit, the broker may transfer more -messages than the application consumes before the call returns. Those surplus -messages sit in the client-side buffer ``_received_messages``. Prior to the fix, -``close``/``close_async`` cleared that buffer without settling it, leaving the -messages locked at the broker until lock expiry (delaying redelivery and -inflating the delivery count). - -These tests assert the drain-on-close behavior directly against the pyamqp -``ReceiveClient`` / ``ReceiveClientAsync`` close paths, with no network: - -* PEEK_LOCK (``ReceiverSettleMode.Second``): every buffered message is released - (``released`` disposition) before the buffer is cleared. -* RECEIVE_AND_DELETE (``ReceiverSettleMode.First``): messages are already settled - on transfer, so nothing is released (asymmetry — the SAME buffered input yields - releases in one mode and none in the other). -* A failure while releasing must never block close. +messages than the application consumes before the call returns. Those messages +end up either in the client-side buffer ``_received_messages`` or still in flight +on the wire (leftover credit). Prior to the fix, closing the receiver cleared the +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). + +The fix lives in the Service Bus transport layer (keeping the vendored ``_pyamqp`` +package untouched): + + * ``PyamqpTransport.drain_receive_link_and_release_messages`` (and the async + twin) drains the link -- ``flow(link_credit=0, drain=True)`` + pumping 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 (parity with the .NET SDK). + +These tests exercise the transport method directly and the receiver gating, with +no network. """ import queue @@ -30,19 +34,17 @@ import pytest from unittest.mock import AsyncMock, MagicMock, patch -from azure.servicebus._pyamqp.client import ReceiveClient, AMQPClient -from azure.servicebus._pyamqp.aio._client_async import ReceiveClientAsync, AMQPClientAsync -from azure.servicebus._pyamqp.constants import ReceiverSettleMode +from azure.servicebus import ServiceBusReceiveMode +from azure.servicebus._servicebus_receiver import ServiceBusReceiver +from azure.servicebus._base_handler import BaseHandler +from azure.servicebus._transport._pyamqp_transport import PyamqpTransport +from azure.servicebus.aio._servicebus_receiver_async import ServiceBusReceiver as ServiceBusReceiverAsync +from azure.servicebus.aio._base_handler_async import BaseHandler as BaseHandlerAsync +from azure.servicebus.aio._transport._pyamqp_transport_async import PyamqpTransportAsync from azure.servicebus._pyamqp.performatives import TransferFrame def _buffer_with(*deliveries): - """Build a _received_messages queue holding (TransferFrame, message) tuples. - - Mirrors what the receive callback puts on the buffer: the transfer frame - (which carries delivery_id at index 1 and delivery_tag at index 2) paired - with the decoded message. - """ buf = queue.Queue() for delivery_id, delivery_tag in deliveries: frame = TransferFrame(handle=0, delivery_id=delivery_id, delivery_tag=delivery_tag) @@ -50,122 +52,230 @@ def _buffer_with(*deliveries): return buf -def _make_sync_client(settle_mode, buffer): - """Create an uninitialized ReceiveClient with only the attributes close() touches.""" - client = ReceiveClient.__new__(ReceiveClient) - client._received_messages = buffer - client._receive_settle_mode = settle_mode - client.settle_messages = MagicMock(name="settle_messages") - return client - - -def _make_async_client(settle_mode, buffer): - """Create an uninitialized ReceiveClientAsync with only the attributes close_async() touches.""" - client = ReceiveClientAsync.__new__(ReceiveClientAsync) - client._received_messages = buffer - client._receive_settle_mode = settle_mode - client.settle_messages_async = AsyncMock(name="settle_messages_async") - return client - - -class TestSyncReceiverDrainOnClose: - def test_peek_lock_releases_all_buffered_messages(self): - buffer = _buffer_with((10, b"tag-10"), (11, b"tag-11"), (12, b"tag-12")) - client = _make_sync_client(ReceiverSettleMode.Second, buffer) - - with patch.object(AMQPClient, "close") as super_close: - client.close() - - # Every buffered delivery is released with the 'released' disposition, - # keyed by (delivery_id, delivery_tag) from the transfer frame. - assert client.settle_messages.call_count == 3 - released = [ - (args[0], args[1], args[2]) for args, _ in client.settle_messages.call_args_list - ] - assert released == [ - (10, b"tag-10", "released"), - (11, b"tag-11", "released"), - (12, b"tag-12", "released"), - ] - # Buffer is drained and the parent close still runs. - assert client._received_messages.empty() - super_close.assert_called_once() - - def test_receive_and_delete_does_not_release(self): - # SAME buffered input as the PEEK_LOCK case — asymmetry is driven only by mode. - buffer = _buffer_with((10, b"tag-10"), (11, b"tag-11"), (12, b"tag-12")) - client = _make_sync_client(ReceiverSettleMode.First, buffer) - - with patch.object(AMQPClient, "close") as super_close: - client.close() - - client.settle_messages.assert_not_called() - super_close.assert_called_once() - - def test_empty_buffer_releases_nothing(self): - client = _make_sync_client(ReceiverSettleMode.Second, queue.Queue()) +def _handler(*, is_closed=False, credit=5, buffer=None, connection=None, is_async=False): + h = MagicMock(name="handler") + h._link = MagicMock(name="link") + h._link._is_closed = is_closed + h._link.current_link_credit = credit + h._link.flow = AsyncMock(name="flow") if is_async else MagicMock(name="flow") + if connection is None: + connection = MagicMock(name="connection") + connection.listen = AsyncMock(name="listen") if is_async else MagicMock(name="listen") + h._connection = connection + h._received_messages = buffer if buffer is not None else queue.Queue() + h._socket_timeout = 0.2 + if is_async: + h.settle_messages_async = AsyncMock(name="settle_messages_async") + else: + h.settle_messages = MagicMock(name="settle_messages") + return h + + +# --------------------------------------------------------------------------- # +# Transport method: drain + release # +# --------------------------------------------------------------------------- # + +class TestPyamqpTransportDrain: + def test_drain_sent_and_buffer_released(self): + h = _handler(credit=5, buffer=_buffer_with((10, b"tag-10"), (11, b"tag-11"))) + PyamqpTransport.drain_receive_link_and_release_messages(h) + + h._link.flow.assert_called_once_with(link_credit=0, drain=True) + assert h._connection.listen.call_count >= 1 + released = [(a[0], a[1], a[2]) for a, _ in h.settle_messages.call_args_list] + assert released == [(10, b"tag-10", "released"), (11, b"tag-11", "released")] + assert h._received_messages.empty() + + def test_pulls_inflight_then_releases(self): + buffer = queue.Queue() + connection = MagicMock(name="connection") + calls = {"n": 0} + + def _listen(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + frame = TransferFrame(handle=0, delivery_id=99, delivery_tag=b"tag-99") + buffer.put((frame, MagicMock())) + + connection.listen = MagicMock(side_effect=_listen) + h = _handler(credit=3, buffer=buffer, connection=connection) + + PyamqpTransport.drain_receive_link_and_release_messages(h) + + h._link.flow.assert_called_once_with(link_credit=0, drain=True) + h.settle_messages.assert_called_once_with(99, b"tag-99", "released") + + def test_skipped_when_no_credit(self): + h = _handler(credit=0, buffer=_buffer_with((10, b"tag-10"))) + PyamqpTransport.drain_receive_link_and_release_messages(h) + h._link.flow.assert_not_called() + h.settle_messages.assert_not_called() + + def test_skipped_when_link_closed(self): + h = _handler(credit=5, is_closed=True, buffer=_buffer_with((10, b"tag-10"))) + PyamqpTransport.drain_receive_link_and_release_messages(h) + h._link.flow.assert_not_called() + + def test_skipped_when_link_none(self): + h = _handler(credit=5) + h._link = None + PyamqpTransport.drain_receive_link_and_release_messages(h) # must not raise + h.settle_messages.assert_not_called() + + def test_failure_does_not_raise(self): + h = _handler(credit=5, buffer=_buffer_with((10, b"tag-10"))) + h._link.flow.side_effect = ValueError("link faulted") + PyamqpTransport.drain_receive_link_and_release_messages(h) # must not raise + # drain aborted before releasing; no exception escapes + h.settle_messages.assert_not_called() + + def test_bounded_when_never_quiescent(self): + buffer = queue.Queue() + connection = MagicMock(name="connection") + seq = {"n": 0} + + def _listen(*args, **kwargs): + seq["n"] += 1 + frame = TransferFrame(handle=0, delivery_id=seq["n"], delivery_tag=bytes(str(seq["n"]), "ascii")) + buffer.put((frame, MagicMock())) # never quiescent + + connection.listen = MagicMock(side_effect=_listen) + h = _handler(credit=3, buffer=buffer, connection=connection) + + with patch("azure.servicebus._transport._pyamqp_transport.RECEIVE_LINK_DRAIN_TIMEOUT", 0.05): + PyamqpTransport.drain_receive_link_and_release_messages(h) # must terminate, not hang + + assert h._connection.listen.call_count >= 1 + + +class TestPyamqpTransportDrainAsync: + @pytest.mark.asyncio + async def test_drain_sent_and_buffer_released(self): + h = _handler(credit=5, buffer=_buffer_with((20, b"tag-20"), (21, b"tag-21")), is_async=True) + await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) - with patch.object(AMQPClient, "close") as super_close: - client.close() + h._link.flow.assert_awaited_once_with(link_credit=0, drain=True) + assert h._connection.listen.await_count >= 1 + released = [(a[0], a[1], a[2]) for a, _ in h.settle_messages_async.call_args_list] + assert released == [(20, b"tag-20", "released"), (21, b"tag-21", "released")] + assert h._received_messages.empty() - client.settle_messages.assert_not_called() - super_close.assert_called_once() + @pytest.mark.asyncio + async def test_pulls_inflight_then_releases(self): + buffer = queue.Queue() + connection = MagicMock(name="connection") + calls = {"n": 0} - def test_release_failure_does_not_block_close(self): - buffer = _buffer_with((10, b"tag-10"), (11, b"tag-11")) - client = _make_sync_client(ReceiverSettleMode.Second, buffer) - # Simulate a faulted / already-detached link raising on disposition. - client.settle_messages.side_effect = ValueError("link detached") + async def _listen(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + frame = TransferFrame(handle=0, delivery_id=88, delivery_tag=b"tag-88") + buffer.put((frame, MagicMock())) - with patch.object(AMQPClient, "close") as super_close: - client.close() # must not raise + connection.listen = AsyncMock(side_effect=_listen) + h = _handler(credit=3, buffer=buffer, connection=connection, is_async=True) - # We stop releasing on the first failure but still tear down cleanly. - assert client.settle_messages.call_count == 1 - assert client._received_messages.empty() - super_close.assert_called_once() + await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) + h._link.flow.assert_awaited_once_with(link_credit=0, drain=True) + h.settle_messages_async.assert_awaited_once_with(88, b"tag-88", "released") -class TestAsyncReceiverDrainOnClose: @pytest.mark.asyncio - async def test_peek_lock_releases_all_buffered_messages(self): - buffer = _buffer_with((20, b"tag-20"), (21, b"tag-21")) - client = _make_async_client(ReceiverSettleMode.Second, buffer) - - with patch.object(AMQPClientAsync, "close_async", new=AsyncMock()) as super_close: - await client.close_async() - - assert client.settle_messages_async.call_count == 2 - released = [ - (args[0], args[1], args[2]) for args, _ in client.settle_messages_async.call_args_list - ] - assert released == [ - (20, b"tag-20", "released"), - (21, b"tag-21", "released"), - ] - assert client._received_messages.empty() - super_close.assert_awaited_once() + async def test_skipped_when_no_credit(self): + h = _handler(credit=0, buffer=_buffer_with((20, b"tag-20")), is_async=True) + await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) + h._link.flow.assert_not_called() + h.settle_messages_async.assert_not_called() @pytest.mark.asyncio - async def test_receive_and_delete_does_not_release(self): - buffer = _buffer_with((20, b"tag-20"), (21, b"tag-21")) - client = _make_async_client(ReceiverSettleMode.First, buffer) + async def test_failure_does_not_raise(self): + h = _handler(credit=5, buffer=_buffer_with((20, b"tag-20")), is_async=True) + h._link.flow.side_effect = ValueError("link faulted") + await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) # must not raise + h.settle_messages_async.assert_not_called() - with patch.object(AMQPClientAsync, "close_async", new=AsyncMock()) as super_close: - await client.close_async() - - client.settle_messages_async.assert_not_called() - super_close.assert_awaited_once() + @pytest.mark.asyncio + async def test_bounded_when_never_quiescent(self): + buffer = queue.Queue() + connection = MagicMock(name="connection") + seq = {"n": 0} + + async def _listen(*args, **kwargs): + seq["n"] += 1 + frame = TransferFrame(handle=0, delivery_id=seq["n"], delivery_tag=bytes(str(seq["n"]), "ascii")) + buffer.put((frame, MagicMock())) + + connection.listen = AsyncMock(side_effect=_listen) + h = _handler(credit=3, buffer=buffer, connection=connection, is_async=True) + + with patch("azure.servicebus.aio._transport._pyamqp_transport_async.RECEIVE_LINK_DRAIN_TIMEOUT", 0.05): + await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) # must terminate + + assert h._connection.listen.await_count >= 1 + + +# --------------------------------------------------------------------------- # +# Receiver gating: only non-session PEEK_LOCK receivers drain on close # +# --------------------------------------------------------------------------- # + +def _sync_receiver(mode, session): + r = ServiceBusReceiver.__new__(ServiceBusReceiver) + r._handler = MagicMock(name="handler") + r._message_iter = None + r._receive_mode = mode + r._session = session + r._amqp_transport = MagicMock(name="transport") + r._amqp_transport.drain_receive_link_and_release_messages = MagicMock(name="drain") + return r + + +def _async_receiver(mode, session): + r = ServiceBusReceiverAsync.__new__(ServiceBusReceiverAsync) + r._handler = MagicMock(name="handler") + r._message_iter = None + r._receive_mode = mode + r._session = session + r._amqp_transport = MagicMock(name="transport") + r._amqp_transport.drain_receive_link_and_release_messages_async = AsyncMock(name="drain") + return r + + +class TestReceiverCloseGating: + def test_peek_lock_non_session_drains(self): + r = _sync_receiver(ServiceBusReceiveMode.PEEK_LOCK, session=None) + with patch.object(BaseHandler, "_close_handler"): + r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages.assert_called_once_with(r._handler) + + def test_receive_and_delete_does_not_drain(self): + r = _sync_receiver(ServiceBusReceiveMode.RECEIVE_AND_DELETE, session=None) + with patch.object(BaseHandler, "_close_handler"): + r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages.assert_not_called() + + def test_session_receiver_does_not_drain(self): + r = _sync_receiver(ServiceBusReceiveMode.PEEK_LOCK, session=MagicMock(name="session")) + with patch.object(BaseHandler, "_close_handler"): + r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages.assert_not_called() @pytest.mark.asyncio - async def test_release_failure_does_not_block_close(self): - buffer = _buffer_with((20, b"tag-20"), (21, b"tag-21")) - client = _make_async_client(ReceiverSettleMode.Second, buffer) - client.settle_messages_async.side_effect = ValueError("link detached") + async def test_peek_lock_non_session_drains_async(self): + r = _async_receiver(ServiceBusReceiveMode.PEEK_LOCK, session=None) + with patch.object(BaseHandlerAsync, "_close_handler", new=AsyncMock()): + await r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages_async.assert_awaited_once_with(r._handler) - with patch.object(AMQPClientAsync, "close_async", new=AsyncMock()) as super_close: - await client.close_async() # must not raise + @pytest.mark.asyncio + async def test_receive_and_delete_does_not_drain_async(self): + r = _async_receiver(ServiceBusReceiveMode.RECEIVE_AND_DELETE, session=None) + with patch.object(BaseHandlerAsync, "_close_handler", new=AsyncMock()): + await r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages_async.assert_not_called() - assert client.settle_messages_async.call_count == 1 - assert client._received_messages.empty() - super_close.assert_awaited_once() + @pytest.mark.asyncio + async def test_session_receiver_does_not_drain_async(self): + r = _async_receiver(ServiceBusReceiveMode.PEEK_LOCK, session=MagicMock(name="session")) + with patch.object(BaseHandlerAsync, "_close_handler", new=AsyncMock()): + await r._close_handler() + r._amqp_transport.drain_receive_link_and_release_messages_async.assert_not_called() From 08574e7fede5189d98e1dfd46a0170a179b9751b Mon Sep 17 00:00:00 2001 From: Pranjal Patel Date: Fri, 10 Jul 2026 15:26:53 -0700 Subject: [PATCH 3/4] fix(servicebus): address review feedback on drain-on-close (#42917) - Always release buffered deliveries when the link is open. The previous early-return on zero link credit also skipped releasing already-buffered messages, re-introducing the bug for prefetch (prefetch_count > 0) receivers whose credit is exhausted while _received_messages still has deliveries. Only the drain flow is now gated on outstanding credit. - Harden drain quiescence: pump with batch=outstanding so a multi-frame (fragmented) transfer is assembled within a listen() cycle, and require two consecutive idle cycles before concluding (still deadline-bounded). pyamqp does not surface the broker drain response, so queue-growth remains the signal. - Tests: add no-credit-still-releases (sync/async) and empty-buffer no-op; assert closed link settles nothing. Fixes #42917 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bef8092b-2a9f-4862-aef1-680eeb2e97f0 --- .../_transport/_pyamqp_transport.py | 34 ++++++++++++------ .../aio/_transport/_pyamqp_transport_async.py | 35 +++++++++++++------ .../unittests/test_receiver_drain_on_close.py | 17 +++++++-- 3 files changed, 61 insertions(+), 25 deletions(-) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py index 12045ea82680..a2af9655159c 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py @@ -814,18 +814,30 @@ def drain_receive_link_and_release_messages(handler: "ReceiveClient") -> None: """ # pylint: disable=protected-access link = handler._link - if link is None or link._is_closed or link.current_link_credit <= 0: - return + if link is None or link._is_closed: + return # cannot drain or settle through a closed/absent link try: - link.flow(link_credit=0, drain=True) - deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT - while time.time() < deadline: - before = handler._received_messages.qsize() - # listen() directly, not do_work(): do_work re-issues credit at 0 - # and would undo the drain. Stop on quiescence (no new transfers). - handler._connection.listen(wait=handler._socket_timeout) - if handler._received_messages.qsize() == before: - break + if link.current_link_credit > 0: + # Outstanding credit: stop the broker and pump in-flight transfers + # into the buffer before releasing them below. + outstanding = link.current_link_credit + link.flow(link_credit=0, drain=True) + deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT + idle_cycles = 0 + while time.time() < deadline: + 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; pyamqp drops the drain echo, so stop after 2 idle cycles. + handler._connection.listen(wait=handler._socket_timeout, batch=outstanding) + if handler._received_messages.qsize() == before: + idle_cycles += 1 + if idle_cycles >= 2: + break + else: + idle_cycles = 0 + # Always release buffered deliveries (incl. the credit==0 prefetch case) so + # they are not left locked until lock expiry. while not handler._received_messages.empty(): frame, _ = handler._received_messages.get_nowait() handler.settle_messages(frame[1], frame[2], "released") diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py index 01a02934add1..e2c26dfd4a27 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py @@ -312,18 +312,31 @@ async def drain_receive_link_and_release_messages_async(handler: "ReceiveClientA """ # pylint: disable=protected-access link = handler._link - if link is None or link._is_closed or link.current_link_credit <= 0: - return + if link is None or link._is_closed: + return # cannot drain or settle through a closed/absent link try: - await link.flow(link_credit=0, drain=True) - deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT - while time.time() < deadline: - before = handler._received_messages.qsize() - # listen() directly, not do_work_async(): do_work re-issues credit - # at 0 and would undo the drain. Stop on quiescence (no new transfers). - await handler._connection.listen(wait=handler._socket_timeout) - if handler._received_messages.qsize() == before: - break + if link.current_link_credit > 0: + # Outstanding credit: stop the broker and pump in-flight transfers + # into the buffer before releasing them below. + outstanding = link.current_link_credit + await link.flow(link_credit=0, drain=True) + deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT + idle_cycles = 0 + while time.time() < deadline: + 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; pyamqp drops the drain echo, so stop after + # 2 idle cycles. + await handler._connection.listen(wait=handler._socket_timeout, batch=outstanding) + if handler._received_messages.qsize() == before: + idle_cycles += 1 + if idle_cycles >= 2: + break + else: + idle_cycles = 0 + # Always release buffered deliveries (incl. the credit==0 prefetch case) so + # they are not left locked until lock expiry. while not handler._received_messages.empty(): frame, _ = handler._received_messages.get_nowait() await handler.settle_messages_async(frame[1], frame[2], "released") diff --git a/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py b/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py index dc3d94a878b3..91a699144724 100644 --- a/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py +++ b/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py @@ -105,16 +105,26 @@ def _listen(*args, **kwargs): h._link.flow.assert_called_once_with(link_credit=0, drain=True) h.settle_messages.assert_called_once_with(99, b"tag-99", "released") - def test_skipped_when_no_credit(self): + def test_no_credit_skips_drain_but_releases_buffer(self): + # A prefetch receiver can hit zero credit with deliveries still buffered: + # skip the drain flow, but the buffered messages must still be released. h = _handler(credit=0, buffer=_buffer_with((10, b"tag-10"))) PyamqpTransport.drain_receive_link_and_release_messages(h) h._link.flow.assert_not_called() + h.settle_messages.assert_called_once_with(10, b"tag-10", "released") + assert h._received_messages.empty() + + def test_no_credit_empty_buffer_is_noop(self): + h = _handler(credit=0) + PyamqpTransport.drain_receive_link_and_release_messages(h) + h._link.flow.assert_not_called() h.settle_messages.assert_not_called() def test_skipped_when_link_closed(self): h = _handler(credit=5, is_closed=True, buffer=_buffer_with((10, b"tag-10"))) PyamqpTransport.drain_receive_link_and_release_messages(h) h._link.flow.assert_not_called() + h.settle_messages.assert_not_called() # cannot settle through a closed link def test_skipped_when_link_none(self): h = _handler(credit=5) @@ -181,11 +191,12 @@ async def _listen(*args, **kwargs): h.settle_messages_async.assert_awaited_once_with(88, b"tag-88", "released") @pytest.mark.asyncio - async def test_skipped_when_no_credit(self): + async def test_no_credit_skips_drain_but_releases_buffer(self): h = _handler(credit=0, buffer=_buffer_with((20, b"tag-20")), is_async=True) await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) h._link.flow.assert_not_called() - h.settle_messages_async.assert_not_called() + h.settle_messages_async.assert_awaited_once_with(20, b"tag-20", "released") + assert h._received_messages.empty() @pytest.mark.asyncio async def test_failure_does_not_raise(self): From 0b45888b8a70a0d24b068d5d8c84cf3b7a6689e8 Mon Sep 17 00:00:00 2001 From: Pranjal Patel Date: Fri, 10 Jul 2026 15:43:25 -0700 Subject: [PATCH 4/4] fix(servicebus): bound receiver-close drain by the deadline (#42917) The drain deadline was only checked between listen() calls, so a large socket_timeout (or a continuously-available batch) could let a single blocking read overrun RECEIVE_LINK_DRAIN_TIMEOUT. Each listen() now waits at most the remaining drain budget (wait = min(socket_timeout, remaining)), so close() is bounded by the cap regardless of socket_timeout. The bounded tests now use a listen mock that actually blocks for its wait argument with a large socket_timeout and assert close() returns within the cap. Fixes #42917 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bef8092b-2a9f-4862-aef1-680eeb2e97f0 --- .../_transport/_pyamqp_transport.py | 11 ++++-- .../aio/_transport/_pyamqp_transport_async.py | 13 +++++-- .../unittests/test_receiver_drain_on_close.py | 37 ++++++++++++------- 3 files changed, 41 insertions(+), 20 deletions(-) diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py index a2af9655159c..bbf084d6f3b7 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py @@ -824,12 +824,17 @@ def drain_receive_link_and_release_messages(handler: "ReceiveClient") -> None: link.flow(link_credit=0, drain=True) deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT idle_cycles = 0 - while time.time() < deadline: + 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; pyamqp drops the drain echo, so stop after 2 idle cycles. - handler._connection.listen(wait=handler._socket_timeout, batch=outstanding) + # 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: diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py index e2c26dfd4a27..4296e3ed41fe 100644 --- a/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py +++ b/sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py @@ -322,13 +322,18 @@ async def drain_receive_link_and_release_messages_async(handler: "ReceiveClientA await link.flow(link_credit=0, drain=True) deadline = time.time() + RECEIVE_LINK_DRAIN_TIMEOUT idle_cycles = 0 - while time.time() < deadline: + 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; pyamqp drops the drain echo, so stop after - # 2 idle cycles. - await handler._connection.listen(wait=handler._socket_timeout, batch=outstanding) + # 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: diff --git a/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py b/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py index 91a699144724..650646502f37 100644 --- a/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py +++ b/sdk/servicebus/azure-servicebus/tests/unittests/test_receiver_drain_on_close.py @@ -30,6 +30,7 @@ """ import queue +import time import pytest from unittest.mock import AsyncMock, MagicMock, patch @@ -139,23 +140,28 @@ def test_failure_does_not_raise(self): # drain aborted before releasing; no exception escapes h.settle_messages.assert_not_called() - def test_bounded_when_never_quiescent(self): + def test_bounded_by_deadline_even_with_blocking_listen(self): + # A large socket_timeout must not let a single blocking read overrun the cap: + # each listen() should wait at most the remaining drain budget. buffer = queue.Queue() connection = MagicMock(name="connection") - seq = {"n": 0} def _listen(*args, **kwargs): - seq["n"] += 1 - frame = TransferFrame(handle=0, delivery_id=seq["n"], delivery_tag=bytes(str(seq["n"]), "ascii")) + time.sleep(kwargs.get("wait", 0)) # block for the requested wait + frame = TransferFrame(handle=0, delivery_id=1, delivery_tag=b"t") buffer.put((frame, MagicMock())) # never quiescent connection.listen = MagicMock(side_effect=_listen) h = _handler(credit=3, buffer=buffer, connection=connection) + h._socket_timeout = 5 # much larger than the drain cap - with patch("azure.servicebus._transport._pyamqp_transport.RECEIVE_LINK_DRAIN_TIMEOUT", 0.05): - PyamqpTransport.drain_receive_link_and_release_messages(h) # must terminate, not hang + start = time.time() + with patch("azure.servicebus._transport._pyamqp_transport.RECEIVE_LINK_DRAIN_TIMEOUT", 0.1): + PyamqpTransport.drain_receive_link_and_release_messages(h) + elapsed = time.time() - start assert h._connection.listen.call_count >= 1 + assert elapsed < 2.0 # bounded by the cap, not the 5s socket timeout class TestPyamqpTransportDrainAsync: @@ -206,23 +212,28 @@ async def test_failure_does_not_raise(self): h.settle_messages_async.assert_not_called() @pytest.mark.asyncio - async def test_bounded_when_never_quiescent(self): + async def test_bounded_by_deadline_even_with_blocking_listen(self): + import asyncio # pylint: disable=do-not-import-asyncio + buffer = queue.Queue() connection = MagicMock(name="connection") - seq = {"n": 0} async def _listen(*args, **kwargs): - seq["n"] += 1 - frame = TransferFrame(handle=0, delivery_id=seq["n"], delivery_tag=bytes(str(seq["n"]), "ascii")) - buffer.put((frame, MagicMock())) + await asyncio.sleep(kwargs.get("wait", 0)) # block for the requested wait + frame = TransferFrame(handle=0, delivery_id=1, delivery_tag=b"t") + buffer.put((frame, MagicMock())) # never quiescent connection.listen = AsyncMock(side_effect=_listen) h = _handler(credit=3, buffer=buffer, connection=connection, is_async=True) + h._socket_timeout = 5 # much larger than the drain cap - with patch("azure.servicebus.aio._transport._pyamqp_transport_async.RECEIVE_LINK_DRAIN_TIMEOUT", 0.05): - await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) # must terminate + start = time.time() + with patch("azure.servicebus.aio._transport._pyamqp_transport_async.RECEIVE_LINK_DRAIN_TIMEOUT", 0.1): + await PyamqpTransportAsync.drain_receive_link_and_release_messages_async(h) + elapsed = time.time() - start assert h._connection.listen.await_count >= 1 + assert elapsed < 2.0 # bounded by the cap, not the 5s socket timeout # --------------------------------------------------------------------------- #