Encapsulate concept of a "stream id" to a "muxed" connection
This commit is contained in:
parent
e29c1507bf
commit
9c5fb4fa5a
|
@ -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
|
|
||||||
|
|
|
@ -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
|
|
||||||
|
|
|
@ -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)
|
||||||
|
|
||||||
|
|
|
@ -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)
|
||||||
|
|
Loading…
Reference in New Issue
Block a user