Add call to wait_closed method of asyncio.StreamWriter

libp2p/network/connection/raw_connection.py

@@ -37,5 +37,6 @@ class RawConnection(IRawConnection):
         """
         return await self.reader.read(n)
 
-    def close(self) -> None:
+    async def close(self) -> None:
         self.writer.close()
+        await self.writer.wait_closed()

libp2p/network/connection/raw_connection_interface.py

@@ -17,5 +17,5 @@ class IRawConnection(ABC):
         pass
 
     @abstractmethod
-    def close(self) -> None:
+    async def close(self) -> None:
         pass

libp2p/network/swarm.py

@@ -118,7 +118,7 @@ class Swarm(INetwork):
                 )
             except SecurityUpgradeFailure as error:
                 # TODO: Add logging to indicate the failure
-                raw_conn.close()
+                await raw_conn.close()
                 raise SwarmException(
                     f"fail to upgrade the connection to a secured connection from {peer_id}"
                 ) from error
@@ -128,7 +128,7 @@ class Swarm(INetwork):
                 )
             except MuxerUpgradeFailure as error:
                 # TODO: Add logging to indicate the failure
-                secured_conn.close()
+                await secured_conn.close()
                 raise SwarmException(
                     f"fail to upgrade the connection to a muxed connection from {peer_id}"
                 ) from error
@@ -215,7 +215,7 @@ class Swarm(INetwork):
                     )
                 except SecurityUpgradeFailure as error:
                     # TODO: Add logging to indicate the failure
-                    raw_conn.close()
+                    await raw_conn.close()
                     raise SwarmException(
                         "fail to upgrade the connection to a secured connection"
                     ) from error
@@ -226,7 +226,7 @@ class Swarm(INetwork):
                     )
                 except MuxerUpgradeFailure as error:
                     # TODO: Add logging to indicate the failure
-                    secured_conn.close()
+                    await secured_conn.close()
                     raise SwarmException(
                         f"fail to upgrade the connection to a muxed connection from {peer_id}"
                     ) from error

libp2p/security/base_session.py

@@ -36,8 +36,8 @@ class BaseSession(ISecureConn):
     async def read(self, n: int = -1) -> bytes:
         return await self.conn.read(n)
 
-    def close(self) -> None:
-        self.conn.close()
+    async def close(self) -> None:
+        await self.conn.close()
 
     def get_local_peer(self) -> ID:
         return self.local_peer

libp2p/stream_muxer/abc.py

@@ -40,7 +40,7 @@ class IMuxedConn(ABC):
         pass
 
     @abstractmethod
-    def close(self) -> None:
+    async def close(self) -> None:
         """
         close connection
         """

libp2p/stream_muxer/mplex/mplex.py

@@ -69,11 +69,11 @@ class Mplex(IMuxedConn):
     def initiator(self) -> bool:
         return self.conn.initiator
 
-    def close(self) -> None:
+    async def close(self) -> None:
         """
         close the stream muxer and underlying raw connection
         """
-        self.conn.close()
+        await self.conn.close()
 
     def is_closed(self) -> bool:
         """