def stream_protocol_messages(
self,
protocol_identifier: Union[ProtocolAPI, Type[ProtocolAPI]],
) -> AsyncIterator[CommandAPI[Any]]:
"""
Stream the messages for the specified protocol.
"""
if isinstance(protocol_identifier, ProtocolAPI):
protocol_class = type(protocol_identifier)
elif isinstance(protocol_identifier, type) and issubclass(
protocol_identifier, ProtocolAPI):
protocol_class = protocol_identifier
else:
raise TypeError("Unknown protocol identifier: {protocol}")
if not self.has_protocol(protocol_class):
raise UnknownProtocol(f"Unknown protocol '{protocol_class}'")
if self._protocol_locks[protocol_class].locked():
raise Exception(
f"Streaming lock for {protocol_class} is not free.")
elif not self._multiplex_lock.locked():
raise Exception("Not multiplexed.")
return self._stream_protocol_messages(protocol_class)
def stream_protocol_messages(self,
protocol_identifier: Union[ProtocolAPI, Type[ProtocolAPI]],
) -> AsyncIterator[Tuple[CommandAPI, Payload]]:
"""
Stream the messages for the specified protocol.
"""
if isinstance(protocol_identifier, Protocol):
protocol_class = type(protocol_identifier)
elif isinstance(protocol_identifier, type) and issubclass(protocol_identifier, Protocol):
protocol_class = protocol_identifier
else:
raise TypeError("Unknown protocol identifier: {protocol}")
if not self.has_protocol(protocol_class):
raise UnknownProtocol(f"Unknown protocol '{protocol_class}'")
if self._protocol_locks.is_locked(protocol_class):
raise Exception(f"Streaming lock for {protocol_class} is not free.")
elif not self._multiplex_lock.locked():
raise Exception("Not multiplexed.")
# Mostly a sanity check but this ensures we do better than accidentally
# raising an attribute error in whatever race conditions or edge cases
# potentially make the `_multiplex_token` unavailable.
if not hasattr(self, '_multiplex_token'):
raise Exception("No cancel token found for multiplexing.")
# We do the wait_iter here so that the call sites in the handshakers
# that use this don't need to be aware of cancellation tokens.
return self.wait_iter(
self._stream_protocol_messages(protocol_class),
token=self._multiplex_token,
)
def get_protocol_by_type(self, protocol_class: Type[TProtocol]) -> TProtocol:
if issubclass(protocol_class, BaseP2PProtocol):
return cast(TProtocol, self._base_protocol)
for protocol in self._protocols:
if type(protocol) is protocol_class:
return cast(TProtocol, protocol)
raise UnknownProtocol(f"No protocol found with type {protocol_class}")
def __init__(self, connection: ConnectionAPI) -> None:
self._connection = connection
available_protocols = self._connection.get_multiplexer().get_protocols()
for attr, exchange_cls in self._exchange_config.items():
if hasattr(self, attr):
raise AttributeError(
f"Unable to set manager on attribute `{attr}` which is already "
f"present on the class: {getattr(self, attr)}"
)
# determine which protocol should be used to issue requests
supported_protocols = tuple(
protocol
for protocol in available_protocols
if protocol.supports_command(exchange_cls.get_request_cmd_type())
)
if len(supported_protocols) == 1:
protocol_type = type(supported_protocols[0])
elif not supported_protocols:
raise UnknownProtocol(
f"Connection does not have any protocols that support the "
f"request command: {exchange_cls.get_request_cmd_type()}"
)
elif len(supported_protocols) > 1:
raise ValidationError(
f"Could not determine appropriate protocol for command: "
f"{exchange_cls.get_request_cmd_type()}. Command was found in the "
f"protocols {supported_protocols}"
)
else:
raise Exception("This code path should be unreachable")
if not protocol_type.supports_command(exchange_cls.get_response_cmd_type()):
raise ValidationError(
f"Could not determine appropriate protocol: "
f"The response command type "
f"{exchange_cls.get_response_cmd_type()} is not supported by the "
f"protocol that matched the request command type: "
f"{protocol_type}"
)
manager: ExchangeManager[Any, Any, Any]
manager = ExchangeManager(
connection=self._connection,
requesting_on=protocol_type,
listening_for=exchange_cls.get_response_cmd_type(),
)
exchange = exchange_cls(manager)
setattr(self, attr, exchange)
def add_protocol_handler(
self,
protocol_class: Type[ProtocolAPI],
handler_fn: HandlerFn,
) -> SubscriptionAPI:
if not self._multiplexer.has_protocol(protocol_class):
raise UnknownProtocol(
f"Protocol {protocol_class} was not found int he connected "
f"protocols: {self._multiplexer.get_protocols()}")
self._protocol_handlers[protocol_class].add(handler_fn)
cancel_fn = functools.partial(
self._protocol_handlers[protocol_class].remove,
handler_fn,
)
return Subscription(cancel_fn)
def __init__(self, connection: ConnectionAPI,
request_protocol_type: Type[ProtocolAPI],
response_cmd_type: Type[CommandAPI]) -> None:
# This style of initialization keeps `mypy` happy.
BaseService.__init__(self, token=connection.cancel_token)
self._connection = connection
self.request_protocol_type = request_protocol_type
try:
self.request_protocol = self._connection.get_multiplexer(
).get_protocol_by_type(request_protocol_type, )
except UnknownProtocol as err:
raise UnknownProtocol(
f"Response candidate stream configured to use "
f"{request_protocol_type} which is not available in the "
f"Multiplexer") from err
self.response_cmd_type = response_cmd_type
self._lock = asyncio.Lock()
def get_protocol_for_command_type(
self, command_type: Type[CommandAPI[Any]]) -> ProtocolAPI:
supported_protocols = tuple(protocol
for protocol in self.get_protocols()
if protocol.supports_command(command_type))
if len(supported_protocols) == 1:
return supported_protocols[0]
elif not supported_protocols:
raise UnknownProtocol(
f"Connection does not have any protocols that support the "
f"request command: {command_type}")
elif len(supported_protocols) > 1:
raise ValidationError(
f"Could not determine appropriate protocol for command: "
f"{command_type}. Command was found in the "
f"protocols {supported_protocols}")
else:
raise Exception("This code path should be unreachable")
async def stream_protocol_messages(self,
protocol_identifier: Union[ProtocolAPI, Type[ProtocolAPI]],
) -> AsyncIterator[CommandAPI[Any]]:
"""
Stream the messages for the specified protocol.
"""
if isinstance(protocol_identifier, ProtocolAPI):
protocol_class = type(protocol_identifier)
elif isinstance(protocol_identifier, type) and issubclass(protocol_identifier, ProtocolAPI):
protocol_class = protocol_identifier
else:
raise TypeError(f"Unknown protocol identifier: {protocol_identifier}")
if not self.has_protocol(protocol_class):
raise UnknownProtocol(f"Unknown protocol '{protocol_class}'")
if self._protocol_locks[protocol_class].locked():
raise Exception(f"Streaming lock for {protocol_class} is not free.")
async with self._protocol_locks[protocol_class]:
self.raise_if_streaming_error()
msg_queue = self._protocol_queues[protocol_class]
while self.is_streaming:
if not msg_queue.empty():
# We use an optimistic strategy here of using
# `get_nowait()` to reduce the number of times we yield to
# the event loop.
yield msg_queue.get_nowait()
# Manually release the event loop to prevent blocking it for too long in case
# the queue has lots of items.
await asyncio.sleep(0)
else:
# Must run this with a timeout so that we return in case the multiplexer
# stops streaming, althouth there's no real need to have this wake up so
# frequently in case there are no messages, so use a long one.
try:
yield await asyncio.wait_for(
msg_queue.get(), timeout=self._stream_idle_timeout)
except asyncio.TimeoutError:
pass
async def stream_protocol_messages(
self,
protocol_identifier: Union[ProtocolAPI, Type[ProtocolAPI]],
) -> AsyncIterator[Tuple[CommandAPI, Payload]]:
"""
Stream the messages for the specified protocol.
"""
if isinstance(protocol_identifier, Protocol):
protocol_class = type(protocol_identifier)
elif isinstance(protocol_identifier, type) and issubclass(
protocol_identifier, Protocol):
protocol_class = protocol_identifier
else:
raise TypeError("Unknown protocol identifier: {protocol}")
if not self.has_protocol(protocol_class):
raise UnknownProtocol(f"Unknown protocol '{protocol_class}'")
if self._protocol_locks.is_locked(protocol_class):
raise Exception(
f"Streaming lock for {protocol_class} is not free.")
async with self._protocol_locks.lock(protocol_class):
msg_queue = self._protocol_queues[protocol_class]
if not hasattr(self, '_multiplex_token'):
raise Exception("Multiplexer is not multiplexed")
token = self._multiplex_token
while not self.is_closing and not token.triggered:
try:
# We use an optimistic strategy here of using
# `get_nowait()` to reduce the number of times we yield to
# the event loop. Since this is an async generator it will
# yield to the loop each time it returns a value so we
# don't have to worry about this blocking other processes.
yield msg_queue.get_nowait()
except asyncio.QueueEmpty:
yield await self.wait(msg_queue.get(), token=token)