[v3-3-test] Report Dag cache metrics under each component's own namespace (#71815)

airflow-core/src/airflow/api_fastapi/common/dagbag.py

@@ -16,43 +16,35 @@
 # under the License.
 from __future__ import annotations
 
-import logging
 from typing import TYPE_CHECKING, Annotated
 
 from fastapi import Depends, HTTPException, Request, status
 from sqlalchemy.orm import Session
 
 from airflow.configuration import conf
-from airflow.models.dagbag import DBDagBag
+from airflow.models.dagbag import CachedDBDagBag, DBDagBag
 from airflow.models.serialized_dag import SerializedDagModel
 
 if TYPE_CHECKING:
     from airflow.models.dagrun import DagRun
     from airflow.serialization.definitions.dag import SerializedDAG
 
-log = logging.getLogger(__name__)
 
-
-def create_dag_bag() -> DBDagBag:
+def create_dag_bag() -> CachedDBDagBag:
     """Create DagBag with configurable LRU+TTL caching for API server usage."""
     cache_size = conf.getint("api", "dag_cache_size", fallback=64)
-    cache_ttl_config = conf.getint("api", "dag_cache_ttl", fallback=3600)
+    cache_ttl = conf.getint("api", "dag_cache_ttl", fallback=3600)
 
     if cache_size < 0:
-        log.warning("dag_cache_size must be >= 0, using unbounded dict")
-        cache_size = 0
-    if cache_ttl_config < 0:
-        log.warning("dag_cache_ttl must be >= 0, disabling TTL")
-        cache_ttl_config = 0
-
-    # Use unbounded dict (no eviction) if cache_size is 0
-    if cache_size <= 0:
-        return DBDagBag(cache_size=0)
-
-    # Disable TTL if cache_ttl is 0
-    cache_ttl: int | None = cache_ttl_config if cache_ttl_config > 0 else None
-
-    return DBDagBag(cache_size=cache_size, cache_ttl=cache_ttl)
+        raise ValueError("[api] dag_cache_size must be greater than or equal to 0")
+    if cache_ttl < 0:
+        raise ValueError("[api] dag_cache_ttl must be greater than or equal to 0")
+
+    return CachedDBDagBag(
+        cache_size=cache_size,
+        cache_ttl=cache_ttl,
+        stats_prefix="api_server.dag_bag",
+    )
 
 
 def dag_bag_from_app(request: Request) -> DBDagBag:

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

@@ -98,7 +98,7 @@
 )
 from airflow.models.dag import DagModel
 from airflow.models.dag_version import DagVersion
-from airflow.models.dagbag import DBDagBag
+from airflow.models.dagbag import CachedDBDagBag, DBDagBag
 from airflow.models.dagbundle import DagBundleModel
 from airflow.models.dagrun import DagRun
 from airflow.models.dagwarning import DagWarning, DagWarningType
@@ -360,7 +360,12 @@ def __init__(
         if log:
             self._log = log
 
-        self.scheduler_dag_bag = DBDagBag(load_op_links=False, cache_size=SCHEDULER_DAG_CACHE_SIZE)
+        self.scheduler_dag_bag = CachedDBDagBag(
+            load_op_links=False,
+            cache_size=SCHEDULER_DAG_CACHE_SIZE,
+            cache_ttl=0,
+            stats_prefix="scheduler.dag_bag",
+        )
 
         # Set of (dag_id, asset_name, asset_uri) tuples for trigger policies that
         # are permanently unreachable for the rollup window's cardinality — the

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

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import hashlib
+import math
 import time
 from collections.abc import MutableMapping
 from contextlib import nullcontext
@@ -62,44 +63,34 @@ class DBDagBag:
     """
     Internal class for retrieving dags from the database.
 
-    Optionally supports LRU+TTL caching when cache_size is provided.
-    The scheduler uses this without caching, while the API server can
-    enable caching via configuration.
+    Deserialized Dags are retained in an unbounded dictionary. Use :class:`CachedDBDagBag` when
+    the caller needs configurable eviction, thread safety, and cache metrics.
 
     :meta private:
     """
 
-    def __init__(
-        self,
-        load_op_links: bool = True,
-        cache_size: int | None = None,
-        cache_ttl: int | None = None,
-    ) -> None:
+    def __init__(self, load_op_links: bool = True) -> None:
         """
         Initialize DBDagBag.
 
         :param load_op_links: Should the extra operator link be loaded when de-serializing the DAG?
-        :param cache_size: Size of LRU cache. If None or 0, uses unbounded dict (no eviction).
-        :param cache_ttl: Time-to-live for cache entries in seconds. If None or 0, no TTL (LRU only).
         """
         self.load_op_links = load_op_links
         self._dags: MutableMapping[UUID | str, _CacheEntry] = {}
-        self._use_cache = False
-
         self._revalidation_interval = conf.getint("core", "min_serialized_dag_update_interval")
+        self._lock: RLock | nullcontext = nullcontext()
+
+    def _on_cache_hit(self) -> None:
+        """Handle a Dag cache hit."""
 
-        # Initialize bounded cache if cache_size is provided and > 0
-        if cache_size and cache_size > 0:
-            if cache_ttl and cache_ttl > 0:
-                self._dags = TTLCache(maxsize=cache_size, ttl=cache_ttl)
-            else:
-                self._dags = LRUCache(maxsize=cache_size)
-            self._use_cache = True
+    def _on_cache_miss(self) -> None:
+        """Handle a Dag cache miss."""
 
-        # Lock required for bounded caches: cachetools caches are NOT thread-safe
-        # (LRU reordering and TTL cleanup mutate internal linked lists).
-        # nullcontext for unbounded dict avoids lock overhead in the scheduler path.
-        self._lock: RLock | nullcontext = RLock() if self._use_cache else nullcontext()
+    def _on_cache_clear(self) -> None:
+        """Handle the Dag cache being cleared."""
+
+    def _on_cache_size(self, *, rate: float = 1.0) -> None:
+        """Handle a change in the Dag cache size."""
 
     def _read_dag(self, serdag: SerializedDagModel) -> SerializedDAG | None:
         """Read and cache a SerializedDAG (with its ``dag_hash`` for staleness detection)."""
@@ -109,9 +100,7 @@ def _read_dag(self, serdag: SerializedDagModel) -> SerializedDAG | None:
             return None
         with self._lock:
             self._dags[serdag.dag_version_id] = _CacheEntry(dag, serdag.dag_hash, time.monotonic())
-            cache_size = len(self._dags)
-        if self._use_cache:
-            stats.gauge("api_server.dag_bag.cache_size", cache_size, rate=0.1)
+        self._on_cache_size(rate=0.1)
         return dag
 
     @staticmethod
@@ -133,8 +122,7 @@ def _get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG |
             # [core] min_serialized_dag_update_interval, so an entry validated within that window
             # cannot have gone stale yet -- serve it without touching the DB.
             if now - cached.last_validated < self._revalidation_interval:
-                if self._use_cache:
-                    stats.incr("api_server.dag_bag.cache_hit")
+                self._on_cache_hit()
                 return cached.dag
             # Past the window: a version may have been updated in place (same dag_version_id, new
             # content + new dag_hash) by SerializedDagModel.write_dag, so confirm the cached copy
@@ -148,8 +136,7 @@ def _get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG |
                     current = self._dags.get(version_id)
                     if current is not None and current.dag_hash == cached.dag_hash:
                         self._dags[version_id] = current._replace(last_validated=now)
-                if self._use_cache:
-                    stats.incr("api_server.dag_bag.cache_hit")
+                self._on_cache_hit()
                 return cached.dag
             # Stale (updated in place) or the version no longer exists: drop and reload below.
             with self._lock:
@@ -166,12 +153,11 @@ def _get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG |
         # served without an extra hash check, consistent with the policy above. Only emit the miss
         # metric after confirming no other thread cached it, to avoid counting a single lookup as
         # both a miss and a hit.
-        if self._use_cache:
-            with self._lock:
-                if (cached := self._dags.get(version_id)) is not None:
-                    stats.incr("api_server.dag_bag.cache_hit")
-                    return cached.dag
-            stats.incr("api_server.dag_bag.cache_miss")
+        with self._lock:
+            if (cached := self._dags.get(version_id)) is not None:
+                self._on_cache_hit()
+                return cached.dag
+        self._on_cache_miss()
         return self._read_dag(serdag)
 
     def get_dag(self, version_id: UUID | str, session: Session) -> SerializedDAG | None:
@@ -202,9 +188,8 @@ def clear_cache(self) -> int:
             count = len(self._dags)
             self._dags.clear()
 
-        if self._use_cache:
-            stats.incr("api_server.dag_bag.cache_clear")
-            stats.gauge("api_server.dag_bag.cache_size", 0)
+        self._on_cache_clear()
+        self._on_cache_size()
         return count
 
     @staticmethod
@@ -243,6 +228,56 @@ def get_latest_version_of_dag(self, dag_id: str, *, session: Session) -> Seriali
         return self._read_dag(serdag)
 
 
+class CachedDBDagBag(DBDagBag):
+    """Retrieve Dags through a configurable, thread-safe cache that emits component metrics."""
+
+    def __init__(
+        self,
+        load_op_links: bool = True,
+        *,
+        cache_size: int,
+        cache_ttl: int,
+        stats_prefix: str,
+    ) -> None:
+        """
+        Initialize CachedDBDagBag.
+
+        :param load_op_links: Should the extra operator link be loaded when de-serializing the DAG?
+        :param cache_size: Maximum cached entries. Zero means no size limit.
+        :param cache_ttl: Seconds until a cached entry expires. Zero disables TTL.
+        :param stats_prefix: Metric namespace for this component's cache.
+        :raises ValueError: If the metrics namespace is empty.
+        """
+        if not stats_prefix:
+            raise ValueError("CachedDBDagBag requires a stats_prefix")
+
+        super().__init__(load_op_links=load_op_links)
+
+        if cache_ttl > 0:
+            self._dags = TTLCache(maxsize=cache_size or math.inf, ttl=cache_ttl)
+        elif cache_size > 0:
+            self._dags = LRUCache(maxsize=cache_size)
+
+        # Configured caches are shared across component threads. cachetools caches need this for
+        # linked-list mutations, and the unbounded dict needs it for the double-checked load path.
+        self._lock = RLock()
+        self._stats_prefix = stats_prefix
+
+    def _on_cache_hit(self) -> None:
+        stats.incr(f"{self._stats_prefix}.cache_hit")
+
+    def _on_cache_miss(self) -> None:
+        stats.incr(f"{self._stats_prefix}.cache_miss")
+
+    def _on_cache_clear(self) -> None:
+        stats.incr(f"{self._stats_prefix}.cache_clear")
+
+    def _on_cache_size(self, *, rate: float = 1.0) -> None:
+        with self._lock:
+            size = len(self._dags)
+        stats.gauge(f"{self._stats_prefix}.cache_size", size, rate=rate)
+
+
 def generate_md5_hash(context):
     bundle_name = context.get_current_parameters()["bundle_name"]
     relative_fileloc = context.get_current_parameters()["relative_fileloc"]

airflow-core/tests/unit/api_fastapi/common/test_dagbag.py

@@ -16,14 +16,19 @@
 # under the License.
 from __future__ import annotations
 
+import math
+import re
 from unittest import mock
 
 import pytest
 from cachetools import LRUCache, TTLCache
 
 from airflow.api_fastapi.app import purge_cached_app
+from airflow.api_fastapi.common.dagbag import create_dag_bag
+from airflow.models.dagbag import CachedDBDagBag
 from airflow.sdk import BaseOperator
 
+from tests_common.test_utils.config import conf_vars
 from tests_common.test_utils.db import clear_db_dags, clear_db_runs, clear_db_serialized_dags
 
 pytestmark = pytest.mark.db_test
@@ -49,13 +54,13 @@ def patch_dagbag_once_before_app(self):
         """Patch DagBag once before app is created, and reset counter."""
         self.dagbag_call_counter["count"] = 0
 
-        from airflow.models.dagbag import DBDagBag as RealDagBag
+        from airflow.models.dagbag import CachedDBDagBag as RealDagBag
 
         def factory(*args, **kwargs):
             self.dagbag_call_counter["count"] += 1
             return RealDagBag(*args, **kwargs)
 
-        with mock.patch("airflow.api_fastapi.common.dagbag.DBDagBag", side_effect=factory):
+        with mock.patch("airflow.api_fastapi.common.dagbag.CachedDBDagBag", side_effect=factory):
             purge_cached_app()
             yield
 
@@ -89,24 +94,48 @@ class TestCreateDagBag:
     """Tests for create_dag_bag() function."""
 
     @pytest.mark.parametrize(
-        ("cache_size", "cache_ttl", "expected_use_cache", "expected_dags_type"),
+        ("cache_size", "cache_ttl", "expected_bag_type", "expected_dags_type", "expected_maxsize"),
         [
-            pytest.param(64, 3600, True, TTLCache, id="default_ttl_cache"),
-            pytest.param(0, 3600, False, dict, id="size_zero_unbounded"),
-            pytest.param(64, 0, True, LRUCache, id="ttl_zero_lru_only"),
+            pytest.param("64", "3600", CachedDBDagBag, TTLCache, 64, id="default_ttl_cache"),
+            pytest.param("0", "3600", CachedDBDagBag, TTLCache, math.inf, id="size_zero_ttl_only"),
+            pytest.param("64", "0", CachedDBDagBag, LRUCache, 64, id="ttl_zero_lru_only"),
+            pytest.param("0", "0", CachedDBDagBag, dict, None, id="both_zero_no_eviction"),
         ],
     )
-    @mock.patch("airflow.api_fastapi.common.dagbag.conf")
     def test_create_dag_bag_cache_modes(
-        self, mock_conf, cache_size, cache_ttl, expected_use_cache, expected_dags_type
+        self,
+        cache_size,
+        cache_ttl,
+        expected_bag_type,
+        expected_dags_type,
+        expected_maxsize,
     ):
-        from airflow.api_fastapi.common.dagbag import create_dag_bag
+        with conf_vars({("api", "dag_cache_size"): cache_size, ("api", "dag_cache_ttl"): cache_ttl}):
+            dag_bag = create_dag_bag()
 
-        mock_conf.getint.side_effect = lambda section, key, fallback: {
-            "dag_cache_size": cache_size,
-            "dag_cache_ttl": cache_ttl,
-        }.get(key, fallback)
-
-        dag_bag = create_dag_bag()
-        assert dag_bag._use_cache is expected_use_cache
+        assert type(dag_bag) is expected_bag_type
         assert isinstance(dag_bag._dags, expected_dags_type)
+        if expected_maxsize is not None:
+            assert dag_bag._dags.maxsize == expected_maxsize
+
+    @pytest.mark.parametrize(
+        ("cache_size", "cache_ttl", "expected_message"),
+        [
+            pytest.param(
+                "-1",
+                "3600",
+                "[api] dag_cache_size must be greater than or equal to 0",
+                id="negative_size",
+            ),
+            pytest.param(
+                "64",
+                "-1",
+                "[api] dag_cache_ttl must be greater than or equal to 0",
+                id="negative_ttl",
+            ),
+        ],
+    )
+    def test_create_dag_bag_rejects_negative_config(self, cache_size, cache_ttl, expected_message):
+        with conf_vars({("api", "dag_cache_size"): cache_size, ("api", "dag_cache_ttl"): cache_ttl}):
+            with pytest.raises(ValueError, match=re.escape(expected_message)):
+                create_dag_bag()

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

@@ -77,6 +77,7 @@
 )
 from airflow.models.dag import DagModel, get_last_dagrun, infer_automated_data_interval
 from airflow.models.dag_version import DagVersion
+from airflow.models.dagbag import CachedDBDagBag
 from airflow.models.dagbundle import DagBundleModel
 from airflow.models.dagrun import DagRun
 from airflow.models.dagwarning import DagWarning
@@ -416,8 +417,11 @@ def test_scheduler_dag_bag_is_bounded(self):
 
         job_runner = SchedulerJobRunner(Job())
 
+        assert isinstance(job_runner.scheduler_dag_bag, CachedDBDagBag)
         assert isinstance(job_runner.scheduler_dag_bag._dags, LRUCache)
         assert job_runner.scheduler_dag_bag._dags.maxsize == SCHEDULER_DAG_CACHE_SIZE
+        # Reported separately from the API server's cache, not folded into it.
+        assert job_runner.scheduler_dag_bag._stats_prefix == "scheduler.dag_bag"
 
     @pytest.mark.parametrize(
         "heartrate",