[ServiceBus] Initial commit of auto-auto-lock-renewal. (Azure#14956)

* Initial commit of auto-auto-lock-renewal.  Adds parameter to get_*_receiver calls, and autoregisters at message receipt (and session at open). 
* Add tests for auto autolockrenewal
* Shorten runtime of long-running autolockrenew tests
* make clock skew margin slightly smaller as a lazy workaround to not break tests.
* add changelog note for auto-autolockrenewer
* add changelog stanza for max_lock_renew_duration change.
* Make more clear error for when non-renewable messages (e.g. ReceiveAndDelete) are registered to autorenewer, and add validation for this scenario into tests.
* Add check during receiver-getting to prevent attempting to autolockrenew a ReceiveAndDelete receiver (which would not work properly as those messages can't be renewed) and unit tests to validate this.
* Normalize this failure mode to ValueError as technically the Message type is correct, and add a changelog note documenting this.

Co-authored-by: Adam Ling (MSFT) <adam_ling@outlook.com>

Author: KieranBrantnerMagee

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -10,6 +10,7 @@
   - `ServiceBusSession`: `get_state`, `set_state` and `renew_lock`
 * `azure.servicebus.exceptions.ServiceBusError` now inherits from `azure.core.exceptions.AzureError`.
 * Added a `parse_connection_string` method which parses a connection string into a properties bag containing its component parts
+* Add support for `auto_lock_renewer` parameter on `get_queue_receiver` and `get_subscription_receiver` calls to allow auto-registration of messages and sessions for auto-renewal.
 
 **Breaking Changes**
 
@@ -55,6 +56,8 @@ now raise more concrete exception other than `MessageSettleFailed` and `ServiceB
 * Sub-client (`ServiceBusSender` and `ServiceBusReceiver`) `from_connection_string` initializers have been made internal until needed.  Clients should be initialized from root `ServiceBusClient`.
 * `ServiceBusMessage.label` has been renamed to `ServiceBusMessage.subject`.
 * `ServiceBusMessage.amqp_annotated_message` has had its type renamed from `AMQPMessage` to `AMQPAnnotatedMessage`
+* `AutoLockRenewer` `timeout` parameter is renamed to `max_lock_renew_duration`
+* Attempting to autorenew a non-renewable message, such as one received in `ReceiveAndDelete` mode, or configure auto-autorenewal on a `ReceiveAndDelete` receiver, will raise a `ValueError`.
 
 **BugFixes**
 

sdk/servicebus/azure-servicebus/README.md

@@ -358,7 +358,8 @@ session_id = os.environ['SERVICE_BUS_SESSION_ID']
 renewer = AutoLockRenewer()
 with ServiceBusClient.from_connection_string(connstr) as client:
     with client.get_queue_receiver(session_queue_name, session_id=session_id) as receiver:
-        renewer.register(receiver, receiver.session, timeout=300) # Timeout for how long to maintain the lock for, in seconds.
+        renewer.register(receiver, receiver.session, max_lock_renewal_duration=300) # Duration for how long to maintain the lock for, in seconds.
+
         for msg in receiver.receive_messages():
             # Do your application logic here
             receiver.complete_message(msg)

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

@@ -15,20 +15,29 @@
 from .._servicebus_session import ServiceBusSession
 from .message import ServiceBusReceivedMessage
 from ..exceptions import AutoLockRenewFailed, AutoLockRenewTimeout, ServiceBusError
-from .utils import renewable_start_time, utc_now
+from .utils import get_renewable_start_time, utc_now, get_renewable_lock_duration
 
 if TYPE_CHECKING:
-    from typing import Callable, Union, Optional, Awaitable
-    LockRenewFailureCallback = Callable[[Union[ServiceBusSession, ServiceBusReceivedMessage],
-                                         Optional[Exception]], None]
+    from typing import Callable, Union, Optional
     Renewable = Union[ServiceBusSession, ServiceBusReceivedMessage]
+    LockRenewFailureCallback = Callable[[Renewable,
+                                         Optional[Exception]], None]
 
 _log = logging.getLogger(__name__)
 
+SHORT_RENEW_OFFSET = .5 # Seconds that if a renew period is longer than lock duration + offset, it's "too long"
+SHORT_RENEW_SCALING_FACTOR = .75 # In this situation we need a "Short renew" and should scale by this factor.
+
 
 class AutoLockRenewer(object):
     """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
+     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).
+    :type on_lock_renew_failure: Optional[LockRenewFailureCallback]
     :param executor: A user-specified thread pool. This cannot be combined with
      setting `max_workers`.
     :type executor: ~concurrent.futures.ThreadPoolExecutor
@@ -55,10 +64,16 @@ class AutoLockRenewer(object):
 
     """
 
-    def __init__(self, executor=None, max_workers=None):
-        # type: (ThreadPoolExecutor, int) -> None
+    def __init__(self, max_lock_renewal_duration=300, on_lock_renew_failure=None, executor=None, max_workers=None):
+        # type: (float, Optional[LockRenewFailureCallback], ThreadPoolExecutor, int) -> None
         """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
+          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).
+        :type on_lock_renew_failure: Optional[LockRenewFailureCallback]
         :param executor: A user-specified thread pool. This cannot be combined with
          setting `max_workers`.
         :type executor: ~concurrent.futures.ThreadPoolExecutor
@@ -71,6 +86,8 @@ def __init__(self, executor=None, max_workers=None):
         self._shutdown = threading.Event()
         self._sleep_time = 1
         self._renew_period = 10
+        self._max_lock_renewal_duration = max_lock_renewal_duration
+        self._on_lock_renew_failure = on_lock_renew_failure
 
     def __enter__(self):
         if self._shutdown.is_set():
@@ -93,18 +110,26 @@ def _renewable(self, renewable):
             return False
         return True
 
-    def _auto_lock_renew(self, receiver, renewable, starttime, timeout, on_lock_renew_failure=None):
+    def _auto_lock_renew(self,
+                         receiver,
+                         renewable,
+                         starttime,
+                         max_lock_renewal_duration,
+                         on_lock_renew_failure=None,
+                         renew_period_override=None):
         # pylint: disable=protected-access
-        _log.debug("Running lock auto-renew thread for %r seconds", timeout)
+        _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 (utc_now() - starttime) >= datetime.timedelta(seconds=timeout):
-                    _log.debug("Reached auto lock renew timeout - letting lock expire.")
-                    raise AutoLockRenewTimeout("Auto-renew period ({} seconds) elapsed.".format(timeout))
-                if (renewable.locked_until_utc - utc_now()) <= datetime.timedelta(seconds=self._renew_period):
-                    _log.debug("%r seconds or less until lock expires - auto renewing.", self._renew_period)
+                if (utc_now() - starttime) >= datetime.timedelta(seconds=max_lock_renewal_duration):
+                    _log.debug("Reached max auto lock renew duration - letting lock expire.")
+                    raise AutoLockRenewTimeout("Auto-renew period ({} seconds) elapsed.".format(
+                        max_lock_renewal_duration))
+                if (renewable.locked_until_utc - utc_now()) <= datetime.timedelta(seconds=renew_period):
+                    _log.debug("%r seconds or less until lock expires - auto renewing.", renew_period)
                     try:
                         # Renewable is a session
                         renewable.renew_lock()  # type: ignore
@@ -127,17 +152,18 @@ def _auto_lock_renew(self, receiver, renewable, starttime, timeout, on_lock_rene
             if on_lock_renew_failure and not clean_shutdown:
                 on_lock_renew_failure(renewable, error)
 
-    def register(self, receiver, renewable, timeout=300, on_lock_renew_failure=None):
-        # type: (ServiceBusReceiver, Renewable, float, Optional[LockRenewFailureCallback]) -> None
+    def register(self, receiver, renewable, max_lock_renewal_duration=None, on_lock_renew_failure=None):
+        # type: (ServiceBusReceiver, Renewable, Optional[float], Optional[LockRenewFailureCallback]) -> None
         """Register a renewable entity for automatic lock renewal.
 
         :param receiver: The ServiceBusReceiver instance that is associated with the message or the session to
          be auto-lock-renewed.
         :type receiver: ~azure.servicebus.ServiceBusReceiver
         :param renewable: A locked entity that needs to be renewed.
         :type renewable: Union[~azure.servicebus.ServiceBusReceivedMessage, ~azure.servicebus.ServiceBusSession]
-        :param timeout: A time in seconds that the lock should be maintained for. Default value is 300 (5 minutes).
-        :type timeout: float
+        :param max_lock_renewal_duration: A time in seconds that the lock should be maintained for.
+          Default value is 300 (5 minutes).
+        :type max_lock_renewal_duration: Optional[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).
         :type on_lock_renew_failure: Optional[LockRenewFailureCallback]
@@ -151,8 +177,27 @@ def register(self, receiver, renewable, timeout=300, on_lock_renew_failure=None)
         if self._shutdown.is_set():
             raise ServiceBusError("The AutoLockRenewer has already been shutdown. Please create a new instance for"
                                   " auto lock renewing.")
-        starttime = renewable_start_time(renewable)
-        self._executor.submit(self._auto_lock_renew, receiver, renewable, starttime, timeout, on_lock_renew_failure)
+        if renewable.locked_until_utc is None:
+            raise ValueError("Only azure.servicebus.ServiceBusReceivedMessage objects in PeekLock receive mode may"
+                            "be lock-renewed.  (E.g. only messages received via receive() or the receiver iterator,"
+                            "not using ReceiveAndDelete receive mode, and not returned from Peek)")
+
+        starttime = get_renewable_start_time(renewable)
+
+        # This is a heuristic to compensate if it appears the user has a lock duration less than our base renew period
+        time_until_expiry = get_renewable_lock_duration(renewable)
+        renew_period_override = None
+        # Default is 10 seconds, but let's leave ourselves a small margin of error because clock skew is a real problem
+        if time_until_expiry <= datetime.timedelta(seconds=self._renew_period + SHORT_RENEW_OFFSET):
+            renew_period_override = 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)
 
     def close(self, wait=True):
         """Cease autorenewal by shutting down the thread pool to clean up any remaining lock renewal threads.

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

@@ -74,6 +74,13 @@ def _populate_attributes(self, **kwargs):
             raise ValueError("The max_wait_time must be greater than 0.")
         self._max_wait_time = max_wait_time
 
+        self._auto_lock_renewer = kwargs.get("auto_lock_renewer", None)
+        if self._auto_lock_renewer \
+                and self._receive_mode == ReceiveMode.ReceiveAndDelete \
+                and self._session_id is None:
+            raise ValueError("Messages received in ReceiveAndDelete receive mode cannot have their locks removed "
+                             "as they have been deleted, providing an AutoLockRenewer in this mode is invalid.")
+
     def _build_message(self, received, message_type=ServiceBusReceivedMessage):
         message = message_type(message=received, receive_mode=self._receive_mode, receiver=self)
         self._last_received_sequenced_number = message.sequence_number

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

@@ -9,7 +9,7 @@
 import logging
 import functools
 import platform
-from typing import Optional, Dict, Tuple
+from typing import Optional, Dict, Tuple, Union, TYPE_CHECKING
 from msrest.serialization import UTC
 
 try:
@@ -29,6 +29,10 @@
     USER_AGENT_PREFIX,
 )
 
+if TYPE_CHECKING:
+    from .message import ServiceBusReceivedMessage
+    from .._servicebus_session import BaseSession
+
 _log = logging.getLogger(__name__)
 
 
@@ -81,15 +85,26 @@ def create_properties(user_agent=None):
     return properties
 
 
-def renewable_start_time(renewable):
+def get_renewable_start_time(renewable):
     try:
         return renewable._received_timestamp_utc  # pylint: disable=protected-access
     except AttributeError:
         pass
     try:
         return renewable._session_start  # pylint: disable=protected-access
     except AttributeError:
-        raise TypeError("Registered object is not renewable.")
+        raise TypeError("Registered object is not renewable, renewable must be" +
+                        "a ServiceBusReceivedMessage or a ServiceBusSession from a sessionful ServiceBusReceiver.")
+
+
+def get_renewable_lock_duration(renewable):
+    # type: (Union[ServiceBusReceivedMessage, BaseSession]) -> datetime.timedelta
+    # pylint: disable=protected-access
+    try:
+        return max(renewable.locked_until_utc - utc_now(), datetime.timedelta(seconds=0))
+    except AttributeError:
+        raise TypeError("Registered object is not renewable, renewable must be" +
+                        "a ServiceBusReceivedMessage or a ServiceBusSession from a sessionful ServiceBusReceiver.")
 
 
 def create_authentication(client):

sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_client.py

@@ -20,6 +20,7 @@
 
 if TYPE_CHECKING:
     from azure.core.credentials import TokenCredential
+    from ._common.auto_lock_renewer import AutoLockRenewer
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -223,6 +224,9 @@ def get_queue_receiver(self, queue_name, **kwargs):
         :paramtype receive_mode: ~azure.servicebus.ReceiveMode
         :keyword Optional[float] max_wait_time: The timeout in seconds between received messages after which the
          receiver will automatically stop receiving. The default value is None, meaning no timeout.
+        :keyword Optional[AutoLockRenewer] auto_lock_renewer: An AutoLockRenewer can be provided such that messages are
+         automatically registered on receipt.  If the receiver is a session receiver, it will apply to the session
+         instead.
         :keyword int prefetch_count: The maximum number of messages to cache with each request to the service.
          This setting is only for advanced performance tuning. Increasing this value will improve message throughput
          performance but increase the chance that messages will expire while they are cached if they're not
@@ -329,6 +333,9 @@ def get_subscription_receiver(self, topic_name, subscription_name, **kwargs):
         :paramtype receive_mode: ~azure.servicebus.ReceiveMode
         :keyword Optional[float] max_wait_time: The timeout in seconds between received messages after which the
          receiver will automatically stop receiving. The default value is None, meaning no timeout.
+        :keyword Optional[AutoLockRenewer] auto_lock_renewer: An AutoLockRenewer can be provided such that messages are
+         automatically registered on receipt.  If the receiver is a session receiver, it will apply to the session
+         instead.
         :keyword int prefetch_count: The maximum number of messages to cache with each request to the service.
          This setting is only for advanced performance tuning. Increasing this value will improve message throughput
          performance but increase the chance that messages will expire while they are cached if they're not

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

@@ -49,6 +49,7 @@
 if TYPE_CHECKING:
     import datetime
     from azure.core.credentials import TokenCredential
+    from ._common.auto_lock_renewer import AutoLockRenewer
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -94,6 +95,9 @@ class ServiceBusReceiver(BaseHandler, ReceiverMixin):  # pylint: disable=too-man
      keys: `'proxy_hostname'` (str value) and `'proxy_port'` (int value).
      Additionally the following keys may also be present: `'username', 'password'`.
     :keyword str user_agent: If specified, this will be added in front of the built-in user agent string.
+    :keyword Optional[AutoLockRenewer] auto_lock_renewer: An AutoLockRenewer can be provided such that messages are
+     automatically registered on receipt.  If the receiver is a session receiver, it will apply to the session
+     instead.
     :keyword int prefetch_count: The maximum number of messages to cache with each request to the service.
      This setting is only for advanced performance tuning. Increasing this value will improve message throughput
      performance but increase the chance that messages will expire while they are cached if they're not
@@ -187,6 +191,8 @@ def _iter_next(self):
             self._message_iter = self._handler.receive_messages_iter()
         uamqp_message = next(self._message_iter)
         message = self._build_message(uamqp_message)
+        if self._auto_lock_renewer and not self._session:
+            self._auto_lock_renewer.register(self, message)
         return message
 
     def _create_handler(self, auth):
@@ -227,6 +233,9 @@ def _open(self):
             self.close()
             raise
 
+        if self._auto_lock_renewer and self._session:
+            self._auto_lock_renewer.register(self, self.session)
+
     def _receive(self, max_message_count=None, timeout=None):
         # type: (Optional[int], Optional[float]) -> List[ServiceBusReceivedMessage]
         # pylint: disable=protected-access
@@ -510,12 +519,16 @@ def receive_messages(self, max_message_count=None, max_wait_time=None):
         if max_wait_time is not None and max_wait_time <= 0:
             raise ValueError("The max_wait_time must be greater than 0.")
         self._check_live()
-        return self._do_retryable_operation(
+        messages = self._do_retryable_operation(
             self._receive,
             max_message_count=max_message_count,
             timeout=max_wait_time,
             operation_requires_timeout=True
         )
+        if self._auto_lock_renewer and not self._session:
+            for message in messages:
+                self._auto_lock_renewer.register(self, message)
+        return messages
 
     def receive_deferred_messages(self, sequence_numbers, **kwargs):
         # type: (Union[int,List[int]], Any) -> List[ServiceBusReceivedMessage]
@@ -567,6 +580,9 @@ def receive_deferred_messages(self, sequence_numbers, **kwargs):
             handler,
             timeout=timeout
         )
+        if self._auto_lock_renewer and not self._session:
+            for message in messages:
+                self._auto_lock_renewer.register(self, message)
         return messages
 
     def peek_messages(self, max_message_count=1, **kwargs):