Merge pull request #263 from ralexstokes/remove-stream-from-connection

Encapsulate the concept of a stream to the stream multiplexer

libp2p/network/connection/raw_connection.py

@@ -4,32 +4,23 @@ from .raw_connection_interface import IRawConnection
 
 
 class RawConnection(IRawConnection):
-
-    conn_ip: str
-    conn_port: str
     reader: asyncio.StreamReader
     writer: asyncio.StreamWriter
     initiator: bool
 
     _drain_lock: asyncio.Lock
-    _next_id: int
 
     def __init__(
         self,
-        ip: str,
-        port: str,
         reader: asyncio.StreamReader,
         writer: asyncio.StreamWriter,
         initiator: bool,
     ) -> None:
-        self.conn_ip = ip
-        self.conn_port = port
         self.reader = reader
         self.writer = writer
         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)
@@ -48,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/network/swarm.py

@@ -204,13 +204,7 @@ class Swarm(INetwork):
             ) -> None:
                 # Upgrade reader/write to a net_stream and pass \
                 # to appropriate stream handler (using multiaddr)
-                raw_conn = RawConnection(
-                    maddr.value_for_protocol("ip4"),
-                    maddr.value_for_protocol("tcp"),
-                    reader,
-                    writer,
-                    False,
-                )
+                raw_conn = RawConnection(reader, writer, False)
 
                 # Per, https://discuss.libp2p.io/t/multistream-security/130, we first secure
                 # the conn and then mux the conn

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

@@ -3,7 +3,6 @@ from typing import Dict, Optional, Tuple
 
 from multiaddr import Multiaddr
 
-from libp2p.network.connection.raw_connection_interface import IRawConnection
 from libp2p.network.typing import GenericProtocolHandlerFn
 from libp2p.peer.id import ID
 from libp2p.security.secure_conn_interface import ISecureConn
@@ -24,13 +23,13 @@ class Mplex(IMuxedConn):
     """
 
     secured_conn: ISecureConn
-    raw_conn: IRawConnection
     peer_id: ID
     # TODO: `dataIn` in go implementation. Should be size of 8.
     # TODO: Also, `dataIn` is closed indicating EOF in Go. We don't have similar strategies
     #   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,
@@ -47,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
 
@@ -100,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
@@ -110,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)

libp2p/transport/tcp/tcp.py

@@ -70,12 +70,12 @@ class TCP(ITransport):
         :param self_id: peer_id of the dialer (to send to receiver)
         :return: `RawConnection` if successful
         """
-        host = maddr.value_for_protocol("ip4")
-        port = maddr.value_for_protocol("tcp")
+        self.host = maddr.value_for_protocol("ip4")
+        self.port = int(maddr.value_for_protocol("tcp"))
 
-        reader, writer = await asyncio.open_connection(host, int(port))
+        reader, writer = await asyncio.open_connection(self.host, self.port)
 
-        return RawConnection(host, port, reader, writer, True)
+        return RawConnection(reader, writer, True)
 
     def create_listener(self, handler_function: THandler) -> TCPListener:
         """