Ignore stale executor success with per-invocation workload_run_id

airflow-core/docs/migrations-ref.rst

@@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | Revision ID             | Revises ID       | Airflow Version   | Description                                                  |
 +=========================+==================+===================+==============================================================+
-| ``b6a9c2e7d410`` (head) | ``f8c2a1d94e03`` | ``3.4.0``         | Add draining state to DagModel.                              |
+| ``c7d4e8f1a203`` (head) | ``b6a9c2e7d410`` | ``3.4.0``         | Add workload_run_id to task_instance and                     |
+|                         |                  |                   | task_instance_history.                                       |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``b6a9c2e7d410``        | ``f8c2a1d94e03`` | ``3.4.0``         | Add draining state to DagModel.                              |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``f8c2a1d94e03``        | ``8d3f1a6b2c47`` | ``3.4.0``         | Add team_name and bundle_names scope columns to job table.   |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+

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

@@ -87,8 +87,8 @@ def get_execution_api_server_url(conf_source: AirflowConfigParser | ExecutorConf
     from airflow.models.taskinstance import TaskInstance
 
     # Event_buffer dict value type
-    # Tuple of: state, info
-    EventBufferValueType = tuple[str | None, Any]
+    # Tuple of: state, info, workload_run_id (optional; None when unknown/legacy)
+    EventBufferValueType = tuple[Any, Any, str | None]
 
 
 log = logging.getLogger(__name__)
@@ -309,6 +309,10 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None)
             dict
         )
         self.running: set[WorkloadKey] = set()
+        # FIFO of workload_run_id values per TaskInstanceKey, appended on queue and
+        # consumed when a terminal SUCCESS/FAILED event is recorded. Survives a
+        # re-enqueue of the same key so a stale defer-exit SUCCESS keeps its old id.
+        self._workload_run_ids: dict[TaskInstanceKey, deque[str]] = defaultdict(deque)
         self.event_buffer: dict[WorkloadKey, EventBufferValueType] = {}
         self._task_event_logs: deque[Log] = deque()
         self.conf = ExecutorConf(team_name)
@@ -388,14 +392,90 @@ def log_task_event(self, *, event: str, extra: str, ti_key: WorkloadKey):
         self._task_event_logs.append(Log(event=event, task_instance=ti_key, extra=extra))
 
     def queue_workload(self, workload: ExecutorWorkload, session: Session) -> None:
+
+        if isinstance(workload, workloads.ExecuteTask):
+            ti = workload.ti
+            self.queued_tasks[ti.key] = workload
+            run_id = getattr(ti, "workload_run_id", None)
+            if isinstance(run_id, str):
+                self._workload_run_ids[ti.key].append(run_id)
+        elif isinstance(workload, workloads.ExecuteCallback):
+            if not self.supports_callbacks:
+                raise NotImplementedError(
+                    f"{type(self).__name__} does not support ExecuteCallback workloads. "
+                    f"Set supports_callbacks = True and implement callback handling in _process_workloads(). "
+                    f"See LocalExecutor or CeleryExecutor for reference implementation."
+                )
+            self.queued_callbacks[workload.key] = workload
+        elif isinstance(workload, workloads.TestConnection):
+            if not self.supports_connection_test:
+                raise NotImplementedError(
+                    f"{type(self).__name__} does not support TestConnection workloads. "
+                    f"Set supports_connection_test = True and implement connection test handling "
+                    f"in _process_workloads(). See LocalExecutor for reference implementation."
+                )
+            self.queued_connection_tests[workload.key] = workload
+        else:
+            raise ValueError(
+                f"Un-handled workload type {type(workload).__name__!r} in {type(self).__name__}. "
+                f"Workload must be one of: ExecuteTask, ExecuteCallback, TestConnection."
+
         if workload.type not in self.supported_workload_types:
             raise NotImplementedError(
                 f"{type(self).__name__} does not support {workload.type.value} workloads. "
                 f"Add WorkloadType.{workload.type.name} to supported_workload_types and implement handling "
                 f"in _process_workloads()."
+
             )
         self.executor_queues[workload.type][workload.key] = workload
 
+    def _pop_workload_run_id(self, key: WorkloadKey) -> str | None:
+        """Consume the oldest queued workload_run_id for a task key, if any."""
+        if not isinstance(key, TaskInstanceKey):
+            return None
+        pending = self._workload_run_ids.get(key)
+        if not pending:
+            return None
+        run_id = pending.popleft()
+        if not pending:
+            del self._workload_run_ids[key]
+        return run_id
+
+    @staticmethod
+    def unpack_event(value: EventBufferValueType | tuple[Any, ...]) -> tuple[Any, Any, str | None]:
+        """Normalize event buffer values to ``(state, info, workload_run_id)``."""
+        if len(value) >= 3:
+            return value[0], value[1], value[2]
+        return value[0], value[1], None
+
+    def record_event(
+        self,
+        key: WorkloadKey,
+        state: WorkloadState,
+        info=None,
+        *,
+        workload_run_id: str | None = None,
+        consume_run_id: bool | None = None,
+    ) -> None:
+        """
+        Write an executor event, optionally attaching a workload_run_id.
+
+        Terminal task SUCCESS/FAILED events consume the next queued run id when
+        one was not supplied explicitly. Intermediate states (QUEUED/RUNNING)
+        must not consume it — those events are drained separately.
+        """
+        from airflow.utils.state import TaskInstanceState
+
+        if consume_run_id is None:
+            consume_run_id = isinstance(key, TaskInstanceKey) and state in (
+                TaskInstanceState.SUCCESS,
+                TaskInstanceState.FAILED,
+            )
+        popped_run_id = self._pop_workload_run_id(key) if consume_run_id else None
+        if workload_run_id is None:
+            workload_run_id = popped_run_id
+        self.event_buffer[key] = (state, info, workload_run_id)
+
     def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]:
         """
         Select and return the next batch of workloads to schedule, respecting priority policy.
@@ -571,40 +651,51 @@ def order_queued_tasks_by_priority(self) -> list:
     # TODO: This should not be using `TaskInstanceState` here, this is just "did the process complete, or did
     # it die". It is possible for the task itself to finish with success, but the state of the task to be set
     # to FAILED. By using TaskInstanceState enum here it confuses matters!
-    def change_state(self, key: WorkloadKey, state: WorkloadState, info=None, remove_running=True) -> None:
+    def change_state(
+        self,
+        key: WorkloadKey,
+        state: WorkloadState,
+        info=None,
+        remove_running=True,
+        workload_run_id: str | None = None,
+    ) -> None:
         """
         Change state of the task.
 
         :param key: Unique key for the task instance
         :param state: State to set for the task.
         :param info: Executor information for the task instance
         :param remove_running: Whether or not to remove the TI key from running set
+        :param workload_run_id: Invocation id for the finished worker; when omitted for
+            terminal task events, the next id queued for this key is consumed.
         """
         self.log.debug("Changing state: %s", key)
         if remove_running:
             try:
                 self.running.remove(key)
             except KeyError:
                 self.log.debug("Could not find key: %s", key)
-        self.event_buffer[key] = state, info
+        self.record_event(key, state, info, workload_run_id=workload_run_id)
 
-    def fail(self, key: WorkloadKey, info=None) -> None:
+    def fail(self, key: WorkloadKey, info=None, workload_run_id: str | None = None) -> None:
         """
         Set fail state for the event.
 
         :param info: Executor information for the task instance
         :param key: Unique key for the task instance
+        :param workload_run_id: Invocation id for the finished worker, if known
         """
-        self.change_state(key, state_class_for_key(key).FAILED, info)
+        self.change_state(key, state_class_for_key(key).FAILED, info, workload_run_id=workload_run_id)
 
-    def success(self, key: WorkloadKey, info=None) -> None:
+    def success(self, key: WorkloadKey, info=None, workload_run_id: str | None = None) -> None:
         """
         Set success state for the event.
 
         :param info: Executor information for the task instance
         :param key: Unique key for the task instance
+        :param workload_run_id: Invocation id for the finished worker, if known
         """
-        self.change_state(key, state_class_for_key(key).SUCCESS, info)
+        self.change_state(key, state_class_for_key(key).SUCCESS, info, workload_run_id=workload_run_id)
 
     def queued(self, key: WorkloadKey, info=None) -> None:
         """

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

@@ -98,7 +98,7 @@ def _run_worker(
             unread_messages.value -= 1
 
         if workload.running_state is not None:
-            output.put((workload.key, workload.running_state, None))
+            output.put((workload.key, workload.running_state, None, None))
 
         try:
             BaseExecutor.run_workload(
@@ -107,10 +107,12 @@ def _run_worker(
                 proctitle=f"{_get_executor_process_title_prefix(team_conf.team_name)} {workload.display_name}",
                 subprocess_logs_to_stdout=True,
             )
-            output.put((workload.key, workload.success_state, None))
+            run_id = getattr(getattr(workload, "ti", None), "workload_run_id", None)
+            output.put((workload.key, workload.success_state, None, run_id))
         except Exception as e:
             log.exception("Workload execution failed.", workload_type=type(workload).__name__)
-            output.put((workload.key, workload.failure_state, e))
+            run_id = getattr(getattr(workload, "ti", None), "workload_run_id", None)
+            output.put((workload.key, workload.failure_state, e, run_id))
 
 
 class LocalExecutor(BaseExecutor):
@@ -250,8 +252,15 @@ def sync(self) -> None:
     def _read_results(self):
         try:
             while not self.result_queue.empty():
-                key, state, exc = self.result_queue.get()
-                self.change_state(key, state)
+                result = self.result_queue.get()
+                # Support legacy 3-tuples (key, state, exc) and current 4-tuples
+                # (key, state, exc, workload_run_id) from workers / tests.
+                if len(result) == 3:
+                    key, state, _exc = result
+                    workload_run_id = None
+                else:
+                    key, state, _exc, workload_run_id = result
+                self.change_state(key, state, workload_run_id=workload_run_id)
         except (OSError, EOFError):
             self.log.exception("Error reading from result queue")
 

airflow-core/src/airflow/executors/workloads/task.py

@@ -46,6 +46,7 @@ class TaskInstanceDTO(TaskInstance):
     priority_weight: int
 
     external_executor_id: str | None = Field(default=None, exclude=True)
+    workload_run_id: str | None = Field(default=None, exclude=True)
     executor_config: dict | None = Field(default=None, exclude=True)
 
     # TODO: Task-SDK: Can we replace TaskInstanceKey with just the uuid across the codebase?

airflow-core/src/airflow/executors/workloads/types.py

@@ -32,7 +32,8 @@
 
 if TYPE_CHECKING:
     # Type alias for executor workload results (used by executor implementations)
-    WorkloadResultType: TypeAlias = tuple[WorkloadKey, WorkloadState, Exception | None]
+    # key, state, exception, optional workload_run_id from the finished invocation
+    WorkloadResultType: TypeAlias = tuple[WorkloadKey, WorkloadState, Exception | None, str | None]
 
 # Type alias for scheduler workloads (ORM models that can be routed to executors)
 # Must be outside TYPE_CHECKING for use in function signatures

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

@@ -1081,6 +1081,9 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
                 "state": TaskInstanceState.QUEUED,
                 "queued_dttm": timezone.utcnow(),
                 "queued_by_job_id": self.job.id,
+                # Per-invocation id so a stale executor SUCCESS from a previous worker
+                # (e.g. defer exit) cannot be matched to a later enqueue of the same key.
+                "workload_run_id": random_db_uuid(),
             }
 
             # Pre-assign external_executor_id atomically with the QUEUED state so it
@@ -1116,25 +1119,26 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
                 .execution_options(synchronize_session=False)
             )
 
+            # Always read workload_run_id (and external_executor_id when pre-assigned)
+            # back onto in-memory objects so ExecuteTask.make carries them through
+            # make_transient. Use RETURNING on PostgreSQL; SELECT elsewhere.
+            returning_cols = [TI.id, TI.workload_run_id]
             if pre_assign_executors:
-                # Read the DB-generated UUIDs back onto the in-memory objects so the
-                # workload DTO carries them through to send_workload_to_executor (the
-                # objects are about to be detached by make_transient). Use RETURNING
-                # where supported (PostgreSQL); fall back to a SELECT for MySQL and
-                # SQLite (RETURNING requires SQLite 3.35+ which isn't guaranteed).
-                if get_dialect_name(session) == "postgresql":
-                    result = session.execute(queued_update.returning(TI.id, TI.external_executor_id))
-                    id_map = {row[0]: row[1] for row in result}
-                else:
-                    session.execute(queued_update)
-                    id_rows = session.execute(
-                        select(TI.id, TI.external_executor_id).where(filter_for_tis)
-                    ).all()
-                    id_map = {row[0]: row[1] for row in id_rows}
-                for ti in executable_tis:
-                    ti.external_executor_id = id_map.get(ti.id)
+                returning_cols.append(TI.external_executor_id)
+
+            if get_dialect_name(session) == "postgresql":
+                result = session.execute(queued_update.returning(*returning_cols))
+                rows = list(result)
             else:
                 session.execute(queued_update)
+                rows = list(session.execute(select(*returning_cols).where(filter_for_tis)).all())
+
+            workload_run_id_map = {row[0]: row[1] for row in rows}
+            external_id_map = {row[0]: row[2] for row in rows} if pre_assign_executors else {}
+            for ti in executable_tis:
+                ti.workload_run_id = workload_run_id_map.get(ti.id)
+                if pre_assign_executors:
+                    ti.external_executor_id = external_id_map.get(ti.id)
 
             for ti in executable_tis:
                 ti.emit_state_change_metric(TaskInstanceState.QUEUED)
@@ -1412,8 +1416,14 @@ def process_executor_events(
         tis_with_right_state: list[TaskInstanceKey] = []
         callback_keys_with_events: list[CallbackKey] = []
 
+        def _unpack_event(value: tuple) -> tuple[Any, Any, str | None]:
+            if len(value) >= 3:
+                return value[0], value[1], value[2]
+            return value[0], value[1], None
+
         # Report execution - handle both task and callback events
-        for key, (state, _) in event_buffer.items():
+        for key, event_value in event_buffer.items():
+            state, _, _ = _unpack_event(event_value)
             if isinstance(key, TaskInstanceKey):
                 existing_try = ti_primary_key_to_try_number_map.get(key.primary)
                 if existing_try is not None and existing_try != key.try_number:
@@ -1447,7 +1457,7 @@ def process_executor_events(
 
         # Handle callback state events
         for callback_id in callback_keys_with_events:
-            state, info = event_buffer.pop(callback_id)
+            state, info, _ = _unpack_event(event_buffer.pop(callback_id))
             callback = session.get(Callback, UUID(str(callback_id)))
             if not callback:
                 # This should not normally happen - we just received an event for this callback.
@@ -1510,7 +1520,7 @@ def process_executor_events(
                     ti.state,
                     job_id,
                 )
-            state, info = event_buffer.pop(buffer_key)
+            state, info, event_workload_run_id = _unpack_event(event_buffer.pop(buffer_key))
 
             if state in (TaskInstanceState.QUEUED, TaskInstanceState.RUNNING):
                 ti.external_executor_id = info
@@ -1559,22 +1569,27 @@ def process_executor_events(
             # from the worker exit after defer() has not been processed yet - should not fail it.
             # 4) the trigger already put the TI back to queued (resume after defer) but the executor success
             # from the worker exit after defer() has not been processed yet - should not fail it.
-
-            # All of this could also happen if the state is "running",
-            # but that is handled by the scheduler detecting task instances without heartbeats.
+            # 5) the resumed attempt is already RUNNING when the stale defer-exit SUCCESS arrives —
+            # workload_run_id mismatch identifies this without relying on next_method/state heuristics.
 
             ti_queued = ti.try_number == buffer_key.try_number and ti.state in (
                 TaskInstanceState.SCHEDULED,
                 TaskInstanceState.QUEUED,
                 TaskInstanceState.RUNNING,
                 TaskInstanceState.RESTARTING,
             )
+            stale_workload_run = (
+                event_workload_run_id is not None
+                and ti.workload_run_id is not None
+                and event_workload_run_id != ti.workload_run_id
+            )
             ti_requeued = (
                 ti.queued_by_job_id != job_id  # Another scheduler has queued this task again
                 or executor.has_task(ti)  # This scheduler has this task already
+                or stale_workload_run
                 or (
-                    # Resume-after-defer: trigger moved TI to scheduled or queued (next_method set)
-                    # before we saw the executor success from the defer exit for the same try_number.
+                    # Defense in depth for older events without workload_run_id: resume-after-defer
+                    # while next_method is still set (SCHEDULED/QUEUED variants from #66431/#68741).
                     ti.state in (TaskInstanceState.SCHEDULED, TaskInstanceState.QUEUED)
                     and state == TaskInstanceState.SUCCESS
                     and ti.next_method is not None
@@ -3545,6 +3560,7 @@ def adopt_or_reset_orphaned_tasks(self, *, session: Session = NEW_SESSION) -> in
                         ti.state = None
                         ti.queued_by_job_id = None
                         ti.external_executor_id = None
+                        ti.workload_run_id = None
                         ti.clear_next_method_args()
 
                     for ti in set(tis_to_adopt_or_reset) - set(to_reset):