Merge branch 'staging' into dvirdukhan/mcp-t4-index-repo

Author: DvirDukhan

api/analyzers/source_analyzer.py

@@ -123,13 +123,47 @@ def first_pass(self, path: Path, files: list[Path], ignore: list[str], graph: Gr
 
     def second_pass(self, graph: Graph, files: list[Path], path: Path) -> None:
         """
-        Recursively analyze the contents of a directory.
-
-        Args:
-            base (str): The base directory for analysis.
-            root (str): The current directory being analyzed.
-            executor (concurrent.futures.Executor): The executor to run tasks concurrently.
+        Resolve symbol references across the codebase via LSP and write the
+        resulting edges (CALLS / EXTENDS / IMPLEMENTS / RETURNS / PARAMETERS)
+        into the graph.
+
+        Symbol resolution dominates index wall-time on large repos: every
+        file's entities trigger several `lsp.request_definition` calls and
+        most of them are I/O-bound waiting on the language server.
+        multilspy's SyncLanguageServer schedules each request onto a single
+        asyncio loop running in a daemon thread (via
+        `asyncio.run_coroutine_threadsafe`), which makes concurrent calls
+        from multiple worker threads safe and lets us pipeline them.
+
+        We therefore split second_pass into two phases:
+
+          A. Parallel resolution. A bounded thread pool processes files in
+             parallel, calling `entity.resolved_symbol(...)` per entity so
+             each `Symbol.resolved_symbol` set gets populated. No graph
+             writes happen here.
+
+          B. Serial edge writes. The main thread iterates the same files
+             in their original order and emits the EXTENDS / CALLS / ...
+             edges. Keeping graph writes on one thread avoids contending on
+             FalkorDB MERGE locks and produces a deterministic edge order
+             matching the pre-parallel implementation.
+
+        Pool size is controlled by `CODE_GRAPH_INDEX_WORKERS` (default 4),
+        so resolution runs in parallel by default. This is an intentional
+        behaviour change, but the edge output is identical to the
+        single-worker path (verified on a 204-file repo): phase B always
+        writes in the original file order regardless of how phase A
+        interleaves. Set the var to 1 to run a single resolver thread (useful
+        when multilspy/jedi misbehaves under concurrency); note this still
+        dispatches through the pool rather than the main thread.
+
+        Files whose resolution raises are logged with their traceback and
+        excluded from phase B, so one bad file degrades to a logged skip
+        instead of a partial or aborted graph -- and that behaviour no longer
+        depends on the worker count.
         """
+        import os
+        from concurrent.futures import ThreadPoolExecutor, as_completed
 
         logger = MultilspyLogger()
         logger.logger.setLevel(logging.ERROR)
@@ -176,8 +210,23 @@ def second_pass(self, graph: Graph, files: list[Path], path: Path) -> None:
         lsps[".kts"] = NullLanguageServer()
         lsps[".js"] = NullLanguageServer()
         with lsps[".java"].start_server(), lsps[".py"].start_server(), lsps[".cs"].start_server(), lsps[".js"].start_server(), lsps[".kt"].start_server(), lsps[".kts"].start_server():
-            files_len = len(self.files)
-            for i, file_path in enumerate(files):
+            try:
+                n_workers = max(1, int(os.environ.get("CODE_GRAPH_INDEX_WORKERS", "4")))
+            except ValueError:
+                n_workers = 4
+
+            # Drop files we don't actually have an entry for and skip files
+            # whose language has no real LSP (NullLanguageServer provides
+            # no symbol info, so resolution would be a no-op). De-duplicate
+            # while preserving order so a path that appears twice in `files`
+            # isn't resolved concurrently by two workers racing on the same
+            # entity.resolved_symbols sets.
+            resolvable: list[Path] = []
+            seen: set[Path] = set()
+            for file_path in files:
+                if file_path in seen:
+                    continue
+                seen.add(file_path)
                 if file_path not in self.files:
                     # first_pass skipped this file (e.g. parse error, empty,
                     # untracked, or ignored after entering the candidate list).
@@ -188,13 +237,64 @@ def second_pass(self, graph: Graph, files: list[Path], path: Path) -> None:
                         file_path,
                     )
                     continue
-                # Skip symbol resolution when no real LSP is available
                 if isinstance(lsps.get(file_path.suffix), NullLanguageServer):
                     continue
+                resolvable.append(file_path)
+
+            total = len(resolvable)
+            logging.info(
+                "second_pass: resolving symbols in %d files with %d worker(s)",
+                total, n_workers,
+            )
+
+            def _resolve_file(file_path: Path) -> Path:
+                # Populate Symbol.resolved_symbol sets for every entity in
+                # this file. Pure LSP work, safe to run from worker threads
+                # because SyncLanguageServer multiplexes requests through a
+                # single asyncio loop.
+                file = self.files[file_path]
+                for _, entity in file.entities.items():
+                    entity.resolved_symbol(
+                        lambda key, symbol, fp=file_path: analyzers[fp.suffix].resolve_symbol(
+                            self.files, lsps[fp.suffix], fp, path, key, symbol
+                        )
+                    )
+                return file_path
+
+            # Phase A: resolve symbols. A single code path for every worker
+            # count keeps the failure policy identical regardless of
+            # CODE_GRAPH_INDEX_WORKERS -- a ThreadPoolExecutor with
+            # max_workers=1 simply processes one file at a time.
+            failed: set[Path] = set()
+            done = 0
+            log_every = max(1, total // 50) if total else 1
+            with ThreadPoolExecutor(max_workers=n_workers, thread_name_prefix="sa-resolve") as ex:
+                futures = {ex.submit(_resolve_file, fp): fp for fp in resolvable}
+                for fut in as_completed(futures):
+                    fp = futures[fut]
+                    try:
+                        fut.result()
+                    except Exception:
+                        # Exclude this file from phase B so we never persist a
+                        # partially resolved file; keep going so one bad file
+                        # doesn't abort the whole index.
+                        failed.add(fp)
+                        logging.warning(
+                            "second_pass: resolution failed for %s; excluding from edge writes",
+                            fp, exc_info=True,
+                        )
+                    done += 1
+                    if done % log_every == 0 or done == total:
+                        logging.info("second_pass: resolved %d/%d files", done, total)
+
+            # Phase B: serial edge writes, in the original file order so
+            # the graph is bit-identical to the single-threaded path. Files
+            # whose resolution failed are skipped (see phase A).
+            for file_path in resolvable:
+                if file_path in failed:
+                    continue
                 file = self.files[file_path]
-                logging.info(f'Processing file ({i + 1}/{files_len}): {file_path}')
                 for _, entity in file.entities.items():
-                    entity.resolved_symbol(lambda key, symbol, fp=file_path: analyzers[fp.suffix].resolve_symbol(self.files, lsps[fp.suffix], fp, path, key, symbol))
                     for key, resolved_set in entity.resolved_symbols.items():
                         for resolved in resolved_set:
                             if key == "base_class":

tests/test_py_analyzer.py

@@ -1,10 +1,22 @@
 import os
 import unittest
+from collections import Counter
 from pathlib import Path
 
 from api import SourceAnalyzer, Graph
 
 
+def _edge_snapshot(g: Graph) -> Counter:
+    """Return a multiset of all relationships in the graph keyed by
+    (rel_type, src_path, src_name, dst_path, dst_name) so two indexing runs
+    can be compared independent of node ids or write order."""
+    rows = g._query(
+        "MATCH (a)-[r]->(b) "
+        "RETURN type(r), a.path, a.name, b.path, b.name"
+    ).result_set
+    return Counter(tuple(row) for row in rows)
+
+
 class Test_PY_Analyzer(unittest.TestCase):
     def test_analyzer(self):
         path = Path(__file__).parent
@@ -73,3 +85,44 @@ def test_analyzer(self):
         self.assertIn('__init__', callers)
         self.assertIn('log', callers)
 
+    def test_index_workers_edges_are_deterministic(self):
+        """second_pass must produce identical edges regardless of
+        CODE_GRAPH_INDEX_WORKERS, since phase A only parallelises symbol
+        resolution while phase B writes edges in a fixed order.
+
+        Regression guard for the parallel-resolution path (#688): runs the
+        same fixture serially (workers=1) and in parallel (workers=4) and
+        asserts the full relationship multiset matches.
+        """
+        path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'source_files', 'py')
+
+        prev = os.environ.get("CODE_GRAPH_INDEX_WORKERS")
+        g1 = Graph("py_workers1")
+        g4 = Graph("py_workers4")
+        try:
+            os.environ["CODE_GRAPH_INDEX_WORKERS"] = "1"
+            SourceAnalyzer().analyze_local_folder(path, g1)
+            serial = _edge_snapshot(g1)
+
+            os.environ["CODE_GRAPH_INDEX_WORKERS"] = "4"
+            SourceAnalyzer().analyze_local_folder(path, g4)
+            parallel = _edge_snapshot(g4)
+
+            self.assertEqual(
+                serial, parallel,
+                "edge multiset differs between workers=1 and workers=4",
+            )
+            # Sanity: the fixture has resolvable CALLS, so guard against a
+            # vacuous all-empty comparison when an LSP is available.
+            self.assertGreater(
+                sum(serial.values()), 0,
+                "expected at least one resolved edge in the fixture",
+            )
+        finally:
+            if prev is None:
+                os.environ.pop("CODE_GRAPH_INDEX_WORKERS", None)
+            else:
+                os.environ["CODE_GRAPH_INDEX_WORKERS"] = prev
+            g1.delete()
+            g4.delete()
+