Report Dag cache metrics under each component's own namespace

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

@@ -22,15 +22,15 @@
 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
 
 
-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 = conf.getint("api", "dag_cache_ttl", fallback=3600)
@@ -40,7 +40,11 @@ def create_dag_bag() -> DBDagBag:
     if cache_ttl < 0:
         raise ValueError("[api] dag_cache_ttl must be greater than or equal to 0")
 
-    return DBDagBag(cache_size=cache_size, cache_ttl=cache_ttl)
+    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

@@ -99,7 +99,7 @@
 )
 from airflow.models.dag import DagModel
 from airflow.models.dag_version import DagVersion, _resolve_version_data
-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
@@ -382,7 +382,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

@@ -63,55 +63,34 @@ class DBDagBag:
     """
     Internal class for retrieving dags from the database.
 
-    Optionally caches deserialized dags: a size limit enables LRU eviction, and a TTL enables
-    age-based eviction with or without a size limit. Callers that pass neither get a plain dict
-    that never evicts.
+    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: Max cached entries. 0 or None means no size limit.
-        :param cache_ttl: Seconds until a cached entry expires, applied with or without a size limit.
-            0 or None disables TTL. With neither a size limit nor a TTL the cache never evicts.
-        :raises ValueError: If ``cache_size`` or ``cache_ttl`` is negative.
         """
-        # Callers should reject negative values with their own context; validate again defensively.
-        if cache_size is not None and cache_size < 0:
-            raise ValueError("cache_size must be greater than or equal to 0")
-        if cache_ttl is not None and cache_ttl < 0:
-            raise ValueError("cache_ttl must be greater than or equal to 0")
-
         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."""
+
+    def _on_cache_miss(self) -> None:
+        """Handle a Dag cache miss."""
 
-        # A TTL applies with or without a size limit: an uncapped TTLCache is what lets
-        # ``dag_cache_size = 0`` mean "no size limit" rather than "no eviction at all".
-        size = cache_size or 0
-        ttl = cache_ttl or 0
-        if ttl > 0:
-            self._dags = TTLCache(maxsize=size or math.inf, ttl=ttl)
-            self._use_cache = True
-        elif size > 0:
-            self._dags = LRUCache(maxsize=size)
-            self._use_cache = True
-
-        # Lock required for bounded caches: cachetools caches are NOT thread-safe
-        # (LRU reordering and TTL cleanup mutate internal linked lists). A plain dict needs no
-        # lock, so it uses nullcontext.
-        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)."""
@@ -121,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
@@ -145,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
@@ -160,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:
@@ -178,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:
@@ -214,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
@@ -258,6 +231,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

@@ -25,6 +25,7 @@
 
 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
@@ -53,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
 
@@ -93,20 +94,27 @@ class TestCreateDagBag:
     """Tests for create_dag_bag() function."""
 
     @pytest.mark.parametrize(
-        ("cache_size", "cache_ttl", "expected_dags_type", "expected_maxsize"),
+        ("cache_size", "cache_ttl", "expected_bag_type", "expected_dags_type", "expected_maxsize"),
         [
-            pytest.param("64", "3600", TTLCache, 64, id="default_ttl_cache"),
-            pytest.param("0", "3600", TTLCache, math.inf, id="size_zero_ttl_only"),
-            pytest.param("64", "0", LRUCache, 64, id="ttl_zero_lru_only"),
-            pytest.param("0", "0", dict, None, id="both_zero_no_eviction"),
+            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"),
         ],
     )
-    def test_create_dag_bag_cache_modes(self, cache_size, cache_ttl, expected_dags_type, expected_maxsize):
+    def test_create_dag_bag_cache_modes(
+        self,
+        cache_size,
+        cache_ttl,
+        expected_bag_type,
+        expected_dags_type,
+        expected_maxsize,
+    ):
         with conf_vars({("api", "dag_cache_size"): cache_size, ("api", "dag_cache_ttl"): cache_ttl}):
             dag_bag = create_dag_bag()
 
+        assert type(dag_bag) is expected_bag_type
         assert isinstance(dag_bag._dags, expected_dags_type)
-        assert dag_bag._use_cache is (expected_dags_type is not dict)
         if expected_maxsize is not None:
             assert dag_bag._dags.maxsize == expected_maxsize
 

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

@@ -84,6 +84,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
@@ -420,8 +421,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",