From 5b32bfdd3f030eece7b0357c263cd17cb7d94b1e Mon Sep 17 00:00:00 2001
From: Alex Stokes <r.alex.stokes@gmail.com>
Date: Sat, 24 Aug 2019 22:06:24 +0200
Subject: [PATCH] Add call to `wait_closed` method of asyncio.StreamWriter

---
 libp2p/network/connection/raw_connection.py           | 3 ++-
 libp2p/network/connection/raw_connection_interface.py | 2 +-
 libp2p/network/swarm.py                               | 8 ++++----
 libp2p/security/base_session.py                       | 4 ++--
 libp2p/stream_muxer/abc.py                            | 2 +-
 libp2p/stream_muxer/mplex/mplex.py                    | 4 ++--
 6 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/libp2p/network/connection/raw_connection.py b/libp2p/network/connection/raw_connection.py
index a2cff49..fe09c6f 100644
--- a/libp2p/network/connection/raw_connection.py
+++ b/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()
diff --git a/libp2p/network/connection/raw_connection_interface.py b/libp2p/network/connection/raw_connection_interface.py
index 4c3f82e..25b90ae 100644
--- a/libp2p/network/connection/raw_connection_interface.py
+++ b/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
diff --git a/libp2p/network/swarm.py b/libp2p/network/swarm.py
index 57fc955..07b5e14 100644
--- a/libp2p/network/swarm.py
+++ b/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
diff --git a/libp2p/security/base_session.py b/libp2p/security/base_session.py
index 19fe139..54562b1 100644
--- a/libp2p/security/base_session.py
+++ b/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
diff --git a/libp2p/stream_muxer/abc.py b/libp2p/stream_muxer/abc.py
index 7270c15..245a739 100644
--- a/libp2p/stream_muxer/abc.py
+++ b/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
         """
diff --git a/libp2p/stream_muxer/mplex/mplex.py b/libp2p/stream_muxer/mplex/mplex.py
index 3f3e938..7aefe3f 100644
--- a/libp2p/stream_muxer/mplex/mplex.py
+++ b/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:
         """