Feature/add max new dagruns to schedule

airflow-core/newsfragments/64294.feature.rst

@@ -0,0 +1 @@
+Add a new ``max_new_dagruns_per_loop_to_schedule`` configuration to control how many new dagruns are scheduled each scheduler iteration

airflow-core/newsfragments/64294.significant.rst

@@ -0,0 +1 @@
+``get_running_dag_runs_to_examine`` now returns a ``Sequence[DagRun]`` type instead of ``ScalarResult[Dagrun]``

airflow-core/src/airflow/config_templates/config.yml

@@ -2688,6 +2688,19 @@ scheduler:
       type: integer
       default: "20"
       see_also: ":ref:`scheduler:ha:tunables`"
+    max_new_dagruns_per_loop_to_schedule:
+      description: |
+        When set > 0, the scheduler runs a second query per loop that fetches up to this
+        many dagruns that have never been examined (``last_scheduling_decision IS NULL``),
+        in addition to the ``max_dagruns_per_loop_to_schedule`` already-examined ones.
+
+        This prevents starvation of older dagruns when large batches of new dagruns are
+        created at once (for example, via ``TriggerDagRunOperator``). Note that the total
+        number of dagruns locked and scheduled per loop becomes the sum of both limits.
+      example: ~
+      version_added: 3.3.0
+      type: integer
+      default: "0"
     use_job_schedule:
       description: |
         Turn off scheduler use of cron intervals by setting this to ``False``.

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

@@ -40,6 +40,7 @@
     Index,
     Integer,
     PrimaryKeyConstraint,
+    SQLColumnExpression,
     String,
     Text,
     UniqueConstraint,
@@ -57,7 +58,15 @@
 from sqlalchemy.ext.associationproxy import association_proxy
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.ext.mutable import MutableDict
-from sqlalchemy.orm import Mapped, declared_attr, joinedload, mapped_column, relationship, synonym, validates
+from sqlalchemy.orm import (
+    Mapped,
+    declared_attr,
+    joinedload,
+    mapped_column,
+    relationship,
+    synonym,
+    validates,
+)
 from sqlalchemy.orm.exc import StaleDataError
 from sqlalchemy.sql.expression import false, select
 from sqlalchemy.sql.functions import coalesce
@@ -319,6 +328,11 @@ class DagRun(Base, LoggingMixin):
         "max_dagruns_per_loop_to_schedule",
         fallback=20,
     )
+    DEFAULT_NEW_DAGRUNS_TO_EXAMINE = airflow_conf.getint(
+        "scheduler",
+        "max_new_dagruns_per_loop_to_schedule",
+        fallback=0,
+    )
     _ti_dag_versions = association_proxy("task_instances", "dag_version")
     _tih_dag_versions = association_proxy("task_instances_histories", "dag_version")
 
@@ -623,40 +637,77 @@ def active_runs_of_dags(
 
     @classmethod
     @retry_db_transaction
-    def get_running_dag_runs_to_examine(cls, session: Session) -> ScalarResult[DagRun]:
+    def get_running_dag_runs_to_examine(cls, session: Session) -> Sequence[DagRun]:
         """
-        Return the next DagRuns that the scheduler should attempt to schedule.
+        Return the next DagRuns (as a list) that the scheduler should attempt to schedule.
 
-        This will return zero or more DagRun rows that are row-level-locked with a "SELECT ... FOR UPDATE"
+        This will return zero or more DagRuns that are row-level-locked with a "SELECT ... FOR UPDATE"
         query, you should ensure that any scheduling decisions are made in a single transaction -- as soon as
         the transaction is committed it will be unlocked.
+        With max_new_dagruns_per_loop_to_schedule > 0, this runs 2 queries, one for new dagruns (where
+        last_scheduling_decision is None) and for old (already examined) dagruns, cleared / requeued dagruns
+        will appear in the new dagruns query.
 
         :meta private:
         """
         from airflow.models.backfill import BackfillDagRun
         from airflow.models.dag import DagModel
 
-        query = (
-            select(cls)
-            .with_hint(cls, "USE INDEX (idx_dag_run_running_dags)", dialect_name="mysql")
-            .where(cls.state == DagRunState.RUNNING)
-            .join(DagModel, DagModel.dag_id == cls.dag_id)
-            .join(BackfillDagRun, BackfillDagRun.dag_run_id == DagRun.id, isouter=True)
-            .where(
-                DagModel.is_paused == false(),
-                DagModel.is_stale == false(),
-            )
-            .order_by(
-                nulls_first(cast("ColumnElement[Any]", BackfillDagRun.sort_ordinal), session=session),
-                nulls_first(cast("ColumnElement[Any]", cls.last_scheduling_decision), session=session),
-                cls.run_after,
+        def _get_dagrun_query(
+            filters: list[ColumnElement[bool]], order_by: list[SQLColumnExpression[Any]], limit: int
+        ):
+            return (
+                select(DagRun)
+                .with_hint(DagRun, "USE INDEX (idx_dag_run_running_dags)", dialect_name="mysql")
+                .where(DagRun.state == DagRunState.RUNNING)
+                .join(DagModel, DagModel.dag_id == cls.dag_id)
+                .join(BackfillDagRun, BackfillDagRun.dag_run_id == DagRun.id, isouter=True)
+                .where(*filters)
+                .order_by(*order_by)
+                .limit(limit)
             )
-            .limit(cls.DEFAULT_DAGRUNS_TO_EXAMINE)
+
+        filters = [
+            DagRun.run_after <= func.now(),
+            DagModel.is_paused == false(),
+            DagModel.is_stale == false(),
+        ]
+
+        order = [
+            nulls_first(cast("ColumnElement[Any]", BackfillDagRun.sort_ordinal), session=session),
+            nulls_first(cast("ColumnElement[Any]", DagRun.last_scheduling_decision), session=session),
+            DagRun.run_after,
+        ]
+
+        new_dagruns_to_examine = max(cls.DEFAULT_NEW_DAGRUNS_TO_EXAMINE, 0)
+        dagruns_to_examine = cls.DEFAULT_DAGRUNS_TO_EXAMINE
+
+        old_filters = (
+            [*filters, DagRun.last_scheduling_decision.is_not(None)]
+            if new_dagruns_to_examine > 0
+            else filters
         )
+        query = _get_dagrun_query(filters=old_filters, order_by=order, limit=dagruns_to_examine)
 
-        query = query.where(DagRun.run_after <= func.now())
+        result: list[DagRun] = cast(
+            "list[DagRun]",
+            session.scalars(with_row_locks(query, of=cls, session=session, skip_locked=True)).unique().all(),
+        )
+
+        if new_dagruns_to_examine > 0:
+            new_dagruns_query = _get_dagrun_query(
+                filters=[*filters, DagRun.last_scheduling_decision.is_(None)],
+                order_by=order,
+                limit=new_dagruns_to_examine,
+            )
+            new_dagruns: Sequence[DagRun] = (
+                session.scalars(with_row_locks(new_dagruns_query, of=cls, session=session, skip_locked=True))
+                .unique()
+                .all()
+            )
+
+            result.extend(new_dagruns)
 
-        result = session.scalars(with_row_locks(query, of=cls, session=session, skip_locked=True)).unique()
         return result
 
     @classmethod