Shared connection

sdk/eventhub/azure-eventhubs/azure/eventhub/_connection_manager.py

@@ -3,14 +3,14 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-import threading
-from uamqp import Connection, TransportType
+from threading import RLock
+from uamqp import Connection, TransportType, c_uamqp
 
 
 class _SharedConnectionManager(object):
     def __init__(self, **kwargs):
-        self._lock = threading.Lock()
-        self._conn = None
+        self._lock = RLock()
+        self._conn = None  # type: Connection
 
         self._container_id = kwargs.get("container_id")
         self._debug = kwargs.get("debug")
@@ -42,12 +42,22 @@ def get_connection(self, host, auth):
                     encoding=self._encoding)
             return self._conn
 
-    def close_connection(self, conn=None):
+    def close_connection(self):
         with self._lock:
             if self._conn:
                 self._conn.destroy()
             self._conn = None
 
+    def reset_connection_if_broken(self):
+        with self._lock:
+            if self._conn and self._conn._state in (
+                c_uamqp.ConnectionState.CLOSE_RCVD,
+                c_uamqp.ConnectionState.CLOSE_SENT,
+                c_uamqp.ConnectionState.DISCARDING,
+                c_uamqp.ConnectionState.END,
+            ):
+                self._conn = None
+
 
 class _SeparateConnectionManager(object):
     def __init__(self, **kwargs):
@@ -59,6 +69,9 @@ def get_connection(self, host, auth):
     def close_connection(self):
         pass
 
+    def reset_connection_if_broken(self):
+        pass
+
 
 def get_connection_manager(**kwargs):
-    return _SeparateConnectionManager(**kwargs)
+    return _SharedConnectionManager(**kwargs)

sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_connection_manager_async.py

@@ -0,0 +1,77 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from asyncio import Lock
+from uamqp import TransportType, c_uamqp
+from uamqp.async_ops import ConnectionAsync
+
+
+class _SharedConnectionManager(object):
+    def __init__(self, **kwargs):
+        self._lock = Lock()
+        self._conn = None
+
+        self._container_id = kwargs.get("container_id")
+        self._debug = kwargs.get("debug")
+        self._error_policy = kwargs.get("error_policy")
+        self._properties = kwargs.get("properties")
+        self._encoding = kwargs.get("encoding") or "UTF-8"
+        self._transport_type = kwargs.get('transport_type') or TransportType.Amqp
+        self._http_proxy = kwargs.get('http_proxy')
+        self._max_frame_size = kwargs.get("max_frame_size")
+        self._channel_max = kwargs.get("channel_max")
+        self._idle_timeout = kwargs.get("idle_timeout")
+        self._remote_idle_timeout_empty_frame_send_ratio = kwargs.get("remote_idle_timeout_empty_frame_send_ratio")
+
+    async def get_connection(self, host, auth):
+        # type: (...) -> ConnectionAsync
+        async with self._lock:
+            if self._conn is None:
+                self._conn = ConnectionAsync(
+                    host,
+                    auth,
+                    container_id=self._container_id,
+                    max_frame_size=self._max_frame_size,
+                    channel_max=self._channel_max,
+                    idle_timeout=self._idle_timeout,
+                    properties=self._properties,
+                    remote_idle_timeout_empty_frame_send_ratio=self._remote_idle_timeout_empty_frame_send_ratio,
+                    error_policy=self._error_policy,
+                    debug=self._debug,
+                    encoding=self._encoding)
+            return self._conn
+
+    async def close_connection(self):
+        async with self._lock:
+            if self._conn:
+                await self._conn.destroy_async()
+            self._conn = None
+
+    def reset_connection_if_broken(self):
+        with self._lock:
+            if self._conn and self._conn._state in (
+                c_uamqp.ConnectionState.CLOSE_RCVD,
+                c_uamqp.ConnectionState.CLOSE_SENT,
+                c_uamqp.ConnectionState.DISCARDING,
+                c_uamqp.ConnectionState.END,
+            ):
+                self._conn = None
+
+class _SeparateConnectionManager(object):
+    def __init__(self, **kwargs):
+        pass
+
+    async def get_connection(self, host, auth):
+        pass  # return None
+
+    async def close_connection(self):
+        pass
+
+    def reset_connection_if_broken(self):
+        pass
+
+
+def get_connection_manager(**kwargs):
+    return _SharedConnectionManager(**kwargs)

sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py

@@ -22,6 +22,8 @@
 
 from .producer_async import EventHubProducer
 from .consumer_async import EventHubConsumer
+from ._connection_manager_async import get_connection_manager
+from .error_async import _handle_exception
 
 
 log = logging.getLogger(__name__)
@@ -42,6 +44,16 @@ class EventHubClient(EventHubClientAbstract):
 
     """
 
+    def __init__(self, host, event_hub_path, credential, **kwargs):
+        super(EventHubClient, self).__init__(host, event_hub_path, credential, **kwargs)
+        self._conn_manager = get_connection_manager(**kwargs)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
     def _create_auth(self, username=None, password=None):
         """
         Create an ~uamqp.authentication.cbs_auth_async.SASTokenAuthAsync instance to authenticate
@@ -85,33 +97,31 @@ def _create_auth(self, username=None, password=None):
                                                 get_jwt_token, http_proxy=http_proxy,
                                                 transport_type=transport_type)
 
+    async def _handle_exception(self, exception, retry_count, max_retries):
+        await _handle_exception(exception, retry_count, max_retries, self, log)
+
+    async def _close_connection(self):
+        self._conn_manager.reset_connection_if_broken()
+
     async def _management_request(self, mgmt_msg, op_type):
-        alt_creds = {
-            "username": self._auth_config.get("iot_username"),
-            "password": self._auth_config.get("iot_password")}
-        connect_count = 0
+        max_retries = self.config.max_retries
+        retry_count = 0
         while True:
-            connect_count += 1
-            mgmt_auth = self._create_auth(**alt_creds)
+            mgmt_auth = self._create_auth()
             mgmt_client = AMQPClientAsync(self.mgmt_target, auth=mgmt_auth, debug=self.config.network_tracing)
             try:
-                await mgmt_client.open_async()
+                conn = await self._conn_manager.get_connection(self.host, mgmt_auth)
+                await mgmt_client.open_async(connection=conn)
                 response = await mgmt_client.mgmt_request_async(
                     mgmt_msg,
                     constants.READ_OPERATION,
                     op_type=op_type,
                     status_code_field=b'status-code',
                     description_fields=b'status-description')
                 return response
-            except (errors.AMQPConnectionError, errors.TokenAuthFailure, compat.TimeoutException) as failure:
-                if connect_count >= self.config.max_retries:
-                    err = ConnectError(
-                        "Can not connect to EventHubs or get management info from the service. "
-                        "Please make sure the connection string or token is correct and retry. "
-                        "Besides, this method doesn't work if you use an IoT connection string.",
-                        failure
-                    )
-                    raise err
+            except Exception as exception:
+                await self._handle_exception(exception, retry_count, max_retries)
+                retry_count += 1
             finally:
                 await mgmt_client.close_async()
 
@@ -263,3 +273,6 @@ def create_producer(
         handler = EventHubProducer(
             self, target, partition=partition_id, send_timeout=send_timeout, loop=loop)
         return handler
+
+    async def close(self):
+        await self._conn_manager.close_connection()