diff --git a/airflow-core/newsfragments/72134.bugfix.rst b/airflow-core/newsfragments/72134.bugfix.rst new file mode 100644 index 0000000000000..fddca02ef5ca6 --- /dev/null +++ b/airflow-core/newsfragments/72134.bugfix.rst @@ -0,0 +1 @@ +MySQL deadlocks between the scheduler timeout sweep and triggerer unused-trigger cleanup no longer take down the triggerer. Deferred tasks whose triggerer is still alive are not timed out by the scheduler. diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 123f56c1f10c9..883536ba6a444 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -879,7 +879,7 @@ def ti_skip_downstream( tuple_(TI.task_id, TI.map_index).in_(task_ids), skippable_state_clause, ) - .values(state=TaskInstanceState.SKIPPED, start_date=now, end_date=now) + .values(state=TaskInstanceState.SKIPPED, start_date=now, end_date=now, trigger_id=None) .execution_options(synchronize_session=False) ) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index f8b2df4aac2ca..5087740371c3c 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -333,6 +333,8 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): """ job_type = "SchedulerJob" + # One timeout-fallback batch per scheduler tick (HITL sweep shape). + _TRIGGER_TIMEOUT_BATCH_SIZE = 100 def __init__( self, @@ -2974,6 +2976,7 @@ def _schedule_dag_run( ) for task_instance in unfinished_task_instances: task_instance.state = TaskInstanceState.SKIPPED + task_instance.trigger_id = None session.merge(task_instance) session.flush() self.log.info("Run %s of %s has timed-out", dag_run.run_id, dag_run.dag_id) @@ -3498,22 +3501,52 @@ def adopt_or_reset_orphaned_tasks(self, *, session: Session = NEW_SESSION) -> in def check_trigger_timeouts( self, max_retries: int = MAX_DB_RETRIES, *, session: Session = NEW_SESSION ) -> None: - """Mark any "deferred" task as failed if the trigger or execution timeout has passed.""" + """ + Time out deferred tasks whose triggerer is gone or never assigned. + + A healthy assigned triggerer already maps cancel-past-timeout to + ``submit_failure``. This sweep is a mixed-version / orphan fallback. + """ + now = timezone.utcnow() + threshold = conf.getint("triggerer", "triggerer_health_check_threshold") + alive_triggerer_ids = select(Job.id).where( + Job.end_date.is_(None), + Job.latest_heartbeat > now - timedelta(seconds=threshold), + Job.job_type == "TriggererJob", + ) for attempt in run_with_db_retries(max_retries, logger=self.log): with attempt: - result = session.execute( - update(TI) + query = ( + select(TI.id) + .outerjoin(Trigger, TI.trigger_id == Trigger.id) .where( TI.state == TaskInstanceState.DEFERRED, - TI.trigger_timeout < timezone.utcnow(), + TI.trigger_timeout < now, + or_( + TI.trigger_id.is_(None), + Trigger.id.is_(None), + Trigger.triggerer_id.is_(None), + ~Trigger.triggerer_id.in_(alive_triggerer_ids), + ), ) + .order_by(TI.id) + .limit(self._TRIGGER_TIMEOUT_BATCH_SIZE) + ) + query = with_row_locks(query, of=TI, session=session, skip_locked=True) + timed_out_ids = list(session.scalars(query).all()) + if not timed_out_ids: + return + result = session.execute( + update(TI) + .where(TI.id.in_(timed_out_ids)) .values( state=TaskInstanceState.SCHEDULED, next_method=TRIGGER_FAIL_REPR, next_kwargs={"error": TriggerFailureReason.TRIGGER_TIMEOUT}, - scheduled_dttm=timezone.utcnow(), + scheduled_dttm=now, trigger_id=None, ) + .execution_options(synchronize_session=False) ) num_timed_out_tasks = getattr(result, "rowcount", 0) if num_timed_out_tasks: diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index 46fbebb080038..9464336f9bf3f 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -1461,6 +1461,7 @@ def _filter_tis_and_exclude_removed(dag: SerializedDAG, tis: list[TI]) -> Iterab if ti.state != TaskInstanceState.REMOVED: self.log.error("Failed to get task for ti %s. Marking it as removed.", ti) ti.state = TaskInstanceState.REMOVED + ti.trigger_id = None session.flush() else: yield ti @@ -1916,6 +1917,7 @@ def _check_for_removed_or_restored_tasks( tags={**self.stats_tags, "dag_id": dag.dag_id}, ) ti.state = TaskInstanceState.REMOVED + ti.trigger_id = None continue try: @@ -1933,6 +1935,7 @@ def _check_for_removed_or_restored_tasks( "Removing the unmapped TI '%s' as the mapping can't be resolved yet", ti ) ti.state = TaskInstanceState.REMOVED + ti.trigger_id = None continue # Upstreams finished, check there aren't any extras if ti.map_index >= total_length: @@ -1942,6 +1945,7 @@ def _check_for_removed_or_restored_tasks( total_length, ) ti.state = TaskInstanceState.REMOVED + ti.trigger_id = None else: # Check if the number of mapped literals has changed, and we need to mark this TI as removed. if ti.map_index >= num_mapped_tis: @@ -1951,9 +1955,11 @@ def _check_for_removed_or_restored_tasks( num_mapped_tis, ) ti.state = TaskInstanceState.REMOVED + ti.trigger_id = None elif ti.map_index < 0: self.log.debug("Removing the unmapped TI '%s' as the mapping can now be performed", ti) ti.state = TaskInstanceState.REMOVED + ti.trigger_id = None return task_ids @@ -2141,7 +2147,7 @@ def _revise_map_indexes_if_mapped( TI.run_id == self.run_id, TI.map_index.in_(removed_indexes), ) - .values(state=TaskInstanceState.REMOVED) + .values(state=TaskInstanceState.REMOVED, trigger_id=None) ) session.flush() diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index ff610f2146588..ca4e1787b345f 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -443,6 +443,7 @@ def clear_task_instances( ti.max_tries = max(ti.max_tries, ti.try_number) ti.state = None ti.external_executor_id = None + ti.trigger_id = None ti.clear_next_method_args() # Match DagVersion to latest serialized DAG when running on the latest version. if use_latest_version: @@ -1052,6 +1053,8 @@ def set_state(self, state: str | None, *, session: Session = NEW_SESSION) -> boo self.log.debug("Setting task state for %s to %s", self, state) if self not in session: self.refresh_from_db(session=session) + if self.state == TaskInstanceState.DEFERRED: + self.trigger_id = None self.state = state self.start_date = self.start_date or current_time if self.state in State.finished or self.state == TaskInstanceState.UP_FOR_RETRY: diff --git a/airflow-core/src/airflow/models/taskmap.py b/airflow-core/src/airflow/models/taskmap.py index 96b3c0831cbcb..4ada36dfaead4 100644 --- a/airflow-core/src/airflow/models/taskmap.py +++ b/airflow-core/src/airflow/models/taskmap.py @@ -190,6 +190,7 @@ def expand_mapped_task( # are not done yet, so the task can't fail yet. if not task.dag or not task.dag.partial: unmapped_ti.state = TaskInstanceState.UPSTREAM_FAILED + unmapped_ti.trigger_id = None elif total_length < 1: # If the upstream maps this to a zero-length value, simply mark # the unmapped task instance as SKIPPED (if needed). @@ -199,6 +200,7 @@ def expand_mapped_task( total_length, ) unmapped_ti.state = TaskInstanceState.SKIPPED + unmapped_ti.trigger_id = None else: dr = unmapped_ti.dag_run zero_index_ti_exists = exists_query( @@ -289,5 +291,6 @@ def expand_mapped_task( to_update = session.scalars(with_row_locks(query, of=TaskInstance, session=session, skip_locked=True)) for ti in to_update: ti.state = TaskInstanceState.REMOVED + ti.trigger_id = None session.flush() return all_expanded_tis, total_expanded_ti_count - 1 diff --git a/airflow-core/src/airflow/models/trigger.py b/airflow-core/src/airflow/models/trigger.py index 75578974db826..c5e033ce7f569 100644 --- a/airflow-core/src/airflow/models/trigger.py +++ b/airflow-core/src/airflow/models/trigger.py @@ -37,7 +37,6 @@ from airflow.models.taskinstance import TaskInstance from airflow.serialization.enums import stringify_encoding_keys from airflow.triggers.base import BaseTaskEndEvent -from airflow.utils.retries import run_with_db_retries from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime, get_dialect_name, with_row_locks from airflow.utils.state import TaskInstanceState @@ -236,20 +235,12 @@ def clean_unused(cls, *, session: Session = NEW_SESSION) -> None: """ Delete all triggers that have no tasks dependent on them and are not associated to an asset. - Triggers have a one-to-many relationship to task instances, so we need to clean those up first. Afterward we can drop the triggers not referenced by anyone. - """ - # Update all task instances with trigger IDs that are not DEFERRED to remove them - for attempt in run_with_db_retries(): - with attempt: - session.execute( - update(TaskInstance) - .where( - TaskInstance.state != TaskInstanceState.DEFERRED, TaskInstance.trigger_id.is_not(None) - ) - .values(trigger_id=None) - ) + Deferred-exit paths must NULL ``task_instance.trigger_id`` themselves. + This method no longer bulk-updates task instances (that UPDATE deadlocked + with scheduler timeout scans on MySQL). + """ # Get all triggers that have no task instances, assets, or callbacks depending on them and delete them ids = select(cls.id).where( ~cls.assets.any(), @@ -258,10 +249,20 @@ def clean_unused(cls, *, session: Session = NEW_SESSION) -> None: ) ids = with_row_locks(ids, session, of=cls, skip_locked=True, key_share=False) if get_dialect_name(session) == "mysql": - # MySQL doesn't support DELETE with JOIN, so we need to do it in two steps + # MySQL doesn't support a DELETE whose subquery selects from the target table, + # so materialize the ids first. The DELETE re-checks the reference predicates: + # a task can defer onto one of these triggers in between, and deleting it + # would cascade-delete the task instance row. ids_list = list(session.scalars(ids).all()) session.execute( - delete(Trigger).where(Trigger.id.in_(ids_list)).execution_options(synchronize_session=False) + delete(Trigger) + .where( + Trigger.id.in_(ids_list), + ~cls.assets.any(), + ~cls.callback.has(), + ~cls.task_instance.has(), + ) + .execution_options(synchronize_session=False) ) else: session.execute( @@ -277,12 +278,15 @@ def submit_event(cls, trigger_id, event: TriggerEvent, *, session: Session = NEW Resume all tasks that were in deferred state. 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 - ) - ): + # Resume deferred tasks. SKIP LOCKED: if the scheduler fallback or another + # triggerer already took the row, this is a no-op. + query = ( + select(TaskInstance) + .where(TaskInstance.trigger_id == trigger_id, TaskInstance.state == TaskInstanceState.DEFERRED) + .order_by(TaskInstance.id) + ) + query = with_row_locks(query, of=TaskInstance, session=session, skip_locked=True) + for task_instance in session.scalars(query): handle_event_submit(event, task_instance=task_instance, session=session) # Send an event to assets @@ -317,11 +321,13 @@ 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 - ) - ): + query = ( + select(TaskInstance) + .where(TaskInstance.trigger_id == trigger_id, TaskInstance.state == TaskInstanceState.DEFERRED) + .order_by(TaskInstance.id) + ) + query = with_row_locks(query, of=TaskInstance, session=session, skip_locked=True) + for task_instance in session.scalars(query): # 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/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 80392fd17e206..cac52bf7738d5 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -2828,6 +2828,37 @@ def test_skip_downstream_still_skips_none_state_ti(self, client, session, dag_ma ti_downstream = dr.get_task_instance("downstream") assert ti_downstream.state == State.SKIPPED + def test_skip_downstream_nulls_trigger_id_on_deferred_ti(self, client, session, dag_maker): + """DEFERRED-exit must NULL trigger_id; skip-downstream does not go through set_state.""" + with dag_maker("skip_race_dag_deferred", session=session): + branch = EmptyOperator(task_id="branch") + downstream = EmptyOperator(task_id="downstream") + branch >> downstream + dr = dag_maker.create_dagrun(run_id="run") + + ti_branch = dr.get_task_instance("branch") + ti_branch.set_state(State.SUCCESS) + + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + session.add(trigger) + session.flush() + + ti_downstream = dr.get_task_instance("downstream") + ti_downstream.state = TaskInstanceState.DEFERRED + ti_downstream.trigger_id = trigger.id + session.commit() + + response = client.patch( + f"/execution/task-instances/{ti_branch.id}/skip-downstream", + json={"tasks": ["downstream"]}, + ) + assert response.status_code == 204 + + session.expire_all() + ti_downstream = dr.get_task_instance("downstream") + assert ti_downstream.state == State.SKIPPED + assert ti_downstream.trigger_id is None + class TestTIHealthEndpoint: def setup_method(self): diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index a7b332e763f1c..0804508d5871c 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -8143,6 +8143,142 @@ def test_timeout_triggers(self, dag_maker): assert ti1.next_method == "__fail__" assert ti2.state == State.DEFERRED + def test_timeout_triggers_does_not_flip_healthy_assigned(self, dag_maker): + """A past trigger_timeout on a TI whose triggerer is alive must not be swept.""" + session = settings.Session() + with dag_maker( + dag_id="test_timeout_triggers_healthy_assigned", + start_date=DEFAULT_DATE, + schedule="@once", + session=session, + ): + EmptyOperator(task_id="dummy1") + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("dummy1", session=session) + triggerer_job = Job(heartrate=10, state=State.RUNNING) + triggerer_job.job_type = "TriggererJob" + triggerer_job.latest_heartbeat = timezone.utcnow() + session.add(triggerer_job) + session.flush() + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + trigger.triggerer_id = triggerer_job.id + session.add(trigger) + session.flush() + ti.state = State.DEFERRED + ti.trigger_timeout = timezone.utcnow() - datetime.timedelta(seconds=60) + ti.trigger_id = trigger.id + session.flush() + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job) + self.job_runner.check_trigger_timeouts(session=session) + + session.refresh(ti) + assert ti.state == State.DEFERRED + assert ti.trigger_id == trigger.id + + def test_timeout_triggers_falls_back_for_dead_triggerer(self, dag_maker): + """Orphaned deferred TIs whose triggerer heartbeat is dead are still timed out.""" + session = settings.Session() + with dag_maker( + dag_id="test_timeout_triggers_dead_triggerer", + start_date=DEFAULT_DATE, + schedule="@once", + session=session, + ): + EmptyOperator(task_id="dummy1") + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("dummy1", session=session) + triggerer_job = Job(heartrate=10, state=State.RUNNING) + triggerer_job.job_type = "TriggererJob" + triggerer_job.latest_heartbeat = timezone.utcnow() - datetime.timedelta(hours=1) + session.add(triggerer_job) + session.flush() + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + trigger.triggerer_id = triggerer_job.id + session.add(trigger) + session.flush() + ti.state = State.DEFERRED + ti.trigger_timeout = timezone.utcnow() - datetime.timedelta(seconds=60) + ti.trigger_id = trigger.id + session.flush() + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job) + self.job_runner.check_trigger_timeouts(session=session) + + session.refresh(ti) + assert ti.state == State.SCHEDULED + assert ti.next_method == "__fail__" + assert ti.trigger_id is None + + @pytest.mark.backend("mysql") + def test_check_trigger_timeouts_skips_locked_rows(self, dag_maker, session): + """The timeout fallback must SKIP LOCKED rows another session already holds.""" + with dag_maker( + dag_id="test_timeout_triggers_skip_locked", + start_date=DEFAULT_DATE, + schedule="@once", + session=session, + ): + EmptyOperator(task_id="dummy1") + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("dummy1", session=session) + ti.state = State.DEFERRED + ti.trigger_timeout = timezone.utcnow() - datetime.timedelta(seconds=60) + session.commit() + + with create_session(scoped=False) as competing_session: + locked = competing_session.scalars( + with_row_locks( + select(TaskInstance).where(TaskInstance.id == ti.id), + of=TaskInstance, + session=competing_session, + skip_locked=True, + ) + ).all() + assert len(locked) == 1 + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job) + self.job_runner.check_trigger_timeouts(session=session) + + session.refresh(ti) + assert ti.state == State.DEFERRED + + def test_timeout_triggers_is_pk_ordered_and_bounded(self, dag_maker, monkeypatch): + """Fallback must ORDER BY id and stop at the batch size (one tick).""" + monkeypatch.setattr(SchedulerJobRunner, "_TRIGGER_TIMEOUT_BATCH_SIZE", 1) + session = settings.Session() + with dag_maker( + dag_id="test_timeout_triggers_pk_order", + start_date=DEFAULT_DATE, + schedule="@once", + session=session, + ): + EmptyOperator(task_id="dummy1") + EmptyOperator(task_id="dummy2") + dr = dag_maker.create_dagrun() + ti1 = dr.get_task_instance("dummy1", session=session) + ti2 = dr.get_task_instance("dummy2", session=session) + past = timezone.utcnow() - datetime.timedelta(seconds=60) + for ti in (ti1, ti2): + ti.state = State.DEFERRED + ti.trigger_timeout = past + session.flush() + first, second = sorted((ti1, ti2), key=lambda ti: ti.id) + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job) + self.job_runner.check_trigger_timeouts(session=session) + + session.refresh(first) + session.refresh(second) + assert first.state == State.SCHEDULED + assert first.next_method == "__fail__" + assert first.trigger_id is None + assert second.state == State.DEFERRED + def test_awaiting_input_timeout_with_defaults_resumes(self, dag_maker): """ A parked ``awaiting_input`` task past its deadline with defaults is resumed to SCHEDULED by diff --git a/airflow-core/tests/unit/models/test_cleartasks.py b/airflow-core/tests/unit/models/test_cleartasks.py index 407b51bbec23e..67175b2359b3d 100644 --- a/airflow-core/tests/unit/models/test_cleartasks.py +++ b/airflow-core/tests/unit/models/test_cleartasks.py @@ -28,6 +28,7 @@ from airflow.models.taskinstance import TaskInstance, TaskInstance as TI, clear_task_instances from airflow.models.taskinstancehistory import TaskInstanceHistory from airflow.models.taskreschedule import TaskReschedule +from airflow.models.trigger import Trigger from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.sensors.python import PythonSensor from airflow.serialization.definitions.dag import SerializedDAG @@ -135,9 +136,13 @@ def test_clear_task_instances_next_method(self, dag_maker, session): EmptyOperator(task_id="task0") ti0 = dag_maker.create_dagrun().task_instances[0] + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + session.add(trigger) + session.flush() ti0.state = State.DEFERRED ti0.next_method = "next_method" ti0.next_kwargs = {} + ti0.trigger_id = trigger.id session.add(ti0) session.commit() @@ -148,6 +153,7 @@ def test_clear_task_instances_next_method(self, dag_maker, session): assert ti0.next_method is None assert ti0.next_kwargs is None + assert ti0.trigger_id is None @pytest.mark.parametrize( ("state", "last_scheduling"), [(DagRunState.QUEUED, None), (DagRunState.RUNNING, DEFAULT_DATE)] diff --git a/airflow-core/tests/unit/models/test_mappedoperator.py b/airflow-core/tests/unit/models/test_mappedoperator.py index 3539428be43c3..eb0db8ed88da9 100644 --- a/airflow-core/tests/unit/models/test_mappedoperator.py +++ b/airflow-core/tests/unit/models/test_mappedoperator.py @@ -31,6 +31,7 @@ from airflow.models.dag_version import DagVersion from airflow.models.taskinstance import TaskInstance from airflow.models.taskmap import TaskMap +from airflow.models.trigger import Trigger from airflow.providers.standard.operators.python import PythonOperator from airflow.sdk import DAG, BaseOperator, TaskGroup, setup, task, task_group, teardown from airflow.serialization.definitions.baseoperator import SerializedBaseOperator @@ -248,6 +249,59 @@ def test_expand_mapped_task_instance_skipped_on_zero(dag_maker, session): assert indices == [(-1, TaskInstanceState.SKIPPED)] +def test_expand_mapped_task_nulls_trigger_id_on_removed_extra(dag_maker, session): + """Shrinking a map must NULL trigger_id; extras do not go through set_state.""" + with dag_maker(session=session, serialized=True) as dag: + task1 = BaseOperator(task_id="op1") + mapped = MockOperator.partial(task_id="task_2").expand(arg2=task1.output) + + mapped_deser = dag.task_dict[mapped.task_id] + dr = dag_maker.create_dagrun() + session.add( + TaskMap( + dag_id=dr.dag_id, + task_id=task1.task_id, + run_id=dr.run_id, + map_index=-1, + length=1, + keys=None, + ) + ) + session.execute( + delete(TaskInstance).where( + TaskInstance.dag_id == mapped.dag_id, + TaskInstance.task_id == mapped.task_id, + TaskInstance.run_id == dr.run_id, + ) + ) + dag_version = DagVersion.get_latest_version(dr.dag_id) + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + session.add(trigger) + session.flush() + kept = TaskInstance( + mapped_deser, + run_id=dr.run_id, + map_index=0, + state=TaskInstanceState.SUCCESS, + dag_version_id=dag_version.id, + ) + extra = TaskInstance( + mapped_deser, + run_id=dr.run_id, + map_index=1, + state=TaskInstanceState.DEFERRED, + dag_version_id=dag_version.id, + ) + extra.trigger_id = trigger.id + session.add_all([kept, extra]) + session.flush() + + TaskMap.expand_mapped_task(mapped_deser, dr.run_id, session=session) + session.refresh(extra) + assert extra.state == TaskInstanceState.REMOVED + assert extra.trigger_id is None + + @pytest.mark.parametrize( ("num_existing_tis", "expected"), ( diff --git a/airflow-core/tests/unit/models/test_trigger.py b/airflow-core/tests/unit/models/test_trigger.py index 0ab5a89e7ca03..0b37c7e61173c 100644 --- a/airflow-core/tests/unit/models/test_trigger.py +++ b/airflow-core/tests/unit/models/test_trigger.py @@ -153,8 +153,9 @@ def test_clean_unused(session, dag_maker): tis = {ti.task_id: ti for ti in dr.task_instances} tis["fake0"].state = State.DEFERRED tis["fake0"].trigger_id = trigger1.id + # DEFERRED-exit must NULL trigger_id; leftover FKs are no longer bulk-unlinked. tis["fake1"].state = State.SUCCESS - tis["fake1"].trigger_id = trigger2.id + tis["fake1"].trigger_id = None tis["fake2"].state = State.SUCCESS tis["fake2"].trigger_id = trigger4.id session.flush() @@ -180,6 +181,66 @@ def test_clean_unused(session, dag_maker): assert {result.id for result in results} == {trigger1.id, trigger4.id, trigger5.id, trigger6.id} +def test_clean_unused_does_not_update_task_instance(session, dag_maker): + """clean_unused must not bulk-UPDATE task_instance (lock-order inversion vs timeout).""" + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + session.add(trigger) + session.flush() + with dag_maker(session=session): + EmptyOperator(task_id="fake") + dr = dag_maker.create_dagrun() + ti = dr.task_instances[0] + ti.state = State.SUCCESS + ti.trigger_id = trigger.id + session.flush() + + from sqlalchemy.sql.dml import Update + + updates: list[object] = [] + original_execute = session.execute + + def _spy(stmt, *args, **kwargs): + if isinstance(stmt, Update): + updates.append(stmt) + return original_execute(stmt, *args, **kwargs) + + session.execute = _spy # type: ignore[method-assign] + Trigger.clean_unused(session=session) + assert updates == [] + session.refresh(ti) + assert ti.trigger_id == trigger.id + + +def test_clean_unused_keeps_triggers_referenced_after_candidate_select(session, dag_maker, monkeypatch): + """MySQL two-step DELETE must re-check references so a deferral between SELECT and DELETE survives.""" + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + session.add(trigger) + session.flush() + with dag_maker(session=session): + EmptyOperator(task_id="fake") + dr = dag_maker.create_dagrun() + ti = dr.task_instances[0] + + original_scalars = session.scalars + seen = {"done": False} + + def _scalars(statement, *args, **kwargs): + result = original_scalars(statement, *args, **kwargs) + if not seen["done"]: + ti.state = State.DEFERRED + ti.trigger_id = trigger.id + session.flush() + seen["done"] = True + return result + + monkeypatch.setattr(session, "scalars", _scalars) + monkeypatch.setattr("airflow.models.trigger.get_dialect_name", lambda _session: "mysql") + Trigger.clean_unused(session=session) + assert session.get(Trigger, trigger.id) is not None + session.refresh(ti) + assert ti.trigger_id == trigger.id + + @patch.object(TriggererCallback, "handle_event") def test_submit_event(mock_callback_handle_event, session, create_task_instance): """