def frame_split_size(frames, n=BIG_BYTES_SHARD_SIZE):
"""
Split a list of frames into a list of frames of maximum size
This helps us to avoid passing around very large bytestrings.
Examples
--------
>>> frame_split_size([b'12345', b'678'], n=3) # doctest: +SKIP
[b'123', b'45', b'678']
"""
if not frames:
return frames
if max(map(nbytes, frames)) <= n:
return frames
out = []
for frame in frames:
if nbytes(frame) > n:
if isinstance(frame, (bytes, bytearray)):
frame = memoryview(frame)
try:
itemsize = frame.itemsize
except AttributeError:
itemsize = 1
for i in range(0, nbytes(frame) // itemsize, n // itemsize):
out.append(frame[i:i + n // itemsize])
else:
out.append(frame)
return out
def write(self, msg, serializers=None, on_error="message"):
"""
Write a message (a Python object).
This method is a coroutine.
Parameters
----------
msg :
on_error : Optional[str]
The behavior when serialization fails. See
``distributed.protocol.core.dumps`` for valid values.
"""
stream = self.stream
bytes_since_last_yield = 0
if stream is None:
raise CommClosedError
frames = yield to_frames(
msg,
serializers=serializers,
on_error=on_error,
context={"sender": self._local_addr, "recipient": self._peer_addr},
)
try:
lengths = [nbytes(frame) for frame in frames]
#print("Size of entire payload: ", sum(lengths))
length_bytes = [struct.pack("Q", len(frames))] + [
struct.pack("Q", x) for x in lengths
]
# sys.version_info[0] == 3 means Python3
if (sys.version_info[0] == 3) and sum(lengths) < 2 ** 17: # 128kiB
b = b"".join(length_bytes + frames) # small enough, send in one go
stream.write(b)
else:
stream.write(b"".join(length_bytes)) # avoid large memcpy, send in many
for frame in frames:
# Can't wait for the write() Future as it may be lost
# ("If write is called again before that Future has resolved,
# the previous future will be orphaned and will never resolve")
if not self._iostream_allows_memoryview:
frame = ensure_bytes(frame)
future = stream.write(frame)
bytes_since_last_yield += nbytes(frame)
if bytes_since_last_yield > 32e6:
yield future
bytes_since_last_yield = 0
except StreamClosedError as e:
stream = None
convert_stream_closed_error("Lambda", e)
except TypeError as e:
if stream._write_buffer is None:
print("tried to write message {} on closed stream".format(msg))
else:
raise
raise gen.Return(sum(map(nbytes, frames)))
def merge_frames(header, frames):
""" Merge frames into original lengths
Examples
--------
>>> merge_frames({'lengths': [3, 3]}, [b'123456'])
[b'123', b'456']
>>> merge_frames({'lengths': [6]}, [b'123', b'456'])
[b'123456']
"""
lengths = list(header["lengths"])
if not frames:
return frames
assert sum(lengths) == sum(map(nbytes, frames))
if all(len(f) == l for f, l in zip(frames, lengths)):
return frames
frames = frames[::-1]
lengths = lengths[::-1]
out = []
while lengths:
l = lengths.pop()
L = []
while l:
frame = frames.pop()
if nbytes(frame) <= l:
L.append(frame)
l -= nbytes(frame)
else:
mv = memoryview(frame)
L.append(mv[:l])
frames.append(mv[l:])
l = 0
out.append(b"".join(map(ensure_bytes, L)))
return out
def pack_frames_prelude(frames):
lengths = [len(f) for f in frames]
lengths = [struct.pack("Q", len(frames))
] + [struct.pack("Q", nbytes(frame)) for frame in frames]
return b"".join(lengths)