diff --git a/folio/graph.py b/folio/graph.py index a88b165..c8a8d5a 100644 --- a/folio/graph.py +++ b/folio/graph.py @@ -21,7 +21,7 @@ import time import traceback from enum import Enum -from functools import cache +from functools import cache, lru_cache from pathlib import Path from typing import Dict, List, Literal, Optional, Tuple @@ -124,9 +124,21 @@ class FOLIOTypes(Enum): # minimum length for prefix search MIN_PREFIX_LENGTH: int = 3 +# Bound query-result caches so long-running processes cannot retain every unique search forever. +DEFAULT_SEARCH_CACHE_SIZE: int = 128 + # Set up logger LOGGER = get_logger(__name__) + +def _store_prefix_cache( + cache_map: Dict[str, List[OWLClass]], key: str, classes: List[OWLClass] +) -> None: + """Store one prefix-search result and evict the oldest entry at capacity.""" + cache_map[key] = classes + while len(cache_map) > DEFAULT_SEARCH_CACHE_SIZE: + cache_map.pop(next(iter(cache_map))) + # try to import rapidfuzz and marisa_trie with importlib; log if not able to. try: if importlib.util.find_spec("rapidfuzz") is not None: @@ -1384,7 +1396,7 @@ def _search_by_prefix_sensitive(self, prefix: str) -> List[OWLClass]: iri_list.append(idx) classes = [self[index] for index in iri_list] - self._prefix_cache[prefix] = classes # type: ignore[assignment] + _store_prefix_cache(self._prefix_cache, prefix, classes) return classes def _search_by_prefix_insensitive(self, prefix: str) -> List[OWLClass]: @@ -1436,11 +1448,11 @@ def _search_by_prefix_insensitive(self, prefix: str) -> List[OWLClass]: iri_list.append(idx) classes = [self[index] for index in iri_list] - self._ci_prefix_cache[folded] = classes # type: ignore[assignment] + _store_prefix_cache(self._ci_prefix_cache, folded, classes) return classes @staticmethod - @cache + @lru_cache(maxsize=DEFAULT_SEARCH_CACHE_SIZE) def _basic_search( query: str, search_list: Tuple[str], diff --git a/tests/test_folio.py b/tests/test_folio.py index f62d83e..0763916 100644 --- a/tests/test_folio.py +++ b/tests/test_folio.py @@ -345,6 +345,30 @@ def test_search_prefix_fallback_parity(folio_graph, monkeypatch): ) +def test_search_caches_are_bounded(folio_graph): + """Unique search queries must not grow process-lifetime caches without bound.""" + FOLIO._basic_search.cache_clear() + folio_graph._prefix_cache = {} + folio_graph._ci_prefix_cache = {} + + try: + for index in range(129): + query = f"no-match-{index}" + FOLIO._basic_search(query, ("alpha", "beta"), limit=1) + folio_graph.search_by_prefix(query, case_sensitive=True) + folio_graph.search_by_prefix(query) + + cache_info = FOLIO._basic_search.cache_info() + assert cache_info.maxsize == 128 + assert cache_info.currsize == 128 + assert len(folio_graph._prefix_cache) == 128 + assert len(folio_graph._ci_prefix_cache) == 128 + finally: + FOLIO._basic_search.cache_clear() + folio_graph._prefix_cache = {} + folio_graph._ci_prefix_cache = {} + + def test_search_label(folio_graph): for c, score in folio_graph.search_by_label("Georgia"): assert "Georgia" in c.label