[EventHubs] add fixed backoff retry mode (Azure#21884)

* sync retry mode impl

* async impl

* add test

* pylint

* lint

* adams comments

* lint

* adams comments

* fix test

Author: swathipil

sdk/eventhub/azure-eventhub/CHANGELOG.md

@@ -4,6 +4,10 @@
 
 ### Features Added
 
+- Added support for fixed (linear) retry backoff:
+  - Sync/async `EventHubProducerClient` and `EventHubConsumerClient` constructors and `from_connection_string` take `retry_mode` as a keyword argument.
+  - `RetryMode` enum has been added to `azure.eventhub`, with values `FIXED` and `EXPONENTIAL`.
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/eventhub/azure-eventhub/azure/eventhub/__init__.py

@@ -18,6 +18,7 @@
     parse_connection_string,
     EventHubConnectionStringProperties
 )
+from ._retry import RetryMode
 
 TransportType = constants.TransportType
 
@@ -33,5 +34,6 @@
     "LoadBalancingStrategy",
     "PartitionContext",
     "parse_connection_string",
-    "EventHubConnectionStringProperties"
+    "EventHubConnectionStringProperties",
+    "RetryMode"
 ]

sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py

@@ -20,12 +20,17 @@
 
 from uamqp import AMQPClient, Message, authentication, constants, errors, compat, utils
 import six
-from azure.core.credentials import AccessToken, AzureSasCredential, AzureNamedKeyCredential
+from azure.core.credentials import (
+    AccessToken,
+    AzureSasCredential,
+    AzureNamedKeyCredential,
+)
 from azure.core.utils import parse_connection_string as core_parse_connection_string
 
 
 from .exceptions import _handle_exception, ClientClosedError, ConnectError
 from ._configuration import Configuration
+from ._retry import RetryMode
 from ._utils import utc_from_timestamp, parse_sas_credential
 from ._connection_manager import get_connection_manager
 from ._constants import (
@@ -34,7 +39,7 @@
     MGMT_OPERATION,
     MGMT_PARTITION_OPERATION,
     MGMT_STATUS_CODE,
-    MGMT_STATUS_DESC
+    MGMT_STATUS_DESC,
 )
 
 if TYPE_CHECKING:
@@ -52,9 +57,11 @@ def _parse_conn_str(conn_str, **kwargs):
     entity_path = None  # type: Optional[str]
     shared_access_signature = None  # type: Optional[str]
     shared_access_signature_expiry = None
-    eventhub_name = kwargs.pop("eventhub_name", None) # type: Optional[str]
-    check_case = kwargs.pop("check_case", False) # type: bool
-    conn_settings = core_parse_connection_string(conn_str, case_sensitive_keys=check_case)
+    eventhub_name = kwargs.pop("eventhub_name", None)  # type: Optional[str]
+    check_case = kwargs.pop("check_case", False)  # type: bool
+    conn_settings = core_parse_connection_string(
+        conn_str, case_sensitive_keys=check_case
+    )
     if check_case:
         shared_access_key = conn_settings.get("SharedAccessKey")
         shared_access_key_name = conn_settings.get("SharedAccessKeyName")
@@ -79,7 +86,7 @@ def _parse_conn_str(conn_str, **kwargs):
         try:
             # Expiry can be stored in the "se=<timestamp>" clause of the token. ('&'-separated key-value pairs)
             shared_access_signature_expiry = int(
-                shared_access_signature.split("se=")[1].split("&")[0]   # type: ignore
+                shared_access_signature.split("se=")[1].split("&")[0]  # type: ignore
             )
         except (
             IndexError,
@@ -117,12 +124,14 @@ def _parse_conn_str(conn_str, **kwargs):
             "At least one of the SharedAccessKey or SharedAccessSignature must be present."
         )
 
-    return (host,
-            str(shared_access_key_name) if shared_access_key_name else None,
-            str(shared_access_key) if shared_access_key else None,
-            entity,
-            str(shared_access_signature) if shared_access_signature else None,
-            shared_access_signature_expiry)
+    return (
+        host,
+        str(shared_access_key_name) if shared_access_key_name else None,
+        str(shared_access_key) if shared_access_key else None,
+        entity,
+        str(shared_access_signature) if shared_access_signature else None,
+        shared_access_signature_expiry,
+    )
 
 
 def _generate_sas_token(uri, policy, key, expiry=None):
@@ -154,6 +163,14 @@ def _build_uri(address, entity):
     return address
 
 
+def _get_backoff_time(retry_mode, backoff_factor, backoff_max, retried_times):
+    if retry_mode == RetryMode.FIXED:
+        backoff_value = backoff_factor
+    else:
+        backoff_value = backoff_factor * (2 ** retried_times)
+    return min(backoff_max, backoff_value)
+
+
 class EventHubSharedKeyCredential(object):
     """The shared access key credential used for authentication.
 
@@ -200,6 +217,7 @@ class EventHubSASTokenCredential(object):
     :param str token: The shared access token string
     :param int expiry: The epoch timestamp
     """
+
     def __init__(self, token, expiry):
         # type: (str, int) -> None
         """
@@ -225,6 +243,7 @@ class EventhubAzureSasTokenCredential(object):
     :param azure_sas_credential: The credential to be used for authentication.
     :type azure_sas_credential: ~azure.core.credentials.AzureSasCredential
     """
+
     def __init__(self, azure_sas_credential):
         # type: (AzureSasCredential) -> None
         """The shared access token credential used for authentication
@@ -257,9 +276,9 @@ def __init__(self, fully_qualified_namespace, eventhub_name, credential, **kwarg
         if isinstance(credential, AzureSasCredential):
             self._credential = EventhubAzureSasTokenCredential(credential)
         elif isinstance(credential, AzureNamedKeyCredential):
-            self._credential = EventhubAzureNamedKeyTokenCredential(credential) # type: ignore
+            self._credential = EventhubAzureNamedKeyTokenCredential(credential)  # type: ignore
         else:
-            self._credential = credential #type: ignore
+            self._credential = credential  # type: ignore
         self._keep_alive = kwargs.get("keep_alive", 30)
         self._auto_reconnect = kwargs.get("auto_reconnect", True)
         self._mgmt_target = "amqps://{}/{}".format(
@@ -274,7 +293,9 @@ def __init__(self, fully_qualified_namespace, eventhub_name, credential, **kwarg
     @staticmethod
     def _from_connection_string(conn_str, **kwargs):
         # type: (str, Any) -> Dict[str, Any]
-        host, policy, key, entity, token, token_expiry = _parse_conn_str(conn_str, **kwargs)
+        host, policy, key, entity, token, token_expiry = _parse_conn_str(
+            conn_str, **kwargs
+        )
         kwargs["fully_qualified_namespace"] = host
         kwargs["eventhub_name"] = entity
         if token and token_expiry:
@@ -291,7 +312,7 @@ def _create_auth(self):
         """
         try:
             # ignore mypy's warning because token_type is Optional
-            token_type = self._credential.token_type    # type: ignore
+            token_type = self._credential.token_type  # type: ignore
         except AttributeError:
             token_type = b"jwt"
         if token_type == b"servicebus.windows.net:sastoken":
@@ -305,7 +326,7 @@ def _create_auth(self):
                 transport_type=self._config.transport_type,
                 custom_endpoint_hostname=self._config.custom_endpoint_hostname,
                 port=self._config.connection_port,
-                verify=self._config.connection_verify
+                verify=self._config.connection_verify,
             )
             auth.update_token()
             return auth
@@ -319,7 +340,7 @@ def _create_auth(self):
             transport_type=self._config.transport_type,
             custom_endpoint_hostname=self._config.custom_endpoint_hostname,
             port=self._config.connection_port,
-            verify=self._config.connection_verify
+            verify=self._config.connection_verify,
         )
 
     def _close_connection(self):
@@ -331,7 +352,12 @@ def _backoff(
     ):
         # type: (int, Exception, Optional[int], Optional[str]) -> None
         entity_name = entity_name or self._container_id
-        backoff = self._config.backoff_factor * 2 ** retried_times
+        backoff = _get_backoff_time(
+            self._config.retry_mode,
+            self._config.backoff_factor,
+            self._config.backoff_max,
+            retried_times,
+        )
         if backoff <= self._config.backoff_max and (
             timeout_time is None or time.time() + backoff <= timeout_time
         ):  # pylint:disable=no-else-return
@@ -360,7 +386,7 @@ def _management_request(self, mgmt_msg, op_type):
                 self._mgmt_target, auth=mgmt_auth, debug=self._config.network_tracing
             )
             try:
-                conn = self._conn_manager.get_connection(
+                conn = self._conn_manager.get_connection(  # pylint:disable=assignment-from-none
                     self._address.hostname, mgmt_auth
                 )
                 mgmt_client.open(connection=conn)
@@ -373,29 +399,28 @@ def _management_request(self, mgmt_msg, op_type):
                     description_fields=MGMT_STATUS_DESC,
                 )
                 status_code = int(response.application_properties[MGMT_STATUS_CODE])
-                description = response.application_properties.get(MGMT_STATUS_DESC)  # type: Optional[Union[str, bytes]]
+                description = response.application_properties.get(
+                    MGMT_STATUS_DESC
+                )  # type: Optional[Union[str, bytes]]
                 if description and isinstance(description, six.binary_type):
-                    description = description.decode('utf-8')
+                    description = description.decode("utf-8")
                 if status_code < 400:
                     return response
                 if status_code in [401]:
                     raise errors.AuthenticationException(
                         "Management authentication failed. Status code: {}, Description: {!r}".format(
-                            status_code,
-                            description
+                            status_code, description
                         )
                     )
                 if status_code in [404]:
                     raise ConnectError(
                         "Management connection failed. Status code: {}, Description: {!r}".format(
-                            status_code,
-                            description
+                            status_code, description
                         )
                     )
                 raise errors.AMQPConnectionError(
                     "Management request error. Status code: {}, Description: {!r}".format(
-                        status_code,
-                        description
+                        status_code, description
                     )
                 )
             except Exception as exception:  # pylint: disable=broad-except
@@ -491,9 +516,7 @@ def _check_closed(self):
             )
 
     def _open(self):
-        """Open the EventHubConsumer/EventHubProducer using the supplied connection.
-
-        """
+        """Open the EventHubConsumer/EventHubProducer using the supplied connection."""
         # pylint: disable=protected-access
         if not self.running:
             if self._handler:

sdk/eventhub/azure-eventhub/azure/eventhub/_common.py

@@ -182,6 +182,7 @@ def _from_message(cls, message, raw_amqp_message=None):
         """
         event_data = cls(body="")
         event_data.message = message
+        # pylint: disable=protected-access
         event_data._raw_amqp_message = raw_amqp_message if raw_amqp_message else AmqpAnnotatedMessage(message=message)
         return event_data
 

sdk/eventhub/azure-eventhub/azure/eventhub/_configuration.py

@@ -10,13 +10,15 @@
     from urllib.parse import urlparse
 
 from uamqp.constants import TransportType, DEFAULT_AMQPS_PORT, DEFAULT_AMQP_WSS_PORT
+from ._retry import RetryMode
 
 
 class Configuration(object):  # pylint:disable=too-many-instance-attributes
     def __init__(self, **kwargs):
         self.user_agent = kwargs.get("user_agent")  # type: Optional[str]
         self.retry_total = kwargs.get("retry_total", 3)  # type: int
         self.max_retries = self.retry_total  # type: int
+        self.retry_mode = kwargs.get("retry_mode", RetryMode.EXPONENTIAL)
         self.backoff_factor = kwargs.get("retry_backoff_factor", 0.8)  # type: float
         self.backoff_max = kwargs.get("retry_backoff_max", 120)  # type: int
         self.network_tracing = kwargs.get("network_tracing", False)  # type: bool

sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py

@@ -72,6 +72,15 @@ class EventHubConsumerClient(ClientBase):
      The failed internal partition consumer will be closed (`on_partition_close` will be called if provided) and
      new internal partition consumer will be created (`on_partition_initialize` will be called if provided) to resume
      receiving.
+    :keyword float retry_backoff_factor: A backoff factor to apply between attempts after the second try
+     (most errors are resolved immediately by a second try without a delay).
+     In fixed mode, retry policy will always sleep for {backoff factor}.
+     In 'exponential' mode, retry policy will sleep for: `{backoff factor} * (2 ** ({number of total retries} - 1))`
+     seconds. If the backoff_factor is 0.1, then the retry will sleep
+     for [0.0s, 0.2s, 0.4s, ...] between retries. The default value is 0.8.
+    :keyword float retry_backoff_max: The maximum back off time. Default value is 120 seconds (2 minutes).
+    :keyword retry_mode: Fixed or exponential delay between attempts, default is exponential.
+    :paramtype retry_mode: ~azure.eventhub.RetryMode
     :keyword float idle_timeout: Timeout, in seconds, after which this client will close the underlying connection
      if there is no further activity. By default the value is None, meaning that the client will not shutdown due to
      inactivity unless initiated by the service.
@@ -219,6 +228,15 @@ def from_connection_string(cls, conn_str, consumer_group, **kwargs):
          information. The failed internal partition consumer will be closed (`on_partition_close` will be called
          if provided) and new internal partition consumer will be created (`on_partition_initialize` will be called if
          provided) to resume receiving.
+        :keyword float retry_backoff_factor: A backoff factor to apply between attempts after the second try
+         (most errors are resolved immediately by a second try without a delay).
+         In fixed mode, retry policy will always sleep for {backoff factor}.
+         In 'exponential' mode, retry policy will sleep for: `{backoff factor} * (2 ** ({number of total retries} - 1))`
+         seconds. If the backoff_factor is 0.1, then the retry will sleep
+         for [0.0s, 0.2s, 0.4s, ...] between retries. The default value is 0.8.
+        :keyword float retry_backoff_max: The maximum back off time. Default value is 120 seconds (2 minutes).
+        :keyword retry_mode: Fixed or exponential delay between attempts, default is exponential.
+        :paramtype retry_mode: ~azure.eventhub.RetryMode
         :keyword float idle_timeout: Timeout, in seconds, after which this client will close the underlying connection
          if there is no furthur activity. By default the value is None, meaning that the client will not shutdown due
          to inactivity unless initiated by the service.

sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py

@@ -47,6 +47,15 @@ class EventHubProducerClient(ClientBase):
     :keyword str user_agent: If specified, this will be added in front of the user agent string.
     :keyword int retry_total: The total number of attempts to redo a failed operation when an error occurs. Default
      value is 3.
+    :keyword float retry_backoff_factor: A backoff factor to apply between attempts after the second try
+     (most errors are resolved immediately by a second try without a delay).
+     In fixed mode, retry policy will always sleep for {backoff factor}.
+     In 'exponential' mode, retry policy will sleep for: `{backoff factor} * (2 ** ({number of total retries} - 1))`
+     seconds. If the backoff_factor is 0.1, then the retry will sleep
+     for [0.0s, 0.2s, 0.4s, ...] between retries. The default value is 0.8.
+    :keyword float retry_backoff_max: The maximum back off time. Default value is 120 seconds (2 minutes).
+    :keyword retry_mode: Fixed or exponential delay between attempts, default is exponential.
+    :paramtype retry_mode: ~azure.eventhub.RetryMode
     :keyword float idle_timeout: Timeout, in seconds, after which this client will close the underlying connection
      if there is no activity. By default the value is None, meaning that the client will not shutdown due to inactivity
      unless initiated by the service.
@@ -185,6 +194,15 @@ def from_connection_string(cls, conn_str, **kwargs):
         :keyword str user_agent: If specified, this will be added in front of the user agent string.
         :keyword int retry_total: The total number of attempts to redo a failed operation when an error occurs.
          Default value is 3.
+        :keyword float retry_backoff_factor: A backoff factor to apply between attempts after the second try
+         (most errors are resolved immediately by a second try without a delay).
+         In fixed mode, retry policy will always sleep for {backoff factor}.
+         In 'exponential' mode, retry policy will sleep for: `{backoff factor} * (2 ** ({number of total retries} - 1))`
+         seconds. If the backoff_factor is 0.1, then the retry will sleep
+         for [0.0s, 0.2s, 0.4s, ...] between retries. The default value is 0.8.
+        :keyword float retry_backoff_max: The maximum back off time. Default value is 120 seconds (2 minutes).
+        :keyword retry_mode: Fixed or exponential delay between attempts, default is exponential.
+        :paramtype retry_mode: ~azure.eventhub.RetryMode
         :keyword float idle_timeout: Timeout, in seconds, after which this client will close the underlying connection
          if there is no activity. By default the value is None, meaning that the client will not shutdown due to
          inactivity unless initiated by the service.

sdk/eventhub/azure-eventhub/azure/eventhub/_retry.py

@@ -0,0 +1,11 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+from typing import Optional, Dict, Any
+
+from enum import Enum
+
+class RetryMode(str, Enum):
+    EXPONENTIAL = 'exponential'
+    FIXED = 'fixed'

sdk/eventhub/azure-eventhub/azure/eventhub/_utils.py

@@ -311,14 +311,14 @@ def decode_with_recurse(data, encoding="UTF-8"):
         return data
     if isinstance(data, six.binary_type):
         return data.decode(encoding)
-    if isinstance(data, Mapping):
+    if isinstance(data, Mapping): # pylint:disable=isinstance-second-argument-not-valid-type
         decoded_mapping = {}
         for k, v in data.items():
             decoded_key = decode_with_recurse(k, encoding)
             decoded_val = decode_with_recurse(v, encoding)
             decoded_mapping[decoded_key] = decoded_val
         return decoded_mapping
-    if isinstance(data, Iterable):
+    if isinstance(data, Iterable): # pylint:disable=isinstance-second-argument-not-valid-type
         decoded_list = []
         for d in data:
             decoded_list.append(decode_with_recurse(d, encoding))

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_async_utils.py

@@ -6,7 +6,6 @@
 
 import sys
 
-
 def get_dict_with_loop_if_needed(loop):
     if sys.version_info >= (3, 10):
         if loop: