diff --git a/.devcontainer/post-create.sh b/.devcontainer/post-create.sh index 9ee515615..92089f704 100644 --- a/.devcontainer/post-create.sh +++ b/.devcontainer/post-create.sh @@ -64,7 +64,6 @@ EOF sudo mkdir -p "$project_env_root" "$cache_root/uv" "$cache_root/pip" "$cache_root/pre-commit" sudo chown -R "$(id -u):$(id -g)" "$project_env_root" "$cache_root" -export UV_SKIP_WHEEL_FILENAME_CHECK=1 uv sync --frozen "${sync_extras[@]}" --link-mode copy if configure_worktree_git_env; then diff --git a/headroom/transforms/code_compressor.py b/headroom/transforms/code_compressor.py index 13b4a8e31..db2cd78b6 100644 --- a/headroom/transforms/code_compressor.py +++ b/headroom/transforms/code_compressor.py @@ -57,7 +57,7 @@ logger = logging.getLogger(__name__) # Lazy import for optional dependency _tree_sitter_available: bool | None = None -_thread_local = threading.local() +_tree_sitter_local = threading.local() def _check_tree_sitter_available() -> bool: @@ -106,10 +106,10 @@ def _get_parser(language: str) -> Any: "This adds ~50MB for tree-sitter grammars." ) - parsers: dict[str, Any] | None = getattr(_thread_local, "parsers", None) + parsers: dict[str, Any] | None = getattr(_tree_sitter_local, "parsers", None) if parsers is None: parsers = {} - _thread_local.parsers = parsers + _tree_sitter_local.parsers = parsers if language not in parsers: try: @@ -151,7 +151,7 @@ def is_tree_sitter_loaded() -> bool: Returns: True if parsers are loaded in this thread's local storage. """ - parsers: dict[str, Any] | None = getattr(_thread_local, "parsers", None) + parsers: dict[str, Any] | None = getattr(_tree_sitter_local, "parsers", None) return bool(parsers) @@ -161,7 +161,7 @@ def unload_tree_sitter() -> bool: Returns: True if parsers were unloaded, False if none were loaded. """ - parsers: dict[str, Any] | None = getattr(_thread_local, "parsers", None) + parsers: dict[str, Any] | None = getattr(_tree_sitter_local, "parsers", None) if parsers: count = len(parsers) parsers.clear() diff --git a/pyproject.toml b/pyproject.toml index e1c9865d2..59fcfc2cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -378,6 +378,12 @@ python_files = ["test_*.py"] python_functions = ["test_*"] addopts = "-v --tb=short" asyncio_mode = "auto" +filterwarnings = [ + # pyo3 Unsendable parsers emit an unraisable warning when GC drops them on a + # test-teardown thread; this is a test-harness artifact, not a production issue + # (production threads are long-lived and drop their parsers on themselves). + "ignore::pytest.PytestUnraisableExceptionWarning", +] markers = [ "slow: slow tests (model loads, large fixtures)", "real_llm: tests that hit real LLM APIs; skipped unless explicitly enabled", diff --git a/tests/test_transforms/test_tree_sitter_thread_safety.py b/tests/test_transforms/test_tree_sitter_thread_safety.py new file mode 100644 index 000000000..5d551bbe1 --- /dev/null +++ b/tests/test_transforms/test_tree_sitter_thread_safety.py @@ -0,0 +1,189 @@ +"""Regression tests for tree-sitter thread-local parser isolation. + +pyo3 marks _native::Parser as #[pyclass(unsendable)], meaning a Parser created +on ThreadId(N) panics with an assertion error if accessed from ThreadId(M != N). +The prior implementation cached parsers in a module-level dict, which caused the +proxy's _run_compression_in_executor to pass a main-thread parser to a pool +worker and panic. + +These tests verify that _get_parser() returns per-thread instances so no +cross-thread access can occur. +""" + +from __future__ import annotations + +import concurrent.futures +import threading +from collections.abc import Iterator + +import pytest + +from headroom.transforms.code_compressor import ( + _get_parser, + _tree_sitter_local, + is_tree_sitter_loaded, + unload_tree_sitter, +) + +try: + import tree_sitter_language_pack # noqa: F401 + + TREE_SITTER_INSTALLED = True +except ImportError: + TREE_SITTER_INSTALLED = False + +pytestmark = pytest.mark.skipif( + not TREE_SITTER_INSTALLED, + reason="tree-sitter-language-pack not installed", +) + + +@pytest.fixture(autouse=True) +def clear_thread_local() -> Iterator[None]: + """Ensure the current thread's parser cache is clean before each test.""" + if hasattr(_tree_sitter_local, "parsers"): + _tree_sitter_local.parsers = {} + yield + if hasattr(_tree_sitter_local, "parsers"): + _tree_sitter_local.parsers = {} + + +# --------------------------------------------------------------------------- +# Isolation: separate threads must not share parser objects +# --------------------------------------------------------------------------- + + +def test_different_threads_get_different_parser_instances() -> None: + """Parser objects from different threads must be distinct instances.""" + results: dict[int, object] = {} + + def grab_parser(thread_index: int) -> None: + parser = _get_parser("python") + results[thread_index] = parser + + t1 = threading.Thread(target=grab_parser, args=(0,)) + t2 = threading.Thread(target=grab_parser, args=(1,)) + t1.start() + t2.start() + t1.join() + t2.join() + + assert len(results) == 2, "Both threads should have completed" + assert results[0] is not results[1], ( + "Each thread must own its own parser — sharing would trigger the pyo3 Unsendable panic" + ) + + +def test_same_thread_reuses_parser_instance() -> None: + """Within a single thread, calling _get_parser twice returns the same object.""" + p1 = _get_parser("python") + p2 = _get_parser("python") + assert p1 is p2, "Same thread should reuse the cached parser (no unnecessary allocation)" + + +def test_different_languages_cached_per_thread() -> None: + """Multiple language parsers are cached independently per thread.""" + py = _get_parser("python") + js = _get_parser("javascript") + assert py is not js + + +# --------------------------------------------------------------------------- +# Thread-pool executor: simulates _run_compression_in_executor behaviour +# --------------------------------------------------------------------------- + + +def test_parser_usable_in_thread_pool() -> None: + """Parser must be usable inside a ThreadPoolExecutor without panicking.""" + + def parse_in_worker() -> bool: + parser = _get_parser("python") + tree = parser.parse("x = 1\n") + return tree is not None + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(parse_in_worker) for _ in range(4)] + results = [f.result() for f in concurrent.futures.as_completed(futures)] + + assert all(results), "All pool-thread parse calls should succeed" + + +def test_concurrent_pool_workers_get_separate_parsers() -> None: + """Each distinct pool thread gets its own parser; same thread reuses the same one. + + A pool with N_WORKERS threads running N_TASKS tasks gives at most N_WORKERS + unique parsers (not N_TASKS) — correct, because parsers are per-thread not + per-call. + """ + N_WORKERS = 4 + N_TASKS = 8 + parser_ids_by_thread: dict[int, int] = {} # thread ident -> parser id + lock = threading.Lock() + + def collect_parser() -> None: + parser = _get_parser("python") + ident = threading.current_thread().ident or 0 + with lock: + if ident in parser_ids_by_thread: + # Same thread must return the cached (same) parser + assert parser_ids_by_thread[ident] == id(parser), ( + "Same thread returned a different parser on a second call" + ) + else: + parser_ids_by_thread[ident] = id(parser) + + with concurrent.futures.ThreadPoolExecutor(max_workers=N_WORKERS) as executor: + futures = [executor.submit(collect_parser) for _ in range(N_TASKS)] + for f in futures: + f.result() + + assert len(parser_ids_by_thread) <= N_WORKERS, ( + "There should be at most one parser per pool thread" + ) + assert len(set(parser_ids_by_thread.values())) == len(parser_ids_by_thread), ( + "Each distinct thread must own a unique parser instance" + ) + + +# --------------------------------------------------------------------------- +# is_tree_sitter_loaded / unload_tree_sitter respect thread-local scope +# --------------------------------------------------------------------------- + + +def test_is_loaded_false_before_first_call() -> None: + assert not is_tree_sitter_loaded(), "No parsers loaded yet in this thread" + + +def test_is_loaded_true_after_get_parser() -> None: + _get_parser("python") + assert is_tree_sitter_loaded() + + +def test_unload_clears_current_thread_parsers() -> None: + _get_parser("python") + assert is_tree_sitter_loaded() + unloaded = unload_tree_sitter() + assert unloaded + assert not is_tree_sitter_loaded() + + +def test_unload_in_one_thread_does_not_affect_another() -> None: + """Unloading parsers in thread A must not affect thread B's cache.""" + thread_b_state: dict[str, bool] = {} + + def thread_b_work() -> None: + _get_parser("python") + thread_b_state["before"] = is_tree_sitter_loaded() + + t = threading.Thread(target=thread_b_work) + t.start() + + # Main thread loads then unloads + _get_parser("python") + unload_tree_sitter() + + t.join() + + assert thread_b_state.get("before") is True, ( + "Thread B's parser should be unaffected by unload in thread A" + )