Refactor mplex and start to add close detection

libp2p/stream_muxer/abc.py

@@ -1,15 +1,8 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
 
 from libp2p.io.abc import ReadWriteCloser
 from libp2p.peer.id import ID
 from libp2p.security.secure_conn_interface import ISecureConn
-from libp2p.stream_muxer.mplex.constants import HeaderTags
-from libp2p.stream_muxer.mplex.datastructures import StreamID
-
-if TYPE_CHECKING:
-    # Prevent GenericProtocolHandlerFn introducing circular dependencies
-    from libp2p.network.typing import GenericProtocolHandlerFn  # noqa: F401
 
 
 class IMuxedConn(ABC):
@@ -20,16 +13,10 @@ class IMuxedConn(ABC):
     peer_id: ID
 
     @abstractmethod
-    def __init__(
-        self,
-        conn: ISecureConn,
-        generic_protocol_handler: "GenericProtocolHandlerFn",
-        peer_id: ID,
-    ) -> None:
+    def __init__(self, conn: ISecureConn, peer_id: ID) -> None:
         """
         create a new muxed connection
         :param conn: an instance of secured connection
-        :param generic_protocol_handler: generic protocol handler
         for new muxed streams
         :param peer_id: peer_id of peer the connection is to
         """
@@ -60,22 +47,11 @@ class IMuxedConn(ABC):
         """
 
     @abstractmethod
-    async def accept_stream(self, stream_id: StreamID, name: str) -> None:
+    async def accept_stream(self) -> "IMuxedStream":
         """
         accepts a muxed stream opened by the other end
         """
 
-    @abstractmethod
-    async def send_message(
-        self, flag: HeaderTags, data: bytes, stream_id: StreamID
-    ) -> int:
-        """
-        sends a message over the connection
-        :param header: header to use
-        :param data: data to send in the message
-        :param stream_id: stream the message is in
-        """
-
 
 class IMuxedStream(ReadWriteCloser):
 

libp2p/stream_muxer/mplex/mplex.py

@@ -2,7 +2,6 @@ import asyncio
 from typing import Any  # noqa: F401
 from typing import Dict, List, Optional, Tuple
 
-from libp2p.network.typing import GenericProtocolHandlerFn
 from libp2p.peer.id import ID
 from libp2p.security.secure_conn_interface import ISecureConn
 from libp2p.stream_muxer.abc import IMuxedConn, IMuxedStream
@@ -34,17 +33,13 @@ class Mplex(IMuxedConn):
     next_channel_id: int
     streams: Dict[StreamID, MplexStream]
     streams_lock: asyncio.Lock
+    new_stream_queue: "asyncio.Queue[IMuxedStream]"
     shutdown: asyncio.Event
 
     _tasks: List["asyncio.Future[Any]"]
 
     # TODO: `generic_protocol_handler` should be refactored out of mplex conn.
-    def __init__(
-        self,
-        secured_conn: ISecureConn,
-        generic_protocol_handler: GenericProtocolHandlerFn,
-        peer_id: ID,
-    ) -> None:
+    def __init__(self, secured_conn: ISecureConn, peer_id: ID) -> None:
         """
         create a new muxed connection
         :param secured_conn: an instance of ``ISecureConn``
@@ -56,15 +51,13 @@ class Mplex(IMuxedConn):
 
         self.next_channel_id = 0
 
-        # Store generic protocol handler
-        self.generic_protocol_handler = generic_protocol_handler
-
         # Set peer_id
         self.peer_id = peer_id
 
         # Mapping from stream ID -> buffer of messages for that stream
         self.streams = {}
         self.streams_lock = asyncio.Lock()
+        self.new_stream_queue = asyncio.Queue()
         self.shutdown = asyncio.Event()
 
         self._tasks = []
@@ -101,9 +94,10 @@ class Mplex(IMuxedConn):
         return next_id
 
     async def _initialize_stream(self, stream_id: StreamID, name: str) -> MplexStream:
+        stream = MplexStream(name, stream_id, self)
         async with self.streams_lock:
-            stream = MplexStream(name, stream_id, self)
-        self.streams[stream_id] = stream
+            self.streams[stream_id] = stream
+        print(f"!@# _initialize_stream: stream_id={stream_id}, name={name}")
         return stream
 
     async def open_stream(self) -> IMuxedStream:
@@ -119,13 +113,11 @@ class Mplex(IMuxedConn):
         await self.send_message(HeaderTags.NewStream, name.encode(), stream_id)
         return stream
 
-    async def accept_stream(self, stream_id: StreamID, name: str) -> None:
+    async def accept_stream(self) -> IMuxedStream:
         """
         accepts a muxed stream opened by the other end
         """
-        stream = await self._initialize_stream(stream_id, name)
-        # Perform protocol negotiation for the stream.
-        self._tasks.append(asyncio.ensure_future(self.generic_protocol_handler(stream)))
+        return await self.new_stream_queue.get()
 
     async def send_message(
         self, flag: HeaderTags, data: Optional[bytes], stream_id: StreamID
@@ -178,7 +170,11 @@ class Mplex(IMuxedConn):
                         # `NewStream` for the same id is received twice...
                         # TODO: Shutdown
                         pass
-                    await self.accept_stream(stream_id, message.decode())
+                    mplex_stream = await self._initialize_stream(
+                        stream_id, message.decode()
+                    )
+                    # TODO: Check if `self` is shutdown.
+                    await self.new_stream_queue.put(mplex_stream)
                 elif flag in (
                     HeaderTags.MessageInitiator.value,
                     HeaderTags.MessageReceiver.value,

libp2p/stream_muxer/muxer_multistream.py

@@ -2,7 +2,6 @@ from collections import OrderedDict
 from typing import Mapping, Type
 
 from libp2p.network.connection.raw_connection_interface import IRawConnection
-from libp2p.network.typing import GenericProtocolHandlerFn
 from libp2p.peer.id import ID
 from libp2p.protocol_muxer.multiselect import Multiselect
 from libp2p.protocol_muxer.multiselect_client import MultiselectClient
@@ -69,11 +68,6 @@ class MuxerMultistream:
             protocol, _ = await self.multiselect.negotiate(communicator)
         return self.transports[protocol]
 
-    async def new_conn(
-        self,
-        conn: ISecureConn,
-        generic_protocol_handler: GenericProtocolHandlerFn,
-        peer_id: ID,
-    ) -> IMuxedConn:
+    async def new_conn(self, conn: ISecureConn, peer_id: ID) -> IMuxedConn:
         transport_class = await self.select_transport(conn)
-        return transport_class(conn, generic_protocol_handler, peer_id)
+        return transport_class(conn, peer_id)