fixup! Run loop gates and advance iterations atomically

A gate decision that the pinned loop would reject (continue at the iteration limit, or a stop before the limit of a fixed-count loop) was accepted by set_xcom and only rewritten into a failure when the state update arrived. By then the task runner had already run the success callbacks, so a failed or retried gate never saw its failure or retry callbacks (T7.15). The three decision checks now live in validate_loop_decision, shared by complete_loop_gate and set_xcom, so the write is refused with a 400 inside execute() and goes through normal retry and failure handling; the state-update rewrite stays as a backstop and its tests now seed the invalid decision directly.

Reading one slot of a mapped loop terminal went through resolve_current_producers, which hydrated every live producer row of the pass just to detect two producers at one index (T7.18). The candidate query is now built once and select_current_producer_ids fetches only id and region_index, sharing the ambiguity check with the full-row path. The resolver also passes the producer request fields to both functions explicitly instead of through attrs.asdict, so mypy checks them and a one-sided rename cannot slip through (T7.19).

The first region of a loop was inserted with a plain add and flush, so two verify_integrity calls racing on the same run could fail on the unique slot key before the later task instance race handling applied; it now uses DynamicRegion.get_or_create, which reuses the winner's region (T7.21).

count_loaded_task_instances registered list.append as a mapper load listener, which receives (target, context) and raised TypeError on the first loaded row, so its callers only passed while nothing loaded; it now filters in a real listener and has a positive-path test (T7.13). The retry test for the gate now asserts the DeleteXCom for the decision key so skipping the delete is caught (T7.22), and the unsuccessful gate test matches LoopMaxIterationsExceeded instead of a bare Exception (T7.17).

The lock-order test for the loop decision write and gate completion paused the writer only after its first dag_run lock and seeded the same continue the writer wrote, so a set_xcom that locked in the reverse order, or a completion that consumed a stale decision, would still have passed (T7.14). The writer now pauses after its first locking statement of any table, observed at the engine so session.scalar locks are seen too, and sleeps briefly once completion has issued its first lock, so a reversed order deadlocks; the seeded decision is stop, which a fixed-count loop rejects if completion consumes it before the writer replaces it. The test is still timing-based by necessity: the short sleep stands in for observing the lock wait directly.

When a Dag change gives a loop a different gate task id on a run that already has regions, verify_integrity raised from _create_initial_tasks and failed the whole run's integrity check even though only that loop was affected (T7.16). It now logs a warning and skips that loop so the run's other tasks and loops are still created; the loop keeps no gate and cannot advance, which the warning names.

The loop_decision flag on SetXCom and the set_xcom route carried nothing the reserved decision key did not already carry, because the route required key == LOOP_DECISION_KEY inside the flagged branch and rejected every other reserved key outside it (T7.20). The route now takes the gate branch when the key is the decision key, the field is gone from SetXCom, the client, the request handler, the Cadwyn entry and the regenerated Task SDK, Go and TypeScript supervisor schemas. The loop result read test that asserted the pass's producers were loaded now asserts no producer row is loaded, which is what the earlier id-only selection guarantees.

Author: ashb

airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py

@@ -40,7 +40,7 @@
 from airflow.api_fastapi.execution_api.security import CurrentTIToken
 from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyArchivedTaskStateUpdates
 from airflow.exceptions import TaskNotFound
-from airflow.models.dagrun import DagRun
+from airflow.models.dagrun import DagRun, InvalidLoopDecision, validate_loop_decision
 from airflow.models.dynamic_region import (
     LOOP_DECISION_KEY,
     LOOP_XCOM_PREFIX,
@@ -582,7 +582,6 @@ def set_xcom(
         ),
     ] = None,
     map_index: Annotated[int, Query()] = -1,
-    loop_decision: bool = False,
     dag_result: Annotated[bool, Query(description="Whether this XCom is a dag result")] = False,
     mapped_length: Annotated[
         int | None, Query(ge=0, description="Number of mapped tasks this value expands into")
@@ -594,11 +593,10 @@ def set_xcom(
 
     region_id: UUID | None = None
     region_index: int | None = None
-    if loop_decision:
+    if key == LOOP_DECISION_KEY:
         caller = session.get(TaskInstance, token.id)
         if (
-            key != LOOP_DECISION_KEY
-            or value not in ("continue", "stop")
+            value not in ("continue", "stop")
             or mapped_length is not None
             or dag_result
             or map_index != -1
@@ -624,6 +622,10 @@ def set_xcom(
             raise HTTPException(status.HTTP_400_BAD_REQUEST, "Loop decisions require the pinned loop gate")
         if caller.state != TaskInstanceState.RUNNING or caller.working_set is not True:
             raise HTTPException(status.HTTP_409_CONFLICT, "Loop gate is no longer running")
+        try:
+            validate_loop_decision(value, caller.region_index, loop[0])
+        except InvalidLoopDecision as error:
+            raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error
         region_id, region_index = caller.region_id, caller.region_index
     elif key.startswith(LOOP_XCOM_PREFIX):
         raise HTTPException(status.HTTP_400_BAD_REQUEST, "Reserved loop XCom key")

airflow-core/src/airflow/models/dagrun.py

@@ -151,6 +151,17 @@ class InvalidLoopDecision(ValueError):
     """A gate cannot complete successfully with the available decision."""
 
 
+def validate_loop_decision(decision: str | None, region_index: int, group: SerializedLoopTaskGroup) -> None:
+    """Raise :class:`InvalidLoopDecision` unless ``decision`` is valid for the gate of pass ``region_index``."""
+    at_limit = region_index + 1 >= group.max_iterations
+    if decision not in ("continue", "stop"):
+        raise InvalidLoopDecision("Successful loop gate requires a decision")
+    if decision == "continue" and at_limit:
+        raise InvalidLoopDecision("Loop cannot continue beyond its iteration limit")
+    if not group.has_until and (decision == "stop") != at_limit:
+        raise InvalidLoopDecision("Fixed-count loop decision does not match its iteration limit")
+
+
 class TISchedulingDecision(NamedTuple):
     """Type of return for DagRun.task_instance_scheduling_decisions."""
 
@@ -2246,13 +2257,7 @@ def complete_loop_gate(
                 session.delete(signal)
             return
         decision = signal.value if signal is not None else None
-        at_limit = gate.region_index + 1 >= group.max_iterations
-        if decision not in ("continue", "stop"):
-            raise InvalidLoopDecision("Successful loop gate requires a decision")
-        if decision == "continue" and at_limit:
-            raise InvalidLoopDecision("Loop cannot continue beyond its iteration limit")
-        if not group.has_until and (decision == "stop") != at_limit:
-            raise InvalidLoopDecision("Fixed-count loop decision does not match its iteration limit")
+        validate_loop_decision(decision, gate.region_index, group)
         session.delete(signal)
         if decision == "stop":
             return
@@ -2315,10 +2320,16 @@ def _create_initial_tasks(
                     )
                     .limit(1)
                 ):
-                    raise ValueError(f"Loop {group_id!r} has regions but no live gate")
-                region = DynamicRegion(dag_id=self.dag_id, run_id=self.run_id, node_id=group_id)
-                session.add(region)
-                session.flush()
+                    log.warning(
+                        "skipping loop that has regions but no live gate",
+                        dag_id=self.dag_id,
+                        run_id=self.run_id,
+                        loop=group_id,
+                    )
+                    continue
+                region = DynamicRegion.get_or_create(
+                    dag_id=self.dag_id, run_id=self.run_id, node_id=loops[group_id].node_id, session=session
+                )
                 coordinates = [(region.id, 0)]
             for region_id, region_index in coordinates:
                 yield from self._create_tasks(

airflow-core/src/airflow/models/dynamic_region.py

@@ -21,7 +21,7 @@
 from collections.abc import Collection, Iterable
 from datetime import datetime
 from functools import partial
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 from uuid import UUID
 
 import attrs
@@ -531,22 +531,20 @@ def _validate_producer_request(
         raise ValueError("Previous-iteration lookup requires a loop context")
 
 
-def resolve_current_producers(
+def _build_candidate_query(
+    query: Select[Any],
     *,
     dag_id: str,
     run_id: str,
     task_id: str,
     is_mapped: bool,
-    context: ProducerContext | None = None,
-    map_indexes: int | Collection[int] | None = None,
-    region_id: UUID | None = None,
-    region_index: int | None = None,
+    context: ProducerContext | None,
+    map_indexes: int | Collection[int] | None,
+    region_id: UUID | None,
+    region_index: int | None,
     session: Session,
-) -> tuple[TaskInstance, ...]:
-    """Resolve the live producer task instances whose data the caller reads by task instance UUID."""
-    from airflow.models.taskinstance import TaskInstance
-
-    _validate_producer_request(context=context, region_id=region_id, region_index=region_index)
+) -> Select[Any] | None:
+    """Narrow ``query`` to the candidate producers of the caller's scope, or ``None`` if there are none."""
     regions: dict[UUID, DynamicRegion] = {}
     position = None
     if context and region_id is None:
@@ -560,32 +558,74 @@ def resolve_current_producers(
             if context.previous_iteration:
                 position = position[0], position[1] - 1
                 if position[1] < 0:
-                    return ()
+                    return None
+
+    return _filter_producers(
+        query,
+        dag_id=dag_id,
+        run_id=run_id,
+        task_id=task_id,
+        is_mapped=is_mapped,
+        map_indexes=map_indexes,
+        region_id=region_id,
+        region_index=region_index,
+        top_level_only=region_id is None and position is None,
+        loop_pass=_load_loop_pass_regions(position, session=session) if position is not None else None,
+    )
 
-    candidates = session.scalars(
-        _filter_producers(
-            select(TaskInstance),
-            dag_id=dag_id,
-            run_id=run_id,
-            task_id=task_id,
-            is_mapped=is_mapped,
-            map_indexes=map_indexes,
-            region_id=region_id,
-            region_index=region_index,
-            top_level_only=region_id is None and position is None,
-            loop_pass=_load_loop_pass_regions(position, session=session) if position is not None else None,
-        )
-    ).all()
 
-    selected: dict[int, TaskInstance] = {}
-    for ti in candidates:
-        public_index = ti.region_index if is_mapped else -1
-        if public_index in selected:
+def _check_producers_unambiguous(
+    region_indexes: Iterable[int], *, dag_id: str, run_id: str, task_id: str, is_mapped: bool
+) -> None:
+    seen: set[int] = set()
+    for region_index in region_indexes:
+        public_index = region_index if is_mapped else -1
+        if public_index in seen:
             raise AmbiguousProducerError(
                 f"Multiple live producers for {dag_id}/{run_id}/{task_id} index {public_index}"
             )
-        selected[public_index] = ti
-    return tuple(selected[index] for index in sorted(selected))
+        seen.add(public_index)
+
+
+def resolve_current_producers(
+    *,
+    dag_id: str,
+    run_id: str,
+    task_id: str,
+    is_mapped: bool,
+    context: ProducerContext | None = None,
+    map_indexes: int | Collection[int] | None = None,
+    region_id: UUID | None = None,
+    region_index: int | None = None,
+    session: Session,
+) -> tuple[TaskInstance, ...]:
+    """Resolve the live producer task instances whose data the caller reads by task instance UUID."""
+    from airflow.models.taskinstance import TaskInstance
+
+    _validate_producer_request(context=context, region_id=region_id, region_index=region_index)
+    query = _build_candidate_query(
+        select(TaskInstance),
+        dag_id=dag_id,
+        run_id=run_id,
+        task_id=task_id,
+        is_mapped=is_mapped,
+        context=context,
+        map_indexes=map_indexes,
+        region_id=region_id,
+        region_index=region_index,
+        session=session,
+    )
+    if query is None:
+        return ()
+    candidates = session.scalars(query).all()
+    _check_producers_unambiguous(
+        (ti.region_index for ti in candidates),
+        dag_id=dag_id,
+        run_id=run_id,
+        task_id=task_id,
+        is_mapped=is_mapped,
+    )
+    return tuple(sorted(candidates, key=lambda ti: ti.region_index if is_mapped else -1))
 
 
 def select_current_producer_ids(
@@ -632,14 +672,27 @@ def select_current_producer_ids(
             )
             single_region = not session.scalar(select(other_region.exists()))
         if not single_region:
-            producers = resolve_current_producers(
+            candidate_query = _build_candidate_query(
+                select(TaskInstance.id, TaskInstance.region_index),
                 dag_id=dag_id,
                 run_id=run_id,
                 task_id=task_id,
                 is_mapped=is_mapped,
                 context=context,
                 map_indexes=map_indexes,
+                region_id=None,
+                region_index=None,
                 session=session,
             )
-            return select(TaskInstance.id).where(TaskInstance.id.in_([ti.id for ti in producers]))
+            candidates = session.execute(candidate_query).all() if candidate_query is not None else []
+            _check_producers_unambiguous(
+                (candidate.region_index for candidate in candidates),
+                dag_id=dag_id,
+                run_id=run_id,
+                task_id=task_id,
+                is_mapped=is_mapped,
+            )
+            return select(TaskInstance.id).where(
+                TaskInstance.id.in_([candidate.id for candidate in candidates])
+            )
     return query

airflow-core/src/airflow/models/task_coordinates.py

@@ -598,7 +598,17 @@ def resolve(
                     ).order_by(TaskInstance.region_index)
                 )
             )
-        return resolve_current_producers(**attrs.asdict(request, recurse=False), session=self.session)
+        return resolve_current_producers(
+            dag_id=request.dag_id,
+            run_id=request.run_id,
+            task_id=request.task_id,
+            is_mapped=request.is_mapped,
+            context=request.context,
+            map_indexes=request.map_indexes,
+            region_id=request.region_id,
+            region_index=request.region_index,
+            session=self.session,
+        )
 
     def select_skip_target_ids(
         self, *, caller: TaskInstance, task_id: str, map_indexes: int | None = None
@@ -681,4 +691,14 @@ def select_producer_ids(
                 region_index=region_index,
                 map_indexes=map_indexes,
             )
-        return select_current_producer_ids(**attrs.asdict(request, recurse=False), session=self.session)
+        return select_current_producer_ids(
+            dag_id=request.dag_id,
+            run_id=request.run_id,
+            task_id=request.task_id,
+            is_mapped=request.is_mapped,
+            context=request.context,
+            map_indexes=request.map_indexes,
+            region_id=request.region_id,
+            region_index=request.region_index,
+            session=self.session,
+        )