diff --git a/docs/usage.md b/docs/usage.md index e6cb6cb..9f5dd81 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -757,11 +757,24 @@ Measured on one core of a laptop, single process, cold start, with the bundled p Time grows a little faster than linearly with the size of the project, because the interprocedural summaries are iterated to a fixpoint over the module graph. A warm -`--cache` brings an unchanged healthchecks to about one second, `--jobs 4` saves a -quarter of the cold time at the cost of one process's memory per job. The +`--cache` brings an unchanged healthchecks to about one second. The `healthchecks` repository is part of the regression corpus, so the analysis time of a real 45 000-line project is exercised on every change. +`--jobs` pays when a wave of the module graph holds many components: each process +prepares the project once and reads the summaries of the modules a component imports +from a directory of finished modules, the cache or one of the run's own, at the cost of +one process's memory per job. Measured on a laptop with 8 cores (4 performance, 4 +efficiency), macOS, Python 3.14, bundled plugins, cold start (October 2026): + +| Project | Python files | `-j 1` | `-j 8` | +|---|---:|---:|---:| +| healthchecks | 653 | 15.3–15.5 s | 7.9–9.2 s | +| mitmproxy | 583 | 51.6 s | 45.6 s | + +The modules of a strongly connected import cycle are analysed together by one process, +and a wave waits for its slowest component: a project whose cycles are large gains less. + ## Plugins The package ships 35 plugins, loaded by default. Security models for the standard library diff --git a/src/coretrace_python/cache.py b/src/coretrace_python/cache.py index 90ee0af..2414194 100644 --- a/src/coretrace_python/cache.py +++ b/src/coretrace_python/cache.py @@ -35,7 +35,6 @@ Mutation, NonlocalWrite, PriorCall, - SummaryIndex, SymbolRead, Target, UnknownTarget, @@ -155,16 +154,6 @@ def decode(data: Mapping[str, Any]) -> CachedModule: ) -def encode_index(index: SummaryIndex) -> dict[str, Any]: - """A summary index as plain data, to hand a worker process what it imports.""" - - return {str(symbol): _encode_summary(index.summaries[symbol]) for symbol in index.symbols} - - -def decode_index(data: Mapping[str, Any]) -> SummaryIndex: - return SummaryIndex({SymbolId(_string(k)): _decode_summary(v) for k, v in data.items()}) - - def _string(value: Any) -> str: if not isinstance(value, str): raise TypeError(f"expected a string, got {value!r}") diff --git a/src/coretrace_python/engine.py b/src/coretrace_python/engine.py index 9a3a56c..9bd6f2a 100644 --- a/src/coretrace_python/engine.py +++ b/src/coretrace_python/engine.py @@ -10,6 +10,7 @@ import hashlib import multiprocessing import sys +import tempfile from collections import Counter from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass, field, replace @@ -30,10 +31,8 @@ CachedModule, ProjectCache, decode, - decode_index, directory_fingerprint, encode, - encode_index, fingerprint, module_keys, ) @@ -508,9 +507,30 @@ def analyze_project( module_plugins = tuple(p for p in all_plugins if not isinstance(p, ProjectPlugin)) suffixes = suffix_index(affected, analysable) pool = None + spooled: tempfile.TemporaryDirectory[str] | None = None + # Where finished modules are kept for the next waves: the cache, else a directory of + # this run's own. Workers read from it what a component imports (#245). + spool = cache if jobs > 1: + if spool is None: + spooled = tempfile.TemporaryDirectory(prefix="coretrace-") + spool = ProjectCache(Path(spooled.name)) + project = _Project( + root, + tuple(plugin_roots), + tuple(plugins), + advisory_paths, + _encode_routes(routes), + tuple(sorted(escaped)), + template_calls, + tuple(sorted(analysable)), + spool.directory, + ) pool = concurrent.futures.ProcessPoolExecutor( - max_workers=process_pool_size(jobs), mp_context=multiprocessing.get_context("spawn") + max_workers=process_pool_size(jobs), + mp_context=multiprocessing.get_context("spawn"), + initializer=_start_worker, + initargs=(project,), ) try: for wave in graph.schedule(): @@ -529,29 +549,27 @@ def analyze_project( pool.submit( _analyse_batch, _Batch( - root, - tuple(plugin_roots), - tuple(plugins), {name: _path_of(files[name]) for name in sorted(component)}, - encode_index(_seed(results, graph, component)), - advisory_paths, - _encode_routes(routes), - tuple(sorted(escaped)), - template_calls, - project=tuple(sorted(analysable)), + { + name: (import_name, keys[name]) + for name, import_name in _imported_names(graph, component, results).items() + }, ), ) for component in pending ] for future in futures: computed.update({name: decode(data) for name, data in future.result().items()}) + store = cache if pool is None else spool for name, entry in computed.items(): - if cache is not None: - cache.store(keys[name], entry) + if store is not None: + store.store(keys[name], entry) results.update(computed) finally: if pool is not None: pool.shutdown() + if spooled is not None: + spooled.cleanup() index = _index(results, {name: name for name in results}) call_graphs: dict[str, CallGraph] = {} @@ -720,12 +738,20 @@ def located(name: str) -> str: def _seed(results: Mapping[str, CachedModule], graph: ModuleGraph, component: frozenset[str]) -> SummaryIndex: - """The summaries a component starts from: those of every module it imports, - transitively, all final by the time its wave runs, each under the import name the - component's code refers to it by. Of several files with one import name, only the - one a member imports directly, resolved in its own root, goes under that name: the - others are reached through summaries that inline them already, and a name no - member's root resolved is followed into neither file.""" + """The summaries a component starts from (``_imported_names``).""" + + return _index(results, _imported_names(graph, component, results)) + + +def _imported_names( + graph: ModuleGraph, component: frozenset[str], results: Mapping[str, CachedModule] +) -> dict[str, str]: + """The modules whose summaries a component starts from, each with the import name the + component's code refers to it by: every module it imports, transitively, all final + by the time its wave runs. Of several files with one import name, only the one a + member imports directly, resolved in its own root, goes under that name: the others + are reached through summaries that inline them already, and a name no member's root + resolved is followed into neither file.""" wanted: set[str] = set() pending = list(component) @@ -743,7 +769,7 @@ def _seed(results: Mapping[str, CachedModule], graph: ModuleGraph, component: fr chosen = modules if len(modules) == 1 else [name for name in modules if name in direct] if len(chosen) == 1: names[chosen[0]] = import_name - return _index(results, names) + return names def _index(results: Mapping[str, CachedModule], names: Mapping[str, str]) -> SummaryIndex: @@ -797,57 +823,118 @@ def _analyse_managers( @dataclass(frozen=True) -class _Batch: - """One component handed to a worker process, its files under the engine's module - names: it rebuilds the configuration from the project root and the plugin roots, so - only paths and the imported summaries travel.""" +class _Project: + """What every worker process needs of the project, whatever the component it + analyses: it rebuilds the configuration from the project root and the plugin roots + once, when the pool starts it (``_start_worker``), and reads the finished modules + from ``spool``, a ``ProjectCache`` directory the parent writes after each wave.""" root: Path plugin_roots: tuple[Path, ...] plugins: tuple[Plugin, ...] - paths: Mapping[str, Path] - seed: Mapping[str, Any] - advisory_paths: tuple[Path, ...] = () - routes: tuple[tuple[str, str, str, int], ...] = () - escaped: tuple[str, ...] = () - template_calls: TemplateCalls = field(default_factory=TemplateCalls) - # Every analysable module of the project, not only the component's: a curated + advisory_paths: tuple[Path, ...] + routes: tuple[tuple[str, str, str, int], ...] + escaped: tuple[str, ...] + template_calls: TemplateCalls + # Every analysable module of the project, not only a component's: a curated # entry-point suffix matches derived symbols of any project module. - project: tuple[str, ...] = () + modules: tuple[str, ...] + spool: Path -def _analyse_batch(batch: _Batch) -> dict[str, dict[str, Any]]: +@dataclass(frozen=True) +class _Batch: + """One component handed to a worker process: its files under the engine's module + names, and the modules its summaries start from, each with its import name and the + key of its entry in the spool. Only these travel, not the summaries themselves.""" + + paths: Mapping[str, Path] + imported: Mapping[str, tuple[str, str]] + + +@dataclass +class _Worker: + """The configuration a worker process prepared once, and the finished modules it has + read from the spool so far, by key.""" + + managed: Callable[[AnalysisManager], None] + plugins: tuple[Plugin, ...] + affected: Affected + suffixes: SuffixIndex + spool: ProjectCache + read: dict[str, CachedModule] = field(default_factory=dict) + + def entries(self, imported: Mapping[str, tuple[str, str]]) -> dict[str, CachedModule]: + found: dict[str, CachedModule] = {} + for name, (_, key) in imported.items(): + entry = self.read.get(key) + if entry is None: + entry = self.spool.load(key) + if entry is None: + raise RuntimeError(f"the summaries of {name} are missing from {self.spool.directory}") + self.read[key] = entry + found[name] = entry + return found + + +# The state of a worker process, set by ``_start_worker`` when the pool starts it. +_WORKER: _Worker | None = None + + +def _start_worker(project: _Project) -> None: + """Prepare this worker process for every component it will analyse (#245).""" + + global _WORKER sources = SourceManager() - managers = {name: _register_all(build_hir(sources.load_file(path))) for name, path in batch.paths.items()} - registry = load_plugins(batch.plugin_roots, next(iter(managers.values()))) - all_plugins: tuple[Plugin, ...] = (*(loaded.plugin for loaded in registry), *batch.plugins) - dependencies = resolve_dependencies(batch.root, sources) + probe = _register_all(build_hir(sources.add_source("", ""))) + registry = load_plugins(project.plugin_roots, probe) + all_plugins: tuple[Plugin, ...] = (*(loaded.plugin for loaded in registry), *project.plugins) + dependencies = resolve_dependencies(project.root, sources) file_advisories: list[tuple[Path, tuple[Advisory, ...]]] = [] - for path in batch.advisory_paths: + for path in project.advisory_paths: try: file_advisories.append((path, load_advisories(path))) except AdvisoryFileError: continue - advisories, _ = _merge_advisories(_contributions(registry, batch.plugins, batch.root, file_advisories)) + advisories, _ = _merge_advisories(_contributions(registry, project.plugins, project.root, file_advisories)) affected = affected_symbols(dependencies, advisories) - models = plugin_models(all_plugins, batch.root).extended( + models = plugin_models(all_plugins, project.root).extended( *advisory_sinks(affected), suffixes=advisory_suffix_sinks(affected) ) - routes = _decode_routes(batch.routes) - escaped = frozenset(batch.escaped) + routes = _decode_routes(project.routes) + escaped = frozenset(project.escaped) members = models.members_by_class() - for manager in managers.values(): + signatures = models.signatures() + clearing = models.clearing(escaped) + + def managed(manager: AnalysisManager) -> None: manager.provide(SecurityModelAnalysis, models) manager.provide(MembersAnalysis, members) - manager.provide(SignaturesAnalysis, models.signatures()) + manager.provide(SignaturesAnalysis, signatures) manager.provide(DependencyAnalysis, dependencies) manager.provide(RegisteredRoutes, routes) manager.provide(EscapedTemplates, escaped) - manager.provide(ClearingAnalysis, models.clearing(escaped)) - manager.provide(TemplateCallsAnalysis, batch.template_calls) - module_plugins = tuple(p for p in all_plugins if not isinstance(p, ProjectPlugin)) - suffixes = suffix_index(affected, batch.project) - results = _analyse_managers(managers, decode_index(batch.seed), module_plugins, affected, suffixes) + manager.provide(ClearingAnalysis, clearing) + manager.provide(TemplateCallsAnalysis, project.template_calls) + + _WORKER = _Worker( + managed, + tuple(p for p in all_plugins if not isinstance(p, ProjectPlugin)), + affected, + suffix_index(affected, project.modules), + ProjectCache(project.spool), + ) + + +def _analyse_batch(batch: _Batch) -> dict[str, dict[str, Any]]: + worker = _WORKER + assert worker is not None, "the pool starts every worker with _start_worker" + sources = SourceManager() + managers = {name: _register_all(build_hir(sources.load_file(path))) for name, path in batch.paths.items()} + for manager in managers.values(): + worker.managed(manager) + seed = _index(worker.entries(batch.imported), {name: import_name for name, (import_name, _) in batch.imported.items()}) + results = _analyse_managers(managers, seed, worker.plugins, worker.affected, worker.suffixes) return {name: encode(entry) for name, entry in results.items()} diff --git a/tests/test_jobs_workers.py b/tests/test_jobs_workers.py new file mode 100644 index 0000000..b0b4e83 --- /dev/null +++ b/tests/test_jobs_workers.py @@ -0,0 +1,99 @@ +"""Acceptance tests for issue #245: with ``-j N``, a worker process prepares the project +once and reads each imported module's summaries once. + +Every component used to rebuild the plugins, the dependencies, the advisories and the +models in its worker, and to be sent the summaries of every module it imports, +transitively, encoded again by the parent for each component: ``-j N`` was slower than +``-j 1`` on real projects. The report does not depend on the number of processes. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from coretrace_python import engine +from coretrace_python.cache import ProjectCache + +COUNTER_MANIFEST = ( + 'name = "counter"\nversion = "1.0.0"\nplugin_api = ">=1,<2"\nrequires = []\n' + 'provides = []\n\n[entrypoint]\nmodule = "counter"\nclass = "Counter"\n' +) +# Each process that loads the plugins instantiates this one, and says so. +COUNTER = ''' +import os +from pathlib import Path +from typing import ClassVar + +from coretrace_python.plugins import ModelPlugin + + +class Counter(ModelPlugin): + name: ClassVar[str] = "counter" + + def __init__(self) -> None: + super().__init__() + with (Path(__file__).parent / "instances.log").open("a", encoding="utf-8") as log: + log.write(f"{os.getpid()}\\n") +''' + +# Eight modules importing nothing of each other: eight components in one wave. +INDEPENDENT = {f"app/m{index}.py": f"def run{index}(value):\n return value\n" for index in range(8)} + +# ``views`` reaches ``os.system`` only through the summary of ``helpers``, analysed in an +# earlier wave, and ``more`` through the summary of ``views``. +LAYERED = { + "app/__init__.py": "", + "app/helpers.py": "import os\n\n\ndef run(cmd):\n os.system(cmd)\n", + "app/views.py": ( + "from flask import request\n\nfrom app.helpers import run\n\n\n" + "def view():\n run(request.args['c'])\n\n\ndef forward(cmd):\n run(cmd)\n" + ), + "app/more.py": "from flask import request\n\nfrom app.views import forward\n\n\ndef again():\n forward(request.args['c'])\n", +} + + +def write(root: Path, files: dict[str, str]) -> Path: + for name, text in files.items(): + (root / name).parent.mkdir(parents=True, exist_ok=True) + (root / name).write_text(text, encoding="utf-8") + return root + + +def findings(analysis: engine.ProjectAnalysis) -> list[tuple[str, str, int, str | None]]: + return sorted( + (f.rule_id, Path(str(f.span.source_id)).name, f.span.start_line, f.metadata.get("verdict")) + for f in analysis.findings + ) + + +def test_a_worker_prepares_the_project_once(tmp_path: Path) -> None: + plugin = tmp_path / "plugins" / "counter" + plugin.mkdir(parents=True) + (plugin / "plugin.toml").write_text(COUNTER_MANIFEST, encoding="utf-8") + (plugin / "counter.py").write_text(COUNTER, encoding="utf-8") + root = write(tmp_path / "project", INDEPENDENT) + + engine.analyze_project(root, [engine.BUNDLED_PLUGINS, tmp_path / "plugins"], jobs=2) + + loads = (plugin / "instances.log").read_text(encoding="utf-8").split() + in_workers = [pid for pid in loads if pid != str(os.getpid())] + assert in_workers + assert len(in_workers) == len(set(in_workers)) <= 2 + + +@pytest.mark.parametrize("cached", [False, True], ids=["no-cache", "cache"]) +def test_summaries_reach_the_modules_of_later_waves(tmp_path: Path, cached: bool) -> None: + root = write(tmp_path / "project", LAYERED) + cache = ProjectCache(tmp_path / "cache") if cached else None + + sequential = findings(engine.analyze_project(root, [engine.BUNDLED_PLUGINS])) + parallel = findings(engine.analyze_project(root, [engine.BUNDLED_PLUGINS], jobs=2, cache=cache)) + + assert parallel == sequential + assert [f for f in parallel if f[0] == "command-injection"] == [ + ("command-injection", "more.py", 7, "vulnerability"), + ("command-injection", "views.py", 7, "vulnerability"), + ]