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 <[email protected]>

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,7 +107,159 @@ 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):
+        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:
@@ -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:

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()
+        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:

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

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"

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]