Limit TCP writes with Tornado to 2GB
Fix OverflowError resulting from trying to send a ridiculously large numpy array to a dask cluster.
@@ -235,7 +235,7 @@ class TCP(Comm):
range(0, frames_nbytes + OPENSSL_MAX_CHUNKSIZE, OPENSSL_MAX_CHUNKSIZE),
):
chunk = frames[i:j]
- chunk_nbytes = len(chunk)
+ chunk_nbytes = chunk.nbytes
n = await stream.read_into(chunk)
assert n == chunk_nbytes, (n, chunk_nbytes)
except StreamClosedError as e:
@@ -303,16 +303,25 @@ class TCP(Comm):
# trick to enque all frames for writing beforehand
for each_frame_nbytes, each_frame in zip(frames_nbytes, frames):
if each_frame_nbytes:
- if stream._write_buffer is None:
- raise StreamClosedError()
-
- if isinstance(each_frame, memoryview):
# Make sure that `len(data) == data.nbytes`
# See <https://github.com/tornadoweb/tornado/pull/2996>
each_frame = ensure_memoryview(each_frame)
+ for i, j in sliding_window(
+ 2,
+ range(
+ 0,
+ each_frame_nbytes + OPENSSL_MAX_CHUNKSIZE,
+ OPENSSL_MAX_CHUNKSIZE,
+ ),
+ ):
+ chunk = each_frame[i:j]
+ chunk_nbytes = chunk.nbytes
+
+ if stream._write_buffer is None:
+ raise StreamClosedError()
- stream._write_buffer.append(each_frame)
- stream._total_write_index += each_frame_nbytes
+ stream._write_buffer.append(chunk)
+ stream._total_write_index += chunk_nbytes
# start writing frames
stream.write(b"")
