From 8d2415a404904833614fcf0db47981376b3f3fe3 Mon Sep 17 00:00:00 2001
From: mhchia <kevin.mh.chia@gmail.com>
Date: Mon, 23 Sep 2019 15:01:58 +0800
Subject: [PATCH] Move calls to `Notifee` inside `Swarm`

---
 libp2p/network/connection/swarm_connection.py | 17 ++++----
 libp2p/network/stream/net_stream.py           |  2 +-
 libp2p/network/swarm.py                       | 42 ++++++++++++++-----
 3 files changed, 39 insertions(+), 22 deletions(-)

diff --git a/libp2p/network/connection/swarm_connection.py b/libp2p/network/connection/swarm_connection.py
index 6714bb8..cf1dc9e 100644
--- a/libp2p/network/connection/swarm_connection.py
+++ b/libp2p/network/connection/swarm_connection.py
@@ -51,7 +51,7 @@ class SwarmConn(INetConn):
             except asyncio.CancelledError:
                 pass
         # Schedule `self._notify_disconnected` to make it execute after `close` is finished.
-        asyncio.ensure_future(self._notify_disconnected())
+        self._notify_disconnected()
 
     async def _handle_new_streams(self) -> None:
         while True:
@@ -76,20 +76,18 @@ class SwarmConn(INetConn):
             self.remove_stream(net_stream)
 
     async def _handle_muxed_stream(self, muxed_stream: IMuxedStream) -> None:
-        net_stream = await self._add_stream(muxed_stream)
+        net_stream = self._add_stream(muxed_stream)
         if self.swarm.common_stream_handler is not None:
             await self.run_task(self._call_stream_handler(net_stream))
 
-    async def _add_stream(self, muxed_stream: IMuxedStream) -> NetStream:
+    def _add_stream(self, muxed_stream: IMuxedStream) -> NetStream:
         net_stream = NetStream(muxed_stream)
         self.streams.add(net_stream)
-        for notifee in self.swarm.notifees:
-            await notifee.opened_stream(self.swarm, net_stream)
+        self.swarm.notify_opened_stream(net_stream)
         return net_stream
 
-    async def _notify_disconnected(self) -> None:
-        for notifee in self.swarm.notifees:
-            await notifee.disconnected(self.swarm, self)
+    def _notify_disconnected(self) -> None:
+        self.swarm.notify_disconnected(self)
 
     async def start(self) -> None:
         await self.run_task(self._handle_new_streams())
@@ -99,12 +97,11 @@ class SwarmConn(INetConn):
 
     async def new_stream(self) -> NetStream:
         muxed_stream = await self.conn.open_stream()
-        return await self._add_stream(muxed_stream)
+        return self._add_stream(muxed_stream)
 
     async def get_streams(self) -> Tuple[NetStream, ...]:
         return tuple(self.streams)
 
-    # TODO: Called by `Stream` whenever it is time to remove the stream.
     def remove_stream(self, stream: NetStream) -> None:
         if stream not in self.streams:
             return
diff --git a/libp2p/network/stream/net_stream.py b/libp2p/network/stream/net_stream.py
index 018ef6d..625a288 100644
--- a/libp2p/network/stream/net_stream.py
+++ b/libp2p/network/stream/net_stream.py
@@ -68,5 +68,5 @@ class NetStream(INetStream):
         await self.muxed_stream.reset()
 
     # TODO: `remove`: Called by close and write when the stream is in specific states.
-    #   It notify `ClosedStream` after `SwarmConn.remove_stream` is called.
+    #   It notifies `ClosedStream` after `SwarmConn.remove_stream` is called.
     # Reference: https://github.com/libp2p/go-libp2p-swarm/blob/99831444e78c8f23c9335c17d8f7c700ba25ca14/swarm_stream.go  # noqa: E501
diff --git a/libp2p/network/swarm.py b/libp2p/network/swarm.py
index 4fd3bbc..bed204e 100644
--- a/libp2p/network/swarm.py
+++ b/libp2p/network/swarm.py
@@ -223,8 +223,7 @@ class Swarm(INetwork):
                 await listener.listen(maddr)
 
                 # Call notifiers since event occurred
-                for notifee in self.notifees:
-                    await notifee.listen(self, maddr)
+                self.notify_listen(maddr)
 
                 return True
             except IOError:
@@ -234,13 +233,6 @@ class Swarm(INetwork):
         # No maddr succeeded
         return False
 
-    def register_notifee(self, notifee: INotifee) -> None:
-        """
-        :param notifee: object implementing Notifee interface
-        :return: true if notifee registered successfully, false otherwise
-        """
-        self.notifees.append(notifee)
-
     def add_router(self, router: IPeerRouting) -> None:
         self.router = router
 
@@ -279,8 +271,7 @@ class Swarm(INetwork):
         # Store muxed_conn with peer id
         self.connections[muxed_conn.peer_id] = swarm_conn
         # Call notifiers since event occurred
-        for notifee in self.notifees:
-            await notifee.connected(self, swarm_conn)
+        self.notify_connected(swarm_conn)
         await swarm_conn.start()
         return swarm_conn
 
@@ -294,3 +285,32 @@ class Swarm(INetwork):
         # 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
+
+    # 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
+        :return: true if notifee registered successfully, false otherwise
+        """
+        self.notifees.append(notifee)
+
+    def notify_opened_stream(self, stream: INetStream) -> None:
+        asyncio.gather(
+            *[notifee.opened_stream(self, stream) for notifee in self.notifees]
+        )
+
+    # TODO: `notify_closed_stream`
+
+    def notify_connected(self, conn: INetConn) -> None:
+        asyncio.gather(*[notifee.connected(self, conn) for notifee in self.notifees])
+
+    def notify_disconnected(self, conn: INetConn) -> None:
+        asyncio.gather(*[notifee.disconnected(self, conn) for notifee in self.notifees])
+
+    def notify_listen(self, multiaddr: Multiaddr) -> None:
+        asyncio.gather(*[notifee.listen(self, multiaddr) for notifee in self.notifees])
+
+    # TODO: `notify_listen_close`