dabla commented on code in PR #75396:
URL: https://github.com/apache/airflow/pull/75396#discussion_r4554823302
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -178,88 +183,144 @@ def get_order(self, entity: AliasedClass, *, descending:
bool = False) -> tuple[
return tuple(column.desc() if descending else column.asc() for column
in columns)
+def _render_json(content: JsonValue) -> JSONResponse:
+ """
+ Render an XCom read's response on a worker thread.
+
+ The values come from a JSON column, so returning a ``JSONResponse`` skips
FastAPI's re-validation of
+ them as ``JsonValue``, which for large values costs several times more
than encoding them. Cadwyn still
+ decodes and re-encodes a returned ``JSONResponse`` body on the event loop
+ (https://github.com/zmievsa/cadwyn/issues/51), so only the first encoding
runs on the worker thread.
+ """
+ return JSONResponse(content)
+
+
+def _reject_archived_caller(
+ token: ExecutionToken, *, dag_id: str, run_id: str, task_id: str,
map_index: int | None, session: Session
+) -> None:
+ """Refuse a task's read of its own XCom once its attempt has been
archived."""
+ if token.claims.scope != "execution":
+ return
+ caller = session.execute(
+ select(TaskInstance.id, TaskInstance.working_set)
+ .where(
+ TaskInstance.id == token.id,
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.task_id == task_id,
+ *build_coordinate_filters(TaskInstance, map_index=map_index),
+ )
+ .limit(1)
+ .execution_options(include_all_attempts=True)
+ ).first()
+ if caller is None or caller.working_set is not None:
+ return
+ identify_archived = IdentifyArchivedTaskStateUpdates.is_applied
+ raise HTTPException(
+ status_code=status.HTTP_410_GONE if identify_archived else
status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": (
+ "Task Instance not found in the working set; its attempt has
been archived"
+ if identify_archived
+ else "Task Instance not found"
+ ),
+ },
+ )
+
+
def _build_xcom_read(
*,
dag_id: str,
run_id: str,
task_id: str,
key: str,
- session: Session,
dag_bag: DBDagBag,
token: ExecutionToken,
map_index: int | None = None,
include_prior_dates: bool = False,
previous_iteration: bool = False,
sequence: bool = False,
+ reject_archived_caller: bool = False,
) -> _XComRead:
"""
Select the XCom rows of the producers visible to the calling task instance.
A ``sequence`` read returns several values in order. When it targets a
loop task from outside that
loop, it covers every iteration of the task rather than the one a caller
inside the loop would see.
+
+ Routes run this on a worker thread: resolving producers can deserialize a
Dag on a DagBag cache miss,
+ CPU work that would stall the event loop. A session cannot cross threads,
so this opens its own.
+
+ ``reject_archived_caller`` makes a caller reading its own XCom get the
archived-attempt error first. The
+ check runs here rather than on the route's session, so it leaves no route
connection idle during the lookup.
"""
- resolver = TaskCoordinateResolver(dag_bag, session)
- read = partial(
- XComModel.get_many,
- dag_ids=dag_id,
- run_id=run_id,
- task_ids=task_id,
- key=key,
- include_prior_dates=include_prior_dates,
- )
- try:
- if previous_iteration and include_prior_dates:
- raise ValueError("previous_iteration cannot be combined with prior
dates")
- if not previous_iteration and not resolver.has_regions(
- dag_id, None if include_prior_dates else run_id, task_id
- ):
- return _XComRead(read(region_id=SENTINEL_REGION_ID,
map_indexes=map_index))
- caller = session.get(TaskInstance, token.id) if token.claims.scope ==
"execution" else None
- if include_prior_dates:
- task = None
- if session.scalar(select(DagRun.id).where(DagRun.dag_id == dag_id,
DagRun.run_id == run_id)):
- with suppress(TaskNotFound):
- task = resolver.get_task(
- dag_id,
- run_id,
- task_id,
- dag_version_id=caller.dag_version_id
- if caller and (caller.dag_id, caller.run_id) ==
(dag_id, run_id)
- else None,
- )
- if task is not None and enclosing_loop(task) is not None:
- raise ValueError("Prior-date lookups are not supported for
tasks inside a loop")
- return _XComRead(
- read(region_id=SENTINEL_REGION_ID, include_node_regions=True,
map_indexes=map_index)
+ with create_session(scoped=False) as session:
Review Comment:
**[warning]** The common read never touches the DagBag, but it still pays
the worker thread and the second connection.
For a task that is not in a loop or a mapped region (and without
`previous_iteration`), `_build_xcom_read` only does this before it returns:
- `_reject_archived_caller`: one `SELECT` on `task_instance`;
- `resolver.has_regions(...)`: one `SELECT EXISTS` on `task_instance`
(`task_coordinates.py:571-582`), no Dag lookup;
- `XComModel.get_many(...)`: builds a statement, no query.
So on that path nothing can deserialize a Dag, which is the reason given for
the worker thread. The request still does a thread hop, a sync checkout with
its pre-ping, two of its three queries on the worker thread, and then an async
checkout with a second pre-ping for the one `SELECT` that is awaited. That fits
the small-value rows of the benchmark table (own small value on PostgreSQL: 6.2
ms on `main`, 8.7 and 7.3 ms here), and it is the case most XCom reads are. The
goal in #67799 is that database waits leave the threadpool; here most of them
stay in it, and the "two connections per request" note in the description
applies to every read.
Proposal: run the archived-caller check and the `has_regions` query on the
route's `AsyncSession` and go to `run_in_threadpool` only when the task has
regions or `previous_iteration` is set. The common read would then use one
connection and no worker thread, and the bridge stays for the path that needs
it. The same holds for `_find_writer_id` (line 892): a task writing its own
XCom returns after `session.get` and `get_public_map_index`, which runs no
query when the task is not in a region, before the resolver is built.
The cost is an async form of those two queries next to the sync ones the
resolver keeps. I have not measured this; if you prefer to keep this PR a plain
conversion, a follow-up is fine with me, but then I would say in the
description that the common path is the one that pays for the bridge.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -129,8 +131,11 @@ def has_xcom_access(
_team_name_for_ti_stmt,
)
- requester_team = session.scalar(_team_name_for_ti_stmt(token.id))
- target_team = session.scalar(_team_name_for_dag_stmt(dag_id))
+ # Cadwyn solves this router dependency in two passes, and the first pass's
request session would stay
+ # checked out until the response; a session of its own returns the
connection before the producer lookup.
+ async with create_session_async() as session:
Review Comment:
**[nit]** With `[core] multi_team` on, this check still runs in both of
Cadwyn's dependency passes.
The comment above says Cadwyn solves this router dependency twice. Now that
the check opens its own session, each request checks out two async connections
and runs the two team queries twice (your reply on the second thread counts 3
async checkouts per multi-team read). `_resolve_trigger_token` below already
skips the first pass with `Depends(current_dependency_solver)` and
`dependency_solver != "cadwyn"`. The same guard here would bring a multi-team
request back to one access check, and the 403 would still be raised in the pass
the endpoint runs with.
--
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]
