[SB] Missing Link Credit Subtraction (#37427)

* move back down but subtract link credit

* date

* nit

* link credit math

* Update sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/receiver.py

* await

* comments

* pylint

* eventhub + mock test

* pin dev req

* Update sdk/servicebus/azure-servicebus/CHANGELOG.md

* mock only outgoing_flow

Author: l0lawrence

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

@@ -152,7 +152,7 @@ async def _keep_alive_async(self):
                 elapsed_time = current_time - start_time
                 if elapsed_time >= self._keep_alive_interval:
                     await asyncio.shield(self._connection.listen(wait=self._socket_timeout,
-                        batch=self._link.current_link_credit))
+                        batch=self._link.total_link_credit))
                     start_time = current_time
                 await asyncio.sleep(1)
         except Exception as e:  # pylint: disable=broad-except
@@ -759,7 +759,7 @@ async def _client_run_async(self, **kwargs):
         :rtype: bool
         """
         try:
-            if self._link.current_link_credit <= 0:
+            if self._link.total_link_credit <= 0:
                 await self._link.flow(link_credit=self._link_credit)
             await self._connection.listen(wait=self._socket_timeout, **kwargs)
         except ValueError:

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_link_async.py

@@ -101,6 +101,7 @@ def __init__(
         self._on_link_state_change = kwargs.get("on_link_state_change")
         self._on_attach = kwargs.get("on_attach")
         self._error: Optional[AMQPLinkError] = None
+        self.total_link_credit = self.link_credit
 
     async def __aenter__(self) -> "Link":
         await self.attach()
@@ -273,5 +274,19 @@ async def detach(self, close: bool = False, error: Optional[AMQPError] = None) -
             await self._set_state(LinkState.DETACHED)
 
     async def flow(self, *, link_credit: Optional[int] = None, **kwargs) -> None:
-        self.current_link_credit = link_credit if link_credit is not None else self.link_credit
-        await self._outgoing_flow(**kwargs)
+        # Given the desired link credit `link_credit`, the link credit sent via
+        # FlowFrame is calculated as follows: The link credit to flow on the wire
+        # `self.current_link_credit` is the desired link credit `link_credit`
+        # minus the current link credit on the wire `self.total_link_credit`.
+        self.current_link_credit = link_credit - self.total_link_credit \
+            if link_credit is not None else self.link_credit
+
+        # If the link credit to flow is greater than 0 (i.e the desired link credit
+        # is greater than the current link credit on the wire), then we will send a
+        # flow to issue more link credit. Otherwise link credit on the wire is sufficient.
+        if self.current_link_credit > 0:
+            # Calculate the total link credit on the wire, by adding the credit
+            # we will flow to the total link credit.
+            self.total_link_credit = self.current_link_credit + self.total_link_credit \
+                if link_credit is not None else self.link_credit
+            await self._outgoing_flow(**kwargs)

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_receiver_async.py

@@ -62,8 +62,9 @@ async def _incoming_transfer(self, frame):
         self.received_delivery_id = frame[1] # delivery_id
         # If more is false --> this is the last frame of the message
         if not frame[5]:
-            self.current_link_credit -= 1
             self.delivery_count += 1
+            self.current_link_credit -= 1
+            self.total_link_credit -= 1
         if self.received_delivery_id is not None:
             self._first_frame = frame
         if not self.received_delivery_id and not self._received_payload:

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

@@ -239,7 +239,7 @@ def _keep_alive(self):
                 current_time = time.time()
                 elapsed_time = current_time - start_time
                 if elapsed_time >= self._keep_alive_interval:
-                    self._connection.listen(wait=self._socket_timeout, batch=self._link.current_link_credit)
+                    self._connection.listen(wait=self._socket_timeout, batch=self._link.total_link_credit)
                     start_time = current_time
                 time.sleep(1)
         except Exception as e:  # pylint: disable=broad-except
@@ -866,7 +866,7 @@ def _client_run(self, **kwargs):
         :rtype: bool
         """
         try:
-            if self._link.current_link_credit <= 0:
+            if self._link.total_link_credit <= 0:
                 self._link.flow(link_credit=self._link_credit)
             self._connection.listen(wait=self._socket_timeout, **kwargs)
         except ValueError:

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/link.py

@@ -97,6 +97,7 @@ def __init__(
         self._on_link_state_change = kwargs.get("on_link_state_change")
         self._on_attach = kwargs.get("on_attach")
         self._error: Optional[AMQPLinkError] = None
+        self.total_link_credit = self.link_credit
 
     def __enter__(self) -> "Link":
         self.attach()
@@ -268,5 +269,18 @@ def detach(self, close: bool = False, error: Optional[AMQPError] = None) -> None
             self._set_state(LinkState.DETACHED)
 
     def flow(self, *, link_credit: Optional[int] = None, **kwargs: Any) -> None:
-        self.current_link_credit = link_credit if link_credit is not None else self.link_credit
-        self._outgoing_flow(**kwargs)
+        # Given the desired link credit `link_credit`, the link credit sent via
+        # FlowFrame is calculated as follows: The link credit to flow on the wire
+        # `self.current_link_credit` is the desired link credit
+        # `link_credit` minus the current link credit on the wire `self.total_link_credit`.
+        self.current_link_credit = link_credit - self.total_link_credit if link_credit is not None \
+            else self.link_credit
+
+        # If the link credit to flow is greater than 0 (i.e the desired link credit is greater than
+        # the current link credit on the wire), then we will send a flow to issue more link credit.
+        # Otherwise link credit on the wire is sufficient.
+        if self.current_link_credit > 0:
+            # Calculate the total link credit on the wire, by adding the credit we will flow to the total link credit.
+            self.total_link_credit = self.current_link_credit + self.total_link_credit if link_credit is not None \
+                else self.link_credit
+            self._outgoing_flow(**kwargs)

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/receiver.py

@@ -60,6 +60,7 @@ def _incoming_transfer(self, frame):
         # If more is false --> this is the last frame of the message
         if not frame[5]:
             self.current_link_credit -= 1
+            self.total_link_credit -=1
             self.delivery_count += 1
         self.received_delivery_id = frame[1]  # delivery_id
         if self.received_delivery_id is not None:

sdk/eventhub/azure-eventhub/dev_requirements.txt

@@ -1,6 +1,6 @@
 -e ../../../tools/azure-sdk-tools
 ../../core/azure-core
--e ../../identity/azure-identity
+azure-identity~=1.17.0
 azure-mgmt-eventhub<=10.1.0
 azure-mgmt-resource==20.0.0
 aiohttp; python_version < '3.12'

sdk/eventhub/azure-eventhub/tests/pyamqp_tests/unittest/test_link.py

@@ -111,4 +111,45 @@ def test_receive_transfer_continuation_frame():
     assert link.delivery_count == 1
     link._incoming_transfer(transfer_frame_three)
     assert link.current_link_credit == 1
-    assert link.delivery_count == 2
+    assert link.delivery_count == 2
+
+
+def test_receive_transfer_and_flow():
+    def mock_outgoing(): pass
+    session = None
+    link = ReceiverLink(
+        session,
+        3,
+        source_address="test_source",
+        target_address="test_target",
+        network_trace=False,
+        network_trace_params={},
+        on_transfer=Mock(),
+    )
+
+    link._outgoing_flow =  mock_outgoing
+    link.total_link_credit = 0 # Set the total link credit to 0 to start, no credit on the wire
+
+    link.flow(link_credit=100) # Send a flow frame with desired link credit of 100
+
+    # frame: handle, delivery_id, delivery_tag, message_format, settled, more, rcv_settle_mode, state, resume, aborted, batchable, payload
+    transfer_frame_one = [3, 0, b'/blah', 0, True, False, None, None, None, None, False, ""]
+    transfer_frame_two = [3, 1, b'/blah', 0, True, False, None, None, None, None, False, ""]
+    transfer_frame_three = [3, 2, b'/blah', 0, True, False, None, None, None, None, False, ""]
+
+    link._incoming_transfer(transfer_frame_one)
+    assert link.current_link_credit == 99
+    assert link.total_link_credit == 99
+
+    # Only received 1 transfer frame per receive call, we set desired link credit again
+    # this will send a flow of 1
+    link.flow(link_credit=100)
+    assert link.current_link_credit == 1
+    assert link.total_link_credit == 100
+
+    link._incoming_transfer(transfer_frame_two)
+    assert link.current_link_credit == 0
+    assert link.total_link_credit == 99
+    link._incoming_transfer(transfer_frame_three)
+    assert link.current_link_credit == -1
+    assert link.total_link_credit == 98

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -1,6 +1,6 @@
 # Release History
 
-## 7.12.3 (2024-09-17)
+## 7.12.3 (2024-09-19)
 
 ### Bugs Fixed
 

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

@@ -152,7 +152,7 @@ async def _keep_alive_async(self):
                 elapsed_time = current_time - start_time
                 if elapsed_time >= self._keep_alive_interval:
                     await asyncio.shield(self._connection.listen(wait=self._socket_timeout,
-                        batch=self._link.current_link_credit))
+                        batch=self._link.total_link_credit))
                     start_time = current_time
                 await asyncio.sleep(1)
         except Exception as e:  # pylint: disable=broad-except
@@ -759,7 +759,7 @@ async def _client_run_async(self, **kwargs):
         :rtype: bool
         """
         try:
-            if self._link.current_link_credit <= 0:
+            if self._link.total_link_credit <= 0:
                 await self._link.flow(link_credit=self._link_credit)
             await self._connection.listen(wait=self._socket_timeout, **kwargs)
         except ValueError: