[EventHub/ServiceBus] lock pending mgmt op (#39633)

* [SB] Msg IDs on mgmt/cbs ops should be unique

* copy over sb pyamqp to eh

* [ServiceBus] fix autolockrenew sleep/lock pending mgmt ops

* remove alr sleep

* add tests

* add sync alr 300 msg renew test

* remove disable local auth for tests ci run

* add alr renew 300 msgs test for sync

* revert alr sample

* fix stress bicep policy violations

* update changelog

Author: swathipil

sdk/eventhub/azure-eventhub/CHANGELOG.md

@@ -1,6 +1,6 @@
 # Release History
 
-## 5.14.0 (2025-02-11)
+## 5.14.0 (2025-02-13)
 
 ### Features Added
 
@@ -12,6 +12,7 @@
 - Fixed a bug where pyAMQP was doubly retrying, causing latency on reconnect. ([#39037](https://github.com/Azure/azure-sdk-for-python/pull/39037))
 - Fix to handle large messages being sent twice due to incoming flow frames triggering a resend. ([#38067](https://github.com/Azure/azure-sdk-for-python/pull/38067))
 - Missing await in sender async on pyAMQP. ([#39182](https://github.com/Azure/azure-sdk-for-python/pull/39182))
+- Fixed a bug where message IDs in management operation requests were not unique.
 
 ### Other Changes
 

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/management_link.py

@@ -10,6 +10,7 @@
 from functools import partial
 from collections import namedtuple
 from typing import Optional, Union
+from threading import Lock
 
 from .sender import SenderLink
 from .receiver import ReceiverLink
@@ -40,6 +41,7 @@ class ManagementLink(object):  # pylint:disable=too-many-instance-attributes
     def __init__(self, session, endpoint, **kwargs):
         self.state = ManagementLinkState.IDLE
         self._pending_operations = []
+        self.lock = Lock()
         self._session = session
         self._network_trace_params = kwargs.get("network_trace_params")
         self._request_link: SenderLink = session.create_sender_link(
@@ -157,6 +159,7 @@ def _on_message_received(self, _, message):
             to_remove_operation.on_execute_operation_complete(
                 mgmt_result, status_code, status_description, message, response_detail.get(b"error-condition")
             )
+        with self.lock:
             self._pending_operations.remove(to_remove_operation)
 
     def _on_send_complete(self, message_delivery, reason, state):  # todo: reason is never used, should check spec

sdk/eventhub/azure-eventhub/stress/stress-test-resources.bicep

@@ -27,7 +27,9 @@ resource eventHubsNamespace 'Microsoft.EventHub/Namespaces@2017-04-01' = {
     name: 'Standard'
     tier: 'Standard'
   }
-  properties: {}
+  properties: {
+    disableLocalAuth: true
+  }
 }
 
 resource eventHubsNamespace_eventHubName 'Microsoft.EventHub/namespaces/eventhubs@2017-04-01' = {
@@ -36,6 +38,7 @@ resource eventHubsNamespace_eventHubName 'Microsoft.EventHub/namespaces/eventhub
   properties: {
     messageRetentionInDays: 5
     partitionCount: 32
+    disableLocalAuth: true
   }
   dependsOn: [
     eventHubsNamespace
@@ -66,6 +69,7 @@ resource storageAccount 'Microsoft.Storage/storageAccounts@2019-06-01' = {
   kind: 'StorageV2'
   properties: {
     accessTier: 'Hot'
+    allowSharedKeyAccess: false
   }
 }
 

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -1,6 +1,6 @@
 # Release History
 
-## 7.14.0 (2025-02-11)
+## 7.14.0 (2025-02-13)
 
 ### Features Added
 
@@ -15,6 +15,8 @@
 - Fixed a bug where pyAMQP was doubly retrying, causing latency on reconnect. ([#39037](https://github.com/Azure/azure-sdk-for-python/pull/39037))
 - Fix to handle large messages being sent twice due to incoming flow frames triggering a resend. ([#38067](https://github.com/Azure/azure-sdk-for-python/pull/38067))
 - Missing await in sender async on pyAMQP. ([#39182](https://github.com/Azure/azure-sdk-for-python/pull/39182))
+- Improved AutoLockRenewer to renew locks for more registered messages. ([#37340](https://github.com/Azure/azure-sdk-for-python/issues/37340))
+- Fixed a bug where message IDs in management operation requests were not unique.
 
 ### Other Changes
 

sdk/servicebus/azure-servicebus/azure/servicebus/_common/auto_lock_renewer.py

@@ -200,7 +200,6 @@ def _auto_lock_renew_task(
                     except AttributeError:
                         # Renewable is a message
                         receiver.renew_message_lock(renewable)  # type: ignore
-                time.sleep(self._sleep_time)
                 # enqueue a new task, keeping renewing the renewable
                 if self._renewable(renewable):
                     self._renew_tasks.put(

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/management_link.py

@@ -10,6 +10,7 @@
 from functools import partial
 from collections import namedtuple
 from typing import Optional, Union
+from threading import Lock
 
 from .sender import SenderLink
 from .receiver import ReceiverLink
@@ -40,6 +41,7 @@ class ManagementLink(object):  # pylint:disable=too-many-instance-attributes
     def __init__(self, session, endpoint, **kwargs):
         self.state = ManagementLinkState.IDLE
         self._pending_operations = []
+        self.lock = Lock()
         self._session = session
         self._network_trace_params = kwargs.get("network_trace_params")
         self._request_link: SenderLink = session.create_sender_link(
@@ -157,6 +159,7 @@ def _on_message_received(self, _, message):
             to_remove_operation.on_execute_operation_complete(
                 mgmt_result, status_code, status_description, message, response_detail.get(b"error-condition")
             )
+        with self.lock:
             self._pending_operations.remove(to_remove_operation)
 
     def _on_send_complete(self, message_delivery, reason, state):  # todo: reason is never used, should check spec

sdk/servicebus/azure-servicebus/samples/sync_samples/auto_lock_renew.py

@@ -138,4 +138,4 @@ def on_lock_renew_failure_callback(renewable, error):
 
 renew_lock_on_message_received_from_non_sessionful_entity()
 renew_lock_on_session_of_the_sessionful_entity()
-renew_lock_with_lock_renewal_failure_callback()
+renew_lock_with_lock_renewal_failure_callback()

sdk/servicebus/azure-servicebus/stress/stress-test-resources.bicep

@@ -17,6 +17,7 @@ resource servicebus 'Microsoft.ServiceBus/namespaces@2018-01-01-preview' = {
   }
   properties: {
     zoneRedundant: false
+    disableLocalAuth: true
   }
 }
 
@@ -27,6 +28,9 @@ resource servicebusPremium 'Microsoft.ServiceBus/namespaces@2018-01-01-preview'
     name: 'Premium'
     tier: 'Premium'
   }
+  properties: {
+    disableLocalAuth: true
+  }
 }
 
 

sdk/servicebus/azure-servicebus/tests/test_queues.py

@@ -1678,6 +1678,51 @@ def test_queue_by_queue_client_conn_str_receive_handler_with_auto_autolockrenew(
             renewer.close()
             assert len(messages) == 11
 
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest")
+    @CachedServiceBusNamespacePreparer(name_prefix="servicebustest")
+    @ServiceBusQueuePreparer(
+        name_prefix="servicebustest", dead_lettering_on_message_expiration=True, lock_duration="PT55S"
+    )
+    @pytest.mark.parametrize("uamqp_transport", uamqp_transport_params, ids=uamqp_transport_ids)
+    @ArgPasser()
+    def test_queue_receive_handler_with_autolockrenew_many_msgs(
+        self, uamqp_transport, *, servicebus_namespace=None, servicebus_queue=None, **kwargs
+    ):
+        # number of messages to receive and renew locks for
+        max_msg_count = 300
+
+        # only tested on sync since async can concurrently send renew-lock requests without waiting for the transfer frame from the service
+        if sys.platform.startswith("darwin"):
+            pytest.skip("Skipping since ALR renewal is slower, so message locks are expiring. Potentially due to lower # of available threads/workers.")
+
+        fully_qualified_namespace = f"{servicebus_namespace.name}{SERVICEBUS_ENDPOINT_SUFFIX}"
+        credential = get_credential()
+        with ServiceBusClient(
+            fully_qualified_namespace, credential, logging_enable=True, uamqp_transport=uamqp_transport
+        ) as sb_client:
+            with sb_client.get_queue_sender(servicebus_queue.name) as sender:
+                msgs_to_send = [ServiceBusMessage("message: {}".format(i)) for i in range(400)]
+                sender.send_messages(msgs_to_send)
+
+            # Can also be called via "with AutoLockRenewer() as renewer" to automate shutdown.
+            renewer = AutoLockRenewer(max_lock_renewal_duration=60*5)
+            with sb_client.get_queue_receiver(queue_name=servicebus_queue.name, auto_lock_renewer=renewer) as receiver:
+                received_msgs = receiver.receive_messages(max_message_count=max_msg_count)
+                # At least 300 messages should be renewed
+                # TODO: This should be bumped up once sync alr perf is improved
+                while len(received_msgs) < max_msg_count:
+                    count = max_msg_count - len(received_msgs)
+                    received_msgs.extend(receiver.receive_messages(max_message_count=count))
+
+                assert len(received_msgs) == max_msg_count, "Did not receive all messages"
+                time.sleep(100)
+                for msg in received_msgs:
+                    receiver.complete_message(msg)  # Settling the message deregisters it from the AutoLockRenewer
+
+            renewer.close()
+
     @pytest.mark.liveTest
     @pytest.mark.live_test_only
     @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest")