diff --git a/libp2p/protocol_muxer/multiselect.py b/libp2p/protocol_muxer/multiselect.py index 05b112dc..6a68d2e4 100644 --- a/libp2p/protocol_muxer/multiselect.py +++ b/libp2p/protocol_muxer/multiselect.py @@ -50,7 +50,9 @@ class Multiselect(IMultiselectMuxer): # FIXME: Make TProtocol Optional[TProtocol] to keep types consistent async def negotiate( - self, communicator: IMultiselectCommunicator + self, + communicator: IMultiselectCommunicator, + negotiate_timeout: int = DEFAULT_NEGOTIATE_TIMEOUT, ) -> tuple[TProtocol, StreamHandlerFn | None]: """ Negotiate performs protocol selection. @@ -60,7 +62,7 @@ class Multiselect(IMultiselectMuxer): :raise MultiselectError: raised when negotiation failed """ try: - with trio.fail_after(DEFAULT_NEGOTIATE_TIMEOUT): + with trio.fail_after(negotiate_timeout): await self.handshake(communicator) while True: diff --git a/libp2p/protocol_muxer/multiselect_client.py b/libp2p/protocol_muxer/multiselect_client.py index 75d5ca05..aa6db5ba 100644 --- a/libp2p/protocol_muxer/multiselect_client.py +++ b/libp2p/protocol_muxer/multiselect_client.py @@ -51,7 +51,10 @@ class MultiselectClient(IMultiselectClient): raise MultiselectClientError("multiselect protocol ID mismatch") async def select_one_of( - self, protocols: Sequence[TProtocol], communicator: IMultiselectCommunicator + self, + protocols: Sequence[TProtocol], + communicator: IMultiselectCommunicator, + negotitate_timeout: int = DEFAULT_NEGOTIATE_TIMEOUT, ) -> TProtocol: """ For each protocol, send message to multiselect selecting protocol and @@ -64,7 +67,7 @@ class MultiselectClient(IMultiselectClient): :raise MultiselectClientError: raised when protocol negotiation failed """ try: - with trio.fail_after(DEFAULT_NEGOTIATE_TIMEOUT): + with trio.fail_after(negotitate_timeout): await self.handshake(communicator) for protocol in protocols: @@ -81,7 +84,10 @@ class MultiselectClient(IMultiselectClient): raise MultiselectClientError("response timed out") async def query_multistream_command( - self, communicator: IMultiselectCommunicator, command: str + self, + communicator: IMultiselectCommunicator, + command: str, + response_timeout: int = DEFAULT_NEGOTIATE_TIMEOUT, ) -> list[str]: """ Send a multistream-select command over the given communicator and return @@ -89,11 +95,13 @@ class MultiselectClient(IMultiselectClient): :param communicator: communicator to use to communicate with counterparty :param command: supported multistream-select command(e.g., ls) + :param negotiate_timeout: timeout for negotiation, + defaults to DEFAULT_NEGOTIATE_TIMEOUT :raise MultiselectClientError: If the communicator fails to process data. :return: list of strings representing the response from peer. """ try: - with trio.fail_after(DEFAULT_NEGOTIATE_TIMEOUT): + with trio.fail_after(response_timeout): await self.handshake(communicator) if command == "ls":