ataulmujeeb-cyber opened a new pull request, #75525: URL: https://github.com/apache/airflow/pull/75525
A producer can select a pending partitioned asset run, and the scheduler can consume or delete that same row before the producer attaches its event. Nothing currently keeps the row eligible in between. Deletion case — the scheduler's stale-rollup cleanup: 1. A producer looks up the pending run for a consumer key and gets row `R`. 2. The scheduler claims `R`, finds its rollup fingerprint no longer matches the consumer's timetable, and deletes `R` together with its event links. 3. The producer commits its own `partitioned_asset_key_log` row pointing at `R`. `partitioned_asset_key_log` has no foreign key to `asset_partition_dag_run`, so that link simply survives pointing at nothing. Completion case, which is quieter: 1. A producer looks up pending run `R`. 2. The scheduler claims `R`, reads the links that contribute to it, creates the Dag run and stamps `created_dag_run_id`. 3. The producer commits its link to `R` afterwards. Nothing is orphaned here, so an orphan query cannot find it — the event is just missing from the run that was supposed to consume it. Separately, a transaction registering several assets makes several registrations and can take consumer keys in conflicting orders. Sorting within one registration (a separate follow-up) does not order the whole transaction, so that cycle remains. ## What this changes A producer locks the pending row it intends to reuse, by primary key, re-asserts that it is still pending, and refreshes the instance so the value it decides from is the locked value. The lock is **shared** when the incoming event carries no partition date and **exclusive** when it carries one, decided before acquiring so a shared lock is never upgraded mid-acquisition; date reconciliation becomes a single conditional `UPDATE`. The lock is held until the event and its links commit. The scheduler's claim of pending rows takes its row lock even when `[scheduler] use_row_level_locking=False`. That setting is documented as an HA option, but this claim is also what makes the scheduler skip a row a producer is still extending, so it is treated as a correctness lock. Genuine dialect incapability is still honoured. Remaining cross-registration cycles are recovered by bounded replay of the whole transaction at its owner — the public asset-event route and the task-state route — using existing `airflow.utils.retries`. Callbacks run outside the retry boundary, because by the time they run the transaction has committed and a rollback cannot undo them. Related: #71074, which introduced these pending rows and their marker. ## Dependencies This branch contains both independent follow-ups as well as this change, so GitHub's full diff below includes them: | Order | Change | PR | | --- | --- | --- | | 1 | Preserve unrelated errors during pending asset run creation | #75523 | | 2 | Order partition keys before creating pending asset runs | #75524 | | 3 | This change | — | Review range for this change alone, on the current head: `9f35e003ca..de0d81b9a2` (the lifecycle commit plus a one-line test-docstring correction). Once #75523 and #75524 merge, this branch can be rebased so its diff contains only this change. Neither of those should wait on the discussion here. ## Behaviour changes and limits - Public `POST /assets/events` listeners now run after the event commits. A listener failure no longer rolls back the event. - Dated reuse of the same consumer key serialises where it previously did not. Existing contention figures are exploratory only: the comparison reference varied both the locking and the ORM refresh, so it does not isolate lock cost, and production acceptability is not established. No threshold is proposed. - Replay makes three attempts with bounded backoff and then re-raises the original deadlock. It is classified narrowly: PostgreSQL `40P01` and MySQL `1213`, including a `1213` recovered from behind a missing-savepoint `1305`. A bare `1305`, lock-wait timeouts, arbitrary integrity errors and uncertain commits are not replayed. - Listener delivery is not crash-safe exactly-once. Attempt-level metrics and tracing can repeat, because they are emitted inside the replayed body. - Deletes that do not go through the scheduler's claim, mixed-version participants and SQLite concurrency remain outside this guarantee. - No migration and no executable model change. Upstream's backfill policy is untouched: the highest-ID pending row per `(target_dag_id, partition_key)` is marked, older duplicates can remain unmarked, re-upgrade repeats marking, reuse does not promote a marker, and no historical repair is added. #74568 restructures this same acquisition function into a batched helper, so the two will need to be sequenced; its current diff still references `_lock_asset_model` and an `asset_id` argument that #71074 removed, so its rebased shape is not yet known. If these are accepted, maintainers may want to consider the two small fixes for the 3.4 release branch, which already carries the pending-marker migration. No backport PRs are opened here. ## Testing Run for this submission on the branch head, with `BACKEND` set explicitly: ``` pytest airflow-core/tests/unit/assets/ pytest airflow-core/tests/unit/jobs/test_scheduler_job.py -k 'producer_first or scheduler_first' pytest airflow-core/tests/unit/utils/test_deadlock.py airflow-core/tests/unit/utils/test_retries.py airflow-core/tests/unit/utils/test_sqlalchemy.py pytest airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py -k 'DeadlockRecovery or RealDeadlockVictim' pytest airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py -k 'deadlock or post_commit_callback or state_cleanup or task_span' ``` | Selection | PostgreSQL / psycopg3 | MySQL / PyMySQL | | --- | --- | --- | | `tests/unit/assets/` | 130 passed | 130 passed | | lifecycle matrix | 16 passed | 16 passed | | retries + deadlock + sqlalchemy utilities | 67 passed | 67 passed | | public route owner | 4 passed | 4 passed | | task route owner | 18 passed | 18 passed | No skips in any of the above. The lifecycle matrix is 16 cases: marked and unmarked pending rows, `use_row_level_locking` on and off, stale-rollup cleanup and completion, and both arrival orders. Producer-first cases require the first scheduler tick to defer the held row and the next tick to do the real work — cleanup must remove the parent and both links; completion must create the run and consume both event ids, read from `DagRun.consumed_asset_events`. Scheduler-first cases require the producer to block, confirmed from the database's own view by matching the waiting and blocking backend ids, and then to reselect rather than attach to a row that was deleted or already fired. <details> <summary>Detailed validation</summary> Recorded during implementation on the equivalent tree, attributed to that source rather than to this branch head: | Area | Evidence | | --- | --- | | Manager selections | 90 passed on each concurrent backend; 78 passed and 12 skipped on SQLite | | Classifier / deadlock / SQLAlchemy utilities | 11, 19 and 42 passed per backend | | Replay relocated into `airflow.utils.retries` | its 19 replay tests and 6 existing retry tests passed together; function and class ASTs unchanged | | Lifecycle and route-owner selection, PostgreSQL | 29 passed after ordering, module-placement and type corrections | | Genuine public-route victim | real deadlock formed on both PostgreSQL and MySQL; the public worker was the victim and replayed; committed event and link effects checked | | Task-boundary witnesses, PostgreSQL | 6 passed, covering genuine abort and replay, exhaustion, and post-commit callbacks that are not replayed | The weighted MySQL variant of the public-victim schedule records the route's own worker taking `1305 "SAVEPOINT ... does not exist"` with `1213 "Deadlock found"` in its cause chain, then replaying and committing — the case the chain-walking classification exists for. Distinctions that matter when reading the above: - Genuine database deadlocks, injected errors, staged lookup controls, expected mutation failures and environment/setup failures are recorded separately. Only genuine runs are counted. - The task witness named `commit` aborts immediately before the underlying `COMMIT`. It does not establish recovery from an uncertain commit outcome. - These are focused selections, not whole scheduler or API suites. - SQLite skips are the two-transaction concurrency classes; they do not establish concurrent-writer behaviour. - An earlier combined implementation of this work was discarded after review found it dropped the final winner lookup after an insert conflict and placed replay in a new module the repository's module policy rejects. Its suite counts are not carried over, and its defects are not upstream's. - Static checks: Ruff 0.17.0 `check` and `format --check` pass over the 13 changed Python files on this head. Full-distribution `mypy-airflow-core` and `mypy-devel-common` passes are recorded from implementation time; single-file mypy was used later and is not equivalent to either hook, since both enumerate whole distributions. The OpenAPI generator ran directly in the CI container with no generated diff; its wrapper was skipped because it needs a nested image build. - An optional external-agent review gate was requested during development and was rejected before it ran; it is not claimed as a pass. - Earlier migration experiments are historical. This change adds no migration and did not repeat them. </details> --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes — Codex (GPT-6) and Claude (Opus 5) Generated-by: Codex (GPT-6), with review-driven corrections and publication preparation by Claude (Opus 5), 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). --- Drafted-by: Codex (GPT-6) and Claude (Opus 5) (no human review of this text 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]
