Limit TCP writes with Tornado to 2GB (#6557)

Fix OverflowError resulting from trying to send a ridiculously large numpy array to a dask cluster.
dask · Jun 24, 2022 · f129485 · f129485
1 parent becb366
commit f129485
Showing 1 changed file with 20 additions and 11 deletions.
diff --git a/distributed/comm/tcp.py b/distributed/comm/tcp.py
@@ -235,7 +235,7 @@ async def read(self, deserializers=None):
                 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 @@ async def write(self, msg, serializers=None, on_error="message"):
             # 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)
-
-                    stream._write_buffer.append(each_frame)
-                    stream._total_write_index += each_frame_nbytes
+                    # 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(chunk)
+                        stream._total_write_index += chunk_nbytes
 
             # start writing frames
             stream.write(b"")