From 187e3c14c51d3e2a6448d32c556079670361708a Mon Sep 17 00:00:00 2001 From: Marius Ciepluch <11855163+norandom@users.noreply.github.com> Date: Sun, 4 Oct 2026 16:41:22 +0200 Subject: [PATCH] engine: one resident Joern session per pin (frontend + one query JVM instead of four), ActiveProcessorCount=2 on the cluster Implemented with Codex. Co-Authored-By: Claude Opus 5.5 --- benchmarks/learn/engine_trace_worker.py | 171 +++++++++++++++++++++++- src/openultrasast/cpg/session.py | 32 ++++- tests/test_cpg_session.py | 39 ++++++ tests/test_engine_task_entry.py | 2 +- tests/test_engine_trace.py | 121 +++++++++++++++++ 5 files changed, 358 insertions(+), 7 deletions(-) diff --git a/benchmarks/learn/engine_trace_worker.py b/benchmarks/learn/engine_trace_worker.py index 6304a86..1f445e8 100644 --- a/benchmarks/learn/engine_trace_worker.py +++ b/benchmarks/learn/engine_trace_worker.py @@ -21,7 +21,9 @@ from pathlib import Path from engine_trace_parse import parse_trace -from openultrasast.cpg.backend import DATAFLOW_OVERLAY, JoernBackend, _render +from openultrasast.cpg.backend import DATAFLOW_OVERLAY, JoernBackend, _render, extract_payload +from openultrasast.cpg.session import EngineSession +from openultrasast.model.contracts import ExecutionBudget def write_json(path, value): @@ -43,7 +45,7 @@ MIN_CPG_BYTES = 1024 HEAP_MB = 2560 CLUSTER_JVM_FLAGS = ( "-Xms1700m -Xmx1700m -XX:MaxMetaspaceSize=256m -XX:ReservedCodeCacheSize=128m " - "-XX:MaxDirectMemorySize=256m -Xss512k -XX:ActiveProcessorCount=1 -XX:+UseG1GC" + "-XX:MaxDirectMemorySize=256m -Xss512k -XX:ActiveProcessorCount=2 -XX:+UseG1GC" ) ENGINE_BUG = "engine bug: jssrc2cpg ObjectPropertyCallLinker" EXCLUDED_DIRS = frozenset({"tests", "test", "docs", "node_modules", "vendor", "dist", "build", "templates", "migrations"}) @@ -60,7 +62,9 @@ def instrument_failure(proof): if stage["command"] in FRONTENDS: if stage.get("cpg_bytes", 0) <= MIN_CPG_BYTES: return "frontend wrote no CPG" - elif stage["command"] == "joern" and stage["seconds"] < 5: + elif ( + stage["command"] == "joern" or (stage["command"] == "joern-session" and stage.get("step") == "start" and stage.get("success")) + ) and stage["seconds"] < 5: return f"implausibly fast JVM step: {stage['command']} in {stage['seconds']:.3f}s (<5s)" if not proof.get("jvm"): return "no JVM step recorded" @@ -68,7 +72,7 @@ def instrument_failure(proof): class MeasuredBackend(JoernBackend): - """Use the shipped backend with a process watchdog, no concurrent query threads.""" + """Use the shipped backend with process and resident-session watchdogs.""" def __init__(self, deadline, question_deadline, checkpoint, heap_profile="vm"): if heap_profile not in {"vm", "cluster"}: @@ -92,6 +96,7 @@ class MeasuredBackend(JoernBackend): self.labelled_files = [] self.retry_pending = False self.retries = [] + self.session_fallback = "" def _jvm_env(self): env = super()._jvm_env() @@ -102,8 +107,160 @@ class MeasuredBackend(JoernBackend): env.pop("JDK_JAVA_OPTIONS", None) return env + def _session_stage(self, name, action): + started = time.monotonic() + result = None + try: + result = action() + return result + finally: + self.stages.append( + { + "command": "joern-session", + "step": name, + "success": result is not None and result is not False, + "seconds": time.monotonic() - started, + "heap_mb": self._heap_mb(), + "java_tool_options": self._jvm_env()["JAVA_TOOL_OPTIONS"], + } + ) + self.checkpoint() + + def _fallback_session(self, reason): + self.session_fallback = reason + if self._session is not None: + detail = self._session.log_tail() + if "OutOfMemoryError" in detail: + self.fatal = "JVM out of memory" + self.stages[-1]["oom"] = True + self.stages[-1]["output_tail"] = detail + self.close_session() + self.checkpoint() + + def _session_for(self, cpg_path): + if self.session_fallback or self.fatal or self.timed_out: + return None + try: + if self._session is None: + self._session = EngineSession( + cpg_path.parent / "engine-session", + ExecutionBudget(self.deadline, 5), + self._jvm_env(), + startup_allowance_seconds=min(180, max(20, (self.deadline - time.monotonic()) / 3)), + ) + if not self._session_stage("start", self._session.start): + self._fallback_session(self._session.failure or "session_start_failed") + return None + if self._session.loaded_graph != str(cpg_path): + if not self._session_stage("load", lambda: self._session.load(cpg_path, timeout=self.query_timeout)): + self._fallback_session(self._session.failure or "session_load_failed") + return None + self._defined.clear() + return self._session + except (OSError, RuntimeError) as exc: + self._fallback_session(f"session_start_or_load:{exc}") + return None + + def _session_payload(self, cpg_path, query, params): + session = self._session_for(cpg_path) + if session is None: + return None + end = min(self.deadline, time.monotonic() + self.query_timeout) + question_end = end + prior_asked = list(self.asked) + prior_answers = dict(self.answers) + prior_counts = self.completion_counts.copy() + consumed = 0 + streamed = "" + + def observe(output): + nonlocal consumed, question_end, streamed + streamed = output + lines = output[consumed:].splitlines(keepends=True) + for line in lines: + if not line.endswith("\n"): + break + consumed += len(line) + try: + event = json.loads(line) + except ValueError: + continue + if not isinstance(event, dict): + continue + if "__question__" in event: + self.asked.append(event["__question__"]) + question_end = min(end, time.monotonic() + self.question_deadline) + if "id" in event and "rows" in event: + self.answers[event["id"]] = event["rows"] + self.completion_counts[event["id"]] += 1 + question_end = end + self.checkpoint() + if time.monotonic() >= min(end, question_end): + self.timed_out = True + return False + return True + + def run(): + source = (self.queries_dir / f"{query}.sc").read_text() + # Same shipped script/parameters as JoernBackend, using the graph loaded above. + if source.count("importCpg(cpgFile)") != 1: + return None + symbol = f"ousast_{query}" + if query not in self._defined: + source = source.replace("importCpg(cpgFile)", "()", 1) + source = source.replace("@main def exec(", f"def {symbol}(", 1) + if session.define(source, timeout=max(0, end - time.monotonic())) is None: + return None + self._defined.add(query) + arguments = ", ".join(f"{name} = {json.dumps(_render(value))}" for name, value in sorted(params.items())) + session.output_observer = observe + try: + answer = session.evaluate(f"{symbol}({arguments})", timeout=max(0, end - time.monotonic())) + finally: + session.output_observer = None + if answer is not None: + observe(answer.stdout) + return answer.stdout + return None + + failure = "" + try: + body = self._session_stage(query, run) + except (OSError, RuntimeError) as exc: + body = None + failure = f"session_{query}:{exc}" + if self.timed_out: + self.close_session() + return streamed + parsed = extract_payload(body) if body is not None else None + if not isinstance(parsed, dict): + # Discard markers from an unverified attempt before replaying the script. + self.asked[:] = prior_asked + self.answers = prior_answers + self.completion_counts = prior_counts + self._fallback_session(failure or session.failure or f"session_{query}_payload_missing") + return None + if query == "census": + self.stages[-1]["files"] = int(parsed.get("files", 0)) + if self.stages[-1]["files"] <= 0: + self.fatal = "query census read zero files" + return body + def _apply_overlays(self, cpg_path, scratch): - result = super()._apply_overlays(cpg_path, scratch) + session = self._session_for(cpg_path) + result = False + if session is not None: + body = self._session_payload(cpg_path, "overlay", {"cpgFile": str(cpg_path)}) + saved = session.scratch / "workspace" / cpg_path.name / "cpg.bin" + if body is not None and saved.is_file(): + import shutil + + shutil.copyfile(saved, cpg_path) + result = True + elif body is not None: + self._fallback_session("session_overlay_saved_graph_missing") + if not result: + result = super()._apply_overlays(cpg_path, scratch) # An optional overlay failure must not prevent the taint-only attempt. if self.fatal == ENGINE_BUG and not self.timed_out: self.retry_pending = True @@ -300,6 +457,8 @@ def worker_language(pin, root, out, deadline, question_deadline, persist=write_j if "cpg_path" in stage: proof["cpg_path"] = stage["cpg_path"] proof["cpg_bytes"] = stage["cpg_bytes"] + if getattr(backend, "session_fallback", ""): + proof["session_fallback"] = backend.session_fallback for row in record["units"]: rid = row["unit"] row["instrument"] = proof @@ -408,6 +567,8 @@ def worker_language(pin, root, out, deadline, question_deadline, persist=write_j questions_asked=[unit["unit"]] if unit["unit"] in backend.asked else [], ) finally: + if hasattr(backend, "close_session"): + backend.close_session() if cpg and cpg.cleanup: cpg.cleanup() for stage in backend.stages: diff --git a/src/openultrasast/cpg/session.py b/src/openultrasast/cpg/session.py index f421754..95fd1df 100644 --- a/src/openultrasast/cpg/session.py +++ b/src/openultrasast/cpg/session.py @@ -31,10 +31,12 @@ import os import signal import socket import subprocess +import threading import time import urllib.error import urllib.request import uuid +from collections.abc import Callable from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -78,6 +80,7 @@ class EngineSession: # with nothing: measured as a 25.3 s "build" of which only 4.5 s was the frontend. startup_allowance_seconds: float = 20.0 work_reserve_seconds: float = 10.0 + output_observer: Callable[[str], bool] | None = None failure: str = "" startup_seconds: float = 0.0 loaded_graph: str = "" @@ -199,7 +202,34 @@ class EngineSession: target = self.scratch / f"answer-{nonce}.txt" wrapped = _redirected(code, target, nonce) started = time.monotonic() - body = self._post(wrapped, timeout=remaining) + if self.output_observer is None: + body = self._post(wrapped, timeout=remaining) + else: + # Only HTTP runs in the helper; observation and cancellation stay on the caller. + done = threading.Event() + response: list[dict[str, object] | None] = [] + + def post() -> None: + try: + response.append(self._post(wrapped, timeout=remaining)) + finally: + done.set() + + threading.Thread(target=post, daemon=True).start() + while True: + finished = done.wait(0.05) + try: + first, _, output = target.read_text(errors="replace").partition("\n") + except OSError: + first, output = "", "" + if not self.output_observer(output if first.strip() == nonce else ""): + self._poison("session_observer_cancelled") + self.close() + target.unlink(missing_ok=True) + return None + if finished: + break + body = response[0] if response else None seconds = time.monotonic() - started self.requests += 1 if body is None: diff --git a/tests/test_cpg_session.py b/tests/test_cpg_session.py index 9842062..0976b15 100644 --- a/tests/test_cpg_session.py +++ b/tests/test_cpg_session.py @@ -343,3 +343,42 @@ def test_a_startup_wait_is_bounded_by_its_allowance_not_the_deadline(tmp_path, f # Bounded by the allowance, nowhere near the 600-second deadline. assert time.monotonic() - started < 10.0 live.close() + + +@pytest.mark.parametrize("cancel", [False, True]) +def test_live_receipt_observation_without_socket_or_jvm(tmp_path, monkeypatch, cancel): + import threading + from types import SimpleNamespace + + live = EngineSession(tmp_path, ExecutionBudget(time.monotonic() + 10, 1), {}) + live._process = SimpleNamespace(poll=lambda: None) + stopped = threading.Event() + observed = [] + + def close(): + live._process = None + stopped.set() + + def post(code, **kwargs): + target = code.split('java.nio.file.Path.of("', 1)[1].split('"', 1)[0] + nonce = code.split('__ousastPath, "', 1)[1].split("\\n", 1)[0] + Path(target).write_text(nonce + '\n{"__question__":"a"}\n') + if cancel: + assert stopped.wait(2) + return {"success": True} + + def observe(output): + observed.append(output) + return not cancel + + monkeypatch.setattr(live, "_post", post) + monkeypatch.setattr(live, "close", close) + live.output_observer = observe + answer = live.evaluate("work()", timeout=5) + assert any("__question__" in output for output in observed) + if cancel: + assert answer is None and live.failure == "session_observer_cancelled" + assert stopped.is_set() + else: + assert answer.stdout == '{"__question__":"a"}\n' + assert not list(tmp_path.glob("answer-*.txt")) diff --git a/tests/test_engine_task_entry.py b/tests/test_engine_task_entry.py index e5bba32..fe29d40 100644 --- a/tests/test_engine_task_entry.py +++ b/tests/test_engine_task_entry.py @@ -30,7 +30,7 @@ def test_engine_cluster_heap_flags(monkeypatch): backend = worker.MeasuredBackend(30, 10, lambda: None, heap_profile="cluster") flags = ( "-Xms1700m -Xmx1700m -XX:MaxMetaspaceSize=256m -XX:ReservedCodeCacheSize=128m " - "-XX:MaxDirectMemorySize=256m -Xss512k -XX:ActiveProcessorCount=1 -XX:+UseG1GC" + "-XX:MaxDirectMemorySize=256m -Xss512k -XX:ActiveProcessorCount=2 -XX:+UseG1GC" ) assert backend._jvm_env()["JAVA_TOOL_OPTIONS"] == flags assert backend._heap_flag() == "-J-Xmx1700m" diff --git a/tests/test_engine_trace.py b/tests/test_engine_trace.py index 1385490..b46a9af 100644 --- a/tests/test_engine_trace.py +++ b/tests/test_engine_trace.py @@ -1598,3 +1598,124 @@ def test_ax_pending_fallback_resumes_in_docker_and_retains_attempt(tmp_path, mon progress = json.loads((tmp_path / "progress.json").read_text()) assert progress["pins_done"] == 1 and progress["pins_in_flight"] == 0 assert progress["fallbacks_pending"] == 0 and progress["statuses"] == {"asked-nothing": 1} + + +@pytest.mark.parametrize("failure", ["", "start", "taint"]) +def test_worker_resident_session_and_fallback(tmp_path, monkeypatch, failure): + from types import SimpleNamespace + + calls = [] + graph = tmp_path / "cpg.bin" + graph.write_bytes(b"g" * 2048) + payloads = { + "overlay": {"files": "1", "methods": "1"}, + "census": {"files": "1", "file_names": ["app.py"]}, + "taint": {"__census__": [{"files": 1, "file_names": ["app.py"]}], "a": []}, + } + + def output(query): + return BEGIN + "\n" + json.dumps(payloads[query]) + "\n" + END + + class Session: + loaded_graph = "" + failure = "fake failure" + output_observer = None + + def __init__(self, scratch, budget, env, **kwargs): + self.scratch = scratch + assert "ActiveProcessorCount=2" in env["JAVA_TOOL_OPTIONS"] + + def start(self): + calls.append("start") + return failure != "start" + + def load(self, path, **kwargs): + calls.append("load") + assert path.read_bytes() == b"g" * 2048 + self.loaded_graph = str(path) + return True + + def define(self, source, **kwargs): + assert " importCpg(cpgFile)" not in source + return "" + + def evaluate(self, code, **kwargs): + query = code.split("(")[0].removeprefix("ousast_") + calls.append(query) + if query == "overlay": + saved = self.scratch / "workspace" / graph.name / "cpg.bin" + saved.parent.mkdir(parents=True) + saved.write_bytes(b"s" * 2048) + if query == "taint": + if failure == "taint": + return None + self.output_observer('{"__question__":"a"}\n') + self.output_observer('{"__question__":"a"}\n{"id":"a","rows":[]}\n') + return SimpleNamespace(stdout=output(query)) + + def close(self): + calls.append("close") + + def log_tail(self): + return "fake log" + + monkeypatch.setattr(worker, "EngineSession", Session) + backend = worker.MeasuredBackend(900, 120, lambda: None, heap_profile="cluster") + disposable = [] + + def run(command, **kwargs): + query = Path(command[command.index("--script") + 1]).stem + disposable.append(query) + if query == "overlay": + saved = tmp_path / "workspace" / graph.name / "cpg.bin" + saved.parent.mkdir(parents=True, exist_ok=True) + saved.write_bytes(b"s" * 2048) + return subprocess.CompletedProcess(command, 0, output(query), "") + + monkeypatch.setattr(backend, "_run", run) + monkeypatch.setattr("openultrasast.cpg.backend.shutil.which", lambda name: name) + assert backend._apply_overlays(graph, tmp_path) + assert backend.query(graph, "census", {})["files"] == "1" + assert backend._batch_once(graph, "taint", {"a": {}})["a"] == [] + backend.close_session() + assert calls.count("start") == 1 + if failure == "start": + assert disposable == ["overlay", "census", "taint"] + assert backend.session_fallback == "fake failure" + else: + assert calls.count("load") == 1 + assert [stage["step"] for stage in backend.stages] == ["start", "load", "overlay", "census", "taint"] + assert all(stage["seconds"] >= 0 for stage in backend.stages) + assert backend.stages[3]["files"] == 1 + assert disposable == (["taint"] if failure else []) + if not failure: + assert backend.asked == ["a"] + assert backend.answers == {"a": []} + assert backend.completion_counts["a"] == 1 + assert calls.count("close") == 1 + + +def test_worker_session_question_deadline_keeps_markers_and_does_not_replay(tmp_path): + from types import SimpleNamespace + + graph = tmp_path / "cpg.bin" + graph.write_bytes(b"g" * 2048) + backend = worker.MeasuredBackend(900, 0, lambda: None) + closed = [] + + def evaluate(code, **kwargs): + assert not fake.output_observer('{"__question__":"a"}\n') + return None + + fake = SimpleNamespace( + loaded_graph=str(graph), + define=lambda *a, **kw: "", + evaluate=evaluate, + close=lambda: closed.append(True), + ) + backend._session = fake + body = backend._session_payload(graph, "taint", {"cpgFile": str(graph)}) + assert body == '{"__question__":"a"}\n' + assert backend.timed_out and backend.asked == ["a"] + assert not backend.answers and not backend.session_fallback + assert closed == [True]