diff --git a/osprey_worker/src/osprey/engine/ast/grammar.py b/osprey_worker/src/osprey/engine/ast/grammar.py index 34ba793b..56ad31bf 100644 --- a/osprey_worker/src/osprey/engine/ast/grammar.py +++ b/osprey_worker/src/osprey/engine/ast/grammar.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections import defaultdict +from collections import OrderedDict from dataclasses import dataclass, field, replace from enum import Enum from pathlib import Path @@ -12,14 +12,53 @@ # from osprey.worker.ui_api.lib.osprey_shared.logging import get_logger from osprey.engine.utils.types import add_slots, cached_property -# Keep these outside `Source` so we don't break pickling/hashing/eq -# Will this leak memory? Maybe -# TODO(old): put this stuff back in cached_property -parsed_ast_root_cache: Dict['Source', 'Root'] = {} -ast_root_lock_cache: Dict['Source', Semaphore] = defaultdict(lambda: Semaphore()) - # logger = get_logger() +# Maximum number of parsed source ASTs retained at once. A ruleset is in the low thousands of +# files; this leaves headroom for a few reload generations while bounding growth on the per-query +# UI-API parse path, which previously leaked an AST for every distinct query for the process life. +_AST_ROOT_CACHE_MAX_SIZE = 4096 + + +class _BoundedASTRootCache: + """A size-bounded, least-recently-used cache mapping a `Source` to its parsed `Root`. + + Kept outside `Source` so we don't break pickling/hashing/eq on the frozen dataclass. Replaces a + previously unbounded module-global dict (and a per-`Source` lock `defaultdict`) that never + evicted, leaking every distinct source parsed over the process lifetime. + """ + + def __init__(self, max_size: int) -> None: + assert max_size > 0, 'cache max_size must be positive' + self._max_size = max_size + self._entries: 'OrderedDict[Source, Root]' = OrderedDict() + + def get(self, source: 'Source') -> 'Optional[Root]': + root = self._entries.get(source) + if root is not None: + self._entries.move_to_end(source) + return root + + def put(self, source: 'Source', root: 'Root') -> None: + self._entries[source] = root + self._entries.move_to_end(source) + while len(self._entries) > self._max_size: + self._entries.popitem(last=False) + + def __len__(self) -> int: + return len(self._entries) + + def __contains__(self, source: object) -> bool: + return source in self._entries + + +parsed_ast_root_cache = _BoundedASTRootCache(_AST_ROOT_CACHE_MAX_SIZE) + +# A single lock guarding parsing. Parsing is pure CPU work that doesn't yield and only runs on a +# cache miss, so one global lock suffices and -- unlike the previous per-`Source` defaultdict -- +# does not retain a lock object per distinct source forever. +_ast_root_parse_lock = Semaphore() + @add_slots @dataclass(frozen=True) @@ -47,24 +86,15 @@ def lines(self) -> Sequence[str]: @property def ast_root(self) -> 'Root': """Returns the ast of this source.""" - ast_root_lock = ast_root_lock_cache[self] - # If there is lock contention (which there should never be, print out some debug info). - if ast_root_lock.locked(): - import traceback - - # logger.error('CONTENTION FOR AST ROOT LOCK') - traceback.print_stack() - - with ast_root_lock: - if self in parsed_ast_root_cache: - return parsed_ast_root_cache[self] + with _ast_root_parse_lock: + cached = parsed_ast_root_cache.get(self) + if cached is not None: + return cached from .py_ast import transform - # logger.debug(f'transforming {self.path}') parsed_ast_root = transform(self) - parsed_ast_root_cache[self] = parsed_ast_root # hack to get around frozen dataclass for now - # logger.debug(f'finished transforming {self.path}') + parsed_ast_root_cache.put(self, parsed_ast_root) return parsed_ast_root diff --git a/osprey_worker/src/osprey/engine/ast/tests/__init__.py b/osprey_worker/src/osprey/engine/ast/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/osprey_worker/src/osprey/engine/ast/tests/test_ast_root_cache.py b/osprey_worker/src/osprey/engine/ast/tests/test_ast_root_cache.py new file mode 100644 index 00000000..f2c6bc42 --- /dev/null +++ b/osprey_worker/src/osprey/engine/ast/tests/test_ast_root_cache.py @@ -0,0 +1,42 @@ +from osprey.engine.ast import grammar +from osprey.engine.ast.grammar import Root, Source, _BoundedASTRootCache + + +def _source(name: str) -> Source: + return Source(path=f'{name}.sml', contents=f'{name} = 1') + + +def test_bounded_cache_evicts_least_recently_used() -> None: + a, b, c = _source('a'), _source('b'), _source('c') + root_a, root_b, root_c = a.ast_root, b.ast_root, c.ast_root + + cache = _BoundedASTRootCache(max_size=2) + cache.put(a, root_a) + cache.put(b, root_b) + assert len(cache) == 2 + + # Access `a` so that `b` becomes the least-recently-used entry. + assert cache.get(a) is root_a + + # Inserting a third entry must evict `b`, not grow past the bound. + cache.put(c, root_c) + assert len(cache) == 2 + assert cache.get(b) is None + assert cache.get(a) is root_a + assert cache.get(c) is root_c + + +def test_ast_root_is_cached_and_stable() -> None: + source = _source('cache_stable_check') + first = source.ast_root + second = source.ast_root + + assert isinstance(first, Root) + assert first is second + assert source in grammar.parsed_ast_root_cache + + +def test_module_cache_is_bounded() -> None: + # Regression: the process-global cache must be a bounded cache, not an unbounded dict. + assert isinstance(grammar.parsed_ast_root_cache, _BoundedASTRootCache) + assert grammar._AST_ROOT_CACHE_MAX_SIZE > 0