async def reinitialize(self, connection: Connection[SubscribeRequest,
SubscribeResponse]):
self._reinitializing = True
await self._stop_loopers()
await connection.write(SubscribeRequest(initial=self._initial))
response = await connection.read()
if "initial" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid initial response on the subscribe stream."
))
return
if self._last_received_offset is not None:
# Perform a seek to get the next message after the one we received.
await connection.write(
SubscribeRequest(seek=SeekRequest(cursor=Cursor(
offset=self._last_received_offset + 1))))
seek_response = await connection.read()
if "seek" not in seek_response:
self._connection.fail(
FailedPrecondition(
"Received an invalid seek 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 reinitialize(
self, connection: Connection[SubscribeRequest, SubscribeResponse]
):
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._start_loopers()
async def _try_send_tokens(self):
req = self._outstanding_flow_control.release_pending_request()
if req is None:
return
try:
await self._connection.write(SubscribeRequest(flow_control=req))
except GoogleAPICallError:
# May be transient, in which case these tokens will be resent.
pass
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()
def make_initial_subscribe_request(
base: InitialSubscribeRequest, location: SeekRequest
):
initial_subscribe = deepcopy(base)
initial_subscribe.initial_location = location
return SubscribeRequest(initial=initial_subscribe)
def as_request(flow: FlowControlRequest):
req = SubscribeRequest()
req.flow_control = flow
return req
def initial_request():
return SubscribeRequest(initial=InitialSubscribeRequest(subscription="mysub"))
async def test_message_receipt(
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))]
)
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()
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(flow)), call(initial_request)]
)
await write_result_queue.put(None)
await read_called_queue.get()
await read_result_queue.put(SubscribeResponse(initial={}))
# Sends fetch offset seek on the stream, and checks the response.
seek_req = SubscribeRequest(
seek=SeekRequest(cursor=Cursor(offset=message_2.cursor.offset + 1))
)
await write_called_queue.get()
default_connection.write.assert_has_calls(
[
call(initial_request),
call(as_request(flow)),
call(initial_request),
call(seek_req),
]
)
await write_result_queue.put(None)
await read_called_queue.get()
await read_result_queue.put(SubscribeResponse(seek={}))
# 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(initial_request),
call(seek_req),
call(
as_request(
FlowControlRequest(allowed_messages=98, allowed_bytes=85)
)
),
]
)