From c1a9a48b0da70682ecd2ff356462fb9775c9bc6f Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Mon, 27 Jul 2026 17:30:02 +0800 Subject: [PATCH] bench: add workload scheduler A/B harness --- tests/perf/README.md | 50 ++ tests/perf/workload_scheduler_benchmark.py | 536 +++++++++++++++++++++ tests/perf/workload_scheduler_runner.py | 488 +++++++++++++++++++ 3 files changed, 1074 insertions(+) create mode 100644 tests/perf/workload_scheduler_benchmark.py create mode 100644 tests/perf/workload_scheduler_runner.py diff --git a/tests/perf/README.md b/tests/perf/README.md index 91bc6c1637..168e3fe2e0 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -9,6 +9,56 @@ automatic write, flush, SST, and query path. They do not replace controlled encoding experiments that explicitly compare `plain`/dictionary, `no_dictionary`, BYTE_STREAM_SPLIT, or Auto policies. +## Workload scheduler benchmark + +`workload_scheduler_benchmark.py` measures the experimental query/write +scheduler against an otherwise identical scheduler-disabled standalone server: + +```bash +cargo build --release -p cmd --bin greptime +python3 tests/perf/workload_scheduler_benchmark.py \ + --output /tmp/greptime-workload-scheduler.json +``` + +Each sample starts a fresh server and separate, equivalently seeded query and +write tables. Baseline and scheduled modes are interleaved for three iterations, +with four warmup seconds followed by eight measured seconds. The default +4-worker workload covers query-only, write-only, light-write, and saturated +phases. Client-to-response mean, p50, and p95 latency and successful request +throughput are reported alongside scheduler poll shares. In +`scheduled_vs_baseline_percent`, positive throughput is an improvement while +positive latency is a regression. + +Query and write requests do not represent equal work, so the raw sum of their +request rates is not used as the overhead check. The report calibrates each +iteration with its scheduler-disabled query-only and write-only capacities, then +computes: + +```text +capacity_normalized_work_rate = + query_rps / baseline_query_capacity + + write_rps / baseline_write_capacity +``` + +The scheduled and baseline mixed samples are compared using the capacities from +the same iteration. `paired_capacity_normalized.within_five_percent` verifies +that every paired sample stays within the 5% regression budget; the top-level +`verification` object combines that check with the saturated 80% write-poll +share check. This normalization prevents a policy-driven shift between +differently priced request types from being reported as scheduler overhead. + +The saturated defaults are calibrated to keep both classes runnable on a +32-logical-CPU development host. Adjust `--query-workers` and `--write-workers` +when the output's minimum write poll share shows that a different host did not +reach saturation. `workload_scheduler_runner.py` is the shorter correctness +runner for checking the 80/20 share and work-conserving borrowing against an +already-running server. + +When `max_concurrent_polls` is left at zero in GreptimeDB configuration, the +scheduler uses four times `global_rt_size`. The benchmark resolves the same +default explicitly. This keeps Tokio's worker queues fed while retaining a +bounded admission window. + ## Phase 1: direct readable SST fixtures Phase 1 should generate data by writing readable Mito SST files and matching diff --git a/tests/perf/workload_scheduler_benchmark.py b/tests/perf/workload_scheduler_benchmark.py new file mode 100644 index 0000000000..3bfd005b6e --- /dev/null +++ b/tests/perf/workload_scheduler_benchmark.py @@ -0,0 +1,536 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compare end-to-end GreptimeDB performance with the scheduler off and on. + +Every sample starts a fresh standalone server and seeds an equivalent Mito +table. Modes are interleaved to reduce time/order bias. Request latency is +measured at the HTTP client, from request submission through response parsing. +""" + +from __future__ import annotations + +import argparse +import json +import os +import socket +import statistics +import subprocess +import sys +import tempfile +import time +from pathlib import Path +from typing import Any + +import workload_scheduler_runner as workload + + +PHASE_OPTIONS = { + "query_only": ("query",), + "write_only": ("write",), + "light_write": ("query", "write"), + "saturated": ("query", "write"), +} +MAX_CAPACITY_NORMALIZED_REGRESSION_PERCENT = 5.0 + + +def reserve_ports(count: int) -> list[int]: + sockets = [] + try: + for _ in range(count): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("127.0.0.1", 0)) + sockets.append(sock) + return [int(sock.getsockname()[1]) for sock in sockets] + finally: + for sock in sockets: + sock.close() + + +def write_config( + path: Path, + enabled: bool, + ports: list[int], + runtime_size: int, + max_concurrent_polls: int, +) -> None: + path.write_text( + f"""[runtime] +global_rt_size = {runtime_size} +compact_rt_size = 1 +query_rt_size = {runtime_size} +ingest_rt_size = {runtime_size} + +[runtime.experimental_workload_scheduler] +enable = {str(enabled).lower()} +max_concurrent_polls = {max_concurrent_polls} +query_weight = 2 +write_weight = 8 + +[http] +addr = "127.0.0.1:{ports[0]}" + +[grpc] +bind_addr = "127.0.0.1:{ports[1]}" + +[mysql] +enable = false +addr = "127.0.0.1:{ports[2]}" + +[postgres] +enable = false +addr = "127.0.0.1:{ports[3]}" +""" + ) + + +def wait_for_server( + client: workload.SqlClient, + process: subprocess.Popen[bytes], + log_path: Path, + timeout: float, +) -> None: + deadline = time.monotonic() + timeout + last_error: Any = None + while time.monotonic() < deadline: + if process.poll() is not None: + tail = log_path.read_text(errors="replace")[-8_000:] + raise RuntimeError( + f"GreptimeDB exited with {process.returncode} during startup:\n{tail}" + ) + ok, _, last_error = client.sql("SELECT 1") + if ok: + return + time.sleep(0.2) + raise TimeoutError(f"GreptimeDB did not become ready: {last_error}") + + +def stop_server(process: subprocess.Popen[bytes]) -> None: + if process.poll() is not None: + return + process.terminate() + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + + +def phase_workers(args: argparse.Namespace, phase: str) -> tuple[int, int, float]: + if phase == "query_only": + return args.query_workers, 0, 0 + if phase == "write_only": + return 0, args.write_workers, 0 + if phase == "light_write": + return args.query_workers, 1, args.light_write_delay + return args.query_workers, args.write_workers, 0 + + +def run_sample( + args: argparse.Namespace, + root: Path, + phase: str, + enabled: bool, + sample_number: int, +) -> dict[str, Any]: + mode = "scheduled" if enabled else "baseline" + sample_root = root / f"{phase}-{mode}-{sample_number}" + data_home = sample_root / "data" + log_dir = sample_root / "logs" + config_path = sample_root / "config.toml" + process_log_path = sample_root / "process.log" + sample_root.mkdir(parents=True) + data_home.mkdir() + log_dir.mkdir() + ports = reserve_ports(4) + write_config( + config_path, + enabled, + ports, + args.runtime_size, + args.max_concurrent_polls, + ) + + environment = os.environ.copy() + for variable in ( + "ALL_PROXY", + "HTTPS_PROXY", + "HTTP_PROXY", + "all_proxy", + "https_proxy", + "http_proxy", + ): + environment.pop(variable, None) + + command = [ + str(args.binary), + "standalone", + "start", + "--config-file", + str(config_path), + "--data-home", + str(data_home), + "--log-dir", + str(log_dir), + "--log-level", + "warn", + ] + with process_log_path.open("wb") as process_log: + process = subprocess.Popen( + command, + stdout=process_log, + stderr=subprocess.STDOUT, + env=environment, + ) + try: + client = workload.SqlClient( + f"http://127.0.0.1:{ports[0]}", "public", args.timeout + ) + wait_for_server(client, process, process_log_path, args.start_timeout) + workload.setup_table(client, args.seed_rows, args.seed_batch_size) + query_workers, write_workers, write_delay = phase_workers(args, phase) + result = workload.run_phase( + client, + phase, + args.duration, + args.warmup, + query_workers, + write_workers, + args.write_batch_size, + write_delay, + workload.Sequence(1_800_000_000_000), + "required" if enabled else "disabled", + ) + finally: + stop_server(process) + + result["mode"] = mode + result["sample"] = sample_number + for workload_name in PHASE_OPTIONS[phase]: + requests = result["requests"][workload_name] + if requests["requests"] and requests["failures"] / requests["requests"] >= 0.01: + raise AssertionError( + f"{mode} {phase} {workload_name} failure rate was at least 1%: " + f"{requests}" + ) + if enabled and phase == "saturated": + result["saturation_verified"] = result["poll_share"]["write"] >= 0.799 + return result + + +def median_request_metric( + samples: list[dict[str, Any]], workload_name: str, metric: str +) -> float | None: + values = [ + sample["requests"][workload_name].get(metric) + for sample in samples + if sample["requests"][workload_name].get(metric) is not None + ] + return statistics.median(values) if values else None + + +def percent_change(current: float | None, baseline: float | None) -> float | None: + if current is None or baseline in (None, 0): + return None + return (current / baseline - 1.0) * 100.0 + + +def summarize(samples: list[dict[str, Any]], phases: list[str]) -> dict[str, Any]: + report: dict[str, Any] = {} + for phase in phases: + phase_samples = [sample for sample in samples if sample["name"] == phase] + modes: dict[str, Any] = {} + for mode in ("baseline", "scheduled"): + mode_samples = [ + sample for sample in phase_samples if sample["mode"] == mode + ] + query_rps = median_request_metric( + mode_samples, "query", "successful_rps" + ) + write_rps = median_request_metric( + mode_samples, "write", "successful_rps" + ) + modes[mode] = { + "query_rps": query_rps, + "write_rps": write_rps, + "total_request_rps": (query_rps or 0) + (write_rps or 0), + "query_p50_ms": median_request_metric( + mode_samples, "query", "p50_ms" + ), + "query_mean_ms": median_request_metric( + mode_samples, "query", "mean_ms" + ), + "query_p95_ms": median_request_metric( + mode_samples, "query", "p95_ms" + ), + "write_p50_ms": median_request_metric( + mode_samples, "write", "p50_ms" + ), + "write_mean_ms": median_request_metric( + mode_samples, "write", "mean_ms" + ), + "write_p95_ms": median_request_metric( + mode_samples, "write", "p95_ms" + ), + } + if mode == "scheduled": + shares = [ + sample["poll_share"]["write"] + for sample in mode_samples + if sample["poll_share"] is not None + ] + modes[mode]["write_poll_share"] = ( + statistics.median(shares) if shares else None + ) + modes[mode]["minimum_write_poll_share"] = ( + min(shares) if shares else None + ) + + baseline = modes["baseline"] + scheduled = modes["scheduled"] + report[phase] = { + "baseline": baseline, + "scheduled": scheduled, + "scheduled_vs_baseline_percent": { + "query_throughput": percent_change( + scheduled["query_rps"], baseline["query_rps"] + ), + "write_throughput": percent_change( + scheduled["write_rps"], baseline["write_rps"] + ), + "total_request_throughput": percent_change( + scheduled["total_request_rps"], baseline["total_request_rps"] + ), + "query_p50_latency": percent_change( + scheduled["query_p50_ms"], baseline["query_p50_ms"] + ), + "query_mean_latency": percent_change( + scheduled["query_mean_ms"], baseline["query_mean_ms"] + ), + "query_p95_latency": percent_change( + scheduled["query_p95_ms"], baseline["query_p95_ms"] + ), + "write_p50_latency": percent_change( + scheduled["write_p50_ms"], baseline["write_p50_ms"] + ), + "write_mean_latency": percent_change( + scheduled["write_mean_ms"], baseline["write_mean_ms"] + ), + "write_p95_latency": percent_change( + scheduled["write_p95_ms"], baseline["write_p95_ms"] + ), + }, + } + + if "query_only" in report and "write_only" in report: + query_capacity = report["query_only"]["baseline"]["query_rps"] + write_capacity = report["write_only"]["baseline"]["write_rps"] + if query_capacity and write_capacity: + for phase in phases: + modes = report[phase] + for mode in ("baseline", "scheduled"): + values = modes[mode] + values["capacity_normalized_work_rate"] = ( + values["query_rps"] / query_capacity + + values["write_rps"] / write_capacity + ) + modes["scheduled_vs_baseline_percent"][ + "capacity_normalized_work_rate" + ] = percent_change( + modes["scheduled"]["capacity_normalized_work_rate"], + modes["baseline"]["capacity_normalized_work_rate"], + ) + + samples_by_key = { + (sample["sample"], sample["name"], sample["mode"]): sample + for sample in samples + } + sample_numbers = sorted({sample["sample"] for sample in samples}) + for phase in phases: + paired = [] + for sample_number in sample_numbers: + query_only = samples_by_key.get( + (sample_number, "query_only", "baseline") + ) + write_only = samples_by_key.get( + (sample_number, "write_only", "baseline") + ) + baseline = samples_by_key.get( + (sample_number, phase, "baseline") + ) + scheduled = samples_by_key.get( + (sample_number, phase, "scheduled") + ) + if not all((query_only, write_only, baseline, scheduled)): + continue + query_capacity = query_only["requests"]["query"]["successful_rps"] + write_capacity = write_only["requests"]["write"]["successful_rps"] + if not query_capacity or not write_capacity: + continue + + def normalized(sample: dict[str, Any]) -> float: + requests = sample["requests"] + return ( + requests["query"]["successful_rps"] / query_capacity + + requests["write"]["successful_rps"] / write_capacity + ) + + baseline_rate = normalized(baseline) + scheduled_rate = normalized(scheduled) + paired.append( + { + "sample": sample_number, + "baseline": baseline_rate, + "scheduled": scheduled_rate, + "scheduled_vs_baseline_percent": percent_change( + scheduled_rate, baseline_rate + ), + } + ) + if paired: + changes = [ + sample["scheduled_vs_baseline_percent"] for sample in paired + ] + report[phase]["paired_capacity_normalized"] = { + "samples": paired, + "median_scheduled_vs_baseline_percent": statistics.median( + changes + ), + "worst_scheduled_vs_baseline_percent": min(changes), + "within_five_percent": min(changes) + >= -MAX_CAPACITY_NORMALIZED_REGRESSION_PERCENT, + } + return report + + +def parse_args() -> argparse.Namespace: + default_binary = Path(__file__).resolve().parents[2] / "target/release/greptime" + parser = argparse.ArgumentParser() + parser.add_argument("--binary", type=Path, default=default_binary) + parser.add_argument("--iterations", type=int, default=3) + parser.add_argument("--duration", type=float, default=8.0) + parser.add_argument("--warmup", type=float, default=4.0) + parser.add_argument("--runtime-size", type=int, default=4) + parser.add_argument( + "--max-concurrent-polls", + type=int, + default=0, + help="scheduler in-flight poll limit; zero uses 4 * --runtime-size", + ) + parser.add_argument("--query-workers", type=int, default=2) + parser.add_argument("--write-workers", type=int, default=1152) + parser.add_argument("--seed-rows", type=int, default=10_000) + parser.add_argument("--seed-batch-size", type=int, default=500) + parser.add_argument("--write-batch-size", type=int, default=32) + parser.add_argument("--light-write-delay", type=float, default=0.1) + parser.add_argument("--timeout", type=float, default=60.0) + parser.add_argument("--start-timeout", type=float, default=60.0) + parser.add_argument( + "--phases", + nargs="+", + choices=tuple(PHASE_OPTIONS), + default=list(PHASE_OPTIONS), + ) + parser.add_argument( + "--output", + type=Path, + help="optional path for the JSON report; stdout is always populated", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + args.binary = args.binary.resolve() + if not args.binary.is_file(): + raise FileNotFoundError(args.binary) + if args.iterations <= 0: + raise ValueError("--iterations must be greater than zero") + if args.max_concurrent_polls == 0: + args.max_concurrent_polls = args.runtime_size * 4 + + samples = [] + with tempfile.TemporaryDirectory(prefix="greptime-workload-benchmark-") as temp: + root = Path(temp) + for iteration in range(args.iterations): + for phase_index, phase in enumerate(args.phases): + order = ( + (False, True) + if (iteration + phase_index) % 2 == 0 + else (True, False) + ) + for enabled in order: + mode = "scheduled" if enabled else "baseline" + print( + f"running iteration={iteration + 1} phase={phase} mode={mode}", + file=sys.stderr, + flush=True, + ) + samples.append( + run_sample(args, root, phase, enabled, iteration + 1) + ) + + summary = summarize(samples, args.phases) + normalized_checks = [ + values["paired_capacity_normalized"]["within_five_percent"] + for values in summary.values() + if "paired_capacity_normalized" in values + ] + saturated_checks = [ + sample["saturation_verified"] + for sample in samples + if sample["name"] == "saturated" + and sample["mode"] == "scheduled" + and "saturation_verified" in sample + ] + regression_verified = all(normalized_checks) if normalized_checks else None + saturation_verified = all(saturated_checks) if saturated_checks else None + applicable_checks = [ + check + for check in (regression_verified, saturation_verified) + if check is not None + ] + report = { + "configuration": { + "binary": str(args.binary), + "iterations": args.iterations, + "duration_s": args.duration, + "warmup_s": args.warmup, + "runtime_size": args.runtime_size, + "max_concurrent_polls": args.max_concurrent_polls, + "query_workers": args.query_workers, + "write_workers": args.write_workers, + "seed_rows": args.seed_rows, + "write_batch_size": args.write_batch_size, + }, + "verification": { + "capacity_normalized_regression_budget_percent": ( + MAX_CAPACITY_NORMALIZED_REGRESSION_PERCENT + ), + "capacity_normalized_regression_verified": regression_verified, + "saturated_write_poll_share_at_least_80_percent": saturation_verified, + "passed": all(applicable_checks) if applicable_checks else None, + }, + "summary": summary, + "samples": samples, + } + rendered = json.dumps(report, indent=2, sort_keys=True) + if args.output: + args.output.write_text(rendered + "\n") + print(rendered) + + +if __name__ == "__main__": + main() diff --git a/tests/perf/workload_scheduler_runner.py b/tests/perf/workload_scheduler_runner.py new file mode 100644 index 0000000000..e49ff00aa5 --- /dev/null +++ b/tests/perf/workload_scheduler_runner.py @@ -0,0 +1,488 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Exercise GreptimeDB's experimental query/write workload scheduler. + +Start a standalone server, then run this script against its HTTP port. The +script creates and seeds a real Mito table and executes concurrent request +phases. By default it also verifies the experimental workload scheduler's +Prometheus poll-admission counters; metrics can be disabled when benchmarking +against a scheduler-disabled baseline. +""" + +from __future__ import annotations + +import argparse +import concurrent.futures +import dataclasses +import json +import math +import re +import statistics +import threading +import time +import urllib.error +import urllib.parse +import urllib.request +from typing import Any + + +QUERY_TABLE = "catio_scheduler_query_load" +WRITE_TABLE = "catio_scheduler_write_load" +SHARDS = 64 +QUERY_PARTITIONS = 32 +WRITE_PARTITIONS = 64 +METRIC_RE = re.compile( + r'^greptime_workload_scheduler_polls\{workload="(query|write)"\}\s+([0-9.eE+-]+)$' +) + + +@dataclasses.dataclass +class RequestStats: + requests: int = 0 + failures: int = 0 + latencies_ms: list[float] = dataclasses.field(default_factory=list) + failure_samples: list[str] = dataclasses.field(default_factory=list) + + def merge(self, other: RequestStats) -> None: + self.requests += other.requests + self.failures += other.failures + self.latencies_ms.extend(other.latencies_ms) + self.failure_samples.extend(other.failure_samples[: 3 - len(self.failure_samples)]) + + def summary(self, duration: float) -> dict[str, Any]: + successful = self.requests - self.failures + latencies = sorted(self.latencies_ms) + return { + "requests": self.requests, + "failures": self.failures, + "failure_samples": self.failure_samples, + "successful_rps": successful / duration, + "mean_ms": statistics.fmean(latencies) if latencies else None, + "p50_ms": percentile(latencies, 0.50), + "p95_ms": percentile(latencies, 0.95), + } + + +@dataclasses.dataclass +class PhaseClock: + warmup: float + duration: float + measurement_start: float = 0.0 + deadline: float = 0.0 + + def start(self) -> None: + started = time.monotonic() + self.measurement_start = started + self.warmup + self.deadline = self.measurement_start + self.duration + + +def percentile(values: list[float], quantile: float) -> float | None: + if not values: + return None + index = min(math.ceil(len(values) * quantile) - 1, len(values) - 1) + return values[max(index, 0)] + + +class SqlClient: + def __init__(self, base_url: str, database: str, timeout: float) -> None: + self.base_url = base_url.rstrip("/") + self.database = database + self.timeout = timeout + # Validation targets a local standalone process; inherited development + # proxies can otherwise turn overload into unrelated HTTP 502 errors. + self.opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + + def sql(self, sql: str) -> tuple[bool, float, Any]: + data = urllib.parse.urlencode( + {"sql": sql, "db": self.database, "format": "json"} + ).encode() + request = urllib.request.Request( + f"{self.base_url}/v1/sql", data=data, method="POST" + ) + started = time.monotonic() + try: + with self.opener.open(request, timeout=self.timeout) as response: + body = json.loads(response.read().decode()) + ok = response.status < 400 and not response_has_error(body) + return ok, (time.monotonic() - started) * 1_000, body + except urllib.error.HTTPError as error: + body = error.read().decode(errors="replace") + return ( + False, + (time.monotonic() - started) * 1_000, + f"HTTP {error.code}: {body}", + ) + except (OSError, ValueError, urllib.error.URLError) as error: + return False, (time.monotonic() - started) * 1_000, str(error) + + def scheduler_polls(self, required: bool = True) -> dict[str, int] | None: + with self.opener.open(f"{self.base_url}/metrics", timeout=self.timeout) as response: + text = response.read().decode() + result: dict[str, int] = {} + for line in text.splitlines(): + match = METRIC_RE.match(line) + if match: + result[match.group(1)] = int(float(match.group(2))) + if result.keys() != {"query", "write"}: + if required: + raise RuntimeError( + "scheduler metrics are missing; is " + "runtime.experimental_workload_scheduler.enable=true?" + ) + return None + return result + + +def response_has_error(body: Any) -> bool: + if not isinstance(body, dict): + return False + if body.get("error") or body.get("err_msg") or body.get("error_msg"): + return True + code = str(body.get("code", "")).lower() + return "output" not in body and code not in ("", "0", "success") + + +def setup_table(client: SqlClient, seed_rows: int, batch_size: int) -> None: + def create_table(table: str, partitions: int) -> str: + partition_width = SHARDS // partitions + partition_predicates = [f"shard < {partition_width}"] + partition_predicates.extend( + f"shard >= {lower} AND shard < {lower + partition_width}" + for lower in range( + partition_width, SHARDS - partition_width, partition_width + ) + ) + partition_predicates.append(f"shard >= {SHARDS - partition_width}") + return ( + f"CREATE TABLE {table} (" + "host STRING, shard INT, val DOUBLE, ts TIMESTAMP TIME INDEX, " + "PRIMARY KEY(host, shard)) " + f"PARTITION ON COLUMNS(shard) ({','.join(partition_predicates)}) " + "ENGINE=mito" + ) + + statements = [ + f"DROP TABLE IF EXISTS {QUERY_TABLE}", + f"DROP TABLE IF EXISTS {WRITE_TABLE}", + create_table(QUERY_TABLE, QUERY_PARTITIONS), + create_table(WRITE_TABLE, WRITE_PARTITIONS), + ] + for statement in statements: + ok, _, body = client.sql(statement) + if not ok: + raise RuntimeError(f"setup failed for {statement!r}: {body}") + + timestamp = 1_700_000_000_000 + for offset in range(0, seed_rows, batch_size): + count = min(batch_size, seed_rows - offset) + values = ",".join( + f"('host-{(offset + row) % 64}',{(offset + row) % 64}," + f"{offset + row},{timestamp + offset + row})" + for row in range(count) + ) + ok, _, body = client.sql( + f"INSERT INTO {QUERY_TABLE} (host,shard,val,ts) VALUES {values}" + ) + if not ok: + raise RuntimeError(f"seed insert at row {offset} failed: {body}") + + +def record_request( + stats: RequestStats, + started: float, + completed: float, + clock: PhaseClock, + ok: bool, + latency: float, + body: Any, +) -> None: + if started < clock.measurement_start or completed > clock.deadline: + return + stats.requests += 1 + stats.failures += not ok + stats.latencies_ms.append(latency) + if not ok and len(stats.failure_samples) < 3: + stats.failure_samples.append(str(body)[:500]) + + +def query_worker( + client: SqlClient, clock: PhaseClock, start: threading.Barrier +) -> RequestStats: + stats = RequestStats() + query = ( + f"SELECT host, count(*), sum(val), avg(val) FROM {QUERY_TABLE} " + "GROUP BY host ORDER BY host" + ) + start.wait() + while time.monotonic() < clock.deadline: + started = time.monotonic() + ok, latency, body = client.sql(query) + record_request( + stats, started, time.monotonic(), clock, ok, latency, body + ) + return stats + + +def write_worker( + client: SqlClient, + clock: PhaseClock, + start: threading.Barrier, + sequence: "Sequence", + batch_size: int, + delay: float, +) -> RequestStats: + stats = RequestStats() + start.wait() + while time.monotonic() < clock.deadline: + started = time.monotonic() + offset = sequence.take(batch_size) + values = ",".join( + f"('writer-{(offset + row) % 64}',{(offset + row) % 64}," + f"{offset + row},{offset + row})" + for row in range(batch_size) + ) + ok, latency, body = client.sql( + f"INSERT INTO {WRITE_TABLE} (host,shard,val,ts) VALUES {values}" + ) + record_request( + stats, started, time.monotonic(), clock, ok, latency, body + ) + if delay: + time.sleep(delay) + return stats + + +class Sequence: + def __init__(self, initial: int) -> None: + self.value = initial + self.lock = threading.Lock() + + def take(self, count: int) -> int: + with self.lock: + value = self.value + self.value += count + return value + + +def run_phase( + client: SqlClient, + name: str, + duration: float, + warmup: float, + query_workers: int, + write_workers: int, + write_batch_size: int, + write_delay: float, + sequence: Sequence, + scheduler_metrics: str, +) -> dict[str, Any]: + worker_count = query_workers + write_workers + clock = PhaseClock(warmup, duration) + start = threading.Barrier(worker_count + 1, action=clock.start) + query_stats = RequestStats() + write_stats = RequestStats() + + with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as executor: + query_futures = [ + executor.submit(query_worker, client, clock, start) + for _ in range(query_workers) + ] + write_futures = [ + executor.submit( + write_worker, + client, + clock, + start, + sequence, + write_batch_size, + write_delay, + ) + for _ in range(write_workers) + ] + start.wait() + remaining = clock.measurement_start - time.monotonic() + if remaining > 0: + time.sleep(remaining) + before = ( + client.scheduler_polls(required=scheduler_metrics == "required") + if scheduler_metrics != "disabled" + else None + ) + remaining = clock.deadline - time.monotonic() + if remaining > 0: + time.sleep(remaining) + # Scrape at the phase boundary, while the last requests are still + # backlogged. Waiting for slow queries to drain would incorrectly count + # their post-load polls as part of the saturated interval. + at_deadline = ( + client.scheduler_polls(required=True) if before is not None else None + ) + for future in query_futures: + query_stats.merge(future.result()) + for future in write_futures: + write_stats.merge(future.result()) + + poll_delta = ( + { + workload: at_deadline[workload] - before[workload] + for workload in before + } + if before is not None and at_deadline is not None + else None + ) + total_polls = sum(poll_delta.values()) if poll_delta is not None else 0 + shares = ( + { + workload: (polls / total_polls if total_polls else 0.0) + for workload, polls in poll_delta.items() + } + if poll_delta is not None + else None + ) + return { + "name": name, + "duration_s": duration, + "warmup_s": warmup, + "workers": {"query": query_workers, "write": write_workers}, + "requests": { + "query": query_stats.summary(duration), + "write": write_stats.summary(duration), + }, + "polls": poll_delta, + "poll_share": shares, + } + + +def verify(phases: list[dict[str, Any]]) -> None: + by_name = {phase["name"]: phase for phase in phases} + query_only = by_name["query_only"] + light_write = by_name["light_write"] + saturated = by_name["saturated"] + + if any(phase["polls"] is None for phase in phases): + raise AssertionError("scheduler metrics are required for verification") + if query_only["polls"]["query"] <= 0 or query_only["polls"]["write"] != 0: + raise AssertionError(f"query-only phase did not borrow all capacity: {query_only}") + if light_write["poll_share"]["query"] <= 0.20: + raise AssertionError(f"query did not borrow unused write share: {light_write}") + if saturated["polls"]["query"] < 100 or saturated["polls"]["write"] < 100: + raise AssertionError(f"saturated phase did not generate enough work: {saturated}") + if saturated["poll_share"]["write"] < 0.799: + raise AssertionError(f"write admission share is below 80% (0.1% tolerance): {saturated}") + + for phase in phases: + for workload in ("query", "write"): + request = phase["requests"][workload] + if request["requests"] and request["failures"] / request["requests"] >= 0.01: + raise AssertionError( + f"{phase['name']} {workload} failure rate was at least 1%: {request}" + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--url", default="http://127.0.0.1:4000") + parser.add_argument("--database", default="public") + parser.add_argument("--duration", type=float, default=10.0) + parser.add_argument("--warmup", type=float, default=2.0) + parser.add_argument("--query-workers", type=int, default=2) + parser.add_argument("--write-workers", type=int, default=1152) + parser.add_argument("--seed-rows", type=int, default=10_000) + parser.add_argument("--seed-batch-size", type=int, default=500) + parser.add_argument("--write-batch-size", type=int, default=32) + parser.add_argument("--light-write-delay", type=float, default=0.1) + parser.add_argument("--timeout", type=float, default=60.0) + parser.add_argument( + "--phase", + choices=("all", "query_only", "write_only", "light_write", "saturated"), + default="all", + ) + parser.add_argument( + "--scheduler-metrics", + choices=("required", "optional", "disabled"), + default="required", + help="whether scheduler Prometheus metrics must be collected", + ) + parser.add_argument("--skip-setup", action="store_true") + parser.add_argument("--no-verify", action="store_true") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + client = SqlClient(args.url, args.database, args.timeout) + if not args.skip_setup: + setup_table(client, args.seed_rows, args.seed_batch_size) + + sequence = Sequence(1_800_000_000_000) + phase_options = { + "query_only": (args.query_workers, 0, 0), + "write_only": (0, args.write_workers, 0), + "light_write": (args.query_workers, 1, args.light_write_delay), + "saturated": (args.query_workers, args.write_workers, 0), + } + selected = ( + { + name: phase_options[name] + for name in ("query_only", "light_write", "saturated") + } + if args.phase == "all" + else {args.phase: phase_options[args.phase]} + ) + phases = [] + for name, (query_workers, write_workers, write_delay) in selected.items(): + phases.append( + run_phase( + client, + name, + args.duration, + args.warmup, + query_workers, + write_workers, + args.write_batch_size, + write_delay, + sequence, + args.scheduler_metrics, + ) + ) + if not args.no_verify: + if args.phase != "all": + raise ValueError("--phase requires --no-verify unless all phases are selected") + if args.scheduler_metrics != "required": + raise ValueError("verification requires --scheduler-metrics=required") + verify(phases) + + result = { + "verified": not args.no_verify, + "mean_write_share_saturated": statistics.fmean( + [ + phase["poll_share"]["write"] + for phase in phases + if phase["name"] == "saturated" and phase["poll_share"] is not None + ] + ) + if any( + phase["name"] == "saturated" and phase["poll_share"] is not None + for phase in phases + ) + else None, + "phases": phases, + } + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main()