varun-r-mallya/py-libp2p
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/py-libp2p.git synced 2025-12-31 20:36:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix Mplex and Swarm

Browse Source
This commit is contained in:
mhchia
2019-11-29 19:09:56 +08:00
parent ec43c25b45
commit 1e600ea7e0
13 changed files with 232 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,3 +1,4 @@
import math
import asyncio
import logging
from typing import Any # noqa: F401
@ -18,7 +19,6 @@ from libp2p.utils import (
encode_uvarint,
encode_varint_prefixed,
read_varint_prefixed_bytes,
TrioQueue,
)
from .constants import HeaderTags
@ -41,7 +41,10 @@ class Mplex(IMuxedConn, Service):
next_channel_id: int
streams: Dict[StreamID, MplexStream]
streams_lock: trio.Lock
new_stream_queue: "TrioQueue[IMuxedStream]"
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
@ -64,7 +67,10 @@ class Mplex(IMuxedConn, Service):
# Mapping from stream ID -> buffer of messages for that stream
self.streams = {}
self.streams_lock = trio.Lock()
self.new_stream_queue = TrioQueue()
self.streams_msg_channels = {}
send_channel, receive_channel = trio.open_memory_channel(math.inf)
self.new_stream_send_channel = send_channel
self.new_stream_receive_channel = receive_channel
self.event_shutting_down = trio.Event()
self.event_closed = trio.Event()
@ -105,9 +111,13 @@ class Mplex(IMuxedConn, Service):
return next_id
async def _initialize_stream(self, stream_id: StreamID, name: str) -> MplexStream:
stream = MplexStream(name, stream_id, self)
# Use an unbounded buffer, to avoid `handle_incoming` being blocked when doing
# `send_channel.send`.
send_channel, receive_channel = trio.open_memory_channel(math.inf)
stream = MplexStream(name, stream_id, self, receive_channel)
async with self.streams_lock:
self.streams[stream_id] = stream
self.streams_msg_channels[stream_id] = send_channel
return stream
async def open_stream(self) -> IMuxedStream:
@ -126,7 +136,10 @@ class Mplex(IMuxedConn, Service):
async def accept_stream(self) -> IMuxedStream:
"""accepts a muxed stream opened by the other end."""
return await self.new_stream_queue.get()
try:
return await self.new_stream_receive_channel.receive()
except (trio.ClosedResourceError, trio.EndOfChannel):
raise MplexUnavailable
async def send_message(
self, flag: HeaderTags, data: Optional[bytes], stream_id: StreamID
@ -138,6 +151,9 @@ class Mplex(IMuxedConn, Service):
:param data: data to send in the message
:param stream_id: stream the message is in
"""
print(
f"!@# send_message: {self._id}: flag={flag}, data={data}, stream_id={stream_id}"
)
# << by 3, then or with flag
header = encode_uvarint((stream_id.channel_id << 3) | flag.value)
@ -162,14 +178,21 @@ class Mplex(IMuxedConn, Service):
"""Read a message off of the secured connection and add it to the
corresponding message buffer."""
while True:
while self.manager.is_running:
try:
print(
f"!@# handle_incoming: {self._id}: before _handle_incoming_message"
)
await self._handle_incoming_message()
print(
f"!@# handle_incoming: {self._id}: after _handle_incoming_message"
)
except MplexUnavailable as e:
logger.debug("mplex unavailable while waiting for incoming: %s", e)
print(f"!@# handle_incoming: {self._id}: MplexUnavailable: {e}")
break
# Force context switch
await trio.sleep(0)
print(f"!@# handle_incoming: {self._id}: leaving")
# If we enter here, it means this connection is shutting down.
# We should clean things up.
await self._cleanup()
@ -181,51 +204,73 @@ class Mplex(IMuxedConn, Service):
:return: stream_id, flag, message contents
"""
# FIXME: No timeout is used in Go implementation.
try:
header = await decode_uvarint_from_stream(self.secured_conn)
except (ParseError, RawConnError, IncompleteReadError) as error:
raise MplexUnavailable(
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(
"failed to read messages correctly from the underlying connection"
) from error
except asyncio.TimeoutError as error:
raise MplexUnavailable(
"failed to read more message body within the timeout"
) from error
"failed to read the message body correctly from the underlying connection: "
f"{error}"
)
flag = header & 0x07
channel_id = header >> 3
return channel_id, flag, message
@property
def _id(self) -> int:
return 0 if self.is_initiator else 1
async def _handle_incoming_message(self) -> None:
"""
Read and handle a new incoming message.
:raise MplexUnavailable: `Mplex` encounters fatal error or is shutting down.
"""
print(f"!@# _handle_incoming_message: {self._id}: before reading")
channel_id, flag, message = await self.read_message()
print(
f"!@# _handle_incoming_message: {self._id}: channel_id={channel_id}, flag={flag}, message={message}"
)
stream_id = StreamID(channel_id=channel_id, is_initiator=bool(flag & 1))
print(f"!@# _handle_incoming_message: {self._id}: 2")
if flag == HeaderTags.NewStream.value:
print(f"!@# _handle_incoming_message: {self._id}: 3")
await self._handle_new_stream(stream_id, message)
print(f"!@# _handle_incoming_message: {self._id}: 4")
elif flag in (
HeaderTags.MessageInitiator.value,
HeaderTags.MessageReceiver.value,
):
print(f"!@# _handle_incoming_message: {self._id}: 5")
await self._handle_message(stream_id, message)
print(f"!@# _handle_incoming_message: {self._id}: 6")
elif flag in (HeaderTags.CloseInitiator.value, HeaderTags.CloseReceiver.value):
print(f"!@# _handle_incoming_message: {self._id}: 7")
await self._handle_close(stream_id)
print(f"!@# _handle_incoming_message: {self._id}: 8")
elif flag in (HeaderTags.ResetInitiator.value, HeaderTags.ResetReceiver.value):
print(f"!@# _handle_incoming_message: {self._id}: 9")
await self._handle_reset(stream_id)
print(f"!@# _handle_incoming_message: {self._id}: 10")
else:
print(f"!@# _handle_incoming_message: {self._id}: 11")
# Receives messages with an unknown flag
# TODO: logging
async with self.streams_lock:
print(f"!@# _handle_incoming_message: {self._id}: 12")
if stream_id in self.streams:
print(f"!@# _handle_incoming_message: {self._id}: 13")
stream = self.streams[stream_id]
await stream.reset()
print(f"!@# _handle_incoming_message: {self._id}: 14")
async def _handle_new_stream(self, stream_id: StreamID, message: bytes) -> None:
async with self.streams_lock:
@ -235,43 +280,65 @@ class Mplex(IMuxedConn, Service):
f"received NewStream message for existing stream: {stream_id}"
)
mplex_stream = await self._initialize_stream(stream_id, message.decode())
await self.new_stream_queue.put(mplex_stream)
try:
await self.new_stream_send_channel.send(mplex_stream)
except (trio.BrokenResourceError, trio.EndOfChannel):
raise MplexUnavailable
async def _handle_message(self, stream_id: StreamID, message: bytes) -> None:
print(
f"!@# _handle_message: {self._id}: stream_id={stream_id}, message={message}"
)
async with self.streams_lock:
print(f"!@# _handle_message: {self._id}: 1")
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.
print(f"!@# _handle_message: {self._id}: 2")
return
print(f"!@# _handle_message: {self._id}: 3")
stream = self.streams[stream_id]
send_channel = self.streams_msg_channels[stream_id]
async with stream.close_lock:
print(f"!@# _handle_message: {self._id}: 4")
if stream.event_remote_closed.is_set():
print(f"!@# _handle_message: {self._id}: 5")
# TODO: Warn "Received data from remote after stream was closed by them. (len = %d)" # noqa: E501
return
await stream.incoming_data.put(message)
print(f"!@# _handle_message: {self._id}: 6")
await send_channel.send(message)
print(f"!@# _handle_message: {self._id}: 7")
async def _handle_close(self, stream_id: StreamID) -> None:
print(f"!@# _handle_close: {self._id}: step=0")
async with self.streams_lock:
if stream_id not in self.streams:
# Ignore unmatched messages for now.
return
stream = self.streams[stream_id]
send_channel = self.streams_msg_channels[stream_id]
print(f"!@# _handle_close: {self._id}: step=1")
await send_channel.aclose()
print(f"!@# _handle_close: {self._id}: step=2")
# 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
print(f"!@# _handle_close: {self._id}: step=3")
is_local_closed: bool
async with stream.close_lock:
stream.event_remote_closed.set()
is_local_closed = stream.event_local_closed.is_set()
print(f"!@# _handle_close: {self._id}: step=4")
# 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:
if stream_id in self.streams:
del self.streams[stream_id]
print(f"!@# _handle_close: {self._id}: step=5")
async def _handle_reset(self, stream_id: StreamID) -> None:
async with self.streams_lock:
@ -279,11 +346,11 @@ class Mplex(IMuxedConn, Service):
# This is *ok*. We forget the stream on reset.
return
stream = self.streams[stream_id]
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():
@ -291,16 +358,21 @@ class Mplex(IMuxedConn, Service):
async with self.streams_lock:
if stream_id in self.streams:
del self.streams[stream_id]
del self.streams_msg_channels[stream_id]
async def _cleanup(self) -> None:
if not self.event_shutting_down.is_set():
self.event_shutting_down.set()
async with self.streams_lock:
for stream in self.streams.values():
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()
send_channel = self.streams_msg_channels[stream_id]
await send_channel.aclose()
self.streams = None
self.event_closed.set()
await self.new_stream_send_channel.aclose()
await self.new_stream_receive_channel.aclose()