Merge branch 'main' into add-ws-transport

libp2p/network/connection/swarm_connection.py

@@ -23,7 +23,8 @@ if TYPE_CHECKING:
 
 
 """
-Reference: https://github.com/libp2p/go-libp2p-swarm/blob/04c86bbdafd390651cb2ee14e334f7caeedad722/swarm_conn.go
+Reference: https://github.com/libp2p/go-libp2p-swarm/blob/
+04c86bbdafd390651cb2ee14e334f7caeedad722/swarm_conn.go
 """
 
 
@@ -43,6 +44,21 @@ class SwarmConn(INetConn):
         self.streams = set()
         self.event_closed = trio.Event()
         self.event_started = trio.Event()
+        # Provide back-references/hooks expected by NetStream
+        try:
+            setattr(self.muxed_conn, "swarm", self.swarm)
+
+            # NetStream expects an awaitable remove_stream hook
+            async def _remove_stream_hook(stream: NetStream) -> None:
+                self.remove_stream(stream)
+
+            setattr(self.muxed_conn, "remove_stream", _remove_stream_hook)
+        except Exception as e:
+            logging.warning(
+                f"Failed to set optional conveniences on muxed_conn "
+                f"for peer {muxed_conn.peer_id}: {e}"
+            )
+            # optional conveniences
         if hasattr(muxed_conn, "on_close"):
             logging.debug(f"Setting on_close for peer {muxed_conn.peer_id}")
             setattr(muxed_conn, "on_close", self._on_muxed_conn_closed)

libp2p/network/swarm.py

@@ -1,3 +1,7 @@
+from collections.abc import (
+    Awaitable,
+    Callable,
+)
 import logging
 
 from multiaddr import (
@@ -333,8 +337,16 @@ class Swarm(Service, INetworkService):
 
             # Close all listeners
             if hasattr(self, "listeners"):
-                for listener in self.listeners.values():
+                for maddr_str, listener in self.listeners.items():
                     await listener.close()
+                    # Notify about listener closure
+                    try:
+                        multiaddr = Multiaddr(maddr_str)
+                        await self.notify_listen_close(multiaddr)
+                    except Exception as e:
+                        logger.warning(
+                            f"Failed to notify listen_close for {maddr_str}: {e}"
+                        )
                 self.listeners.clear()
 
             # Close the transport if it exists and has a close method
@@ -418,7 +430,17 @@ class Swarm(Service, INetworkService):
                 nursery.start_soon(notifee.listen, self, multiaddr)
 
     async def notify_closed_stream(self, stream: INetStream) -> None:
-        raise NotImplementedError
+        async with trio.open_nursery() as nursery:
+            for notifee in self.notifees:
+                nursery.start_soon(notifee.closed_stream, self, stream)
 
     async def notify_listen_close(self, multiaddr: Multiaddr) -> None:
-        raise NotImplementedError
+        async with trio.open_nursery() as nursery:
+            for notifee in self.notifees:
+                nursery.start_soon(notifee.listen_close, self, multiaddr)
+
+    # Generic notifier used by NetStream._notify_closed
+    async def notify_all(self, notifier: Callable[[INotifee], Awaitable[None]]) -> None:
+        async with trio.open_nursery() as nursery:
+            for notifee in self.notifees:
+                nursery.start_soon(notifier, notifee)