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)