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

This commit is contained in:
Alex Stokes 2019-08-24 21:12:08 +02:00
parent e29c1507bf
commit 9c5fb4fa5a
No known key found for this signature in database
GPG Key ID: 51CE1721B245C086
4 changed files with 16 additions and 20 deletions

View File

@ -9,7 +9,6 @@ class RawConnection(IRawConnection):
initiator: bool initiator: bool
_drain_lock: asyncio.Lock _drain_lock: asyncio.Lock
_next_id: int
def __init__( def __init__(
self, self,
@ -22,7 +21,6 @@ class RawConnection(IRawConnection):
self.initiator = initiator self.initiator = initiator
self._drain_lock = asyncio.Lock() self._drain_lock = asyncio.Lock()
self._next_id = 0 if initiator else 1
async def write(self, data: bytes) -> None: async def write(self, data: bytes) -> None:
self.writer.write(data) self.writer.write(data)
@ -41,12 +39,3 @@ class RawConnection(IRawConnection):
def close(self) -> None: def close(self) -> None:
self.writer.close() 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

View File

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

View File

@ -30,10 +30,6 @@ class BaseSession(ISecureConn):
self.initiator = self.conn.initiator 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: async def write(self, data: bytes) -> None:
await self.conn.write(data) await self.conn.write(data)

View File

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