[ServiceBus] Iterator Support

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

@@ -4,7 +4,7 @@
 # license information.
 # -------------------------------------------------------------------------
 import uuid
-
+import time
 from .._pyamqp.endpoints import Source
 from .message import ServiceBusReceivedMessage
 from ..exceptions import _ServiceBusErrorPolicy, MessageAlreadySettled
@@ -129,8 +129,16 @@ def _populate_message_properties(self, message):
 
     def _enhanced_message_received(self, frame, message):
         # pylint: disable=protected-access
-        self._handler._was_message_received = True
+        self._handler._last_activity_timestamp = time.time()
         if self._receive_context.is_set():
             self._handler._received_messages.put((frame, message))
         else:
             self._handler.settle_messages(frame[1], 'released')
+
+    async def _enhanced_message_received_async(self, frame, message):
+        # pylint: disable=protected-access
+        self._handler._last_activity_timestamp = time.time()
+        if self._receive_context.is_set():
+            self._handler._received_messages.put((frame, message))
+        else:
+            await self._handler.settle_messages_async(frame[1], 'released')

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py

@@ -135,6 +135,12 @@ class AMQPClientAsync(AMQPClientSync):
      Default is None in which case `certifi.where()` will be used.
     :paramtype connection_verify: str
     """
+
+    def __init__(self, hostname, **kwargs):
+        self._mgmt_link_lock_async = asyncio.Lock()
+        super().__init__(hostname,**kwargs)
+
+
     async def _keep_alive_async(self):
         start_time = time.time()
         try:
@@ -147,7 +153,8 @@ async def _keep_alive_async(self):
                         self.__class__.__name__,
                         extra=self._network_trace_params
                     )
-                    await asyncio.shield(self._connection.work_async())
+                    await asyncio.shield(self._connection.listen(wait=self._socket_timeout,
+                        batch=self._link.current_link_credit))
                     start_time = current_time
                 await asyncio.sleep(1)
         except Exception as e:  # pylint: disable=broad-except
@@ -263,10 +270,9 @@ async def open_async(self, connection=None):
         self._network_trace_params["amqpConnection"] = self._connection._container_id
         self._network_trace_params["amqpSession"] = self._session.name
         self._shutdown = False
-        # TODO: Looks like this is broken - should re-enable later and test
-        # correct empty frame behaviour
-        # if self._keep_alive_interval:
-        #    self._keep_alive_thread = asyncio.ensure_future(self._keep_alive_async())
+
+        if self._keep_alive_interval:
+            self._keep_alive_thread = asyncio.ensure_future(self._keep_alive_async())
 
     async def close_async(self):
         """Close the client asynchronously. This includes closing the Session
@@ -277,9 +283,6 @@ async def close_async(self):
         self._shutdown = True
         if not self._session:
             return  # already closed.
-        if self._keep_alive_thread:
-            await self._keep_alive_thread
-            self._keep_alive_thread = None
         await self._close_link_async()
         if self._cbs_authenticator:
             await self._cbs_authenticator.close()
@@ -289,6 +292,9 @@ async def close_async(self):
         if not self._external_connection:
             await self._connection.close()
             self._connection = None
+        if self._keep_alive_thread:
+            await self._keep_alive_thread
+            self._keep_alive_thread = None
         self._network_trace_params["amqpConnection"] = None
         self._network_trace_params["amqpSession"] = None
 
@@ -360,15 +366,16 @@ async def mgmt_request_async(self, message, **kwargs):
         operation_type = kwargs.pop("operation_type", None)
         node = kwargs.pop("node", "$management")
         timeout = kwargs.pop('timeout', 0)
-        try:
-            mgmt_link = self._mgmt_links[node]
-        except KeyError:
-            mgmt_link = ManagementOperation(self._session, endpoint=node, **kwargs)
-            self._mgmt_links[node] = mgmt_link
-            await mgmt_link.open()
+        async with self._mgmt_link_lock_async:
+            try:
+                mgmt_link = self._mgmt_links[node]
+            except KeyError:
+                mgmt_link = ManagementOperation(self._session, endpoint=node, **kwargs)
+                self._mgmt_links[node] = mgmt_link
+                await mgmt_link.open()
 
-            while not await mgmt_link.ready():
-                await self._connection.listen(wait=False)
+        while not await mgmt_link.ready():
+            await self._connection.listen(wait=False)
 
         operation_type = operation_type or b'empty'
         status, description, response = await mgmt_link.execute(
@@ -724,7 +731,8 @@ async def _client_run_async(self, **kwargs):
         :rtype: bool
         """
         try:
-            await self._link.flow()
+            if self._link.current_link_credit == 0:
+                await self._link.flow()
             await self._connection.listen(wait=self._socket_timeout, **kwargs)
         except ValueError:
             _logger.info("Timeout reached, closing receiver.", extra=self._network_trace_params)
@@ -741,6 +749,7 @@ async def _message_received_async(self, frame, message):
         :param message: Received message.
         :type message: ~pyamqp.message.Message
         """
+        self._last_activity_timestamp = time.time()
         if self._message_received_callback:
             await self._message_received_callback(message)
         if not self._streaming_receive:
@@ -829,6 +838,47 @@ async def receive_message_batch_async(self, **kwargs):
             **kwargs
         )
 
+    async def receive_messages_iter_async(self, timeout=None, on_message_received=None):
+        """Receive messages by generator. Messages returned in the generator have already been
+        accepted - if you wish to add logic to accept or reject messages based on custom
+        criteria, pass in a callback.
+
+        :param on_message_received: A callback to process messages as they arrive from the
+         service. It takes a single argument, a ~pyamqp.message.Message object.
+        :type on_message_received: callable[~pyamqp.message.Message]
+        """
+        self._message_received_callback = on_message_received
+        return self._message_generator_async(timeout=timeout)
+
+    async def _message_generator_async(self, timeout=None):
+        """Iterate over processed messages in the receive queue.
+
+        :rtype: generator[~pyamqp.message.Message]
+        """
+        receiving = True
+        message = None
+        self._last_activity_timestamp = time.time()
+        self._timeout_reached = False
+        self._timeout = timeout if timeout else self._timeout
+        try:
+            while receiving and not self._timeout_reached:
+                if self._timeout > 0:
+                    if time.time() - self._last_activity_timestamp >= self._timeout:
+                        self._timeout_reached = True
+
+                if not self._timeout_reached:
+                    receiving = await self.do_work_async()
+
+                while not self._received_messages.empty():
+                    message = self._received_messages.get()
+                    self._last_activity_timestamp = time.time()
+                    self._received_messages.task_done()
+                    yield message
+
+        finally:
+            if self._shutdown:
+                await self.close_async()
+
     @overload
     async def settle_messages_async(
         self,

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_connection_async.py

@@ -236,7 +236,7 @@ async def _read_frame(self, wait: Union[bool, int, float] = True, **kwargs) -> b
         """
         timeout: Optional[Union[int, float]] = None
         if wait is False:
-            timeout = 1  # TODO: What should this default be?
+            timeout = .2
         elif wait is True:
             timeout = None
         else:

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_transport_async.py

@@ -341,7 +341,7 @@ def _init_socket(self, socket_settings):
             # For uamqp exception parity. Remove later when resolving issue #27128.
             exc.filename = self.sslopts
             raise exc
-        self.sock.settimeout(1)  # set socket back to non-blocking mode
+        self.sock.settimeout(.2)  # set socket back to non-blocking mode
 
     def _get_tcp_socket_defaults(self, sock):  # pylint: disable=no-self-use
         tcp_opts = {}

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py

@@ -10,6 +10,7 @@
 # pylint: disable=too-many-lines
 # TODO: Check types of kwargs (issue exists for this)
 import logging
+import threading
 import queue
 import time
 import uuid
@@ -163,8 +164,9 @@ def __init__(self, hostname, **kwargs):
         self._cbs_authenticator = None
         self._auth_timeout = kwargs.pop("auth_timeout", DEFAULT_AUTH_TIMEOUT)
         self._mgmt_links = {}
+        self._mgmt_link_lock = threading.Lock()
         self._retry_policy = kwargs.pop("retry_policy", RetryPolicy())
-        self._keep_alive_interval = int(kwargs.get("keep_alive_interval") or 0)
+        self._keep_alive_interval = int(kwargs.get("keep_alive_interval", 0))
         self._keep_alive_thread = None
 
         # Connection settings
@@ -217,6 +219,20 @@ def __exit__(self, *args):
         """Close and destroy Client on exiting a context manager."""
         self.close()
 
+    def _keep_alive(self):
+        start_time = time.time()
+        try:
+            while self._connection and not self._shutdown:
+                current_time = time.time()
+                elapsed_time = current_time - start_time
+                if elapsed_time >= self._keep_alive_interval:
+                    _logger.debug("Keeping %r connection alive.", self.__class__.__name__)
+                    self._connection.listen(wait=self._socket_timeout, batch=self._link.current_link_credit)
+                    start_time = current_time
+                time.sleep(1)
+        except Exception as e:  # pylint: disable=broad-except
+            _logger.debug("Connection keep-alive for %r failed: %r.", self.__class__.__name__, e)
+
     def _client_ready(self):  # pylint: disable=no-self-use
         """Determine whether the client is ready to start sending and/or
         receiving messages. To be ready, the connection must be open and
@@ -306,6 +322,10 @@ def open(self, connection=None):
                 outgoing_window=self._outgoing_window,
             )
             self._session.begin()
+        if self._keep_alive_interval:
+            self._keep_alive_thread = threading.Thread(target=self._keep_alive)
+            self._keep_alive_thread.daemon = True
+            self._keep_alive_thread.start()
         if self._auth.auth_type == AUTH_TYPE_CBS:
             self._cbs_authenticator = CBSAuthenticator(
                 session=self._session, auth=self._auth, auth_timeout=self._auth_timeout
@@ -339,6 +359,12 @@ def close(self):
         if not self._external_connection:
             self._connection.close()
             self._connection = None
+        if self._keep_alive_thread:
+            try:
+                self._keep_alive_thread.join()
+            except RuntimeError:  # Probably thread failed to start in .open()
+                logging.debug("Keep alive thread failed to join.", exc_info=True)
+            self._keep_alive_thread = None
         self._network_trace_params["amqpConnection"] = None
         self._network_trace_params["amqpSession"] = None
 
@@ -409,16 +435,16 @@ def mgmt_request(self, message, **kwargs):
         operation_type = kwargs.pop("operation_type", None)
         node = kwargs.pop("node", "$management")
         timeout = kwargs.pop("timeout", 0)
-        try:
-            mgmt_link = self._mgmt_links[node]
-        except KeyError:
-            mgmt_link = ManagementOperation(self._session, endpoint=node, **kwargs)
-            self._mgmt_links[node] = mgmt_link
-            mgmt_link.open()
-
-            while not mgmt_link.ready():
-                self._connection.listen(wait=False)
-
+        with self._mgmt_link_lock:
+            try:
+                mgmt_link = self._mgmt_links[node]
+            except KeyError:
+                mgmt_link = ManagementOperation(self._session, endpoint=node, **kwargs)
+                self._mgmt_links[node] = mgmt_link
+                mgmt_link.open()
+
+        while not mgmt_link.ready():
+            self._connection.listen(wait=False)
         operation_type = operation_type or b"empty"
         status, description, response = mgmt_link.execute(
             message, operation=operation, operation_type=operation_type, timeout=timeout
@@ -659,7 +685,7 @@ def send_message(self, message, **kwargs):
         self._do_retryable_operation(self._send_message_impl, message=message, **kwargs)
 
 
-class ReceiveClient(AMQPClient):
+class ReceiveClient(AMQPClient): # pylint:disable=too-many-instance-attributes
     """
     An AMQP client for receiving messages.
     :param source: The source AMQP service endpoint. This can either be the URI as
@@ -762,6 +788,12 @@ def __init__(self, hostname, source, **kwargs):
         self._max_message_size = kwargs.pop("max_message_size", MAX_FRAME_SIZE_BYTES)
         self._link_properties = kwargs.pop("link_properties", None)
         self._link_credit = kwargs.pop("link_credit", 300)
+
+        # Iterator
+        self._timeout = kwargs.pop("timeout", 0)
+        self._timeout_reached = False
+        self._last_activity_timestamp = time.time()
+
         super(ReceiveClient, self).__init__(hostname, **kwargs)
 
     def _client_ready(self):
@@ -799,7 +831,8 @@ def _client_run(self, **kwargs):
         :rtype: bool
         """
         try:
-            self._link.flow()
+            if self._link.current_link_credit == 0:
+                self._link.flow()
             self._connection.listen(wait=self._socket_timeout, **kwargs)
         except ValueError:
             _logger.info("Timeout reached, closing receiver.", extra=self._network_trace_params)
@@ -816,6 +849,7 @@ def _message_received(self, frame, message):
         :param message: Received message.
         :type message: ~pyamqp.message.Message
         """
+        self._last_activity_timestamp = time.time()
         if self._message_received_callback:
             self._message_received_callback(message)
         if not self._streaming_receive:
@@ -895,6 +929,47 @@ def receive_message_batch(self, **kwargs):
         """
         return self._do_retryable_operation(self._receive_message_batch_impl, **kwargs)
 
+    def receive_messages_iter(self, timeout=None, on_message_received=None):
+        """Receive messages by generator. Messages returned in the generator have already been
+        accepted - if you wish to add logic to accept or reject messages based on custom
+        criteria, pass in a callback.
+
+        :param on_message_received: A callback to process messages as they arrive from the
+         service. It takes a single argument, a ~pyamqp.message.Message object.
+        :type on_message_received: callable[~pyamqp.message.Message]
+        """
+        self._message_received_callback = on_message_received
+        return self._message_generator(timeout=timeout)
+
+    def _message_generator(self, timeout=None):
+        """Iterate over processed messages in the receive queue.
+
+        :rtype: generator[~pyamqp.message.Message]
+        """
+        self._timeout_reached = False
+        receiving = True
+        message = None
+        self._last_activity_timestamp = time.time()
+        self._timeout = timeout if timeout else self._timeout
+        try:
+            while receiving and not self._timeout_reached:
+                if self._timeout > 0:
+                    if time.time() - self._last_activity_timestamp >= self._timeout:
+                        self._timeout_reached = True
+
+                if not self._timeout_reached:
+                    receiving = self.do_work()
+
+                while not self._received_messages.empty():
+                    message = self._received_messages.get()
+                    self._last_activity_timestamp = time.time()
+                    self._received_messages.task_done()
+                    yield message
+
+        finally:
+            if self._shutdown:
+                self.close()
+
     @overload
     def settle_messages(
         self,