seanghaeli commented on code in PR #76423:
URL: https://github.com/apache/airflow/pull/76423#discussion_r4485564879
##########
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:
My agent flagged this one: if a handler raises anything other than the three
errors caught here, for example an `OSError` writing a trigger's log file on a
full disk, `_send_frame` exits with only part of the frame written.
`handle_requests` then catches it and sends an error frame right after the
partial one. The runner reads the length of the first frame and waits for bytes
that never come, so it never gets another reply. We reproduced it with a small
test: a 4 MiB reply, a log handler that raises on its third call, and the
runner got 219 KB of a 4,194,345 byte frame before the decode failed. On main
the same error would have stopped the triggerer instead. Could this finish
sending the frame before letting the error out, or catch it, log it, and stop
reading that socket for the rest of the send?
##########
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:
Small thing: this only registers a log socket, so the test would still pass
if stdout and stderr weren't drained during a send. A trigger that prints a lot
blocks on stdout the same way. Could it be parametrized over both?
--
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]
