matthieuauger opened a new pull request, #71488: URL: https://github.com/apache/airflow/pull/71488
`TaskInstance.set_state` is `@provide_session`: called without a session, every pending task instance gets its own session — with the documented `sql_alchemy_pool_enabled = False` + pgbouncer setup, that is a fresh DB connection, a refresh SELECT, a merge SELECT, an UPDATE and a COMMIT **per task instance**. Marking a dag run failed on a run with 3,917 mapped task instances took **280 s** server-side (~19,500 statements, ~3,900 connections) on an otherwise idle PostgreSQL 17 with 2.5 ms RTT and every involved query shape indexed. Because this exceeds typical proxy timeouts, UI users see a failure while the backend keeps looping, and aborted requests leave `idle in transaction` sessions holding the `SELECT … FOR UPDATE` locks taken on running TIs. This change passes the ambient `session` through, keeping the loop in the request's existing transaction. The dag-run state update (`_set_dag_run_state(..., session)`) and the running-TI path (`set_state(..., session=session)`) in the same function already use that session, so transactional semantics are unchanged and existing tests cover the behavior. A/B measured on the same 71-TI run, same machine (Airflow 3.3.0, PostgreSQL 17 behind pgbouncer in transaction mode): | | unpatched | patched | |---|---:|---:| | PATCH `state=failed` | 5.61 s | **0.53 s** | | per-TI cost | 79 ms | **7.5 ms** | Extrapolated to the 3,917-TI run above: 280 s → ~29 s. A follow-up could replace the loop with a set-based UPDATE entirely (like #68666 did for the mark-success path) — kept out of this PR to stay minimal. `clear_task_instances` has the same per-TI shape and would benefit from the same treatment. --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes (please specify the tool below) Generated-by: Claude Code 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]
