[v3-4-test] Fix scheduler crash when retrying a task that timed out or was killed externally (#75472)

airflow-core/src/airflow/callbacks/database_callback_sink.py

@@ -19,9 +19,9 @@
 
 from typing import TYPE_CHECKING
 
+from airflow import settings
 from airflow.callbacks.base_callback_sink import BaseCallbackSink
 from airflow.models.db_callback_request import DbCallbackRequest
-from airflow.utils.session import NEW_SESSION, provide_session
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
@@ -32,8 +32,14 @@
 class DatabaseCallbackSink(BaseCallbackSink):
     """Sends callbacks to database."""
 
-    @provide_session
-    def send(self, callback: CallbackRequest, *, session: Session = NEW_SESSION) -> None:
-        """Send callback for execution."""
+    def send(self, callback: CallbackRequest, *, session: Session | None = None) -> None:
+        """
+        Send callback for execution.
+
+        The request is added to ``session``, or to the thread's scoped session when none is given,
+        without committing: the caller commits it with the rest of its transaction.
+        """
+        if session is None:
+            session = settings.get_session()()
         db_callback = DbCallbackRequest(callback=callback, priority_weight=1)
         session.add(db_callback)

airflow-core/src/airflow/executors/base_executor.py

@@ -836,6 +836,8 @@ def send_callback(self, request: CallbackRequest) -> None:
         Send callback for execution.
 
         Provides a default implementation which sends the callback to the `callback_sink` object.
+        The scheduler's database sink adds the request to the calling thread's scoped session
+        without committing, so the caller must commit that session.
 
         :param request: Callback request to be executed.
         """

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -1706,7 +1706,7 @@ def process_executor_events(
                         ti.handle_failure(error=msg, session=session)
                         continue
                     cls.logger().exception("Marking task instance %s as %s", ti, state)
-                    ti.set_state(state)
+                    ti.set_state(state, session=session)
                     continue
                 ti.task = task
                 if task.has_on_retry_callback or task.has_on_failure_callback:
@@ -2178,8 +2178,10 @@ def _do_scheduling(self, session: Session) -> int:
         for dag_run, callback_to_run in callback_tuples:
             dag = cached_get_dag(dag_run)
             if dag:
-                # Sending callbacks to the database, so it must be done outside of prohibit_commit.
                 self._send_dag_callbacks_to_processor(dag, callback_to_run)
+                # The sink only adds the request to this session. Commit each one so a later Dag
+                # lookup error, or the critical section below rolling back, can't lose it.
+                session.commit()
             else:
                 self.log.error("DAG '%s' not found in serialized_dag table", dag_run.dag_id)
 

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -224,6 +224,16 @@ def get_callback_ti(ti):
     )
 
 
+def _attempt_states(ti, session):
+    attempts = session.scalars(
+        select(TaskInstance)
+        .where(TaskInstance.dag_id == ti.dag_id, TaskInstance.task_id == ti.task_id)
+        .execution_options(include_all_attempts=True)
+        .order_by(TaskInstance.try_number)
+    )
+    return [(attempt.try_number, attempt.state, attempt.working_set) for attempt in attempts]
+
+
 @pytest.fixture(scope="class")
 def disable_load_example():
     with conf_vars({("core", "load_examples"): "false"}):
@@ -1246,6 +1256,84 @@ def test_process_executor_events_with_callback(
             },
         )
 
+    @pytest.mark.parametrize(
+        "callback_kwargs",
+        [
+            pytest.param({"on_failure_callback": lambda context: None}, id="on_failure_callback"),
+            pytest.param({"email": "test@example.com"}, id="email"),
+        ],
+    )
+    def test_process_executor_events_callback_for_running_ti_with_retries_moves_it_to_retry(
+        self, callback_kwargs, dag_maker, session
+    ):
+        """Sending the callback must not close the session and detach the TI before its retry is set up."""
+        with dag_maker(dag_id="executor_event_running_ti_with_retries", fileloc="/test_path1/"):
+            EmptyOperator(task_id="task", retries=1, **callback_kwargs)
+        dr = dag_maker.create_dagrun()
+        ti = dr.task_instances[0]
+        executor = MockExecutor(do_update=False)
+        executor.callback_sink = DatabaseCallbackSink()
+        scheduler_job = Job()
+        session.add(scheduler_job)
+        session.flush()
+        self.job_runner = SchedulerJobRunner(scheduler_job, executors=[executor])
+        ti.state = State.RUNNING
+        ti.try_number = 1
+        ti.queued_by_job_id = scheduler_job.id
+        session.merge(ti)
+        session.commit()
+        executor.event_buffer[TaskInstanceUuid(ti.id)] = State.SUCCESS, None
+
+        self.job_runner._process_executor_events(executor=executor, session=session)
+        session.commit()
+
+        session.expire_all()
+        assert _attempt_states(ti, session) == [
+            (1, State.FAILED, None),
+            (2, State.UP_FOR_RETRY, True),
+        ]
+        assert session.scalar(select(func.count()).select_from(DbCallbackRequest)) == 1
+
+    def test_process_executor_events_dag_lookup_failure_keeps_session_for_rest_of_batch(
+        self, dag_maker, session
+    ):
+        """A failed Dag lookup must not close the session and detach the rest of the batch."""
+        with dag_maker(dag_id="executor_event_dag_lookup_failure"):
+            EmptyOperator(task_id="first", retries=1)
+            EmptyOperator(task_id="second", retries=1)
+        dr = dag_maker.create_dagrun()
+        executor = MockExecutor(do_update=False)
+        scheduler_job = Job()
+        session.add(scheduler_job)
+        session.flush()
+        self.job_runner = SchedulerJobRunner(scheduler_job, executors=[executor])
+        for ti in dr.task_instances:
+            ti.state = State.RUNNING
+            ti.try_number = 1
+            ti.queued_by_job_id = scheduler_job.id
+            session.merge(ti)
+            executor.event_buffer[TaskInstanceUuid(ti.id)] = State.SUCCESS, None
+        session.commit()
+        real_get_dag_for_run = self.job_runner.scheduler_dag_bag.get_dag_for_run
+        lookup_errors = [AirflowException("Dag lookup failed")]
+
+        def get_dag_for_run(*, dag_run, session):
+            if lookup_errors:
+                raise lookup_errors.pop()
+            return real_get_dag_for_run(dag_run=dag_run, session=session)
+
+        with mock.patch.object(
+            self.job_runner.scheduler_dag_bag, "get_dag_for_run", autospec=True, side_effect=get_dag_for_run
+        ):
+            self.job_runner._process_executor_events(executor=executor, session=session)
+        session.commit()
+
+        session.expire_all()
+        live_states = sorted(
+            session.scalars(select(TaskInstance.state).where(TaskInstance.dag_id == dr.dag_id))
+        )
+        assert live_states == [State.SUCCESS, State.UP_FOR_RETRY]
+
     @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest")
     def test_process_executor_events_callback_for_removed_mapped_task(
         self, mock_task_callback, dag_maker, session
@@ -10462,6 +10550,33 @@ def test_heartbeat_timeout_sets_callback_type_by_retry_eligibility(
         )
         assert dispatched_callback == expected_dispatched_callback
 
+    def test_heartbeat_timeout_with_retries_moves_ti_to_retry(self, dag_maker, session):
+        """Sending the heartbeat-timeout callback must not close the purge's session and detach the TI."""
+        with dag_maker(dag_id="hb_timeout_with_retries", session=session):
+            EmptyOperator(task_id="test_task", retries=1)
+        dag_run = dag_maker.create_dagrun(run_id="test_run", state=DagRunState.RUNNING)
+        executor = MockExecutor(do_update=False)
+        executor.callback_sink = DatabaseCallbackSink()
+        scheduler_job = Job()
+        session.add(scheduler_job)
+        session.flush()
+        self.job_runner = SchedulerJobRunner(scheduler_job, executors=[executor])
+        ti = dag_run.get_task_instance(task_id="test_task", session=session)
+        ti.state = TaskInstanceState.RUNNING
+        ti.try_number = 1
+        ti.queued_by_job_id = scheduler_job.id
+        ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
+        session.commit()
+
+        self.job_runner._find_and_purge_task_instances_without_heartbeats()
+
+        session.expire_all()
+        assert _attempt_states(ti, session) == [
+            (1, TaskInstanceState.FAILED, None),
+            (2, TaskInstanceState.UP_FOR_RETRY, True),
+        ]
+        assert session.scalar(select(func.count()).select_from(DbCallbackRequest)) == 1
+
     @pytest.mark.parametrize(
         ("retries", "callback_kind", "expected"),
         [