diff --git a/.tekton/on-cm-runner.yaml b/.tekton/on-cm-runner.yaml index feecd3e74..5a22261b8 100644 --- a/.tekton/on-cm-runner.yaml +++ b/.tekton/on-cm-runner.yaml @@ -11,6 +11,14 @@ metadata: pipelinesascode.tekton.dev/task: "[git-clone, ./.tekton/post-integration-evaluation.yaml]" pipelinesascode.tekton.dev/max-keep-runs: "100" spec: + # Cluster-specific: requires p4d.24xlarge GPU nodes with matching taint; update if migrating to a different cluster. + podTemplate: + tolerations: + - key: "p4d-gpu" + operator: "Exists" + effect: "NoSchedule" + nodeSelector: + node.kubernetes.io/instance-type: p4d.24xlarge timeouts: pipeline: 10h30m0s # Timeout for the entire PipelineRun params: @@ -124,6 +132,14 @@ spec: - name: source - name: basic-auth - name: exploit-iq-data + # Cluster-specific: requires p4d.24xlarge GPU nodes with matching taint; update if migrating to a different cluster. + podTemplate: + tolerations: + - key: "p4d-gpu" + operator: "Exists" + effect: "NoSchedule" + nodeSelector: + node.kubernetes.io/instance-type: p4d.24xlarge volumes: - name: google-creds-volume secret: @@ -137,7 +153,7 @@ spec: items: - key: service-ca.crt path: service-ca.crt # Mounts as a file named service-ca.crt - + # >>> THE SERVER (Sidecar) <<< sidecars: - name: server-application @@ -148,10 +164,10 @@ spec: resources: requests: cpu: "3000m" # CPU request (3 cores) - memory: "12Gi" # Memory request (8 gigabytes) + memory: "32Gi" # Memory request limits: cpu: "3000m" # CPU limit (3 cores) - memory: "32Gi" # Memory limit (16 gigabytes) + memory: "64Gi" # Memory limit volumeMounts: - name: google-creds-volume diff --git a/kustomize/base/exploit_iq_service.yaml b/kustomize/base/exploit_iq_service.yaml index 381fda706..ea4700017 100644 --- a/kustomize/base/exploit_iq_service.yaml +++ b/kustomize/base/exploit_iq_service.yaml @@ -78,10 +78,10 @@ spec: resources: limits: memory: "8Gi" - cpu: "1000m" + cpu: "2000m" requests: memory: "1Gi" - cpu: "1000m" + cpu: "2000m" env: - name: SERPAPI_API_KEY valueFrom: diff --git a/src/exploit_iq_commons/utils/dep_tree.py b/src/exploit_iq_commons/utils/dep_tree.py index ac1672b67..3773afb23 100644 --- a/src/exploit_iq_commons/utils/dep_tree.py +++ b/src/exploit_iq_commons/utils/dep_tree.py @@ -56,60 +56,7 @@ from exploit_iq_commons.logging.loggers_factory import LoggingFactory -def _available_cpus() -> int: - """Return the number of CPUs available to this process. - - Respects container CPU limits from cgroup v2/v1 before falling back to - process affinity and Python CPU APIs. - """ - - # cgroup v2 - try: - with open("/sys/fs/cgroup/cpu.max", encoding="utf-8") as f: - quota_s, period_s = f.read().strip().split() - - if quota_s != "max": - quota = int(quota_s) - period = int(period_s) - - if quota > 0 and period > 0: - return max(1, math.ceil(quota / period)) - except (OSError, ValueError): - pass - - # cgroup v1 - for base in ( - "/sys/fs/cgroup/cpu", - "/sys/fs/cgroup/cpu,cpuacct", - ): - try: - with open(f"{base}/cpu.cfs_quota_us", encoding="utf-8") as f: - quota = int(f.read().strip()) - - with open(f"{base}/cpu.cfs_period_us", encoding="utf-8") as f: - period = int(f.read().strip()) - - # cgroup v1 uses quota == -1 to mean "no CPU quota". - if quota > 0 and period > 0: - return max(1, math.ceil(quota / period)) - except (OSError, ValueError): - pass - - # CPU affinity fallback. - try: - return max(1, len(os.sched_getaffinity(0))) - except (AttributeError, OSError): - pass - - # Python 3.13+ fallback. - # Keep this after cgroup checks: in OpenShift/Kubernetes it may expose - # the full node-visible CPU set rather than the pod CPU limit. - if hasattr(os, "process_cpu_count"): - count = os.process_cpu_count() - if count: - return max(1, count) - - return os.cpu_count() or 4 +from exploit_iq_commons.utils.system_utils import available_cpus as _available_cpus def _extract_source_jar(jar: Path, dest: Path) -> None: """Extract a single source JAR into dest directory.""" diff --git a/src/exploit_iq_commons/utils/source_code_git_loader.py b/src/exploit_iq_commons/utils/source_code_git_loader.py index 040b6bf06..1ff416d41 100644 --- a/src/exploit_iq_commons/utils/source_code_git_loader.py +++ b/src/exploit_iq_commons/utils/source_code_git_loader.py @@ -16,6 +16,7 @@ import contextlib import os import tempfile +import time import threading import typing import urllib.parse @@ -79,6 +80,8 @@ logger = LoggingFactory.get_agent_logger(__name__) _gitconfig_lock = threading.Lock() +_MAX_GITCONFIG_RETRIES = 5 +_GITCONFIG_RETRY_DELAY = 0.2 def _ssh_known_hosts_path() -> Path: """Resolve the managed known_hosts path; fail closed if missing or empty. @@ -263,11 +266,28 @@ def load_repo(self): # After "RUN git config --global --add safe.directory '*'" was removed from Dockerfile, need to flag the directory that passed all the # security checks as a safe directory. + # resolve() follows symlinks so the path matches what git sees internally + # (e.g. .cache/am_cache -> /exploit-iq-data symlink in OpenShift). + resolved_path = str(self.repo_path.resolve()) + # _gitconfig_lock serializes within this process; the retry loop handles + # cross-process contention (worker subprocesses each have their own lock). with _gitconfig_lock: - global_config = git.config.get_config_path("global") - config = git.GitConfigParser(global_config, read_only=False) - with config: - config.add_value("safe", "directory", str(self.repo_path.resolve())) + for attempt in range(_MAX_GITCONFIG_RETRIES): + try: + global_config = git.config.get_config_path("global") + config = git.GitConfigParser(global_config, read_only=False) + with config: + # Dedup: skip if path already registered to prevent unbounded growth. + existing = config.get_values("safe", "directory") if config.has_section("safe") else [] + if resolved_path not in existing: + config.add_value("safe", "directory", resolved_path) + break + except OSError: + # GitPython raises OSError when another process holds .gitconfig.lock. + if attempt < _MAX_GITCONFIG_RETRIES - 1: + time.sleep(_GITCONFIG_RETRY_DELAY * (attempt + 1)) + else: + raise if not os.path.exists(self.repo_path) and self.clone_url is None: raise ValueError(f"Path {self.repo_path} does not exist") diff --git a/src/exploit_iq_commons/utils/system_utils.py b/src/exploit_iq_commons/utils/system_utils.py new file mode 100644 index 000000000..8066d2bbe --- /dev/null +++ b/src/exploit_iq_commons/utils/system_utils.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import os + + +def available_cpus() -> int: + """Return the number of CPUs available to this process. + + Respects container CPU limits from cgroup v2/v1 before falling back to + process affinity and Python CPU APIs. + """ + + # cgroup v2 + try: + with open("/sys/fs/cgroup/cpu.max", encoding="utf-8") as f: + quota_s, period_s = f.read().strip().split() + + if quota_s != "max": + quota = int(quota_s) + period = int(period_s) + + if quota > 0 and period > 0: + return max(1, math.ceil(quota / period)) + except (OSError, ValueError): + pass + + # cgroup v1 + for base in ( + "/sys/fs/cgroup/cpu", + "/sys/fs/cgroup/cpu,cpuacct", + ): + try: + with open(f"{base}/cpu.cfs_quota_us", encoding="utf-8") as f: + quota = int(f.read().strip()) + + with open(f"{base}/cpu.cfs_period_us", encoding="utf-8") as f: + period = int(f.read().strip()) + + # cgroup v1 uses quota == -1 to mean "no CPU quota". + if quota > 0 and period > 0: + return max(1, math.ceil(quota / period)) + except (OSError, ValueError): + pass + + # CPU affinity fallback. + try: + return max(1, len(os.sched_getaffinity(0))) + except (AttributeError, OSError): + pass + + # Python 3.13+ fallback. + # Keep this after cgroup checks: in OpenShift/Kubernetes it may expose + # the full node-visible CPU set rather than the pod CPU limit. + if hasattr(os, "process_cpu_count"): + count = os.process_cpu_count() + if count: + return max(1, count) + + return os.cpu_count() or 4 \ No newline at end of file diff --git a/src/exploit_iq_commons/utils/tests/test_source_code_git_loader.py b/src/exploit_iq_commons/utils/tests/test_source_code_git_loader.py new file mode 100644 index 000000000..8000de346 --- /dev/null +++ b/src/exploit_iq_commons/utils/tests/test_source_code_git_loader.py @@ -0,0 +1,248 @@ +"""Tests for safe.directory handling in SourceCodeGitLoader.load_repo.""" + +import threading +from unittest.mock import patch + +import git.config +import pytest + +from exploit_iq_commons.utils.source_code_git_loader import SourceCodeGitLoader + + +@pytest.fixture +def fake_repos(tmp_path): + """Create multiple fake repo directories with .git inside.""" + repos = [] + for i in range(5): + repo_dir = tmp_path / f"repo-{i}" + repo_dir.mkdir() + (repo_dir / ".git").mkdir() + repos.append(repo_dir) + return repos + + +@pytest.fixture +def fake_gitconfig(tmp_path, monkeypatch): + """Redirect the global gitconfig to a temp file so we don't touch real ~/.gitconfig.""" + config_path = str(tmp_path / ".gitconfig") + monkeypatch.setattr(git.config, "get_config_path", + lambda scope: config_path) + return config_path + + +def test_concurrent_load_repo_gitconfig_lock(fake_repos, fake_gitconfig): + """Concurrent load_repo calls must not crash on .gitconfig lock contention. + + Reproduces the integration test failure where multiple requests arrive + simultaneously and GitPython's GitConfigParser raises: + OSError: Lock for file '~/.gitconfig' did already exist + """ + errors = [] + barrier = threading.Barrier(len(fake_repos)) + + def load_one(repo_path): + loader = SourceCodeGitLoader( + repo_path=repo_path, + clone_url=None, + ref="main", + ) + try: + barrier.wait(timeout=5) + loader.load_repo() + except OSError as e: + if "Lock" in str(e) or "lock" in str(e): + errors.append(e) + else: + raise + except Exception: + pass + + threads = [threading.Thread(target=load_one, args=(r,)) + for r in fake_repos] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + + assert not errors, ( + f"{len(errors)} thread(s) hit .gitconfig lock contention: {errors[0]}" + ) + + +def test_concurrent_load_repo_preserves_all_safe_directories(fake_repos, fake_gitconfig): + """All repos must have their own safe.directory entry after concurrent load_repo calls. + + Reproduces the bug where set_value overwrote previous safe.directory entries, + leaving only the last repo's path. Subsequent git operations on earlier repos + failed with 'dubious ownership'. + """ + barrier = threading.Barrier(len(fake_repos)) + + def load_one(repo_path): + loader = SourceCodeGitLoader( + repo_path=repo_path, + clone_url=None, + ref="main", + ) + try: + barrier.wait(timeout=5) + loader.load_repo() + except Exception: + pass + + threads = [threading.Thread(target=load_one, args=(r,)) + for r in fake_repos] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + + config = git.GitConfigParser(fake_gitconfig, read_only=True) + with config: + entries = config.get_values("safe", "directory") + + for repo_path in fake_repos: + resolved = str(repo_path.resolve()) + assert resolved in entries, ( + f"safe.directory missing for {resolved}; entries: {entries}" + ) + + +def test_safe_directory_resolves_symlinks(tmp_path, fake_gitconfig): + """safe.directory must use the resolved (real) path, not the symlink path. + + Reproduces the 'dubious ownership' failure where .cache/am_cache was a + symlink to /exploit-iq-data. load_repo wrote the symlink path but git + checked the resolved path, so the safe.directory entry didn't match. + """ + real_dir = tmp_path / "real-data" / "git" / "repo" + real_dir.mkdir(parents=True) + (real_dir / ".git").mkdir() + + symlink_base = tmp_path / "cache" + symlink_base.symlink_to(tmp_path / "real-data") + symlink_path = symlink_base / "git" / "repo" + + loader = SourceCodeGitLoader( + repo_path=symlink_path, + clone_url=None, + ref="main", + ) + try: + loader.load_repo() + except Exception: + pass + + config = git.GitConfigParser(fake_gitconfig, read_only=True) + with config: + entries = config.get_values("safe", "directory") + + resolved = str(real_dir.resolve()) + assert resolved in entries, ( + f"safe.directory should contain resolved path {resolved}, " + f"not symlink path; entries: {entries}" + ) + + +def test_dedup_prevents_duplicate_safe_directory_entries(tmp_path, fake_gitconfig): + """Calling load_repo multiple times with the same path must not duplicate the safe.directory entry. + + Reproduces the unbounded gitconfig growth where every load_repo call + appended the same path without checking for existing entries. + """ + repo_dir = tmp_path / "repo" + repo_dir.mkdir() + (repo_dir / ".git").mkdir() + + for _ in range(3): + loader = SourceCodeGitLoader( + repo_path=repo_dir, + clone_url=None, + ref="main", + ) + try: + loader.load_repo() + except Exception: + pass + + config = git.GitConfigParser(fake_gitconfig, read_only=True) + with config: + entries = config.get_values("safe", "directory") + + resolved = str(repo_dir.resolve()) + count = entries.count(resolved) + assert count == 1, ( + f"safe.directory should contain exactly 1 entry for {resolved}, " + f"found {count}; entries: {entries}" + ) + + +@patch("time.sleep") +def test_retry_recovers_from_transient_gitconfig_lock(mock_sleep, tmp_path, fake_gitconfig): + """load_repo retries when GitConfigParser raises OSError from cross-process lock contention. + + Reproduces the worker subprocess scenario where another process holds + .gitconfig.lock temporarily, causing OSError on the first attempts + before releasing the lock. + """ + repo_dir = tmp_path / "repo" + repo_dir.mkdir() + (repo_dir / ".git").mkdir() + + real_parser = git.GitConfigParser + attempts = [] + + def flaky_parser(*args, **kwargs): + attempts.append(1) + if len(attempts) <= 2: + raise OSError("Lock for file '~/.gitconfig' did already exist") + return real_parser(*args, **kwargs) + + with patch("git.GitConfigParser", side_effect=flaky_parser): + loader = SourceCodeGitLoader( + repo_path=repo_dir, + clone_url=None, + ref="main", + ) + try: + loader.load_repo() + except OSError: + raise + except Exception: + pass + + assert len(attempts) == 3, ( + f"Expected 3 GitConfigParser attempts (2 OSError + 1 success), got {len(attempts)}" + ) + assert mock_sleep.call_count == 2, ( + f"Expected 2 sleep calls between retries, got {mock_sleep.call_count}" + ) + + +@patch("time.sleep") +def test_retry_exhaustion_reraises_oserror_after_max_attempts(mock_sleep, tmp_path, fake_gitconfig): + """load_repo re-raises OSError after exhausting all retry attempts. + + Reproduces the scenario where .gitconfig.lock is permanently held + (e.g., stale lock file from a crashed process) and all retries fail. + """ + repo_dir = tmp_path / "repo" + repo_dir.mkdir() + (repo_dir / ".git").mkdir() + + with patch( + "git.GitConfigParser", + side_effect=OSError("Lock for file '~/.gitconfig' did already exist"), + ): + loader = SourceCodeGitLoader( + repo_path=repo_dir, + clone_url=None, + ref="main", + ) + with pytest.raises(OSError, match="Lock"): + loader.load_repo() + + from exploit_iq_commons.utils.source_code_git_loader import _MAX_GITCONFIG_RETRIES + assert mock_sleep.call_count == _MAX_GITCONFIG_RETRIES - 1, ( + f"Expected {_MAX_GITCONFIG_RETRIES - 1} sleep calls, got {mock_sleep.call_count}" + ) \ No newline at end of file diff --git a/src/vuln_analysis/data_models/state.py b/src/vuln_analysis/data_models/state.py index 5264f9a69..143743d78 100644 --- a/src/vuln_analysis/data_models/state.py +++ b/src/vuln_analysis/data_models/state.py @@ -26,7 +26,6 @@ class ExploitIqEngineState(BaseModel): doc_vdb_path: str | None = None code_index_path: str | None = None cve_intel: list[CveIntel] - transitive_code_searcher: typing.Any | None = None original_input: ExploitIqEngineInput | None = None checklist_plans: dict[str, list[str]] = {} checklist_results: dict[str, list[dict[str, typing.Any]]] = {} @@ -38,4 +37,7 @@ class ExploitIqEngineState(BaseModel): current_vuln_id: str | None = None patch_results: dict[str, typing.Any] = {} uber_jar_file_threshold: int = 600 + # Store key identifying the cached TransitiveCodeSearcher in the worker process. + # Set by _ensure_searcher_built, cleared by _invalidate_and_rebuild on eviction. + worker_store_key: tuple | None = None diff --git a/src/vuln_analysis/functions/cve_clone_and_deps.py b/src/vuln_analysis/functions/cve_clone_and_deps.py index e5ac07775..29ec8ed00 100644 --- a/src/vuln_analysis/functions/cve_clone_and_deps.py +++ b/src/vuln_analysis/functions/cve_clone_and_deps.py @@ -20,6 +20,7 @@ the codebase. Expensive VDB/indexing work is deferred to cve_segmentation, which can be skipped if verify_vuln_package determines no CVE is vulnerable. """ +import asyncio from aiq.builder.builder import Builder from aiq.builder.framework_enum import LLMFrameworkEnum @@ -118,7 +119,10 @@ async def _arun(message: ExploitIqInput) -> ExploitIqEngineInput: if message.image.analysis_type == AnalysisType.SOURCE: manifest_relative_path = message.image.manifest_path ecosystem = message.image.ecosystem - embedder.clone_and_install_dependencies(si, manifest_relative_path, ecosystem) + await asyncio.to_thread( + embedder.clone_and_install_dependencies, + si, manifest_relative_path, ecosystem, + ) except Exception as e: logger.warning( "Error cloning/installing for source %s: %s", diff --git a/src/vuln_analysis/functions/cve_generate_vdbs.py b/src/vuln_analysis/functions/cve_generate_vdbs.py index a3ffe2281..1f946013f 100644 --- a/src/vuln_analysis/functions/cve_generate_vdbs.py +++ b/src/vuln_analysis/functions/cve_generate_vdbs.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import json import os import time @@ -228,7 +229,8 @@ async def _arun(message: ExploitIqInput) -> ExploitIqEngineInput: if config.ignore_code_embedding else source_infos ) - vdb_code_path, vdb_doc_path = embedder.build_vdbs( + vdb_code_path, vdb_doc_path = await asyncio.to_thread( + embedder.build_vdbs, vdb_source_infos, config.ignore_code_embedding, manifest_relative_path=manifest_relative_path, @@ -255,8 +257,11 @@ async def _arun(message: ExploitIqInput) -> ExploitIqEngineInput: image = f"{message.image.name}:{message.image.tag}" RPMDependencyManager.get_instance().container_image = image - code_index_path = _build_code_index(source_infos, manifest_relative_path=manifest_relative_path, - ecosystem=ecosystem) + code_index_path = await asyncio.to_thread( + _build_code_index, source_infos, + manifest_relative_path=manifest_relative_path, + ecosystem=ecosystem, + ) if code_index_path is None: logger.warning(("Failed to generate code index for image '%s'. " diff --git a/src/vuln_analysis/functions/cve_segmentation.py b/src/vuln_analysis/functions/cve_segmentation.py index 0907eb312..de75f94c8 100644 --- a/src/vuln_analysis/functions/cve_segmentation.py +++ b/src/vuln_analysis/functions/cve_segmentation.py @@ -21,6 +21,7 @@ cve_clone_and_deps. """ +import asyncio import json import os import time @@ -36,7 +37,6 @@ from exploit_iq_commons.logging.loggers_factory import LoggingFactory, trace_id from exploit_iq_commons.utils.credential_client import credential_context -from exploit_iq_commons.utils.dep_tree import Ecosystem from vuln_analysis.tools.tool_names import ToolNames logger = LoggingFactory.get_agent_logger(__name__) @@ -225,11 +225,12 @@ async def _arun(state: ExploitIqEngineInput) -> ExploitIqEngineInput: ) with credential_context(message.credential_id): - vdb_code_path, vdb_doc_path = embedder.build_vdbs( + vdb_code_path, vdb_doc_path = await asyncio.to_thread( + embedder.build_vdbs, source_infos, config.ignore_code_embedding, manifest_relative_path=manifest_relative_path, - ecosystem=ecosystem + ecosystem=ecosystem, ) if vdb_code_path is None and not config.ignore_code_embedding: @@ -260,7 +261,9 @@ async def _arun(state: ExploitIqEngineInput) -> ExploitIqEngineInput: try: if not config.ignore_code_index: - code_index_path = _build_code_index(source_infos, manifest_relative_path) + code_index_path = await asyncio.to_thread( + _build_code_index, source_infos, manifest_relative_path, + ) if code_index_path is None: logger.warning( "Failed to generate code index for image '%s'", diff --git a/src/vuln_analysis/tools/tests/test_concurrency.py b/src/vuln_analysis/tools/tests/test_concurrency.py index 65893eae4..59782a846 100644 --- a/src/vuln_analysis/tools/tests/test_concurrency.py +++ b/src/vuln_analysis/tools/tests/test_concurrency.py @@ -14,17 +14,16 @@ from exploit_iq_commons.data_models.input import SourceDocumentsInfo from exploit_iq_commons.utils.chain_of_calls_retriever import ChainOfCallsRetriever from exploit_iq_commons.utils.java_chain_of_calls_retriever import JavaChainOfCallsRetriever -from exploit_iq_commons.utils.transitive_code_searcher_tool import TransitiveCodeSearcher from vuln_analysis.data_models.state import ExploitIqEngineState from vuln_analysis.tools.transitive_code_search import ( _build_or_get_cached, - _build_searcher, get_git_and_pickle_base_dirs, - _searcher_cache, + _worker_built_searcher_keys, _searcher_building, _repo_build_locks, ) +from vuln_analysis.tools.worker_functions import _build_searcher _DEFAULT_THRESHOLD = ExploitIqEngineState.model_fields["uber_jar_file_threshold"].default @@ -115,42 +114,34 @@ def _make_si(git_repo: str, ref: str = "main"): def _clear_caches(): """Reset module-level caches between tests.""" - _searcher_cache.clear() + _worker_built_searcher_keys.clear() _searcher_building.clear() _repo_build_locks.clear() -def _make_java_searcher(): - """Create a mock TransitiveCodeSearcher with a Java retriever.""" - mock = MagicMock() - mock.chain_of_calls_retriever = MagicMock(spec=JavaChainOfCallsRetriever) - return mock +def _make_slow_worker_mock(build_log, sleep_secs=0.2, java=True): + """Create an async mock for run_in_cpu_process that records timing. - -def _make_nonjava_searcher(): - """Create a mock TransitiveCodeSearcher with a non-Java retriever.""" - mock = MagicMock() - mock.chain_of_calls_retriever = MagicMock(spec=ChainOfCallsRetriever) - return mock - - -def _make_slow_builder(build_log, sleep_secs=0.2, java=True): - """Create a synchronous _build_searcher replacement that records timing. - - Runs inside asyncio.to_thread (in a real thread), so uses time.sleep. + Simulates a worker build with configurable delay. Returns a store_key + (full_key for Java, repo_key for non-Java). """ lock = threading.Lock() - def slow_build(si, query, uber_jar_file_threshold=_DEFAULT_THRESHOLD, ecosystem=None, manifest_path=None, base_dirs=()): + async def mock_run(git_repo, ref, func, *args, **kwargs): + query = args[1] if len(args) > 1 else "" tag = query.split(",")[0] start = time.monotonic() - time.sleep(sleep_secs) + await asyncio.sleep(sleep_secs) end = time.monotonic() with lock: build_log.append((tag, start, end)) - return _make_java_searcher() if java else _make_nonjava_searcher() + if java: + package = query.split(",")[0].strip() + return (git_repo, ref, package) + else: + return (git_repo, ref) - return slow_build + return mock_run # --------------------------------------------------------------------------- @@ -234,8 +225,8 @@ async def test_java_same_repo_different_packages_are_serialized(): si = _make_si("https://github.com/example/repo") build_log = [] - with patch("vuln_analysis.tools.transitive_code_search._build_searcher", - side_effect=_make_slow_builder(build_log, java=True)): + with patch("vuln_analysis.tools.transitive_code_search.run_in_cpu_process", + side_effect=_make_slow_worker_mock(build_log, java=True)): task1 = asyncio.create_task(_build_or_get_cached(si=si, query="pkg-a:art-a:1.0,ClassA.foo", uber_jar_file_threshold=_DEFAULT_THRESHOLD )) task2 = asyncio.create_task(_build_or_get_cached(si=si, query="pkg-b:art-b:2.0,ClassB.bar", uber_jar_file_threshold=_DEFAULT_THRESHOLD)) await asyncio.gather(task1, task2) @@ -258,16 +249,16 @@ async def test_different_repos_can_build_concurrently(): si_b = _make_si("https://github.com/example/repo-b") build_log = [] - def slow_build(si, query, uber_jar_file_threshold=_DEFAULT_THRESHOLD, ecosystem=None, manifest_path=None, base_dirs=()): - tag = si[0].git_repo.split("/")[-1] + async def mock_run(git_repo, ref, func, *args, **kwargs): + tag = git_repo.split("/")[-1] start = time.monotonic() - time.sleep(0.2) + await asyncio.sleep(0.2) end = time.monotonic() build_log.append((tag, start, end)) - return _make_nonjava_searcher() + return (git_repo, ref) - with patch("vuln_analysis.tools.transitive_code_search._build_searcher", - side_effect=slow_build): + with patch("vuln_analysis.tools.transitive_code_search.run_in_cpu_process", + side_effect=mock_run): task1 = asyncio.create_task(_build_or_get_cached(si_a, query="pkg-a:art-a:1.0,Foo.bar", uber_jar_file_threshold=_DEFAULT_THRESHOLD, base_dirs=())) task2 = asyncio.create_task(_build_or_get_cached(si_b, query="pkg-b:art-b:2.0,Baz.qux", uber_jar_file_threshold=_DEFAULT_THRESHOLD, base_dirs=())) await asyncio.gather(task1, task2) @@ -290,21 +281,21 @@ async def test_same_key_deduplicates_build(): build_count = 0 count_lock = threading.Lock() - def counting_build(build_si, q, uber_jar_file_threshold=_DEFAULT_THRESHOLD, ecosystem=None, manifest_path=None, base_dirs=()): + async def counting_mock(git_repo, ref, func, *args, **kwargs): nonlocal build_count with count_lock: build_count += 1 - time.sleep(0.1) - return _make_nonjava_searcher() + await asyncio.sleep(0.1) + return (git_repo, ref) - with patch("vuln_analysis.tools.transitive_code_search._build_searcher", - side_effect=counting_build): + with patch("vuln_analysis.tools.transitive_code_search.run_in_cpu_process", + side_effect=counting_mock): task1 = asyncio.create_task(_build_or_get_cached(si, query=query, uber_jar_file_threshold=_DEFAULT_THRESHOLD, base_dirs=())) task2 = asyncio.create_task(_build_or_get_cached(si, query=query, uber_jar_file_threshold=_DEFAULT_THRESHOLD, base_dirs=())) results = await asyncio.gather(task1, task2) assert build_count == 1, f"Expected 1 build (deduplicated), got {build_count}" - assert results[0] is results[1], "Both tasks should return the same cached searcher" + assert results[0] == results[1], "Both tasks should return the same store key" @pytest.mark.asyncio @@ -313,25 +304,23 @@ async def test_cache_hit_skips_build(): _clear_caches() si = _make_si("https://github.com/example/repo") query = "pkg-a:art-a:1.0,ClassA.foo" - # Non-Java cache uses repo-level key (git_repo, ref) repo_key = ("https://github.com/example/repo", "main") - pre_cached = _make_nonjava_searcher() - _searcher_cache[repo_key] = pre_cached + _worker_built_searcher_keys.add(repo_key) build_count = 0 - def counting_build(build_si, q, uber_jar_file_threshold=_DEFAULT_THRESHOLD, ecosystem=None, manifest_path=None, base_dirs=()): + async def counting_mock(git_repo, ref, func, *args, **kwargs): nonlocal build_count build_count += 1 - return _make_nonjava_searcher() + return (git_repo, ref) - with patch("vuln_analysis.tools.transitive_code_search._build_searcher", - side_effect=counting_build): + with patch("vuln_analysis.tools.transitive_code_search.run_in_cpu_process", + side_effect=counting_mock): result = await _build_or_get_cached(si, query=query, uber_jar_file_threshold=_DEFAULT_THRESHOLD, base_dirs=()) assert build_count == 0, "Build should not run when cache hit exists" - assert result is pre_cached, "Should return the pre-cached searcher" + assert result == repo_key, "Should return the cached repo key" @pytest.mark.asyncio @@ -342,44 +331,43 @@ async def test_java_cache_hit_skips_build(): query = "pkg-a:art-a:1.0,ClassA.foo" full_key = ("https://github.com/example/repo", "main", "pkg-a:art-a:1.0") - pre_cached = _make_java_searcher() - _searcher_cache[full_key] = pre_cached + _worker_built_searcher_keys.add(full_key) build_count = 0 - def counting_build(build_si, q, uber_jar_file_threshold=_DEFAULT_THRESHOLD, ecosystem=None, manifest_path=None, base_dirs=()): + async def counting_mock(git_repo, ref, func, *args, **kwargs): nonlocal build_count build_count += 1 - return _make_java_searcher() + return full_key - with patch("vuln_analysis.tools.transitive_code_search._build_searcher", - side_effect=counting_build): + with patch("vuln_analysis.tools.transitive_code_search.run_in_cpu_process", + side_effect=counting_mock): result = await _build_or_get_cached(si, query=query, uber_jar_file_threshold=_DEFAULT_THRESHOLD, base_dirs=()) assert build_count == 0, "Build should not run when Java cache hit exists" - assert result is pre_cached, "Should return the pre-cached Java searcher" + assert result == full_key, "Should return the cached Java full key" @pytest.mark.asyncio async def test_build_failure_cleans_up_building_marker(): - """If _build_searcher raises, the building marker must be cleaned up.""" + """If run_in_cpu_process raises, the building marker must be cleaned up.""" _clear_caches() si = _make_si("https://github.com/example/repo") query = "pkg-a:art-a:1.0,ClassA.foo" full_key = ("https://github.com/example/repo", "main", "pkg-a:art-a:1.0") repo_key = ("https://github.com/example/repo", "main") - def failing_build(build_si, q, uber_jar_file_threshold=_DEFAULT_THRESHOLD, ecosystem=None, manifest_path=None, base_dirs=()): + async def failing_mock(git_repo, ref, func, *args, **kwargs): raise RuntimeError("Maven failed") - with patch("vuln_analysis.tools.transitive_code_search._build_searcher", - side_effect=failing_build): + with patch("vuln_analysis.tools.transitive_code_search.run_in_cpu_process", + side_effect=failing_mock): with pytest.raises(RuntimeError, match="Maven failed"): await _build_or_get_cached(si, query=query, uber_jar_file_threshold=_DEFAULT_THRESHOLD, base_dirs=()) assert full_key not in _searcher_building, "Building marker not cleaned up after failure" - assert full_key not in _searcher_cache, "Failed build should not be cached" - assert repo_key not in _searcher_cache, "Failed build should not be cached" + assert full_key not in _worker_built_searcher_keys, "Failed build should not be cached" + assert repo_key not in _worker_built_searcher_keys, "Failed build should not be cached" @pytest.mark.asyncio @@ -396,17 +384,19 @@ async def test_java_repo_lock_recheck_avoids_redundant_build(): build_count = 0 count_lock = threading.Lock() - def build_that_precaches_b(build_si, q, uber_jar_file_threshold=_DEFAULT_THRESHOLD, ecosystem=None, manifest_path=None, base_dirs=()): + async def build_that_precaches_b(git_repo, ref, func, *args, **kwargs): nonlocal build_count with count_lock: build_count += 1 current = build_count if current == 1: - time.sleep(0.1) - _searcher_cache[full_key_b] = _make_java_searcher() - return _make_java_searcher() + await asyncio.sleep(0.1) + _worker_built_searcher_keys.add(full_key_b) + query = args[1] if len(args) > 1 else "" + package = query.split(",")[0].strip() + return (git_repo, ref, package) - with patch("vuln_analysis.tools.transitive_code_search._build_searcher", + with patch("vuln_analysis.tools.transitive_code_search.run_in_cpu_process", side_effect=build_that_precaches_b): task1 = asyncio.create_task(_build_or_get_cached(si, query=query_a, uber_jar_file_threshold=_DEFAULT_THRESHOLD, base_dirs=())) await asyncio.sleep(0.01) @@ -431,15 +421,15 @@ async def test_nonjava_same_repo_different_packages_share_cache(): build_count = 0 count_lock = threading.Lock() - def counting_build(build_si, q, uber_jar_file_threshold=_DEFAULT_THRESHOLD, ecosystem=None, manifest_path=None, base_dirs=()): + async def counting_mock(git_repo, ref, func, *args, **kwargs): nonlocal build_count with count_lock: build_count += 1 - time.sleep(0.1) - return _make_nonjava_searcher() + await asyncio.sleep(0.1) + return (git_repo, ref) - with patch("vuln_analysis.tools.transitive_code_search._build_searcher", - side_effect=counting_build): + with patch("vuln_analysis.tools.transitive_code_search.run_in_cpu_process", + side_effect=counting_mock): task1 = asyncio.create_task( _build_or_get_cached(si, query="crypto/x509,ParsePKCS1PrivateKey", uber_jar_file_threshold=_DEFAULT_THRESHOLD, base_dirs=())) task2 = asyncio.create_task(_build_or_get_cached(si, query="net/http,ListenAndServe", uber_jar_file_threshold=_DEFAULT_THRESHOLD, base_dirs=())) @@ -451,7 +441,7 @@ def counting_build(build_si, q, uber_jar_file_threshold=_DEFAULT_THRESHOLD, ecos assert build_count == 1, ( f"Expected 1 build for non-Java (shared repo_key), got {build_count}" ) - assert results[0] is results[1], "Both tasks should return the same cached searcher" + assert results[0] == results[1], "Both tasks should return the same store key" @pytest.mark.asyncio @@ -462,12 +452,15 @@ async def test_nonjava_caches_under_repo_key(): repo_key = ("https://github.com/example/repo", "main") full_key = ("https://github.com/example/repo", "main", "crypto/x509") - with patch("vuln_analysis.tools.transitive_code_search._build_searcher", - return_value=_make_nonjava_searcher()): + async def mock_run(git_repo, ref, func, *args, **kwargs): + return (git_repo, ref) + + with patch("vuln_analysis.tools.transitive_code_search.run_in_cpu_process", + side_effect=mock_run): await _build_or_get_cached(si, query="crypto/x509,ParsePKCS1PrivateKey", uber_jar_file_threshold=_DEFAULT_THRESHOLD, base_dirs=()) - assert repo_key in _searcher_cache, "Non-Java should cache under repo_key" - assert full_key not in _searcher_cache, "Non-Java should NOT cache under full_key" + assert repo_key in _worker_built_searcher_keys, "Non-Java should cache under repo_key" + assert full_key not in _worker_built_searcher_keys, "Non-Java should NOT cache under full_key" # --------------------------------------------------------------------------- @@ -511,8 +504,8 @@ class TestBuildSearcherBaseDirs: def test_empty_base_dirs_uses_defaults(self): """Empty tuple should create DocumentEmbedding with default dirs.""" si = [SourceDocumentsInfo(git_repo="https://github.com/example/repo", ref="main", type="code")] - with patch("vuln_analysis.tools.transitive_code_search.get_call_of_chains_retriever") as mock_get_coc, \ - patch("vuln_analysis.tools.transitive_code_search.DocumentEmbedding") as mock_de: + with patch("vuln_analysis.tools.worker_functions._get_call_of_chains_retriever") as mock_get_coc, \ + patch("vuln_analysis.tools.worker_functions.DocumentEmbedding") as mock_de: mock_get_coc.return_value = MagicMock() _build_searcher(si, "pkg,Func", _DEFAULT_THRESHOLD, base_dirs=()) mock_de.assert_called_once_with(embedding=None) @@ -520,8 +513,8 @@ def test_empty_base_dirs_uses_defaults(self): def test_valid_base_dirs_passed_to_document_embedding(self): """Two-element tuple should pass git_directory and pickle_cache_directory.""" si = [SourceDocumentsInfo(git_repo="https://github.com/example/repo", ref="main", type="code")] - with patch("vuln_analysis.tools.transitive_code_search.get_call_of_chains_retriever") as mock_get_coc, \ - patch("vuln_analysis.tools.transitive_code_search.DocumentEmbedding") as mock_de: + with patch("vuln_analysis.tools.worker_functions._get_call_of_chains_retriever") as mock_get_coc, \ + patch("vuln_analysis.tools.worker_functions.DocumentEmbedding") as mock_de: mock_get_coc.return_value = MagicMock() _build_searcher(si, "pkg,Func", _DEFAULT_THRESHOLD, base_dirs=("/custom/git", "/custom/pickle")) mock_de.assert_called_once_with( @@ -533,8 +526,8 @@ def test_valid_base_dirs_passed_to_document_embedding(self): def test_single_element_tuple_falls_back_to_defaults(self): """Tuple with wrong length should fall back to defaults, not crash.""" si = [SourceDocumentsInfo(git_repo="https://github.com/example/repo", ref="main", type="code")] - with patch("vuln_analysis.tools.transitive_code_search.get_call_of_chains_retriever") as mock_get_coc, \ - patch("vuln_analysis.tools.transitive_code_search.DocumentEmbedding") as mock_de: + with patch("vuln_analysis.tools.worker_functions._get_call_of_chains_retriever") as mock_get_coc, \ + patch("vuln_analysis.tools.worker_functions.DocumentEmbedding") as mock_de: mock_get_coc.return_value = MagicMock() _build_searcher(si, "pkg,Func", _DEFAULT_THRESHOLD, base_dirs=("/only/one",)) mock_de.assert_called_once_with(embedding=None) @@ -542,8 +535,8 @@ def test_single_element_tuple_falls_back_to_defaults(self): def test_default_parameter_uses_defaults(self): """Omitting base_dirs entirely should use defaults.""" si = [SourceDocumentsInfo(git_repo="https://github.com/example/repo", ref="main", type="code")] - with patch("vuln_analysis.tools.transitive_code_search.get_call_of_chains_retriever") as mock_get_coc, \ - patch("vuln_analysis.tools.transitive_code_search.DocumentEmbedding") as mock_de: + with patch("vuln_analysis.tools.worker_functions._get_call_of_chains_retriever") as mock_get_coc, \ + patch("vuln_analysis.tools.worker_functions.DocumentEmbedding") as mock_de: mock_get_coc.return_value = MagicMock() _build_searcher(si, "pkg,Func", _DEFAULT_THRESHOLD) mock_de.assert_called_once_with(embedding=None) @@ -700,3 +693,75 @@ def test_route_to_llm_engine_when_empty_vuln_deps(self): route = "segmentation" if any_vulnerable else "llm_engine" assert route == "llm_engine", "Empty vuln_deps should skip segmentation" + + +# --------------------------------------------------------------------------- +# Searcher eviction recovery tests +# --------------------------------------------------------------------------- + +class TestSearcherEvictionRecovery: + """Tests that SearcherEvictedError from the worker triggers a rebuild + instead of crashing. Reproduces the bug where the main process kept + a stale key in _worker_built_searcher_keys after the worker evicted + the searcher from its LRU cache.""" + + def test_worker_function_raises_searcher_evicted_on_missing_key(self): + """Worker functions raise SearcherEvictedError when store_key was + evicted from the LRU cache (e.g., after processing 7+ repos).""" + from vuln_analysis.tools.worker_functions import ( + SearcherEvictedError, _worker_searcher_cache, worker_run_method, + ) + + _worker_searcher_cache.clear() + missing_key = ("https://github.com/example/repo", "main") + + with pytest.raises(SearcherEvictedError): + worker_run_method(missing_key, lambda s, q: None, "pkg,func") + + @pytest.mark.asyncio + async def test_invalidate_and_rebuild_triggers_new_build(self): + """_invalidate_and_rebuild removes the stale key and triggers a fresh + build via _ensure_searcher_built, which calls run_in_cpu_process.""" + from vuln_analysis.tools.transitive_code_search import _invalidate_and_rebuild + from vuln_analysis.runtime_context import ctx_state + from exploit_iq_commons.data_models.input import ( + ExploitIqEngineInput, ExploitIqInput, ImageInfoInput, + ManualSBOMInfoInput, SBOMPackage, ScanInfoInput, VulnInfo, ExploitIqInfo, + ) + from exploit_iq_commons.data_models.common import AnalysisType + + _clear_caches() + stale_key = ("https://github.com/example/repo", "main") + _worker_built_searcher_keys.add(stale_key) + + si = _make_si("https://github.com/example/repo") + state = ExploitIqEngineState( + original_input=ExploitIqEngineInput( + input=ExploitIqInput( + image=ImageInfoInput( + source_info=si, + sbom_info=ManualSBOMInfoInput(packages=[SBOMPackage(name="a", version="1.0", system="x")]), + analysis_type=AnalysisType.IMAGE, + ), + scan=ScanInfoInput(vulns=[VulnInfo(vuln_id="CVE-2025-0001")]), + ), + info=ExploitIqInfo(), + ), + code_vdb_path="", doc_vdb_path="", code_index_path="", cve_intel=[], + ) + state.worker_store_key = stale_key + ctx_state.set(state) + + build_called = False + + async def mock_run(git_repo, ref, func, *args, **kwargs): + nonlocal build_called + build_called = True + return ("https://github.com/example/repo", "main") + + with patch("vuln_analysis.tools.transitive_code_search.run_in_cpu_process", + side_effect=mock_run): + new_key, _, _ = await _invalidate_and_rebuild(stale_key, "pkg,Func", ()) + + assert build_called, "Should have triggered a new build via run_in_cpu_process" + assert state.worker_store_key is not None, "state.worker_store_key should be set after rebuild" \ No newline at end of file diff --git a/src/vuln_analysis/tools/tests/test_transitive_code_search.py b/src/vuln_analysis/tools/tests/test_transitive_code_search.py index e8d458be2..4b89440ae 100644 --- a/src/vuln_analysis/tools/tests/test_transitive_code_search.py +++ b/src/vuln_analysis/tools/tests/test_transitive_code_search.py @@ -6,7 +6,8 @@ from exploit_iq_commons.data_models.common import AnalysisType from vuln_analysis.data_models.state import ExploitIqEngineState -from vuln_analysis.tools.transitive_code_search import transitive_search, TransitiveCodeSearchToolConfig, _searcher_cache +from vuln_analysis.tools.transitive_code_search import transitive_search, TransitiveCodeSearchToolConfig, _worker_built_searcher_keys +from vuln_analysis.tools import worker_functions as _wf from exploit_iq_commons.data_models.input import (ExploitIqEngineInput, ExploitIqInput, ImageInfoInput, SourceDocumentsInfo, ManualSBOMInfoInput , SBOMPackage, ScanInfoInput, VulnInfo, ExploitIqInfo) @@ -24,6 +25,30 @@ from exploit_iq_commons.utils.git_utils import sanitize_git_url_for_path, get_repo_path_with_ref from pathlib import Path + +@pytest.fixture(autouse=True) +def inline_worker_execution(monkeypatch): + """Execute worker functions in the test process instead of spawning subprocesses. + + Tests that mock internal dependencies (retrieve_from_cache, run_command, etc.) + need those mocks to be effective in the same process where the functions run. + + Also clears both the main-process tracking set and the worker-process searcher + cache between tests so stale searchers from a previous test case (built with + different mock documents) don't leak into the next one. + """ + _worker_built_searcher_keys.clear() + _wf._worker_searcher_cache.clear() + + async def _run_inline(git_repo, ref, func, *args, **kwargs): + return func(*args, **kwargs) + + monkeypatch.setattr( + "vuln_analysis.tools.transitive_code_search.run_in_cpu_process", + _run_inline, + ) + + @pytest.fixture(autouse=True) def patch_repo_path_with_fallback(): """ @@ -254,7 +279,7 @@ def mock_file_open(*args, **kwargs): @patch('builtins.open', side_effect=mock_file_open) async def test_transitive_search_python_parameterized(mock_open, mock_run_command,test_case): """Parameterized test that runs all existing test cases with their respective configurations.""" - _searcher_cache.clear() + _worker_built_searcher_keys.clear() transitive_code_search_runner_coroutine = await get_transitive_code_runner_function() logging.basicConfig(level=logging.DEBUG) diff --git a/src/vuln_analysis/tools/transitive_code_search.py b/src/vuln_analysis/tools/transitive_code_search.py index 45b877184..c1beb334f 100644 --- a/src/vuln_analysis/tools/transitive_code_search.py +++ b/src/vuln_analysis/tools/transitive_code_search.py @@ -13,10 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import asyncio -import os -from collections import OrderedDict -from exploit_iq_commons.utils.git_utils import resolve_path_to_manifest from vuln_analysis.runtime_context import ctx_state from exploit_iq_commons.utils.transitive_code_searcher_tool import TransitiveCodeSearcher from aiq.builder.builder import Builder @@ -28,17 +25,22 @@ from langchain.docstore.document import Document from vuln_analysis.data_models.state import ExploitIqEngineState -from exploit_iq_commons.utils.document_embedding import DocumentEmbedding -from exploit_iq_commons.data_models.input import SourceDocumentsInfo -from exploit_iq_commons.utils.chain_of_calls_retriever_base import ChainOfCallsRetrieverBase -from exploit_iq_commons.utils.chain_of_calls_retriever_factory import get_chain_of_calls_retriever from exploit_iq_commons.utils.dep_tree import Ecosystem -from ..utils.error_handling_decorator import catch_pipeline_errors_async, catch_tool_errors -from ..utils.function_name_extractor import FunctionNameExtractor -from ..utils.function_name_locator import FunctionNameLocator +from ..utils.error_handling_decorator import catch_tool_errors +from ..utils.cpu_worker import run_in_cpu_process +from .worker_functions import ( + _compute_cache_keys, + worker_build, + worker_run_method, + worker_search_and_summarize, + worker_find_callers, + worker_locate_functions, + worker_find_version, + worker_check_package_in_tree, + SearcherEvictedError, +) from exploit_iq_commons.logging.loggers_factory import LoggingFactory -from exploit_iq_commons.utils.java_chain_of_calls_retriever import JavaChainOfCallsRetriever, _release_repo_data from ..functions.cve_segmentation import CVESegmentationConfig @@ -55,22 +57,17 @@ TRANSITIVE_CODE_SEARCH_TOOL_NAME = "transitive_code_search" logger = LoggingFactory.get_agent_logger(__name__) -# Maximum number of entries in the searcher cache. -# Evicts oldest entry when exceeded to bound memory consumption. -_SEARCHER_CACHE_MAX_SIZE = 6 - -# Module-level cache: avoids rebuilding the TransitiveCodeSearcher for the same -# repo+ref(+package for Java) across concurrent requests within the same process. -# Key is (git_repo, ref) for non-Java, (git_repo, ref, package_name) for Java -# because JavaDependencyTreeBuilder.build_tree uses the query's package to filter the dep graph. -_searcher_cache: OrderedDict[tuple, TransitiveCodeSearcher] = OrderedDict() +# ── Main-process coordination state ── +# Tracks which searchers have been built in worker processes, +# without holding the heavy objects in the main process. +_worker_built_searcher_keys: set[tuple] = set() +# Guards reads/writes to _worker_built_searcher_keys and _searcher_building. _searcher_cache_lock = asyncio.Lock() -# Per-key build coordination: tasks waiting for the same key await its Event -# instead of holding the global lock during the entire build. +# Per-key build dedup: tasks waiting for the same key await its Event +# instead of submitting redundant builds to the worker. _searcher_building: dict[tuple, asyncio.Event] = {} -# Per-repo build serialization: builds for the same (git_repo, ref) share filesystem -# resources (git checkout, install_dependencies, document creation). -# Concurrent builds on the same repo corrupt each other, so we serialize them. +# Per-repo build serialization: builds for the same (git_repo, ref) share +# filesystem resources. Concurrent builds on the same repo corrupt each other. _repo_build_locks: dict[tuple, asyncio.Lock] = {} _repo_locks_lock = asyncio.Lock() @@ -161,40 +158,12 @@ class FunctionLibraryVersionFinderToolConfig(FunctionBaseConfig, name=FUNCTION_L -def get_call_of_chains_retriever(documents_embedder, si, query: str, uber_jar_file_threshold: int, ecosystem: Ecosystem | None, manifest_relative_path : str | None): - documents: list[Document] - git_repo = None - code_source_info: SourceDocumentsInfo +def _extract_repo_ref(si) -> tuple[str, str]: + """Extract (git_repo, ref) from source_info list for worker routing.""" for source_info in si: if source_info.type == "code": - code_source_info = source_info - git_repo = documents_embedder.get_repo_path(source_info) - documents = documents_embedder.collect_documents(source_info) - if git_repo is None: - raise ValueError("No code source info found") - - if not ecosystem: - path_to_ecosystem_data_file = git_repo - if manifest_relative_path: - path_to_ecosystem_data_file = os.path.join(git_repo, manifest_relative_path) - with open(os.path.join(path_to_ecosystem_data_file, 'ecosystem_data.txt'), 'r', encoding='utf-8') as file: - ecosystem = file.read() - ecosystem = Ecosystem[ecosystem.upper()] - - git_repo_with_manifest = resolve_path_to_manifest(git_repo, manifest_relative_path) - logger.debug("Manifest relative path: %s", manifest_relative_path) - coc_retriever = get_chain_of_calls_retriever(ecosystem=ecosystem, - documents=documents, - manifest_path=git_repo_with_manifest, - query=query, - code_source_info=code_source_info, - uber_jar_file_threshold=uber_jar_file_threshold, - manifest_relative_path=manifest_relative_path) - # Release the raw documents list — the retriever has classified and indexed - # them into its own structures. Dropping this reference allows the GC to - # reclaim any Document objects not retained by the retriever's indexes. - del documents - return coc_retriever + return source_info.git_repo, source_info.ref + raise ValueError("No code source info found") def _get_repo_key(si) -> tuple | None: @@ -214,86 +183,28 @@ async def _get_repo_lock(si) -> asyncio.Lock: return _repo_build_locks[repo_key] -def _get_cache_keys(si, query: str, manifest_relative_path: str | None = None) -> tuple[tuple | None, tuple | None]: - """Derive cache keys from source_info list. - - Returns (repo_key, full_key): - - repo_key = (git_repo, ref) — used for non-Java ecosystems where - build_tree produces the same result regardless of package. - - full_key = (git_repo, ref, package_name) — used for Java where - JavaDependencyTreeBuilder.build_tree uses -DtargetIncludes to - filter the dependency graph per package. - """ - for source_info in si: - if source_info.type == "code": - if not manifest_relative_path: - git_repo_path = source_info.git_repo - else: - git_repo_path = f"{source_info.git_repo}/{manifest_relative_path}" - repo_key = (git_repo_path, source_info.ref) - package_name = query.split(",")[0].strip() if query else "" - full_key = (git_repo_path, source_info.ref, package_name) - return repo_key, full_key - - return None, None - - -def _build_searcher(si, query: str, uber_jar_file_threshold: int, ecosystem: Ecosystem | None = None, manifest_relative_path : str | None = None, base_dirs : tuple = ()) -> TransitiveCodeSearcher: - """Synchronous helper that builds a TransitiveCodeSearcher. - - Separated so it can be offloaded to a thread via asyncio.to_thread(), - keeping the event loop free for other tasks. - For Java, repo-level data (documents, type maps) is shared across - package-specific retrievers via _JavaRepoData (in-memory cache keyed - by (git_repo, ref), backed by pickle sub-caches on disk). - """ - if len(base_dirs) == 2: - git_base_dir_config, pickle_base_dir_config = base_dirs - documents_embedder = DocumentEmbedding( - embedding=None, - pickle_cache_directory=pickle_base_dir_config, - git_directory=git_base_dir_config - ) - else: - documents_embedder = DocumentEmbedding(embedding=None) - - coc_retriever = get_call_of_chains_retriever(documents_embedder, si, query, uber_jar_file_threshold, ecosystem, manifest_relative_path) - return TransitiveCodeSearcher(chain_of_calls_retriever=coc_retriever) - - -async def _build_or_get_cached(si, query: str, uber_jar_file_threshold: int, ecosystem: Ecosystem | None = None, manifest_relative_path : str | None = None, base_dirs: tuple = ()) -> TransitiveCodeSearcher: - """Build a TransitiveCodeSearcher, or return a cached one. - - Cache keys: - - Non-Java: (git_repo, ref) — build_tree is package-independent. - - Java: (git_repo, ref, package_name) — build_tree uses -DtargetIncludes. +async def _build_or_get_cached(si, query: str, uber_jar_file_threshold: int, ecosystem: Ecosystem | None = None, manifest_relative_path: str | None = None, base_dirs: tuple = ()) -> tuple: + """Ensure a TransitiveCodeSearcher is built in the worker process. - Since the ecosystem isn't known until after the first build, lookups check - both the repo-level key and the full (package-inclusive) key. After building, - the result is cached under the appropriate key based on the retriever type. + Returns the store_key (a tuple) that identifies the cached searcher in + the worker. The heavy TransitiveCodeSearcher object stays in the worker + process — only the lightweight key crosses the process boundary. - Coordination: + Coordination (all preserved from the original implementation): - Same-key dedup via asyncio.Event (avoids redundant builds). - - Per-repo serialization via asyncio.Lock (protects shared filesystem: - git checkout, install_dependencies, document creation). - - The expensive build runs in a thread via asyncio.to_thread(). - :param base_dirs: a tuple contains base dirs of cache - (git_base_dir, pickle_base_dir) + - Per-repo serialization via asyncio.Lock (protects shared filesystem). + - The expensive build runs in a worker process via run_in_cpu_process(). """ - repo_key, full_key = _get_cache_keys(si, query, manifest_relative_path) + repo_key, full_key = _compute_cache_keys(si, query, manifest_relative_path) + git_repo, ref = _extract_repo_ref(si) while True: async with _searcher_cache_lock: - # Check both cache keys: repo-level (non-Java) and full (Java) - if repo_key and repo_key in _searcher_cache: - logger.info("Reusing cached TransitiveCodeSearcher for %s", repo_key) - _searcher_cache.move_to_end(repo_key) - return _searcher_cache[repo_key] - if full_key and full_key in _searcher_cache: - logger.info("Reusing cached TransitiveCodeSearcher for %s", full_key) - _searcher_cache.move_to_end(full_key) - return _searcher_cache[full_key] - - # Check if another task is already building this exact key + if repo_key and repo_key in _worker_built_searcher_keys: + return repo_key + if full_key and full_key in _worker_built_searcher_keys: + return full_key + if full_key and full_key in _searcher_building: event = _searcher_building[full_key] else: @@ -302,64 +213,42 @@ async def _build_or_get_cached(si, query: str, uber_jar_file_threshold: int, eco _searcher_building[full_key] = asyncio.Event() if event: - # Another task is building this key; await it, then re-check cache logger.info("Waiting for another task to build TransitiveCodeSearcher for %s", full_key) await event.wait() continue - # This task is responsible for building. - # Acquire the per-repo lock to serialize builds that share the same - # filesystem (git checkout, install_dependencies, document creation). repo_lock = await _get_repo_lock(si) try: async with repo_lock: - # Re-check cache after acquiring repo lock — another build for a - # different package on the same repo may have finished while we waited. async with _searcher_cache_lock: - if repo_key and repo_key in _searcher_cache: - logger.info("Reusing cached TransitiveCodeSearcher for %s (after repo lock)", repo_key) - _searcher_cache.move_to_end(repo_key) + if repo_key and repo_key in _worker_built_searcher_keys: if full_key in _searcher_building: _searcher_building[full_key].set() del _searcher_building[full_key] - return _searcher_cache[repo_key] - if full_key and full_key in _searcher_cache: - logger.info("Reusing cached TransitiveCodeSearcher for %s (after repo lock)", full_key) - _searcher_cache.move_to_end(full_key) + return repo_key + if full_key and full_key in _worker_built_searcher_keys: if full_key in _searcher_building: _searcher_building[full_key].set() del _searcher_building[full_key] - return _searcher_cache[full_key] - - # Build outside _searcher_cache_lock but inside repo_lock — - # other tasks can read/write the cache while this build runs, - # but no concurrent build on the same repo's filesystem. - logger.info("Building TransitiveCodeSearcher for %s", full_key) - searcher = await asyncio.to_thread(_build_searcher, si, query, uber_jar_file_threshold, ecosystem, manifest_relative_path, base_dirs) + return full_key + + logger.info("Building TransitiveCodeSearcher for %s (in worker process)", full_key) + ecosystem_name = ecosystem.name if ecosystem else None + store_key = await run_in_cpu_process( + git_repo, ref, worker_build, + [s.model_dump() for s in si], query, + uber_jar_file_threshold, ecosystem_name, + manifest_relative_path, base_dirs, + ) async with _searcher_cache_lock: - # Cache under the appropriate key based on ecosystem - if isinstance(searcher.chain_of_calls_retriever, JavaChainOfCallsRetriever): - store_key = full_key - else: - store_key = repo_key - if store_key: - if len(_searcher_cache) >= _SEARCHER_CACHE_MAX_SIZE: - evicted_key, evicted_searcher = _searcher_cache.popitem(last=False) - logger.info("Evicted oldest TransitiveCodeSearcher cache entry: %s", evicted_key) - if (hasattr(evicted_searcher, 'chain_of_calls_retriever') - and isinstance(evicted_searcher.chain_of_calls_retriever, JavaChainOfCallsRetriever) - and hasattr(evicted_searcher.chain_of_calls_retriever, '_repo_data')): - _release_repo_data(evicted_searcher.chain_of_calls_retriever._repo_data) - _searcher_cache[store_key] = searcher + _worker_built_searcher_keys.add(store_key) if full_key and full_key in _searcher_building: _searcher_building[full_key].set() del _searcher_building[full_key] - return searcher + return store_key except BaseException: - # Clean up the building marker so waiting tasks don't hang forever. - # BaseException catches KeyboardInterrupt, SystemExit, etc. async with _searcher_cache_lock: if full_key and full_key in _searcher_building: _searcher_building[full_key].set() @@ -367,31 +256,52 @@ async def _build_or_get_cached(si, query: str, uber_jar_file_threshold: int, eco raise -async def get_transitive_code_searcher(query: str, base_dirs: tuple): +async def _ensure_searcher_built(query: str, base_dirs: tuple) -> tuple[tuple, str, str]: + """Ensure a searcher is built in the worker process. Returns (store_key, git_repo, ref). + + Uses the pipeline state to extract source_info and ecosystem. Handles + the Java package-change case by re-checking the cache. + """ state: ExploitIqEngineState = ctx_state.get() si = state.original_input.input.image.source_info threshold = state.uber_jar_file_threshold ecosystem = state.original_input.input.image.ecosystem manifest_relative_path = state.original_input.input.image.manifest_path + git_repo, ref = _extract_repo_ref(si) - if state.transitive_code_searcher is None: - state.transitive_code_searcher = await _build_or_get_cached(si, query, threshold, ecosystem, manifest_relative_path, base_dirs) - elif isinstance(state.transitive_code_searcher.chain_of_calls_retriever, JavaChainOfCallsRetriever): - # Java: different queries produce different dep trees (build_tree uses - # -DtargetIncludes for GAV queries), so rebuild when the package changes. - # Only re-check the cache when the package actually changed to avoid - # unnecessary lock acquisition on every tool call within the same request. - _, full_key = _get_cache_keys(si, query, manifest_relative_path) + if state.worker_store_key is None: + state.worker_store_key = await _build_or_get_cached( + si, query, threshold, ecosystem, manifest_relative_path, base_dirs, + ) + else: + repo_key, full_key = _compute_cache_keys(si, query, manifest_relative_path) + needs_build = False async with _searcher_cache_lock: - cached = _searcher_cache.get(full_key) - if cached is not None and cached is state.transitive_code_searcher: - pass # Same searcher, no change needed - else: - state.transitive_code_searcher = await _build_or_get_cached(si, query, threshold, ecosystem, manifest_relative_path, base_dirs) + if repo_key and repo_key in _worker_built_searcher_keys: + state.worker_store_key = repo_key + elif full_key and full_key in _worker_built_searcher_keys: + state.worker_store_key = full_key + elif full_key: + needs_build = True + if needs_build: + state.worker_store_key = await _build_or_get_cached( + si, query, threshold, ecosystem, manifest_relative_path, base_dirs, + ) + + return state.worker_store_key, git_repo, ref - # Both Java and non-Java retrievers use per-search context objects (_JavaSearchCtx / _SearchCtx) - # for mutable state, so the retriever instance is immutable after __init__ and needs no deep copy. - return state.transitive_code_searcher + +async def _invalidate_and_rebuild(store_key: tuple, query: str, base_dirs: tuple) -> tuple[tuple, str, str]: + """Remove a stale key from the main-process cache and force a rebuild. + + Called when a worker raises SearcherEvictedError (LRU eviction or restart). + """ + state: ExploitIqEngineState = ctx_state.get() + async with _searcher_cache_lock: + _worker_built_searcher_keys.discard(store_key) + state.worker_store_key = None + logger.info("Invalidated stale store_key %s, triggering rebuild", store_key) + return await _ensure_searcher_built(query, base_dirs) # The below is attached in the prompting section prior to where the tools are introduced, to clarify the tool selection strategy @@ -416,23 +326,33 @@ async def _arun(query: str) -> tuple: is_valid, validation_result = _validate_query_format(query) if not is_valid: return False, [validation_result] - transitive_code_searcher: TransitiveCodeSearcher base_dirs = await get_git_and_pickle_base_dirs(builder) - transitive_code_searcher = await get_transitive_code_searcher(validation_result, base_dirs) - found_path, call_hierarchy_list = transitive_code_searcher.search(validation_result) - # Return concise call chain summary instead of full Document objects - # to avoid blowing up the agent's context window with source code. - path_summary = _summarize_call_chain(call_hierarchy_list) - # When a package isn't in the dependency tree (e.g. stdlib like crypto/x509), - # the retriever falls back to a "dummy package" branch that synthesizes a - # function document and searches imports instead of walking the real call chain. - # This can produce misleading results — warn the agent so it doesn't treat - # import-based scanning as equivalent to full CCA reachability proof. + store_key, git_repo, ref = await _ensure_searcher_built(validation_result, base_dirs) + try: + found_path, path_summary = await run_in_cpu_process( + git_repo, ref, worker_search_and_summarize, + store_key, validation_result, + ) + except SearcherEvictedError: + store_key, git_repo, ref = await _invalidate_and_rebuild(store_key, validation_result, base_dirs) + found_path, path_summary = await run_in_cpu_process( + git_repo, ref, worker_search_and_summarize, + store_key, validation_result, + ) package_name = validation_result.split(",")[0].strip() - coc = transitive_code_searcher.chain_of_calls_retriever - if package_name not in coc.supported_packages and not any( - coc.language_parser.is_same_package(package_name, sp) for sp in coc.supported_packages): + try: + supported_packages, _ = await run_in_cpu_process( + git_repo, ref, worker_check_package_in_tree, + store_key, package_name, + ) + except SearcherEvictedError: + store_key, git_repo, ref = await _invalidate_and_rebuild(store_key, validation_result, base_dirs) + supported_packages, _ = await run_in_cpu_process( + git_repo, ref, worker_check_package_in_tree, + store_key, package_name, + ) + if not supported_packages: if path_summary: path_summary[0] = ("NOTE (import-scan fallback, package not in dependency tree): " + path_summary[0]) @@ -468,13 +388,19 @@ async def functions_usage_search(config: CallingFunctionNameExtractorToolConfig, """ @catch_tool_errors(FUNCTION_NAME_EXTRACTOR_TOOL_NAME) async def _arun(query: str) -> list: - coc_retriever: ChainOfCallsRetrieverBase - transitive_code_searcher: TransitiveCodeSearcher base_dirs = await get_git_and_pickle_base_dirs(builder) - transitive_code_searcher = await get_transitive_code_searcher(query, base_dirs) - coc_retriever = transitive_code_searcher.chain_of_calls_retriever - function_name_extractor = FunctionNameExtractor(coc_retriever) - result = function_name_extractor.fetch_list(query) + store_key, git_repo, ref = await _ensure_searcher_built(query, base_dirs) + try: + result = await run_in_cpu_process( + git_repo, ref, worker_find_callers, + store_key, query, + ) + except SearcherEvictedError: + store_key, git_repo, ref = await _invalidate_and_rebuild(store_key, query, base_dirs) + result = await run_in_cpu_process( + git_repo, ref, worker_find_callers, + store_key, query, + ) return result yield FunctionInfo.from_fn( @@ -500,15 +426,21 @@ async def _arun(query: str) -> dict: is_valid, validation_result = _validate_query_format(query) if not is_valid: return {"error": validation_result} - coc_retriever: ChainOfCallsRetrieverBase - transitive_code_searcher: TransitiveCodeSearcher base_dirs = await get_git_and_pickle_base_dirs(builder) - transitive_code_searcher = await get_transitive_code_searcher(validation_result, base_dirs) - coc_retriever = transitive_code_searcher.chain_of_calls_retriever - locator = FunctionNameLocator(coc_retriever) - result = await locator.locate_functions(validation_result) + store_key, git_repo, ref = await _ensure_searcher_built(validation_result, base_dirs) + try: + worker_result, is_package_valid, is_std_package = await run_in_cpu_process( + git_repo, ref, worker_locate_functions, + store_key, validation_result, + ) + except SearcherEvictedError: + store_key, git_repo, ref = await _invalidate_and_rebuild(store_key, validation_result, base_dirs) + worker_result, is_package_valid, is_std_package = await run_in_cpu_process( + git_repo, ref, worker_locate_functions, + store_key, validation_result, + ) package_name = validation_result.split(",", 1)[0].strip() - if locator.is_package_valid or locator.is_std_package: + if is_package_valid or is_std_package: pkg_msg = f"{FL_PACKAGE_PRESENT_PREFIX} {package_name}" package_status = "present" else: @@ -516,10 +448,10 @@ async def _arun(query: str) -> dict: package_status = "absent" return { - "ecosystem": coc_retriever.ecosystem.name, + "ecosystem": worker_result["ecosystem"], "package_msg": pkg_msg, "status": package_status, - "result": result + "result": worker_result["result"] } yield FunctionInfo.from_fn( @@ -542,36 +474,19 @@ async def library_version_finder(config: FunctionLibraryVersionFinderToolConfig, @catch_tool_errors(FUNCTION_LIBRARY_VERSION_FINDER_TOOL_NAME) async def _arun(query: str) -> dict: base_dirs = await get_git_and_pickle_base_dirs(builder) - transitive_code_searcher = await get_transitive_code_searcher(query, base_dirs) - coc_retriever = transitive_code_searcher.chain_of_calls_retriever - - # Clean the query: strip whitespace, trailing junk after newlines, then quotes (including unicode smart quotes) - cleaned_query = query.strip().split("\n")[0].strip().strip("'\"\u2018\u2019\u201c\u201d").strip() - # Search for matching packages in the dependency tree - search_term = cleaned_query.lower() - matching_packages = [] - for package in coc_retriever.supported_packages: - # Match against full GAV or individual segments - # Java GAV format: "groupId:artifactId:version" - package_lower = package.lower() - parts = package_lower.split(":") - if search_term in package_lower: - matching_packages.append(package) - - if not matching_packages: - return { - "ecosystem": coc_retriever.ecosystem.name, - "found": False, - "message": f"No package matching '{cleaned_query}' found in the dependency tree.", - "matching_packages": [] - } - - return { - "ecosystem": coc_retriever.ecosystem.name, - "found": True, - "message": f"Found {len(matching_packages)} matching package(s) in the dependency tree.", - "matching_packages": matching_packages - } + store_key, git_repo, ref = await _ensure_searcher_built(query, base_dirs) + try: + result = await run_in_cpu_process( + git_repo, ref, worker_find_version, + store_key, query, + ) + except SearcherEvictedError: + store_key, git_repo, ref = await _invalidate_and_rebuild(store_key, query, base_dirs) + result = await run_in_cpu_process( + git_repo, ref, worker_find_version, + store_key, query, + ) + return result yield FunctionInfo.from_fn( _arun, diff --git a/src/vuln_analysis/tools/worker_functions.py b/src/vuln_analysis/tools/worker_functions.py new file mode 100644 index 000000000..00e02b68e --- /dev/null +++ b/src/vuln_analysis/tools/worker_functions.py @@ -0,0 +1,272 @@ +"""Functions that run exclusively in CPU worker processes. + +Everything in this module executes inside worker subprocesses managed by +CpuWorkerPool. The main process never calls these functions directly — +it submits them via run_in_cpu_process() and receives lightweight results +via IPC. + +Module-level state (_worker_searcher_cache, _worker_cache_lock) is +per-process: each worker has its own copy, shared by its internal threads. +""" + +import asyncio +import os +import threading +from collections import OrderedDict + +from langchain.docstore.document import Document + +from exploit_iq_commons.data_models.input import SourceDocumentsInfo +from exploit_iq_commons.logging.loggers_factory import LoggingFactory +from exploit_iq_commons.utils.chain_of_calls_retriever_factory import get_chain_of_calls_retriever +from exploit_iq_commons.utils.dep_tree import Ecosystem +from exploit_iq_commons.utils.document_embedding import DocumentEmbedding +from exploit_iq_commons.utils.git_utils import resolve_path_to_manifest +from exploit_iq_commons.utils.java_chain_of_calls_retriever import JavaChainOfCallsRetriever, _release_repo_data +from exploit_iq_commons.utils.transitive_code_searcher_tool import TransitiveCodeSearcher +from vuln_analysis.utils.function_name_extractor import FunctionNameExtractor +from vuln_analysis.utils.function_name_locator import FunctionNameLocator + +logger = LoggingFactory.get_agent_logger(__name__) + +# Maximum cached searchers per worker — matches the main process limit. +_SEARCHER_CACHE_MAX_SIZE = 6 + +# Per-worker in-memory cache. Each worker process gets its own copy of these +# globals (separate address space). Internal threads within the same worker +# share them, protected by _worker_cache_lock. +_worker_searcher_cache: OrderedDict[tuple, TransitiveCodeSearcher] = OrderedDict() +_worker_cache_lock = threading.Lock() + + +class SearcherEvictedError(Exception): + """Raised when a store_key is no longer in the worker cache (LRU eviction or restart).""" + + +def _get_cached_searcher(store_key: tuple) -> TransitiveCodeSearcher: + """Look up a cached searcher, raising SearcherEvictedError if evicted.""" + with _worker_cache_lock: + try: + searcher = _worker_searcher_cache[store_key] + except KeyError: + raise SearcherEvictedError(store_key) + _worker_searcher_cache.move_to_end(store_key) + return searcher + + +def _compute_cache_keys(si, query: str, manifest_relative_path: str | None = None) -> tuple[tuple | None, tuple | None]: + """Derive (repo_key, full_key) from source_info list. + + repo_key = (git_repo, ref) — used for non-Java ecosystems. + full_key = (git_repo, ref, package_name) — used for Java where + build_tree filters the dependency graph per package. + """ + for source_info in si: + if source_info.type == "code": + if not manifest_relative_path: + git_repo_path = source_info.git_repo + else: + git_repo_path = f"{source_info.git_repo}/{manifest_relative_path}" + repo_key = (git_repo_path, source_info.ref) + package_name = query.split(",")[0].strip() if query else "" + full_key = (git_repo_path, source_info.ref, package_name) + return repo_key, full_key + return None, None + + +def _get_call_of_chains_retriever(documents_embedder, si, query: str, + uber_jar_file_threshold: int, + ecosystem: Ecosystem | None, + manifest_relative_path: str | None): + documents: list[Document] + git_repo = None + code_source_info: SourceDocumentsInfo + for source_info in si: + if source_info.type == "code": + code_source_info = source_info + git_repo = documents_embedder.get_repo_path(source_info) + documents = documents_embedder.collect_documents(source_info) + if git_repo is None: + raise ValueError("No code source info found") + + if not ecosystem: + path_to_ecosystem_data_file = git_repo + if manifest_relative_path: + path_to_ecosystem_data_file = os.path.join(git_repo, manifest_relative_path) + with open(os.path.join(path_to_ecosystem_data_file, 'ecosystem_data.txt'), + 'r', encoding='utf-8') as file: + ecosystem = file.read() + ecosystem = Ecosystem[ecosystem.upper()] + + git_repo_with_manifest = resolve_path_to_manifest(git_repo, manifest_relative_path) + logger.debug("Manifest relative path: %s", manifest_relative_path) + coc_retriever = get_chain_of_calls_retriever( + ecosystem=ecosystem, documents=documents, + manifest_path=git_repo_with_manifest, query=query, + code_source_info=code_source_info, + uber_jar_file_threshold=uber_jar_file_threshold, + manifest_relative_path=manifest_relative_path, + ) + del documents + return coc_retriever + + +def _build_searcher(si, query: str, uber_jar_file_threshold: int, + ecosystem: Ecosystem | None = None, + manifest_relative_path: str | None = None, + base_dirs: tuple = ()) -> TransitiveCodeSearcher: + """Build a TransitiveCodeSearcher from source info.""" + if len(base_dirs) == 2: + git_base_dir_config, pickle_base_dir_config = base_dirs + documents_embedder = DocumentEmbedding( + embedding=None, + pickle_cache_directory=pickle_base_dir_config, + git_directory=git_base_dir_config, + ) + else: + documents_embedder = DocumentEmbedding(embedding=None) + + coc_retriever = _get_call_of_chains_retriever( + documents_embedder, si, query, uber_jar_file_threshold, + ecosystem, manifest_relative_path, + ) + return TransitiveCodeSearcher(chain_of_calls_retriever=coc_retriever) + + +# ── Public worker functions (submitted via run_in_cpu_process) ── + + +def worker_build(source_infos_dicts: list[dict], query: str, + uber_jar_threshold: int, ecosystem_name: str | None, + manifest_path: str | None, base_dirs: tuple) -> tuple: + """Build a TransitiveCodeSearcher in the worker and cache it. + + Returns the store_key so the main process can track what's been built. + """ + source_infos = [SourceDocumentsInfo(**d) for d in source_infos_dicts] + ecosystem = Ecosystem[ecosystem_name.upper()] if ecosystem_name else None + + with _worker_cache_lock: + repo_key, full_key = _compute_cache_keys(source_infos, query, manifest_path) + if repo_key and repo_key in _worker_searcher_cache: + _worker_searcher_cache.move_to_end(repo_key) + return repo_key + if full_key and full_key in _worker_searcher_cache: + _worker_searcher_cache.move_to_end(full_key) + return full_key + + searcher = _build_searcher( + source_infos, query, uber_jar_threshold, ecosystem, + manifest_path, base_dirs, + ) + + if isinstance(searcher.chain_of_calls_retriever, JavaChainOfCallsRetriever): + store_key = full_key + else: + store_key = repo_key + + with _worker_cache_lock: + if store_key: + if len(_worker_searcher_cache) >= _SEARCHER_CACHE_MAX_SIZE: + evicted_key, evicted = _worker_searcher_cache.popitem(last=False) + logger.info("Worker evicted searcher cache entry: %s", evicted_key) + if (hasattr(evicted, 'chain_of_calls_retriever') + and isinstance(evicted.chain_of_calls_retriever, JavaChainOfCallsRetriever) + and hasattr(evicted.chain_of_calls_retriever, '_repo_data')): + _release_repo_data(evicted.chain_of_calls_retriever._repo_data) + _worker_searcher_cache[store_key] = searcher + return store_key + + +def worker_run_method(store_key: tuple, method, *args, **kwargs): + """Look up a cached searcher and call an unbound method on it. + + Accepts unbound method references (picklable in Python 3): + worker_run_method(key, TransitiveCodeSearcher.search, query) + """ + searcher = _get_cached_searcher(store_key) + return method(searcher, *args, **kwargs) + + +def _summarize_docs(docs: list) -> list[str]: + """Summarize Documents into concise strings (source :: first_line). + + Runs in the worker to avoid pickling full Document objects across IPC. + """ + summary = [] + for doc in docs: + source = doc.metadata.get('source', 'unknown') + content = doc.page_content.strip() + first_line = content.split('\n')[0].strip() if content else '' + if len(first_line) > 150: + first_line = first_line[:150] + '...' + summary.append(f"{source} :: {first_line}") + return summary + + +def worker_search_and_summarize(store_key: tuple, query: str) -> tuple[bool, list[str]]: + """CCA search + summarization in one call. Returns (found_path, summarized_chain). + + Avoids pickling full Document objects across the process boundary. + """ + searcher = _get_cached_searcher(store_key) + found_path, call_hierarchy_list = searcher.search(query) + return found_path, _summarize_docs(call_hierarchy_list) + + +def worker_find_callers(store_key: tuple, query: str) -> list: + """FCF tool — runs in worker process.""" + searcher = _get_cached_searcher(store_key) + extractor = FunctionNameExtractor(searcher.chain_of_calls_retriever) + return extractor.fetch_list(query) + + +def worker_locate_functions(store_key: tuple, query: str) -> tuple[dict, bool, bool]: + """FL tool — runs in worker process. Returns (result_dict, is_package_valid, is_std_package).""" + searcher = _get_cached_searcher(store_key) + coc_retriever = searcher.chain_of_calls_retriever + locator = FunctionNameLocator(coc_retriever) + # asyncio.run() provides an event loop for the conditional inner aiohttp call + # in quick_standard_lib_check. Safe in worker threads (no outer shared loop state). + result = asyncio.run(locator.locate_functions(query)) + return { + "ecosystem": coc_retriever.ecosystem.name, + "result": result, + }, locator.is_package_valid, locator.is_std_package + + +def worker_find_version(store_key: tuple, query: str) -> dict: + """FLVF tool — runs in worker process.""" + searcher = _get_cached_searcher(store_key) + coc_retriever = searcher.chain_of_calls_retriever + cleaned_query = query.strip().split("\n")[0].strip().strip("'\"‘’“”").strip() + search_term = cleaned_query.lower() + matching_packages = [ + pkg for pkg in coc_retriever.supported_packages + if search_term in pkg.lower() + ] + if not matching_packages: + return { + "ecosystem": coc_retriever.ecosystem.name, + "found": False, + "message": f"No package matching '{cleaned_query}' found in the dependency tree.", + "matching_packages": [] + } + return { + "ecosystem": coc_retriever.ecosystem.name, + "found": True, + "message": f"Found {len(matching_packages)} matching package(s) in the dependency tree.", + "matching_packages": matching_packages + } + + +def worker_check_package_in_tree(store_key: tuple, package_name: str) -> tuple[bool, bool]: + """Check if a package is in the dependency tree (runs in worker).""" + searcher = _get_cached_searcher(store_key) + coc = searcher.chain_of_calls_retriever + in_supported = package_name in coc.supported_packages + is_same = any( + coc.language_parser.is_same_package(package_name, sp) + for sp in coc.supported_packages + ) + return (in_supported or is_same), is_same diff --git a/src/vuln_analysis/utils/cpu_worker.py b/src/vuln_analysis/utils/cpu_worker.py new file mode 100644 index 000000000..1aeca1342 --- /dev/null +++ b/src/vuln_analysis/utils/cpu_worker.py @@ -0,0 +1,463 @@ +import asyncio +import atexit +import logging +import multiprocessing +import os +import signal +import threading +import traceback +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from queue import Empty +from typing import Any, Callable +from uuid import uuid4 + + +from exploit_iq_commons.utils.system_utils import available_cpus as _available_cpus + +logger = logging.getLogger(__name__) + +# Maximum concurrent in-flight tasks per worker before backpressure kicks in. +# Prevents unbounded memory growth when requests arrive faster than workers +# can process them. Based on pending asyncio Futures (not queue depth, +# since the worker drains the command queue instantly). +_DEFAULT_MAX_PENDING = 64 + +# Raised when a worker pool is at capacity and cannot accept more work. +class WorkerPoolBusyError(Exception): + pass + +@dataclass +class WorkerCommand: + request_id: str + func: Callable + args: tuple = () + kwargs: dict = field(default_factory=dict) + trace_id: str = "" + + +class CpuWorker: + """Runs in a dedicated subprocess. Generic executor that dispatches + incoming commands to an internal thread pool. + + The worker has its own GIL, so CPU-intensive work here does not block + the main process event loop. Internal threads allow concurrent tool + calls for CVEs routed to this worker (GIL-interleaved, same model as + the current asyncio.to_thread approach but isolated from the event loop). + """ + + def __init__(self, cmd_queue: multiprocessing.Queue, + result_queue: multiprocessing.Queue, + worker_id: int): + self._cmd_queue = cmd_queue + self._result_queue = result_queue + self._worker_id = worker_id + self._thread_pool = ThreadPoolExecutor( + max_workers=max(2, _available_cpus()), + thread_name_prefix=f"cpu-worker-{worker_id}", + ) + + def run(self) -> None: + logger.info("CpuWorker-%d started (pid=%d, threads=%d)", + self._worker_id, os.getpid(), + self._thread_pool._max_workers) + while True: + try: + cmd = self._cmd_queue.get() + except (EOFError, OSError): + break + if cmd is None: + # Shutdown sentinel received — stop accepting new work. + break + self._thread_pool.submit(self._handle_command, cmd) + # Wait for all in-flight thread pool tasks to finish before exiting. + self._thread_pool.shutdown(wait=True) + logger.info("CpuWorker-%d shutting down", self._worker_id) + + def _handle_command(self, cmd: WorkerCommand) -> None: + try: + if cmd.trace_id: + from exploit_iq_commons.logging.loggers_factory import trace_id + trace_id.set(cmd.trace_id) + result = cmd.func(*cmd.args, **cmd.kwargs) + self._result_queue.put((cmd.request_id, "ok", result)) + except Exception as exc: + tb = traceback.format_exc() + self._result_queue.put((cmd.request_id, "error", exc, tb)) + + +def _worker_entry(cmd_queue: multiprocessing.Queue, + result_queue: multiprocessing.Queue, + worker_id: int, + log_level: int = logging.WARNING) -> None: + # Workers ignore SIGINT — the main process handles it. + signal.signal(signal.SIGINT, signal.SIG_IGN) + # Inherit the main process logging configuration so worker log messages + # use the same format (traceId, timestamps) and go to the same stream. + from exploit_iq_commons.logging.loggers_factory import LoggingFactory + logging.root.setLevel(log_level) + LoggingFactory() + worker = CpuWorker(cmd_queue, result_queue, worker_id) + worker.run() + + +class CpuWorkerProxy: + """Main-process interface to a single worker subprocess. + + Bridges asyncio Futures to the multiprocessing Queue IPC: + - submit() pickles (func, args, kwargs) onto the command queue + - A background daemon thread reads the result queue and resolves + the corresponding asyncio Future via call_soon_threadsafe + """ + + def __init__(self, worker_id: int, max_pending: int = _DEFAULT_MAX_PENDING): + self._worker_id = worker_id + self._ctx = multiprocessing.get_context("spawn") + self._cmd_queue: multiprocessing.Queue = self._ctx.Queue() + self._result_queue: multiprocessing.Queue = self._ctx.Queue() + self._pending: dict[str, asyncio.Future] = {} + self._pending_lock = threading.Lock() + self._process: multiprocessing.Process | None = None + self._reader_thread: threading.Thread | None = None + self._shutdown = False + self._restarting = False + self._restart_lock = threading.Lock() + # Backpressure: reject new work when this many tasks are in-flight. + self._max_pending = max_pending + + @property + def pending_count(self) -> int: + with self._pending_lock: + return len(self._pending) + + def start(self) -> None: + self._shutdown = False + self._process = self._ctx.Process( + target=_worker_entry, + args=(self._cmd_queue, self._result_queue, self._worker_id, + logging.root.level), + daemon=True, + name=f"CpuWorker-{self._worker_id}", + ) + self._process.start() + logger.info("CpuWorkerProxy-%d: started worker pid=%d", + self._worker_id, self._process.pid) + self._reader_thread = threading.Thread( + target=self._result_reader, + daemon=True, + name=f"cpu-worker-{self._worker_id}-reader", + ) + self._reader_thread.start() + + async def submit(self, func: Callable, *args: Any, + timeout: int = 5400, **kwargs: Any) -> Any: + """Submit a function to the worker process and await its result. + + Args: + func: Top-level module function (must be picklable). + timeout: Max seconds to wait. Defaults to 5400 (CVE analysis timeout). + On timeout the Future is cancelled; the worker keeps running. + + Raises: + RuntimeError: If the pool is shutting down. + WorkerPoolBusyError: If the worker's command queue is full (backpressure). + """ + # Guard: reject new work after shutdown has started. + if self._shutdown: + raise RuntimeError( + f"CpuWorker-{self._worker_id} is shutting down; cannot accept new work" + ) + + # Backpressure: reject if too many tasks are already in-flight. + with self._pending_lock: + if len(self._pending) >= self._max_pending: + raise WorkerPoolBusyError( + f"CpuWorker-{self._worker_id} has {len(self._pending)} tasks " + f"in-flight (max {self._max_pending}); try again later" + ) + + request_id = uuid4().hex + loop = asyncio.get_running_loop() + future: asyncio.Future = loop.create_future() + with self._pending_lock: + self._pending[request_id] = future + + from exploit_iq_commons.logging.loggers_factory import trace_id as _trace_id_ctx + cmd = WorkerCommand( + request_id=request_id, func=func, args=args, kwargs=kwargs, + trace_id=_trace_id_ctx.get(), + ) + # Queue.put can block if full — offload to thread so we don't stall + # the event loop. + await asyncio.to_thread(self._cmd_queue.put, cmd) + + try: + return await asyncio.wait_for(future, timeout=timeout) + except asyncio.TimeoutError: + with self._pending_lock: + self._pending.pop(request_id, None) + # A task that ran for the full timeout (default 90 min) is stuck. + # Kill the worker and start fresh — the stuck task holds a worker + # slot hostage and likely blocks all other tasks on the same worker. + logger.error( + "CpuWorkerProxy-%d: task %s timed out after %ds, restarting worker", + self._worker_id, request_id, timeout, + ) + # Offload to thread — _restart_worker calls Process.join() and + # Process.start() which block and would stall the event loop. + await asyncio.to_thread(self._restart_worker) + raise + + def shutdown(self) -> None: + """Initiate graceful shutdown. Stops accepting new work, waits for + in-flight tasks to complete (up to 10s), then kills the worker. + + Safe to call from any thread. For async contexts, wrap in + asyncio.to_thread(proxy.shutdown) to avoid blocking the event loop. + """ + self._shutdown = True + + # Cancel all pending futures so callers don't hang. + with self._pending_lock: + for req_id, fut in list(self._pending.items()): + if not fut.done(): + try: + fut.get_loop().call_soon_threadsafe(fut.cancel) + except RuntimeError: + pass + self._pending.clear() + + # Send shutdown sentinel to worker. + try: + self._cmd_queue.put_nowait(None) + except Exception: + pass + + if self._process and self._process.is_alive(): + # Wait for worker to finish in-flight tasks. + self._process.join(timeout=10) + if self._process.is_alive(): + logger.warning("CpuWorkerProxy-%d: worker did not exit in 10s, killing", + self._worker_id) + self._process.kill() + self._process.join(timeout=5) + + def _result_reader(self) -> None: + """Background thread: bridges worker results to asyncio Futures. + + Blocks on Queue.get() — this thread's only job is to shuttle results + from the multiprocessing Queue to asyncio Futures via + call_soon_threadsafe. Exits when _shutdown is set or when + _restarting is set (a new reader will be started with fresh queues). + """ + while not self._shutdown: + if self._restarting: + return + try: + # timeout=5 is a health-check interval, not a task timeout. + # If no result arrives in 5s we check whether the worker is alive. + item = self._result_queue.get(timeout=5) + except Empty: + if self._restarting: + return + if not self._shutdown and self._process and not self._process.is_alive(): + logger.error("CpuWorkerProxy-%d: worker died, restarting", + self._worker_id) + self._restart_worker() + return + continue + except (EOFError, OSError): + if self._restarting or self._shutdown: + return + logger.error("CpuWorkerProxy-%d: result queue broken, restarting", + self._worker_id) + self._restart_worker() + return + + if len(item) == 4: + request_id, status, payload, tb_str = item + elif len(item) == 3: + request_id, status, payload = item + tb_str = None + else: + logger.error("CpuWorkerProxy-%d: unexpected result tuple length %d: %r", + self._worker_id, len(item), item) + continue + + with self._pending_lock: + future = self._pending.pop(request_id, None) + # Discard late results from timed-out or cancelled tasks. + if future is None or future.done(): + continue + + loop = future.get_loop() + if status == "ok": + loop.call_soon_threadsafe(future.set_result, payload) + else: + if tb_str: + logger.error("CpuWorkerProxy-%d: worker exception:\n%s", + self._worker_id, tb_str) + loop.call_soon_threadsafe(future.set_exception, payload) + + def _restart_worker(self) -> None: + # Non-blocking acquire: if timeout thread and reader thread both detect + # a crash, only one restarts; the other returns immediately. + if not self._restart_lock.acquire(blocking=False): + return + try: + self._restart_worker_locked() + finally: + self._restart_lock.release() + + def _restart_worker_locked(self) -> None: + logger.info("CpuWorkerProxy-%d: restarting worker process", + self._worker_id) + self._restarting = True + + # Fail all pending futures so callers don't hang. + with self._pending_lock: + for req_id, fut in list(self._pending.items()): + if not fut.done(): + exc = RuntimeError( + f"CpuWorker-{self._worker_id} crashed; request {req_id} lost" + ) + fut.get_loop().call_soon_threadsafe(fut.set_exception, exc) + self._pending.clear() + + # Reap the dead process to prevent zombies. + if self._process is not None: + self._process.join(timeout=0) + + # Recreate queues — old ones may be corrupted after a crash. + try: + self._cmd_queue.close() + self._result_queue.close() + except Exception: + pass + + self._cmd_queue = self._ctx.Queue() + self._result_queue = self._ctx.Queue() + + self._process = self._ctx.Process( + target=_worker_entry, + args=(self._cmd_queue, self._result_queue, self._worker_id, + logging.root.level), + daemon=True, + name=f"CpuWorker-{self._worker_id}", + ) + self._process.start() + logger.info("CpuWorkerProxy-%d: restarted worker pid=%d", + self._worker_id, self._process.pid) + + self._restarting = False + # Start a new reader thread for the new result queue. + self._reader_thread = threading.Thread( + target=self._result_reader, + daemon=True, + name=f"cpu-worker-{self._worker_id}-reader", + ) + self._reader_thread.start() + + +class CpuWorkerPool: + """Manages N worker processes with repo-based routing. + + CVEs are routed by (git_repo, ref) so that all tool calls for the + same image land on the same worker (sharing the in-memory searcher + cache). Different images go to different workers for true + process-level parallelism. + """ + + def __init__(self, num_workers: int | None = None): + if num_workers is None: + num_workers = max(2, _available_cpus()) + self._workers = [CpuWorkerProxy(worker_id=i) + for i in range(num_workers)] + self._assignments: dict[tuple, int] = {} + self._lock = asyncio.Lock() + self._started = False + self._shutting_down = False + + def start(self) -> None: + for w in self._workers: + w.start() + self._started = True + logger.info("CpuWorkerPool started with %d workers", len(self._workers)) + + async def submit(self, git_repo: str, ref: str, + func: Callable, *args: Any, **kwargs: Any) -> Any: + """Submit work to the routed worker. Raises RuntimeError if shutting down, + WorkerPoolBusyError if the target worker's queue is full.""" + if self._shutting_down: + raise RuntimeError("CpuWorkerPool is shutting down; cannot accept new work") + if not self._started: + self.start() + worker = await self._assign(git_repo, ref) + return await worker.submit(func, *args, **kwargs) + + async def _assign(self, git_repo: str, ref: str) -> CpuWorkerProxy: + key = (git_repo, ref) + async with self._lock: + if key not in self._assignments: + # Assign to the least-loaded worker. When loads are equal, + # distribute across workers by cycling through candidates. + loads = [w.pending_count for w in self._workers] + min_load = min(loads) + candidates = [i for i, load in enumerate(loads) if load == min_load] + idx = candidates[len(self._assignments) % len(candidates)] + self._assignments[key] = idx + logger.info("CpuWorkerPool: routing (%s, %s) → worker-%d", + git_repo, ref, idx) + return self._workers[self._assignments[key]] + + def shutdown(self) -> None: + """Graceful shutdown: stop accepting work, wait for in-flight tasks, + terminate workers. Blocks the calling thread — for async contexts, + wrap in asyncio.to_thread(pool.shutdown).""" + self._shutting_down = True + for w in self._workers: + w.shutdown() + self._started = False + logger.info("CpuWorkerPool shut down") + + async def async_shutdown(self) -> None: + """Non-blocking shutdown for async contexts. Wraps the blocking + shutdown() in a thread so the event loop stays responsive.""" + await asyncio.to_thread(self.shutdown) + + +_pool: CpuWorkerPool | None = None +_pool_lock = asyncio.Lock() + + +async def run_in_cpu_process(git_repo: str, ref: str, + func: Callable, *args: Any, **kwargs: Any) -> Any: + """Submit a top-level function to the routed worker process and await result.""" + global _pool + if _pool is None: + async with _pool_lock: + if _pool is None: + _pool = CpuWorkerPool() + _pool.start() + return await _pool.submit(git_repo, ref, func, *args, **kwargs) + + +def shutdown_pool() -> None: + """Synchronous shutdown — safe to call from atexit or signal handlers.""" + global _pool + if _pool is not None: + _pool.shutdown() + _pool = None + + +# Register atexit handler so workers are cleaned up on normal exit. +atexit.register(shutdown_pool) + + +def _sigterm_handler(signum, frame): + """Ensure worker cleanup on SIGTERM (pod rolling update / scale-down).""" + shutdown_pool() + raise SystemExit(0) + + +signal.signal(signal.SIGTERM, _sigterm_handler) \ No newline at end of file diff --git a/tests/test_cpu_worker.py b/tests/test_cpu_worker.py new file mode 100644 index 000000000..deb6a55a1 --- /dev/null +++ b/tests/test_cpu_worker.py @@ -0,0 +1,369 @@ +import asyncio +import os +import signal +import time + +import pytest + +from vuln_analysis.utils.cpu_worker import ( + CpuWorkerPool, + CpuWorkerProxy, + WorkerPoolBusyError, + run_in_cpu_process, + shutdown_pool, +) + + +def _add(a, b): + return a + b + + +def _sleep_and_return(seconds, value): + time.sleep(seconds) + return value + + +def _raise_error(msg): + raise ValueError(msg) + + +def _cpu_busy(duration): + end = time.monotonic() + duration + while time.monotonic() < end: + pass + return "done" + + +def _get_pid(): + return os.getpid() + + +_WORKER_STATE = {} + + +def _set_state(key, value): + _WORKER_STATE[key] = value + return True + + +def _get_state(key): + return _WORKER_STATE.get(key) + + +@pytest.fixture +def pool(): + p = CpuWorkerPool(num_workers=2) + p.start() + yield p + p.shutdown() + + +@pytest.fixture(autouse=True) +def cleanup_global_pool(): + yield + shutdown_pool() + + +class TestWorkerPoolLifecycle: + + @pytest.mark.asyncio + async def test_start_and_shutdown(self, pool): + result = await pool.submit("repo", "ref", _add, 2, 3) + assert result == 5 + pool.shutdown() + + @pytest.mark.asyncio + async def test_submit_and_receive_result(self, pool): + result = await pool.submit("repo", "ref", _add, 10, 20) + assert result == 30 + + @pytest.mark.asyncio + async def test_exception_propagates(self, pool): + with pytest.raises(ValueError, match="test error"): + await pool.submit("repo", "ref", _raise_error, "test error") + + @pytest.mark.asyncio + async def test_submit_timeout(self, pool): + with pytest.raises(asyncio.TimeoutError): + await pool.submit("repo", "ref", _sleep_and_return, 10, "late", + timeout=1) + + +class TestRepoRouting: + + @pytest.mark.asyncio + async def test_same_repo_goes_to_same_worker(self, pool): + pid1 = await pool.submit("repoA", "ref1", _get_pid) + pid2 = await pool.submit("repoA", "ref1", _get_pid) + assert pid1 == pid2 + + @pytest.mark.asyncio + async def test_different_repos_may_go_to_different_workers(self, pool): + pid1 = await pool.submit("repoA", "ref1", _get_pid) + pid2 = await pool.submit("repoB", "ref2", _get_pid) + assert pid1 != pid2 + + @pytest.mark.asyncio + async def test_least_loaded_assignment(self, pool): + await pool.submit("repo1", "ref", _add, 1, 1) + await pool.submit("repo2", "ref", _add, 1, 1) + + assignments = pool._assignments + workers_used = set(assignments.values()) + assert len(workers_used) == 2 + + +class TestConcurrency: + + @pytest.mark.asyncio + async def test_concurrent_submissions(self, pool): + tasks = [ + pool.submit("repoA", "ref", _sleep_and_return, 0.5, i) + for i in range(5) + ] + results = await asyncio.gather(*tasks) + assert sorted(results) == [0, 1, 2, 3, 4] + + @pytest.mark.asyncio + async def test_event_loop_responsive_during_cpu_work(self, pool): + cpu_task = asyncio.create_task( + pool.submit("repo", "ref", _cpu_busy, 2) + ) + + async def check_responsive(): + for _ in range(10): + await asyncio.sleep(0.1) + return "responsive" + + responsive_task = asyncio.create_task(check_responsive()) + responsive_result = await responsive_task + assert responsive_result == "responsive" + + cpu_result = await cpu_task + assert cpu_result == "done" + + +class TestWorkerCachePersistence: + + @pytest.mark.asyncio + async def test_cache_persists_across_calls(self, pool): + await pool.submit("repo", "ref", _set_state, "key1", "value1") + result = await pool.submit("repo", "ref", _get_state, "key1") + assert result == "value1" + + @pytest.mark.asyncio + async def test_different_repo_has_separate_cache(self, pool): + await pool.submit("repoA", "ref", _set_state, "shared_key", "A") + await pool.submit("repoB", "ref", _set_state, "shared_key", "B") + + result_a = await pool.submit("repoA", "ref", _get_state, "shared_key") + result_b = await pool.submit("repoB", "ref", _get_state, "shared_key") + assert result_a == "A" + assert result_b == "B" + + +class TestWorkerCrashRecovery: + + @pytest.mark.asyncio + async def test_worker_restarts_after_crash(self): + proxy = CpuWorkerProxy(worker_id=99) + proxy.start() + try: + pid_before = await proxy.submit(_get_pid) + os.kill(pid_before, signal.SIGKILL) + # Retry until the restarted worker responds. + # The reader thread polls every 5s, and spawn-context process + # creation can take several seconds with heavy imports. + pid_after = None + for attempt in range(10): + await asyncio.sleep(3) + try: + pid_after = await proxy.submit(_get_pid, timeout=10) + break + except (RuntimeError, asyncio.TimeoutError): + continue + assert pid_after is not None, "Worker did not restart within 30s" + assert pid_after != pid_before + finally: + proxy.shutdown() + + +class TestGlobalPool: + + @pytest.mark.asyncio + async def test_run_in_cpu_process(self): + result = await run_in_cpu_process("repo", "ref", _add, 100, 200) + assert result == 300 + + +class _MockSearcher: + def search(self, query): + return f"found:{query}" + + def locate(self, func, pkg): + return [f"{pkg}.{func}"] + + +_mock_cache: dict = {} + + +def _store_mock_searcher(key): + _mock_cache[key] = _MockSearcher() + return True + + +def _run_on_mock_searcher(key, method_name, *args): + searcher = _mock_cache[key] + return getattr(searcher, method_name)(*args) + + +class TestWorkerRunMethod: + + @pytest.mark.asyncio + async def test_generic_method_dispatch(self, pool): + store_key = ("test_repo", "test_ref") + await pool.submit("test_repo", "test_ref", _store_mock_searcher, store_key) + + search_result = await pool.submit( + "test_repo", "test_ref", _run_on_mock_searcher, + store_key, "search", "my_query", + ) + assert search_result == "found:my_query" + + locate_result = await pool.submit( + "test_repo", "test_ref", _run_on_mock_searcher, + store_key, "locate", "func1", "pkg1", + ) + assert locate_result == ["pkg1.func1"] + + +# ── Helper for backpressure test ── + +def _sleep_forever(): + """Blocks the worker indefinitely — used to fill the queue.""" + import threading + threading.Event().wait() + + +class TestGracefulShutdown: + + @pytest.mark.asyncio + async def test_submit_after_shutdown_raises(self): + proxy = CpuWorkerProxy(worker_id=50) + proxy.start() + proxy.shutdown() + with pytest.raises(RuntimeError, match="shutting down"): + await proxy.submit(_add, 1, 2) + + @pytest.mark.asyncio + async def test_shutdown_cancels_pending_futures(self): + proxy = CpuWorkerProxy(worker_id=51) + proxy.start() + try: + # Submit a long task then immediately shut down. + task = asyncio.create_task( + proxy.submit(_sleep_and_return, 60, "late") + ) + await asyncio.sleep(0.5) + proxy.shutdown() + with pytest.raises((asyncio.CancelledError, RuntimeError)): + await task + finally: + try: + proxy.shutdown() + except Exception: + pass + + @pytest.mark.asyncio + async def test_pool_submit_after_shutdown_raises(self): + p = CpuWorkerPool(num_workers=1) + p.start() + p.shutdown() + with pytest.raises(RuntimeError, match="shutting down"): + await p.submit("repo", "ref", _add, 1, 2) + + +class TestBackpressure: + + @pytest.mark.asyncio + async def test_too_many_pending_raises_busy_error(self): + # max_pending=2 means the 3rd concurrent submit is rejected. + proxy = CpuWorkerProxy(worker_id=60, max_pending=2) + proxy.start() + try: + # Submit 2 slow tasks — both become pending futures. + tasks = [ + asyncio.create_task(proxy.submit(_sleep_and_return, 30, i)) + for i in range(2) + ] + await asyncio.sleep(0.5) + # The 3rd submit exceeds max_pending. + with pytest.raises(WorkerPoolBusyError, match="in-flight"): + await proxy.submit(_add, 1, 2) + for t in tasks: + t.cancel() + finally: + proxy.shutdown() + + +class TestTimeoutRestartsWorker: + + @pytest.mark.asyncio + async def test_timeout_kills_and_restarts_worker(self): + proxy = CpuWorkerProxy(worker_id=70) + proxy.start() + try: + pid_before = await proxy.submit(_get_pid) + # Submit a stuck task with a short timeout. + with pytest.raises(asyncio.TimeoutError): + await proxy.submit(_sleep_and_return, 60, "stuck", timeout=2) + # Worker should have been restarted — wait for the new one. + await asyncio.sleep(3) + pid_after = await proxy.submit(_get_pid, timeout=15) + assert pid_after != pid_before + finally: + proxy.shutdown() + + +class TestMalformedResultTuple: + + @pytest.mark.asyncio + async def test_malformed_result_tuple_is_logged_and_skipped(self): + proxy = CpuWorkerProxy(worker_id=85) + proxy.start() + try: + # Inject a malformed 2-element tuple directly into the result queue. + # The reader thread should log an error and skip it without crashing. + proxy._result_queue.put(("bad_id", "ok")) + # A subsequent valid submission should still succeed. + result = await proxy.submit(_add, 1, 2, timeout=10) + assert result == 3 + finally: + proxy.shutdown() + + +class TestZombieReap: + + @pytest.mark.asyncio + async def test_restart_reaps_dead_process(self): + proxy = CpuWorkerProxy(worker_id=80) + proxy.start() + try: + pid = await proxy.submit(_get_pid) + os.kill(pid, signal.SIGKILL) + # Wait for detection + restart. + for _ in range(10): + await asyncio.sleep(3) + try: + new_pid = await proxy.submit(_get_pid, timeout=10) + break + except (RuntimeError, asyncio.TimeoutError): + continue + else: + pytest.fail("Worker did not restart after crash") + assert new_pid != pid + # The old process should be reaped (no zombie). + # On Linux, a reaped process's /proc entry disappears. + assert not os.path.exists(f"/proc/{pid}") + finally: + proxy.shutdown()