ci(backend): Copy per-worker Snuba databases instead of bootstrapping each

.github/workflows/backend.yml

@@ -352,6 +352,10 @@ jobs:
           name: selected-tests-${{ github.run_id }}
           path: .artifacts/
 
+      - name: Attribute resource usage to the tests from here
+        continue-on-error: true
+        run: python3 .github/workflows/scripts/monitor-resources.py --phase tests
+
       - name: Run backend test (${{ steps.setup.outputs.matrix-instance-number }} of ${{ steps.setup.outputs.matrix-instance-total }})
         env:
           SELECTED_TESTS_FILE: ${{ needs.select-tests.outputs.has-selected-tests == 'true' && '.artifacts/selected-tests.txt' || '' }}
@@ -374,6 +378,13 @@ jobs:
         continue-on-error: true
         run: python3 .github/workflows/scripts/monitor-resources.py --stop
 
+      # bootstrap-snuba.py turns off ClickHouse's query log for the per-worker
+      # databases, so they should log no queries.
+      - name: ClickHouse queries logged
+        if: ${{ !cancelled() && env.XDIST_PER_WORKER_SNUBA == '1' }}
+        continue-on-error: true
+        run: python3 .github/workflows/scripts/bootstrap-snuba.py --query-log
+
       # pytest-selected.py balances full-run shards with these. See
       # .github/workflows/scripts/backend-test-durations.py to refresh them.
       - name: Extract test durations

.github/workflows/scripts/bootstrap-snuba.py

@@ -1,14 +1,17 @@
 #!/usr/bin/env python3
 """Bootstrap per-worker Snuba instances for CI.
 
-Phase 1: Create per-worker databases and run ``snuba bootstrap --force``
-  for each of them.
+Before phase 1: Load the CI ClickHouse config (clickhouse/), which stops
+  ClickHouse logging and profiling every query.
+Phase 1: Run ``snuba bootstrap --force`` for the first worker's database,
+  then copy it to the other workers with BACKUP/RESTORE.
 Phase 2: Stop snuba-snuba-1 and start per-worker API containers.
 
 Usage:
-  bootstrap-snuba.py           Run in the foreground, after devservices is up.
-  bootstrap-snuba.py --start   Run in the background.
-  bootstrap-snuba.py --wait    Wait for the background run and exit with its status.
+  bootstrap-snuba.py              Run in the foreground, after devservices is up.
+  bootstrap-snuba.py --start      Run in the background.
+  bootstrap-snuba.py --wait       Wait for the background run and exit with its status.
+  bootstrap-snuba.py --query-log  Print how many queries ClickHouse logged per database.
 
 With DEVSERVICES_STATUS_FILE set (the exit status of a background
 ``devservices up``), phase 1 starts once ClickHouse, Kafka and Redis are
@@ -43,6 +46,20 @@
 
 ENV_ARGS = [flag for k, v in SNUBA_ENV.items() for flag in ("-e", f"{k}={v}")]
 
+SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__))
+MONITOR = os.path.join(SCRIPTS_DIR, "monitor-resources.py")
+CLICKHOUSE = "snuba-clickhouse-1"
+CLICKHOUSE_URL = "http://localhost:8123/"
+# zz- so they merge after the sentry.xml that Snuba's devservices config mounts.
+CLICKHOUSE_CI_CONFIG = {
+    "clickhouse/ci-profile.xml": "/etc/clickhouse-server/users.d/zz-ci-profile.xml",
+    "clickhouse/ci-logger.xml": "/etc/clickhouse-server/config.d/zz-ci-logger.xml",
+}
+
+
+def worker_db(worker_id: int) -> str:
+    return f"default_gw{worker_id}"
+
 
 def redis_db_args(worker_id: int) -> list[str]:
     # Each test's Snuba reset (/tests/<dataset>/drop) flushes Snuba's Redis DB. On a
@@ -53,7 +70,7 @@ def redis_db_args(worker_id: int) -> list[str]:
 
 
 # What `snuba bootstrap` talks to. In early mode phase 1 only waits for these.
-BOOTSTRAP_DEPENDENCIES = ("snuba-clickhouse-1", "kafka-kafka-1", "redis-redis-1")
+BOOTSTRAP_DEPENDENCIES = (CLICKHOUSE, "kafka-kafka-1", "redis-redis-1")
 
 # Set by the workflow when setup-sentry starts devservices in the background.
 DEVSERVICES_STATUS_FILE = os.environ.get("DEVSERVICES_STATUS_FILE", "")
@@ -97,6 +114,17 @@ def http_ok(url: str) -> bool:
         return False
 
 
+def clickhouse_query(sql: str, *, session: str | None = None, timeout: int = 30) -> str:
+    url = f"{CLICKHOUSE_URL}?session_id={session}" if session else CLICKHOUSE_URL
+    with urlopen(url, sql.encode(), timeout=timeout) as r:
+        return r.read().decode()
+
+
+def set_monitor_phase(phase: str) -> None:
+    """Attribute resource-monitor samples to a phase. Best effort."""
+    subprocess.run([sys.executable, MONITOR, "--phase", phase], check=False)
+
+
 def docker(
     *args: str, check: bool = False, timeout: int | None = None
 ) -> subprocess.CompletedProcess[str]:
@@ -204,15 +232,52 @@ def wait_for_devservices(timeout: int = PREREQ_TIMEOUT) -> None:
     log(f"devservices up finished ({time.monotonic() - start:.0f}s)")
 
 
+def configure_clickhouse() -> None:
+    """Load the CI ClickHouse config. Best effort, since it only saves CPU.
+
+    Profile settings apply when a session starts, so this runs before any
+    per-worker Snuba connects. snuba-snuba-1 may have connected already, but
+    phase 2 stops it.
+    """
+    try:
+        for src, dest in CLICKHOUSE_CI_CONFIG.items():
+            docker("cp", os.path.join(SCRIPTS_DIR, src), f"{CLICKHOUSE}:{dest}", check=True)
+        clickhouse_query("SYSTEM RELOAD CONFIG")
+        log_queries = clickhouse_query("SELECT getSetting('log_queries')").strip()
+    except (subprocess.CalledProcessError, URLError, OSError) as e:
+        log(f"::warning::Could not load the CI ClickHouse config: {e}")
+        return
+    log(f"Loaded the CI ClickHouse config (log_queries={log_queries})")
+
+
+def clone_databases(source: int, targets: range) -> None:
+    """Copy a bootstrapped worker database, schema and migration state, to others.
+
+    `snuba bootstrap` runs about 700 DDL statements per database, mostly ALTERs
+    of wide tables, for about 15 CPU-seconds of ClickHouse. RESTORE ... AS
+    rewrites the database in materialized views and Merge tables and copies
+    them all in under a second.
+    """
+    src = worker_db(source)
+    # Memory backups only live as long as the session that made them.
+    session = f"bootstrap-snuba-{os.getpid()}"
+    clickhouse_query(f"BACKUP DATABASE {src} TO Memory('{src}')", session=session)
+    for i in targets:
+        clickhouse_query(f"DROP DATABASE IF EXISTS {worker_db(i)} SYNC", session=session)
+        clickhouse_query(
+            f"RESTORE DATABASE {src} AS {worker_db(i)} FROM Memory('{src}')",
+            session=session,
+            timeout=120,
+        )
+        log(f"Copied {src} to {worker_db(i)}")
+
+
 def bootstrap_worker(worker_id: int, *, image: str, network: str) -> None:
     """Create a ClickHouse database and run snuba bootstrap."""
-    db = f"default_gw{worker_id}"
+    db = worker_db(worker_id)
 
     def create_db() -> None:
-        with urlopen(
-            "http://localhost:8123/", f"CREATE DATABASE IF NOT EXISTS {db}".encode(), timeout=30
-        ):
-            pass
+        clickhouse_query(f"CREATE DATABASE IF NOT EXISTS {db}")
 
     retry(create_db, label=f"CREATE DATABASE {db}")
 
@@ -240,7 +305,7 @@ def run_bootstrap() -> None:
 
 def start_worker_container(worker_id: int, *, image: str, network: str) -> None:
     """Start a per-worker Snuba API container and wait for health."""
-    db = f"default_gw{worker_id}"
+    db = worker_db(worker_id)
     port = 1230 + worker_id
     name = f"snuba-gw{worker_id}"
 
@@ -288,9 +353,19 @@ def main() -> None:
 
     wait_for_prerequisites()
     image, network = inspect_snuba_container()
+    configure_clickhouse()
 
     log("Phase 1: bootstrapping ClickHouse databases")
-    run_parallel(partial(bootstrap_worker, image=image, network=network), workers)
+    set_monitor_phase("snuba-bootstrap")
+    bootstrap = partial(bootstrap_worker, image=image, network=network)
+    run_parallel(bootstrap, workers[:1])
+    if len(workers) > 1:
+        try:
+            clone_databases(workers[0], workers[1:])
+        except (URLError, OSError) as e:
+            log(f"::warning::Copying {worker_db(workers[0])} failed ({e}), bootstrapping each")
+            run_parallel(bootstrap, workers[1:])
+    set_monitor_phase("setup")
     log(f"Phase 1 done ({time.monotonic() - start:.0f}s)")
 
     wait_for_devservices()
@@ -312,6 +387,19 @@ def main() -> None:
     sys.exit(rc)
 
 
+def report_query_log() -> None:
+    """Shows whether the CI profile reached Snuba's connections: the
+    default_gw* databases should log no queries."""
+    clickhouse_query("SYSTEM FLUSH LOGS", timeout=60)
+    log(
+        clickhouse_query(
+            "SELECT current_database, count() AS logged_queries FROM system.query_log "
+            "WHERE type = 'QueryFinish' GROUP BY current_database "
+            "ORDER BY current_database FORMAT PrettyCompactMonoBlock"
+        )
+    )
+
+
 def start_background() -> None:
     os.makedirs(STATE_DIR, exist_ok=True)
     with open(LOG_FILE, "w") as log_file:
@@ -361,5 +449,7 @@ def wait_background(timeout: int = WAIT_TIMEOUT) -> None:
             run_background()
         case ["--wait"]:
             wait_background()
+        case ["--query-log"]:
+            report_query_log()
         case _:
             fail(f"Unknown arguments: {sys.argv[1:]}")

.github/workflows/scripts/clickhouse/ci-logger.xml

@@ -0,0 +1,9 @@
+<!--
+CI only, installed by bootstrap-snuba.py: the image logs at trace level, to a
+file and to system.text_log. The logger level reloads without a restart.
+-->
+<clickhouse>
+    <logger>
+        <level>warning</level>
+    </logger>
+</clickhouse>

.github/workflows/scripts/clickhouse/ci-profile.xml

@@ -0,0 +1,16 @@
+<!--
+CI only, installed by bootstrap-snuba.py: don't log or sample-profile every
+query. Profile settings apply when a session starts, so this has to be loaded
+before the per-worker Snuba containers connect.
+-->
+<clickhouse>
+    <profiles>
+        <default>
+            <log_queries>0</log_queries>
+            <log_processors_profiles>0</log_processors_profiles>
+            <query_profiler_real_time_period_ns>0</query_profiler_real_time_period_ns>
+            <query_profiler_cpu_time_period_ns>0</query_profiler_cpu_time_period_ns>
+            <memory_profiler_step>0</memory_profiler_step>
+        </default>
+    </profiles>
+</clickhouse>

.github/workflows/scripts/monitor-resources.py

@@ -15,6 +15,10 @@
 Each sample also attributes CPU to process groups: docker containers by name,
 xdist workers, the pytest controller, and other host processes by command name.
 CPU that exited between samples shows up as `unattributed`.
+
+Steps can name the phase the job is in with `--phase NAME` (`setup` until the
+first one). The sampler takes a sample as soon as the phase changes, so each
+sample, and the summary's `cpu_s_by_phase`, covers exactly one phase.
 """
 
 from __future__ import annotations
@@ -28,6 +32,7 @@
 import sys
 import time
 from pathlib import Path
+from typing import Any
 
 INTERVAL_S = 5.0
 STOP_TIMEOUT_S = 30
@@ -36,6 +41,9 @@
 PID_FILE = STATE_DIR / "pid"
 STOP_FILE = STATE_DIR / "stop"
 LOG_FILE = STATE_DIR / "log"
+PHASE_FILE = STATE_DIR / "phase"
+DEFAULT_PHASE = "setup"
+PHASE_POLL_S = 0.5
 
 CPU_HEADROOM_P95 = 75.0
 CPU_HEADROOM_MAX = 90.0
@@ -360,6 +368,25 @@ def cpu_s_by_group(samples: list[dict]) -> dict[str, float]:
     return {k: round(v, 1) for k, v in sorted(totals.items(), key=lambda kv: -kv[1])}
 
 
+def cpu_s_by_phase(samples: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
+    by_phase: dict[str, list[dict[str, Any]]] = {}
+    durations: dict[str, float] = {}
+    prev_ts = None
+    for s in samples:
+        phase = s.get("phase", DEFAULT_PHASE)
+        by_phase.setdefault(phase, []).append(s)
+        if prev_ts is not None:
+            durations[phase] = durations.get(phase, 0.0) + s["ts"] - prev_ts
+        prev_ts = s["ts"]
+    return {
+        phase: {
+            "duration_s": round(durations.get(phase, 0.0), 1),
+            "cpu_s_by_group": cpu_s_by_group(phase_samples),
+        }
+        for phase, phase_samples in by_phase.items()
+    }
+
+
 def summarize(samples: list[dict]) -> dict:
     cpus = [s["cpu_pct"] for s in samples if "cpu_pct" in s]
     mems = [s["mem_used_pct"] for s in samples if "mem_used_pct" in s]
@@ -416,6 +443,7 @@ def summarize(samples: list[dict]) -> dict:
         "cpu_s_by_group_during_tests": cpu_s_by_group(
             [s for s in samples if "pytest:worker" in s.get("cpu_s_by_group", {})]
         ),
+        "cpu_s_by_phase": cpu_s_by_phase(samples),
         "headroom": {
             "cpu": cpu_verdict,
             "memory": mem_verdict,
@@ -425,23 +453,48 @@ def summarize(samples: list[dict]) -> dict:
     return summary
 
 
+def read_phase() -> str:
+    try:
+        return PHASE_FILE.read_text().strip() or DEFAULT_PHASE
+    except OSError:
+        return DEFAULT_PHASE
+
+
+def set_phase(phase: str) -> None:
+    STATE_DIR.mkdir(parents=True, exist_ok=True)
+    tmp = PHASE_FILE.with_suffix(".tmp")
+    tmp.write_text(phase)
+    tmp.replace(PHASE_FILE)
+
+
+def wait_for_next_sample(phase: str, interval: float) -> str:
+    """Waits out the interval, or until the phase changes. Returns the phase."""
+    deadline = time.monotonic() + interval
+    while (current := read_phase()) == phase and time.monotonic() < deadline:
+        time.sleep(min(PHASE_POLL_S, interval))
+    return current
+
+
 def run_sampler(output: Path, interval: float) -> None:
     output.parent.mkdir(parents=True, exist_ok=True)
     tracker = ProcessCpuTracker()
     tracker.sample()
     prev_cpu = read_cpu_times()
+    phase = read_phase()
     with open(output, "a") as f:
         while not STOP_FILE.exists():
-            time.sleep(interval)
+            next_phase = wait_for_next_sample(phase, interval)
             entry, prev_cpu = sample(prev_cpu, tracker)
             if len(entry) > 1:
+                entry["phase"] = phase
                 f.write(json.dumps(entry) + "\n")
                 f.flush()
+            phase = next_phase
 
 
 def start_background(output: Path, interval: float) -> None:
     STATE_DIR.mkdir(parents=True, exist_ok=True)
-    for stale in (PID_FILE, STOP_FILE):
+    for stale in (PID_FILE, STOP_FILE, PHASE_FILE):
         try:
             stale.unlink()
         except OSError:
@@ -514,6 +567,12 @@ def stop_background(output: Path, summary_path: Path | None) -> int:
     print(line)
     top = list(summary["cpu_s_by_group_during_tests"].items())[:10]
     print("cpu-seconds during tests: " + ", ".join(f"{k}={v}" for k, v in top))
+    for phase, usage in summary["cpu_s_by_phase"].items():
+        top = list(usage["cpu_s_by_group"].items())[:5]
+        print(
+            f"cpu-seconds in phase {phase} ({usage['duration_s']}s): "
+            + ", ".join(f"{k}={v}" for k, v in top)
+        )
     if step_summary := os.environ.get("GITHUB_STEP_SUMMARY"):
         with open(step_summary, "a") as f:
             f.write(f"### Resource usage ({summary['samples']} samples)\n\n")
@@ -532,16 +591,19 @@ def main() -> int:
     parser.add_argument("--start", action="store_true")
     parser.add_argument("--stop", action="store_true")
     parser.add_argument("--run-background", action="store_true")
+    parser.add_argument("--phase", help="attribute the following samples to this phase")
     args = parser.parse_args()
 
-    if args.start:
+    if args.phase:
+        set_phase(args.phase)
+    elif args.start:
         start_background(args.output, args.interval)
     elif args.run_background:
         run_background(args.output, args.interval)
     elif args.stop:
         stop_background(args.output, args.summary)
     else:
-        parser.error("expected one of --start, --stop, --run-background")
+        parser.error("expected one of --start, --stop, --run-background, --phase")
     return 0