Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 14 additions & 11 deletions airflow-core/src/airflow/models/dagbag.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from __future__ import annotations

import hashlib
import math
import time
from collections.abc import MutableMapping
from contextlib import nullcontext
Expand Down Expand Up @@ -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:
"""
Expand All @@ -79,21 +80,23 @@ 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] = {}
self._use_cache = False

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
Expand Down
34 changes: 30 additions & 4 deletions airflow-core/tests/unit/models/test_dagbag.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand Down