def on_messages(self, messages: List[SequencedMessage]):
byte_size = 0
for message in messages:
byte_size += message.size_bytes
self._client_tokens += FlowControlRequest(
allowed_bytes=-byte_size, allowed_messages=-len(messages)
)
def to_optional(self) -> Optional[FlowControlRequest]:
allowed_messages = _clamp(self._request.allowed_messages)
allowed_bytes = _clamp(self._request.allowed_bytes)
if allowed_messages == 0 and allowed_bytes == 0:
return None
request = FlowControlRequest()
request._pb.allowed_messages = allowed_messages
request._pb.allowed_bytes = allowed_bytes
return request
async def test_handle_reset(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
):
ack_called_queue = asyncio.Queue()
ack_result_queue = asyncio.Queue()
ack_set_tracker.ack.side_effect = make_queue_waiter(
ack_called_queue, ack_result_queue
)
async with subscriber:
message_1 = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message_1
read_1: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
assert read_1.message_id == "1"
assert read_1.ack_id == ack_id(0, 1)
await subscriber.handle_reset()
ack_set_tracker.clear_and_commit.assert_called_once()
# Message ACKed after reset. Its flow control tokens are refilled
# but offset not committed (verified below after message 2).
read_1.ack()
message_2 = SequencedMessage(cursor=Cursor(offset=2), size_bytes=10)
underlying.read.return_value = message_2
read_2: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1), call(2)])
assert read_2.message_id == "2"
assert read_2.ack_id == ack_id(1, 2)
read_2.ack()
await ack_called_queue.get()
await ack_result_queue.put(None)
underlying.allow_flow.assert_has_calls(
[
call(FlowControlRequest(allowed_messages=1000, allowed_bytes=1000,)),
call(FlowControlRequest(allowed_messages=1, allowed_bytes=5,)),
call(FlowControlRequest(allowed_messages=1, allowed_bytes=10,)),
]
)
ack_set_tracker.ack.assert_has_calls([call(2)])
async def __aenter__(self):
await self._ack_set_tracker.__aenter__()
await self._underlying.__aenter__()
self._looper_future = asyncio.ensure_future(self._looper())
await self._underlying.allow_flow(
FlowControlRequest(
allowed_messages=self._flow_control_settings.messages_outstanding,
allowed_bytes=self._flow_control_settings.bytes_outstanding,
)
)
return self
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
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
async def test_out_of_order_receipt_failure(
subscriber: Subscriber,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
flow = FlowControlRequest(allowed_messages=100, allowed_bytes=100)
message_1 = SequencedMessage(cursor=Cursor(offset=3), size_bytes=5)
message_2 = SequencedMessage(cursor=Cursor(offset=5), size_bytes=10)
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(SubscribeResponse(initial={}))
write_result_queue.put_nowait(None)
async with subscriber:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Send tokens.
flow_fut = asyncio.ensure_future(subscriber.allow_flow(flow))
assert not flow_fut.done()
# Handle the inline write since initial tokens are 100% of outstanding.
await write_called_queue.get()
await write_result_queue.put(None)
await flow_fut
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(flow))]
)
read_fut = asyncio.ensure_future(subscriber.read())
# Send out of order messages to the subscriber.
await read_result_queue.put(as_response([message_2, message_1]))
# Wait for the next read call
await read_called_queue.get()
try:
await read_fut
assert False
except GoogleAPICallError as e:
assert e.grpc_status_code == StatusCode.FAILED_PRECONDITION
pass
async def _handle_ack(self, message: requests.AckRequest):
offset = int(message.ack_id)
await self._underlying.allow_flow(
FlowControlRequest(
allowed_messages=1,
allowed_bytes=self._messages_by_offset[offset].size_bytes,
)
)
del self._messages_by_offset[offset]
try:
await self._ack_set_tracker.ack(offset)
except GoogleAPICallError as e:
self.fail(e)
def _handle_ack(self, message: requests.AckRequest):
flow_control = FlowControlRequest()
flow_control._pb.allowed_messages = 1
flow_control._pb.allowed_bytes = self._messages_by_ack_id[
message.ack_id].size_bytes
self._underlying.allow_flow(flow_control)
del self._messages_by_ack_id[message.ack_id]
# Always refill flow control tokens, but do not commit offsets from outdated generations.
ack_id = _AckId.parse(message.ack_id)
if ack_id.generation == self._ack_generation_id:
try:
self._ack_set_tracker.ack(ack_id.offset)
except GoogleAPICallError as e:
self.fail(e)
async def stop_processing(self, error: GoogleAPICallError):
await self._stop_loopers()
if is_reset_signal(error):
# Discard undelivered messages and refill flow control tokens.
while not self._message_queue.empty():
batch: List[SequencedMessage.meta.pb] = self._message_queue.get_nowait()
allowed_bytes = sum(message.size_bytes for message in batch)
self._outstanding_flow_control.add(
FlowControlRequest(
allowed_messages=len(batch), allowed_bytes=allowed_bytes,
)
)
await self._reset_handler.handle_reset()
self._last_received_offset = None
async def reinitialize(
self,
connection: Connection[SubscribeRequest, SubscribeResponse],
last_error: Optional[GoogleAPICallError],
):
self._reinitializing = True
await self._stop_loopers()
if last_error and is_reset_signal(last_error):
# Discard undelivered messages and refill flow control tokens.
while not self._message_queue.empty():
msg = self._message_queue.get_nowait()
self._outstanding_flow_control.add(
FlowControlRequest(
allowed_messages=1,
allowed_bytes=msg.size_bytes,
))
await self._reset_handler.handle_reset()
self._last_received_offset = None
initial = deepcopy(self._base_initial)
if self._last_received_offset is not None:
initial.initial_location = SeekRequest(cursor=Cursor(
offset=self._last_received_offset + 1))
else:
initial.initial_location = SeekRequest(
named_target=SeekRequest.NamedTarget.COMMITTED_CURSOR)
await connection.write(SubscribeRequest(initial=initial))
response = await connection.read()
if "initial" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid initial response on the subscribe stream."
))
return
tokens = self._outstanding_flow_control.request_for_restart()
if tokens is not None:
await connection.write(SubscribeRequest(flow_control=tokens))
self._reinitializing = False
self._start_loopers()
async def test_message_receipt(
subscriber: Subscriber,
default_connection,
base_initial_subscribe,
initial_request,
asyncio_sleep,
sleep_queues,
):
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
flow = FlowControlRequest(allowed_messages=100, allowed_bytes=100)
message_1 = SequencedMessage(cursor=Cursor(offset=3), size_bytes=5)
message_2 = SequencedMessage(cursor=Cursor(offset=5), size_bytes=10)
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(SubscribeResponse(initial={}))
write_result_queue.put_nowait(None)
async with subscriber:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Send tokens.
flow_fut = asyncio.ensure_future(subscriber.allow_flow(flow))
assert not flow_fut.done()
# Handle the inline write since initial tokens are 100% of outstanding.
await write_called_queue.get()
await write_result_queue.put(None)
await flow_fut
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(flow))]
)
message1_fut = asyncio.ensure_future(subscriber.read())
# Send messages to the subscriber.
await read_result_queue.put(as_response([message_1, message_2]))
# Wait for the next read call
await read_called_queue.get()
assert (await message1_fut) == message_1
assert (await subscriber.read()) == message_2
# Fail the connection with a retryable error
await read_called_queue.get()
await read_result_queue.put(InternalServerError("retryable"))
await sleep_queues[_MIN_BACKOFF_SECS].called.get()
await sleep_queues[_MIN_BACKOFF_SECS].results.put(None)
# Reinitialization
await write_called_queue.get()
seek_to_cursor_request = make_initial_subscribe_request(
base_initial_subscribe,
SeekRequest(cursor=Cursor(offset=message_2.cursor.offset + 1)),
)
default_connection.write.assert_has_calls(
[
call(initial_request),
call(as_request(flow)),
call(seek_to_cursor_request),
]
)
await write_result_queue.put(None)
await read_called_queue.get()
await read_result_queue.put(SubscribeResponse(initial={}))
# Re-sending flow tokens on the new stream.
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls(
[
call(initial_request),
call(as_request(flow)),
call(seek_to_cursor_request),
call(
as_request(
FlowControlRequest(allowed_messages=98, allowed_bytes=85)
)
),
]
)
def __init__(self):
self.request = FlowControlRequest()
async def test_basic_flow_control_after_timeout(
subscriber: Subscriber,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
sleep_called = sleep_queues[FLUSH_SECONDS].called
sleep_results = sleep_queues[FLUSH_SECONDS].results
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
flow_1 = FlowControlRequest(allowed_messages=100, allowed_bytes=100)
flow_2 = FlowControlRequest(allowed_messages=5, allowed_bytes=10)
flow_3 = FlowControlRequest(allowed_messages=10, allowed_bytes=5)
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(SubscribeResponse(initial={}))
write_result_queue.put_nowait(None)
async with subscriber:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Send tokens.
flow_fut1 = asyncio.ensure_future(subscriber.allow_flow(flow_1))
assert not flow_fut1.done()
# Handle the inline write since initial tokens are 100% of outstanding.
await write_called_queue.get()
await write_result_queue.put(None)
await flow_fut1
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(flow_1))]
)
# Should complete without writing to the connection
await subscriber.allow_flow(flow_2)
await subscriber.allow_flow(flow_3)
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_called_with(FLUSH_SECONDS)
# Handle the connection write
await sleep_results.put(None)
await write_called_queue.get()
await write_result_queue.put(None)
# Called with aggregate
default_connection.write.assert_has_calls(
[
call(initial_request),
call(as_request(flow_1)),
call(
as_request(
FlowControlRequest(allowed_messages=15, allowed_bytes=15)
)
),
]
)
async def test_flow_resent_on_restart(
subscriber: Subscriber,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
flow_1 = FlowControlRequest(allowed_messages=100, allowed_bytes=100)
flow_2 = FlowControlRequest(allowed_messages=5, allowed_bytes=10)
flow_3 = FlowControlRequest(allowed_messages=10, allowed_bytes=5)
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(SubscribeResponse(initial={}))
write_result_queue.put_nowait(None)
async with subscriber:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Send tokens.
flow_fut1 = asyncio.ensure_future(subscriber.allow_flow(flow_1))
assert not flow_fut1.done()
# Handle the inline write since initial tokens are 100% of outstanding.
await write_called_queue.get()
await write_result_queue.put(None)
await flow_fut1
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(flow_1))]
)
# Should complete without writing to the connection
await subscriber.allow_flow(flow_2)
await subscriber.allow_flow(flow_3)
# Fail the connection with a retryable error
await read_called_queue.get()
await read_result_queue.put(InternalServerError("retryable"))
await sleep_queues[_MIN_BACKOFF_SECS].called.get()
await sleep_queues[_MIN_BACKOFF_SECS].results.put(None)
# Reinitialization
await write_called_queue.get()
await write_result_queue.put(None)
await read_called_queue.get()
await read_result_queue.put(SubscribeResponse(initial={}))
# Re-sending flow tokens on the new stream
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls(
[
call(initial_request),
call(as_request(flow_1)),
call(initial_request),
call(
as_request(
FlowControlRequest(allowed_messages=115, allowed_bytes=115)
)
),
]
)
def initial_flow_request(flow_control_settings):
return FlowControlRequest(
allowed_messages=flow_control_settings.messages_outstanding,
allowed_bytes=flow_control_settings.bytes_outstanding,
)
async def test_handle_reset_signal(
subscriber: Subscriber,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
reset_handler,
):
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
flow = FlowControlRequest(allowed_messages=100, allowed_bytes=100)
message_1 = SequencedMessage(cursor=Cursor(offset=2), size_bytes=5)
message_2 = SequencedMessage(cursor=Cursor(offset=4), size_bytes=10)
# Ensure messages with earlier offsets can be handled post-reset.
message_3 = SequencedMessage(cursor=Cursor(offset=1), size_bytes=20)
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(SubscribeResponse(initial={}))
write_result_queue.put_nowait(None)
async with subscriber:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Send tokens.
flow_fut = asyncio.ensure_future(subscriber.allow_flow(flow))
assert not flow_fut.done()
# Handle the inline write since initial tokens are 100% of outstanding.
await write_called_queue.get()
await write_result_queue.put(None)
await flow_fut
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(flow))]
)
# Send messages to the subscriber.
await read_result_queue.put(as_response([message_1, message_2]))
# Read one message.
await read_called_queue.get()
assert (await subscriber.read()) == message_1
# Fail the connection with an error containing the RESET signal.
await read_called_queue.get()
await read_result_queue.put(make_reset_signal())
await sleep_queues[_MIN_BACKOFF_SECS].called.get()
await sleep_queues[_MIN_BACKOFF_SECS].results.put(None)
# Reinitialization.
await write_called_queue.get()
await write_result_queue.put(None)
await read_called_queue.get()
await read_result_queue.put(SubscribeResponse(initial={}))
# Re-sending flow tokens on the new stream.
await write_called_queue.get()
await write_result_queue.put(None)
reset_handler.handle_reset.assert_called_once()
default_connection.write.assert_has_calls(
[
call(initial_request),
call(as_request(flow)),
call(initial_request),
call(
as_request(
# Tokens for undelivered message_2 refilled.
FlowControlRequest(allowed_messages=99, allowed_bytes=95)
)
),
]
)
# Ensure the subscriber accepts an earlier message.
await read_result_queue.put(as_response([message_3]))
await read_called_queue.get()
assert (await subscriber.read()) == message_3