py-libp2p/libp2p/stream_muxer/mplex/mplex.py

338 lines
13 KiB
Python
Raw Normal View History

2019-11-08 12:57:43 +08:00
import logging
2019-12-06 17:06:37 +08:00
from typing import Dict, Optional, Tuple
2019-08-02 18:28:04 +08:00
import trio
from libp2p.exceptions import ParseError
from libp2p.io.exceptions import IncompleteReadError
from libp2p.network.connection.exceptions import RawConnError
from libp2p.peer.id import ID
2019-08-05 11:22:44 +08:00
from libp2p.security.secure_conn_interface import ISecureConn
from libp2p.stream_muxer.abc import IMuxedConn, IMuxedStream
from libp2p.typing import TProtocol
from libp2p.utils import (
decode_uvarint_from_stream,
encode_uvarint,
encode_varint_prefixed,
read_varint_prefixed_bytes,
)
2019-01-10 02:38:56 +08:00
2019-08-02 17:14:43 +08:00
from .constants import HeaderTags
2019-08-28 21:43:34 +08:00
from .datastructures import StreamID
from .exceptions import MplexUnavailable
2019-08-03 13:36:19 +08:00
from .mplex_stream import MplexStream
2018-11-01 05:31:00 +08:00
MPLEX_PROTOCOL_ID = TProtocol("/mplex/6.7.0")
# Ref: https://github.com/libp2p/go-mplex/blob/414db61813d9ad3e6f4a7db5c1b1612de343ace9/multiplex.go#L115 # noqa: E501
MPLEX_MESSAGE_CHANNEL_SIZE = 8
2019-11-08 12:57:43 +08:00
logger = logging.getLogger("libp2p.stream_muxer.mplex.mplex")
2019-08-02 17:53:51 +08:00
class Mplex(IMuxedConn):
2018-11-01 05:31:00 +08:00
"""
reference: https://github.com/libp2p/go-mplex/blob/master/multiplex.go
"""
secured_conn: ISecureConn
peer_id: ID
2019-08-28 21:43:34 +08:00
next_channel_id: int
2019-09-05 22:29:33 +08:00
streams: Dict[StreamID, MplexStream]
streams_lock: trio.Lock
2019-11-29 19:09:56 +08:00
streams_msg_channels: Dict[StreamID, "trio.MemorySendChannel[bytes]"]
new_stream_send_channel: "trio.MemorySendChannel[IMuxedStream]"
new_stream_receive_channel: "trio.MemoryReceiveChannel[IMuxedStream]"
event_shutting_down: trio.Event
event_closed: trio.Event
event_started: trio.Event
2019-08-02 17:53:51 +08:00
def __init__(self, secured_conn: ISecureConn, peer_id: ID) -> None:
2018-11-01 05:31:00 +08:00
"""
create a new muxed connection.
:param secured_conn: an instance of ``ISecureConn``
:param generic_protocol_handler: generic protocol handler
for new muxed streams
[WIP] PubSub and FloodSub development (#133) * Add notifee interface * Add notify function to network interface * Implement notify feature * Add tests for notify * Make notifee functions all async * Fix linting issue * Fix linting issue * Scaffold pubsub router interface * Scaffold pubsub directory * Store peer_id in muxed connection * Implement pubsub notifee * Remove outdated files * Implement pubsub first attempt * Prepare pubsub for floodsub * Add mplex conn to net stream and add conn in notify tests * Implement floodsub * Use NetStream in generic protocol handler * Debugging async issues * Modify test to perform proper assert. Test passes * Remove callbacks. Reduce sleep time * Add simple three node test * Clean up code. Add message classes * Add test for two topics * Add conn to net stream and conn tests * Refactor test setup to remove duplicate code * Fix linting issues * Fix linting issue * Fix linting issue * Fix outstanding unrelated lint issue in multiselect_client * Add connect function * Remove debug prints * Remove debug prints from floodsub * Use MessageTalk in place of direct message breakdown * Remove extra prints * Remove outdated function * Add message to queues for all topics in message * Debugging * Add message self delivery * Increase read timeout to 5 to get pubsub tests passing * Refactor testing helper func. Add tests * Add tests and increase timeout to get tests passing * Add dummy account demo scaffolding * Attempt to use threads. Test fails * Implement basic dummy node tests using threads * Add generic testing function * Add simple seven node tree test * Add more complex seven node tree tests * Add five node ring tests * Remove unnecessary get_message_type func * Add documentation to classes * Add message id to messages * Add documentation to test helper func * Add docs to dummy account node helper func * Add more docs to dummy account node test helper func * fixed linting errors in floodsub * small notify bugfix * move pubsub into libp2p * fixed pubsub linting * fixing pubsub test failures * linting
2019-03-24 01:52:02 +08:00
:param peer_id: peer_id of peer the connection is to
2018-11-01 05:31:00 +08:00
"""
self.secured_conn = secured_conn
2018-11-26 00:05:56 +08:00
2019-08-28 21:43:34 +08:00
self.next_channel_id = 0
[WIP] PubSub and FloodSub development (#133) * Add notifee interface * Add notify function to network interface * Implement notify feature * Add tests for notify * Make notifee functions all async * Fix linting issue * Fix linting issue * Scaffold pubsub router interface * Scaffold pubsub directory * Store peer_id in muxed connection * Implement pubsub notifee * Remove outdated files * Implement pubsub first attempt * Prepare pubsub for floodsub * Add mplex conn to net stream and add conn in notify tests * Implement floodsub * Use NetStream in generic protocol handler * Debugging async issues * Modify test to perform proper assert. Test passes * Remove callbacks. Reduce sleep time * Add simple three node test * Clean up code. Add message classes * Add test for two topics * Add conn to net stream and conn tests * Refactor test setup to remove duplicate code * Fix linting issues * Fix linting issue * Fix linting issue * Fix outstanding unrelated lint issue in multiselect_client * Add connect function * Remove debug prints * Remove debug prints from floodsub * Use MessageTalk in place of direct message breakdown * Remove extra prints * Remove outdated function * Add message to queues for all topics in message * Debugging * Add message self delivery * Increase read timeout to 5 to get pubsub tests passing * Refactor testing helper func. Add tests * Add tests and increase timeout to get tests passing * Add dummy account demo scaffolding * Attempt to use threads. Test fails * Implement basic dummy node tests using threads * Add generic testing function * Add simple seven node tree test * Add more complex seven node tree tests * Add five node ring tests * Remove unnecessary get_message_type func * Add documentation to classes * Add message id to messages * Add documentation to test helper func * Add docs to dummy account node helper func * Add more docs to dummy account node test helper func * fixed linting errors in floodsub * small notify bugfix * move pubsub into libp2p * fixed pubsub linting * fixing pubsub test failures * linting
2019-03-24 01:52:02 +08:00
# Set peer_id
self.peer_id = peer_id
2018-11-26 00:05:56 +08:00
# Mapping from stream ID -> buffer of messages for that stream
2019-09-05 22:29:33 +08:00
self.streams = {}
self.streams_lock = trio.Lock()
2019-11-29 19:09:56 +08:00
self.streams_msg_channels = {}
channels = trio.open_memory_channel[IMuxedStream](0)
2019-12-06 17:06:37 +08:00
self.new_stream_send_channel, self.new_stream_receive_channel = channels
self.event_shutting_down = trio.Event()
self.event_closed = trio.Event()
self.event_started = trio.Event()
async def start(self) -> None:
await self.handle_incoming()
2018-11-01 05:31:00 +08:00
@property
2019-10-25 01:25:34 +08:00
def is_initiator(self) -> bool:
2019-10-25 01:28:19 +08:00
return self.secured_conn.is_initiator
async def close(self) -> None:
"""close the stream muxer and underlying secured connection."""
if self.event_shutting_down.is_set():
return
# Set the `event_shutting_down`, to allow graceful shutdown.
self.event_shutting_down.set()
await self.secured_conn.close()
# Blocked until `close` is finally set.
await self.event_closed.wait()
2018-11-01 05:31:00 +08:00
@property
2019-08-02 17:53:51 +08:00
def is_closed(self) -> bool:
2018-11-01 05:31:00 +08:00
"""
check connection is fully closed.
2018-11-01 05:31:00 +08:00
:return: true if successful
"""
return self.event_closed.is_set()
2018-11-01 05:31:00 +08:00
2019-08-28 21:43:34 +08:00
def _get_next_channel_id(self) -> int:
"""
Get next available stream id.
:return: next available stream id for the connection
"""
2019-08-28 21:43:34 +08:00
next_id = self.next_channel_id
self.next_channel_id += 1
return next_id
2019-09-05 22:29:33 +08:00
async def _initialize_stream(self, stream_id: StreamID, name: str) -> MplexStream:
send_channel, receive_channel = trio.open_memory_channel[bytes](
MPLEX_MESSAGE_CHANNEL_SIZE
)
stream = MplexStream(name, stream_id, self, receive_channel)
2019-09-05 22:29:33 +08:00
async with self.streams_lock:
self.streams[stream_id] = stream
self.streams_msg_channels[stream_id] = send_channel
2019-09-05 22:29:33 +08:00
return stream
async def open_stream(self) -> IMuxedStream:
2018-11-01 05:31:00 +08:00
"""
creates a new muxed_stream.
:return: a new ``MplexStream``
2018-11-01 05:31:00 +08:00
"""
2019-08-28 21:43:34 +08:00
channel_id = self._get_next_channel_id()
stream_id = StreamID(channel_id=channel_id, is_initiator=True)
# Default stream name is the `channel_id`
2019-09-05 22:29:33 +08:00
name = str(channel_id)
stream = await self._initialize_stream(stream_id, name)
await self.send_message(HeaderTags.NewStream, name.encode(), stream_id)
return stream
async def accept_stream(self) -> IMuxedStream:
"""accepts a muxed stream opened by the other end."""
2019-11-29 19:09:56 +08:00
try:
return await self.new_stream_receive_channel.receive()
except trio.EndOfChannel:
2019-11-29 19:09:56 +08:00
raise MplexUnavailable
2019-08-28 21:43:34 +08:00
async def send_message(
self, flag: HeaderTags, data: Optional[bytes], stream_id: StreamID
2019-08-28 21:43:34 +08:00
) -> int:
"""
sends a message over the connection.
:param flag: header to use
:param data: data to send in the message
2018-11-12 06:38:11 +08:00
:param stream_id: stream the message is in
"""
2018-11-12 06:38:11 +08:00
# << by 3, then or with flag
header = encode_uvarint((stream_id.channel_id << 3) | flag.value)
2018-11-19 00:22:17 +08:00
if data is None:
data = b""
_bytes = header + encode_varint_prefixed(data)
2018-11-13 00:00:43 +08:00
2019-11-19 14:01:12 +08:00
return await self.write_to_stream(_bytes)
2018-11-12 06:38:11 +08:00
async def write_to_stream(self, _bytes: bytes) -> None:
"""
writes a byte array to a secured connection.
:param _bytes: byte array to write
:return: length written
"""
try:
await self.secured_conn.write(_bytes)
except RawConnError as e:
raise MplexUnavailable(
"failed to write message to the underlying connection"
) from e
2019-08-02 17:53:51 +08:00
async def handle_incoming(self) -> None:
"""Read a message off of the secured connection and add it to the
corresponding message buffer."""
self.event_started.set()
while True:
try:
await self._handle_incoming_message()
2019-11-08 12:57:43 +08:00
except MplexUnavailable as e:
logger.debug("mplex unavailable while waiting for incoming: %s", e)
break
# If we enter here, it means this connection is shutting down.
# We should clean things up.
await self._cleanup()
2019-08-02 17:53:51 +08:00
async def read_message(self) -> Tuple[int, int, bytes]:
"""
Read a single message off of the secured connection.
:return: stream_id, flag, message contents
"""
try:
header = await decode_uvarint_from_stream(self.secured_conn)
except (ParseError, RawConnError, IncompleteReadError) as error:
raise MplexUnavailable(
2019-11-29 19:09:56 +08:00
f"failed to read the header correctly from the underlying connection: {error}"
)
try:
message = await read_varint_prefixed_bytes(self.secured_conn)
except (ParseError, RawConnError, IncompleteReadError) as error:
raise MplexUnavailable(
2019-11-29 19:09:56 +08:00
"failed to read the message body correctly from the underlying connection: "
f"{error}"
)
flag = header & 0x07
2019-08-28 21:43:34 +08:00
channel_id = header >> 3
2019-08-28 21:43:34 +08:00
return channel_id, flag, message
async def _handle_incoming_message(self) -> None:
"""
Read and handle a new incoming message.
:raise MplexUnavailable: `Mplex` encounters fatal error or is shutting down.
"""
2019-11-19 14:01:12 +08:00
channel_id, flag, message = await self.read_message()
stream_id = StreamID(channel_id=channel_id, is_initiator=bool(flag & 1))
if flag == HeaderTags.NewStream.value:
await self._handle_new_stream(stream_id, message)
elif flag in (
HeaderTags.MessageInitiator.value,
HeaderTags.MessageReceiver.value,
):
await self._handle_message(stream_id, message)
elif flag in (HeaderTags.CloseInitiator.value, HeaderTags.CloseReceiver.value):
await self._handle_close(stream_id)
elif flag in (HeaderTags.ResetInitiator.value, HeaderTags.ResetReceiver.value):
await self._handle_reset(stream_id)
else:
# Receives messages with an unknown flag
# TODO: logging
async with self.streams_lock:
if stream_id in self.streams:
stream = self.streams[stream_id]
await stream.reset()
async def _handle_new_stream(self, stream_id: StreamID, message: bytes) -> None:
async with self.streams_lock:
if stream_id in self.streams:
# `NewStream` for the same id is received twice...
raise MplexUnavailable(
f"received NewStream message for existing stream: {stream_id}"
)
mplex_stream = await self._initialize_stream(stream_id, message.decode())
2019-11-29 19:09:56 +08:00
try:
await self.new_stream_send_channel.send(mplex_stream)
except trio.ClosedResourceError:
2019-11-29 19:09:56 +08:00
raise MplexUnavailable
async def _handle_message(self, stream_id: StreamID, message: bytes) -> None:
async with self.streams_lock:
if stream_id not in self.streams:
# We receive a message of the stream `stream_id` which is not accepted
# before. It is abnormal. Possibly disconnect?
# TODO: Warn and emit logs about this.
return
stream = self.streams[stream_id]
2019-11-29 19:09:56 +08:00
send_channel = self.streams_msg_channels[stream_id]
async with stream.close_lock:
if stream.event_remote_closed.is_set():
# TODO: Warn "Received data from remote after stream was closed by them. (len = %d)" # noqa: E501
return
try:
send_channel.send_nowait(message)
except (trio.BrokenResourceError, trio.ClosedResourceError):
raise MplexUnavailable
except trio.WouldBlock:
# `send_channel` is full, reset this stream.
logger.warning(
"message channel of stream %s is full: stream is reset", stream_id
)
await stream.reset()
async def _handle_close(self, stream_id: StreamID) -> None:
async with self.streams_lock:
if stream_id not in self.streams:
# Ignore unmatched messages for now.
return
stream = self.streams[stream_id]
2019-11-29 19:09:56 +08:00
send_channel = self.streams_msg_channels[stream_id]
await send_channel.aclose()
# NOTE: If remote is already closed, then return: Technically a bug
# on the other side. We should consider killing the connection.
async with stream.close_lock:
if stream.event_remote_closed.is_set():
return
is_local_closed: bool
async with stream.close_lock:
stream.event_remote_closed.set()
is_local_closed = stream.event_local_closed.is_set()
# If local is also closed, both sides are closed. Then, we should clean up
# the entry of this stream, to avoid others from accessing it.
if is_local_closed:
async with self.streams_lock:
self.streams.pop(stream_id, None)
async def _handle_reset(self, stream_id: StreamID) -> None:
async with self.streams_lock:
if stream_id not in self.streams:
# This is *ok*. We forget the stream on reset.
return
stream = self.streams[stream_id]
2019-11-29 19:09:56 +08:00
send_channel = self.streams_msg_channels[stream_id]
await send_channel.aclose()
async with stream.close_lock:
if not stream.event_remote_closed.is_set():
stream.event_reset.set()
stream.event_remote_closed.set()
# If local is not closed, we should close it.
if not stream.event_local_closed.is_set():
stream.event_local_closed.set()
async with self.streams_lock:
self.streams.pop(stream_id, None)
self.streams_msg_channels.pop(stream_id, None)
async def _cleanup(self) -> None:
if not self.event_shutting_down.is_set():
self.event_shutting_down.set()
async with self.streams_lock:
2019-11-29 19:09:56 +08:00
for stream_id, stream in self.streams.items():
async with stream.close_lock:
if not stream.event_remote_closed.is_set():
stream.event_remote_closed.set()
stream.event_reset.set()
stream.event_local_closed.set()
2019-11-29 19:09:56 +08:00
send_channel = self.streams_msg_channels[stream_id]
await send_channel.aclose()
self.event_closed.set()
2019-11-29 19:09:56 +08:00
await self.new_stream_send_channel.aclose()