diff --git a/src/openultrasast/cpg/backend.py b/src/openultrasast/cpg/backend.py index 5c765e3..9f7d82d 100644 --- a/src/openultrasast/cpg/backend.py +++ b/src/openultrasast/cpg/backend.py @@ -297,6 +297,11 @@ class CpgResult: cpg_path: Path run: Callable[[str, Mapping[str, object]], object | None] run_batch: Callable[[str, Mapping[str, Mapping[str, object]]], dict[str, list[object]] | None] | None = None + # A SINGLE batch attempt: no internal split-on-failure, so its wall time is the cost of exactly one + # invocation and a killed one returns the answers it streamed before the kill. The scan-level portioner + # uses this because it is itself the retry layer -- timing `run_batch`, which splits internally, once + # read one portion at 953 s (its own retries folded in) and collapsed the adaptive sizer. + run_batch_once: Callable[[str, Mapping[str, Mapping[str, object]]], dict[str, list[object]] | None] | None = None # Files the frontend read and could not turn into a graph. A CPG is not all-or-nothing: `php2cpg` logs a # warning per file it drops and still exits 0 with a graph that is missing them, so a scan over the # remainder is a scan of LESS CODE THAN IT WAS ASKED ABOUT, and silence about those files is the same @@ -607,6 +612,7 @@ class JoernBackend: cpg_path=shards[0], run=lambda query, params: self.query_across(shards, query, params), run_batch=lambda query, requests: self.query_batch_across(shards, query, requests), + run_batch_once=lambda query, requests: self.query_batch_across(shards, query, requests, once=True), unparsed=unparsed, cleanup=dispose, execution_diagnostics=lambda: tuple(self._diagnostics), @@ -639,6 +645,7 @@ class JoernBackend: cpg_path=cpg_path, run=lambda query, params: self.query(cpg_path, query, params), run_batch=lambda query, requests: self.query_batch(cpg_path, query, requests), + run_batch_once=lambda query, requests: self._batch_once(cpg_path, query, requests), unparsed=unparsed, cleanup=dispose, execution_diagnostics=lambda: tuple(self._diagnostics), @@ -931,12 +938,14 @@ class JoernBackend: return merged if answered else None def query_batch_across( - self, shards: Sequence[Path], query: str, requests: Mapping[str, Mapping[str, object]] + self, shards: Sequence[Path], query: str, requests: Mapping[str, Mapping[str, object]], *, once: bool = False ) -> dict[str, list[object]] | None: """The batch over every shard, rows concatenated per request id. The census is summed rather than taken from one shard, and carries the shard count, because `cpg_empty` has to mean "no graph anywhere" and not "the first of two graphs was small". + + ``once`` runs a single attempt per shard, for the scan-level portioner that retries itself. """ merged: dict[str, list[object]] = {} methods = 0 @@ -945,7 +954,7 @@ class JoernBackend: file_names: set[str] = set() named_census = True for shard in shards: - one = self.query_batch(shard, query, requests) + one = self._batch_once(shard, query, requests) if once else self.query_batch(shard, query, requests) if one is None: named_census = False if self.execution_budget is not None: diff --git a/src/openultrasast/model/scan.py b/src/openultrasast/model/scan.py index 4da8bcf..dd42160 100644 --- a/src/openultrasast/model/scan.py +++ b/src/openultrasast/model/scan.py @@ -427,8 +427,13 @@ def _scan_repository_impl( kind_started = time.monotonic() batch = getattr(cpg, "run_batch", None) if callable(batch): + # The portioner is the retry layer, so it asks in SINGLE attempts: a portion timed with + # `run_batch` folds that method's own split-retry into the number, which once read 953 s for one + # portion and collapsed the sizer. `run_batch_once` is one invocation whose kill returns what it + # streamed. A backend without it falls back to `run_batch`, unchanged. + once = getattr(cpg, "run_batch_once", None) answered_batch = _batched( - batch, + once if callable(once) else batch, kind, requests, weights=_sink_weights(work, evidence_by_pair), @@ -910,7 +915,12 @@ def _batched( chunk = sizer.fit(chunk, pending=chunks[index + 1 :]) started = time.monotonic() answer = batch(kind, chunk) - sizer.observe(chunk, time.monotonic() - started, answered=bool(answer)) + # WHOLE, not merely non-empty. A killed portion streams some answers and returns them, which is the + # point -- but its time is the ceiling, not the cost of the work, so fitting a per-visit rate from it + # is meaningless. Only a portion that answered every id it was asked gives a clean sample; a short + # one halves the budget, because it was too big for the ceiling. + whole = answer is not None and all(rid in answer for rid in chunk) + sizer.observe(chunk, time.monotonic() - started, answered=whole) if answer is None: continue answered_any = True diff --git a/tests/test_cpg_backend.py b/tests/test_cpg_backend.py index 8892e30..7d12bde 100644 --- a/tests/test_cpg_backend.py +++ b/tests/test_cpg_backend.py @@ -1269,3 +1269,45 @@ def test_a_killed_batch_retries_only_the_remainder(tmp_path: Path) -> None: assert answered is not None and answered["a"] == [1] and answered["b"] == ["b"] and answered["c"] == ["c"] assert calls[0] == ["a", "b", "c"] assert all("a" not in retry for retry in calls[1:]), calls + + +def test_the_portioner_asks_single_attempts_and_halves_on_a_short_answer(tmp_path: Path) -> None: + """A killed portion returns what it streamed; the sizer must read that as too big, not as slow. + + Timing `run_batch` folds its internal split-retry into the number -- one portion read 953 s that way and + the sizer fitted a per-visit rate from it and collapsed. The portioner asks `run_batch_once` instead, and + a portion that came back short of the ids it asked halves the budget rather than fitting from a ceiling. + """ + from openultrasast.cpg.backend import CpgResult + from openultrasast.model import scan + + once_calls: list[list[str]] = [] + + def once(kind: str, requests: dict[str, dict[str, object]]) -> dict[str, list[object]] | None: + once_calls.append(sorted(requests)) + # Answer everything except a designated heavy id, as a kill would: it streamed the rest and died. + return {rid: [] for rid in requests if rid != "heavy"} + + result = CpgResult(cpg_path=tmp_path / "c.bin", run=lambda k, p: None, run_batch=lambda k, r: None, run_batch_once=once) + + seen: list[bool] = [] + real_observe = scan._PortionSizer.observe + + def spy(self, chunk, seconds, *, answered): # type: ignore[no-untyped-def] + seen.append(answered) + return real_observe(self, chunk, seconds, answered=answered) + + # Two portions over the budget so the portioner actually splits and observes. + weights = {"heavy": 500, "a": 400, "b": 400} + requests = {rid: {} for rid in weights} + monkey = scan._PortionSizer.observe + scan._PortionSizer.observe = spy # type: ignore[method-assign] + try: + merged = scan._batched(result.run_batch_once, "taint", requests, weights=weights) + finally: + scan._PortionSizer.observe = monkey # type: ignore[method-assign] + + assert merged is not None and "a" in merged and "b" in merged and "heavy" not in merged + assert once_calls, "the single-attempt path was not used" + # The portion that came back short of its ids was reported as not-whole (False) to the sizer. + assert False in seen, seen