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: """