feat: implement Flow control batcher (#15)

* feat: Implement FlowControlBatcher This handles aggregating flow control requests without allowing them to get above the max int64 value. * Use correct request for comparison.

Author: dpcollins-google

google/cloud/pubsublite/internal/wire/flow_control_batcher.py

@@ -0,0 +1,67 @@
+from typing import NamedTuple, List, Optional
+
+from google.cloud.pubsublite_v1 import FlowControlRequest, SequencedMessage
+
+_EXPEDITE_BATCH_REQUEST_RATIO = 0.5
+_MAX_INT64 = 0x7FFFFFFFFFFFFFFF
+
+
+class _AggregateRequest:
+ request: FlowControlRequest
+
+ def __init__(self):
+ self.request = FlowControlRequest()
+
+ def __add__(self, other: FlowControlRequest):
+ self.request.allowed_bytes += other.allowed_bytes
+ self.request.allowed_bytes = min(self.request.allowed_bytes, _MAX_INT64)
+ self.request.allowed_messages += other.allowed_messages
+ self.request.allowed_messages = min(self.request.allowed_messages, _MAX_INT64)
+ return self
+
+
+def _exceeds_expedite_ratio(pending: int, client: int):
+ if client <= 0:
+ return False
+ return (pending/client) >= _EXPEDITE_BATCH_REQUEST_RATIO
+
+
+def _to_optional(req: FlowControlRequest) -> Optional[FlowControlRequest]:
+ if req.allowed_messages == 0 and req.allowed_bytes == 0:
+ return None
+ return req
+
+
+class FlowControlBatcher:
+ _client_tokens: _AggregateRequest
+ _pending_tokens: _AggregateRequest
+
+ def __init__(self):
+ self._client_tokens = _AggregateRequest()
+ self._pending_tokens = _AggregateRequest()
+
+ def add(self, request: FlowControlRequest):
+ self._client_tokens += request
+ self._pending_tokens += request
+
+ def on_messages(self, messages: List[SequencedMessage]):
+ byte_size = sum(message.size_bytes for message in messages)
+ self._client_tokens += FlowControlRequest(allowed_bytes=-byte_size, allowed_messages=-len(messages))
+
+ def request_for_restart(self) -> Optional[FlowControlRequest]:
+ self._pending_tokens = _AggregateRequest()
+ return _to_optional(self._client_tokens.request)
+
+ def release_pending_request(self) -> Optional[FlowControlRequest]:
+ request = self._pending_tokens.request
+ self._pending_tokens = _AggregateRequest()
+ return _to_optional(request)
+
+ def should_expedite(self):
+ pending_request = self._pending_tokens.request
+ client_request = self._client_tokens.request
+ if _exceeds_expedite_ratio(pending_request.allowed_bytes, client_request.allowed_bytes):
+ return True
+ if _exceeds_expedite_ratio(pending_request.allowed_messages, client_request.allowed_messages):
+ return True
+ return False

tests/unit/pubsublite/internal/wire/flow_control_batcher_test.py

@@ -0,0 +1,28 @@
+from google.cloud.pubsublite.internal.wire.flow_control_batcher import FlowControlBatcher
+from google.cloud.pubsublite_v1 import FlowControlRequest, SequencedMessage
+
+
+def test_restart_clears_send():
+ batcher = FlowControlBatcher()
+ batcher.add(FlowControlRequest(allowed_bytes=10, allowed_messages=3))
+ assert batcher.should_expedite()
+ to_send = batcher.release_pending_request()
+ assert to_send.allowed_bytes == 10
+ assert to_send.allowed_messages == 3
+ restart_1 = batcher.request_for_restart()
+ assert restart_1.allowed_bytes == 10
+ assert restart_1.allowed_messages == 3
+ assert not batcher.should_expedite()
+ assert batcher.release_pending_request() is None
+
+
+def test_add_remove():
+ batcher = FlowControlBatcher()
+ batcher.add(FlowControlRequest(allowed_bytes=10, allowed_messages=3))
+ restart_1 = batcher.request_for_restart()
+ assert restart_1.allowed_bytes == 10
+ assert restart_1.allowed_messages == 3
+ batcher.on_messages([SequencedMessage(size_bytes=2), SequencedMessage(size_bytes=3)])
+ restart_2 = batcher.request_for_restart()
+ assert restart_2.allowed_bytes == 5
+ assert restart_2.allowed_messages == 1