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

236 lines
7.9 KiB
Python
Raw Normal View History

2019-09-09 22:48:49 +08:00
from typing import TYPE_CHECKING
2018-11-12 07:03:04 +08:00
import trio
from libp2p.stream_muxer.abc import IMuxedStream
from libp2p.stream_muxer.exceptions import MuxedConnUnavailable
2018-11-01 05:31:00 +08:00
from .constants import HeaderTags
2019-08-28 21:43:34 +08:00
from .datastructures import StreamID
2019-09-09 15:45:35 +08:00
from .exceptions import MplexStreamClosed, MplexStreamEOF, MplexStreamReset
2019-08-02 18:28:04 +08:00
if TYPE_CHECKING:
from libp2p.stream_muxer.mplex.mplex import Mplex
2019-08-02 18:28:04 +08:00
2018-11-21 09:28:41 +08:00
class MplexStream(IMuxedStream):
2018-11-01 06:02:00 +08:00
"""
reference: https://github.com/libp2p/go-mplex/blob/master/stream.go
"""
name: str
2019-08-28 21:43:34 +08:00
stream_id: StreamID
muxed_conn: "Mplex"
2019-08-05 17:02:18 +08:00
read_deadline: int
write_deadline: int
2019-11-29 19:09:56 +08:00
# TODO: Add lock for read/write to avoid interleaving receiving messages?
2019-11-19 14:01:12 +08:00
close_lock: trio.Lock
# NOTE: `dataIn` is size of 8 in Go implementation.
2019-11-29 19:09:56 +08:00
incoming_data_channel: "trio.MemoryReceiveChannel[bytes]"
2019-09-05 22:29:33 +08:00
2019-11-19 14:01:12 +08:00
event_local_closed: trio.Event
event_remote_closed: trio.Event
event_reset: trio.Event
2019-08-02 18:28:04 +08:00
_buf: bytearray
2019-08-07 15:23:20 +08:00
2019-11-29 19:09:56 +08:00
def __init__(
self,
name: str,
stream_id: StreamID,
muxed_conn: "Mplex",
incoming_data_channel: "trio.MemoryReceiveChannel[bytes]",
) -> None:
"""
create new MuxedStream in muxer.
2019-08-28 21:43:34 +08:00
:param stream_id: stream id of this stream
:param muxed_conn: muxed connection of this muxed_stream
2018-11-12 07:03:04 +08:00
"""
self.name = name
2018-11-13 02:02:49 +08:00
self.stream_id = stream_id
self.muxed_conn = muxed_conn
2018-11-13 02:02:49 +08:00
self.read_deadline = None
self.write_deadline = None
2019-11-19 14:01:12 +08:00
self.event_local_closed = trio.Event()
self.event_remote_closed = trio.Event()
self.event_reset = trio.Event()
self.close_lock = trio.Lock()
2019-11-29 19:09:56 +08:00
self.incoming_data_channel = incoming_data_channel
self._buf = bytearray()
2018-11-12 07:03:04 +08:00
2019-08-28 21:43:34 +08:00
@property
def is_initiator(self) -> bool:
return self.stream_id.is_initiator
2019-11-29 19:09:56 +08:00
async def _read_until_eof(self) -> bytes:
async for data in self.incoming_data_channel:
self._buf.extend(data)
payload = self._buf
self._buf = self._buf[len(payload) :]
return bytes(payload)
def _read_return_when_blocked(self) -> bytes:
buf = bytearray()
while True:
try:
data = self.incoming_data_channel.receive_nowait()
buf.extend(data)
except (trio.WouldBlock, trio.EndOfChannel):
break
return buf
async def read(self, n: int = None) -> bytes:
"""
Read up to n bytes. Read possibly returns fewer than `n` bytes, if
there are not enough bytes in the Mplex buffer. If `n is None`, read
until EOF.
2019-08-07 15:23:20 +08:00
:param n: number of bytes to read
:return: bytes actually read
2018-11-12 07:03:04 +08:00
"""
if n is not None and n < 0:
2019-08-16 10:21:51 +08:00
raise ValueError(
f"the number of bytes to read `n` must be non-negative or "
"`None` to indicate read until EOF"
2019-08-16 10:21:51 +08:00
)
2019-09-06 17:26:40 +08:00
if self.event_reset.is_set():
raise MplexStreamReset
if n is None:
2019-11-29 19:09:56 +08:00
return await self._read_until_eof()
if len(self._buf) == 0:
data: bytes
# Peek whether there is data available. If yes, we just read until there is no data,
# and then return.
try:
data = self.incoming_data_channel.receive_nowait()
2019-12-01 19:17:32 +08:00
self._buf.extend(data)
2019-11-29 19:09:56 +08:00
except trio.EndOfChannel:
raise MplexStreamEOF
except trio.WouldBlock:
# We know `receive` will be blocked here. Wait for data here with `receive` and
# catch all kinds of errors here.
try:
data = await self.incoming_data_channel.receive()
2019-12-01 19:17:32 +08:00
self._buf.extend(data)
2019-11-29 19:09:56 +08:00
except trio.EndOfChannel:
if self.event_reset.is_set():
raise MplexStreamReset
if self.event_remote_closed.is_set():
raise MplexStreamEOF
except trio.ClosedResourceError as error:
# Probably `incoming_data_channel` is closed in `reset` when we are waiting
# for `receive`.
if self.event_reset.is_set():
raise MplexStreamReset
raise Exception(
"`incoming_data_channel` is closed but stream is not reset. "
"This should never happen."
) from error
self._buf.extend(self._read_return_when_blocked())
payload = self._buf[:n]
self._buf = self._buf[len(payload) :]
return bytes(payload)
2018-11-01 05:31:00 +08:00
2019-08-02 18:28:04 +08:00
async def write(self, data: bytes) -> int:
"""
write to stream.
2018-11-12 07:03:04 +08:00
:return: number of bytes written
"""
2019-09-09 15:45:35 +08:00
if self.event_local_closed.is_set():
2019-10-24 14:53:19 +08:00
raise MplexStreamClosed(f"cannot write to closed stream: data={data!r}")
flag = (
HeaderTags.MessageInitiator
2019-08-28 21:43:34 +08:00
if self.is_initiator
else HeaderTags.MessageReceiver
)
return await self.muxed_conn.send_message(flag, data, self.stream_id)
2018-11-01 05:31:00 +08:00
async def close(self) -> None:
"""Closing a stream closes it for writing and closes the remote end for
reading but allows writing in the other direction."""
2018-11-21 13:41:13 +08:00
# TODO error handling with timeout
async with self.close_lock:
if self.event_local_closed.is_set():
return
2019-08-28 21:43:34 +08:00
flag = (
HeaderTags.CloseInitiator if self.is_initiator else HeaderTags.CloseReceiver
)
# TODO: Raise when `muxed_conn.send_message` fails and `Mplex` isn't shutdown.
await self.muxed_conn.send_message(flag, None, self.stream_id)
2018-11-12 07:03:04 +08:00
_is_remote_closed: bool
async with self.close_lock:
self.event_local_closed.set()
_is_remote_closed = self.event_remote_closed.is_set()
2018-11-12 07:03:04 +08:00
if _is_remote_closed:
# Both sides are closed, we can safely remove the buffer from the dict.
async with self.muxed_conn.streams_lock:
self.muxed_conn.streams.pop(self.stream_id, None)
2018-11-12 07:03:04 +08:00
async def reset(self) -> None:
"""closes both ends of the stream tells this remote side to hang up."""
async with self.close_lock:
# Both sides have been closed. No need to event_reset.
if self.event_remote_closed.is_set() and self.event_local_closed.is_set():
return
if self.event_reset.is_set():
return
self.event_reset.set()
2018-11-21 10:46:18 +08:00
if not self.event_remote_closed.is_set():
flag = (
HeaderTags.ResetInitiator
2019-08-28 21:43:34 +08:00
if self.is_initiator
else HeaderTags.ResetReceiver
)
# Try to send reset message to the other side. Ignore if there is anything wrong.
try:
await self.muxed_conn.send_message(flag, None, self.stream_id)
except MuxedConnUnavailable:
pass
2018-11-21 10:46:18 +08:00
self.event_local_closed.set()
self.event_remote_closed.set()
2018-11-21 10:46:18 +08:00
2019-11-29 19:09:56 +08:00
await self.incoming_data_channel.aclose()
async with self.muxed_conn.streams_lock:
if self.muxed_conn.streams is not None:
self.muxed_conn.streams.pop(self.stream_id, None)
2018-11-01 05:31:00 +08:00
2018-11-12 07:03:04 +08:00
# TODO deadline not in use
2019-08-05 17:02:18 +08:00
def set_deadline(self, ttl: int) -> bool:
"""
set deadline for muxed stream.
2018-11-12 07:03:04 +08:00
:return: True if successful
2018-11-01 05:31:00 +08:00
"""
2018-11-12 07:03:04 +08:00
self.read_deadline = ttl
self.write_deadline = ttl
return True
2019-08-05 17:02:18 +08:00
def set_read_deadline(self, ttl: int) -> bool:
"""
set read deadline for muxed stream.
2018-11-12 07:03:04 +08:00
:return: True if successful
"""
self.read_deadline = ttl
return True
2019-08-05 17:02:18 +08:00
def set_write_deadline(self, ttl: int) -> bool:
"""
set write deadline for muxed stream.
2018-11-12 07:03:04 +08:00
:return: True if successful
"""
self.write_deadline = ttl
return True