mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
test(code_compressor): add unsendable-panic repro and thread-local parser tests
This commit is contained in:
parent
38aefc1d34
commit
6a70ea653c
2 changed files with 273 additions and 0 deletions
112
tests/repro_unsendable_panic.py
Normal file
112
tests/repro_unsendable_panic.py
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Reproducer for tree-sitter Parser unsendable panic (issue #562).
|
||||
|
||||
Demonstrates that tree-sitter ≥ 0.23 marks Parser as PyO3
|
||||
#[pyclass(unsendable)] — it hard-panics if accessed from a different
|
||||
thread than its creator.
|
||||
|
||||
The OLD pattern (shared dict + lock) triggers the panic.
|
||||
The NEW pattern (thread-local storage) works correctly.
|
||||
|
||||
Usage:
|
||||
python repro_unsendable_panic.py
|
||||
"""
|
||||
|
||||
import threading
|
||||
import concurrent.futures
|
||||
|
||||
CODE = b"def hello():\n return 42\n"
|
||||
|
||||
|
||||
def test_old_pattern_shared_dict():
|
||||
"""OLD pattern: shared parser dict with a lock — PANICS.
|
||||
|
||||
The bug is triggered by tree_sitter_language_pack.get_parser(), which
|
||||
returns the Rust/PyO3 #[pyclass(unsendable)] parser (module `_native`).
|
||||
NOT tree_sitter.Parser — that is a C extension with no thread affinity,
|
||||
so sharing it across threads works fine and never reproduces the panic.
|
||||
"""
|
||||
print("=== OLD pattern: shared dict + lock ===")
|
||||
from tree_sitter_language_pack import get_parser
|
||||
|
||||
lock = threading.Lock()
|
||||
shared_parsers: dict = {}
|
||||
|
||||
def get_shared_parser(lang: str):
|
||||
with lock:
|
||||
if lang not in shared_parsers:
|
||||
shared_parsers[lang] = get_parser(lang)
|
||||
return shared_parsers[lang]
|
||||
|
||||
# Create the parser on the main thread
|
||||
parser = get_shared_parser("python")
|
||||
print(f" Created parser on {threading.current_thread().name}")
|
||||
|
||||
# Access it from a pool thread — this triggers the panic
|
||||
def use_parser():
|
||||
thread = threading.current_thread().name
|
||||
try:
|
||||
tree = parser.parse(CODE)
|
||||
print(f" Parsed on {thread}: {tree.root_node.child_count} children")
|
||||
except BaseException as e:
|
||||
# PyO3's PanicException is a BaseException, not an Exception.
|
||||
print(f" {type(e).__name__} on {thread}: {e}")
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
|
||||
futures = [pool.submit(use_parser) for _ in range(4)]
|
||||
for f in concurrent.futures.as_completed(futures):
|
||||
try:
|
||||
f.result()
|
||||
except BaseException as e:
|
||||
print(f" Future {type(e).__name__}: {e}")
|
||||
|
||||
|
||||
def test_new_pattern_thread_local():
|
||||
"""NEW pattern: thread-local parsers — works correctly."""
|
||||
print("\n=== NEW pattern: thread-local storage ===")
|
||||
_local = threading.local()
|
||||
|
||||
def get_thread_local_parser(lang: str):
|
||||
from tree_sitter import Parser
|
||||
from tree_sitter_language_pack import get_language
|
||||
|
||||
parsers = getattr(_local, "parsers", None)
|
||||
if parsers is None:
|
||||
parsers = {}
|
||||
_local.parsers = parsers
|
||||
if lang not in parsers:
|
||||
p = Parser()
|
||||
p.language = get_language(lang)
|
||||
parsers[lang] = p
|
||||
print(f" Created parser on {threading.current_thread().name}")
|
||||
return parsers[lang]
|
||||
|
||||
def use_parser():
|
||||
thread = threading.current_thread().name
|
||||
parser = get_thread_local_parser("python")
|
||||
tree = parser.parse(CODE)
|
||||
print(f" Parsed on {thread}: {tree.root_node.child_count} children")
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
|
||||
futures = [pool.submit(use_parser) for _ in range(4)]
|
||||
for f in concurrent.futures.as_completed(futures):
|
||||
f.result()
|
||||
|
||||
print(" All tasks completed successfully!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Reproducer for tree-sitter Parser unsendable panic\n")
|
||||
|
||||
test_new_pattern_thread_local()
|
||||
|
||||
print("\nAbout to run the OLD pattern.")
|
||||
print("PyO3 raises a PanicException on the worker thread (surfaced via")
|
||||
print("future.result()) rather than killing the process. Seeing")
|
||||
print("'PanicException: _native::Parser is unsendable, but sent to")
|
||||
print("another thread' confirms the bug.\n")
|
||||
|
||||
try:
|
||||
test_old_pattern_shared_dict()
|
||||
except BaseException as e:
|
||||
print(f"\nCaught {type(e).__name__}: {e}")
|
||||
161
tests/test_code_compressor_thread_safety.py
Normal file
161
tests/test_code_compressor_thread_safety.py
Normal file
|
|
@ -0,0 +1,161 @@
|
|||
"""Tests for tree-sitter Parser thread safety (issue #562).
|
||||
|
||||
Verifies that code-aware compression works correctly when parsers are
|
||||
used from ThreadPoolExecutor workers, which is how the proxy runs
|
||||
compression in production.
|
||||
"""
|
||||
|
||||
import concurrent.futures
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.transforms.code_compressor import (
|
||||
_check_tree_sitter_available,
|
||||
_get_parser,
|
||||
is_tree_sitter_loaded,
|
||||
unload_tree_sitter,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not _check_tree_sitter_available(),
|
||||
reason="tree-sitter not installed (pip install headroom-ai[code])",
|
||||
)
|
||||
|
||||
PYTHON_CODE = "def hello():\n return 42\n"
|
||||
JS_CODE = "function hello() {\n return 42;\n}\n"
|
||||
|
||||
|
||||
class TestParserThreadLocal:
|
||||
"""Verify that _get_parser returns thread-local instances."""
|
||||
|
||||
def test_same_thread_returns_same_parser(self):
|
||||
"""Calling _get_parser twice on the same thread returns the same object."""
|
||||
p1 = _get_parser("python")
|
||||
p2 = _get_parser("python")
|
||||
assert p1 is p2
|
||||
|
||||
def test_different_languages_return_different_parsers(self):
|
||||
"""Different languages get distinct parser instances."""
|
||||
py = _get_parser("python")
|
||||
js = _get_parser("javascript")
|
||||
assert py is not js
|
||||
|
||||
def test_different_threads_get_different_parsers(self):
|
||||
"""Each thread must get its own Parser to avoid the unsendable panic."""
|
||||
main_parser = _get_parser("python")
|
||||
worker_parser = [None]
|
||||
|
||||
def grab_parser():
|
||||
worker_parser[0] = _get_parser("python")
|
||||
|
||||
t = threading.Thread(target=grab_parser)
|
||||
t.start()
|
||||
t.join()
|
||||
|
||||
assert worker_parser[0] is not None
|
||||
assert worker_parser[0] is not main_parser
|
||||
|
||||
def test_is_tree_sitter_loaded_per_thread(self):
|
||||
"""is_tree_sitter_loaded reflects the current thread's state."""
|
||||
# Ensure current thread has a parser
|
||||
_get_parser("python")
|
||||
assert is_tree_sitter_loaded() is True
|
||||
|
||||
# A fresh thread should report not loaded
|
||||
result = [None]
|
||||
|
||||
def check():
|
||||
result[0] = is_tree_sitter_loaded()
|
||||
|
||||
t = threading.Thread(target=check)
|
||||
t.start()
|
||||
t.join()
|
||||
assert result[0] is False
|
||||
|
||||
def test_unload_tree_sitter_per_thread(self):
|
||||
"""unload_tree_sitter only affects the calling thread."""
|
||||
_get_parser("python")
|
||||
assert is_tree_sitter_loaded() is True
|
||||
|
||||
# Load a parser on a worker, then unload on main — worker unaffected
|
||||
worker_loaded_after = [None]
|
||||
|
||||
def worker():
|
||||
_get_parser("python")
|
||||
# Wait for main thread to unload
|
||||
event.wait()
|
||||
worker_loaded_after[0] = is_tree_sitter_loaded()
|
||||
|
||||
event = threading.Event()
|
||||
t = threading.Thread(target=worker)
|
||||
t.start()
|
||||
|
||||
unload_tree_sitter()
|
||||
assert is_tree_sitter_loaded() is False
|
||||
|
||||
event.set()
|
||||
t.join()
|
||||
assert worker_loaded_after[0] is True
|
||||
|
||||
|
||||
class TestParserCrossThreadParsing:
|
||||
"""Verify that parsing works from ThreadPoolExecutor workers."""
|
||||
|
||||
def test_parse_from_single_worker(self):
|
||||
"""A parser created and used on the same worker thread works."""
|
||||
|
||||
def parse_on_worker():
|
||||
parser = _get_parser("python")
|
||||
tree = parser.parse(bytes(PYTHON_CODE, "utf-8"))
|
||||
return tree.root_node.child_count
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
||||
result = pool.submit(parse_on_worker).result()
|
||||
assert result > 0
|
||||
|
||||
def test_parse_from_multiple_workers(self):
|
||||
"""Multiple workers can parse concurrently without panics."""
|
||||
results = []
|
||||
errors = []
|
||||
|
||||
def parse_on_worker(code: str, lang: str):
|
||||
parser = _get_parser(lang)
|
||||
tree = parser.parse(bytes(code, "utf-8"))
|
||||
return tree.root_node.child_count
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
|
||||
futures = []
|
||||
for _ in range(8):
|
||||
futures.append(pool.submit(parse_on_worker, PYTHON_CODE, "python"))
|
||||
futures.append(pool.submit(parse_on_worker, JS_CODE, "javascript"))
|
||||
|
||||
for f in concurrent.futures.as_completed(futures):
|
||||
try:
|
||||
results.append(f.result())
|
||||
except BaseException as e:
|
||||
# PyO3's PanicException is a BaseException, not an Exception.
|
||||
errors.append(e)
|
||||
|
||||
assert not errors, f"Cross-thread parsing errors: {errors}"
|
||||
assert len(results) == 16
|
||||
assert all(r > 0 for r in results)
|
||||
|
||||
def test_repeated_parse_same_worker(self):
|
||||
"""The same worker can parse repeatedly (parser reuse works)."""
|
||||
|
||||
def parse_many():
|
||||
counts = []
|
||||
for _ in range(10):
|
||||
parser = _get_parser("python")
|
||||
tree = parser.parse(bytes(PYTHON_CODE, "utf-8"))
|
||||
counts.append(tree.root_node.child_count)
|
||||
return counts
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
||||
counts = pool.submit(parse_many).result()
|
||||
|
||||
assert len(counts) == 10
|
||||
assert all(c > 0 for c in counts)
|
||||
# All parses of the same code should give the same result
|
||||
assert len(set(counts)) == 1
|
||||
Loading…
Add table
Add a link
Reference in a new issue