This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch v3-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-3-test by this push:
new 0c8354ca9d8 [v3-3-test] Fix short-read handling in task-sdk IPC
framing (#69253) (#71609)
0c8354ca9d8 is described below
commit 0c8354ca9d8c56d1837c424fe8caca320d63fd28
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Mon Aug 17 18:49:41 2026 +0200
[v3-3-test] Fix short-read handling in task-sdk IPC framing (#69253)
(#71609)
Stream sockets are not required to return the full requested count from
recv(), but the 4-byte length prefix on the supervisor↔subprocess socket
was read with a single sock.recv(4). On a short read, the parsed length
is garbage, which either raises msgspec.DecodeError in the subprocess or
silently deadlocks the supervisor's selector reader (length_needed=0
makes the falsy payload check skip forever, so no frame is ever
dispatched to handle_requests).
Accumulate the header bytes on both sides, mirroring the payload
accumulation that was already correct. Add regression tests using a
socket that returns 2 bytes at a time to reliably reproduce the short
read.
(cherry picked from commit f72e6c9befa7494467fbc8ceb883148dd9cb61ec)
Co-authored-by: Andrew Chang <[email protected]>
---
task-sdk/src/airflow/sdk/execution_time/comms.py | 12 +++++-
.../src/airflow/sdk/execution_time/supervisor.py | 16 +++++---
.../tests/task_sdk/execution_time/test_comms.py | 39 ++++++++++++++++++
.../task_sdk/execution_time/test_supervisor.py | 48 ++++++++++++++++++++++
4 files changed, 108 insertions(+), 7 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py
b/task-sdk/src/airflow/sdk/execution_time/comms.py
index a97df1df8b0..cc314607131 100644
--- a/task-sdk/src/airflow/sdk/execution_time/comms.py
+++ b/task-sdk/src/airflow/sdk/execution_time/comms.py
@@ -331,10 +331,18 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]):
else:
len_bytes = self.socket.recv(4)
- if len_bytes == b"":
+ if not len_bytes:
raise EOFError("Request socket closed before length")
- length = int.from_bytes(len_bytes, byteorder="big")
+ # Stream sockets may return fewer bytes than requested; accumulate the
header.
+ len_buf = bytearray(len_bytes)
+ while len(len_buf) < 4:
+ chunk = self.socket.recv(4 - len(len_buf))
+ if not chunk:
+ raise EOFError(f"Request socket closed mid-length after
{len(len_buf)} of 4 bytes")
+ len_buf.extend(chunk)
+
+ length = int.from_bytes(len_buf, byteorder="big")
buffer = bytearray(length)
mv = memoryview(buffer)
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index ae2b075d8e1..db1e3488d36 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -2321,6 +2321,9 @@ def length_prefixed_frame_reader(
gen: Generator[None, _RequestFrame, None], on_close: Callable[[socket],
None]
):
length_needed: int | None = None
+ # Accumulates the 4-byte length header across selector callbacks; stream
+ # sockets may return fewer than the requested 4 bytes in a single recv.
+ header_buffer = bytearray()
# This will hold our accumulated/partial binary frame if it doesn't come
in a single read
buffer: memoryview | None = None
# position in the buffer to store next read
@@ -2331,16 +2334,19 @@ def length_prefixed_frame_reader(
next(gen)
def cb(sock: socket):
- nonlocal buffer, length_needed, pos
+ nonlocal buffer, length_needed, pos, header_buffer
if length_needed is None:
- # Read the 32bit length of the frame
- bytes = sock.recv(4)
- if bytes == b"":
+ chunk = sock.recv(4 - len(header_buffer))
+ if not chunk:
return False
+ header_buffer.extend(chunk)
+ if len(header_buffer) < 4:
+ return True
- length_needed = int.from_bytes(bytes, byteorder="big")
+ length_needed = int.from_bytes(header_buffer, byteorder="big")
buffer = memoryview(bytearray(length_needed))
+ header_buffer = bytearray()
if length_needed and buffer:
n = sock.recv_into(buffer[pos:])
if n == 0:
diff --git a/task-sdk/tests/task_sdk/execution_time/test_comms.py
b/task-sdk/tests/task_sdk/execution_time/test_comms.py
index 05dd682d890..665e5418f5b 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_comms.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_comms.py
@@ -340,3 +340,42 @@ class TestCommsDecoder:
server2.join(timeout=2)
assert result is not None
+
+ def test_read_frame_recovers_from_short_read_on_header(self):
+ msg = VariableResult(key="k", value="v", type="VariableResult")
+ payload = msgspec.msgpack.encode(_ResponseFrame(0, msg.model_dump(),
None))
+ wire = len(payload).to_bytes(4, byteorder="big") + payload
+
+ class ChunkedSocket:
+ def __init__(self, data: bytes, chunk_size: int):
+ self._data = data
+ self._chunk_size = chunk_size
+ self._pos = 0
+
+ def setblocking(self, flag):
+ pass
+
+ def recv(self, n):
+ remaining = self._data[self._pos :]
+ if not remaining:
+ return b""
+ chunk = remaining[: min(n, self._chunk_size)]
+ self._pos += len(chunk)
+ return chunk
+
+ def recv_into(self, buf):
+ remaining = self._data[self._pos :]
+ if not remaining:
+ return 0
+ take = min(len(buf), self._chunk_size, len(remaining))
+ buf[:take] = remaining[:take]
+ self._pos += take
+ return take
+
+ sock = ChunkedSocket(wire, chunk_size=2)
+ decoder = CommsDecoder(socket=sock, log=None)
+
+ result = decoder._get_response()
+ assert isinstance(result, VariableResult)
+ assert result.key == "k"
+ assert result.value == "v"
diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index f8756a293c6..f79e10e364a 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -4533,3 +4533,51 @@ class TestMakeBufferedSocketReader:
finally:
r.close()
w.close()
+
+
+class TestLengthPrefixedFrameReader:
+ def test_recovers_from_short_read_on_header(self):
+ received: list[_RequestFrame] = []
+
+ def collecting_gen():
+ while True:
+ frame = yield
+ received.append(frame)
+
+ payload = msgspec.msgpack.encode(_RequestFrame(id=42, body={"key":
"foo"}))
+ wire = len(payload).to_bytes(4, byteorder="big") + payload
+
+ class ChunkedSocket:
+ def __init__(self, data: bytes, chunk_size: int):
+ self._data = data
+ self._chunk_size = chunk_size
+ self._pos = 0
+
+ def recv(self, n):
+ remaining = self._data[self._pos :]
+ if not remaining:
+ return b""
+ chunk = remaining[: min(n, self._chunk_size)]
+ self._pos += len(chunk)
+ return chunk
+
+ def recv_into(self, buf):
+ remaining = self._data[self._pos :]
+ if not remaining:
+ return 0
+ take = min(len(buf), self._chunk_size, len(remaining))
+ buf[:take] = remaining[:take]
+ self._pos += take
+ return take
+
+ sock = ChunkedSocket(wire, chunk_size=2)
+ on_close = MagicMock()
+ cb, _ = supervisor.length_prefixed_frame_reader(collecting_gen(),
on_close=on_close)
+
+ for _ in range(len(wire) + 1):
+ if not cb(sock):
+ break
+ if received:
+ break
+
+ assert received == [_RequestFrame(id=42, body={"key": "foo"})]
