From 1ef87ddd17d9899431cd4f6a0395c2fb1a47368e Mon Sep 17 00:00:00 2001 From: hussein-awala Date: Thu, 10 Sep 2026 00:53:18 +0200 Subject: [PATCH] Preserve concurrent task changes when triggers report results A late trigger event or failure could overwrite a task that another transaction had cleared, failed, resumed, or attached to a different trigger. Terminal events could also send callbacks and XComs for a transition that had already lost the race. --- airflow-core/src/airflow/models/trigger.py | 24 +-- .../tests/unit/models/test_trigger.py | 150 +++++++++++++++++- 2 files changed, 162 insertions(+), 12 deletions(-) diff --git a/airflow-core/src/airflow/models/trigger.py b/airflow-core/src/airflow/models/trigger.py index 75578974db826..18723151877f6 100644 --- a/airflow-core/src/airflow/models/trigger.py +++ b/airflow-core/src/airflow/models/trigger.py @@ -268,6 +268,18 @@ def clean_unused(cls, *, session: Session = NEW_SESSION) -> None: delete(Trigger).where(Trigger.id.in_(ids)).execution_options(synchronize_session=False) ) + @staticmethod + def _get_deferred_task_instances(trigger_id: int, *, session: Session) -> Iterable[TaskInstance]: + """Lock and refresh tasks still waiting on this trigger before applying its result.""" + query = ( + select(TaskInstance) + .where(TaskInstance.trigger_id == trigger_id, TaskInstance.state == TaskInstanceState.DEFERRED) + .order_by(TaskInstance.id) + .execution_options(populate_existing=True) + ) + # Wait for concurrent transitions; skipping a locked row would lose this result. + return session.scalars(with_row_locks(query, session, of=TaskInstance)) + @classmethod @provide_session def submit_event(cls, trigger_id, event: TriggerEvent, *, session: Session = NEW_SESSION) -> None: @@ -278,11 +290,7 @@ def submit_event(cls, trigger_id, event: TriggerEvent, *, session: Session = NEW Send an event to all assets associated to the trigger. """ # Resume deferred tasks - for task_instance in session.scalars( - select(TaskInstance).where( - TaskInstance.trigger_id == trigger_id, TaskInstance.state == TaskInstanceState.DEFERRED - ) - ): + for task_instance in cls._get_deferred_task_instances(trigger_id, session=session): handle_event_submit(event, task_instance=task_instance, session=session) # Send an event to assets @@ -317,11 +325,7 @@ def submit_failure(cls, trigger_id, exc=None, *, session: Session = NEW_SESSION) the runtime code understands as immediate-fail, and pack the error into next_kwargs. """ - for task_instance in session.scalars( - select(TaskInstance).where( - TaskInstance.trigger_id == trigger_id, TaskInstance.state == TaskInstanceState.DEFERRED - ) - ): + for task_instance in cls._get_deferred_task_instances(trigger_id, session=session): # Add the error and set the next_method to the fail state if isinstance(exc, BaseException): traceback = format_exception(type(exc), exc, exc.__traceback__) diff --git a/airflow-core/tests/unit/models/test_trigger.py b/airflow-core/tests/unit/models/test_trigger.py index 0ab5a89e7ca03..27a91a8e68b05 100644 --- a/airflow-core/tests/unit/models/test_trigger.py +++ b/airflow-core/tests/unit/models/test_trigger.py @@ -19,6 +19,8 @@ import datetime import json from collections.abc import AsyncIterator +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError +from threading import Event from typing import TYPE_CHECKING, Any from unittest.mock import patch @@ -26,12 +28,13 @@ import pytest import pytz from cryptography.fernet import Fernet -from sqlalchemy import delete, func, select +from sqlalchemy import delete, event as sqlalchemy_event, func, select, update +from sqlalchemy.exc import OperationalError from airflow._shared.timezones import timezone from airflow.jobs.job import Job from airflow.jobs.triggerer_job_runner import TriggererJobRunner -from airflow.models import TaskInstance, Trigger +from airflow.models import DagRun, TaskInstance, Trigger from airflow.models.asset import AssetEvent, AssetModel, AssetWatcherModel from airflow.models.callback import Callback, TriggererCallback from airflow.models.taskinstancehistory import TaskInstanceHistory @@ -344,6 +347,149 @@ def test_submit_failure(session, create_task_instance): assert updated_task_instance.next_method == "__fail__" +@pytest.mark.parametrize( + "trigger_event", + [TriggerEvent("payload"), None, TaskSuccessEvent(), TaskFailedEvent(), TaskSkippedEvent()], + ids=["event", "failure", "success", "failed", "skipped"], +) +def test_submit_trigger_result_locks_only_dependent_tasks(session, create_task_instance, trigger_event): + if session.bind.dialect.name == "sqlite": + pytest.skip("SQLite does not support row locks") + + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + session.add(trigger) + task_instance = create_task_instance(session=session, state=State.DEFERRED) + task_instance.trigger_id = trigger.id + session.commit() + ti_id = task_instance.id + dag_run_id = task_instance.dag_run.id + checked_locks = [] + + def check_locks(orm_execute_state): + if ( + not orm_execute_state.is_select + or orm_execute_state.is_relationship_load + or orm_execute_state.is_column_load + or orm_execute_state.bind_mapper.class_ is not TaskInstance + ): + return None + result = orm_execute_state.invoke_statement() + with pytest.raises(OperationalError): + with create_session(scoped=False) as competing_session: + competing_session.execute( + select(TaskInstance.id).where(TaskInstance.id == ti_id).with_for_update(nowait=True) + ) + with create_session(scoped=False) as competing_session: + competing_session.execute( + select(DagRun.id).where(DagRun.id == dag_run_id).with_for_update(nowait=True) + ) + checked_locks.append(ti_id) + return result + + sqlalchemy_event.listen(session, "do_orm_execute", check_locks) + try: + if trigger_event is None: + Trigger.submit_failure(trigger.id, session=session) + else: + Trigger.submit_event(trigger.id, trigger_event, session=session) + finally: + sqlalchemy_event.remove(session, "do_orm_execute", check_locks) + + assert checked_locks == [ti_id] + + +def test_submit_event_refreshes_cached_resume_kwargs(session, create_task_instance): + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + session.add(trigger) + task_instance = create_task_instance(session=session, state=State.DEFERRED) + task_instance.trigger_id = trigger.id + task_instance.next_kwargs = {"version": "old"} + session.commit() + ti_id, trigger_id = task_instance.id, trigger.id + + with create_session(scoped=False) as competing_session: + competing_session.execute( + update(TaskInstance).where(TaskInstance.id == ti_id).values(next_kwargs={"version": "new"}) + ) + + assert task_instance.next_kwargs == {"version": "old"} + Trigger.submit_event(trigger_id, TriggerEvent("payload"), session=session) + assert task_instance.next_kwargs == {"version": "new", "event": "payload"} + + +@pytest.mark.parametrize("transition", ["fail", "clear", "redeferral", "resume"]) +@pytest.mark.parametrize( + "trigger_event", + [TriggerEvent("late"), None, TaskSuccessEvent(xcoms={"return_value": "late"})], + ids=["event", "failure", "terminal-event"], +) +@patch("airflow.callbacks.database_callback_sink.DatabaseCallbackSink.send", autospec=True) +def test_submit_trigger_result_preserves_concurrent_transition( + mock_send, session, create_task_instance, transition, trigger_event +): + if session.bind.dialect.name == "sqlite": + pytest.skip("SQLite does not support concurrent row locking") + + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + replacement_trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + session.add_all([trigger, replacement_trigger]) + task_instance = create_task_instance(session=session, state=State.DEFERRED) + task_instance.trigger_id = trigger.id + task_instance.next_method = "execute_complete" + task_instance.next_kwargs = {"version": "initial"} + session.commit() + ti_id, trigger_id = task_instance.id, trigger.id + query_started = Event() + + def submit_result(): + with create_session(scoped=False) as result_session: + sqlalchemy_event.listen( + result_session, "do_orm_execute", lambda _: query_started.set(), once=True + ) + if trigger_event is None: + Trigger.submit_failure(trigger_id, session=result_session) + else: + Trigger.submit_event(trigger_id, trigger_event, session=result_session) + + expected_state = { + "fail": State.FAILED, + "clear": None, + "redeferral": State.DEFERRED, + "resume": State.SCHEDULED, + }[transition] + expected_trigger_id = replacement_trigger.id if transition == "redeferral" else None + expected_kwargs = {"version": "current"} + + with ThreadPoolExecutor(max_workers=1) as executor: + with create_session(scoped=False) as competing_session: + competing_session.execute( + update(TaskInstance) + .where(TaskInstance.id == ti_id) + .values( + state=expected_state, + trigger_id=expected_trigger_id, + next_kwargs=expected_kwargs, + ) + ) + pending_result = executor.submit(submit_result) + assert query_started.wait(timeout=10) + with pytest.raises(FutureTimeoutError): + pending_result.result(timeout=0.2) + pending_result.result(timeout=10) + + session.refresh(task_instance) + assert task_instance.state == expected_state + assert task_instance.trigger_id == expected_trigger_id + assert task_instance.next_kwargs == expected_kwargs + assert task_instance.next_method == "execute_complete" + mock_send.assert_not_called() + assert not session.scalars( + XComModel.get_many( + dag_ids=[task_instance.dag_id], task_ids=[task_instance.task_id], run_id=task_instance.run_id + ) + ).all() + + @pytest.mark.parametrize( ("event_cls", "expected"), [