Fix deadlock when concurrent tasks queue the same asset's downstream dags

[Andrushika]

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.

The script for reproduce:

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()

• Yes (please specify the tool below)

Generated-by: Claude Code Opus 4.8 following the guidelines

---

• Read the 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) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.