music-assistant-server
24.2 KB•PY
controller.py
24.2 KB • 584 lines • python
1"""Cache controller implementation."""
2
3from __future__ import annotations
4
5import asyncio
6import os
7import time
8from collections.abc import AsyncGenerator, Mapping
9from contextlib import asynccontextmanager
10from pathlib import Path
11from typing import TYPE_CHECKING, Any
12
13from music_assistant_models.background_task import TaskSchedule
14from music_assistant_models.config_entries import ConfigActionResult, ConfigEntry
15from music_assistant_models.enums import ConfigEntryType
16
17from music_assistant.constants import (
18 DB_TABLE_CACHE,
19 DB_TABLE_SETTINGS,
20 VACUUM_MIN_RECLAIM_RATIO,
21)
22from music_assistant.controllers.cache.constants import (
23 BYPASS_CACHE,
24 CACHE_DATABASE_CLEANUP_TASK_ID,
25 CONF_CLEAR_CACHE,
26 DB_SCHEMA_VERSION,
27 DEFAULT_CACHE_EXPIRATION,
28 LOGGER,
29 MAX_CACHE_DB_SIZE_MB,
30 SWR_FALLBACK_MAX_AGE,
31)
32from music_assistant.controllers.tasks.context import (
33 update_current_task_progress_text,
34)
35from music_assistant.helpers.database import DatabaseConnection
36from music_assistant.helpers.datetime import local_clock_time_to_utc
37from music_assistant.helpers.json import SerializableType, async_json_loads, json_dumps, json_loads
38from music_assistant.models.core_controller import CoreController
39
40if TYPE_CHECKING:
41 from music_assistant_models.config_entries import CoreConfig
42
43 from music_assistant import MusicAssistant
44
45
46class CacheController(CoreController):
47 """Controller handling caching of data throughout the application."""
48
49 domain: str = "cache"
50
51 def __init__(self, mass: MusicAssistant) -> None:
52 """Initialize core controller."""
53 super().__init__(mass)
54 self.database: DatabaseConnection | None = None
55 self.manifest.name = "Cache controller"
56 self.manifest.description = (
57 "Music Assistant's core controller for caching data throughout the application."
58 )
59 self.manifest.icon = "memory"
60
61 async def get_config_entries(self) -> tuple[ConfigEntry, ...]:
62 """Return all Config Entries for this core module (if any)."""
63 return (
64 ConfigEntry(
65 key=CONF_CLEAR_CACHE,
66 type=ConfigEntryType.ACTION,
67 ),
68 )
69
70 async def handle_config_action(
71 self, action: str
72 ) -> tuple[ConfigEntry, ...] | ConfigActionResult | None:
73 """Handle a one-shot action button press and report its outcome."""
74 if action == CONF_CLEAR_CACHE:
75 await self.clear()
76 return ConfigActionResult(translation_key=f"{CONF_CLEAR_CACHE}.result")
77 return await super().handle_config_action(action)
78
79 async def setup(self, config: CoreConfig) -> None:
80 """Async initialize of cache module."""
81 self.logger.info("Initializing cache controller...")
82 await self._setup_database()
83
84 async def post_setup(self) -> None:
85 """Handle logic after all core controllers have been set up."""
86 self._register_cleanup_task()
87
88 async def close(self) -> None:
89 """Cleanup on exit."""
90 if self.database:
91 await self.database.close()
92
93 async def get_diagnostics(self) -> dict[str, SerializableType]:
94 """Return diagnostics info for this controller to include in diagnostics reports."""
95 return {
96 "db_schema_version": DB_SCHEMA_VERSION,
97 "db_size_mb": round(await self._get_cache_db_size_mb(), 1),
98 "entries": await self.database.get_count(DB_TABLE_CACHE) if self.database else None,
99 }
100
101 async def get(
102 self,
103 key: str,
104 provider: str = "default",
105 category: int = 0,
106 checksum: str | int | None = None,
107 default: Any = None,
108 allow_bypass: bool | None = None,
109 base_class: Any = None,
110 allow_expired_cache: bool = False,
111 ) -> Any:
112 """
113 Get data from cache.
114
115 Returns JSON-deserialized data (dicts, lists, strings, numbers, booleans, None).
116
117 If base_class is provided, the raw data is automatically reconstructed using
118 its from_dict() method. If the cached data is a list of dicts, each item is
119 reconstructed individually.
120
121 :param key: The (unique) lookup key of the cache object.
122 :param provider: Provider id to group cache objects.
123 :param category: Category to group cache objects.
124 :param checksum: If provided, only return data if the stored checksum matches.
125 :param default: Value to return if no cache object is found.
126 :param allow_bypass: Whether to respect the BYPASS_CACHE context variable.
127 :param base_class: If provided, reconstruct data using base_class.from_dict().
128 :param allow_expired_cache: If True, also return entries past their expiration
129 time instead of treating them as cache misses.
130 """
131 data, _, found = await self.get_with_freshness(
132 key,
133 provider=provider,
134 category=category,
135 checksum=checksum,
136 allow_bypass=allow_bypass,
137 base_class=base_class,
138 include_expired=allow_expired_cache,
139 )
140 return data if found else default
141
142 async def get_with_freshness(
143 self,
144 key: str,
145 provider: str = "default",
146 category: int = 0,
147 checksum: str | int | None = None,
148 allow_bypass: bool | None = None,
149 base_class: Any = None,
150 include_expired: bool = False,
151 ) -> tuple[Any, bool, bool]:
152 """
153 Get data from cache together with the freshness and presence of the entry.
154
155 Returns a (data, is_fresh, found) tuple. found is False when there is no usable
156 entry, in which case data is None; is_fresh is False when the entry is expired.
157 Because a stored None value is returned as-is, use the found flag to tell a cache
158 miss from a cached None.
159
160 :param key: The (unique) lookup key of the cache object.
161 :param provider: Provider id to group cache objects.
162 :param category: Category to group cache objects.
163 :param checksum: If provided, only return data if the stored checksum matches.
164 :param allow_bypass: Whether to respect the BYPASS_CACHE context variable.
165 :param base_class: If provided, reconstruct data using base_class.from_dict().
166 :param include_expired: If False (default), an expired entry is reported as not found
167 and is not deserialized; set True to also return expired entries as stale data.
168 """
169 assert self.database is not None
170 assert key, "No key provided"
171 if allow_bypass and BYPASS_CACHE.get():
172 return None, False, False
173 cur_time = int(time.time())
174 if checksum is not None and not isinstance(checksum, str):
175 checksum = str(checksum)
176 if (
177 db_row := await self.database.get_row(
178 DB_TABLE_CACHE, {"category": category, "provider": provider, "key": key}
179 )
180 ) and (not checksum or db_row["checksum"] == checksum):
181 # if allow_bypass is not explicitly set,
182 # determine it based on the 'persistent' flag of the cache entry
183 if allow_bypass is None:
184 allow_bypass = not bool(db_row["persistent"])
185 if allow_bypass and BYPASS_CACHE.get():
186 return None, False, False
187 is_fresh = bool(db_row["expires"] >= cur_time)
188 # skip deserialization for an expired entry the caller will not use
189 if not is_fresh and not include_expired:
190 return None, False, False
191 try:
192 data = await async_json_loads(db_row["data"])
193 except Exception as exc:
194 LOGGER.error(
195 "Error parsing cache data for %s/%s/%s: %s",
196 provider,
197 category,
198 key,
199 str(exc),
200 exc_info=exc if self.logger.isEnabledFor(10) else None,
201 )
202 else:
203 if base_class is not None and data is not None:
204 if isinstance(data, list):
205 return [base_class.from_dict(item) for item in data], is_fresh, True
206 return base_class.from_dict(data), is_fresh, True
207 return data, is_fresh, True
208 return None, False, False
209
210 async def get_all(
211 self,
212 provider: str = "default",
213 category: int = 0,
214 base_class: Any = None,
215 ) -> dict[str, Any]:
216 """
217 Return every non-expired cache entry for a provider/category as a key -> data mapping.
218
219 Use this instead of many individual :meth:`get` calls when a caller needs to check a
220 large number of keys against the cache at once (e.g. while scanning a whole library),
221 since it issues a single query and a single deserialization batch rather than one of
222 each per key.
223
224 :param provider: Provider id to group cache objects.
225 :param category: Category to group cache objects.
226 :param base_class: If provided, reconstruct each entry using base_class.from_dict().
227 """
228 assert self.database is not None
229 cur_time = int(time.time())
230 rows = await self.database.get_rows_from_query(
231 f"SELECT key, data FROM {DB_TABLE_CACHE} "
232 "WHERE category = :category AND provider = :provider AND expires >= :cur_time",
233 {"category": category, "provider": provider, "cur_time": cur_time},
234 limit=0,
235 )
236 # deserialize every row in one thread-pool submission instead of one per row, which
237 # otherwise means thousands of thread submissions/context switches on a large result
238 result = await asyncio.to_thread(self._deserialize_rows, rows, provider, category)
239 if base_class is not None:
240 for key, data in result.items():
241 if data is None:
242 continue
243 result[key] = (
244 [base_class.from_dict(item) for item in data]
245 if isinstance(data, list)
246 else base_class.from_dict(data)
247 )
248 return result
249
250 async def get_expiration(
251 self,
252 key: str,
253 provider: str = "default",
254 category: int = 0,
255 ) -> int | None:
256 """
257 Return the expiration timestamp (epoch seconds) of a cache entry, if any.
258
259 Cheap existence/freshness probe: only the expiration column is read, the
260 stored data is not. Returns None when no entry exists for the given key.
261
262 :param key: The (unique) lookup key of the cache object.
263 :param provider: Provider id to group cache objects.
264 :param category: Category to group cache objects.
265 """
266 assert self.database is not None
267 assert key, "No key provided"
268 rows = await self.database.get_rows_from_query(
269 f"SELECT expires FROM {DB_TABLE_CACHE} "
270 "WHERE category = :category AND provider = :provider AND key = :key",
271 {"category": category, "provider": provider, "key": key},
272 limit=1,
273 )
274 return int(rows[0]["expires"]) if rows else None
275
276 async def set(
277 self,
278 key: str,
279 data: SerializableType,
280 expiration: int = DEFAULT_CACHE_EXPIRATION,
281 provider: str = "default",
282 category: int = 0,
283 checksum: str | None = None,
284 persistent: bool = False,
285 allow_expired_cache: bool = False,
286 ) -> None:
287 """
288 Store data in cache.
289
290 Data must be JSON-serializable (str, int, float, bool, None, list, dict).
291 Do not pass model objects directly â use .to_dict() first.
292 Non-serializable data will raise TypeError.
293
294 :param key: The (unique) lookup key of the cache object.
295 :param data: JSON-serializable data to store.
296 :param expiration: Time in seconds the cache object should be valid.
297 :param provider: Provider id to group cache objects.
298 :param category: Category to group cache objects.
299 :param checksum: Optional checksum to store with the cache object.
300 :param persistent: If True, the entry survives cache clears.
301 :param allow_expired_cache: If True, the entry survives the auto-cleanup task
302 after it expires, so it can still be served as fallback data by the
303 stale-while-revalidate path of `@use_cache`.
304 """
305 assert self.database is not None
306 if not key:
307 return
308 if checksum is not None:
309 checksum = str(checksum)
310 expires = int(time.time() + expiration)
311 # always serialize to JSON to ensure data is serializable
312 # this raises if the data contains non-serializable objects
313 data = await asyncio.to_thread(json_dumps, data)
314 # upsert (update in place on the UNIQUE(category, key, provider) conflict) instead of
315 # INSERT OR REPLACE, which deletes and re-inserts the row and so rewrites every index
316 await self.database.upsert(
317 DB_TABLE_CACHE,
318 {
319 "category": category,
320 "provider": provider,
321 "key": key,
322 "expires": expires,
323 "checksum": checksum,
324 "data": data,
325 "persistent": persistent,
326 "allow_expired_cache": allow_expired_cache,
327 },
328 )
329
330 async def delete(
331 self, key: str | None, category: int | None = None, provider: str | None = None
332 ) -> None:
333 """Delete data from cache."""
334 assert self.database is not None
335 match: dict[str, str | int] = {}
336 if key is not None:
337 match["key"] = key
338 if category is not None:
339 match["category"] = category
340 if provider is not None:
341 match["provider"] = provider
342 await self.database.delete(DB_TABLE_CACHE, match)
343
344 async def clear(
345 self,
346 key_filter: str | None = None,
347 category_filter: int | None = None,
348 provider_filter: str | None = None,
349 include_persistent: bool = False,
350 ) -> None:
351 """Clear all/partial items from cache."""
352 assert self.database is not None
353 self.logger.info("Clearing database...")
354 query_parts: list[str] = []
355 if category_filter is not None:
356 query_parts.append(f"category = {category_filter}")
357 if provider_filter is not None:
358 query_parts.append(f"provider LIKE '%{provider_filter}%'")
359 if key_filter is not None:
360 query_parts.append(f"key LIKE '%{key_filter}%'")
361 if not include_persistent:
362 query_parts.append("persistent = 0")
363 query = "WHERE " + " AND ".join(query_parts) if query_parts else None
364 await self.database.delete(DB_TABLE_CACHE, query=query)
365 self.logger.info("Clearing database DONE")
366
367 async def auto_cleanup(self) -> None:
368 """Run scheduled auto cleanup task."""
369 assert self.database is not None
370 self.logger.debug("Running automatic cleanup...")
371 update_current_task_progress_text("Removing expired cache records")
372 cur_timestamp = int(time.time())
373 # remove expired entries; allow_expired_cache entries are kept as stale-while-revalidate
374 # fallback, but only until they are expired beyond SWR_FALLBACK_MAX_AGE - past that their
375 # key is clearly no longer requested and the row would otherwise live forever
376 swr_cutoff = cur_timestamp - SWR_FALLBACK_MAX_AGE
377 cursor = await self.database.execute(
378 f"DELETE FROM {DB_TABLE_CACHE} WHERE "
379 "(expires < :timestamp AND allow_expired_cache = 0) "
380 "OR (expires < :swr_cutoff AND allow_expired_cache = 1)",
381 {"timestamp": cur_timestamp, "swr_cutoff": swr_cutoff},
382 )
383 await self.database.commit()
384 cleaned_records = cursor.rowcount
385 update_current_task_progress_text(f"Cleaned up {cleaned_records} expired cache record(s)")
386 self.logger.debug("Automatic cleanup finished (cleaned up %s records)", cleaned_records)
387
388 @asynccontextmanager
389 async def handle_refresh(self, bypass: bool) -> AsyncGenerator[None]:
390 """Handle the cache bypass."""
391 try:
392 token = BYPASS_CACHE.set(bypass)
393 yield None
394 finally:
395 BYPASS_CACHE.reset(token)
396
397 async def _check_oversized_cache(self) -> None:
398 """Warn if the cache database exceeds the recommended max size."""
399 db_size_mb = await self._get_cache_db_size_mb()
400 if db_size_mb > MAX_CACHE_DB_SIZE_MB:
401 self.logger.warning(
402 "Cache database size %.2f MB exceeds recommended maximum of %d MB",
403 db_size_mb,
404 MAX_CACHE_DB_SIZE_MB,
405 )
406
407 def _deserialize_rows(
408 self, rows: list[Mapping[str, Any]], provider: str, category: int
409 ) -> dict[str, Any]:
410 """
411 JSON-deserialize a batch of raw cache rows synchronously, skipping unparsable ones.
412
413 :param rows: Raw ``key``/``data`` rows selected from the cache table.
414 :param provider: Provider id the rows were selected for, used only for error logging.
415 :param category: Category the rows were selected for, used only for error logging.
416 """
417 result: dict[str, Any] = {}
418 for row in rows:
419 try:
420 result[row["key"]] = json_loads(row["data"])
421 except ValueError as exc:
422 LOGGER.error(
423 "Error parsing cache data for %s/%s/%s: %s",
424 provider,
425 category,
426 row["key"],
427 str(exc),
428 exc_info=exc if self.logger.isEnabledFor(10) else None,
429 )
430 return result
431
432 async def _get_cache_db_size_mb(self) -> float:
433 """Return the on-disk size of the cache database (in MB)."""
434 db_path = os.path.join(self.mass.cache_path, "cache.db")
435 # also include the write ahead log and shared memory db files
436 db_files = [db_path + suffix for suffix in ("", "-wal", "-shm")]
437
438 def _get_db_size() -> float:
439 total = 0
440 for path in db_files:
441 if Path(path).exists():
442 total += Path(path).stat().st_size
443 return total / (1024 * 1024)
444
445 return await asyncio.to_thread(_get_db_size)
446
447 async def _setup_database(self) -> None:
448 """Initialize database."""
449 await self._check_oversized_cache()
450 db_path = os.path.join(self.mass.cache_path, "cache.db")
451 self.database = DatabaseConnection(db_path)
452 await self.database.setup()
453
454 # always create db tables if they don't exist to prevent errors trying to access them later
455 await self.__create_database_tables()
456
457 try:
458 if db_row := await self.database.get_row(DB_TABLE_SETTINGS, {"key": "version"}):
459 prev_version = int(db_row["value"])
460 else:
461 prev_version = 0
462 except KeyError, ValueError:
463 prev_version = 0
464
465 if prev_version not in (0, DB_SCHEMA_VERSION):
466 LOGGER.warning(
467 "Performing database migration from %s to %s",
468 prev_version,
469 DB_SCHEMA_VERSION,
470 )
471 try:
472 await self.__migrate_database(prev_version)
473 except Exception as err:
474 LOGGER.warning("Cache database migration failed: %s, resetting cache", err)
475 await self.database.execute(f"DROP TABLE IF EXISTS {DB_TABLE_CACHE}")
476 await self.__create_database_tables()
477
478 # store current schema version
479 await self.database.insert_or_replace(
480 DB_TABLE_SETTINGS,
481 {"key": "version", "value": str(DB_SCHEMA_VERSION), "type": "str"},
482 )
483 await self.__create_database_indexes()
484
485 # Skip the full rebuild unless a meaningful share of the file can be reclaimed.
486 try:
487 reclaimable_ratio = await self.database.get_reclaimable_ratio()
488 if reclaimable_ratio < VACUUM_MIN_RECLAIM_RATIO:
489 self.logger.debug(
490 "Skipping database compaction (only %.1f%% reclaimable)",
491 reclaimable_ratio * 100,
492 )
493 else:
494 self.logger.debug(
495 "Compacting database (%.1f%% reclaimable)...", reclaimable_ratio * 100
496 )
497 await self.database.vacuum()
498 self.logger.debug("Compacting database done")
499 except Exception as err:
500 self.logger.warning("Database vacuum failed: %s", str(err))
501
502 async def __create_database_tables(self) -> None:
503 """Create database table(s)."""
504 assert self.database is not None
505 await self.database.execute(
506 f"""CREATE TABLE IF NOT EXISTS {DB_TABLE_SETTINGS}(
507 key TEXT PRIMARY KEY,
508 value TEXT,
509 type TEXT
510 );"""
511 )
512 await self.database.execute(
513 f"""CREATE TABLE IF NOT EXISTS {DB_TABLE_CACHE}(
514 [id] INTEGER PRIMARY KEY AUTOINCREMENT,
515 [category] INTEGER NOT NULL DEFAULT 0,
516 [key] TEXT NOT NULL,
517 [provider] TEXT NOT NULL,
518 [expires] INTEGER NOT NULL,
519 [data] TEXT NULL,
520 [checksum] TEXT NULL,
521 [persistent] INTEGER NOT NULL DEFAULT 0,
522 [allow_expired_cache] INTEGER NOT NULL DEFAULT 0,
523 UNIQUE(category, key, provider)
524 )"""
525 )
526
527 await self.database.commit()
528
529 async def __create_database_indexes(self) -> None:
530 """Create database indexes."""
531 assert self.database is not None
532 # The UNIQUE(category, key, provider) constraint already provides an index that serves
533 # every point lookup (get() matches exactly those three columns) and any delete that
534 # includes the category. The only access pattern its column order cannot serve is a
535 # delete that filters by (key, provider) without a category, so that is the single
536 # secondary index kept here.
537 await self.database.execute(
538 f"CREATE INDEX IF NOT EXISTS {DB_TABLE_CACHE}_key_provider_idx "
539 f"ON {DB_TABLE_CACHE}(key,provider);"
540 )
541 await self.database.commit()
542
543 async def __migrate_database(self, prev_version: int) -> None:
544 """Perform a database migration."""
545 assert self.database is not None
546 if prev_version <= 6:
547 # clear spotify cache entries to fix bloated cache from playlist pagination bug
548 await self.database.delete(DB_TABLE_CACHE, query="WHERE provider LIKE '%spotify%'")
549 if prev_version <= 7:
550 await self.database.execute(
551 f"ALTER TABLE {DB_TABLE_CACHE} "
552 "ADD COLUMN allow_expired_cache INTEGER NOT NULL DEFAULT 0"
553 )
554 if prev_version <= 8:
555 # drop the redundant secondary indexes: they either duplicate the
556 # UNIQUE(category, key, provider) autoindex or are a left-prefix of it, so the
557 # autoindex already serves their lookups. The (key, provider) index is (re)created
558 # by __create_database_indexes and intentionally kept.
559 for index_name in (
560 "category_idx",
561 "key_idx",
562 "provider_idx",
563 "category_key_idx",
564 "category_provider_idx",
565 "category_key_provider_idx",
566 ):
567 await self.database.execute(f"DROP INDEX IF EXISTS {DB_TABLE_CACHE}_{index_name}")
568 await self.database.commit()
569
570 def _register_cleanup_task(self) -> None:
571 """Register the recurring cache database cleanup task."""
572 utc_hour, utc_minute = local_clock_time_to_utc(4, 0)
573 desired_schedule = TaskSchedule.daily(hour=utc_hour, minute=utc_minute)
574 self.mass.tasks.register_scheduled_task(
575 task_id=CACHE_DATABASE_CLEANUP_TASK_ID,
576 name="Cache database cleanup",
577 handler=self.auto_cleanup,
578 schedule=desired_schedule,
579 translation_key="cache_database_cleanup",
580 translation_owner=self.translation_owner,
581 metadata={"task_domain": "cache_database_cleanup"},
582 allow_retry=True,
583 )
584