def request_response(self, payload: Payload) -> Awaitable[Payload]:
logger().debug('%s: request-response: %s', self._log_identifier(),
payload)
requester = RequestResponseRequester(self, payload)
self.register_new_stream(requester).setup()
return requester.run()
def frame_received(self, frame: Frame):
if isinstance(frame, CancelFrame):
self.subscriber.subscription.cancel()
self._finish_stream()
elif isinstance(frame, RequestNFrame):
if self.subscriber.subscription is not None:
self.subscriber.subscription.request(frame.request_n)
else:
logger().warning(
'%s: Received request_n but no publisher provided',
self.__class__.__name__)
elif isinstance(frame, PayloadFrame):
if frame.flags_next:
self.remote_subscriber.on_next(
payload_from_frame(frame),
is_complete=frame.flags_complete)
elif frame.flags_complete:
self.remote_subscriber.on_complete()
if frame.flags_complete:
self.mark_completed_and_finish(received=True)
elif isinstance(frame, ErrorFrame):
self.remote_subscriber.on_error(error_frame_to_exception(frame))
self.mark_completed_and_finish(received=True)
async def feed_subscriber(self):
try:
while True:
payload, is_complete = await self._queue.get()
if self._fragment_size is None:
self._send_to_subscriber(payload, is_complete)
else:
async for fragment in payload_to_n_size_fragments(
BytesIO(payload.data), BytesIO(payload.metadata),
self._fragment_size):
self._send_to_subscriber(
fragment, is_complete and fragment.is_last)
await asyncio.sleep(
self._delay_between_messages.total_seconds())
self._queue.task_done()
if is_complete:
if self._on_complete is not None:
self._on_complete()
break
except asyncio.CancelledError:
logger().debug('Asyncio task canceled: stream_from_generator')
finally:
self._cancel_n_feeder()
async def receive_data(self,
data: bytes,
header_length=3) -> AsyncGenerator[Frame, None]:
self._buffer.extend(data)
total = len(self._buffer)
frame_length_byte_count = header_length
while total >= frame_length_byte_count:
if header_length > 0:
length = struct.unpack(
'>I', b'\x00' + self._buffer[:frame_length_byte_count])[0]
else:
length = len(data)
if total < length + frame_length_byte_count:
return
try:
new_frame = frame.parse_or_ignore(
self._buffer[frame_length_byte_count:length +
frame_length_byte_count])
if new_frame is not None:
yield new_frame
except Exception:
logger().error('Error parsing frame', exc_info=True)
yield InvalidFrame()
self._buffer = self._buffer[length + frame_length_byte_count:]
total -= length + frame_length_byte_count
def request_stream(self,
payload: Payload) -> Union[BackpressureApi, Publisher]:
logger().debug('%s: request-stream: %s', self._log_identifier(),
payload)
requester = RequestStreamRequester(self, payload)
return self.register_new_stream(requester)
async def _sender(self):
try:
try:
transport = await self._current_transport()
self._before_sender()
while self.is_server_alive():
frame = await self._send_queue.get()
await transport.send_frame(frame)
log_frame(frame, self._log_identifier(), 'Sent')
self._send_queue.task_done()
if frame.sent_future is not None:
frame.sent_future.set_result(None)
if self._send_queue.empty():
await transport.on_send_queue_empty()
except RSocketTransportError as exception:
await self._on_connection_lost(exception)
except asyncio.CancelledError:
logger().debug('%s: Asyncio task canceled: sender',
self._log_identifier())
except RSocketTransportError as exception:
await self._on_connection_lost(exception)
except Exception:
logger().error('%s: RSocket error',
self._log_identifier(),
exc_info=True)
raise
finally:
await self._finally_sender()
def metadata_push(self, metadata: bytes) -> Awaitable[None]:
logger().debug('%s: metadata-push: %s', self._log_identifier(),
metadata)
frame = to_metadata_push_frame(metadata)
self.send_frame(frame)
return frame.sent_future
async def _keepalive_send_task(self):
try:
while True:
await asyncio.sleep(self._keep_alive_period.total_seconds())
self._send_new_keepalive()
except asyncio.CancelledError:
logger().debug('%s: Asyncio task canceled: keepalive_send',
self._log_identifier())
async def _close(self, reconnect=False):
if not reconnect:
await cancel_if_task_exists(self._reconnect_task)
else:
logger().debug('%s: Closing before reconnect',
self._log_identifier())
await super().close()
async def close(self):
logger().debug('%s: Closing', self._log_identifier())
self._is_closing = True
await cancel_if_task_exists(self._sender_task)
await cancel_if_task_exists(self._receiver_task)
await self._close_transport()
async def _trigger_next_request_n(subscriber, limit_rate):
try:
while True:
await subscriber.get_next_n.wait()
subscriber.subscription.request(limit_rate)
subscriber.get_next_n.clear()
except asyncio.CancelledError:
logger().debug('Asyncio task canceled: trigger_next_request_n')
async def handle_incoming_ws_messages(self):
try:
async for message in self._message_generator():
async for frame in self._frame_parser.receive_data(message, 0):
self._incoming_frame_queue.put_nowait(frame)
except asyncio.CancelledError:
logger().debug(
'Asyncio task canceled: aiohttp_handle_incoming_ws_messages')
def fire_and_forget(self, payload: Payload) -> Awaitable[None]:
logger().debug('%s: fire-and-forget: %s', self._log_identifier(),
payload)
stream_id = self._allocate_stream()
frame = to_fire_and_forget_frame(stream_id, payload)
self.send_request(frame)
self.finish_stream(stream_id)
return frame.sent_future
async def transport_provider():
while True:
try:
nonlocal client_connection
client_connection = await asyncio.open_connection(host, port)
yield TransportTCP(*client_connection)
except Exception:
logger().error('Client connection error', exc_info=True)
raise
async def handle_incoming_ws_messages(self):
try:
while True:
data = await websocket.receive()
async for frame in self._frame_parser.receive_data(data, 0):
self._incoming_frame_queue.put_nowait(frame)
except asyncio.CancelledError:
logger().debug(
'Asyncio task canceled: quart_handle_incoming_ws_messages')
def request_channel(
self,
payload: Payload,
local_publisher: Optional[Publisher] = None
) -> Union[BackpressureApi, Publisher]:
logger().debug('%s: request-channel: %s', self._log_identifier(),
payload)
requester = RequestChannelRequester(self, payload, local_publisher)
return self.register_new_stream(requester)
async def transport_provider():
try:
for i in range(3):
client_connection = await asyncio.open_connection(
'localhost', unused_tcp_port)
yield TransportTCP(*client_connection)
except Exception:
logger().error('Client connection error', exc_info=True)
raise
def quic_event_received(self, event: QuicEvent) -> None:
logger().debug('Quic event received: %s', event)
if isinstance(event, ConnectionTerminated):
self.frame_queue.put_nowait(RSocketTransportError())
elif isinstance(event, StreamDataReceived):
self.frame_queue.put_nowait(event.data)
super().quic_event_received(event)
async def transport_provider():
try:
yield TransportAioHttpClient(url)
yield FailingTransport()
yield TransportAioHttpClient(url)
except Exception:
logger().error('Client connection error', exc_info=True)
raise
async def transport_provider():
try:
client_connection = await asyncio.open_connection('localhost', port)
yield TransportTCP(*client_connection)
yield FailingTransport()
client_connection = await asyncio.open_connection('localhost', port)
yield TransportTCP(*client_connection)
except Exception:
logger().error('Client connection error', exc_info=True)
raise
async def cancel_if_task_exists(task: Optional[Task]):
if task is not None and not task.done():
task.cancel()
try:
await task
except asyncio.CancelledError:
logger().debug('Asyncio task cancellation error: %s', task)
except RuntimeError:
logger().warning('Runtime error canceling task: %s',
task,
exc_info=True)
async def handle_incoming_ws_messages(self):
await self._connection_ready.wait()
try:
async for msg in self._ws:
if msg.type == aiohttp.WSMsgType.BINARY:
async for frame in self._frame_parser.receive_data(
msg.data, 0):
self._incoming_frame_queue.put_nowait(frame)
except asyncio.CancelledError:
logger().debug('Asyncio task canceled: incoming_data_listener')
except Exception:
self._incoming_frame_queue.put_nowait(RSocketTransportError())
async def _receiver(self):
try:
await self._receiver_listen()
except asyncio.CancelledError:
logger().debug('%s: Asyncio task canceled: receiver',
self._log_identifier())
except RSocketTransportError as exception:
await self._on_connection_lost(exception)
except Exception:
logger().error('%s: Unknown error',
self._log_identifier(),
exc_info=True)
raise
async def _keepalive_timeout_task(self):
try:
while True:
await asyncio.sleep(self._max_lifetime_period.total_seconds())
now = datetime.now()
time_since_last_keepalive = now - self._last_server_keepalive
if time_since_last_keepalive > self._max_lifetime_period:
self._is_server_alive = False
await self._handler.on_keepalive_timeout(
time_since_last_keepalive, self)
except asyncio.CancelledError:
logger().debug('%s: Asyncio task canceled: keepalive_timeout',
self._log_identifier())
async def _close_transport(self):
if self._current_transport().done():
logger().debug('%s: Closing transport', self._log_identifier())
try:
transport = await self._current_transport()
except asyncio.CancelledError:
raise RSocketTransportError()
if transport is not None:
try:
await transport.close()
except Exception:
logger().debug(
'Transport already closed or failed to close',
exc_info=True)
async def _aio_next():
try:
event = await iterator.__anext__()
if isinstance(event, OnNext):
observer.on_next(event.value)
elif isinstance(event, OnError):
observer.on_error(event.exception)
elif isinstance(event, OnCompleted):
observer.on_completed()
except StopAsyncIteration:
pass
except Exception as exception:
logger().error(str(exception), exc_info=True)
observer.on_error(exception)
async def queue_next_n(self):
try:
await self._start_generator()
while True:
n = await self._request_n_queue.get()
async for payload, is_complete in self._generate_next_n(n):
await self._queue.put((payload, is_complete))
except FinishedIterator:
self._queue.put_nowait((Payload(), True))
except asyncio.CancelledError:
logger().debug('Asyncio task canceled: queue_next_n')
except Exception as exception:
self._subscriber.on_error(exception)
self._cancel_feeders()
async def connect(self):
logger().debug('%s: connecting', self._log_identifier())
self._is_closing = False
self._reset_internals()
self._start_tasks()
try:
await self._connect_new_transport()
except Exception as exception:
logger().error('%s: Connection error',
self._log_identifier(),
exc_info=True)
await self._on_connection_lost(exception)
return
return await super().connect()
async def transport_provider():
try:
logging.info('Quic connection lost valid connection 1')
async with rsocket_connect('localhost', port,
configuration=client_configuration) as transport:
yield transport
logging.info('Quic connection lost invalid connection')
yield FailingTransport()
logging.info('Quic connection lost valid connection 2')
async with rsocket_connect('localhost', port,
configuration=client_configuration) as transport:
yield transport
except Exception:
logger().error('Client connection error', exc_info=True)
raise
async def incoming_data_listener(self):
try:
await self._quic_protocol.wait_connected()
while True:
data = await self._incoming_bytes_queue.get()
if isinstance(data, Exception):
self._incoming_frame_queue.put_nowait(data)
return
else:
async for frame in self._frame_parser.receive_data(data, 0):
self._incoming_frame_queue.put_nowait(frame)
except asyncio.CancelledError:
logger().debug('Asyncio task canceled: incoming_data_listener')
except Exception:
self._incoming_frame_queue.put_nowait(RSocketTransportError())
