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

Handle exceptions inside read_message

Browse Source 
And remove the need of checking `None` for every read messages.
This commit is contained in:
mhchia 2019-09-14 14:16:40 +08:00
parent f62f07bb9f
commit b51c2939a8
No known key found for this signature in database
GPG Key ID: 389EFBEA1362589A
1 changed files with 88 additions and 90 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -178,21 +178,15 @@ class Mplex(IMuxedConn):
""" """
Read a message off of the secured connection and add it to the corresponding message buffer Read a message off of the secured connection and add it to the corresponding message buffer
""" """
# TODO Deal with other types of messages using flag (currently _)
while True: while True:
try: try:
channel_id, flag, message = await self._wait_until_shutting_down_or_closed( channel_id, flag, message = await self._wait_until_shutting_down_or_closed(
self.read_message() self.read_message()
) )
except ( except MplexUnavailable as error:
MplexUnavailable,
ConnectionResetError,
IncompleteReadError,
) as error:
print(f"!@# handle_incoming: read_message: exception={error}") print(f"!@# handle_incoming: read_message: exception={error}")
break break
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)) stream_id = StreamID(channel_id=channel_id, is_initiator=bool(flag & 1))
is_stream_id_seen: bool is_stream_id_seen: bool
stream: MplexStream stream: MplexStream
@ -200,8 +194,6 @@ class Mplex(IMuxedConn):
is_stream_id_seen = stream_id in self.streams is_stream_id_seen = stream_id in self.streams
if is_stream_id_seen: if is_stream_id_seen:
stream = self.streams[stream_id] stream = self.streams[stream_id]
# Other consequent stream message should wait until the stream get accepted
# TODO: Handle more tags, and refactor `HeaderTags`
if flag == HeaderTags.NewStream.value: if flag == HeaderTags.NewStream.value:
if is_stream_id_seen: if is_stream_id_seen:
# `NewStream` for the same id is received twice... # `NewStream` for the same id is received twice...
@ -279,6 +271,8 @@ class Mplex(IMuxedConn):
# Force context switch # Force context switch
await asyncio.sleep(0) await asyncio.sleep(0)
# If we enter here, it means this connection is shutting down.
# We should clean the things up.
await self._cleanup() await self._cleanup()
async def read_message(self) -> Tuple[int, int, bytes]: async def read_message(self) -> Tuple[int, int, bytes]:
@ -290,15 +284,19 @@ class Mplex(IMuxedConn):
# FIXME: No timeout is used in Go implementation. # FIXME: No timeout is used in Go implementation.
# Timeout is set to a relatively small value to alleviate wait time to exit # Timeout is set to a relatively small value to alleviate wait time to exit
# loop in handle_incoming # loop in handle_incoming
header = await decode_uvarint_from_stream(self.secured_conn)
# TODO: Handle the case of EOF and other exceptions?
try: try:
header = await decode_uvarint_from_stream(self.secured_conn)
message = await asyncio.wait_for( message = await asyncio.wait_for(
read_varint_prefixed_bytes(self.secured_conn), timeout=5 read_varint_prefixed_bytes(self.secured_conn), timeout=5
) )
except asyncio.TimeoutError: except (ConnectionResetError, IncompleteReadError) as error:
# TODO: Investigate what we should do if time is out. raise MplexUnavailable(
return None, None, None "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
flag = header & 0x07 flag = header & 0x07
channel_id = header >> 3 channel_id = header >> 3