Merge pull request #950 from boffin-dmytro/fix/dir-size-gb-ttl-caching

fix(helpers): add TTL cache to dir_size_gb to prevent blocking rglob walks

Author: Lightheartdevs

dream-server/extensions/services/dashboard-api/helpers.py

@@ -17,6 +17,31 @@
 from config import SERVICES, INSTALL_DIR, DATA_DIR, LLM_BACKEND
 from models import ServiceStatus, DiskUsage, ModelInfo, BootstrapStatus
 
+
+class _DirSizeCache:
+    """Per-path TTL cache for dir_size_gb to avoid repeated rglob walks."""
+
+    def __init__(self, ttl: float = 60.0):
+        self._ttl = ttl
+        self._store: dict[str, tuple[float, float]] = {}
+
+    def get(self, path: Path) -> float | None:
+        key = str(path.resolve())
+        entry = self._store.get(key)
+        if entry is None:
+            return None
+        expires_at, value = entry
+        if time.monotonic() > expires_at:
+            del self._store[key]
+            return None
+        return value
+
+    def set(self, path: Path, value: float):
+        self._store[str(path.resolve())] = (time.monotonic() + self._ttl, value)
+
+
+_dir_size_cache = _DirSizeCache()
+
 # Lemonade serves at /api/v1 instead of llama.cpp's /v1
 _LLM_API_PREFIX = "/api/v1" if LLM_BACKEND == "lemonade" else "/v1"
 
@@ -282,8 +307,13 @@ def dir_size_gb(path: Path) -> float:
     """Calculate total size of a directory in GB. Returns 0.0 if path doesn't exist.
 
     Skips symlinks to avoid following links outside DATA_DIR and double-counting.
+    Results are cached for 60 seconds to avoid repeated expensive rglob walks.
     """
+    cached = _dir_size_cache.get(path)
+    if cached is not None:
+        return cached
     if not path.exists():
+        _dir_size_cache.set(path, 0.0)
         return 0.0
     total = 0
     try:
@@ -295,7 +325,19 @@ def dir_size_gb(path: Path) -> float:
                 pass
     except (PermissionError, OSError):
         pass
-    return round(total / (1024**3), 2)
+    result = round(total / (1024**3), 2)
+    _dir_size_cache.set(path, result)
+    return result
+
+
+def invalidate_dir_size_cache(path: Path):
+    """Remove cached size for a specific path after it has been modified."""
+    _dir_size_cache._store.pop(str(path.resolve()), None)
+
+
+def clear_dir_size_cache():
+    """Clear the entire dir_size_gb cache (e.g. after bulk operations)."""
+    _dir_size_cache._store.clear()
 
 
 def get_disk_usage() -> DiskUsage:

dream-server/extensions/services/dashboard-api/routers/extensions.py

@@ -1411,10 +1411,11 @@ def purge_extension_data(
         if not body.confirm:
             raise HTTPException(status_code=400, detail="Confirmation required: set confirm=true")
 
-        from helpers import dir_size_gb  # noqa: PLC0415
+        from helpers import dir_size_gb, invalidate_dir_size_cache  # noqa: PLC0415
         size_gb = dir_size_gb(data_path)
 
         shutil.rmtree(data_path, ignore_errors=True)
+        invalidate_dir_size_cache(data_path)
 
         if data_path.exists():
             raise HTTPException(status_code=500, detail=f"Could not fully remove data/{service_id}. Some files may be owned by root.")

dream-server/extensions/services/dashboard-api/tests/test_helpers.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import json
+from pathlib import Path
 from unittest.mock import AsyncMock, MagicMock
 
 import aiohttp
@@ -13,7 +14,7 @@
     get_uptime, get_cpu_metrics, get_ram_metrics,
     check_service_health, get_all_services,
     get_llama_metrics, get_loaded_model, get_llama_context_size,
-    get_disk_usage, dir_size_gb,
+    get_disk_usage, dir_size_gb, invalidate_dir_size_cache, clear_dir_size_cache,
     _get_aio_session, set_services_cache, get_cached_services,
     _get_lifetime_tokens,
 )
@@ -821,14 +822,17 @@ def test_invalid_eta_string(self, data_dir):
 class TestDirSizeGb:
 
     def test_nonexistent_path_returns_zero(self, tmp_path):
+        clear_dir_size_cache()
         assert dir_size_gb(tmp_path / "does-not-exist") == 0.0
 
     def test_empty_directory_returns_zero(self, tmp_path):
+        clear_dir_size_cache()
         empty = tmp_path / "empty"
         empty.mkdir()
         assert dir_size_gb(empty) == 0.0
 
     def test_directory_with_files(self, tmp_path):
+        clear_dir_size_cache()
         d = tmp_path / "data"
         d.mkdir()
         # Write 100 MiB (avoids allocating 1 GiB in CI)
@@ -837,6 +841,7 @@ def test_directory_with_files(self, tmp_path):
         assert dir_size_gb(d) == 0.1
 
     def test_symlinks_are_skipped(self, tmp_path):
+        clear_dir_size_cache()
         d = tmp_path / "withlinks"
         d.mkdir()
         real = d / "real.bin"
@@ -846,3 +851,40 @@ def test_symlinks_are_skipped(self, tmp_path):
         # Only real.bin should be counted (1024 B ≈ 0.0 GB when rounded to 2dp)
         result = dir_size_gb(d)
         assert result == 0.0  # 1024 bytes rounds to 0.0 GB
+
+    def test_uses_cached_value_until_invalidated(self, tmp_path, monkeypatch):
+        clear_dir_size_cache()
+        d = tmp_path / "cached"
+        d.mkdir()
+        (d / "data.bin").write_bytes(b"\x00" * 1024)
+
+        assert dir_size_gb(d) == 0.0
+
+        def _unexpected_rglob(self, pattern):
+            raise AssertionError("dir_size_gb unexpectedly walked the filesystem")
+
+        monkeypatch.setattr(Path, "rglob", _unexpected_rglob)
+        assert dir_size_gb(d) == 0.0
+
+    def test_invalidate_dir_size_cache_forces_refresh(self, tmp_path, monkeypatch):
+        clear_dir_size_cache()
+        d = tmp_path / "refresh"
+        d.mkdir()
+        (d / "data.bin").write_bytes(b"\x00" * 1024)
+
+        assert dir_size_gb(d) == 0.0
+
+        original_rglob = Path.rglob
+        calls = {"count": 0}
+
+        def _tracking_rglob(self, pattern):
+            calls["count"] += 1
+            return original_rglob(self, pattern)
+
+        monkeypatch.setattr(Path, "rglob", _tracking_rglob)
+        assert dir_size_gb(d) == 0.0
+        assert calls["count"] == 0
+
+        invalidate_dir_size_cache(d)
+        assert dir_size_gb(d) == 0.0
+        assert calls["count"] == 1