diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6a597a8730e..c5b2fbd1e3e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -224,6 +224,40 @@ jobs: uv run --no-project --with pytest --with xxhash \ pytest scripts/tests/test_measurement_id.py + benchmark-ingest: + name: "Benchmark ingestion" + runs-on: ubuntu-latest + timeout-minutes: 10 + services: + postgres: + image: postgres:16-alpine + env: + POSTGRES_PASSWORD: postgres + ports: + - 5432:5432 + options: >- + --health-cmd pg_isready + --health-interval 5s + --health-timeout 5s + --health-retries 10 + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + with: + repository: vortex-data/benchmarks-website + ref: 8115d47800fc1e9e1c70fdb8022cf6120a4fad09 + path: benchmarks-website + - uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 + with: + sync: false + - name: Pytest - benchmark writer against website migrations + env: + BENCH_TEST_POSTGRES_DSN: postgresql://postgres:postgres@localhost:5432/postgres + BENCH_WEBSITE_DIR: ${{ github.workspace }}/benchmarks-website + run: | + uv run --no-project --with pytest --with xxhash --with 'psycopg[binary]>=3.2' \ + pytest scripts/tests/test_post_ingest.py + benchmark-reporting: name: "Benchmark reporting" runs-on: ubuntu-latest diff --git a/scripts/post-ingest.py b/scripts/post-ingest.py index 120bb8d3134..658b88c6184 100755 --- a/scripts/post-ingest.py +++ b/scripts/post-ingest.py @@ -27,8 +27,10 @@ import json import math import os +import random import subprocess import sys +import time import urllib.request from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -419,7 +421,7 @@ def _memory_quartet_consistent(r: dict) -> bool: def _validate_record_values(record: dict, kind: str, index: int) -> None: """Validate every field's type/range before the Postgres write. - Runs in `ingest_postgres`'s loop, where the record index is known. It drives + Runs before any transaction, where the record index is known. It drives type/range checks from `_FIELD_TYPES`, then applies semantic checks the type alone does not cover (the storage enum + memory quartet for query_measurements). """ @@ -736,60 +738,82 @@ def _upsert_commit(conn, commit: dict) -> None: ) -# CI runs concurrent writers whose upserts can touch commits and dimensions in -# conflicting orders. Retrying transaction-level deadlocks and serialization -# failures keeps each JSONL file all-or-nothing. -_WRITE_CONFLICT_ATTEMPTS = 128 +# Eight attempts cap repeated conflicts. The elapsed budget includes transaction work and backoff, +# but does not cancel an active transaction. Per-statement and lock timeouts bound individual waits. +_WRITE_CONFLICT_ATTEMPTS = 8 +_WRITE_CONFLICT_BUDGET_SECONDS = 120.0 +_WRITE_CONFLICT_BACKOFF_SECONDS = 10.0 def _retry_write_conflicts(op): - """Retry `op` on a Postgres write conflict. - - Row-level `ON CONFLICT DO UPDATE` upserts touching the same commits or - dimensions in conflicting orders can deadlock. The retryable Postgres errors are deadlock - (`SQLSTATE 40P01`) and serialization failure (`40001`); both abort one transaction cleanly, - so re-running the whole transaction is safe. A non-retryable error (e.g. a validation - `SystemExit`) propagates immediately. Returns `op`'s value on the first success. - """ - from psycopg import errors as pg_errors + """Retry whole transactions only for deadlocks (40P01) and serialization failures (40001).""" + import psycopg + started = time.monotonic() for attempt in range(1, _WRITE_CONFLICT_ATTEMPTS + 1): try: - return op() - except (pg_errors.DeadlockDetected, pg_errors.SerializationFailure): - # The failing `op`'s `with conn.transaction()` block already rolled back, so the - # connection is idle and the whole transaction can be retried. Re-raise on the - # final attempt. - if attempt >= _WRITE_CONFLICT_ATTEMPTS: + result = op() + except psycopg.Error as exc: + elapsed = time.monotonic() - started + retryable = exc.sqlstate in ("40P01", "40001") + remaining = _WRITE_CONFLICT_BUDGET_SECONDS - elapsed + if not retryable or attempt == _WRITE_CONFLICT_ATTEMPTS or remaining <= 0: + outcome = "exhausted" if retryable else "failed" + print( + f"ingest attempt={attempt} elapsed={elapsed:.3f}s sqlstate={exc.sqlstate} outcome={outcome}", + file=sys.stderr, + ) + raise + + ceiling = min(2 ** (attempt - 1), _WRITE_CONFLICT_BACKOFF_SECONDS) + delay = min(random.uniform(ceiling / 2, ceiling), remaining) + print( + f"ingest attempt={attempt} elapsed={elapsed:.3f}s sqlstate={exc.sqlstate} " + f"outcome=retry delay={delay:.3f}s", + file=sys.stderr, + ) + time.sleep(delay) + # A delayed wakeup must not start another transaction after the retry budget expires. + elapsed = time.monotonic() - started + if elapsed >= _WRITE_CONFLICT_BUDGET_SECONDS: + print( + f"ingest attempt={attempt} elapsed={elapsed:.3f}s sqlstate={exc.sqlstate} outcome=exhausted", + file=sys.stderr, + ) raise + else: + elapsed = time.monotonic() - started + print(f"ingest attempt={attempt} elapsed={elapsed:.3f}s sqlstate=none outcome=committed", file=sys.stderr) + return result raise AssertionError("unreachable: _retry_write_conflicts exited without return or raise") def ingest_postgres(conn, commit: dict, records: list[dict]) -> tuple[int, int]: - """Upsert a commit and its records into Postgres, retrying on write conflicts.""" + """Validate the entire file before atomically upserting it, with bounded conflict retries.""" + sha = commit["sha"] + if not isinstance(sha, str) or len(sha) != 40 or any(c not in "0123456789abcdef" for c in sha): + raise SystemExit(f"commit SHA must be 40-hex lowercase, got: {sha!r}") + for idx, record in enumerate(records): + kind = _validate_record_fields(record, idx) + _validate_record_values(record, kind, idx) + if record["commit_sha"] != sha: + raise SystemExit( + f"record {idx} ({kind}): commit_sha {record['commit_sha']!r} does not " + f"match the requested commit SHA {sha!r}" + ) + mid_mod = _measurement_id_module() return _retry_write_conflicts(lambda: _ingest_postgres_once(conn, commit, records, mid_mod)) def _ingest_postgres_once(conn, commit: dict, records: list[dict], mid_mod) -> tuple[int, int]: - """Upsert a commit and its records in one transaction (a single attempt). - - Upsert `commits` first, then each fact record while classifying it as inserted - or updated. Any validation failure rolls the whole transaction back. - """ + """Upsert the validated commit and records in one transaction, rolling back on any error.""" inserted = 0 updated = 0 with conn.transaction(): _upsert_commit(conn, commit) - for idx, record in enumerate(records): - kind = _validate_record_fields(record, idx) - if record["commit_sha"] != commit["sha"]: - raise SystemExit( - f"record {idx} ({kind}): commit_sha {record['commit_sha']!r} does not " - f"match the requested commit SHA {commit['sha']!r}" - ) - _validate_record_values(record, kind, idx) - if _APPLY_RECORD[kind](conn, mid_mod, record): + for record in records: + if _APPLY_RECORD[record["kind"]](conn, mid_mod, record): updated += 1 else: inserted += 1 @@ -887,7 +911,10 @@ def connect_postgres(dsn: str, region: str | None): # left-to-right (last wins), so appending ours last makes it authoritative even if the DSN # already set one. existing_options = params.get("options") or "" - params["options"] = f"{existing_options} -c search_path=public".strip() + params["options"] = ( + f"{existing_options} -c search_path=public -c statement_timeout=30000 -c lock_timeout=10000" + ).strip() + params["connect_timeout"] = 10 conn = psycopg.connect(**params) # Verify the RESOLVED transport actually used TLS, not merely that the DSN requested @@ -965,9 +992,7 @@ def refresh_site_cache(base_url: str, token: str, timeout: float) -> None: def _main_postgres(args: argparse.Namespace) -> int: records = read_records(args.jsonl_path) # `build_commit` runs `git show `, so the SHA must be in the runner's local git - # history. The v4 ingest step inherits the v3 `--server` step's checkout assumption (the default - # checkout provides the head SHA); a shallow checkout missing the SHA fails loud here, and the - # v4 step is best-effort (continue-on-error), so it never fails the job. + # history. A shallow checkout missing the SHA fails before opening a database connection. commit = build_commit(args.commit_sha, args.repo_url, args.git_dir) conn = connect_postgres(args.postgres, args.region) try: diff --git a/scripts/tests/test_post_ingest.py b/scripts/tests/test_post_ingest.py new file mode 100644 index 00000000000..024d2c6ae62 --- /dev/null +++ b/scripts/tests/test_post_ingest.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Writer regressions, including real PostgreSQL transactions against the website migrations. + +Set BENCH_TEST_POSTGRES_DSN to a disposable superuser database and BENCH_WEBSITE_DIR to a website +checkout for integration tests. The fixture creates and drops its own database and invokes the +website's migration runner. CI pins that checkout in .github/workflows/ci.yml. +""" + +import importlib.util +import os +import subprocess +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Event +from types import SimpleNamespace + +import psycopg +import pytest +from psycopg import errors, sql +from psycopg.conninfo import make_conninfo + +REPO_ROOT = Path(__file__).resolve().parents[2] +TABLES = ("query_measurements", "compression_times", "compression_sizes", "random_access_times", "vector_search_runs") + + +@pytest.fixture +def writer(): + spec = importlib.util.spec_from_file_location("post_ingest", REPO_ROOT / "scripts/post-ingest.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def commit(): + return { + "sha": "a" * 40, + "timestamp": "2026-09-08T12:00:00Z", + "message": "benchmark", + "author_name": "Author", + "author_email": "author@example.com", + "committer_name": "Committer", + "committer_email": "committer@example.com", + "tree_sha": "b" * 40, + "url": "https://github.com/vortex-data/vortex/commit/" + "a" * 40, + } + + +@pytest.fixture +def records(commit): + common = {"commit_sha": commit["sha"], "dataset": "test", "format": "vortex"} + timing = {"value_ns": 100, "all_runtimes_ns": [99, 101]} + return [ + {**common, **timing, "kind": "query_measurement", "query_idx": 1, "storage": "nvme", "engine": "datafusion"}, + {**common, **timing, "kind": "compression_time", "op": "encode"}, + {**common, "kind": "compression_size", "value_bytes": 10, "uncompressed_bytes": 100}, + {**common, **timing, "kind": "random_access_time", "open_mode": "cached"}, + { + "kind": "vector_search_run", + "commit_sha": commit["sha"], + "dataset": "test", + "layout": "flat", + "flavor": "exact", + "threshold": 0.5, + **timing, + "matches": 1, + "rows_scanned": 10, + "bytes_scanned": 100, + "iterations": 2, + }, + ] + + +@pytest.fixture +def clock(monkeypatch, writer): + state = SimpleNamespace(now=0.0, sleeps=[]) + + def sleep(delay): + state.sleeps.append(delay) + state.now += delay + + monkeypatch.setattr(writer.time, "monotonic", lambda: state.now) + monkeypatch.setattr(writer.time, "sleep", sleep) + monkeypatch.setattr(writer.random, "uniform", lambda lower, upper: upper) + return state + + +@pytest.mark.parametrize( + "invalid", [None, {"extra": 1}, {"iterations": True}, {"threshold": float("inf")}, {"commit_sha": "c" * 40}] +) +def test_late_invalid_record_never_starts_transaction(writer, commit, records, invalid): + records[-1] = None if invalid is None else {**records[-1], **invalid} + # A connection with no methods proves validation precedes even BEGIN, including the commit upsert. + with pytest.raises(SystemExit, match="record 4"): + writer.ingest_postgres(object(), commit, records) + + +def test_invalid_commit_sha_never_starts_transaction(writer, commit): + commit["sha"] = "not-a-sha" + with pytest.raises(SystemExit, match="commit SHA"): + writer.ingest_postgres(object(), commit, []) + + +@pytest.mark.parametrize("error", [errors.DeadlockDetected, errors.SerializationFailure]) +def test_conflict_attempts_and_backoff_are_bounded(writer, clock, error, capsys): + calls = 0 + + def fail(): + nonlocal calls + calls += 1 + raise error("sensitive server detail") + + with pytest.raises(error): + writer._retry_write_conflicts(fail) + assert calls == 8 + assert clock.sleeps == [1, 2, 4, 8, 10, 10, 10] + diagnostics = capsys.readouterr().err + assert "attempt=8 elapsed=45.000s" in diagnostics + assert f"sqlstate={error.sqlstate} outcome=exhausted" in diagnostics + assert "sensitive server detail" not in diagnostics + + +@pytest.mark.parametrize("elapsed, sleeps", [(119.5, [0.5]), (121, [])]) +def test_elapsed_budget_prevents_another_attempt(writer, clock, elapsed, sleeps): + calls = 0 + + def fail(): + nonlocal calls + calls += 1 + clock.now += elapsed + raise errors.DeadlockDetected() + + with pytest.raises(errors.DeadlockDetected): + writer._retry_write_conflicts(fail) + assert calls == 1 + assert clock.sleeps == sleeps + + +def test_budget_does_not_cancel_successful_transaction(writer, clock): + def succeed(): + clock.now += 121 + return (5, 0) + + assert writer._retry_write_conflicts(succeed) == (5, 0) + assert clock.sleeps == [] + + +@pytest.mark.parametrize( + "error", [errors.CheckViolation, errors.LockNotAvailable, errors.QueryCanceled, psycopg.OperationalError] +) +def test_other_database_errors_are_not_retried(writer, clock, error, capsys): + def fail(): + raise error("sensitive server detail") + + with pytest.raises(error): + writer._retry_write_conflicts(fail) + assert clock.sleeps == [] + assert "outcome=failed" in capsys.readouterr().err + + +def test_connection_timeouts_override_dsn_options(writer, monkeypatch): + captured = {} + connection = SimpleNamespace(pgconn=SimpleNamespace(ssl_in_use=True)) + + def connect(**kwargs): + captured.update(kwargs) + return connection + + monkeypatch.setattr(psycopg, "connect", connect) + assert ( + writer.connect_postgres( + "host=example.com user=bench_ingest password=secret connect_timeout=0 " + "options='-c statement_timeout=0 -c lock_timeout=0 -c search_path=other'", + None, + ) + is connection + ) + assert captured["connect_timeout"] == 10 + assert captured["sslmode"] == "verify-full" + assert captured["options"].endswith("-c search_path=public -c statement_timeout=30000 -c lock_timeout=10000") + + +@pytest.fixture(scope="module") +def database(): + dsn = os.environ.get("BENCH_TEST_POSTGRES_DSN") + website = os.environ.get("BENCH_WEBSITE_DIR") + if not dsn or not website: + pytest.skip("set BENCH_TEST_POSTGRES_DSN and BENCH_WEBSITE_DIR to run PostgreSQL writer tests") + name = "writer_test_" + uuid.uuid4().hex + with psycopg.connect(dsn, autocommit=True) as admin: + admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name))) + target = make_conninfo(dsn, dbname=name) + try: + subprocess.run( + [sys.executable, str(Path(website) / "scripts/migrate-schema.py"), "apply", "--target", target], + check=True, + capture_output=True, + text=True, + ) + yield target + finally: + admin.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name))) + + +@pytest.fixture +def conn(database): + with psycopg.connect(database, autocommit=True) as connection: + connection.execute( + sql.SQL("TRUNCATE commits, {} CASCADE").format(sql.SQL(", ").join(map(sql.Identifier, TABLES))) + ) + connection.execute("SET ROLE bench_ingest") + yield connection + + +def stored_rows(conn): + return [ + conn.execute(sql.SQL("SELECT * FROM {} ORDER BY measurement_id").format(sql.Identifier(table))).fetchall() + for table in TABLES + ] + + +def test_replay_updates_every_family_without_duplicate_ids(writer, conn, commit, records): + assert writer.ingest_postgres(conn, commit, records) == (5, 0) + before = stored_rows(conn) + for record in records: + if "value_ns" in record: + record.update(value_ns=200, all_runtimes_ns=[199, 201]) + else: + record.update(value_bytes=20, uncompressed_bytes=200) + assert writer.ingest_postgres(conn, commit, records + records) == (0, 10) + after = stored_rows(conn) + assert [rows[0][0] for rows in after] == [rows[0][0] for rows in before] + assert all(len(rows) == 1 for rows in after) + assert all(old != new for old, new in zip(before, after, strict=True)) + assert conn.execute("SELECT count(*) FROM commits").fetchone() == (1,) + assert writer.ingest_postgres(conn, commit, records) == (0, 5) + assert stored_rows(conn) == after + + +def test_late_database_error_rolls_back_entire_file(writer, database, conn, commit, records, capsys): + assert writer.ingest_postgres(conn, commit, records) == (5, 0) + before = stored_rows(conn) + commit["message"] = "must roll back" + records[0]["value_ns"] = 999 + records.insert(1, {**records[0], "query_idx": 2}) + records[-1]["value_ns"] = 1337 + with psycopg.connect(database, autocommit=True) as admin: + admin.execute("ALTER TABLE vector_search_runs ADD CONSTRAINT reject_test_value CHECK (value_ns <> 1337)") + try: + with pytest.raises(errors.CheckViolation): + writer.ingest_postgres(conn, commit, records) + finally: + admin.execute("ALTER TABLE vector_search_runs DROP CONSTRAINT reject_test_value") + assert stored_rows(conn) == before + assert conn.execute("SELECT message FROM commits").fetchone() == ("benchmark",) + assert "attempt=1" in capsys.readouterr().err + assert conn.info.transaction_status == psycopg.pq.TransactionStatus.IDLE + + +@pytest.mark.parametrize("conflict", ["deadlock", "serialization"]) +def test_real_transaction_conflict_retries_the_entire_file( + writer, database, conn, commit, records, monkeypatch, conflict, capsys +): + other_commit = {**commit, "sha": "c" * 40} + writer.ingest_postgres(conn, other_commit, []) + writer.ingest_postgres(conn, commit, []) + reached = Event() + original = writer._APPLY_RECORD["query_measurement"] + attempts = 0 + with psycopg.connect(database, autocommit=True) as other: + if conflict == "serialization": + conn.execute("SET default_transaction_isolation = 'serializable'") + else: + # The writer detects the cycle first and becomes the deadlock victim deterministically. + conn.execute("RESET ROLE") + conn.execute("SET deadlock_timeout = '100ms'") + conn.execute("SET ROLE bench_ingest") + other.execute("SET deadlock_timeout = '5s'") + other.execute("BEGIN") + other.execute("SELECT 1 FROM commits WHERE commit_sha = %s FOR UPDATE", (other_commit["sha"],)) + + def insert_and_conflict(connection, mid_mod, record): + nonlocal attempts + attempts += 1 + result = original(connection, mid_mod, record) + if attempts == 1: + if conflict == "serialization": + other.execute( + "UPDATE commits SET message = 'concurrent' WHERE commit_sha = %s", (other_commit["sha"],) + ) + reached.set() + connection.execute("SELECT 1 FROM commits WHERE commit_sha = %s FOR UPDATE", (other_commit["sha"],)) + return result + + def complete_deadlock(): + assert reached.wait(timeout=10) + other.execute("UPDATE commits SET message = 'concurrent' WHERE commit_sha = %s", (commit["sha"],)) + other.execute("COMMIT") + + monkeypatch.setitem(writer._APPLY_RECORD, "query_measurement", insert_and_conflict) + monkeypatch.setattr(writer.random, "uniform", lambda lower, upper: lower) + with ThreadPoolExecutor(max_workers=1) as pool: + blocker = pool.submit(complete_deadlock) if conflict == "deadlock" else None + assert writer.ingest_postgres(conn, commit, records) == (5, 0) + if blocker: + blocker.result(timeout=10) + assert attempts == 2 + assert all(len(rows) == 1 for rows in stored_rows(conn)) + assert conn.execute("SELECT message FROM commits WHERE commit_sha = %s", (commit["sha"],)).fetchone() == ( + "benchmark", + ) + diagnostics = capsys.readouterr().err + sqlstate = "40P01" if conflict == "deadlock" else "40001" + assert f"sqlstate={sqlstate} outcome=retry" in diagnostics + assert "attempt=2" in diagnostics + + +def test_refresh_failure_keeps_successful_ingest(writer, database, conn, commit, records, monkeypatch, capsys): + monkeypatch.setattr(writer, "read_records", lambda path: records) + monkeypatch.setattr(writer, "build_commit", lambda *args: commit) + monkeypatch.setattr(writer, "connect_postgres", lambda *args: psycopg.connect(database)) + monkeypatch.setenv("BENCH_SITE_BASE_URL", "https://bench.example.com") + monkeypatch.setenv("BENCH_REVALIDATE_TOKEN", "secret") + + def fail_refresh(*args): + raise TimeoutError("refresh timed out") + + monkeypatch.setattr(writer, "_http", fail_refresh) + args = SimpleNamespace( + jsonl_path=None, commit_sha=commit["sha"], repo_url=None, git_dir=None, postgres=None, region=None, timeout=1 + ) + assert writer._main_postgres(args) == 0 + assert all(len(rows) == 1 for rows in stored_rows(conn)) + output = capsys.readouterr() + assert '"inserted":5,"updated":0' in output.out + assert "warning: cache revalidate failed" in output.err