Add detection for disconnections in mplex

libp2p/network/connection/swarm_connection.py

@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Awaitable, List, Set, Tuple
 from libp2p.network.connection.net_connection_interface import INetConn
 from libp2p.network.stream.net_stream import NetStream
 from libp2p.stream_muxer.abc import IMuxedConn, IMuxedStream
+from libp2p.stream_muxer.exceptions import MuxedConnUnavailable
 
 if TYPE_CHECKING:
     from libp2p.network.swarm import Swarm  # noqa: F401
@@ -34,17 +35,27 @@ class SwarmConn(INetConn):
         if self.event_closed.is_set():
             return
         self.event_closed.set()
+        self.swarm.remove_conn(self)
+
         await self.conn.close()
+
+        # This is just for cleaning up state. The connection has already been closed.
+        # We *could* optimize this but it really isn't worth it.
+        for stream in self.streams:
+            await stream.reset()
+        # Schedule `self._notify_disconnected` to make it execute after `close` is finished.
+        asyncio.ensure_future(self._notify_disconnected())
+
         for task in self._tasks:
             task.cancel()
 
-        # TODO: Reset streams for local.
-        # TODO: Notify closed.
-
     async def _handle_new_streams(self) -> None:
         # TODO: Break the loop when anything wrong in the connection.
         while True:
-            stream = await self.conn.accept_stream()
+            try:
+                stream = await self.conn.accept_stream()
+            except MuxedConnUnavailable:
+                break
             # Asynchronously handle the accepted stream, to avoid blocking the next stream.
             await self.run_task(self._handle_muxed_stream(stream))
 
@@ -57,11 +68,16 @@ class SwarmConn(INetConn):
 
     async def _add_stream(self, muxed_stream: IMuxedStream) -> NetStream:
         net_stream = NetStream(muxed_stream)
+        self.streams.add(net_stream)
         # Call notifiers since event occurred
         for notifee in self.swarm.notifees:
             await notifee.opened_stream(self.swarm, net_stream)
         return net_stream
 
+    async def _notify_disconnected(self) -> None:
+        for notifee in self.swarm.notifees:
+            await notifee.disconnected(self.swarm, self.conn)
+
     async def start(self) -> None:
         await self.run_task(self._handle_new_streams())
 

libp2p/network/swarm.py

@@ -262,7 +262,6 @@ class Swarm(INetwork):
         if peer_id not in self.connections:
             return
         connection = self.connections[peer_id]
-        del self.connections[peer_id]
         await connection.close()
 
         logger.debug("successfully close the connection to peer %s", peer_id)
@@ -277,3 +276,10 @@ class Swarm(INetwork):
             await notifee.connected(self, muxed_conn)
         await swarm_conn.start()
         return swarm_conn
+
+    def remove_conn(self, swarm_conn: SwarmConn) -> None:
+        print(f"!@# remove_conn: {swarm_conn}")
+        peer_id = swarm_conn.conn.peer_id
+        # TODO: Should be changed to remove the exact connection,
+        #   if we have several connections per peer in the future.
+        del self.connections[peer_id]