close spawned children

Author: yunhaoling

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## 7.0.0b6 (Unreleased)
 
+**Breaking Changes**
+
+* `ServiceBusClient.close()` now closes spawned senders and receivers.
 
 ## 7.0.0b5 (2020-08-10)
 

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

@@ -69,6 +69,7 @@ def __init__(
             self._auth_uri = "{}/{}".format(self._auth_uri, self._entity_name)
         # Internal flag for switching whether to apply connection sharing, pending fix in uamqp library
         self._connection_sharing = False
+        self._handlers = []
 
     def __enter__(self):
         if self._connection_sharing:
@@ -89,10 +90,15 @@ def _create_uamqp_connection(self):
     def close(self):
         # type: () -> None
         """
-        Close down the ServiceBus client and the underlying connection.
+        Close down the ServiceBus client.
+        All spawned senders, receivers and underlying connection will be shutdown.
 
         :return: None
         """
+        for handler in self._handlers:
+            handler.close()
+        self._handlers.clear()
+
         if self._connection_sharing and self._connection:
             self._connection.destroy()
 
@@ -157,7 +163,7 @@ def get_queue_sender(self, queue_name, **kwargs):
 
         """
         # pylint: disable=protected-access
-        return ServiceBusSender(
+        handler = ServiceBusSender(
             fully_qualified_namespace=self.fully_qualified_namespace,
             queue_name=queue_name,
             credential=self._credential,
@@ -168,6 +174,8 @@ def get_queue_sender(self, queue_name, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_queue_receiver(self, queue_name, **kwargs):
         # type: (str, Any) -> ServiceBusReceiver
@@ -205,7 +213,7 @@ def get_queue_receiver(self, queue_name, **kwargs):
 
         """
         # pylint: disable=protected-access
-        return ServiceBusReceiver(
+        handler = ServiceBusReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             queue_name=queue_name,
             credential=self._credential,
@@ -216,6 +224,8 @@ def get_queue_receiver(self, queue_name, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_queue_deadletter_receiver(self, queue_name, **kwargs):
         # type: (str, Any) -> ServiceBusReceiver
@@ -265,7 +275,7 @@ def get_queue_deadletter_receiver(self, queue_name, **kwargs):
             queue_name=queue_name,
             transfer_deadletter=kwargs.get('transfer_deadletter', False)
         )
-        return ServiceBusReceiver(
+        handler = ServiceBusReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             entity_name=entity_name,
             credential=self._credential,
@@ -277,6 +287,8 @@ def get_queue_deadletter_receiver(self, queue_name, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_topic_sender(self, topic_name, **kwargs):
         # type: (str, Any) -> ServiceBusSender
@@ -300,7 +312,7 @@ def get_topic_sender(self, topic_name, **kwargs):
                 :caption: Create a new instance of the ServiceBusSender from ServiceBusClient.
 
         """
-        return ServiceBusSender(
+        handler = ServiceBusSender(
             fully_qualified_namespace=self.fully_qualified_namespace,
             topic_name=topic_name,
             credential=self._credential,
@@ -311,6 +323,8 @@ def get_topic_sender(self, topic_name, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_subscription_receiver(self, topic_name, subscription_name, **kwargs):
         # type: (str, str, Any) -> ServiceBusReceiver
@@ -353,7 +367,7 @@ def get_subscription_receiver(self, topic_name, subscription_name, **kwargs):
 
         """
         # pylint: disable=protected-access
-        return ServiceBusReceiver(
+        handler = ServiceBusReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             topic_name=topic_name,
             subscription_name=subscription_name,
@@ -365,6 +379,8 @@ def get_subscription_receiver(self, topic_name, subscription_name, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_subscription_deadletter_receiver(self, topic_name, subscription_name, **kwargs):
         # type: (str, str, Any) -> ServiceBusReceiver
@@ -416,7 +432,7 @@ def get_subscription_deadletter_receiver(self, topic_name, subscription_name, **
             subscription_name=subscription_name,
             transfer_deadletter=kwargs.get('transfer_deadletter', False)
         )
-        return ServiceBusReceiver(
+        handler = ServiceBusReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             entity_name=entity_name,
             credential=self._credential,
@@ -428,6 +444,8 @@ def get_subscription_deadletter_receiver(self, topic_name, subscription_name, **
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_subscription_session_receiver(self, topic_name, subscription_name, session_id=None, **kwargs):
         # type: (str, str, str, Any) -> ServiceBusReceiver
@@ -473,7 +491,7 @@ def get_subscription_session_receiver(self, topic_name, subscription_name, sessi
 
         """
         # pylint: disable=protected-access
-        return ServiceBusSessionReceiver(
+        handler = ServiceBusSessionReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             topic_name=topic_name,
             subscription_name=subscription_name,
@@ -486,6 +504,8 @@ def get_subscription_session_receiver(self, topic_name, subscription_name, sessi
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_queue_session_receiver(self, queue_name, session_id=None, **kwargs):
         # type: (str, str, Any) -> ServiceBusSessionReceiver
@@ -526,7 +546,7 @@ def get_queue_session_receiver(self, queue_name, session_id=None, **kwargs):
 
         """
         # pylint: disable=protected-access
-        return ServiceBusSessionReceiver(
+        handler = ServiceBusSessionReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             queue_name=queue_name,
             credential=self._credential,
@@ -538,3 +558,5 @@ def get_queue_session_receiver(self, queue_name, session_id=None, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler

sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_client_async.py

@@ -71,6 +71,7 @@ def __init__(
             self._auth_uri = "{}/{}".format(self._auth_uri, self._entity_name)
         # Internal flag for switching whether to apply connection sharing, pending fix in uamqp library
         self._connection_sharing = False
+        self._handlers = []
 
     async def __aenter__(self):
         if self._connection_sharing:
@@ -133,9 +134,14 @@ async def close(self):
         # type: () -> None
         """
         Close down the ServiceBus client.
+        All spawned senders, receivers and underlying connection will be shutdown.
 
         :return: None
         """
+        for handler in self._handlers:
+            await handler.close()
+        self._handlers.clear()
+
         if self._connection_sharing and self._connection:
             await self._connection.destroy_async()
 
@@ -159,7 +165,7 @@ def get_queue_sender(self, queue_name, **kwargs):
 
         """
         # pylint: disable=protected-access
-        return ServiceBusSender(
+        handler = ServiceBusSender(
             fully_qualified_namespace=self.fully_qualified_namespace,
             queue_name=queue_name,
             credential=self._credential,
@@ -170,6 +176,8 @@ def get_queue_sender(self, queue_name, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_queue_receiver(self, queue_name, **kwargs):
         # type: (str, Any) -> ServiceBusReceiver
@@ -206,7 +214,7 @@ def get_queue_receiver(self, queue_name, **kwargs):
 
         """
         # pylint: disable=protected-access
-        return ServiceBusReceiver(
+        handler = ServiceBusReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             queue_name=queue_name,
             credential=self._credential,
@@ -217,6 +225,8 @@ def get_queue_receiver(self, queue_name, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_queue_deadletter_receiver(self, queue_name, **kwargs):
         # type: (str, Any) -> ServiceBusReceiver
@@ -266,7 +276,7 @@ def get_queue_deadletter_receiver(self, queue_name, **kwargs):
             queue_name=queue_name,
             transfer_deadletter=kwargs.get('transfer_deadletter', False)
         )
-        return ServiceBusReceiver(
+        handler = ServiceBusReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             entity_name=entity_name,
             credential=self._credential,
@@ -278,6 +288,8 @@ def get_queue_deadletter_receiver(self, queue_name, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_topic_sender(self, topic_name, **kwargs):
         # type: (str, Any) -> ServiceBusSender
@@ -301,7 +313,7 @@ def get_topic_sender(self, topic_name, **kwargs):
                 :caption: Create a new instance of the ServiceBusSender from ServiceBusClient.
 
         """
-        return ServiceBusSender(
+        handler = ServiceBusSender(
             fully_qualified_namespace=self.fully_qualified_namespace,
             topic_name=topic_name,
             credential=self._credential,
@@ -312,6 +324,8 @@ def get_topic_sender(self, topic_name, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_subscription_receiver(self, topic_name, subscription_name, **kwargs):
         # type: (str, str, Any) -> ServiceBusReceiver
@@ -354,7 +368,7 @@ def get_subscription_receiver(self, topic_name, subscription_name, **kwargs):
 
         """
         # pylint: disable=protected-access
-        return ServiceBusReceiver(
+        handler = ServiceBusReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             topic_name=topic_name,
             subscription_name=subscription_name,
@@ -366,6 +380,8 @@ def get_subscription_receiver(self, topic_name, subscription_name, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_subscription_deadletter_receiver(self, topic_name, subscription_name, **kwargs):
         # type: (str, str, Any) -> ServiceBusReceiver
@@ -417,7 +433,7 @@ def get_subscription_deadletter_receiver(self, topic_name, subscription_name, **
             subscription_name=subscription_name,
             transfer_deadletter=kwargs.get('transfer_deadletter', False)
         )
-        return ServiceBusReceiver(
+        handler = ServiceBusReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             entity_name=entity_name,
             credential=self._credential,
@@ -429,6 +445,8 @@ def get_subscription_deadletter_receiver(self, topic_name, subscription_name, **
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_subscription_session_receiver(self, topic_name, subscription_name, session_id=None, **kwargs):
         # type: (str, str, str, Any) -> ServiceBusReceiver
@@ -474,7 +492,7 @@ def get_subscription_session_receiver(self, topic_name, subscription_name, sessi
 
         """
         # pylint: disable=protected-access
-        return ServiceBusSessionReceiver(
+        handler = ServiceBusSessionReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             topic_name=topic_name,
             subscription_name=subscription_name,
@@ -487,6 +505,8 @@ def get_subscription_session_receiver(self, topic_name, subscription_name, sessi
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler
 
     def get_queue_session_receiver(self, queue_name, session_id=None, **kwargs):
         # type: (str, str, Any) -> ServiceBusSessionReceiver
@@ -526,7 +546,7 @@ def get_queue_session_receiver(self, queue_name, session_id=None, **kwargs):
 
         """
         # pylint: disable=protected-access
-        return ServiceBusSessionReceiver(
+        handler = ServiceBusSessionReceiver(
             fully_qualified_namespace=self.fully_qualified_namespace,
             queue_name=queue_name,
             credential=self._credential,
@@ -538,3 +558,5 @@ def get_queue_session_receiver(self, queue_name, session_id=None, **kwargs):
             user_agent=self._config.user_agent,
             **kwargs
         )
+        self._handlers.append(handler)
+        return handler

sdk/servicebus/azure-servicebus/tests/async_tests/test_sb_client_async.py

@@ -0,0 +1,58 @@
+#--------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+#--------------------------------------------------------------------------
+
+
+import logging
+import pytest
+
+from azure.servicebus.aio import ServiceBusClient
+from devtools_testutils import AzureMgmtTestCase, CachedResourceGroupPreparer
+from servicebus_preparer import CachedServiceBusNamespacePreparer, CachedServiceBusQueuePreparer
+from utilities import get_logger
+
+_logger = get_logger(logging.DEBUG)
+
+
+class ServiceBusClientAsyncTests(AzureMgmtTestCase):
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedResourceGroupPreparer()
+    @CachedServiceBusNamespacePreparer(name_prefix='servicebustest')
+    @CachedServiceBusQueuePreparer(name_prefix='servicebustest', dead_lettering_on_message_expiration=True)
+    async def test_async_sb_client_close_spawned_handlers(self, servicebus_namespace_connection_string, servicebus_queue, **kwargs):
+        client = ServiceBusClient.from_connection_string(servicebus_namespace_connection_string)
+
+        # context manager
+        async with client:
+            assert len(client._handlers) == 0
+            sender = client.get_queue_sender(servicebus_queue.name)
+            receiver = client.get_queue_receiver(servicebus_queue.name)
+            await sender._open()
+            await receiver._open()
+
+            assert sender._handler and sender._running
+            assert receiver._handler and receiver._running
+            assert len(client._handlers) == 2
+
+        assert not sender._handler and not sender._running
+        assert not receiver._handler and not receiver._running
+        assert len(client._handlers) == 0
+
+        # close operation
+        sender = client.get_queue_sender(servicebus_queue.name)
+        receiver = client.get_queue_receiver(servicebus_queue.name)
+        await sender._open()
+        await receiver._open()
+
+        assert sender._handler and sender._running
+        assert receiver._handler and receiver._running
+        assert len(client._handlers) == 2
+
+        await client.close()
+
+        assert not sender._handler and not sender._running
+        assert not receiver._handler and not receiver._running
+        assert len(client._handlers) == 0