diff --git a/airflow-core/src/airflow/models/dagbag.py b/airflow-core/src/airflow/models/dagbag.py index c4bd8eceea102..d36dbe5d3824f 100644 --- a/airflow-core/src/airflow/models/dagbag.py +++ b/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,9 +63,9 @@ 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. + Optionally caches deserialized dags. A size cap enables LRU eviction and a TTL + enables age-based eviction, with or without a size cap. The scheduler uses this + without caching, while the API server can enable caching via configuration. :meta private: """ @@ -79,8 +80,9 @@ def __init__( 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). + :param cache_size: Max cached entries; 0 or None means no size cap. + :param cache_ttl: Seconds until a cached entry expires. If > 0, entries are evicted by + age regardless of ``cache_size`` (with no size cap this gives TTL-only eviction). """ self.load_op_links = load_op_links self._dags: MutableMapping[UUID | str, _CacheEntry] = {} @@ -88,12 +90,13 @@ def __init__( self._revalidation_interval = conf.getint("core", "min_serialized_dag_update_interval") - # 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) + # Initialize a TTL cache if configured, otherwise a bounded LRU cache. + if cache_ttl and cache_ttl > 0: + maxsize = cache_size if cache_size and cache_size > 0 else math.inf + self._dags = TTLCache(maxsize=maxsize, ttl=cache_ttl) + self._use_cache = True + elif cache_size and cache_size > 0: + self._dags = LRUCache(maxsize=cache_size) self._use_cache = True # Lock required for bounded caches: cachetools caches are NOT thread-safe diff --git a/airflow-core/tests/unit/models/test_dagbag.py b/airflow-core/tests/unit/models/test_dagbag.py index 79668d4fe54f4..770cc3b6980d4 100644 --- a/airflow-core/tests/unit/models/test_dagbag.py +++ b/airflow-core/tests/unit/models/test_dagbag.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import math import time from concurrent.futures import ThreadPoolExecutor from unittest.mock import MagicMock, patch @@ -259,14 +260,24 @@ def test_lru_cache_enabled_with_cache_size(self): assert isinstance(dag_bag._dags, LRUCache) def test_ttl_cache_enabled_with_cache_size_and_ttl(self): - """Test that TTL cache is enabled when both cache_size and cache_ttl are provided.""" + """Test that a bounded TTL cache is used when both cache_size and cache_ttl are provided.""" dag_bag = DBDagBag(cache_size=10, cache_ttl=60) assert dag_bag._use_cache is True assert isinstance(dag_bag._dags, TTLCache) + assert dag_bag._dags.maxsize == 10 - def test_zero_cache_size_uses_unbounded_dict(self): - """Test that cache_size=0 uses unbounded dict (same as no caching).""" - dag_bag = DBDagBag(cache_size=0, cache_ttl=60) + @pytest.mark.parametrize("cache_size", [0, None]) + def test_ttl_only_without_size_cap(self, cache_size): + """Test that a positive cache_ttl with no size cap gives a TTL cache with maxsize=inf.""" + dag_bag = DBDagBag(cache_size=cache_size, cache_ttl=60) + assert dag_bag._use_cache is True + assert isinstance(dag_bag._dags, TTLCache) + assert dag_bag._dags.maxsize == math.inf + + @pytest.mark.parametrize("cache_ttl", [None, 0]) + def test_zero_cache_size_uses_unbounded_dict(self, cache_ttl): + """Test that cache_size=0 without a TTL uses an unbounded dict (same as no caching).""" + dag_bag = DBDagBag(cache_size=0, cache_ttl=cache_ttl) assert dag_bag._use_cache is False assert isinstance(dag_bag._dags, dict) @@ -310,6 +321,21 @@ def test_ttl_cache_expiry(self): with time_machine.travel("2025-01-01 00:00:02", tick=False): assert dag_bag._dags.get("test_version_id") is None + def test_ttl_only_evicts_by_ttl_not_size(self): + """An unbounded (maxsize=inf) TTL cache keeps every entry until it expires by age.""" + dag_bag = DBDagBag(cache_size=0, cache_ttl=1) + assert dag_bag._dags.maxsize == math.inf + dag_bag._dags = TTLCache(maxsize=math.inf, ttl=1, timer=time.time) + + with time_machine.travel("2025-01-01 00:00:00", tick=False): + for i in range(500): + dag_bag._dags[f"version_{i}"] = MagicMock() + assert len(dag_bag._dags) == 500 + + with time_machine.travel("2025-01-01 00:00:02", tick=False): + assert dag_bag._dags.get("version_0") is None + assert len(dag_bag._dags) == 0 + def test_lru_eviction(self): """Test that LRU eviction works when cache is full.""" dag_bag = DBDagBag(cache_size=2)