mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(compression): use thread-local tree-sitter parsers in code handler (#893)
## 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 # <!-- compression-handler review --> ## 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 <noreply@anthropic.com> Co-authored-by: JD Davis <mxjerrett@gmail.com>
This commit is contained in:
parent
3b0bceecf4
commit
6cdb846200
2 changed files with 42 additions and 9 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue