Encapsulate concept of a "stream id" to a "muxed" connection

libp2p/network/connection/raw_connection.py

@@ -9,7 +9,6 @@ class RawConnection(IRawConnection):
     initiator: bool
 
     _drain_lock: asyncio.Lock
-    _next_id: int
 
     def __init__(
         self,
@@ -22,7 +21,6 @@ class RawConnection(IRawConnection):
         self.initiator = initiator
 
         self._drain_lock = asyncio.Lock()
-        self._next_id = 0 if initiator else 1
 
     async def write(self, data: bytes) -> None:
         self.writer.write(data)
@@ -41,12 +39,3 @@ class RawConnection(IRawConnection):
 
     def close(self) -> None:
         self.writer.close()
-
-    def next_stream_id(self) -> int:
-        """
-        Get next available stream id
-        :return: next available stream id for the connection
-        """
-        next_id = self._next_id
-        self._next_id += 2
-        return next_id

libp2p/network/connection/raw_connection_interface.py

@@ -19,7 +19,3 @@ class IRawConnection(ABC):
     @abstractmethod
     def close(self) -> None:
         pass
-
-    @abstractmethod
-    def next_stream_id(self) -> int:
-        pass

libp2p/security/base_session.py

@@ -30,10 +30,6 @@ class BaseSession(ISecureConn):
 
         self.initiator = self.conn.initiator
 
-    # TODO clean up how this is passed around?
-    def next_stream_id(self) -> int:
-        return self.conn.next_stream_id()
-
     async def write(self, data: bytes) -> None:
         await self.conn.write(data)
 

libp2p/stream_muxer/mplex/mplex.py

@@ -29,6 +29,7 @@ class Mplex(IMuxedConn):
     #   to let the `MplexStream`s know that EOF arrived (#235).
     buffers: Dict[int, "asyncio.Queue[bytes]"]
     stream_queue: "asyncio.Queue[int]"
+    next_stream_id: int
 
     def __init__(
         self,
@@ -45,6 +46,11 @@ class Mplex(IMuxedConn):
         """
         self.conn = secured_conn
 
+        if self.conn.initiator:
+            self.next_stream_id = 0
+        else:
+            self.next_stream_id = 1
+
         # Store generic protocol handler
         self.generic_protocol_handler = generic_protocol_handler
 
@@ -98,6 +104,15 @@ class Mplex(IMuxedConn):
             return None
         return await self.buffers[stream_id].get()
 
+    def _get_next_stream_id(self) -> int:
+        """
+        Get next available stream id
+        :return: next available stream id for the connection
+        """
+        next_id = self.next_stream_id
+        self.next_stream_id += 2
+        return next_id
+
     # FIXME: Remove multiaddr from being passed into muxed_conn
     async def open_stream(
         self, protocol_id: str, multi_addr: Multiaddr
@@ -108,7 +123,7 @@ class Mplex(IMuxedConn):
         :param multi_addr: multi_addr that stream connects to
         :return: a new stream
         """
-        stream_id = self.conn.next_stream_id()
+        stream_id = self._get_next_stream_id()
         stream = MplexStream(stream_id, True, self)
         self.buffers[stream_id] = asyncio.Queue()
         await self.send_message(HeaderTags.NewStream, None, stream_id)