diff --git a/benchmarks/unseen/README.md b/benchmarks/unseen/README.md index 9af5c38..1e432f3 100644 --- a/benchmarks/unseen/README.md +++ b/benchmarks/unseen/README.md @@ -143,6 +143,27 @@ and sleepers. A stale candidate 404 is counted; other service failures fail the ### Checkpoint and resume +Replay reads bounded JSONL records into a disposable SQLite working store, with a 2 MiB +page cache. Advisory descriptions, non-fix references and vulnerable-version payloads +are discarded; package identities, CWEs, status, publication dates and every fix link +are retained. Repository metadata keeps canonical/fork identities and eligibility fields. +Advisory ordering, fix indexes and checkpoint lookups stay on disk, independent of the +number of journal records. The original seeded ordering and ordinary draws are preserved. +Per-repository history and exclusion sets also spill to private temporary storage, which +is removed on success or failure. Budget disk space for the journal, working store and +spooled Git output as well as the clone. + +To bound individual reads, journal records are limited to 16 MiB, decoded Git command +output to 8 MiB, history messages to 1 MiB and source blobs to 512 KiB. Exceeding a limit +fails the instrument explicitly; it never counts as an empty history or a clean result. +The synthetic 20,000-advisory regression checks a 64 MiB traced-memory ceiling across +replay, enumeration, ordering and fix indexing; this is not a live coordinator RSS measurement. + +The memory fix can migrate the exact preceding design fingerprints automatically, +keeping `inputs-before-memory-fix.json` and leaving the JSONL journal untouched. All +other input/design mismatches still refuse resume. This compatibility exception must +be removed if a later change alters selection or sampling semantics. + The command above uses journal directory `/home/mc/ousast-results/plane/unseen-draft/journal`. Run the same command to resume. `--max-hours` is a per-invocation budget; changing it does not invalidate checkpoints. diff --git a/benchmarks/unseen/draft.py b/benchmarks/unseen/draft.py index 6597699..23aca1f 100644 --- a/benchmarks/unseen/draft.py +++ b/benchmarks/unseen/draft.py @@ -24,6 +24,7 @@ from pathlib import Path from . import eligibility, extract, source from .extract import Git, InstrumentFailure from .journal import Journal, ResolvedGitHub, ResumeMismatch, Stopped, atomic, digest, external +from .storage import AdvisoryRows, DiskSet, FixIndex, working_store def canonical(value) -> bytes: @@ -173,41 +174,45 @@ def prove_source(git, revision: str) -> dict: raise InstrumentFailure("zero_source_bytes") -def advisory_index(rows: list[dict], github) -> dict[str, set[str]]: +def advisory_index(rows, github, *, db): """Include ALL enumerated fixes, even old, unmapped, multi-fix and other ecosystems. Canonicalize link names as well, so an old advisory URL survives repository renames. A 404 is not treated as proof that an advisory is unrelated to a selected repository. """ - index = defaultdict(set) + index = FixIndex(db) for row in rows: for name, sha in source.fix_links(row): try: resolved = eligibility.resolve(name, github) except eligibility.RepositoryNotFound: - index[name].add(sha) + index.add(name, sha) continue - index[resolved.full_name].add(sha) - return dict(index) + index.add(resolved.full_name, sha) + return index def repository_changes(git, candidate, fixes: set[str], *, seed: int): proof = prove_source(git, candidate.parent) change = extract.introducing(git, candidate) - known = set() - expanded = set() - for fix in sorted(fixes | {candidate.fix}): - # Explicitly fetch detached advisory commits unavailable on default-branch history. - # All git including this possible network call still goes through the injected client. - git.run("fetch", "--quiet", "--filter=blob:none", "origin", fix) - full = git.run("rev-parse", fix).strip() - expanded.add(full) - parents = git.parents(full) - for parent in parents: - sites = extract.fix_sites(git, full, parent, global_ok=True) - known.update((site["path"], site["function"]) for site in sites) - rows = extract.ordinary_history(git, "HEAD") - ordinary, exclusions = extract.draw_ordinary(rows, expanded, known, seed=seed) + with working_store() as db: + known = DiskSet(db, "known_functions") + expanded = DiskSet(db, "expanded_fixes") + pending = DiskSet(db, "pending_fixes") + pending.update(fixes) + pending.add(candidate.fix) + for fix in pending: + # Explicitly fetch detached advisory commits unavailable on default-branch history. + # All git including this possible network call still goes through the injected client. + git.run("fetch", "--quiet", "--filter=blob:none", "origin", fix) + full = git.run("rev-parse", fix).strip() + expanded.add(full) + parents = git.parents(full) + for parent in parents: + sites = extract.fix_sites(git, full, parent, global_ok=True) + known.update((site["path"], site["function"]) for site in sites) + rows = extract.ordinary_history(git, "HEAD") + ordinary, exclusions = extract.draw_ordinary(rows, expanded, known, seed=seed) return [change, *ordinary], exclusions, proof @@ -222,14 +227,28 @@ def inputs(used, seed, known_empty, clones): "journal_version": 1, } # Hardcoded design thresholds and algorithms are covered as well as taxonomy data. - for name in ("draft", "source", "extract", "eligibility", "journal"): + for name in ("draft", "source", "extract", "eligibility", "journal", "storage"): parameters["design_" + name] = hashlib.sha256(Path(__file__).with_name(name + ".py").read_bytes()).hexdigest() parameters["design_taxonomy"] = hashlib.sha256(DEFAULT_FAMILIES_PATH.read_bytes()).hexdigest() return parameters +# The pre-memory-fix implementation has identical ordering/selection semantics. Only +# this exact fingerprint set may migrate; changed seeds, taxonomy or used sets still fail. +LEGACY_DESIGNS = { + "design_draft": "906bf36de11761931b2c1b62a2ea8da3c8124f9d754455f1a0acf29a6db576a3", + "design_source": "01c896bd0418faf0d073dc2d5cb6dbc5d7903b95b5e8216c9022bc9985326a9c", + "design_extract": "593d954d9c22b835e1d349df6b575707a39098dca2eff4425f67a0de14c6525e", + "design_journal": "c97ad0761ea5bbfe65db816b13dfa79d93608ea03bcd5848f70686ed9157bdf7", +} + + +def legacy_inputs(parameters): + return {key: LEGACY_DESIGNS.get(key, value) for key, value in parameters.items() if key != "design_storage"} + + def enumerate_advisories(github, journal): - page, records = 1, {} + page, records = 1, AdvisoryRows(journal.db, "advisory_records") while True: key = str(page) if key not in journal.rows["advisories"]: @@ -237,13 +256,14 @@ def enumerate_advisories(github, journal): result = github.advisory_page(page) journal.put("advisories", key, result) result = journal.rows["advisories"][key] - records.update((row["ghsa_id"], row) for row in result["rows"]) + for row in result["rows"]: + records[row["ghsa_id"]] = row if result["done"]: break page += 1 if not records: raise InstrumentFailure("empty_advisories") - return list(records.values()) + return records def decide(row, github, osv, used, selected_networks): @@ -335,7 +355,8 @@ def run( checkpoint = Journal( path, max_hours=max_hours, clock=clock, rate_limit=lambda: getattr(getattr(github, "transport", None), "rate_limit", {}) ) - with checkpoint.open(inputs(used, seed, known_empty, clones)): + parameters = inputs(used, seed, known_empty, clones) + with checkpoint.open(parameters, compatible_parameters=legacy_inputs(parameters)): checkpoint.boundary() api = ResolvedGitHub(github, checkpoint) used = copy.deepcopy(used) # canonicalization must not mutate the resume input @@ -343,7 +364,7 @@ def run( report["used_set"] = used.report(known_empty) rows = enumerate_advisories(api, checkpoint) report["advisories"] = len(rows) - fixes_by_repo = advisory_index(rows, api) + fixes_by_repo = advisory_index(rows, api, db=checkpoint.db) selected_networks, entries = set(), [] for row in source.ordered_advisories(rows, seed): checkpoint.boundary() diff --git a/benchmarks/unseen/extract.py b/benchmarks/unseen/extract.py index 0e7a5a5..58f9b53 100644 --- a/benchmarks/unseen/extract.py +++ b/benchmarks/unseen/extract.py @@ -11,9 +11,12 @@ import ast import json import re import subprocess +import tempfile +from contextlib import contextmanager from pathlib import Path from .source import Candidate, Rejected +from .storage import working_store LANGUAGES = {".py", ".js", ".jsx", ".mjs", ".cjs", ".ts", ".tsx", ".php", ".java"} HUNK = re.compile(r"^@@ -(\d+)(?:,(\d+))? \+(\d+)(?:,(\d+))? @@", re.M) @@ -37,24 +40,54 @@ class Git: self.last_blob_bytes = 0 self.limit_bytes = limit_bytes + @contextmanager + def output(self, *args: str, allowed=(0,)): + # Git can expand a tiny pack into enormous stdout. Spool before decoding, including + # stderr, so neither subprocess.communicate nor an error can consume the VM's RAM. + with tempfile.TemporaryFile() as stdout, tempfile.TemporaryFile() as stderr: + try: + done = self.runner( + ["git", "-c", "core.quotePath=false", "-C", str(self.path), *args], + stdin=subprocess.DEVNULL, + stdout=stdout, + stderr=stderr, + timeout=180, + check=False, + ) + except Exception: + raise InstrumentFailure("git_launch_failed") from None + if done.returncode not in allowed: + raise InstrumentFailure("git_exit_" + str(done.returncode)) + if self.limit_bytes is not None and self.disk_bytes() > self.limit_bytes: + raise InstrumentFailure("clone_size_cap") + # Injected runners may return bytes instead of writing the supplied file. + if done.stdout is not None: + stdout.write(done.stdout) + stdout.seek(0) + yield stdout + def run(self, *args: str, allowed=(0,)) -> str: - try: - done = self.runner( - ["git", "-c", "core.quotePath=false", "-C", str(self.path), *args], - stdin=subprocess.DEVNULL, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - timeout=180, - check=False, - ) - except Exception: - raise InstrumentFailure("git_launch_failed") from None - if done.returncode not in allowed: - raise InstrumentFailure("git_exit_" + str(done.returncode)) - if self.limit_bytes is not None and self.disk_bytes() > self.limit_bytes: - raise InstrumentFailure("clone_size_cap") - self.last_blob_bytes = len(done.stdout) - return done.stdout.decode("utf-8", errors="replace") + with self.output(*args, allowed=allowed) as stdout: + raw = stdout.read(8 * 1024 * 1024 + 1) + if len(raw) > 8 * 1024 * 1024: + raise InstrumentFailure("git_output_size_cap") + self.last_blob_bytes = len(raw) + return raw.decode("utf-8", errors="replace") + + def history_entries(self, head): + with self.output("log", "--first-parent", "--format=%H%x00%P%x00%ct%x00%an%x00%B%x1e", head) as stdout: + pending = b"" + while chunk := stdout.read(64 * 1024): + entries = (pending + chunk).split(b"\x1e") + pending = entries.pop() + for entry in entries: + if len(entry) > 1024 * 1024: + raise InstrumentFailure("git_history_record_size_cap") + yield entry.decode("utf-8", errors="replace") + if len(pending) > 1024 * 1024: + raise InstrumentFailure("git_history_record_size_cap") + if pending.strip(): + raise InstrumentFailure("invalid_history") def disk_bytes(self) -> int: return sum(p.stat().st_size for p in self.path.rglob("*") if p.is_file() and not p.is_symlink()) @@ -66,6 +99,9 @@ class Git: # Missing is established by a successful tree read, never by swallowing a show error. if not self.run("ls-tree", sha, "--", path).strip(): return None + size = int(self.run("cat-file", "-s", f"{sha}:{path}").strip()) + if size > 512 * 1024: + raise InstrumentFailure("source_blob_size_cap") text = self.run("show", f"{sha}:{path}") if Path(path).suffix in LANGUAGES: self.files_read += 1 @@ -378,10 +414,9 @@ def generated(git, revision: str, path: str, text: str) -> bool: return result -def ordinary_history(git, head: str) -> list[dict]: - raw = git.run("log", "--first-parent", "--format=%H%x00%P%x00%ct%x00%an%x00%B%x1e", head) - rows = [] - for entry in raw.split("\x1e"): +def ordinary_history(git, head: str): + """Yield one commit at a time; no repository-wide blob or diff retention.""" + for entry in git.history_entries(head): if not entry.strip(): continue fields = entry.strip().split("\x00") @@ -389,8 +424,8 @@ def ordinary_history(git, head: str) -> list[dict]: raise InstrumentFailure("invalid_history") sha, parents, timestamp, author, message = fields row = {"head": sha, "parents": parents.split(), "timestamp": int(timestamp), "author": author, "message": message, "files": []} - rows.append(row) if len(row["parents"]) != 1: + yield row continue base = row["parents"][0] for path, added, removed in git.stats(base, sha): @@ -417,58 +452,68 @@ def ordinary_history(git, head: str) -> list[dict]: info["functions"].append(name) info["generated"] |= generated(git, revision, path, text) row["files"].append(info) - return rows + yield row -def draw_ordinary(rows: list[dict], fixes: set[str], known_functions: set[tuple[str, str]], *, seed: int) -> tuple[list[dict], dict]: +def draw_ordinary(rows, fixes: set[str], known_functions: set[tuple[str, str]], *, seed: int) -> tuple[list[dict], dict]: + """Spill eligible history to disk; keep the original seeded stratified draw exactly.""" + with working_store() as db: + return _draw_ordinary(db, rows, fixes, known_functions, seed=seed) + + +def _draw_ordinary(db, rows, fixes, known_functions, *, seed): import random - from collections import Counter, defaultdict + from collections import Counter counts = Counter() - eligible = [] + earliest = latest = None + db.execute("CREATE TABLE eligible (head TEXT, timestamp INTEGER, bucket INTEGER, value TEXT, third INTEGER)") for row in rows: + timestamp = row["timestamp"] + earliest = timestamp if earliest is None else min(earliest, timestamp) + latest = timestamp if latest is None else max(latest, timestamp) reason = ordinary_rejection(row, fixes, known_functions) counts[reason or "eligible"] += 1 - if reason is None: - eligible.append(row) - if len(eligible) < 40: + if reason is not None: + continue + size = sum(f["added"] + f["removed"] for f in row["files"] if Path(f["path"]).suffix in LANGUAGES and not manifest(f["path"])) + bucket = 0 if size <= 10 else 1 if size <= 50 else 2 if size <= 200 else 3 + value = { + "kind": "ordinary", + "base": row["parents"][0], + "head": row["head"], + "method": "first_parent", + "size_bucket": bucket, + "files": [{"path": f["path"], "head_lines": f["head_lines"]} for f in row["files"]], + } + db.execute("INSERT INTO eligible VALUES (?, ?, ?, ?, NULL)", (row["head"], timestamp, bucket, json.dumps(value))) + if counts["eligible"] < 40: failure = Rejected("ordinary_floor") failure.counts = dict(counts) raise failure - earliest, latest = min(r["timestamp"] for r in rows), max(r["timestamp"] for r in rows) - cells = defaultdict(list) - for row in eligible: - third = min(2, 3 * (row["timestamp"] - earliest) // max(1, latest - earliest + 1)) - size = sum(f["added"] + f["removed"] for f in row["files"] if Path(f["path"]).suffix in LANGUAGES and not manifest(f["path"])) - bucket = 0 if size <= 10 else 1 if size <= 50 else 2 if size <= 200 else 3 - cells[third, bucket].append(row) + db.execute("UPDATE eligible SET third = min(2, 3 * (timestamp - ?) / ?)", (earliest, max(1, latest - earliest + 1))) + cells = { + (third, bucket): count for third, bucket, count in db.execute("SELECT third, bucket, count(*) FROM eligible GROUP BY third, bucket") + } if {t for t, _ in cells} != {0, 1, 2}: failure = Rejected("time_thirds") failure.counts = dict(counts) raise failure - # Largest-deficit rounding with the mandatory minimum of one per time third. - # No cell can receive more draws than it has eligible commits. allocation = Counter() - targets = {cell: 12 * len(items) / len(eligible) for cell, items in cells.items()} + targets = {cell: 12 * count / counts["eligible"] for cell, count in cells.items()} for third in range(3): cell = max(sorted(c for c in cells if c[0] == third), key=lambda c: targets[c]) allocation[cell] = 1 while sum(allocation.values()) < 12: - cell = max(sorted(c for c in cells if allocation[c] < len(cells[c])), key=lambda c: targets[c] - allocation[c]) + cell = max(sorted(c for c in cells if allocation[c] < cells[c]), key=lambda c: targets[c] - allocation[c]) allocation[cell] += 1 + db.execute("CREATE INDEX eligible_order ON eligible(third, bucket, head)") rng = random.Random(seed) drawn = [] for cell in sorted(cells): - for row in rng.sample(sorted(cells[cell], key=lambda r: r["head"]), allocation[cell]): - drawn.append( - { - "kind": "ordinary", - "base": row["parents"][0], - "head": row["head"], - "method": "first_parent", - "time_third": cell[0], - "size_bucket": cell[1], - "files": [{"path": f["path"], "head_lines": f["head_lines"]} for f in row["files"]], - } - ) + for position in rng.sample(range(cells[cell]), allocation[cell]): + (value,) = db.execute( + "SELECT value FROM eligible WHERE third=? AND bucket=? ORDER BY head LIMIT 1 OFFSET ?", (*cell, position) + ).fetchone() + drawn.append({**json.loads(value), "time_third": cell[0]}) return drawn, dict(counts) diff --git a/benchmarks/unseen/journal.py b/benchmarks/unseen/journal.py index 9938ead..aefb262 100644 --- a/benchmarks/unseen/journal.py +++ b/benchmarks/unseen/journal.py @@ -11,6 +11,20 @@ import time from contextlib import contextmanager from pathlib import Path +from .source import compact_advisory, compact_repository +from .storage import DiskMap, working_store + +MAX_RECORD_BYTES = 16 * 1024 * 1024 + + +def compact(stage, value): + if stage == "advisories": + return {"rows": [compact_advisory(row) for row in value["rows"]], "done": value["done"]} + if stage == "resolution" and value is not None: + return compact_repository(value) + return value + + STAGES = ("advisories", "resolution", "eligibility", "extraction", "completion") @@ -68,13 +82,13 @@ class Journal: self.started = self.clock() self.max_seconds = None if max_hours is None else max_hours * 3600 self.rate_limit = rate_limit or (lambda: {}) - self.rows = {stage: {} for stage in STAGES} + self.rows = {} self.elapsed_before = 0.0 self.last = None self.previous_rate = {} @contextmanager - def open(self, parameters): + def open(self, parameters, *, compatible_parameters=None): self.path.mkdir(parents=True, exist_ok=True, mode=0o700) with (self.path / "LOCK").open("a") as lock: try: @@ -82,6 +96,7 @@ class Journal: except BlockingIOError: raise ValueError("journal_busy") from None metadata = self.path / "inputs.json" + migrating = False if metadata.exists(): old = json.loads(metadata.read_bytes()) # Only current, fixed parameter labels can enter stdout. @@ -89,30 +104,42 @@ class Journal: if old.keys() != parameters.keys(): changed.append("parameter_schema") if changed: - raise ResumeMismatch(changed) + if old != compatible_parameters: + raise ResumeMismatch(changed) + migrating = True else: if any(self.path.glob("*.jsonl")): raise ValueError("journal_missing_inputs") atomic(metadata, encoded(parameters)) - for stage in STAGES: - path = self.path / (stage + ".jsonl") - if not path.exists(): - continue - with path.open("rb") as handle: - for line in handle: - # Fail closed on torn/corrupt writes; never silently lose a completed unit. - if not line.endswith(b"\n"): - raise ValueError("journal_incomplete_record") - row = json.loads(line) - if row["id"] in self.rows[stage]: - raise ValueError("journal_duplicate_unit") - self.rows[stage][row["id"]] = row["value"] - if row["elapsed_seconds"] >= self.elapsed_before: - self.elapsed_before = row["elapsed_seconds"] - self.last = {"stage": stage, "id": row["id"]} - self.previous_rate = row["rate_limit"] - self.progress("sourcing") - yield self + with working_store(self.path) as self.db: + self.rows = {stage: DiskMap(self.db, stage) for stage in STAGES} + self._replay() + if migrating: + atomic(self.path / "inputs-before-memory-fix.json", encoded(old)) + atomic(metadata, encoded(parameters)) + self.progress("sourcing") + yield self + + def _replay(self): + for stage in STAGES: + path = self.path / (stage + ".jsonl") + if not path.exists(): + continue + with path.open("rb") as handle: + while line := handle.readline(MAX_RECORD_BYTES + 1): + if len(line) > MAX_RECORD_BYTES: + raise ValueError("journal_record_size_cap") + # Fail closed on torn/corrupt writes; never silently lose a completed unit. + if not line.endswith(b"\n"): + raise ValueError("journal_incomplete_record") + row = json.loads(line) + if row["id"] in self.rows[stage]: + raise ValueError("journal_duplicate_unit") + self.rows[stage][row["id"]] = compact(stage, row["value"]) + if row["elapsed_seconds"] >= self.elapsed_before: + self.elapsed_before = row["elapsed_seconds"] + self.last = {"stage": stage, "id": row["id"]} + self.previous_rate = row["rate_limit"] def boundary(self): reason = None @@ -145,11 +172,15 @@ class Journal: def put(self, stage, key, value): if key in self.rows[stage]: raise ValueError("journal_duplicate_unit") + value = compact(stage, value) row = {"id": key, "value": value, "elapsed_seconds": self.elapsed(), "rate_limit": self.rate_limit() or self.previous_rate} path = self.path / (stage + ".jsonl") with path.open("ab") as handle: os.chmod(path, 0o600) - handle.write(encoded(row) + b"\n") + payload = encoded(row) + b"\n" + if len(payload) > MAX_RECORD_BYTES: + raise ValueError("journal_record_size_cap") + handle.write(payload) handle.flush() os.fsync(handle.fileno()) sync_directory(self.path) diff --git a/benchmarks/unseen/source.py b/benchmarks/unseen/source.py index 08dac91..cd7ceba 100644 --- a/benchmarks/unseen/source.py +++ b/benchmarks/unseen/source.py @@ -6,6 +6,7 @@ import base64 import json import random import re +from collections import OrderedDict from dataclasses import dataclass from datetime import date from urllib.parse import urlencode, urlsplit @@ -42,6 +43,23 @@ class Candidate: license_path: str +def compact_advisory(row: dict) -> dict: + """Keep exactly the input fields used by selection and the all-advisory fix index.""" + return { + **{key: row[key] for key in ("ghsa_id", "published_at", "type", "withdrawn_at") if key in row}, + "cwes": [{"cwe_id": c["cwe_id"]} for c in row.get("cwes", [])], + "references": [f"https://github.com/{name}/commit/{sha}" for name, sha in sorted(fix_links(row))], + "vulnerabilities": [{"package": {key: v["package"][key] for key in ("ecosystem", "name")}} for v in row.get("vulnerabilities", [])], + } + + +def compact_repository(row: dict) -> dict: + return { + **{key: row[key] for key in ("full_name", "private", "size", "clone_url", "former_names") if key in row}, + **{key: {"full_name": row[key]["full_name"]} for key in ("parent", "source") if row.get(key)}, + } + + def post_cutoff(row: dict) -> bool: return date.fromisoformat(row["published_at"][:10]) > date(2026, 6, 1) @@ -63,6 +81,10 @@ def ecosystem(row: dict) -> str | None: def ordered_advisories(rows: list[dict], seed: int) -> list[dict]: """Stable seeded ordering, post-cutoff first in each ecosystem; round-robin strata.""" + from .storage import AdvisoryRows + + if isinstance(rows, AdvisoryRows): + return rows.ordered(seed) rng = random.Random(seed) groups = [] for eco in ECOSYSTEMS: @@ -218,7 +240,7 @@ class GitHub: raise ValueError("missing_gh_token") self.transport = transport or JSONTransport() self.headers = {"Authorization": "Bearer " + token, "Accept": "application/vnd.github+json"} - self.repos: dict[str, dict] = {} + self.repos: OrderedDict[str, dict] = OrderedDict() self.calls = 0 def get(self, path: str): @@ -228,7 +250,10 @@ class GitHub: def get_repo(self, name: str) -> dict: name = normalize(name) if name not in self.repos: - self.repos[name] = self.get("/repos/" + name) + self.repos[name] = compact_repository(self.get("/repos/" + name)) + if len(self.repos) > 128: + self.repos.popitem(last=False) + self.repos.move_to_end(name) return self.repos[name] def get_commit(self, name: str, sha: str) -> dict: diff --git a/benchmarks/unseen/storage.py b/benchmarks/unseen/storage.py new file mode 100644 index 0000000..0673192 --- /dev/null +++ b/benchmarks/unseen/storage.py @@ -0,0 +1,170 @@ +"""Disposable private SQLite working sets; JSONL remains the durable source of truth.""" + +from __future__ import annotations + +import json +import sqlite3 +import tempfile +from collections.abc import Mapping, Sequence, Set +from contextlib import contextmanager +from pathlib import Path + + +@contextmanager +def working_store(directory=None): + # TemporaryDirectory is mode 0700. No identity-bearing scratch data enters the checkout. + with tempfile.TemporaryDirectory(prefix="sourcing-", dir=directory) as temporary: + db = sqlite3.connect(Path(temporary) / "working.sqlite") + try: + db.execute("PRAGMA journal_mode=OFF") + db.execute("PRAGMA synchronous=OFF") + db.execute("PRAGMA temp_store=FILE") + db.execute("PRAGMA cache_size=-2048") + db.execute("PRAGMA mmap_size=0") + db.execute("CREATE TABLE records (kind TEXT, key TEXT, value TEXT, PRIMARY KEY (kind, key)) WITHOUT ROWID") + db.execute("CREATE TABLE sequences (kind TEXT, position INTEGER, value TEXT, PRIMARY KEY (kind, position)) WITHOUT ROWID") + yield db + finally: + db.close() + + +class DiskMap(Mapping): + def __init__(self, db, kind): + self.db, self.kind = db, kind + self.count = 0 + + def __len__(self): + return self.count + + def __iter__(self): + for (key,) in self.db.execute("SELECT key FROM records WHERE kind=? ORDER BY key", (self.kind,)): + yield key + + def __contains__(self, key): + return self.db.execute("SELECT 1 FROM records WHERE kind=? AND key=?", (self.kind, key)).fetchone() is not None + + def __getitem__(self, key): + row = self.db.execute("SELECT value FROM records WHERE kind=? AND key=?", (self.kind, key)).fetchone() + if row is None: + raise KeyError(key) + return json.loads(row[0]) + + def __setitem__(self, key, value): + if key not in self: + self.count += 1 + self.db.execute("INSERT OR REPLACE INTO records VALUES (?, ?, ?)", (self.kind, key, json.dumps(value))) + + def values(self): + for (value,) in self.db.execute("SELECT value FROM records WHERE kind=? ORDER BY key", (self.kind,)): + yield json.loads(value) + + +class DiskSequence(Sequence): + def __init__(self, db, kind): + self.db, self.kind = db, kind + self.count = 0 + + def __len__(self): + return self.count + + def __getitem__(self, position): + if position < 0: + position += self.count + row = self.db.execute("SELECT value FROM sequences WHERE kind=? AND position=?", (self.kind, position)).fetchone() + if row is None: + raise IndexError(position) + return json.loads(row[0]) + + def __setitem__(self, position, value): + self.db.execute("INSERT OR REPLACE INTO sequences VALUES (?, ?, ?)", (self.kind, position, json.dumps(value))) + + def append(self, value): + self[self.count] = value + self.count += 1 + + def __iter__(self): + for (value,) in self.db.execute("SELECT value FROM sequences WHERE kind=? ORDER BY position", (self.kind,)): + yield json.loads(value) + + +class DiskSet(Set): + def __init__(self, db, kind): + self.rows = DiskMap(db, kind) + + def __len__(self): + return len(self.rows) + + def __contains__(self, value): + return json.dumps(value) in self.rows + + def __iter__(self): + for key in self.rows: + value = json.loads(key) + yield tuple(value) if isinstance(value, list) else value + + def add(self, value): + self.rows[json.dumps(value)] = None + + def update(self, values): + for value in values: + self.add(value) + + +class AdvisoryRows(DiskMap): + def __iter__(self): + return self.values() + + def ordered(self, seed): + import random + + from .source import ECOSYSTEMS, ecosystem + + rng = random.Random(seed) + groups = [DiskSequence(self.db, "order_" + eco) for eco in ECOSYSTEMS] + for row in self.values(): # key order is the original ghsa_id sort + eco = ecosystem(row) + if eco in ECOSYSTEMS: + groups[ECOSYSTEMS.index(eco)].append(row) + for group in groups: + rng.shuffle(group) # Same Fisher-Yates draws as the in-memory implementation. + cursors = [ + self.db.execute( + "SELECT value FROM sequences WHERE kind=? " + "ORDER BY substr(json_extract(value, '$.published_at'), 1, 10) <= '2026-06-01', position", + (group.kind,), + ) + for group in groups + ] + for _ in range(max(map(len, groups), default=0)): + for cursor in cursors: + row = cursor.fetchone() + if row is not None: + yield json.loads(row[0]) + + +class FixIndex: + def __init__(self, db): + self.db = db + db.execute("CREATE TABLE fixes (repository TEXT, sha TEXT, PRIMARY KEY (repository, sha)) WITHOUT ROWID") + + def add(self, repository, sha): + self.db.execute("INSERT OR IGNORE INTO fixes VALUES (?, ?)", (repository, sha)) + + def get(self, repository, default=None): + fixes = RepositoryFixes(self.db, repository) + return fixes if fixes else default + + +class RepositoryFixes(Set): + def __init__(self, db, repository): + self.db, self.repository = db, repository + + def __len__(self): + return self.db.execute("SELECT count(*) FROM fixes WHERE repository=?", (self.repository,)).fetchone()[0] + + def __iter__(self): + for (sha,) in self.db.execute("SELECT sha FROM fixes WHERE repository=? ORDER BY sha", (self.repository,)): + yield sha + + def __contains__(self, sha): + return self.db.execute("SELECT 1 FROM fixes WHERE repository=? AND sha=?", (self.repository, sha)).fetchone() is not None diff --git a/tests/test_unseen_extract.py b/tests/test_unseen_extract.py index be81071..04b8515 100644 --- a/tests/test_unseen_extract.py +++ b/tests/test_unseen_extract.py @@ -289,7 +289,7 @@ def test_draw_seed_proportional_cells_floor_and_thirds(): def test_git_ordinary_collection_generated_attributes_and_headers(history): git, commit, _, _, _, _, _ = history commit("# Generated by fixture\nx = 2\n", extra={".gitattributes": "other.py linguist-generated=true\n", "other.py": "x = 1\n"}) - rows = e.ordinary_history(git, "HEAD") + rows = list(e.ordinary_history(git, "HEAD")) assert rows and rows[0]["files"] files = {f["path"]: f for f in rows[0]["files"]} assert files["app.py"]["generated"] and files["other.py"]["generated"] diff --git a/tests/test_unseen_memory.py b/tests/test_unseen_memory.py new file mode 100644 index 0000000..e85ca62 --- /dev/null +++ b/tests/test_unseen_memory.py @@ -0,0 +1,244 @@ +"""Offline memory and semantic regression checks using synthetic identities only.""" + +import importlib +import json +import sys +import tracemalloc +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +d = importlib.import_module("benchmarks.unseen.draft") +s = importlib.import_module("benchmarks.unseen.source") +j = importlib.import_module("benchmarks.unseen.journal") +storage = importlib.import_module("benchmarks.unseen.storage") + + +def advisory(index): + return { + "ghsa_id": f"GHSA-test-{index:08d}", + "published_at": "2026-06-02T00:00:00Z" if index % 3 else "2025-01-01T00:00:00Z", + "type": "reviewed", + "withdrawn_at": None, + "description": "An affected handler accepts untrusted input. " * 400, + "cwes": [{"cwe_id": "CWE-89", "name": "SQL injection"}], + "references": [ + {"url": f"https://github.com/fixture/project{index}/commit/{index:040x}"}, + *[{"url": f"https://example.invalid/details/{n}"} for n in range(30)], + ], + "vulnerabilities": [ + { + "package": {"ecosystem": s.ECOSYSTEMS[index % 4], "name": f"package-{n}"}, + "vulnerable_version_range": ", ".join(f">= {v}.0, < {v}.9" for v in range(50)), + "first_patched_version": "50.0", + } + for n in range(8) + ], + } + + +class OfflineAPI: + def advisory_page(self, page): + pytest.fail("A complete journal must not fetch advisories") + + def get_repo(self, name): + return {"full_name": name} + + +def test_replay_20000_advisories_has_bounded_peak_memory(tmp_path): + path = tmp_path / "journal" + path.mkdir() + (path / "inputs.json").write_text('{"seed": 19}') + journal_file = path / "advisories.jsonl" + # Write old, unprojected REST pages without ever building the entire fixture in RAM. + with journal_file.open("w") as handle: + for page in range(200): + json.dump( + { + "id": str(page + 1), + "value": {"rows": [advisory(i) for i in range(page * 100, (page + 1) * 100)], "done": page == 199}, + "elapsed_seconds": page, + "rate_limit": {"remaining": "10"}, + }, + handle, + ) + handle.write("\n") + with journal_file.open("rb") as handle: + first_line = handle.readline() + assert len(json.loads(first_line)["value"]["rows"]) == 100 + assert len(first_line) > 2 * 1024 * 1024 + del first_line + input_bytes = journal_file.stat().st_size + assert input_bytes > 400 * 1024 * 1024 + print(f"journal_bytes_read={input_bytes}") + tracemalloc.start() + try: + with j.Journal(path).open({"seed": 19}) as journal: + assert len(journal.rows["advisories"]) == 200 + assert journal.elapsed_before == 199 + rows = d.enumerate_advisories(OfflineAPI(), journal) + assert len(rows) == 20_000 + index = d.advisory_index(rows, OfflineAPI(), db=journal.db) + assert index.get("fixture/project19999") == {f"{19999:040x}"} + count = 0 + for row in s.ordered_advisories(rows, 19): + assert "description" not in row + assert all(set(v) == {"package"} for v in row["vulnerabilities"]) + count += 1 + assert count == 20_000 + assert journal.last == {"stage": "advisories", "id": "200"} + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + print(f"replay_enumeration_order_index_peak_bytes={peak}") + assert peak < 64 * 1024 * 1024 + assert not list(path.glob("sourcing-*")) + + +def test_disk_order_matches_original_seeded_order_and_duplicate_last_wins(tmp_path): + raw = [advisory(i) for i in range(93)] + raw[0]["vulnerabilities"] = [{"package": {"ecosystem": "other", "name": "fixture"}}] + # One advisory lists multiple ecosystems; retain the original ecosystem priority. + raw[1]["vulnerabilities"].append({"package": {"ecosystem": "pip", "name": "extra"}}) + expected = [s.compact_advisory(row) for row in raw] + with storage.working_store(tmp_path) as db: + rows = storage.AdvisoryRows(db, "advisories") + for row in reversed(expected): + rows[row["ghsa_id"]] = row + rows[expected[2]["ghsa_id"]] = expected[2] + assert len(rows) == len(expected) + assert list(s.ordered_advisories(rows, 17)) == s.ordered_advisories(expected, 17) + + +def test_compaction_preserves_candidate_and_all_fix_links(): + fixtures = importlib.import_module("test_unseen_source") + row = fixtures.advisory() + assert s.candidate(s.compact_advisory(row), fixtures.API(), fixtures.OSV()) == s.candidate(row, fixtures.API(), fixtures.OSV()) + row["references"] += [ + "https://github.com/fixture/old/commit/ABCDEF1#diff", + {"url": "https://github.com/fixture/other/commit/" + "b" * 40}, + ] + assert s.fix_links(s.compact_advisory(row)) == s.fix_links(row) + repo = { + "full_name": "fixture/new", + "private": False, + "size": 10, + "clone_url": "https://github.com/fixture/new", + "source": {"full_name": "fixture/source", "description": "large"}, + "parent": {"full_name": "fixture/parent", "description": "large"}, + "former_names": ["fixture/old"], + } + eligibility = importlib.import_module("benchmarks.unseen.eligibility") + + class API: + def __init__(self, data): + self.data = data + + def get_repo(self, name): + return self.data + + assert eligibility.resolve("fixture/old", API(repo)) == eligibility.resolve("fixture/old", API(s.compact_repository(repo))) + + +def test_legacy_journal_migration_keeps_payload_and_refuses_changed_parameters(tmp_path): + fixtures = importlib.import_module("test_unseen_draft") + used = fixtures.e.UsedSet() + parameters = d.inputs(used, 19, set(used.sources), fixtures.FakeClones()) + legacy = d.legacy_inputs(parameters) + path = tmp_path / "journal" + with j.Journal(path).open(legacy) as journal: + journal.put("advisories", "1", {"rows": [advisory(1)], "done": True}) + original = (path / "advisories.jsonl").read_bytes() + for changed in ({**parameters, "seed": 20}, {**parameters, "design_taxonomy": "different"}): + with pytest.raises(j.ResumeMismatch), j.Journal(path).open(changed, compatible_parameters=d.legacy_inputs(changed)): + pass + with j.Journal(path).open(parameters, compatible_parameters=legacy) as journal: + assert journal.rows["advisories"]["1"]["rows"][0]["ghsa_id"] == advisory(1)["ghsa_id"] + assert json.loads((path / "inputs.json").read_text()) == parameters + assert json.loads((path / "inputs-before-memory-fix.json").read_text()) == legacy + assert (path / "advisories.jsonl").read_bytes() == original + + +def test_replay_duplicate_torn_and_oversized_records_fail_closed(tmp_path, monkeypatch): + path = tmp_path / "journal" + with j.Journal(path).open({}) as journal: + journal.put("advisories", "1", {"rows": [], "done": True}) + record = (path / "advisories.jsonl").read_bytes() + for data, reason in ((record * 2, "duplicate_unit"), (record[:-1], "incomplete_record")): + (path / "advisories.jsonl").write_bytes(data) + with pytest.raises(ValueError, match=reason), j.Journal(path).open({}): + pass + assert not list(path.glob("sourcing-*")) + monkeypatch.setattr(j, "MAX_RECORD_BYTES", 16) + (path / "advisories.jsonl").write_bytes(record) + with pytest.raises(ValueError, match="record_size_cap"), j.Journal(path).open({}): + pass + + +def test_ordinary_sampling_streams_payloads_and_preserves_frozen_draw(): + fixtures = importlib.import_module("test_unseen_extract") + extract = fixtures.e + rows = [fixtures.ordinary(i) for i in range(120)] + for i, row in enumerate(rows): + row["files"][0]["added"] = [5, 20, 100, 500][i % 4] + # Captured from the pre-memory-fix implementation, not a second copy of its algorithm. + drawn, counts = extract.draw_ordinary(iter(rows), set(), set(), seed=19) + assert [int(row["head"], 16) for row in drawn] == [0, 33, 6, 35, 52, 65, 62, 75, 96, 117, 90, 119] + assert counts == {"eligible": 120} + + def history(): + for i in range(20_000): + row = fixtures.ordinary(i) + row["message"] = "ordinary change " * 1500 + yield row + + tracemalloc.start() + try: + drawn, counts = extract.draw_ordinary(history(), set(), set(), seed=19) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert counts == {"eligible": 20_000} + assert len(drawn) == 12 + assert peak < 16 * 1024 * 1024 + + +def test_git_spools_large_output_and_refuses_large_blob_before_reading(tmp_path): + import subprocess + + extract = importlib.import_module("benchmarks.unseen.extract") + + def runner(argv, **kwargs): + assert kwargs["stdout"] != subprocess.PIPE + assert kwargs["stderr"] != subprocess.PIPE + for _ in range(9): + kwargs["stdout"].write(b"x" * 1024 * 1024) + return subprocess.CompletedProcess(argv, 0) + + with pytest.raises(extract.InstrumentFailure, match="git_output_size_cap"): + extract.Git(tmp_path, runner=runner).run("log") + calls = [] + git = extract.Git(tmp_path) + + def run(*args): + calls.append(args[0]) + return "app.py" if args[0] == "ls-tree" else str(512 * 1024 + 1) + + git.run = run + with pytest.raises(extract.InstrumentFailure, match="source_blob_size_cap"): + git.blob("a" * 40, "app.py") + assert calls == ["ls-tree", "cat-file"] + assert git.bytes_read == git.files_read == 0 + + +def test_repository_cache_is_bounded_and_drops_payloads(): + class Transport: + def get(self, url, headers): + return {"full_name": url.rsplit("/", 2)[-2] + "/" + url.rsplit("/", 1)[-1], "description": "large" * 1000} + + github = s.GitHub("offline", transport=Transport()) + for i in range(200): + github.get_repo(f"fixture/project{i}") + assert len(github.repos) == 128 + assert all("description" not in row for row in github.repos.values())