Fix all modules except for security

libp2p/io/trio.py

@@ -1,7 +1,6 @@
 import logging
 
 import trio
-from trio import SocketStream
 
 from libp2p.io.abc import ReadWriteCloser
 from libp2p.io.exceptions import IOException
@@ -9,29 +8,48 @@ from libp2p.io.exceptions import IOException
 logger = logging.getLogger("libp2p.io.trio")
 
 
-class TrioReadWriteCloser(ReadWriteCloser):
-    stream: SocketStream
+class TrioTCPStream(ReadWriteCloser):
+    stream: trio.SocketStream
+    # NOTE: Add both read and write lock to avoid `trio.BusyResourceError`
+    read_lock: trio.Lock
+    write_lock: trio.Lock
 
-    def __init__(self, stream: SocketStream) -> None:
+    def __init__(self, stream: trio.SocketStream) -> None:
         self.stream = stream
+        self.read_lock = trio.Lock()
+        self.write_lock = trio.Lock()
 
     async def write(self, data: bytes) -> None:
         """Raise `RawConnError` if the underlying connection breaks."""
-        try:
-            await self.stream.send_all(data)
-        except (trio.ClosedResourceError, trio.BrokenResourceError) as error:
-            raise IOException(error)
+        async with self.write_lock:
+            try:
+                await self.stream.send_all(data)
+            except (trio.ClosedResourceError, trio.BrokenResourceError) as error:
+                raise IOException from error
+            except trio.BusyResourceError as error:
+                # This should never happen, since we already access streams with read/write locks.
+                raise Exception(
+                    "this should never happen "
+                    "since we already access streams with read/write locks."
+                ) from error
 
     async def read(self, n: int = -1) -> bytes:
-        if n == 0:
-            # Check point
-            await trio.sleep(0)
-            return b""
-        max_bytes = n if n != -1 else None
-        try:
-            return await self.stream.receive_some(max_bytes)
-        except (trio.ClosedResourceError, trio.BrokenResourceError) as error:
-            raise IOException(error)
+        async with self.read_lock:
+            if n == 0:
+                # Checkpoint
+                await trio.hazmat.checkpoint()
+                return b""
+            max_bytes = n if n != -1 else None
+            try:
+                return await self.stream.receive_some(max_bytes)
+            except (trio.ClosedResourceError, trio.BrokenResourceError) as error:
+                raise IOException from error
+            except trio.BusyResourceError as error:
+                # This should never happen, since we already access streams with read/write locks.
+                raise Exception(
+                    "this should never happen "
+                    "since we already access streams with read/write locks."
+                ) from error
 
     async def close(self) -> None:
         await self.stream.aclose()