Dev-iL commented on code in PR #75396:
URL: https://github.com/apache/airflow/pull/75396#discussion_r4503664412
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -563,31 +581,34 @@ def head_xcom(
]
),
description="Get a single XCom Value",
+ response_model=XComResponse,
)
-def get_xcom(
+async def get_xcom(
dag_id: str,
run_id: str,
task_id: str,
key: Annotated[str, Path(min_length=1)],
- session: SessionDep,
+ session: AsyncSessionDep,
params: Annotated[GetXcomFilterParams, Query()],
dag_bag: DagBagDep,
token: ExecutionToken = _TriggerAwareToken,
-) -> XComResponse:
+) -> JSONResponse:
"""Get an Airflow XCom from database - not other XCom Backends."""
owner = None
if token.claims.scope == "execution":
- owner = 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=params.map_index),
+ owner = (
+ await 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=params.map_index),
+ )
+ .limit(1)
+ .execution_options(include_all_attempts=True)
)
- .limit(1)
- .execution_options(include_all_attempts=True)
).first()
Review Comment:
Good catch. I measured the alternatives before choosing, in Breeze with pool
checkout and check-in listeners on both engines: 11 sequential requests per
scenario, the first dropped, two runs each (run 1; run 2).
**The owner check.** Both fixes remove the overlap, so what decides between
them is what each costs:
| Owner check | Checkouts per request (async + sync) | Peak connections |
Both pools held, PostgreSQL, Dag cache cleared (median ms) | Request time,
SQLite, Dag cache cleared (median ms) |
| --- | --- | --- | --- | --- |
| On the route session (before) | 1 + 1 | 2 | 155; 151 | 99.8; 97.3 |
| On the route session, then `rollback()` | 2 + 1 | 1 | 0; 0 | 172; 153 |
| In a worker-thread session before the lookup | 1 + 2 | 1 | 0; 0 | 92.9;
98.6 |
The rollback releases the connection, but the route then checks out a second
one, with the pre-ping on PostgreSQL, and cold reads on SQLite got slower.
Moving the check to a worker thread had no such cost. This PR goes one step
further than the measured version and runs the check first inside
`_build_xcom_read`'s own session, so it keeps the same point in the error order
without a second sync checkout: 1 + 1 per request, with no overlap, in the
final benchmarks.
**The two dependency checks.** Measuring this turned up the same pattern
twice more. Cadwyn solves `has_xcom_access` in both of its dependency passes,
and with `[core] multi_team` on, the first pass's session stayed checked out
until the response. `_resolve_trigger_token`, which checks a trigger's claimed
identity, queried the route's session before the lookup. Both now query in a
short `create_session_async()` session of their own:
| Request | Peak connections, before → after | Both pools held, before →
after (median ms, PostgreSQL; SQLite) | Async checkouts per request, before →
after |
| --- | --- | --- | --- |
| Read with `[core] multi_team` on | 3 → 1 | 14.3 and 6.3; 9.9 and 12.0 → 0
| 2 → 3 |
| Trigger read (`GET` with `Airflow-Trigger-Ti-Id`) | 2 → 1 | 15.6 and 14.4;
11.5 and 12.8 → 0 | 1 → 2 |
| Trigger write (`POST` with `Airflow-Trigger-Ti-Id`) | 2 → 1 | 4.3 and 3.9;
1.1 and 1.5 → 0 | 1 → 2 |
The trigger requests go through `InProcessExecutionAPI` with the header, as
the triggerer sends them. The extra checkouts are short and taken before the
lookup. Event-loop stall and request time didn't differ measurably in any of
these cases: the medians moved by less than their run-to-run spread. For a
trigger's write, the deferred-state check and the write now run in separate
transactions. The check never locked anything (it's a plain `SELECT`), so the
task instance's state could already change between them before.
Two tests pin this:
`test_xcom_read_holds_no_route_connection_during_the_producer_lookup` (without
and with multi-team) and
`test_trigger_xcom_request_holds_no_route_connection_during_the_producer_lookup`
(`GET` and `POST`) assert that no async connection is checked out while the
lookup runs. Each case fails without its change.
---
Drafted-by: Claude Code (Opus 5.5) (no human review 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]
