StarMirror/py-libp2p
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Conform stream_id to go-mplex

Browse Source
This commit is contained in:
mhchia 2019-08-28 21:43:34 +08:00
parent 9b60e1757d
commit d35b8ffc64
No known key found for this signature in database
GPG Key ID: 389EFBEA1362589A
5 changed files with 55 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
libp2p/peer/id.py
View File

@ -41,7 +41,7 @@ class ID:
def __str__(self) -> str:
if FRIENDLY_IDS:
return self.to_string()[2:8]
return f"<peer.ID {self.to_string()[2:8]}>"
else:
return self.to_string()

9
libp2p/stream_muxer/abc.py
View File

@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Optional
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
@ -51,7 +52,7 @@ class IMuxedConn(ABC):
"""
@abstractmethod
async def read_buffer(self, stream_id: int) -> bytes:
async def read_buffer(self, stream_id: StreamID) -> bytes:
"""
Read a message from stream_id's buffer, check raw connection for new messages
:param stream_id: stream id of stream to read from
@ -59,7 +60,7 @@ class IMuxedConn(ABC):
"""
@abstractmethod
async def read_buffer_nonblocking(self, stream_id: int) -> Optional[bytes]:
async def read_buffer_nonblocking(self, stream_id: StreamID) -> Optional[bytes]:
"""
Read a message from `stream_id`'s buffer, non-blockingly.
"""
@ -78,7 +79,9 @@ class IMuxedConn(ABC):
"""
@abstractmethod
async def send_message(self, flag: HeaderTags, data: bytes, stream_id: int) -> int:
async def send_message(
self, flag: HeaderTags, data: bytes, stream_id: StreamID
) -> int:
"""
sends a message over the connection
:param header: header to use

48
libp2p/stream_muxer/mplex/mplex.py
View File

@ -13,6 +13,7 @@ from libp2p.utils import (
)
from .constants import HeaderTags
from .datastructures import StreamID
from .exceptions import StreamNotFound
from .mplex_stream import MplexStream
@ -29,9 +30,9 @@ class Mplex(IMuxedConn):
# 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
buffers: Dict[StreamID, "asyncio.Queue[bytes]"]
stream_queue: "asyncio.Queue[StreamID]"
next_channel_id: int
# TODO: `generic_protocol_handler` should be refactored out of mplex conn.
def __init__(
@ -49,10 +50,7 @@ class Mplex(IMuxedConn):
"""
self.secured_conn = secured_conn
if self.secured_conn.initiator:
self.next_stream_id = 0
else:
self.next_stream_id = 1
self.next_channel_id = 0
# Store generic protocol handler
self.generic_protocol_handler = generic_protocol_handler
@ -85,7 +83,7 @@ class Mplex(IMuxedConn):
"""
raise NotImplementedError()
async def read_buffer(self, stream_id: int) -> bytes:
async def read_buffer(self, stream_id: StreamID) -> bytes:
"""
Read a message from buffer of the stream specified by `stream_id`,
check secured connection for new messages.
@ -97,7 +95,7 @@ class Mplex(IMuxedConn):
raise StreamNotFound(f"stream {stream_id} is not found")
return await self.buffers[stream_id].get()
async def read_buffer_nonblocking(self, stream_id: int) -> Optional[bytes]:
async def read_buffer_nonblocking(self, stream_id: StreamID) -> Optional[bytes]:
"""
Read a message from buffer of the stream specified by `stream_id`, non-blockingly.
`StreamNotFound` is raised when stream `stream_id` is not found in `Mplex`.
@ -108,13 +106,13 @@ class Mplex(IMuxedConn):
return None
return await self.buffers[stream_id].get()
def _get_next_stream_id(self) -> int:
def _get_next_channel_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
next_id = self.next_channel_id
self.next_channel_id += 1
return next_id
async def open_stream(self) -> IMuxedStream:
@ -122,11 +120,12 @@ class Mplex(IMuxedConn):
creates a new muxed_stream
:return: a new ``MplexStream``
"""
stream_id = self._get_next_stream_id()
name = str(stream_id)
stream = MplexStream(name, stream_id, True, self)
channel_id = self._get_next_channel_id()
stream_id = StreamID(channel_id=channel_id, is_initiator=True)
name = str(channel_id)
stream = MplexStream(name, stream_id, self)
self.buffers[stream_id] = asyncio.Queue()
# Default stream name is the `stream_id`
# Default stream name is the `channel_id`
await self.send_message(HeaderTags.NewStream, name.encode(), stream_id)
return stream
@ -135,10 +134,12 @@ class Mplex(IMuxedConn):
accepts a muxed stream opened by the other end
"""
stream_id = await self.stream_queue.get()
stream = MplexStream(name, stream_id, False, self)
stream = MplexStream(name, stream_id, self)
asyncio.ensure_future(self.generic_protocol_handler(stream))
async def send_message(self, flag: HeaderTags, data: bytes, stream_id: int) -> int:
async def send_message(
self, flag: HeaderTags, data: bytes, stream_id: StreamID
) -> int:
"""
sends a message over the connection
:param header: header to use
@ -146,7 +147,7 @@ class Mplex(IMuxedConn):
:param stream_id: stream the message is in
"""
# << by 3, then or with flag
header = (stream_id << 3) | flag.value
header = (stream_id.channel_id << 3) | flag.value
header = encode_uvarint(header)
if data is None:
@ -174,9 +175,10 @@ class Mplex(IMuxedConn):
# TODO Deal with other types of messages using flag (currently _)
while True:
stream_id, flag, message = await self.read_message()
channel_id, flag, message = await self.read_message()
if stream_id is not None and flag is not None and message is not None:
if channel_id is not None and flag is not None and message is not None:
stream_id = StreamID(channel_id=channel_id, is_initiator=bool(flag & 1))
if stream_id not in self.buffers:
self.buffers[stream_id] = asyncio.Queue()
await self.stream_queue.put(stream_id)
@ -214,6 +216,6 @@ class Mplex(IMuxedConn):
return None, None, None
flag = header & 0x07
stream_id = header >> 3
channel_id = header >> 3
return stream_id, flag, message
return channel_id, flag, message

24
libp2p/stream_muxer/mplex/mplex_stream.py
View File

@ -3,6 +3,7 @@ import asyncio
from libp2p.stream_muxer.abc import IMuxedConn, IMuxedStream
from .constants import HeaderTags
from .datastructures import StreamID
class MplexStream(IMuxedStream):
@ -11,8 +12,7 @@ class MplexStream(IMuxedStream):
"""
name: str
stream_id: int
initiator: bool
stream_id: StreamID
mplex_conn: IMuxedConn
read_deadline: int
write_deadline: int
@ -22,18 +22,14 @@ class MplexStream(IMuxedStream):
_buf: bytearray
def __init__(
self, name: str, stream_id: int, initiator: bool, mplex_conn: IMuxedConn
) -> None:
def __init__(self, name: str, stream_id: StreamID, mplex_conn: IMuxedConn) -> None:
"""
create new MuxedStream in muxer
:param stream_id: stream stream id
:param initiator: boolean if this is an initiator
:param stream_id: stream id of this stream
:param mplex_conn: muxed connection of this muxed_stream
"""
self.name = name
self.stream_id = stream_id
self.initiator = initiator
self.mplex_conn = mplex_conn
self.read_deadline = None
self.write_deadline = None
@ -42,6 +38,10 @@ class MplexStream(IMuxedStream):
self.stream_lock = asyncio.Lock()
self._buf = bytearray()
@property
def is_initiator(self) -> bool:
return self.stream_id.is_initiator
async def read(self, n: int = -1) -> bytes:
"""
Read up to n bytes. Read possibly returns fewer than `n` bytes,
@ -85,7 +85,7 @@ class MplexStream(IMuxedStream):
"""
flag = (
HeaderTags.MessageInitiator
if self.initiator
if self.is_initiator
else HeaderTags.MessageReceiver
)
return await self.mplex_conn.send_message(flag, data, self.stream_id)
@ -98,7 +98,9 @@ class MplexStream(IMuxedStream):
"""
# TODO error handling with timeout
# TODO understand better how mutexes are used from go repo
flag = HeaderTags.CloseInitiator if self.initiator else HeaderTags.CloseReceiver
flag = (
HeaderTags.CloseInitiator if self.is_initiator else HeaderTags.CloseReceiver
)
await self.mplex_conn.send_message(flag, None, self.stream_id)
remote_lock = False
@ -131,7 +133,7 @@ class MplexStream(IMuxedStream):
if not self.remote_closed:
flag = (
HeaderTags.ResetInitiator
if self.initiator
if self.is_initiator
else HeaderTags.ResetReceiver
)
await self.mplex_conn.send_message(flag, None, self.stream_id)

18
tests/pubsub/conftest.py
View File

@ -23,14 +23,16 @@ async def hosts(num_hosts):
await asyncio.gather(
*[_host.get_network().listen(LISTEN_MADDR) for _host in _hosts]
)
yield _hosts
# Clean up
listeners = []
for _host in _hosts:
for listener in _host.get_network().listeners.values():
listener.server.close()
listeners.append(listener)
await asyncio.gather(*[listener.server.wait_closed() for listener in listeners])
try:
yield _hosts
finally:
# Clean up
listeners = []
for _host in _hosts:
for listener in _host.get_network().listeners.values():
listener.server.close()
listeners.append(listener)
await asyncio.gather(*[listener.server.wait_closed() for listener in listeners])
@pytest.fixture