uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397328232
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -480,6 +544,222 @@ def _handle_set_task_state_store(
)
return super()._handle_set_task_state_store(msg, log, req_id)
+ def _monitor_subprocess(self):
+ try:
+ super()._monitor_subprocess()
+ finally:
+ if self._xcom_backend is not None:
+ self._xcom_backend.stop()
+
+ def _start_xcom_backend_call(
+ self, call: XComBackendCall, finish: Callable[[XComBackendResult],
None]
+ ) -> None:
+ """Send ``call`` to the XCom backend process, which starts on the
first call it handles."""
+ if self._xcom_backend is None:
+ self._xcom_backend = XComBackendProcess.start(
+ id=self.id,
+ client=self.client,
+ logger=self.process_log,
+ selector=self.selector,
+ on_result=self._finish_xcom_backend_call,
+ )
+ self._pending_xcom_calls.append(finish)
+ self._xcom_backend.send_call(call)
+
+ def _finish_xcom_backend_call(self, result: XComBackendResult) -> None:
+ """Finish the oldest call the XCom backend process has not answered
yet."""
+ if not self._pending_xcom_calls:
+ return
+ self._pending_xcom_calls.popleft()(result)
+
+ def _call_xcom_backend(
+ self,
+ call: XComBackendCall,
+ then: Callable[[XComBackendResult], RequestResult],
+ req_id: int,
+ *,
+ on_error: Callable[[str], RequestResult] | None = None,
+ ) -> ResponseSent:
+ """
+ Answer request ``req_id`` with ``then(result)`` once the XCom backend
has run ``call``.
+
+ The request loop goes on servicing the task and the backend process
while the call runs, so it
+ keeps heartbeating and stops the task when it reaches its
``execution_timeout``.
+
+ ``on_error`` answers a failed call instead; without it, a failed call
fails the request.
+ """
+
+ def finish(result: XComBackendResult) -> None:
+ try:
+ if result.error is None:
+ resp, dump_opts = then(result)
+ elif on_error is not None:
+ resp, dump_opts = on_error(result.error)
+ else:
+ raise XComBackendError(result.error)
+ except Exception as e:
+ self.process_log.exception("Failed to handle the XCom backend
result", request_id=req_id)
+ self.send_msg(None, request_id=req_id,
error=_build_request_error(e))
+ else:
+ self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
+
+ self._start_xcom_backend_call(call, finish)
+ return ResponseSent.ALREADY_SENT
+
+ def _run_xcom_backend_call(self, call: XComBackendCall) -> None:
+ """
+ Run ``call`` and wait for it, for a caller with no request loop to
return to.
+
+ :raises XComBackendError: when the call fails or is not answered in
time.
+ """
+ answer: list[XComBackendResult] = []
+ self._start_xcom_backend_call(call, answer.append)
+ deadline = time.monotonic() + _XCOM_BACKEND_CALL_TIMEOUT
+ # The task process has not been told to start, so servicing it here
reads no request of its own.
+ while not answer and (now := time.monotonic()) < deadline:
+ self._service_subprocess(max_wait_time=deadline - now)
+ self._pending_xcom_calls.clear()
Review Comment:
This clears the callbacks on timeout, but the helper is still running the
call, and XComBackendProcess still counts it as unanswered. When the late
result arrives, `popleft()` hands it to the next call, and every later answer
is off by one.
A timed-out helper is in an unknown state anyway. Kill it and start fresh on
the next call.
--
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]
