[ServiceBus] Improve auto lock renewer implementation (Azure#20071)

* imrpove auto lock renewer implementation

* small fixes

* fix bug

* update renewer so that it would not be blocking the main thread

* update changelog and test

* update code to infer max worker

* detect max worker

* small fix

* update comment

* improvement max_workers value infering

* fix pylint

* revert time setting in test

* review feedback

* small fix

* PR review feedback

Author: yunhaoling

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -8,6 +8,7 @@
 
 ### Bugs Fixed
 - Fixed a bug that `azure.servicebus.aio.AutoLockRenewer` crashes on disposal if no messages have been registered (#19642).
+- Fixed a bug that `azure.servicebus.AutoLockRenewer` only supports auto lock renewal for `max_workers` amount of messages/sessions at a time (#19362).
 
 ### Other Changes
 

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

@@ -8,7 +8,7 @@
 import logging
 import threading
 import time
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
 from typing import TYPE_CHECKING
 
 from .._servicebus_receiver import ServiceBusReceiver
@@ -23,6 +23,11 @@
     Renewable = Union[ServiceBusSession, ServiceBusReceivedMessage]
     LockRenewFailureCallback = Callable[[Renewable, Optional[Exception]], None]
 
+try:
+    import queue
+except ImportError:
+    import Queue as queue  # type: ignore
+
 _log = logging.getLogger(__name__)
 
 SHORT_RENEW_OFFSET = 0.5  # Seconds that if a renew period is longer than lock duration + offset, it's "too long"
@@ -31,7 +36,7 @@
 )
 
 
-class AutoLockRenewer(object):
+class AutoLockRenewer(object):  # pylint:disable=too-many-instance-attributes
     """Auto renew locks for messages and sessions using a background thread pool.
 
     :param max_lock_renewal_duration: A time in seconds that locks registered to this renewer
@@ -74,10 +79,12 @@ def __init__(
         max_workers=None,
     ):
         # type: (float, Optional[LockRenewFailureCallback], Optional[ThreadPoolExecutor], Optional[int]) -> None
-        """Auto renew locks for messages and sessions using a background thread pool.
+        """Auto renew locks for messages and sessions using a background thread pool. It is recommended
+        setting max_worker to a large number or passing ThreadPoolExecutor of large max_workers number when
+        AutoLockRenewer is supposed to deal with multiple messages or sessions simultaneously.
 
         :param max_lock_renewal_duration: A time in seconds that locks registered to this renewer
-          should be maintained for. Default value is 300 (5 minutes).
+         should be maintained for. Default value is 300 (5 minutes).
         :type max_lock_renewal_duration: float
         :param on_lock_renew_failure: A callback may be specified to be called when the lock is lost on the renewable
          that is being registered. Default value is None (no callback).
@@ -91,23 +98,54 @@ def __init__(
         :type max_workers: Optional[int]
         """
         self._executor = executor or ThreadPoolExecutor(max_workers=max_workers)
+        # None indicates it's unknown whether the provided executor has max workers > 1
+        self._is_max_workers_greater_than_one = None if executor else (max_workers is None or max_workers > 1)
         self._shutdown = threading.Event()
-        self._sleep_time = 1
+        self._sleep_time = 0.5
         self._renew_period = 10
+        self._running_dispatcher = threading.Event()  # indicate whether the dispatcher is running
+        self._last_activity_timestamp = None  # the last timestamp when the dispatcher is active dealing with tasks
+        self._dispatcher_timeout = 5  # the idle time that dispatcher should exit if there's no activity
         self._max_lock_renewal_duration = max_lock_renewal_duration
         self._on_lock_renew_failure = on_lock_renew_failure
+        self._renew_tasks = queue.Queue()  # type: ignore
+        self._infer_max_workers_time = 1
 
     def __enter__(self):
         if self._shutdown.is_set():
             raise ServiceBusError(
                 "The AutoLockRenewer has already been shutdown. Please create a new instance for"
                 " auto lock renewing."
             )
+
+        self._init_workers()
         return self
 
     def __exit__(self, *args):
         self.close()
 
+    def _init_workers(self):
+        if not self._running_dispatcher.is_set():
+            self._infer_max_workers_greater_than_one_if_needed()
+            self._running_dispatcher.set()
+            self._executor.submit(self._dispatch_worker)
+
+    def _infer_max_workers_greater_than_one_if_needed(self):
+        # infer max_workers value if executor is passed in
+        if self._is_max_workers_greater_than_one is None:
+            max_wokers_checker = self._executor.submit(self._infer_max_workers_value_worker)
+            max_wokers_checker.result()
+
+    def _infer_max_workers_value_worker(self):
+        max_workers_checker = self._executor.submit(pow, 1, 1)
+        # This will never complete because there is only one worker thread and
+        # it is executing this function.
+        try:
+            max_workers_checker.result(timeout=self._infer_max_workers_time)
+            self._is_max_workers_greater_than_one = True
+        except FuturesTimeoutError:
+            self._is_max_workers_greater_than_one = False
+
     def _renewable(self, renewable):
         # pylint: disable=protected-access
         if self._shutdown.is_set():
@@ -120,7 +158,26 @@ def _renewable(self, renewable):
             return False
         return True
 
-    def _auto_lock_renew(
+    def _dispatch_worker(self):
+        self._last_activity_timestamp = time.time()
+        while not self._shutdown.is_set() and self._running_dispatcher.is_set():
+            while not self._renew_tasks.empty():
+                renew_task = self._renew_tasks.get()
+                if self._is_max_workers_greater_than_one:
+                    self._executor.submit(self._auto_lock_renew_task, *renew_task)
+                else:
+                    self._auto_lock_renew_task(*renew_task)
+                self._renew_tasks.task_done()
+                self._last_activity_timestamp = time.time()
+            # If there's no activity in the past self._idle_timeout seconds, exit the method
+            # This ensures the dispatching thread could exit, not blocking the main python thread
+            # the main worker thread could be started again if new tasks get registered
+            if time.time() - self._last_activity_timestamp >= self._dispatcher_timeout:
+                self._running_dispatcher.clear()
+                self._last_activity_timestamp = None
+                return
+
+    def _auto_lock_renew_task(
         self,
         receiver,
         renewable,
@@ -130,14 +187,11 @@ def _auto_lock_renew(
         renew_period_override=None,
     ):
         # pylint: disable=protected-access
-        _log.debug(
-            "Running lock auto-renew thread for %r seconds", max_lock_renewal_duration
-        )
         error = None
         clean_shutdown = False  # Only trigger the on_lock_renew_failure if halting was not expected (shutdown, etc)
         renew_period = renew_period_override or self._renew_period
         try:
-            while self._renewable(renewable):
+            if self._renewable(renewable):
                 if (utc_now() - starttime) >= datetime.timedelta(
                     seconds=max_lock_renewal_duration
                 ):
@@ -163,6 +217,18 @@ def _auto_lock_renew(
                         # 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(
+                        (
+                            receiver,
+                            renewable,
+                            starttime,
+                            max_lock_renewal_duration,
+                            on_lock_renew_failure,
+                            renew_period_override
+                        )
+                    )
             clean_shutdown = not renewable._lock_expired
         except AutoLockRenewTimeout as e:
             error = e
@@ -231,14 +297,21 @@ def register(
                 time_until_expiry.seconds * SHORT_RENEW_SCALING_FACTOR
             )
 
-        self._executor.submit(
-            self._auto_lock_renew,
-            receiver,
-            renewable,
-            starttime,
-            max_lock_renewal_duration or self._max_lock_renewal_duration,
-            on_lock_renew_failure or self._on_lock_renew_failure,
-            renew_period_override,
+        _log.debug(
+            "Running lock auto-renew for %r for %r seconds", renewable, max_lock_renewal_duration
+        )
+
+        self._init_workers()
+
+        self._renew_tasks.put(
+            (
+                receiver,
+                renewable,
+                starttime,
+                max_lock_renewal_duration or self._max_lock_renewal_duration,
+                on_lock_renew_failure or self._on_lock_renew_failure,
+                renew_period_override
+            )
         )
 
     def close(self, wait=True):
@@ -249,5 +322,6 @@ def close(self, wait=True):
 
         :rtype: None
         """
+        self._running_dispatcher.clear()
         self._shutdown.set()
         self._executor.shutdown(wait=wait)

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

@@ -18,7 +18,7 @@ def renew_message_lock(self, message):
 class MockReceivedMessage(ServiceBusReceivedMessage):
     def __init__(self, prevent_renew_lock=False, exception_on_renew_lock=False):
         self._lock_duration = 2
-
+        self._raw_amqp_message = None
         self._received_timestamp_utc = utc_now()
         self.locked_until_utc = self._received_timestamp_utc + timedelta(seconds=self._lock_duration)
         self._settled = False
@@ -27,7 +27,6 @@ def __init__(self, prevent_renew_lock=False, exception_on_renew_lock=False):
         self._prevent_renew_lock = prevent_renew_lock
         self._exception_on_renew_lock = exception_on_renew_lock
 
-
     @property
     def _lock_expired(self):
         if self.locked_until_utc and self.locked_until_utc <= utc_now():

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

@@ -7,6 +7,7 @@
 import logging
 import sys
 import os
+from concurrent.futures import ThreadPoolExecutor
 import types
 import pytest
 import time
@@ -1002,6 +1003,70 @@ def test_queue_by_queue_client_conn_str_receive_handler_with_autolockrenew(self,
             renewer.close()
             assert len(messages) == 11
 
+            renewer = AutoLockRenewer(max_workers=8)
+            with sb_client.get_queue_sender(servicebus_queue.name) as sender:
+                for i in range(10):
+                    message = ServiceBusMessage("{}".format(i))
+                    sender.send_messages(message)
+
+            with sb_client.get_queue_receiver(servicebus_queue.name,
+                                                 max_wait_time=5,
+                                                 receive_mode=ServiceBusReceiveMode.PEEK_LOCK,
+                                                 prefetch_count=10) as receiver:
+                received_msgs = receiver.receive_messages(max_message_count=10, max_wait_time=5)
+                for msg in received_msgs:
+                    renewer.register(receiver, msg, max_lock_renewal_duration=10)
+                time.sleep(10)
+
+                for msg in received_msgs:
+                    receiver.complete_message(msg)
+            assert len(received_msgs) == 10
+            renewer.close()
+
+            executor = ThreadPoolExecutor(max_workers=1)
+            renewer = AutoLockRenewer(executor=executor)
+            with sb_client.get_queue_sender(servicebus_queue.name) as sender:
+                for i in range(2):
+                    message = ServiceBusMessage("{}".format(i))
+                    sender.send_messages(message)
+
+            with sb_client.get_queue_receiver(servicebus_queue.name,
+                                                 max_wait_time=5,
+                                                 receive_mode=ServiceBusReceiveMode.PEEK_LOCK,
+                                                 prefetch_count=3) as receiver:
+                received_msgs = receiver.receive_messages(max_message_count=3, max_wait_time=5)
+                for msg in received_msgs:
+                    renewer.register(receiver, msg, max_lock_renewal_duration=10)
+                time.sleep(10)
+
+                for msg in received_msgs:
+                    receiver.complete_message(msg)
+            assert len(received_msgs) == 2
+            assert not renewer._is_max_workers_greater_than_one
+            renewer.close()
+
+            executor = ThreadPoolExecutor(max_workers=2)
+            renewer = AutoLockRenewer(executor=executor)
+            with sb_client.get_queue_sender(servicebus_queue.name) as sender:
+                for i in range(3):
+                    message = ServiceBusMessage("{}".format(i))
+                    sender.send_messages(message)
+
+            with sb_client.get_queue_receiver(servicebus_queue.name,
+                                                 max_wait_time=5,
+                                                 receive_mode=ServiceBusReceiveMode.PEEK_LOCK,
+                                                 prefetch_count=3) as receiver:
+                received_msgs = receiver.receive_messages(max_message_count=3, max_wait_time=5)
+                for msg in received_msgs:
+                    renewer.register(receiver, msg, max_lock_renewal_duration=10)
+                time.sleep(10)
+
+                for msg in received_msgs:
+                    receiver.complete_message(msg)
+            assert len(received_msgs) == 3
+            assert renewer._is_max_workers_greater_than_one
+            renewer.close()
+
     @pytest.mark.liveTest
     @pytest.mark.live_test_only
     @CachedResourceGroupPreparer(name_prefix='servicebustest')