1fanwang opened a new pull request, #71250: URL: https://github.com/apache/airflow/pull/71250
The scheduler claims task instances with `SELECT ... FOR UPDATE SKIP LOCKED` and relies on the clause to keep concurrent schedulers off the same rows. `SKIP LOCKED` is used at roughly twenty call sites across the scheduler, triggerer, DAG processor and asset manager for exactly this purpose. A server that accepts the clause and then discards it breaks that guarantee without saying anything: no error, no warning, no degraded-mode log line. Two schedulers are handed the same rows and both proceed. `with_row_locks()` already guards the case where a MySQL-family server cannot lock at all; this extends the same idea to a server that claims to skip but does not. TiDB is the case I hit. It has parsed and ignored `SKIP LOCKED` since 2020 ([pingcap/tidb#18207](https://github.com/pingcap/tidb/issues/18207), still open; listed under [unsupported features](https://docs.pingcap.com/tidb/stable/mysql-compatibility#unsupported-features)). Notably it does *not* route this through `tidb_enable_noop_functions`, the way it rejects `FOR SHARE` with error 1235 — so there is nothing for an operator to notice. The fallback is plain blocking `FOR UPDATE`, which is correct on such a server: schedulers serialize on each other instead of skipping ahead. That trades throughput for correctness, and only on servers that were already not providing the guarantee. PostgreSQL and MySQL are untouched. ### Testing Done Cluster: TiDB v8.5.1 (1 PD + 1 TiKV + 1 TiDB, `pingcap/*:v8.5.1`), MySQL 8.4 and PostgreSQL 16 as controls. The clause Airflow emits for the scheduler's critical-section query, compiled through `with_row_locks(..., skip_locked=True)` against each live server: | backend | emits SKIP LOCKED | locking clause | |---|---|---| | TiDB v8.5.1 | no | `FOR UPDATE OF task_instance` | | MySQL 8.4 | yes | `FOR UPDATE OF task_instance SKIP LOCKED` | | PostgreSQL 16 | yes | `FOR NO KEY UPDATE OF task_instance SKIP LOCKED` | <details><summary>Raw logs</summary> **The hazard, before the change.** Two workers running the scheduler's claim shape against TiDB, each holding its read before either commits: ``` claim sets: {'w1': [1, 2, 3], 'w2': [1, 2, 3]} commits : {'w1_commit': 'ok', 'w2_commit': 'ok'} final rows: ((1, 'w2'), (2, None), (3, None)) OVERLAP: [1, 2, 3] -> both workers saw the same rows ``` Both claimed every row, both committed, the later write won. `SHOW WARNINGS` was empty and `tidb_enable_noop_functions` was `OFF`. On MySQL 8 the two claim sets are disjoint. **The same race with the fallback in place** (plain `FOR UPDATE`, READ COMMITTED): ``` w1 claimed: [1, 2, 3] w2 claimed: [4, 5, 6] OVERLAP : [] final : ((1,'queued','w1'),(2,'queued','w1'),(3,'queued','w1'), (4,'queued','w2'),(5,'queued','w2'),(6,'queued','w2')) VERDICT: SAFE - disjoint claims, every row owned exactly once ``` The second worker blocks, re-reads under READ COMMITTED, sees the updated states and takes the next batch. **Emitted SQL per backend**, plus the warning operators get: ``` [warning] Database server reports as '8.0.11-tidb-v8.5.1', which accepts SKIP LOCKED but does not honor it. Falling back to plain FOR UPDATE so concurrent schedulers cannot claim the same rows. Schedulers will block on each other instead of skipping ahead. [airflow.utils.sqlalchemy] backend emits SKIP LOCKED locking clause ---------------------------------------------------------------------------- tidb False LIMIT 512 FOR UPDATE OF task_instance mysql True LIMIT 512 FOR UPDATE OF task_instance SKIP LOCKED postgres True LIMIT 512 FOR NO KEY UPDATE OF task_instance SKIP LOCKED VERDICT: PASS ``` **New tests against the unpatched source** — they fail, which is the point: ``` E AssertionError: expected call not found. E Expected: with_for_update(key_share=True) E Actual: with_for_update(skip_locked=True, key_share=True) FAILED ...test_with_row_locks_skip_locked_only_when_server_honors_it[tidb-ignores-skip-locked] FAILED ...test_with_row_locks_skip_locked_only_when_server_honors_it[tidb-lowercase-banner] 2 failed, 3 passed, 30 deselected ``` With the change: `35 passed` in `tests/unit/utils/test_sqlalchemy.py`. **Two schedulers against one TiDB**, 15 dag runs of a DAG with dynamic task mapping, XCom and a fan-in: ``` dag_run states : {'success': 15} task_instance states: {'success': 105} task instances with try_number > 1: 0 duplicate attempt rows in task_instance_history: 0 live SchedulerJob rows: 2 VERDICT: PASS - every task instance succeeded exactly once under two schedulers ``` **Regressions**: `tests/unit/utils/test_sqlalchemy.py` 35 passed; `tests/unit/jobs/test_scheduler_job.py -k "critical_section or executable_task_instances or row_lock or pool"` 47 passed. `prek` static checks and `mypy-airflow-core` clean. </details> The version banner is read once per engine and cached in a `WeakKeyDictionary`, so this costs a single `SELECT VERSION()` and holds no reference to a disposed engine. If the probe fails for any reason the previous behaviour is kept, so a transient connection problem cannot silently disable `SKIP LOCKED` on a server that supports it. -- 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]
