Make Mplex and SwarmConn not Service

After second thoughts, they seem not a good candidate of `Service`.
The shutdown logic becomes simpler by making them not `Service`.

libp2p/network/connection/net_connection_interface.py

@@ -1,6 +1,8 @@
 from abc import abstractmethod
 from typing import Tuple
 
+import trio
+
 from libp2p.io.abc import Closer
 from libp2p.network.stream.net_stream_interface import INetStream
 from libp2p.stream_muxer.abc import IMuxedConn
@@ -8,11 +10,12 @@ from libp2p.stream_muxer.abc import IMuxedConn
 
 class INetConn(Closer):
     muxed_conn: IMuxedConn
+    event_started: trio.Event
 
     @abstractmethod
     async def new_stream(self) -> INetStream:
         ...
 
     @abstractmethod
-    async def get_streams(self) -> Tuple[INetStream, ...]:
+    def get_streams(self) -> Tuple[INetStream, ...]:
         ...

libp2p/network/connection/swarm_connection.py

@@ -1,6 +1,5 @@
 from typing import TYPE_CHECKING, Set, Tuple
 
-from async_service import Service
 import trio
 
 from libp2p.network.connection.net_connection_interface import INetConn
@@ -17,10 +16,11 @@ Reference: https://github.com/libp2p/go-libp2p-swarm/blob/04c86bbdafd390651cb2ee
 """
 
 
-class SwarmConn(INetConn, Service):
+class SwarmConn(INetConn):
     muxed_conn: IMuxedConn
     swarm: "Swarm"
     streams: Set[NetStream]
+    event_started: trio.Event
     event_closed: trio.Event
 
     def __init__(self, muxed_conn: IMuxedConn, swarm: "Swarm") -> None:
@@ -28,6 +28,7 @@ class SwarmConn(INetConn, Service):
         self.swarm = swarm
         self.streams = set()
         self.event_closed = trio.Event()
+        self.event_started = trio.Event()
 
     @property
     def is_closed(self) -> bool:
@@ -38,8 +39,6 @@ class SwarmConn(INetConn, Service):
             return
         self.event_closed.set()
         await self._cleanup()
-        # Cancel service
-        await self.manager.stop()
 
     async def _cleanup(self) -> None:
         self.swarm.remove_conn(self)
@@ -57,13 +56,14 @@ class SwarmConn(INetConn, Service):
         self._notify_disconnected()
 
     async def _handle_new_streams(self) -> None:
-        while self.manager.is_running:
+        self.event_started.set()
+        while True:
             try:
                 stream = await self.muxed_conn.accept_stream()
             # Asynchronously handle the accepted stream, to avoid blocking the next stream.
             except MuxedConnUnavailable:
                 break
-            self.manager.run_task(self._handle_muxed_stream, stream)
+            self.swarm.manager.run_task(self._handle_muxed_stream, stream)
 
         await self.close()
 
@@ -87,15 +87,14 @@ class SwarmConn(INetConn, Service):
     def _notify_disconnected(self) -> None:
         self.swarm.notify_disconnected(self)
 
-    async def run(self) -> None:
-        self.manager.run_task(self._handle_new_streams)
-        await self.manager.wait_finished()
+    async def start(self) -> None:
+        await self._handle_new_streams()
 
     async def new_stream(self) -> NetStream:
         muxed_stream = await self.muxed_conn.open_stream()
         return self._add_stream(muxed_stream)
 
-    async def get_streams(self) -> Tuple[NetStream, ...]:
+    def get_streams(self) -> Tuple[NetStream, ...]:
         return tuple(self.streams)
 
     def remove_stream(self, stream: NetStream) -> None:

libp2p/network/swarm.py

@@ -2,7 +2,6 @@ import logging
 from typing import Dict, List, Optional
 
 from multiaddr import Multiaddr
-import trio
 
 from libp2p.io.abc import ReadWriteCloser
 from libp2p.network.connection.net_connection_interface import INetConn
@@ -44,7 +43,6 @@ class Swarm(INetworkService):
     common_stream_handler: Optional[StreamHandlerFn]
 
     notifees: List[INotifee]
-    event_closed: trio.Event
 
     def __init__(
         self,
@@ -63,8 +61,6 @@ class Swarm(INetworkService):
         # Create Notifee array
         self.notifees = []
 
-        self.event_closed = trio.Event()
-
         self.common_stream_handler = None
 
     async def run(self) -> None:
@@ -158,13 +154,11 @@ class Swarm(INetworkService):
 
         try:
             muxed_conn = await self.upgrader.upgrade_connection(secured_conn, peer_id)
-            self.manager.run_child_service(muxed_conn)
         except MuxerUpgradeFailure as error:
             error_msg = "fail to upgrade mux for peer %s"
             logger.debug(error_msg, peer_id)
             await secured_conn.close()
             raise SwarmException(error_msg % peer_id) from error
-
         logger.debug("upgraded mux for peer %s", peer_id)
 
         swarm_conn = await self.add_conn(muxed_conn)
@@ -226,7 +220,6 @@ class Swarm(INetworkService):
                     muxed_conn = await self.upgrader.upgrade_connection(
                         secured_conn, peer_id
                     )
-                    self.manager.run_child_service(muxed_conn)
                 except MuxerUpgradeFailure as error:
                     error_msg = "fail to upgrade mux for peer %s"
                     logger.debug(error_msg, peer_id)
@@ -235,8 +228,8 @@ class Swarm(INetworkService):
                 logger.debug("upgraded mux for peer %s", peer_id)
 
                 await self.add_conn(muxed_conn)
-
                 logger.debug("successfully opened connection to peer %s", peer_id)
+
                 # NOTE: This is a intentional barrier to prevent from the handler exiting and
                 #   closing the connection.
                 await self.manager.wait_finished()
@@ -261,26 +254,12 @@ class Swarm(INetworkService):
         return False
 
     async def close(self) -> None:
-        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()
         logger.debug("swarm successfully closed")
 
     async def close_peer(self, peer_id: ID) -> None:
         if peer_id not in self.connections:
             return
-        # TODO: Should be changed to close multisple connections,
-        #   if we have several connections per peer in the future.
         connection = self.connections[peer_id]
         # NOTE: `connection.close` will delete `peer_id` from `self.connections`
         # and `notify_disconnected` for us.
@@ -293,12 +272,14 @@ class Swarm(INetworkService):
         and start to monitor the connection for its new streams and
         disconnection."""
         swarm_conn = SwarmConn(muxed_conn, self)
-        manager = self.manager.run_child_service(swarm_conn)
+        self.manager.run_task(muxed_conn.start)
+        await muxed_conn.event_started.wait()
+        self.manager.run_task(swarm_conn.start)
+        await swarm_conn.event_started.wait()
         # Store muxed_conn with peer id
         self.connections[muxed_conn.peer_id] = swarm_conn
         # Call notifiers since event occurred
         self.notify_connected(swarm_conn)
-        await manager.wait_started()
         return swarm_conn
 
     def remove_conn(self, swarm_conn: SwarmConn) -> None:
@@ -307,8 +288,6 @@ class Swarm(INetworkService):
         peer_id = swarm_conn.muxed_conn.peer_id
         if peer_id not in self.connections:
             return
-        # TODO: Should be changed to remove the exact connection,
-        #   if we have several connections per peer in the future.
         del self.connections[peer_id]
 
     # Notifee