Count post-deferral scheduled tasks toward include_deferred pool slots

airflow-core/newsfragments/69187.bugfix.rst

@@ -0,0 +1 @@
+Fix pool slot undercounting for tasks that return to ``scheduled`` after deferral when ``include_deferred`` is enabled. Those task instances keep occupying a pool slot (identified via ``next_method``) until they are queued or running again, so pools continue to protect external resources across the deferral boundary.

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

@@ -21,7 +21,7 @@
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, TypedDict
 
-from sqlalchemy import Boolean, ForeignKey, Integer, String, Text, func, select
+from sqlalchemy import Boolean, ForeignKey, Integer, String, Text, and_, func, or_, select
 from sqlalchemy.orm import Mapped, mapped_column
 
 from airflow._shared.observability.metrics.stats import normalize_name_for_stats
@@ -36,6 +36,8 @@
     from sqlalchemy.orm.session import Session
     from sqlalchemy.sql import Select
 
+    from airflow.models.taskinstance import TaskInstance as TaskInstanceModel
+
 logger = logging.getLogger(__name__)
 
 
@@ -242,11 +244,32 @@ def slots_stats(
             else:
                 raise AirflowException(f"Unexpected state. Expected values: {allowed_execution_states}.")
 
+        # Scheduled TIs that are resuming after deferral still hold a pool slot when
+        # include_deferred is enabled (see get_occupied_states / occupied_slots).
+        # Distinguish them via next_method so first-lifetime scheduled TIs are not counted
+        # (that would deadlock scheduling when open = total - scheduled).
+        post_deferral_scheduled_by_pool: dict[str, int] = {}
+        if any(pool_includes_deferred.values()):
+            post_deferral_rows = session.execute(
+                select(TaskInstance.pool, func.sum(TaskInstance.pool_slots))
+                .where(
+                    TaskInstance.state == TaskInstanceState.SCHEDULED,
+                    TaskInstance.next_method.is_not(None),
+                )
+                .group_by(TaskInstance.pool)
+            )
+            post_deferral_scheduled_by_pool = {
+                pool_name: int(decimal_count)
+                for pool_name, decimal_count in post_deferral_rows
+                if pool_name is not None
+            }
+
         # calculate open metric
         for pool_name, stats_dict in pools.items():
             stats_dict["open"] = stats_dict["total"] - stats_dict["running"] - stats_dict["queued"]
             if pool_includes_deferred[pool_name]:
                 stats_dict["open"] -= stats_dict["deferred"]
+                stats_dict["open"] -= post_deferral_scheduled_by_pool.get(pool_name, 0)
 
         return pools
 
@@ -270,29 +293,61 @@ def occupied_slots(self, *, session: Session = NEW_SESSION) -> int:
         """
         Get the number of slots used by running/queued tasks at the moment.
 
+        When ``include_deferred`` is enabled, deferred tasks and scheduled tasks that are
+        resuming after deferral (``next_method`` is set) also count as occupied.
+
         :param session: SQLAlchemy ORM Session
         :return: the used number of slots
         """
         from airflow.models.taskinstance import TaskInstance  # Avoid circular import
 
         occupied_states = self.get_occupied_states()
+        occupancy_filter = TaskInstance.state.in_(occupied_states)
+        if self.include_deferred:
+            # Post-deferral scheduled TIs keep holding the slot until they are queued/running.
+            occupancy_filter = or_(
+                occupancy_filter,
+                and_(
+                    TaskInstance.state == TaskInstanceState.SCHEDULED,
+                    TaskInstance.next_method.is_not(None),
+                ),
+            )
 
         return int(
             session.scalar(
                 select(func.sum(TaskInstance.pool_slots))
                 .filter(TaskInstance.pool == self.pool)
-                .filter(TaskInstance.state.in_(occupied_states))
+                .filter(occupancy_filter)
             )
             or 0
         )
 
     def get_occupied_states(self):
+        """
+        Return states that always occupy pool slots for this pool.
+
+        Post-deferral ``scheduled`` occupancy (``next_method`` set) is handled separately in
+        :meth:`occupied_slots` and :meth:`task_instance_occupies_slot` because it cannot be
+        expressed as a state-only set without also counting first-lifetime scheduled tasks.
+        """
         if self.include_deferred:
             return EXECUTION_STATES | {
                 TaskInstanceState.DEFERRED,
             }
         return EXECUTION_STATES
 
+    def task_instance_occupies_slot(self, ti: TaskInstanceModel) -> bool:
+        """
+        Return whether ``ti`` currently occupies a slot in this pool.
+
+        Used by pool availability deps so a TI that already holds a slot does not block itself.
+        """
+        if ti.state in self.get_occupied_states():
+            return True
+        return bool(
+            self.include_deferred and ti.state == TaskInstanceState.SCHEDULED and ti.next_method is not None
+        )
+
     @provide_session
     def running_slots(self, *, session: Session = NEW_SESSION) -> int:
         """

airflow-core/src/airflow/ti_deps/deps/pool_slots_available_dep.py

@@ -71,7 +71,7 @@ def _get_dep_statuses(self, ti, dep_context=None, *, session):
                 return
 
         open_slots = pool.open_slots(session=session)
-        if ti.state in pool.get_occupied_states():
+        if pool.task_instance_occupies_slot(ti):
             open_slots += ti.pool_slots
 
         if open_slots <= (ti.pool_slots - 1):

airflow-core/tests/unit/models/test_pool.py

@@ -178,6 +178,94 @@ def test_open_slots_including_deferred(self, dag_maker):
             },
         }
 
+    def test_open_slots_including_deferred_counts_post_deferral_scheduled(self, dag_maker):
+        """Scheduled TIs resuming after deferral (next_method set) occupy slots when include_deferred."""
+        pool = Pool(pool="test_pool", slots=5, include_deferred=True)
+        with dag_maker(
+            dag_id="test_open_slots_post_deferral_scheduled",
+            start_date=DEFAULT_DATE,
+        ):
+            op1 = EmptyOperator(task_id="dummy1", pool="test_pool")
+            op2 = EmptyOperator(task_id="dummy2", pool="test_pool")
+            op3 = EmptyOperator(task_id="dummy3", pool="test_pool")
+
+        dr = dag_maker.create_dagrun()
+
+        ti1 = dr.get_task_instance(task_id=op1.task_id)
+        ti2 = dr.get_task_instance(task_id=op2.task_id)
+        ti3 = dr.get_task_instance(task_id=op3.task_id)
+        ti1.state = State.RUNNING
+        # Post-deferral resume: scheduled with next_method set.
+        ti2.state = State.SCHEDULED
+        ti2.next_method = "execute_complete"
+        # First-lifetime scheduled: must not occupy a slot.
+        ti3.state = State.SCHEDULED
+        ti3.next_method = None
+
+        session = settings.Session()
+        session.add(pool)
+        session.merge(ti1)
+        session.merge(ti2)
+        session.merge(ti3)
+        session.commit()
+        session.close()
+
+        assert pool.open_slots() == 3
+        assert pool.running_slots() == 1
+        assert pool.scheduled_slots() == 2
+        assert pool.occupied_slots() == 2
+        assert pool.task_instance_occupies_slot(ti2) is True
+        assert pool.task_instance_occupies_slot(ti3) is False
+        assert pool.slots_stats() == {
+            "default_pool": {
+                "open": 128,
+                "queued": 0,
+                "total": 128,
+                "running": 0,
+                "scheduled": 0,
+                "deferred": 0,
+            },
+            "test_pool": {
+                "open": 3,
+                "queued": 0,
+                "running": 1,
+                "deferred": 0,
+                "scheduled": 2,
+                "total": 5,
+            },
+        }
+
+    def test_open_slots_excluding_deferred_ignores_post_deferral_scheduled(self, dag_maker):
+        """Without include_deferred, post-deferral scheduled TIs do not occupy slots."""
+        pool = Pool(pool="test_pool", slots=5, include_deferred=False)
+        with dag_maker(
+            dag_id="test_open_slots_exclude_post_deferral_scheduled",
+            start_date=DEFAULT_DATE,
+        ):
+            op1 = EmptyOperator(task_id="dummy1", pool="test_pool")
+            op2 = EmptyOperator(task_id="dummy2", pool="test_pool")
+
+        dr = dag_maker.create_dagrun()
+
+        ti1 = dr.get_task_instance(task_id=op1.task_id)
+        ti2 = dr.get_task_instance(task_id=op2.task_id)
+        ti1.state = State.RUNNING
+        ti2.state = State.SCHEDULED
+        ti2.next_method = "execute_complete"
+
+        session = settings.Session()
+        session.add(pool)
+        session.merge(ti1)
+        session.merge(ti2)
+        session.commit()
+        session.close()
+
+        assert pool.open_slots() == 4
+        assert pool.occupied_slots() == 1
+        assert pool.task_instance_occupies_slot(ti2) is False
+        assert pool.slots_stats()["test_pool"]["open"] == 4
+        assert pool.slots_stats()["test_pool"]["scheduled"] == 1
+
     def test_infinite_slots(self, dag_maker):
         pool = Pool(pool="test_pool", slots=-1, include_deferred=False)
         with dag_maker(

airflow-core/tests/unit/ti_deps/deps/test_pool_slots_available_dep.py

@@ -70,6 +70,33 @@ def test_deferred_pooled_task_pass(self, mock_open_slots):
         ti_to_fail = Mock(pool="test_pool", state=TaskInstanceState.DEFERRED, pool_slots=1)
         assert not PoolSlotsAvailableDep().is_met(ti=ti_to_fail)
 
+    @patch("airflow.models.Pool.open_slots", return_value=0)
+    def test_post_deferral_scheduled_pooled_task_pass(self, mock_open_slots):
+        """Scheduled TIs resuming after deferral already hold a slot when include_deferred is on."""
+        ti = Mock(
+            pool="test_includes_deferred_pool",
+            state=TaskInstanceState.SCHEDULED,
+            pool_slots=1,
+            next_method="execute_complete",
+        )
+        assert PoolSlotsAvailableDep().is_met(ti=ti)
+        # First-lifetime scheduled (no next_method) must still require an open slot.
+        ti_first_schedule = Mock(
+            pool="test_includes_deferred_pool",
+            state=TaskInstanceState.SCHEDULED,
+            pool_slots=1,
+            next_method=None,
+        )
+        assert not PoolSlotsAvailableDep().is_met(ti=ti_first_schedule)
+        # include_deferred=False never treats scheduled as occupied.
+        ti_excluded = Mock(
+            pool="test_pool",
+            state=TaskInstanceState.SCHEDULED,
+            pool_slots=1,
+            next_method="execute_complete",
+        )
+        assert not PoolSlotsAvailableDep().is_met(ti=ti_excluded)
+
     def test_task_with_nonexistent_pool(self):
         ti = Mock(pool="nonexistent_pool", pool_slots=1)
         assert not PoolSlotsAvailableDep().is_met(ti=ti)