Updated the timeout wrapper for read/write operations

libp2p/protocol_muxer/multiselect.py

@@ -59,19 +59,20 @@ class Multiselect(IMultiselectMuxer):
         :return: selected protocol name, handler function
         :raise MultiselectError: raised when negotiation failed
         """
+        try:
+            with trio.fail_after(DEFAULT_NEGOTIATE_TIMEOUT):
                 await self.handshake(communicator)
 
                 while True:
                     try:
-                with trio.fail_after(DEFAULT_NEGOTIATE_TIMEOUT):
                         command = await communicator.read()
-            except trio.TooSlowError:
-                raise MultiselectError("handshake read timeout")
                     except MultiselectCommunicatorError as error:
                         raise MultiselectError() from error
 
                     if command == "ls":
-                supported_protocols = [p for p in self.handlers.keys() if p is not None]
+                        supported_protocols = [
+                            p for p in self.handlers.keys() if p is not None
+                        ]
                         response = "\n".join(supported_protocols) + "\n"
 
                         try:
@@ -94,6 +95,8 @@ class Multiselect(IMultiselectMuxer):
                             raise MultiselectError() from error
 
                 raise MultiselectError("Negotiation failed: no matching protocol")
+        except trio.TooSlowError:
+            raise MultiselectError("handshake read timeout")
 
     async def handshake(self, communicator: IMultiselectCommunicator) -> None:
         """
@@ -108,10 +111,7 @@ class Multiselect(IMultiselectMuxer):
             raise MultiselectError() from error
 
         try:
-            with trio.fail_after(DEFAULT_NEGOTIATE_TIMEOUT):  # Timeout after 5 seconds
             handshake_contents = await communicator.read()
-        except trio.TooSlowError:
-            raise MultiselectError("protocol selection response timed out")
         except MultiselectCommunicatorError as error:
             raise MultiselectError() from error
 

libp2p/protocol_muxer/multiselect_client.py

@@ -42,10 +42,8 @@ class MultiselectClient(IMultiselectClient):
             raise MultiselectClientError() from error
 
         try:
-            with trio.fail_after(DEFAULT_NEGOTIATE_TIMEOUT):
             handshake_contents = await communicator.read()
-        except trio.TooSlowError:
+
-            raise MultiselectClientError("handshake read timed out")
         except MultiselectCommunicatorError as error:
             raise MultiselectClientError() from error
 
@@ -65,16 +63,22 @@ class MultiselectClient(IMultiselectClient):
         :return: selected protocol
         :raise MultiselectClientError: raised when protocol negotiation failed
         """
+        try:
+            with trio.fail_after(DEFAULT_NEGOTIATE_TIMEOUT):
                 await self.handshake(communicator)
 
                 for protocol in protocols:
                     try:
-                selected_protocol = await self.try_select(communicator, protocol)
+                        selected_protocol = await self.try_select(
+                            communicator, protocol
+                        )
                         return selected_protocol
                     except MultiselectClientError:
                         pass
 
                 raise MultiselectClientError("protocols not supported")
+        except trio.TooSlowError:
+            raise MultiselectClientError("response timed out")
 
     async def query_multistream_command(
         self, communicator: IMultiselectCommunicator, command: str
@@ -88,6 +92,8 @@ class MultiselectClient(IMultiselectClient):
         :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):
                 await self.handshake(communicator)
 
                 if command == "ls":
@@ -99,15 +105,15 @@ class MultiselectClient(IMultiselectClient):
                     raise ValueError("Command not supported")
 
                 try:
-            with trio.fail_after(DEFAULT_NEGOTIATE_TIMEOUT):  # Timeout after 5 seconds
                     response = await communicator.read()
                     response_list = response.strip().splitlines()
-        except trio.TooSlowError:
+
-            raise MultiselectClientError("command response timed out")
                 except MultiselectCommunicatorError as error:
                     raise MultiselectClientError() from error
 
                 return response_list
+        except trio.TooSlowError:
+            raise MultiselectClientError("command response timed out")
 
     async def try_select(
         self, communicator: IMultiselectCommunicator, protocol: TProtocol
@@ -126,10 +132,8 @@ class MultiselectClient(IMultiselectClient):
             raise MultiselectClientError() from error
 
         try:
-            with trio.fail_after(DEFAULT_NEGOTIATE_TIMEOUT):  # Timeout after 5 seconds
             response = await communicator.read()
-        except trio.TooSlowError:
+
-            raise MultiselectClientError("protocol selection response timed out")
         except MultiselectCommunicatorError as error:
             raise MultiselectClientError() from error