Refine Mplex.close and SwarmConn.close

Ensure `close` cleans up things and cancel the service finally.

libp2p/network/connection/swarm_connection.py

@@ -29,10 +29,19 @@ class SwarmConn(INetConn, Service):
         self.streams = set()
         self.event_closed = trio.Event()
 
+    @property
+    def is_closed(self) -> bool:
+        return self.event_closed.is_set()
+
     async def close(self) -> None:
         if self.event_closed.is_set():
             return
         self.event_closed.set()
+        await self._cleanup()
+        # Cancel service
+        await self.manager.stop()
+
+    async def _cleanup(self) -> None:
         self.swarm.remove_conn(self)
 
         await self.muxed_conn.close()
@@ -51,28 +60,23 @@ class SwarmConn(INetConn, Service):
         while self.manager.is_running:
             try:
                 stream = await self.muxed_conn.accept_stream()
-            except MuxedConnUnavailable:
-                # If there is anything wrong in the MuxedConn,
-                # we should break the loop and close the connection.
-                break
             # Asynchronously handle the accepted stream, to avoid blocking the next stream.
+            except MuxedConnUnavailable:
+                break
             self.manager.run_task(self._handle_muxed_stream, stream)
 
         await self.close()
 
-    async def _call_stream_handler(self, net_stream: NetStream) -> None:
-        try:
-            await self.swarm.common_stream_handler(net_stream)
-        # TODO: More exact exceptions
-        except Exception:
-            # TODO: Emit logs.
-            # TODO: Clean up and remove the stream from SwarmConn if there is anything wrong.
-            self.remove_stream(net_stream)
-
     async def _handle_muxed_stream(self, muxed_stream: IMuxedStream) -> None:
         net_stream = await self._add_stream(muxed_stream)
         if self.swarm.common_stream_handler is not None:
-            await self._call_stream_handler(net_stream)
+            try:
+                await self.swarm.common_stream_handler(net_stream)
+            # TODO: More exact exceptions
+            except Exception:
+                # TODO: Emit logs.
+                # TODO: Clean up and remove the stream from SwarmConn if there is anything wrong.
+                self.remove_stream(net_stream)
 
     async def _add_stream(self, muxed_stream: IMuxedStream) -> NetStream:
         net_stream = NetStream(muxed_stream)
@@ -84,7 +88,8 @@ class SwarmConn(INetConn, Service):
         await self.swarm.notify_disconnected(self)
 
     async def run(self) -> None:
-        await self._handle_new_streams()
+        self.manager.run_task(self._handle_new_streams)
+        await self.manager.wait_finished()
 
     async def new_stream(self) -> NetStream:
         muxed_stream = await self.muxed_conn.open_stream()

libp2p/network/swarm.py

@@ -44,6 +44,7 @@ class Swarm(INetwork, Service):
     common_stream_handler: Optional[StreamHandlerFn]
 
     notifees: List[INotifee]
+    event_closed: trio.Event
 
     def __init__(
         self,
@@ -62,6 +63,8 @@ class Swarm(INetwork, Service):
         # Create Notifee array
         self.notifees = []
 
+        self.event_closed = trio.Event()
+
         self.common_stream_handler = None
 
     async def run(self) -> None:
@@ -227,10 +230,19 @@ class Swarm(INetwork, Service):
         return False
 
     async def close(self) -> None:
-        # TODO: Prevent from new listeners and conns being added.
+        if self.event_closed.is_set():
+            return
+        self.event_closed.set()
         #   Reference: https://github.com/libp2p/go-libp2p-swarm/blob/8be680aef8dea0a4497283f2f98470c2aeae6b65/swarm.go#L124-L134  # noqa: E501
+        async with trio.open_nursery() as nursery:
+            for conn in self.connections.values():
+                nursery.start_soon(conn.close)
+        async with trio.open_nursery() as nursery:
+            for listener in self.listeners.values():
+                nursery.start_soon(listener.close)
+
+        # Cancel tasks
         await self.manager.stop()
-        await self.manager.wait_finished()
         logger.debug("swarm successfully closed")
 
     async def close_peer(self, peer_id: ID) -> None:
@@ -270,8 +282,6 @@ class Swarm(INetwork, Service):
 
     # Notifee
 
-    # TODO: Remeber the spawn notifying tasks and clean them up when closing.
-
     def register_notifee(self, notifee: INotifee) -> None:
         """
         :param notifee: object implementing Notifee interface