Rewrite factories, made some of the test running

libp2p/network/connection/swarm_connection.py

@@ -1,6 +1,8 @@
-import asyncio
 from typing import TYPE_CHECKING, Any, Awaitable, List, Set, Tuple
 
+import trio
+from async_service import Service
+
 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
@@ -15,21 +17,17 @@ Reference: https://github.com/libp2p/go-libp2p-swarm/blob/04c86bbdafd390651cb2ee
 """
 
 
-class SwarmConn(INetConn):
+class SwarmConn(INetConn, Service):
     muxed_conn: IMuxedConn
     swarm: "Swarm"
     streams: Set[NetStream]
-    event_closed: asyncio.Event
-
-    _tasks: List["asyncio.Future[Any]"]
+    event_closed: trio.Event
 
     def __init__(self, muxed_conn: IMuxedConn, swarm: "Swarm") -> None:
         self.muxed_conn = muxed_conn
         self.swarm = swarm
         self.streams = set()
-        self.event_closed = asyncio.Event()
-
-        self._tasks = []
+        self.event_closed = trio.Event()
 
     async def close(self) -> None:
         if self.event_closed.is_set():
@@ -45,16 +43,11 @@ class SwarmConn(INetConn):
             await stream.reset()
         # Force context switch for stream handlers to process the stream reset event we just emit
         # before we cancel the stream handler tasks.
-        await asyncio.sleep(0.1)
+        await trio.sleep(0.1)
 
-        for task in self._tasks:
-            task.cancel()
-            try:
-                await task
-            except asyncio.CancelledError:
-                pass
+        # FIXME: Now let `_notify_disconnected` finish first.
         # Schedule `self._notify_disconnected` to make it execute after `close` is finished.
-        self._notify_disconnected()
+        await self._notify_disconnected()
 
     async def _handle_new_streams(self) -> None:
         while True:
@@ -65,7 +58,7 @@ class SwarmConn(INetConn):
                 # we should break the loop and close the connection.
                 break
             # Asynchronously handle the accepted stream, to avoid blocking the next stream.
-            await self.run_task(self._handle_muxed_stream(stream))
+            self.manager.run_task(self._handle_muxed_stream, stream)
 
         await self.close()
 
@@ -79,28 +72,26 @@ class SwarmConn(INetConn):
             self.remove_stream(net_stream)
 
     async def _handle_muxed_stream(self, muxed_stream: IMuxedStream) -> None:
-        net_stream = self._add_stream(muxed_stream)
+        net_stream = await self._add_stream(muxed_stream)
         if self.swarm.common_stream_handler is not None:
-            await self.run_task(self._call_stream_handler(net_stream))
+            await self._call_stream_handler(net_stream)
 
-    def _add_stream(self, muxed_stream: IMuxedStream) -> NetStream:
+    async def _add_stream(self, muxed_stream: IMuxedStream) -> NetStream:
         net_stream = NetStream(muxed_stream)
         self.streams.add(net_stream)
-        self.swarm.notify_opened_stream(net_stream)
+        await self.swarm.notify_opened_stream(net_stream)
         return net_stream
 
-    def _notify_disconnected(self) -> None:
-        self.swarm.notify_disconnected(self)
+    async def _notify_disconnected(self) -> None:
+        await self.swarm.notify_disconnected(self)
 
-    async def start(self) -> None:
-        await self.run_task(self._handle_new_streams())
-
-    async def run_task(self, coro: Awaitable[Any]) -> None:
-        self._tasks.append(asyncio.ensure_future(coro))
+    async def run(self) -> None:
+        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()
-        return self._add_stream(muxed_stream)
+        return await self._add_stream(muxed_stream)
 
     async def get_streams(self) -> Tuple[NetStream, ...]:
         return tuple(self.streams)