rewrite tcp reader/writer interface

libp2p/network/connection/raw_connection.py

@@ -1,42 +1,25 @@
-import asyncio
+import trio
 
+from libp2p.io.exceptions import IOException
 from .exceptions import RawConnError
 from .raw_connection_interface import IRawConnection
+from libp2p.io.abc import ReadWriteCloser
 
 
 class RawConnection(IRawConnection):
-    reader: asyncio.StreamReader
-    writer: asyncio.StreamWriter
+    read_write_closer: ReadWriteCloser
     is_initiator: bool
 
-    _drain_lock: asyncio.Lock
-
-    def __init__(
-        self,
-        reader: asyncio.StreamReader,
-        writer: asyncio.StreamWriter,
-        initiator: bool,
-    ) -> None:
-        self.reader = reader
-        self.writer = writer
+    def __init__(self, read_write_closer: ReadWriteCloser, initiator: bool) -> None:
+        self.read_write_closer = read_write_closer
         self.is_initiator = initiator
 
-        self._drain_lock = asyncio.Lock()
-
     async def write(self, data: bytes) -> None:
         """Raise `RawConnError` if the underlying connection breaks."""
         try:
-            self.writer.write(data)
-        except ConnectionResetError as error:
+            await self.read_write_closer.write(data)
+        except IOException as error:
             raise RawConnError(error)
-        # Reference: https://github.com/ethereum/lahja/blob/93610b2eb46969ff1797e0748c7ac2595e130aef/lahja/asyncio/endpoint.py#L99-L102  # noqa: E501
-        # Use a lock to serialize drain() calls. Circumvents this bug:
-        # https://bugs.python.org/issue29930
-        async with self._drain_lock:
-            try:
-                await self.writer.drain()
-            except ConnectionResetError as error:
-                raise RawConnError(error)
 
     async def read(self, n: int = -1) -> bytes:
         """
@@ -46,10 +29,9 @@ class RawConnection(IRawConnection):
         Raise `RawConnError` if the underlying connection breaks
         """
         try:
-            return await self.reader.read(n)
-        except ConnectionResetError as error:
+            return await self.read_write_closer.read(n)
+        except IOException as error:
             raise RawConnError(error)
 
     async def close(self) -> None:
-        self.writer.close()
-        await self.writer.wait_closed()
+        await self.read_write_closer.close()