Andrushika opened a new pull request, #70128:
URL: https://github.com/apache/airflow/pull/70128
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
https://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.
For user-facing UI changes, please attach before/after screenshots (or a
short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
Fix deadlock when concurrent tasks queue the same asset's downstream dags
---
### Background
When a task produces an asset, the downstream Dags scheduled on that asset
need to run. Airflow records this in the `asset_dag_run_queue` (ADRQ) table,
one row per `(asset_id, target_dag_id)`. So when a task succeeds and updates
asset `S`, `register_asset_change` inserts one ADRQ row for every Dag scheduled
on `S`. Those Dags are collected into a Python `set` (`dags_to_queue`), and the
rows are inserted by looping over that set.
### Why
A Python `set` has no stable order (The essence of `set` is a hash table).
Here, the elements are `DagModel` objects, freshly loaded at different memory
addresses in each process, so two processes can loop over the same Dags in
opposite orders.
When two tasks finish at the same time and both emit to the same asset, they
insert the same ADRQ rows, but maybe in opposite orders. Each insert holds a
row lock until commit. So transaction A can hold row `(S, dag_a)` and wait for
`(S, dag_b)`, while transaction B holds `(S, dag_b)` and waits for `(S,
dag_a)`. **That is a deadlock.** Postgres aborts one side and it retries, which
adds latency and error noise under high asset fan-out.
### What
Insert the rows in a fixed order, sorted by `dag_id`, so every process takes
the row locks in the same order and the cycle cannot form. A small shared
helper `_sorted_by_dag_id` is used in all three insert paths (postgres `ON
CONFLICT`, mysql `ON DUPLICATE KEY`, and the per-row `SAVEPOINT` fallback),
because all three loop over the same set.
I was running the concurrent test for #70078 and found that this problem
exists on `main`.
Then I reproduced it on Postgres with a script that has 8 concurrent
transactions insert one asset's downstream rows in opposite orders (what two
processes' set iteration can produce). Over 20 rounds it hit 658 deadlocks with
the old order and 0 after sorting. I only benchmarked Postgres. MySQL uses
InnoDB row locks with the same inversion, so sorting fixes it there too. The
SQLite path cannot deadlock this way and is sorted only for consistency.
<details>
<summary>The script for reproduce:</summary>
```python
from __future__ import annotations
import threading
import time
from sqlalchemy import delete
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.exc import OperationalError
from airflow import settings
from airflow.models.asset import AssetActive, AssetDagRunQueue, AssetModel,
DagScheduleAssetReference
from airflow.models.dag import DagModel
from airflow.models.dagbundle import DagBundleModel
from airflow.utils.session import create_session
ASSET_ID = 1
DOWNSTREAM_DAGS = ["dag_a", "dag_b", "dag_c", "dag_d"]
N_THREADS = 8
N_ROUNDS = 20
ROW_GAP = 0.004 # small pause between rows to widen the lock window
def _seed():
with create_session() as s:
s.execute(delete(AssetDagRunQueue))
s.execute(delete(DagScheduleAssetReference))
s.query(AssetActive).delete()
s.query(AssetModel).delete()
s.query(DagModel).filter(DagModel.dag_id.in_(DOWNSTREAM_DAGS)).delete(synchronize_session=False)
s.merge(DagBundleModel(name="repro"))
s.flush() # bundle must exist before dags reference it (FK)
asset = AssetModel(id=ASSET_ID, name="S", uri="s3://bucket/S",
group="asset", extra={})
s.add_all([asset, AssetActive.for_asset(asset)])
for d in DOWNSTREAM_DAGS:
s.add(DagModel(dag_id=d, bundle_name="repro", is_stale=False,
fileloc=f"{d}.py"))
s.flush()
asset.scheduled_dags = [DagScheduleAssetReference(dag_id=d) for d in
DOWNSTREAM_DAGS]
def _insert_rows(order, sort_fix, counters, barrier):
dag_ids = sorted(order) if sort_fix else order # the fix: always sort
first
barrier.wait()
while True:
try:
with create_session() as session:
for dag_id in dag_ids:
stmt = (
insert(AssetDagRunQueue)
.values(asset_id=ASSET_ID, target_dag_id=dag_id)
.on_conflict_do_nothing()
)
session.execute(stmt)
time.sleep(ROW_GAP)
return
except OperationalError as e:
if "deadlock detected" in str(e).lower():
counters["deadlocks"] += 1
continue # retry: the losing side re-runs and eventually
wins
raise
def _run(sort_fix):
counters = {"deadlocks": 0}
forward = DOWNSTREAM_DAGS
reverse = list(reversed(DOWNSTREAM_DAGS))
for _ in range(N_ROUNDS):
with create_session() as s: # fresh rows each round so ON CONFLICT
actually inserts
s.execute(delete(AssetDagRunQueue))
barrier = threading.Barrier(N_THREADS)
threads = [
threading.Thread(
target=_insert_rows,
args=(forward if i % 2 == 0 else reverse, sort_fix,
counters, barrier),
)
for i in range(N_THREADS)
]
for t in threads:
t.start()
for t in threads:
t.join()
return counters["deadlocks"]
def main():
if settings.engine.dialect.name != "postgresql":
raise SystemExit("Run under --backend postgres (row locks are needed
to reproduce).")
_seed()
unsorted_deadlocks = _run(sort_fix=False)
sorted_deadlocks = _run(sort_fix=True)
print(
f"\n=== ADRQ insert-order deadlock — {N_THREADS} threads x
{N_ROUNDS} rounds, "
f"{len(DOWNSTREAM_DAGS)} downstream dags ===\n"
f" UNSORTED (insert in set order, forward vs reversed):
{unsorted_deadlocks} deadlocks\n"
f" SORTED (sort by dag_id first, the fix) :
{sorted_deadlocks} deadlocks\n"
)
if __name__ == "__main__":
main()
```
</details>
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [X] Yes (please specify the tool below)
Generated-by: Claude Code Opus 4.8 following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
---
* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)**
for more information. Note: commit author/co-author name and email in commits
become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal
([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
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]
