async def read_message(self) -> Message:
buffer = bytearray()
while len(buffer) < 4:
try:
data = await self._socket.receive_some(4 - len(buffer))
except (trio.ClosedResourceError, trio.BrokenResourceError) as err:
raise RemoteDisconnected from err
if data == b"":
raise RemoteDisconnected()
buffer.extend(data)
t_size = 4 + int.from_bytes(buffer[:4], "little")
while len(buffer) < t_size:
try:
data = await self._socket.receive_some(t_size - len(buffer))
except (trio.ClosedResourceError, trio.BrokenResourceError) as err:
raise RemoteDisconnected from err
if data == b"":
raise RemoteDisconnected()
buffer.extend(data)
msg = cast(Message, pickle.loads(buffer[4:t_size]))
return msg
async def read_message(self) -> Message:
if self.reader.at_eof():
raise RemoteDisconnected()
try:
raw_size = await self.reader.readexactly(SIZE_MARKER_LENGTH)
size = int.from_bytes(raw_size, "little")
message = await self.reader.readexactly(size)
obj = cast(Message, pickle.loads(message))
return obj
except (asyncio.IncompleteReadError, BrokenPipeError,
ConnectionResetError):
raise RemoteDisconnected()
async def send_message(self, message: Msg) -> None:
pickled = pickle.dumps(message)
size = len(pickled)
try:
self.writer.write(
size.to_bytes(SIZE_MARKER_LENGTH, "little") + pickled)
async with self._drain_lock:
# Use a lock to serialize drain() calls. Circumvents this bug:
# https://bugs.python.org/issue29930
await self.writer.drain()
except (BrokenPipeError, ConnectionResetError):
raise RemoteDisconnected()
except RuntimeError as err:
# We don't do a pre-check for a closed writer since this is a performance critical
# path. We also don't want to swallow runtime errors unrelated to closed handlers.
if "handler is closed" in str(err):
self.logger.warning("Failed to send %s. Handler closed.",
message)
raise RemoteDisconnected from err
raise