[ServiceBus] release message in background (Azure#23285)

* release message in background

* decouple keep receiving and fix async

* small fix

* fix mypy and pylint

* review feedback

* fix pylint and update changelog

* update test to reflect the latest change

* add tests and replace uamqp keep alive

* fix impl bug

* add test non-releasing behavior

* update release date

Author: yunhaoling

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -1,15 +1,11 @@
 # Release History
 
-## 7.6.1 (Unreleased)
-
-### Features Added
-
-### Breaking Changes
-
-### Bugs Fixed
+## 7.6.1 (2022-04-06)
 
 ### Other Changes
 
+- Improved receiving by releasing messages from internal buffer when the `prefetch_count` of `ServiceBusReceiver`  is set 0 and there is no active receive call, this helps avoid receiving expired messages and incrementing delivery count of a message.
+
 ## 7.6.0 (2022-02-10)
 
 ### Features Added

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

@@ -176,3 +176,11 @@ def _on_attach(self, source, target, properties, error):
     def _populate_message_properties(self, message):
         if self._session:
             message[MGMT_REQUEST_SESSION_ID] = self._session_id
+
+    def _enhanced_message_received(self, message):
+        # pylint: disable=protected-access
+        self._handler._was_message_received = True
+        if self._receive_context.is_set():
+            self._handler._received_messages.put(message)
+        else:
+            message.release()

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

@@ -2,6 +2,8 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
+# pylint:disable=too-many-lines
+import threading
 import time
 import logging
 import functools
@@ -197,6 +199,7 @@ def __init__(
         self._session = (
             None if self._session_id is None else ServiceBusSession(self._session_id, self)
         )
+        self._receive_context = threading.Event()
 
     def __iter__(self):
         return self._iter_contextual_wrapper()
@@ -238,26 +241,34 @@ def _inner_next(self):
 
     def __next__(self):
         # Normally this would wrap the yield of the iter, but for a direct next call we just trace imperitively.
-        message = self._inner_next()
-        links = get_receive_links(message)
-        with receive_trace_context_manager(self, links=links):
-            return message
+        try:
+            self._receive_context.set()
+            message = self._inner_next()
+            links = get_receive_links(message)
+            with receive_trace_context_manager(self, links=links):
+                return message
+        finally:
+            self._receive_context.clear()
 
     next = __next__  # for python2.7
 
     def _iter_next(self):
-        self._open()
-        if not self._message_iter:
-            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
-            and self._receive_mode != ServiceBusReceiveMode.RECEIVE_AND_DELETE
-        ):
-            self._auto_lock_renewer.register(self, message)
-        return message
+        try:
+            self._receive_context.set()
+            self._open()
+            if not self._message_iter:
+                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
+                and self._receive_mode != ServiceBusReceiveMode.RECEIVE_AND_DELETE
+            ):
+                self._auto_lock_renewer.register(self, message)
+            return message
+        finally:
+            self._receive_context.clear()
 
     @classmethod
     def _from_connection_string(cls, conn_str, **kwargs):
@@ -338,9 +349,12 @@ def _create_handler(self, auth):
             else None,
             timeout=self._max_wait_time * 1000 if self._max_wait_time else 0,
             prefetch=self._prefetch_count,
-            keep_alive_interval=self._config.keep_alive,
+            # If prefetch is 1, then keep_alive coroutine serves as keep receiving for releasing messages
+            keep_alive_interval=self._config.keep_alive if self._prefetch_count != 1 else 5,
             shutdown_after_timeout=False,
         )
+        if self._prefetch_count == 1:
+            self._handler._message_received = self._enhanced_message_received  # pylint: disable=protected-access
 
     def _open(self):
         # pylint: disable=protected-access
@@ -366,65 +380,68 @@ def _open(self):
     def _receive(self, max_message_count=None, timeout=None):
         # type: (Optional[int], Optional[float]) -> List[ServiceBusReceivedMessage]
         # pylint: disable=protected-access
-        self._open()
-
-        amqp_receive_client = self._handler
-        received_messages_queue = amqp_receive_client._received_messages
-        max_message_count = max_message_count or self._prefetch_count
-        timeout_ms = (
-            1000 * (timeout or self._max_wait_time)
-            if (timeout or self._max_wait_time)
-            else 0
-        )
-        abs_timeout_ms = (
-            amqp_receive_client._counter.get_current_ms() + timeout_ms
-            if timeout_ms
-            else 0
-        )
-
-        batch = []  # type: List[Message]
-        while not received_messages_queue.empty() and len(batch) < max_message_count:
-            batch.append(received_messages_queue.get())
-            received_messages_queue.task_done()
-        if len(batch) >= max_message_count:
-            return [self._build_message(message) for message in batch]
-
-        # Dynamically issue link credit if max_message_count > 1 when the prefetch_count is the default value 1
-        if max_message_count and self._prefetch_count == 1 and max_message_count > 1:
-            link_credit_needed = max_message_count - len(batch)
-            amqp_receive_client.message_handler.reset_link_credit(link_credit_needed)
-
-        first_message_received = expired = False
-        receiving = True
-        while receiving and not expired and len(batch) < max_message_count:
-            while receiving and received_messages_queue.qsize() < max_message_count:
-                if (
-                    abs_timeout_ms
-                    and amqp_receive_client._counter.get_current_ms() > abs_timeout_ms
-                ):
-                    expired = True
-                    break
-                before = received_messages_queue.qsize()
-                receiving = amqp_receive_client.do_work()
-                received = received_messages_queue.qsize() - before
-                if (
-                    not first_message_received
-                    and received_messages_queue.qsize() > 0
-                    and received > 0
-                ):
-                    # first message(s) received, continue receiving for some time
-                    first_message_received = True
-                    abs_timeout_ms = (
-                        amqp_receive_client._counter.get_current_ms()
-                        + self._further_pull_receive_timeout_ms
-                    )
-            while (
-                not received_messages_queue.empty() and len(batch) < max_message_count
-            ):
+        try:
+            self._receive_context.set()
+            self._open()
+
+            amqp_receive_client = self._handler
+            received_messages_queue = amqp_receive_client._received_messages
+            max_message_count = max_message_count or self._prefetch_count
+            timeout_ms = (
+                1000 * (timeout or self._max_wait_time)
+                if (timeout or self._max_wait_time)
+                else 0
+            )
+            abs_timeout_ms = (
+                amqp_receive_client._counter.get_current_ms() + timeout_ms
+                if timeout_ms
+                else 0
+            )
+            batch = []  # type: List[Message]
+            while not received_messages_queue.empty() and len(batch) < max_message_count:
                 batch.append(received_messages_queue.get())
                 received_messages_queue.task_done()
+            if len(batch) >= max_message_count:
+                return [self._build_message(message) for message in batch]
+
+            # Dynamically issue link credit if max_message_count > 1 when the prefetch_count is the default value 1
+            if max_message_count and self._prefetch_count == 1 and max_message_count > 1:
+                link_credit_needed = max_message_count - len(batch)
+                amqp_receive_client.message_handler.reset_link_credit(link_credit_needed)
+
+            first_message_received = expired = False
+            receiving = True
+            while receiving and not expired and len(batch) < max_message_count:
+                while receiving and received_messages_queue.qsize() < max_message_count:
+                    if (
+                        abs_timeout_ms
+                        and amqp_receive_client._counter.get_current_ms() > abs_timeout_ms
+                    ):
+                        expired = True
+                        break
+                    before = received_messages_queue.qsize()
+                    receiving = amqp_receive_client.do_work()
+                    received = received_messages_queue.qsize() - before
+                    if (
+                        not first_message_received
+                        and received_messages_queue.qsize() > 0
+                        and received > 0
+                    ):
+                        # first message(s) received, continue receiving for some time
+                        first_message_received = True
+                        abs_timeout_ms = (
+                            amqp_receive_client._counter.get_current_ms()
+                            + self._further_pull_receive_timeout_ms
+                        )
+                while (
+                    not received_messages_queue.empty() and len(batch) < max_message_count
+                ):
+                    batch.append(received_messages_queue.get())
+                    received_messages_queue.task_done()
 
-        return [self._build_message(message) for message in batch]
+            return [self._build_message(message) for message in batch]
+        finally:
+            self._receive_context.clear()
 
     def _settle_message_with_retry(
         self,