buhuipao commented on code in PR #76423:
URL: https://github.com/apache/airflow/pull/76423#discussion_r4533641819
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -835,6 +837,55 @@ def _handle_request(self, msg: ToTriggerSupervisor, log:
FilteringBoundLogger, r
self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
+ def _send_frame(self, data: bytes | bytearray) -> None:
+ """
+ Send ``data`` to the runner, reading its output sockets while the send
is blocked.
+
+ The runner's event loop can be stuck writing to a full log socket, so
it does not read this
+ response until we read its logs. A plain ``sendall()`` would leave
both processes waiting on
+ each other. EOF on an output socket is left for the main select loop
to handle.
+ """
+ view = memoryview(data)
+
+ def send_available() -> None:
+ nonlocal view
+ with suppress(BlockingIOError):
+ while view:
+ view = view[self.stdin.send(view) :]
+
+ # Non-blocking mode rather than MSG_DONTWAIT: macOS ignores that flag
for large AF_UNIX sends.
+ timeout = self.stdin.gettimeout()
+ self.stdin.setblocking(False)
+ try:
+ send_available()
+ if not view:
+ return
+ with selectors.DefaultSelector() as sel:
+ sel.register(self.stdin, selectors.EVENT_WRITE)
+ for key in list(self.selector.get_map().values()):
+ sock: socket = key.fileobj # type: ignore[assignment]
+ if sock is not self.stdin:
+ # The caller's select results may still list this
socket as readable after
+ # we drain it here; non-blocking makes that stale read
return instead of hang.
+ sock.setblocking(False)
+ sel.register(sock, selectors.EVENT_READ, key.data)
+ while view:
+ for key, _ in sel.select():
+ if key.fileobj is self.stdin:
+ send_available()
+ continue
+ socket_handler, _ = key.data
+ try:
+ need_more = socket_handler(key.fileobj)
+ except BlockingIOError:
+ need_more = True
+ except (BrokenPipeError, ConnectionResetError):
+ need_more = False
+ if not need_more:
Review Comment:
You are right, thanks. I reproduced it with a test like yours: a 4 MiB reply
and a log reader that raises on its third line.
I took your first option:
- The reply is now sent in full even when a handler fails.
- For the rest of that send, the supervisor keeps emptying the failed socket
and drops the data. If it stopped reading, a runner blocked on that socket
would never read the reply.
- After the reply, `_service_subprocess` raises the error once. So the
triggerer stops, as it does on main.
The new test is `test_send_msg_completes_response_when_output_handler_fails`.
---
Drafted-by: Claude Code (Fable 5.2); reviewed by @buhuipao before posting
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -659,6 +660,104 @@ def fake_super_start(cls, **kwargs):
fake_proc.send_msg.assert_called_once()
+def test_run_once_drains_all_waiting_runner_logs(jobless_supervisor, mocker):
+ from airflow.sdk.execution_time.supervisor import
make_buffered_socket_reader
+
+ for method in (
+ "load_triggers",
+ "handle_events",
+ "handle_recovered_triggers",
+ "handle_failed_triggers",
+ "clean_unused",
+ "heartbeat",
+ "emit_metrics",
+ "_check_subprocess_exit",
+ ):
+ mocker.patch.object(TriggerRunnerSupervisor, method, autospec=True,
return_value=None)
+
+ lines: list[bytes] = []
+
+ def collect():
+ while True:
+ lines.append((yield))
+
+ log_parent, log_child = socketpair()
+ selector = selectors.DefaultSelector()
+ jobless_supervisor.selector = selector
+ selector.register(
+ log_parent,
+ selectors.EVENT_READ,
+ make_buffered_socket_reader(collect(), on_close=selector.unregister),
+ )
+ try:
+ # More than the 4 KiB one service call reads, but small enough to fit
the 8 KiB macOS buffer.
+ log_child.sendall(b'{"event": "waiting"}\n' * 300)
+
+ jobless_supervisor.run_once()
+
+ assert len(lines) == 300
+ finally:
+ selector.close()
+ log_parent.close()
+ log_child.close()
+
+
+def
test_send_msg_drains_logs_while_runner_is_blocked_on_log_write(jobless_supervisor,
mocker):
+ mocker.patch.object(TriggerRunnerSupervisor, "_check_subprocess_exit",
autospec=True, return_value=None)
+ request_parent, request_child = socketpair()
+ log_parent, log_child = socketpair()
+ jobless_supervisor.stdin = request_parent
+ selector = selectors.DefaultSelector()
+ jobless_supervisor.selector = selector
+ selector.register(request_parent, selectors.EVENT_READ, (lambda sock:
True, selector.unregister))
+ selector.register(
+ log_parent, selectors.EVENT_READ, (lambda sock:
bool(sock.recv(65536)), selector.unregister)
Review Comment:
Done. The test now runs for `logs`, `stdout` and `stderr`, and uses the real
readers from `_register_pipe_readers`.
If `_send_frame` drains only the log socket, the `stdout` and `stderr` cases
fail.
---
Drafted-by: Claude Code (Fable 5.2); reviewed by @buhuipao before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
