chopratejas--headroom/tests/test_transforms/test_tree_sitter_thread_safety.py

"""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"
    )