From 6cdb8462000d9610b5d15f6c7c45adb787bfec1e Mon Sep 17 00:00:00 2001 From: Ashish Date: Mon, 15 Jun 2026 14:38:08 -0700 Subject: [PATCH] fix(compression): use thread-local tree-sitter parsers in code handler (#893) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `CodeStructureHandler` cached tree-sitter parsers in a process-global dict; the lock only guarded creation, while `parse()` ran unlocked on any thread. tree-sitter `Parser` objects are pyo3 `unsendable` — using one from a non-creator thread panics. The proxy invokes handlers from executor pool threads, so a shared parser is an eventual crash. Same class already fixed in `transforms/code_compressor.py` (#604). Stacked on #892. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/compression/handlers/code_handler.py`: one parser per (thread, language) via `threading.local()`, porting the pattern from `transforms/code_compressor.py`. - `tests/test_compression/test_code_handler.py`: regression test parsing from a 4-worker thread pool, asserting every call stays on the tree-sitter path. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ pytest tests/test_compression/ -q 94 passed, 8 skipped ``` ## Real Behavior Proof - Environment: local macOS, Python 3.11, tree-sitter-language-pack 1.8.1, branch `fix/code-threadlocal-parsers` (stacked on #892). - Exact command / steps: `pytest tests/test_compression/ -q`. - Observed result: 16 parses across a 4-worker pool all stay on the tree-sitter path with no pyo3 panic; previously a shared parser would be touched cross-thread. - Not tested: Reproducing the original panic under production concurrency (covered structurally by the thread-pool test). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A — library change. See Test Output. ## Additional Notes Stacked on #892 — review the top commit until that merges. PR 5 of 7. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Fable 5 Co-authored-by: JD Davis --- headroom/compression/handlers/code_handler.py | 26 ++++++++++++------- tests/test_compression/test_code_handler.py | 25 ++++++++++++++++++ 2 files changed, 42 insertions(+), 9 deletions(-) diff --git a/headroom/compression/handlers/code_handler.py b/headroom/compression/handlers/code_handler.py index 3de5a20cb..a15466e65 100644 --- a/headroom/compression/handlers/code_handler.py +++ b/headroom/compression/handlers/code_handler.py @@ -29,8 +29,7 @@ logger = logging.getLogger(__name__) # Lazy-loaded tree-sitter _tree_sitter_available: bool | None = None -_tree_sitter_parsers: dict[str, Any] = {} -_tree_sitter_lock = threading.Lock() +_tree_sitter_local = threading.local() def _check_tree_sitter() -> bool: @@ -47,19 +46,28 @@ def _check_tree_sitter() -> bool: def _get_parser(language: str) -> Any: - """Get tree-sitter parser for language.""" - global _tree_sitter_parsers + """Return a **thread-local** tree-sitter parser for ``language``. + tree-sitter ``Parser`` objects are pyo3 ``unsendable`` — touching + one from a thread other than its creator panics. Handlers run on + executor pool threads in the proxy, so parsers must never be shared + across threads. One parser per (thread, language); same fix as + ``transforms/code_compressor.py`` (#604). + """ if not _check_tree_sitter(): raise ImportError("tree-sitter-language-pack not installed") - with _tree_sitter_lock: - if language not in _tree_sitter_parsers: - from tree_sitter_language_pack import get_parser + cache: dict[str, Any] | None = getattr(_tree_sitter_local, "parsers", None) + if cache is None: + cache = {} + _tree_sitter_local.parsers = cache - _tree_sitter_parsers[language] = get_parser(language) # type: ignore[arg-type] + if language not in cache: + from tree_sitter_language_pack import get_parser - return _tree_sitter_parsers[language] + cache[language] = get_parser(language) # type: ignore[arg-type] + + return cache[language] # tree-sitter API compatibility. tree-sitter-language-pack switched to a diff --git a/tests/test_compression/test_code_handler.py b/tests/test_compression/test_code_handler.py index 50fb3939c..4cebc697e 100644 --- a/tests/test_compression/test_code_handler.py +++ b/tests/test_compression/test_code_handler.py @@ -229,6 +229,31 @@ class TestTreeSitterContainers: result.mask.mask[i] for i in range(start, start + len("let body_line = 5;")) ), "impl method body must be compressible" + def test_concurrent_parsing_uses_tree_sitter(self, handler): + """Parsers must be thread-local. + + Regression: parsers were cached in a process-global dict and + shared across threads. tree-sitter Parser objects are pyo3 + unsendable — touching one from a non-creator thread panics (or + raises, dropping the handler to the regex fallback). Parsing + from a thread pool must succeed on the tree-sitter path in + every thread. + """ + from concurrent.futures import ThreadPoolExecutor + + code = "class Foo:\n def m(self):\n x = 1\n return x\n" + + def work(_: int) -> str: + result = handler.get_mask(code, language="python") + return str(result.metadata["parser"]) + + with ThreadPoolExecutor(max_workers=4) as pool: + parsers = list(pool.map(work, range(16))) + + assert parsers == ["tree-sitter"] * 16, ( + f"all threads must parse via tree-sitter, got: {set(parsers)}" + ) + def test_non_ascii_content_mask_alignment(self, handler): """Byte offsets must be converted to char offsets.