diff --git a/.github/actions/build-linux-artifacts/action.yml b/.github/actions/build-linux-artifacts/action.yml index c5212c28998..26028485061 100644 --- a/.github/actions/build-linux-artifacts/action.yml +++ b/.github/actions/build-linux-artifacts/action.yml @@ -57,6 +57,7 @@ runs: # Without GT_S3_*, every remote object-store case is skipped and the # release only ever exercises the local file backend. env: + RUST_MIN_STACK: 8388608 # 8 MiB, aligned with Rust CI tests. GT_S3_BUCKET: ${{ inputs.aws-ci-test-bucket }} GT_S3_REGION: ${{ inputs.aws-region }} GT_S3_ACCESS_KEY_ID: ${{ inputs.aws-access-key-id }} diff --git a/.github/actions/build-macos-artifacts/action.yml b/.github/actions/build-macos-artifacts/action.yml index 2caefe1c382..2affb5b94f7 100644 --- a/.github/actions/build-macos-artifacts/action.yml +++ b/.github/actions/build-macos-artifacts/action.yml @@ -64,6 +64,7 @@ runs: shell: bash env: CARGO_BUILD_RUSTFLAGS: "-Clink-arg=-Wl,-ld_classic" + RUST_MIN_STACK: 8388608 # 8 MiB, aligned with Rust CI and release tests. SQLNESS_OPTS: "--preserve-state" run: | make test sqlness-test diff --git a/.github/actions/build-windows-artifacts/action.yml b/.github/actions/build-windows-artifacts/action.yml index 2c4ae1f7db7..d25dd27e117 100644 --- a/.github/actions/build-windows-artifacts/action.yml +++ b/.github/actions/build-windows-artifacts/action.yml @@ -22,7 +22,9 @@ inputs: runs: using: composite steps: - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 + with: + tool: protoc@3.23.4 - name: Install rust toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 @@ -48,6 +50,7 @@ runs: run: make test sqlness-test env: RUST_BACKTRACE: 1 + RUST_MIN_STACK: 8388608 # 8 MiB, aligned with Rust CI tests. SQLNESS_OPTS: "--preserve-state" - name: Upload sqlness logs diff --git a/.github/actions/setup-greptimedb-cluster/with-observability.yaml b/.github/actions/setup-greptimedb-cluster/with-observability.yaml index 581f60dea60..6cc923aede3 100644 --- a/.github/actions/setup-greptimedb-cluster/with-observability.yaml +++ b/.github/actions/setup-greptimedb-cluster/with-observability.yaml @@ -1,3 +1,7 @@ +datanode: + storage: + storageClassName: standard + tracing: enabled: true endpoint: "http://my-greptimedb-monitor-standalone.my-greptimedb:4000/v1/otlp/v1/traces" @@ -15,9 +19,9 @@ monitoring: limits: cpu: 500m memory: 256Mi - datanodeStorage: - fs: - storageSize: 5Gi + datanodeStorage: + fs: + storageSize: 5Gi vector: resources: requests: diff --git a/.github/scripts/metric-snapshot-fixture.py b/.github/scripts/metric-snapshot-fixture.py new file mode 100644 index 00000000000..4341bc6f579 --- /dev/null +++ b/.github/scripts/metric-snapshot-fixture.py @@ -0,0 +1,294 @@ +#!/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. + +"""Deterministic wide, heterogeneous Metric tables for snapshot benchmarks.""" + +import argparse +import base64 +import concurrent.futures +import hashlib +import http.client +import json +import os +from pathlib import Path +import threading +import time +import urllib.parse + +GROUPS = 10 +TAG_POOL = 298 +ROWS = 32 +DATABASE = "metrics" + + +class Client: + def __init__(self, addr): + self.addr = addr + self.connections = threading.local() + self.headers = {"Content-Type": "application/x-www-form-urlencoded"} + if auth := os.environ.get("GREPTIME_BENCH_AUTH"): + self.headers["Authorization"] = "Basic " + base64.b64encode( + auth.encode() + ).decode() + + def sql(self, query, database=DATABASE, batch=False): + connection = getattr(self.connections, "connection", None) + if connection is None: + connection = http.client.HTTPConnection(self.addr, timeout=600) + self.connections.connection = connection + body = urllib.parse.urlencode({"db": database, "sql": query}) + try: + connection.request( + "POST", "/v1/ddl/logical-tables" if batch else "/v1/sql", + body, self.headers, + ) + response = connection.getresponse() + raw = response.read() + except (OSError, http.client.HTTPException): + connection.close() + self.connections.connection = None + raise + data = json.loads(raw) + if response.status != 200 or "error" in data: + raise RuntimeError(f"SQL failed (HTTP {response.status}): {data}") + return data["output"] + + def values(self, query, database=DATABASE): + return self.sql(query, database)[0]["records"]["rows"] + + +def tags(index, per_group): + return sorted((index % per_group * 7 + offset) % TAG_POOL for offset in range(10)) + + +def column_names(tag_ids): + return ["ts", *[f"tag_{tag:03d}" for tag in tag_ids], "val"] + + +def columns(tag_ids): + names = [f"tag_{tag:03d}" for tag in tag_ids] + return ( + "ts TIMESTAMP TIME INDEX, " + + ", ".join(name + " STRING" for name in names) + + ", val DOUBLE, PRIMARY KEY(" + ",".join(names) + ")" + ) + + +def table_ddl(index, per_group): + return ( + f"CREATE TABLE m{index:06d} ({columns(tags(index, per_group))}) " + f"ENGINE=metric WITH(on_physical_table='phy_{index // per_group:02d}')" + ) + + +def row(index, row_index, per_group): + values = [] + for tag in tags(index, per_group): + if (row_index + tag) % 29 == 0: + value = None + elif (row_index + tag) % 17 == 0: + value = "" + elif tag % 10 == 0: + value = f"instance-{index:06d}-{row_index % 4}" + elif tag % 10 == 1: + value = "区域-" + str(row_index % 3) + elif tag % 10 == 2: + value = "label-" + "x" * 48 + "-" + str(row_index % 2) + else: + value = f"tag{tag:03d}-v{row_index % (2 + tag % 7)}" + values.append(value) + return [row_index, *values, None if row_index % 7 == 0 else index % 97 + row_index / 4] + + +def sql_value(value): + if value is None: + return "NULL" + if isinstance(value, str): + return "'" + value.replace("'", "''") + "'" + return str(value) + + +def canonical(value): + return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + + +def update_digest(digest, value): + digest.update((canonical(value) + "\n").encode()) + + +def progress(phase, completed, total): + if completed == total or completed % 4096 == 0: + print(f"{phase}: {completed}/{total}", flush=True) + + +def prepare(client, per_group): + stages = {} + started = time.monotonic() + client.sql(f"CREATE DATABASE {DATABASE}", "public") + for group in range(GROUPS): + client.sql( + f"CREATE TABLE phy_{group:02d} ({columns(range(TAG_POOL))}) " + "ENGINE=metric WITH(physical_metric_table='')" + ) + total = GROUPS * per_group + for group in range(GROUPS): + for base in range(group * per_group, (group + 1) * per_group, 128): + end = min((group + 1) * per_group, base + 128) + client.sql(";".join(table_ddl(i, per_group) for i in range(base, end)), batch=True) + progress("create", end, total) + stages["create_seconds"] = time.monotonic() - started + + def insert(index): + rows = ",".join( + "(" + ",".join(sql_value(v) for v in row(index, j, per_group)) + ")" + for j in range(ROWS) + ) + names = ",".join(column_names(tags(index, per_group))) + client.sql(f"INSERT INTO m{index:06d} ({names}) VALUES {rows}") + + started = time.monotonic() + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: + for base in range(0, total, 128): + end = min(total, base + 128) + list(pool.map(insert, range(base, end))) + progress("insert", end, total) + stages["insert_seconds"] = time.monotonic() - started + started = time.monotonic() + for group in range(GROUPS): + client.sql(f"ADMIN FLUSH_TABLE('{DATABASE}.phy_{group:02d}')") + stages["flush_seconds"] = time.monotonic() - started + return stages + + +def check_records(index, per_group, records): + names = column_names(tags(index, per_group)) + types = ["TimestampMillisecond", *["String"] * 10, "Float64"] + expected_columns = [dict(name=name, data_type=kind) for name, kind in zip(names, types)] + actual_columns = records["schema"]["column_schemas"] + if sorted(actual_columns, key=lambda c: c["name"]) != sorted( + expected_columns, key=lambda c: c["name"] + ): + raise RuntimeError(f"logical schema mismatch: m{index:06d}") + # Metric tables canonicalize column order; SELECT * still checks every column. + positions = [names.index(column["name"]) for column in actual_columns] + expected_rows = [ + [values[position] for position in positions] + for values in (row(index, j, per_group) for j in range(ROWS)) + ] + # Canonical JSON also distinguishes integers, floats, booleans, and NULL. + if canonical(records["rows"]) != canonical(expected_rows): + raise RuntimeError(f"typed values/NULL mismatch: m{index:06d}") + + +def validate(client, per_group): + total = GROUPS * per_group + expected_names = {f"m{i:06d}" for i in range(total)} | { + f"phy_{group:02d}" for group in range(GROUPS) + } + table_rows = client.values( + "SELECT table_name, table_id FROM information_schema.tables " + f"WHERE table_schema='{DATABASE}' ORDER BY table_name" + ) + table_ids = dict(table_rows) + if len(table_rows) != len(expected_names) or set(table_ids) != expected_names: + raise RuntimeError("table membership mismatch") + if len(set(table_ids.values())) != len(table_ids): + raise RuntimeError("duplicate table IDs") + digests = {name: hashlib.sha256() for name in ("schema", "data", "ddl")} + + def check_ddl(names): + outputs = client.sql(";".join(f"SHOW CREATE TABLE {name}" for name in names)) + if len(outputs) != len(names): + raise RuntimeError("incomplete SHOW CREATE TABLE response") + for name, output in zip(names, outputs): + records = output["records"]["rows"] + if len(records) != 1 or records[0][0] != name or not records[0][1]: + raise RuntimeError(f"invalid SHOW CREATE TABLE response: {name}") + update_digest(digests["ddl"], records) + + for base in range(0, total, 64): + end = min(total, base + 64) + names = [f"m{i:06d}" for i in range(base, end)] + outputs = client.sql(";".join(f"SELECT * FROM {name} ORDER BY ts" for name in names)) + if len(outputs) != len(names): + raise RuntimeError("incomplete SELECT response") + for index, output in enumerate(outputs, base): + records = output["records"] + check_records(index, per_group, records) + update_digest(digests["schema"], [f"m{index:06d}", records["schema"]]) + update_digest(digests["data"], [f"m{index:06d}", records["rows"]]) + check_ddl(names) + progress("validate", end, total) + check_ddl([f"phy_{group:02d}" for group in range(GROUPS)]) + return {**{name + "_sha256": digest.hexdigest() for name, digest in digests.items()}, + "table_ids": table_ids} + + +def compare_source(result, source): + for field in ("fixture", "schema_sha256", "data_sha256", "ddl_sha256"): + if result[field] != source[field]: + raise RuntimeError(f"source/target mismatch: {field}") + source_ids, target_ids = source["table_ids"], result["table_ids"] + if source_ids.keys() != target_ids.keys(): + raise RuntimeError("source/target table membership mismatch") + changed = sum(source_ids[name] != target_ids[name] for name in source_ids) + if changed != len(source_ids): + raise RuntimeError(f"table IDs not all remapped: {changed}/{len(source_ids)}") + result["remapped_table_ids"] = changed + + +def placeholders(client): + database = "metric_bench_guard" + client.sql(f"CREATE DATABASE {database}", "public") + for index in range(17): + client.sql(f"CREATE TABLE placeholder_{index} (ts TIMESTAMP TIME INDEX)", database) + return {"placeholder_tables": 17} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("action", choices=("prepare", "validate", "placeholders")) + parser.add_argument("--addr", default="127.0.0.1:4000") + parser.add_argument("--tables-per-group", type=int, default=10000) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--source", type=Path, help="Source validation JSON, for target comparison") + args = parser.parse_args() + if not 50 <= args.tables_per_group <= 10000: + parser.error("--tables-per-group must be between 50 and 10000") + if args.source is not None and args.action != "validate": + parser.error("--source only applies to validate") + coverage = {tag for i in range(args.tables_per_group) for tag in tags(i, args.tables_per_group)} + if coverage != set(range(TAG_POOL)): + parser.error("fixture does not cover all 298 tags") + client = Client(args.addr) + started = time.monotonic() + result = {"fixture": {"groups": GROUPS, "tables_per_group": args.tables_per_group, + "rows_per_table": ROWS, "tag_pool": TAG_POOL}} + if args.action == "placeholders": + result.update(placeholders(client)) + elif args.action == "prepare": + result.update(prepare(client, args.tables_per_group)) + else: + result.update(validate(client, args.tables_per_group)) + result["wall_seconds"] = time.monotonic() - started + if args.source is not None: + compare_source(result, json.loads(args.source.read_text())) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n") + print(json.dumps({key: value for key, value in result.items() if key != "table_ids"})) + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/metric-snapshot-s3-bench.py b/.github/scripts/metric-snapshot-s3-bench.py new file mode 100644 index 00000000000..7bb01a410b1 --- /dev/null +++ b/.github/scripts/metric-snapshot-s3-bench.py @@ -0,0 +1,392 @@ +#!/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. + +"""Run one packed Metric snapshot smoke, then one full S3 round trip.""" + +import argparse +from datetime import datetime +import json +import os +from pathlib import Path +import re +import signal +import subprocess +import sys +import threading +import time +import urllib.request + +HERE = Path(__file__).resolve().parent + + +def command(args, **kwargs): + result = subprocess.run(args, text=True, stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, **kwargs) + if result.returncode: + # CLI argv and error output can contain storage credentials. + raise RuntimeError(f"{args[0]} exited with status {result.returncode}") + return result.stdout + + +def write_json(path, value): + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n") + + +def metric_values(text): + values = {} + for line in text.splitlines(): + if line and not line.startswith("#"): + key, value = line.rsplit(" ", 1) + values[key] = float(value) + return values + + +def docker_timestamp(value): + # Docker emits nanoseconds; Python 3.9 accepts at most microseconds. + return datetime.fromisoformat(re.sub(r"(\.\d{6})\d+", r"\1", value).replace("Z", "+00:00")) + + +def host_memory(): + if not Path("/proc/vmstat").exists(): + return {"available": False} + vm = dict(line.split() for line in Path("/proc/vmstat").read_text().splitlines()) + mem = dict(line.split(":", 1) for line in Path("/proc/meminfo").read_text().splitlines()) + return {"available": True, "pswpin_pages": int(vm["pswpin"]), + "pswpout_pages": int(vm["pswpout"]), "oom_kill": int(vm.get("oom_kill", 0)), + **{key + "_kib": int(mem[key].split()[0]) + for key in ("MemTotal", "MemAvailable", "SwapTotal", "SwapFree")}} + + +def process_usage(pid): + try: + stat = Path(f"/proc/{pid}/stat").read_text().rsplit(")", 1)[1].split() + return {"cpu_seconds": (int(stat[11]) + int(stat[12])) / os.sysconf("SC_CLK_TCK"), + "rss_bytes": int(stat[21]) * os.sysconf("SC_PAGE_SIZE")} + except (OSError, IndexError): + return None + + +class Benchmark: + def __init__(self, local_smoke): + self.local_smoke = local_smoke + self.image = os.environ["GREPTIMEDB_IMAGE"] + self.commit = os.environ["EXPECTED_COMMIT"] + if not re.fullmatch(r"greptime/greptimedb(?:-nightly)?@sha256:[0-9a-f]{64}", self.image): + raise ValueError("GREPTIMEDB_IMAGE must be an immutable official image digest") + if not re.fullmatch(r"[0-9a-f]{40}", self.commit): + raise ValueError("EXPECTED_COMMIT must be a full commit SHA") + self.per_group = int(os.environ.get("LOGICAL_TABLES_PER_GROUP", "10000")) + if not 50 <= self.per_group <= 10000: + raise ValueError("LOGICAL_TABLES_PER_GROUP must be in 50..10000") + self.prefix = os.environ["RUN_PREFIX"] + if not re.fullmatch(r"metric-snapshot-bench/[0-9]+-[0-9]+/", self.prefix): + raise ValueError("RUN_PREFIX must be metric-snapshot-bench/-/") + self.results = Path(os.environ["RESULTS"]).resolve() + self.data = Path(os.environ["DATA_ROOT"]).resolve() + self.results.mkdir(parents=True, exist_ok=False) + self.data.mkdir(parents=True, exist_ok=False) + self.network = "metric-snapshot-" + self.prefix.split("/")[1] + self.server = self.network + "-db" + self.cli_name = self.network + "-cli" + self.port = int(os.environ.get("BENCH_HTTP_PORT", "19410")) + self.base_url = f"http://127.0.0.1:{self.port}" + self.endpoint = os.environ.get("S3_ENDPOINT", "") + if self.endpoint and not local_smoke: + raise ValueError("Custom S3 endpoints are restricted to local smoke") + if not local_smoke and os.environ["S3_REGION"] != "ap-southeast-1": + raise ValueError("Cloud benchmark requires ap-southeast-1") + self.containers = set() + self.network_created = False + self.phases = [] + self.server_pid = 0 + self.round_dir = None + + def docker(self, *args): + return command(["docker", *args]) + + def metrics(self): + try: + with urllib.request.urlopen(self.base_url + "/metrics", timeout=5) as response: + return response.read().decode() + except OSError: + return "" + + def phase(self, name, action): + print(f"Starting {self.round_dir.name}/{name}", flush=True) + before = host_memory() + metrics_before = self.metrics() + samples = [] + done = threading.Event() + + def sample(): + while not done.is_set(): + samples.append({"time": time.time(), "server": process_usage(self.server_pid), + "cli": process_usage(getattr(self, "cli_pid", 0)), + "query_memory": {k: v for k, v in metric_values(self.metrics()).items() + if "query_memory" in k} if self.server_pid else {}}) + done.wait(0.5) + + monitor = threading.Thread(target=sample, daemon=True) + monitor.start() + start = time.monotonic() + outcome = "failed" + result = None + try: + result = action() + outcome = "passed" + return result + finally: + elapsed = time.monotonic() - start + done.set() + monitor.join() + after = host_memory() + interfered = before["available"] and after["available"] and ( + after["pswpout_pages"] > before["pswpout_pages"] or after["oom_kill"] > before["oom_kill"]) + if interfered: + outcome = "resource-invalid" + record = {"round": self.round_dir.name, "phase": name, "outcome": outcome, + "wall_seconds": elapsed, "host_before": before, "host_after": after, + "process_samples": samples} + if isinstance(result, dict): + record.update(result) + if before["available"] and after["available"]: + record["new_swapout_pages"] = after["pswpout_pages"] - before["pswpout_pages"] + record["new_oom_kills"] = after["oom_kill"] - before["oom_kill"] + metrics_after = self.metrics() + pre, post = metric_values(metrics_before), metric_values(metrics_after) + record["server_counter_deltas"] = { + key: value - pre.get(key, 0) for key, value in post.items() + if ("opendal" in key and ("_count{" in key or "_sum{" in key)) + or key.startswith("greptime_servers_http_requests_total{")} + write_json(self.round_dir / f"{name}-resources.json", record) + (self.round_dir / f"{name}-before.prom").write_text(metrics_before) + (self.round_dir / f"{name}-after.prom").write_text(metrics_after) + self.phases.append({k: v for k, v in record.items() if k != "process_samples"}) + self.summary() + print(f"Finished {name}: {outcome}, {elapsed:.3f}s", flush=True) + if interfered: + raise RuntimeError(f"Resource interference during {name}; evidence retained") + + def fixture(self, operation, per_group, output, source=None): + args = [sys.executable, str(HERE / "metric-snapshot-fixture.py"), operation, + "--addr", f"127.0.0.1:{self.port}", "--tables-per-group", str(per_group), + "--output", str(output)] + if source: + args += ["--source", str(source)] + with (self.round_dir / f"{operation}-{output.stem}.log").open("w") as log: + status = subprocess.run(args, stdout=log, stderr=subprocess.STDOUT).returncode + if status: + raise RuntimeError(f"Fixture {operation} failed; see its artifact log") + + def start_server(self, role): + root = self.data / self.round_dir.name / role + root.mkdir(parents=True) + evidence = self.round_dir / role + evidence.mkdir() + config = root / "config.toml" + config.write_text('''enable_telemetry = false +[http] +addr = "0.0.0.0:4000" +[grpc] +bind_addr = "0.0.0.0:4001" +[mysql] +enable = false +[postgres] +enable = false +[storage] +type = "File" +data_home = "/data/db" +[query] +parallelism = 4 +memory_pool_size = "2GiB" +[logging] +dir = "/evidence/logs" +level = "info" +''') + self.containers.add(self.server) + self.docker("run", "-d", "--name", self.server, "--network", self.network, + "--network-alias", "metric-db", "--user", f"{os.getuid()}:{os.getgid()}", + "--add-host", "host.docker.internal:host-gateway", + "-p", f"127.0.0.1:{self.port}:4000", + "-v", f"{root}:/data", "-v", f"{evidence}:/evidence", + self.image, "standalone", "start", "-c", "/data/config.toml") + self.server_pid = int(self.docker("inspect", "-f", "{{.State.Pid}}", self.server)) + deadline = time.monotonic() + 120 + while time.monotonic() < deadline: + try: + with urllib.request.urlopen(self.base_url + "/v1/capabilities", timeout=2) as response: + caps = json.load(response) + if any(caps.get(key) != 1 for key in + ("metric_packed_export", "metric_packed_import", "metric_batch_ddl")): + raise RuntimeError("Required packed snapshot capabilities are absent") + write_json(evidence / "capabilities.json", caps) + return + except OSError: + if self.docker("inspect", "-f", "{{.State.Running}}", self.server).strip() != "true": + break + time.sleep(1) + raise RuntimeError("Server did not become ready") + + def stop_server(self, role): + self.docker("stop", "-t", "60", self.server) + state = json.loads(self.docker("inspect", "-f", "{{json .State}}", self.server)) + write_json(self.round_dir / role / "exit.json", state) + (self.round_dir / role / "stdout.log").write_text(self.docker("logs", self.server)) + self.docker("rm", self.server) + self.containers.remove(self.server) + self.server_pid = 0 + if state["OOMKilled"] or state["ExitCode"] != 0: + raise RuntimeError(f"{role} server did not shut down cleanly") + + def cli(self, args, name): + # Expand credentials only inside the container. Never persist argv or env. + script = 'exec greptime "$@" --s3 --s3-region "$S3_REGION" --s3-access-key-id "$AWS_ACCESS_KEY_ID" --s3-secret-access-key "$AWS_SECRET_ACCESS_KEY"' + if self.endpoint: + args += ["--s3-endpoint", self.endpoint] + self.containers.add(self.cli_name) + self.docker("create", "--name", self.cli_name, "--network", self.network, + "--user", f"{os.getuid()}:{os.getgid()}", + "--add-host", "host.docker.internal:host-gateway", + "-e", "AWS_ACCESS_KEY_ID", "-e", "AWS_SECRET_ACCESS_KEY", "-e", "S3_REGION", + "-v", f"{self.round_dir}:/evidence", "--entrypoint", "/bin/sh", self.image, + "-c", script, "sh", "--log-dir", f"/evidence/{name}-logs", "cli", "data", *args) + + def execute(): + self.docker("start", self.cli_name) + self.cli_pid = int(self.docker("inspect", "-f", "{{.State.Pid}}", self.cli_name)) + status = int(self.docker("wait", self.cli_name)) + state = json.loads(self.docker("inspect", "-f", "{{json .State}}", self.cli_name)) + (self.round_dir / f"{name}.log").write_text(self.docker("logs", self.cli_name)) + write_json(self.round_dir / f"{name}-exit.json", state) + if status or state["ExitCode"] or state["OOMKilled"]: + raise RuntimeError(f"{name} failed; see artifact log") + return {"cli_wall_seconds": (docker_timestamp(state["FinishedAt"]) - + docker_timestamp(state["StartedAt"])).total_seconds()} + + try: + self.phase(name, execute) + finally: + self.cli_pid = 0 + if not (self.round_dir / f"{name}-exit.json").exists(): + self.save_container(self.cli_name, name + "-failure") + self.docker("rm", "-f", self.cli_name) + self.containers.remove(self.cli_name) + + def round(self, label, per_group): + self.round_dir = self.results / label + self.round_dir.mkdir() + self.start_server("source") + source = self.round_dir / "source-validation.json" + self.phase("prepare", lambda: self.fixture("prepare", per_group, self.round_dir / "prepare.json")) + self.phase("source-validation", lambda: self.fixture("validate", per_group, source)) + location = f"s3://{os.environ['S3_BUCKET']}/{self.prefix}{label}" + self.cli(["export-v2", "create", "--addr", "metric-db:4000", "--to", location, + "--schemas", "metrics", "--pack-metric-data", "--parallelism", "4", + "--chunk-parallelism", "1", "--start-time", "1970-01-01T00:00:00Z", + "--end-time", "1970-01-01T00:00:01Z", "--timeout", "1h", + "--no-proxy", "--progress", "always"], "export") + if not self.local_smoke: + command(["uv", "run", str(HERE / "metric-snapshot-s3-ops.py"), "inventory", "--round", label, + "--output", str(self.round_dir / "s3-inventory.json")]) + self.stop_server("source") + self.cli(["export-v2", "verify", "--snapshot", location], "verify") + self.start_server("target") + self.fixture("placeholders", per_group, self.round_dir / "placeholders.json") + self.cli(["import-v2", "--addr", "metric-db:4000", "--from", location, + "--task-parallelism", "1", "--state-path", "/evidence/import-state.json", + "--timeout", "1h", "--no-proxy", "--progress", "always"], "import") + imported = json.loads((self.round_dir / "import-resources.json").read_text()) + batch_requests = sum(value for key, value in imported["server_counter_deltas"].items() + if '/v1/ddl/logical-tables' in key and 'code="200"' in key) + if batch_requests <= 0: + raise RuntimeError("Import did not complete any successful batch DDL HTTP requests") + write_json(self.round_dir / "import-completion.json", {"batch_ddl_requests": batch_requests, + "resume_state_removed": not (self.round_dir / "import-state.json").exists()}) + if (self.round_dir / "import-state.json").exists(): + raise RuntimeError("Successful import left its resume state file") + self.phase("target-validation", lambda: self.fixture( + "validate", per_group, self.round_dir / "target-validation.json", source)) + self.stop_server("target") + + def summary(self): + lines = ["# Metric packed snapshot benchmark", "", f"Image: `{self.image}`", + f"Source commit: `{self.commit}`", "", + "Local smoke only; timings are not AWS results." if self.local_smoke else + "AWS S3 in ap-southeast-1; source and target use local File storage.", "", + "| Round | Phase | Seconds | Result |", "| --- | --- | ---: | --- |"] + for phase in self.phases: + lines.append(f"| {phase['round']} | {phase['phase']} | {phase['wall_seconds']:.3f} | {phase['outcome']} |") + lines += ["", "Each logical table has 32 rows, 10 overlapping string tags, a timestamp and a DOUBLE.", + "Phase durations include Docker start, wait and log collection. cli_wall_seconds records Docker StartedAt to FinishedAt separately; setup, loading and full validation are separate.", + "Process CPU/RSS are sampled every 0.5 s on Linux, not exact process peaks. Linux host swap/OOM counters are checked per phase; configured swap is not RAM.", + "Raw server Prometheus snapshots and CLI progress/logs are retained. OpenDAL operation counts are not AWS HTTP request counts: multipart requests, retries and CLI operations are not fully covered. Exact AWS request totals are unavailable."] + (self.results / "summary.md").write_text("\n".join(lines) + "\n") + + def run(self): + self.docker("pull", self.image) + version = self.docker("run", "--rm", self.image, "--version") + (self.results / "version.txt").write_text(version) + if self.commit not in version or "clean: true" not in version: + raise RuntimeError("Image version does not match the expected clean source commit") + for args, required in ((["export-v2", "create"], "--pack-metric-data"), + (["import-v2"], "--task-parallelism")): + help_text = self.docker("run", "--rm", self.image, "cli", "data", *args, "--help") + (self.results / (args[0] + "-help.txt")).write_text(help_text) + if required not in help_text: + raise RuntimeError(f"Image CLI lacks {required}") + write_json(self.results / "host.json", {"memory": host_memory(), "local_smoke": self.local_smoke, + "platform": self.docker("image", "inspect", "-f", "{{.Os}}/{{.Architecture}}", self.image).strip()}) + self.docker("network", "create", self.network) + self.network_created = True + self.round("smoke", 50) + if not self.local_smoke and self.per_group != 50: + self.round("full", self.per_group) + + def save_container(self, name, label): + state = subprocess.run(["docker", "inspect", "-f", "{{json .State}}", name], + capture_output=True, text=True) + (self.round_dir / (label + "-state.json")).write_text(state.stdout) + result = subprocess.run(["docker", "logs", name], capture_output=True, text=True) + (self.round_dir / (label + ".log")).write_text(result.stdout + result.stderr) + + def cleanup(self): + for name in self.containers: + if self.round_dir: + self.save_container(name, name + "-failure") + subprocess.run(["docker", "rm", "-f", name], stdout=subprocess.DEVNULL, check=False) + if self.network_created: + subprocess.run(["docker", "network", "rm", self.network], stdout=subprocess.DEVNULL, check=False) + self.summary() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("command", choices=["run"]) + parser.add_argument("--local-smoke", action="store_true") + args = parser.parse_args() + bench = Benchmark(args.local_smoke) + + def cancelled(signum, frame): + raise RuntimeError(f"Received signal {signum}") + + signal.signal(signal.SIGTERM, cancelled) + try: + bench.run() + finally: + bench.cleanup() + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/metric-snapshot-s3-ops.py b/.github/scripts/metric-snapshot-s3-ops.py new file mode 100644 index 00000000000..f1e34defb89 --- /dev/null +++ b/.github/scripts/metric-snapshot-s3-ops.py @@ -0,0 +1,217 @@ +#!/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. + +# /// script +# requires-python = ">=3.10" +# dependencies = ["boto3>=1.35,<2"] +# /// + +"""Budget preflight and strictly scoped S3 operations for the metric snapshot benchmark.""" + +import argparse +from decimal import Decimal +import json +import os +from pathlib import Path +import re +import time +import urllib.request + +REGION = "ap-southeast-1" +HOURLY_USD = Decimal("0.9408") +OTHER_RESERVE_USD = Decimal("1.50") + + +def owned_prefix(prefix): + if not re.fullmatch(r"metric-snapshot-bench/[0-9]+-[0-9]+/", prefix): + raise ValueError("prefix must be exactly metric-snapshot-bench/-/") + return prefix + + +def plan(env): + image = env["GREPTIMEDB_IMAGE"] + if not re.fullmatch(r"greptime/greptimedb(?:-nightly)?@sha256:[0-9a-f]{64}", image): + raise ValueError("image must be an immutable greptime/greptimedb sha256 reference") + if not re.fullmatch(r"[0-9a-f]{40}", env["EXPECTED_COMMIT"]): + raise ValueError("expected commit must be a full lowercase source SHA") + tables, minutes = int(env["LOGICAL_TABLES_PER_GROUP"]), int(env["BENCHMARK_TIMEOUT_MINUTES"]) + if not 50 <= tables <= 10_000 or not 30 <= minutes <= 240: + raise ValueError("logical tables per group must be 50..10000; timeout must be 30..240 minutes") + if env["S3_REGION"] != REGION or env["RUNNER_REGION"] != REGION: + raise ValueError("bucket and runner must both use ap-southeast-1") + if not env["S3_BUCKET"]: + raise ValueError("S3 bucket is required") + prior = Decimal(env["PRIOR_SPEND_USD"]) + # Reserve 45m provision, 10m queue allowance and 20m cleanup beyond the job timeout. + reservation = HOURLY_USD * Decimal(minutes + 75) / 60 + OTHER_RESERVE_USD + if not prior.is_finite() or prior < 0 or prior >= 15 or prior + reservation > 20: + raise ValueError("campaign budget exhausted: prior spend must be below $15 and leave the full attempt reserve within $20") + return {"image": image, "expected_commit": env["EXPECTED_COMMIT"], + "logical_tables_per_group": tables, "smoke_logical_tables": 500, + "timeout_minutes": minutes, "prefix": owned_prefix(env["RUN_PREFIX"]), + "instance_type": "c6id.4xlarge", "region": REGION, + "prior_estimated_usd": str(prior), "attempt_reserve_usd": str(reservation), + "ec2_hourly_usd": str(HOURLY_USD), "other_reserve_usd": str(OTHER_RESERVE_USD), + "cost_note": "Estimate, not an AWS billing total or a hard spending cap."} + + +def client(region): + import boto3 + from botocore.config import Config + + if region != REGION: + raise ValueError("S3 region must be ap-southeast-1") + return boto3.client("s3", region_name=region, config=Config( + connect_timeout=10, read_timeout=30, retries={"mode": "standard", "max_attempts": 3})) + + +def entries(s3, operation, bucket, prefix, fields, scope=""): + listed = owned_prefix(prefix) + scope + for page in s3.get_paginator(operation).paginate(Bucket=bucket, Prefix=listed): + for field in fields: + for item in page.get(field, []): + if not item["Key"].startswith(listed): + raise ValueError("S3 returned an object outside the owned prefix") + yield item + + +def inventory(s3, bucket, prefix, round=None): + # A round's snapshot lives under /; earlier rounds stay until cleanup. + if round not in (None, "smoke", "full"): + raise ValueError("round must be smoke or full") + scope = round + "/" if round else "" + objects = list(entries(s3, "list_objects_v2", bucket, prefix, ("Contents",), scope)) + return {"prefix": prefix + scope, "objects": len(objects), "bytes": sum(item["Size"] for item in objects), + "measurement": "S3 listing inventory; not request counts or transferred bytes."} + + +def probe(s3, bucket, prefix): + if inventory(s3, bucket, prefix)["objects"]: + raise ValueError("run prefix is already occupied; use a fresh workflow attempt") + key = prefix + "_probe" + s3.put_object(Bucket=bucket, Key=key, Body=b"metric snapshot benchmark") + try: + body = s3.get_object(Bucket=bucket, Key=key)["Body"] + try: + if body.read() != b"metric snapshot benchmark": + raise ValueError("S3 probe returned unexpected content") + finally: + body.close() + finally: + s3.delete_object(Bucket=bucket, Key=key) + return inventory(s3, bucket, prefix) + + +def cleanup(s3, bucket, prefix): + owned_prefix(prefix) + for item in entries(s3, "list_objects_v2", bucket, prefix, ("Contents",)): + s3.delete_object(Bucket=bucket, Key=item["Key"]) + after = inventory(s3, bucket, prefix) + if after["objects"]: + raise RuntimeError("owned S3 prefix is not empty after cleanup") + return after + + +def scrub(root, env): + secrets = [env[name].encode() for name in ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN") + if env.get(name)] + if not secrets: + raise ValueError("credential scan requires the CI credential environment") + overlap = max(map(len, secrets)) - 1 + for result in Path(root).rglob("*"): + if result.is_file(): + with result.open("rb") as stream: + previous = b"" + while chunk := stream.read(1024 * 1024): + data = previous + chunk + if any(secret in data for secret in secrets): + raise ValueError("credential detected in results; artifact upload is blocked") + previous = data[-overlap:] if overlap else b"" + + +def wait_for_benchmark(env, started_at): + """Bound runner allocation even when the benchmark job never leaves its queue.""" + run_id, attempt = owned_prefix(env["RUN_PREFIX"]).split("/")[1].split("-") + repo = env["GITHUB_REPOSITORY"] + if run_id != env["GITHUB_RUN_ID"] or not re.fullmatch(r"[\w.-]+/[\w.-]+", repo): + raise ValueError("workflow identity does not match its owned resource prefix") + minutes = int(env["BENCHMARK_TIMEOUT_MINUTES"]) + if not 30 <= minutes <= 240 or started_at <= 0: + raise ValueError("invalid benchmark allocation deadline") + # Use the original allocation timestamp; a delayed or retried watcher cannot extend it. + deadline = started_at + (45 + minutes + 10) * 60 + request = urllib.request.Request( + f"https://api.github.com/repos/{repo}/actions/runs/{run_id}/attempts/{attempt}/jobs?per_page=100", + headers={"Authorization": f"Bearer {env['GH_TOKEN']}", "Accept": "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28"}) + while True: + remaining = deadline - time.time() + if remaining <= 0: + raise TimeoutError("benchmark allocation deadline reached; proceeding to resource cleanup") + # API errors deliberately fail this step: subsequent always() steps clean up immediately. + with urllib.request.urlopen(request, timeout=min(15, remaining)) as response: + jobs = json.load(response)["jobs"] + benchmark = next((job for job in jobs if job["name"] == "benchmark"), None) + if benchmark and benchmark["status"] == "completed": + return {"benchmark_conclusion": benchmark["conclusion"], "owned_attempt": attempt} + remaining = deadline - time.time() + if remaining > 0: + time.sleep(min(30, remaining)) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("command", choices=("plan", "probe", "inventory", "cleanup", "scrub", "cost", "wait-for-benchmark")) + parser.add_argument("--output") + parser.add_argument("--started-at", type=int) + parser.add_argument("--round") + args = parser.parse_args() + if args.round is not None and args.command != "inventory": + parser.error("--round applies only to inventory") + if args.command == "plan": + result = plan(os.environ) + elif args.command == "scrub": + scrub(os.environ["RESULTS"], os.environ) + return + elif args.command == "wait-for-benchmark": + if args.started_at is None: + parser.error("wait-for-benchmark requires --started-at") + result = wait_for_benchmark(os.environ, args.started_at) + elif args.command == "cost": + if args.started_at is None: + parser.error("cost requires --started-at") + seconds = max(0, int(time.time()) - args.started_at) + estimate = HOURLY_USD * Decimal(seconds) / 3600 + result = {"allocation_window_seconds": seconds, "ec2_estimate_usd": str(round(estimate, 4)), + "campaign_estimate_with_other_reserve_usd": str(round( + Decimal(os.environ["PRIOR_SPEND_USD"]) + estimate + OTHER_RESERVE_USD, 4)), + "note": "Allocation wall-time estimate plus $1.50 reserve for S3, disk and IPv4; not an AWS bill."} + else: + prefix = owned_prefix(os.environ["RUN_PREFIX"]) + s3, bucket = client(os.environ["S3_REGION"]), os.environ["S3_BUCKET"] + if args.command == "inventory": + result = inventory(s3, bucket, prefix, args.round) + else: + result = globals()[args.command](s3, bucket, prefix) + rendered = json.dumps(result, indent=2) + "\n" + if args.output: + destination = Path(args.output) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(rendered) + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/oswal-s3-bench-crash.py b/.github/scripts/oswal-s3-bench-crash.py new file mode 100644 index 00000000000..35b382aa24b --- /dev/null +++ b/.github/scripts/oswal-s3-bench-crash.py @@ -0,0 +1,229 @@ +#!/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. + +"""Writer and oracle of the benchmark's crash cells. + +Every writer owns the stripe `first + index, first + index + writers, ...` of +batch sequence numbers, inserts `rows` rows per batch and stops at the first +statement that is not acknowledged, so a stripe has at most one attempted but +unacknowledged batch: the one in flight at the kill. Every request is logged +with its status and the time it ended, on the same clock as the kill stamps. +With `--seconds` the writers also stop once that long has passed, and the +command fails unless every writer wrote until then without a failure. +""" + +from __future__ import annotations + +import argparse +import collections +import json +import os +import sys +import threading +import time +import urllib.error +import urllib.parse +import urllib.request +from pathlib import Path + + +def insert(url: str, table: str, seq: int, rows: int) -> str: + values = ", ".join(f"({seq}, {seq * rows + i})" for i in range(rows)) + body = urllib.parse.urlencode({"sql": f"INSERT INTO {table} (seq, ts) VALUES {values}"}).encode() + try: + with urllib.request.urlopen(urllib.request.Request(url, data=body), timeout=60) as response: + return str(response.status) + except urllib.error.HTTPError as error: + return str(error.code) + except (OSError, urllib.error.URLError): + # The connection dropped: the request may or may not have been applied. + return "000" + + +def writer(args, index: int, deadline: float, failures: list) -> None: + directory, seq = Path(args.dir), args.first_seq + index + with open(directory / f"attempted-{index}.log", "a") as attempted, \ + open(directory / f"requests-{index}.log", "a") as requests, \ + open(directory / f"acked-{index}.log", "a") as acked: + while time.monotonic() < deadline: + attempted.write(f"{seq}\n") + attempted.flush() + code = insert(args.url, args.table, seq, args.rows) + requests.write(f"{seq} {code} {time.time_ns()}\n") + requests.flush() + if code != "200": + failures.append(f"writer {index}: batch {seq} ended with {code}") + break + acked.write(f"{seq}\n") + acked.flush() + seq += args.writers + (directory / f"exit-{index}.log").write_text(f"{time.time_ns()}\n") + + +def write(args) -> int: + os.makedirs(args.dir, exist_ok=True) + deadline = time.monotonic() + args.seconds if args.seconds else float("inf") + failures = [] + + def run(index): + try: + writer(args, index, deadline, failures) + except Exception as error: + failures.append(f"writer {index}: {error!r}") + raise + + threads = [threading.Thread(target=run, args=(i,)) for i in range(args.writers)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + # Crash cells end their writers with a kill; timed writers must have run to the deadline. + if args.seconds and (failures or not any(lines(Path(args.dir) / f"acked-{i}.log") for i in range(args.writers))): + print("; ".join(failures) or "no batch was acknowledged", file=sys.stderr) + return 1 + return 0 + + +def lines(path: Path) -> list: + return [line.split() for line in path.read_text().splitlines() if line.strip()] if path.exists() else [] + + +def found_batches(path: Path) -> dict: + """GreptimeDB's JSON answer to SELECT seq, count(*) ... GROUP BY seq, as batch -> row count.""" + return {int(seq): int(count) for seq, count in json.loads(path.read_text())["output"][0]["records"]["rows"]} + + +def check_kill(directory: Path, writers: int, before_ns: int, after_ns: int, window_max_ms: int) -> dict: + """A dropped connection is in flight when it ended at or after the stamp taken before the kill.""" + errors, in_flight, running = [], 0, 0 + if after_ns - before_ns > window_max_ms * 1_000_000: + errors.append(f"kill window {after_ns - before_ns}ns exceeds {window_max_ms}ms") + for index in range(writers): + exit_path = directory / f"exit-{index}.log" + if not exit_path.exists(): + errors.append(f"writer {index} left no exit time") + continue + running += int(exit_path.read_text()) >= before_ns + for seq, code, end_ns in lines(directory / f"requests-{index}.log"): + if code == "200": + continue + if code != "000": + errors.append(f"writer {index}: batch {seq} rejected with HTTP {code}") + elif int(end_ns) < before_ns: + errors.append(f"writer {index}: batch {seq} lost its connection before the kill") + else: + in_flight += 1 + if running == 0: + errors.append("no writer was running at the kill") + return {"kill_window_ns": after_ns - before_ns, "in_flight": in_flight, "writers_at_kill": running, + "kill_errors": errors} + + +def verify(run_dir: Path, found: dict, rows: int, enqueued: bool, cycle: int) -> dict: + """Checks the table, read as batch -> row count, against every cycle's writers so far and + against what every earlier restart recovered, and counts what the given cycle's writers + acknowledged, left and lost on their own.""" + acked, stripes, cycle_attempted, cycle_acked = [], collections.OrderedDict(), set(), set() + recovered_before = set() + for cycle_dir in sorted(run_dir.glob("cycle-*"), key=lambda p: int(p.name.split("-")[1])): + current = cycle_dir.name == f"cycle-{cycle}" + if int(cycle_dir.name.split("-")[1]) < cycle and (cycle_dir / "found.json").exists(): + recovered_before.update(found_batches(cycle_dir / "found.json")) + for path in sorted(cycle_dir.glob("attempted-*.log")): + stripes[(cycle_dir.name, path.name)] = [int(seq) for (seq,) in lines(path)] + if current: + cycle_attempted.update(stripes[(cycle_dir.name, path.name)]) + for path in sorted(cycle_dir.glob("acked-*.log")): + acked += [int(seq) for (seq,) in lines(path)] + if current: + cycle_acked.update(int(seq) for (seq,) in lines(path)) + acked_set = set(acked) + attempted = {seq for stripe in stripes.values() for seq in stripe} + errors = [] + duplicated = sorted(seq for seq, count in found.items() if count > rows) + partial = sorted(seq for seq, count in found.items() if count < rows) + if duplicated: + errors.append(f"batches present more than once: {duplicated[:20]}") + if partial: + errors.append(f"batches present in part: {partial[:20]}") + missing = sorted(acked_set - set(found)) + if missing and not enqueued: + errors.append(f"acknowledged but absent: {missing[:20]}") + # A batch an earlier restart recovered was durable then, so in either mode only the + # current cycle's backlog that was not durable at its kill may be missing. + vanished = sorted(recovered_before - set(found)) + if vanished: + errors.append(f"recovered by an earlier restart but absent: {vanished[:20]}") + unattempted = sorted(set(found) - attempted) + if unattempted: + errors.append(f"present but never attempted: {unattempted[:20]}") + # The unacknowledged tail of each stripe was in flight at a kill. + in_flight = set() + for key, stripe in stripes.items(): + prefix = 0 + while prefix < len(stripe) and stripe[prefix] in acked_set: + prefix += 1 + if any(seq in acked_set for seq in stripe[prefix:]): + errors.append(f"{key}: acknowledged after a failure") + in_flight.update(stripe[prefix:]) + extra = sorted(set(found) - acked_set) + stray = [seq for seq in extra if seq not in in_flight] + if stray and not enqueued: + errors.append(f"present, unacknowledged and not in flight at a kill: {stray[:20]}") + if len(acked) != len(acked_set): + errors.append("a batch was acknowledged twice") + present = set(found) & cycle_attempted + return {"cycle_acked_batches": len(cycle_acked), "cycle_present_batches": len(present), + "cycle_lost_batches": len(cycle_acked - present), "cycle_lost_rows": len(cycle_acked - present) * rows, + "cycle_extra_batches": len(present - cycle_acked), + "total_found_batches": len(found), "total_acked_batches": len(acked_set), + "total_lost_batches": len(missing), "total_extra_batches": len(extra), "verify_errors": errors} + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + w = commands.add_parser("write") + w.add_argument("--dir", required=True) + w.add_argument("--url", required=True) + w.add_argument("--table", default="crash") + for name in ("writers", "first-seq", "rows"): + w.add_argument(f"--{name}", type=int, required=True) + w.add_argument("--seconds", type=float, default=0, help="stop after this long; 0 writes until a failure") + c = commands.add_parser("check") + c.add_argument("--run-dir", type=Path, required=True) + c.add_argument("--cycle", type=int, required=True) + c.add_argument("--found", type=Path, required=True, help="GreptimeDB JSON output of SELECT seq, count(*)") + c.add_argument("--mode", choices=("durable", "enqueued"), required=True) + for name in ("writers", "rows", "before-ns", "after-ns", "window-max-ms"): + c.add_argument(f"--{name}", type=int, required=True) + c.add_argument("--pin-rows", type=int, help="rows of the pin table, whose one row was acknowledged first") + args = parser.parse_args() + if args.command == "write": + return write(args) + result = {"cycle": args.cycle} + result |= check_kill(args.run_dir / f"cycle-{args.cycle}", args.writers, args.before_ns, args.after_ns, + args.window_max_ms) + result |= verify(args.run_dir, found_batches(args.found), args.rows, + args.mode == "enqueued", args.cycle) + if args.pin_rows is not None and (args.pin_rows > 1 or args.mode == "durable" and args.pin_rows != 1): + result["verify_errors"].append(f"the pin table holds {args.pin_rows} rows, not 1") + result["passed"] = not result["kill_errors"] and not result["verify_errors"] + print(json.dumps(result)) + return 0 if result["passed"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/scripts/oswal-s3-bench-load.sh b/.github/scripts/oswal-s3-bench-load.sh new file mode 100644 index 00000000000..1748fb40d66 --- /dev/null +++ b/.github/scripts/oswal-s3-bench-load.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# Runs inside the o11ybench runtime image: the load of scripts/agent-observability/greptimedb.sh +# with the same arguments plus the loader's --skip-create-table, for a table the cell created. +set -euo pipefail + +mkdir "/results/$RESULT_NAME" +mkdir -p /results/corpus +cp /work/corpus/summary.json /work/corpus/evidence.json /results/corpus/ +w="$PWD/workloads/logbench-agent-observability" +r="/results/$RESULT_NAME" +prepare_workers=${PREPARE_WORKERS:-$(python3 -c 'import sys; sys.path.insert(0, "scripts"); from benchmark_workers import effective_cpus; print(effective_cpus())')} +o11ybench-greptimedb-loader --workload logbench-agent-observability --prepare-workers "$prepare_workers" \ + --skip-create-table \ + --grpc-endpoint "$GRPC_ENDPOINT" --sql-endpoint "$ENDPOINT" \ + --database public --table agent_observations --batch-size "${GREPTIMEDB_BATCH_SIZE:-10000}" --write-mode stream --max-in-flight "${MAX_IN_FLIGHT:-1}" \ + --flush-table --request-timeout-seconds 900 \ + --schema "$w/schema.json" --physical-index-profiles "$w/physical-index-profiles.json" \ + --data /work/corpus/agent_observations.jsonl \ + --profile-summary /work/corpus/summary.json --evidence /work/corpus/evidence.json \ + --transcript-output "$r/load-transcript.jsonl" \ + --target-metadata-output "$r/target-metadata.json" \ + --greptimedb-image "$TARGET_IMAGE" --greptimedb-version "$TARGET_VERSION" \ + --greptimedb-digest "$TARGET_DIGEST" \ + --startup-input=standalone --startup-input=start \ + --startup-input=--http-addr=0.0.0.0:4000 \ + --startup-input=--grpc-bind-addr=0.0.0.0:4001 \ + > "$r/load-result.json" diff --git a/.github/scripts/oswal-s3-bench-run.sh b/.github/scripts/oswal-s3-bench-run.sh index a9835c1e5f1..d158ab14d60 100644 --- a/.github/scripts/oswal-s3-bench-run.sh +++ b/.github/scripts/oswal-s3-bench-run.sh @@ -10,24 +10,64 @@ runtime() { --security-opt no-new-privileges --cap-drop ALL "$@" } +# Pin the runtime image to a digest without pulling it: ECR Public's anonymous data limit can refuse a +# hosted runner even a manifest read. A reference that names its digest is pinned as given; only a tag +# is resolved from the registry's manifest. The benchmark host pulls the image and checks the contents. check_runtime() { - local image=$1 resolved + local image=$1 repository digest attempt component='[a-z0-9]+(([._]|__|-+)[a-z0-9]+)*' [[ "$image" == public.ecr.aws/* ]] || { echo 'The runtime image must come from AWS ECR Public' >&2; return 1; } - bash "$here/pull-benchmark-image.sh" "$image" - resolved=$(docker image inspect "$image" | jq -er '.[0].RepoDigests[0]') - docker run --rm --entrypoint sh "$resolved" -c \ - 'test -f scripts/agent-observability/greptimedb.sh && test -f scripts/generate_agent_observations.py' \ - || { echo 'The runtime image lacks the agent observability loader or generator' >&2; return 1; } - printf 'runtime_image=%s\n' "$resolved" >> "${GITHUB_OUTPUT:-/dev/stdout}" + # Docker's reference grammar, checked here because a named digest never reaches the registry. + local grammar="^public\.ecr\.aws(/$component)+(:[A-Za-z0-9_][A-Za-z0-9_.-]{0,127})?(@sha256:[0-9a-f]{64})?\$" + [[ "$image" =~ $grammar ]] || { echo "$image is not a full ECR Public image reference" >&2; return 1; } + if [[ "$image" == *@* ]]; then + digest=${image##*@} + else + for attempt in 1 2 3; do + digest=$(docker buildx imagetools inspect --format '{{json .Manifest}}' "$image" | jq -er .digest) && break + ((attempt < 3)) || { echo "Cannot read the manifest of $image" >&2; return 1; } + sleep $((attempt * 10)) + done + fi + [[ "$digest" =~ ^sha256:[0-9a-f]{64}$ ]] || { echo "$image has no sha256 digest" >&2; return 1; } + repository=${image%@*} + [[ "${repository##*/}" != *:* ]] || repository=${repository%:*} + printf 'runtime_image=%s@%s\n' "$repository" "$digest" >> "${GITHUB_OUTPUT:-/dev/stdout}" +} + +# Pull every GreptimeDB image of the run, pin it to its digest and record the version and +# commit it reports, before any instance is allocated. +check_images() { + local entry image resolved info version commit checked=() + while IFS= read -r entry; do + image=$(jq -r .image <<< "$entry") + bash "$here/pull-benchmark-image.sh" "$image" + resolved=$(docker image inspect "$image" | jq -er '[.[0].RepoDigests[]? | select(test("@sha256:[0-9a-f]{64}$"))][0]') \ + || { echo "$image has no sha256 repository digest" >&2; return 1; } + info=$(docker run --rm --entrypoint greptime "$resolved" --version) + version=$(sed -n 's/^version: //p' <<< "$info") commit=$(sed -n 's/^commit: //p' <<< "$info") + [[ -n "$version" && -n "$commit" ]] || { echo "$image reports no version or commit" >&2; return 1; } + echo "${image}: $resolved, version $version, commit $commit" + checked+=("$(jq -c --arg image "$resolved" --arg version "$version" --arg commit "$commit" \ + '{label, image: $image, requested: .image, version: $version, commit: $commit}' <<< "$entry")") + done < <(jq -c '.[]' <<< "$1") + ((${#checked[@]})) || { echo 'No GreptimeDB image to check' >&2; return 1; } + printf 'greptimedb_images=%s\n' "$(jq -sc . <<< "${checked[*]}")" >> "${GITHUB_OUTPUT:-/dev/stdout}" } pull_images() { - : "${RESULTS:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_TAG:?}" + : "${RESULTS:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_IMAGES:?}" + local image images=() mkdir -p "$RESULTS/images" bash "$here/pull-benchmark-image.sh" "$RUNTIME_IMAGE" - bash "$here/pull-benchmark-image.sh" "greptime/greptimedb:$GREPTIMEDB_TAG" + docker run --rm --entrypoint sh "$RUNTIME_IMAGE" -c \ + 'test -f scripts/agent-observability/greptimedb.sh && test -f scripts/generate_agent_observations.py' \ + || { echo 'The runtime image lacks the agent observability loader or generator' >&2; return 1; } + while IFS= read -r image; do + bash "$here/pull-benchmark-image.sh" "$image" + images+=("$image") + done < <(jq -r '.[].image' <<< "$GREPTIMEDB_IMAGES") docker image inspect "$RUNTIME_IMAGE" > "$RESULTS/images/runtime.json" - docker image inspect "greptime/greptimedb:$GREPTIMEDB_TAG" > "$RESULTS/images/greptimedb.json" + docker image inspect "${images[@]}" > "$RESULTS/images/greptimedb.json" } wal_disk() { @@ -46,25 +86,242 @@ wal_disk() { findmnt -no SOURCE,FSTYPE -T "$root" | tee "$RESULTS/wal-disk.txt" } +# One corpus per row count the load cells use; crash cells write their own rows. generate() { : "${PLAN:?}" "${RESULTS:?}" "${CORPUS_DIR:?}" "${OWNER:?}" "${RUNTIME_IMAGE:?}" - local args=() arg - while IFS= read -r arg; do args+=("$arg"); done < <(jq -r '.generator_args[]' "$PLAN") - mkdir -p "$CORPUS_DIR" - runtime --name "$OWNER-generator" --mount "type=bind,src=$CORPUS_DIR,dst=/work/corpus" \ - "$RUNTIME_IMAGE" scripts/generate_agent_observations.py "${args[@]}" \ - --output /work/corpus/agent_observations.jsonl \ - --summary-output /work/corpus/summary.json --evidence-output /work/corpus/evidence.json \ - > "$RESULTS/generator.log" 2>&1 - jq -e --argjson rows "$(jq .rows "$PLAN")" '.line_count == $rows' "$CORPUS_DIR/summary.json" >/dev/null + local corpus args arg dir rows + while IFS= read -r corpus; do + args=() dir="$CORPUS_DIR/$(jq -r .dir <<< "$corpus")" rows=$(jq .rows <<< "$corpus") + while IFS= read -r arg; do args+=("$arg"); done < <(jq -r '.generator_args[]' <<< "$corpus") + mkdir -p "$dir" + runtime --name "$OWNER-generator-$rows" --mount "type=bind,src=$dir,dst=/work/corpus" \ + "$RUNTIME_IMAGE" scripts/generate_agent_observations.py "${args[@]}" \ + --output /work/corpus/agent_observations.jsonl \ + --summary-output /work/corpus/summary.json --evidence-output /work/corpus/evidence.json \ + >> "$RESULTS/generator.log" 2>&1 + jq -e --argjson rows "$rows" '.line_count == $rows' "$dir/summary.json" >/dev/null + done < <(jq -c '.corpora[]' "$PLAN") } -# Scrape once, after the load succeeded or failed, while the server still runs. +# Scrape once per server, after the load succeeded or failed, while it still runs: +# the loaded server into metrics.prom, a restarted one into metrics-restart.prom. scrape_metrics() { curl --noproxy '*' -fsS --max-time 30 http://127.0.0.1:4000/metrics > "$dir/metrics.all" || { rm -f "$dir/metrics.all"; return 1; } - grep -E '^(# (HELP|TYPE) )?greptime_logstore_object_store_wal_' "$dir/metrics.all" > "$dir/metrics.prom" || true + grep -E '^(# (HELP|TYPE) )?(greptime_logstore_object_store_wal_|opendal_operation_(duration_seconds_count|errors_total))' \ + "$dir/metrics.all" > "$dir/metrics$phase.prom" || true rm -f "$dir/metrics.all" - metrics_scraped=true + case "$phase" in "") metrics_scraped=true ;; -restart) restart_scraped=true ;; esac +} + +# A crash cycle cannot repeat a scrape once its server is killed, so it retries. +# The named WAL sample must be in the scrape, or it does not count. +scrape_retry() { + local required=$1 attempt + for attempt in 1 2 3; do + scrape_metrics && grep -q "^greptime_logstore_object_store_wal_$required " "$dir/metrics$phase.prom" && return 0 + sleep 1 + done + return 1 +} + +sql() { curl --noproxy '*' -fsS --max-time 30 --data-urlencode "sql=$1" 'http://127.0.0.1:4000/v1/sql?db=public'; } +now_ns() { python3 -c 'import time; print(time.time_ns())'; } + +now() { date +%s.%N; } +elapsed() { awk -v a="$1" -v b="$(now)" 'BEGIN {printf "%.3f", b - a}'; } +# CPU seconds the server's cgroup used, and busy and total jiffies of the host. +server_cpu() { docker exec "$cid" cat /sys/fs/cgroup/cpu.stat 2>/dev/null | awk '/^usage_usec / {printf "%.3f", $2 / 1e6}' || true; } +host_cpu() { awk '/^cpu / {print $2 + $3 + $4 + $7 + $8 + $9, $2 + $3 + $4 + $5 + $6 + $7 + $8 + $9}' /proc/stat 2>/dev/null || true; } +# The host's 1, 5 and 15 minute load averages as a JSON list. +host_load() { awk '{printf "[%s, %s, %s]", $1, $2, $3}' "${OSWAL_LOADAVG:-/proc/loadavg}" 2>/dev/null || true; } + +# A server's log is named after the phase it started in, whatever phase stops it. +start_server() { + server_phase=$phase + cid=$(docker run -d --init --name "$name$phase" --label "o11ybench.owner=$OWNER" \ + --add-host host.docker.internal:host-gateway \ + -p 127.0.0.1:4000:4000 -p 127.0.0.1:4001:4001 "${docker_env[@]}" ${mounts[@]+"${mounts[@]}"} \ + "$GREPTIMEDB_IMAGE" standalone start --http-addr=0.0.0.0:4000 --grpc-bind-addr=0.0.0.0:4001) + local deadline=$((SECONDS + 300)) + until curl --noproxy '*' -fsS --max-time 5 http://127.0.0.1:4000/health >/dev/null 2>&1; do + [[ $(docker inspect --format '{{.State.Running}}' "$cid") == true ]] || { echo "$id: server$phase exited" >&2; exit 1; } + ((SECONDS < deadline)) || { echo "$id: server$phase not healthy in 300s" >&2; exit 1; } + sleep 0.2 + done +} + +# kill sends SIGKILL, killed collects a server the caller already killed, and +# no argument stops it gracefully. +stop_server() { + case "${1:-}" in + kill) docker kill --signal KILL "$cid" >/dev/null ;; + killed) docker wait "$cid" >/dev/null ;; + *) docker stop -t 120 "$cid" >/dev/null ;; + esac + docker logs "$cid" > "$dir/server$server_phase.log" 2>&1 + docker rm -fv "$cid" >/dev/null + cid= +} + +# Kill the loaded server, start the same configuration on the same prefix and +# data home, and time it until /health and until the loaded table answers. +restart() { + local table count started deadline + table=$(jq -er .table "$dir/greptimedb-1/load-result.json") + scrape_metrics || { echo "$id: metrics scrape failed" >&2; exit 4; } + # The objects recovery has to read are those still indexed at the kill. + indexed_before_kill=$(awk '/^greptime_logstore_object_store_wal_indexed_objects / {print $2}' "$dir/metrics.prom") + stop_server kill + phase=-restart started=$(now) + start_server + restart_seconds=$(elapsed "$started") + deadline=$((SECONDS + 300)) + # Regions may still be opening after /health. + until count=$(curl --noproxy '*' -fsS --max-time 30 --data-urlencode "sql=SELECT count(*) FROM $table" \ + 'http://127.0.0.1:4000/v1/sql?db=public' | jq -er '.output[0].records.rows[0][0]'); do + ((SECONDS < deadline)) || { echo "$id: $table not readable in 300s after restart" >&2; exit 5; } + sleep 0.2 + done + restart_query_seconds=$(elapsed "$started") restart_count=$count + [[ "$count" == "$(jq .rows <<< "$spec")" ]] || { echo "$id: $count rows after restart" >&2; exit 5; } + # The restart was measured, but on fewer objects than the cell asked for. + awk -v n="${indexed_before_kill:-0}" -v min="$(jq -r '.min_objects // 0' <<< "$spec")" 'BEGIN {exit !(n >= min)}' \ + || { echo "$id: ${indexed_before_kill:-no} objects indexed at the kill, fewer than min_objects" >&2; exit 7; } +} + +# Pins the WAL with one row in a second table before the writers start, if asked, +# and samples the object counts every 30 s while they run. +sample_objects() { + set +e + local body + while sleep "${OSWAL_SAMPLE_SECONDS:-30}"; do + # A failed scrape is a missing sample, never zero objects. + if ! body=$(curl --noproxy '*' -fsS --max-time 10 http://127.0.0.1:4000/metrics 2>/dev/null) \ + || ! grep -q '^greptime_logstore_object_store_wal_indexed_objects ' <<< "$body"; then + echo "$(date +%s) missing" >> "$dir/objects.tsv" + continue + fi + awk -v t="$(date +%s)" ' + /^greptime_logstore_object_store_wal_indexed_objects / {i = $2} + /^greptime_logstore_object_store_wal_deleted_objects_total / {d = $2} + /^greptime_logstore_object_store_wal_created_objects_total / {c = $2} + END {print t, i + 0, d + 0, c + 0}' <<< "$body" >> "$dir/objects.tsv" + done +} + +# Writers insert striped batches until a SIGKILL at a random point stops the +# server; it restarts on the same prefix and data home and the table is checked +# against what was acknowledged. Nothing is cleaned between cycles. +crash_cycles() { + local writers rows cycles kill_min kill_max cycle cdir wpid spid= delay before after started restart_s query_s next=0 pin_count pin_args + local load_start complete pin_left pin_base deleted sql_url=${OSWAL_SQL_URL:-http://127.0.0.1:4000/v1/sql?db=public} + writers=$(jq -r .concurrency <<< "$spec") rows=$(jq -r .batch_rows "$PLAN") cycles=$(jq -r .cycles <<< "$spec") + kill_min=$(jq -r '.kill_after[0]' <<< "$spec") kill_max=$(jq -r '.kill_after[1]' <<< "$spec") + if jq -e .pin <<< "$spec" >/dev/null; then + # Only the objects indexed before the pin row lie below it, such as the empty start object that + # the row's write lets the store collect at any later point; the servers may collect no more. + phase=-pin; scrape_retry indexed_objects || { echo "$id: metrics scrape failed" >&2; exit 4; } + read -r pin_left pin_base < <(awk '/^greptime_logstore_object_store_wal_indexed_objects / {i = $2} + /^greptime_logstore_object_store_wal_deleted_objects_total / {d = $2} END {printf "%d %d\n", i, d}' "$dir/metrics-pin.prom") + # The pin table's own auto flush outlasts every cycle a crash group allows, so its row holds the WAL. + sql "CREATE TABLE pin (v BIGINT, ts TIMESTAMP TIME INDEX) WITH ('auto_flush_interval' = '1d')" >/dev/null + sql "INSERT INTO pin VALUES (1, 0)" >/dev/null + fi + sql "CREATE TABLE crash (seq BIGINT, ts TIMESTAMP TIME INDEX) WITH (append_mode = 'true')" >/dev/null + for ((cycle = 1; cycle <= cycles; cycle++)); do + cdir="$dir/cycle-$cycle" complete=true; mkdir -p "$cdir" + # The counters of this server process before its writers start are the cycle's baseline. + phase=-start-$cycle; scrape_retry indexed_objects || complete=false + load_start=$(now_ns) + python3 "$here/oswal-s3-bench-crash.py" write --dir "$cdir" --url "$sql_url" \ + --writers "$writers" --first-seq "$next" --rows "$rows" & + wpid=$! + if jq -e .pin <<< "$spec" >/dev/null; then sample_objects & spid=$!; fi + delay=$(awk -v a="$kill_min" -v b="$kill_max" 'BEGIN {srand(); printf "%.3f", a + rand() * (b - a)}') + sleep "$delay" + phase=-cycle-$cycle; scrape_retry indexed_objects || complete=false + before=$(now_ns); docker kill --signal KILL "$cid" >/dev/null; after=$(now_ns) + wait "$wpid" + if [[ -n "$spid" ]]; then kill "$spid" 2>/dev/null || true; wait "$spid" 2>/dev/null || true; spid=; fi + stop_server killed + phase=-restart-$cycle started=$(now) + start_server + restart_s=$(elapsed "$started") + until sql "SELECT seq, count(*) FROM crash GROUP BY seq" > "$cdir/found.json" 2>/dev/null; do + (($(elapsed "$started" | cut -d. -f1) < 300)) || { echo "$id: crash not readable after restart $cycle" >&2; exit 5; } + sleep 0.2 + done + query_s=$(elapsed "$started") + scrape_retry recovery_seconds_count || complete=false + pin_count=null pin_args=() + if jq -e .pin <<< "$spec" >/dev/null; then + pin_count=$(sql "SELECT count(*) FROM pin" | jq -er '.output[0].records.rows[0][0]') pin_args=(--pin-rows "$pin_count") + fi + local verdict=0 + python3 "$here/oswal-s3-bench-crash.py" check --run-dir "$dir" --cycle "$cycle" --found "$cdir/found.json" \ + --mode "$(jq -r .wal.ACK_MODE <<< "$spec")" --writers "$writers" --rows "$rows" \ + --before-ns "$before" --after-ns "$after" --window-max-ms 2000 ${pin_args[@]+"${pin_args[@]}"} \ + > "$cdir/check.json" || verdict=$? + jq -c --arg delay "$delay" --arg restart "$restart_s" --arg query "$query_s" --argjson pin "$pin_count" \ + --argjson load "$(awk -v a="$load_start" -v b="$before" 'BEGIN {printf "%.3f", (b - a) / 1e9}')" \ + --argjson complete "$complete" \ + '. + {kill_after_seconds: ($delay | tonumber), load_seconds: $load, restart_seconds: ($restart | tonumber), + restart_query_seconds: ($query | tonumber), pin_rows: $pin, metrics_complete: $complete}' \ + "$cdir/check.json" >> "$dir/cycles.jsonl" + ((verdict == 0)) || { echo "$id: cycle $cycle failed its check" >&2; exit 6; } + # A cycle without its scrapes is not a measurement, whatever the table held. + [[ "$complete" == true ]] || { echo "$id: cycle $cycle lost a metrics scrape" >&2; exit 4; } + # Each server counts its collections from its own start: the first one from its pin scrape, + # every restarted one from zero, and an absent counter has not moved. + if jq -e .pin <<< "$spec" >/dev/null; then + deleted=$(awk '/^greptime_logstore_object_store_wal_deleted_objects_total / {d = $2} END {printf "%d", d}' \ + "$dir/metrics-cycle-$cycle.prom") + pin_left=$((pin_left - (deleted - pin_base))) pin_base=0 + ((pin_left >= 0)) || { echo "$id: cycle $cycle collected WAL objects the pin held" >&2; exit 11; } + fi + next=$(cat "$cdir"/attempted-*.log | sort -n | tail -1); next=$((next + 1)) + done +} + +# Writes the table `dropped` for write_seconds, drops it and kills the server right +# away, while most of its objects are still indexed; restarts on the same prefix and +# data home, then writes and flushes the table `kept` and samples the indexed objects +# until the dropped table's objects are collected. +drop_restart() { + local writers rows seconds started deadline indexed recovered sql_url=${OSWAL_SQL_URL:-http://127.0.0.1:4000/v1/sql?db=public} + writers=$(jq -r .concurrency <<< "$spec") rows=$(jq -r .batch_rows "$PLAN") seconds=$(jq -r .write_seconds <<< "$spec") + sql "CREATE TABLE dropped (seq BIGINT, ts TIMESTAMP TIME INDEX) WITH (append_mode = 'true')" >/dev/null + sql "CREATE TABLE kept (seq BIGINT, ts TIMESTAMP TIME INDEX) WITH (append_mode = 'true')" >/dev/null + python3 "$here/oswal-s3-bench-crash.py" write --dir "$dir/dropped" --url "$sql_url" --table dropped \ + --writers "$writers" --first-seq 0 --rows "$rows" --seconds "$seconds" \ + || { echo "$id: the writers of dropped stopped before $seconds s" >&2; exit 10; } + phase=-written; scrape_retry indexed_objects || { echo "$id: metrics scrape failed" >&2; exit 4; } + sql "DROP TABLE dropped" >/dev/null + stop_server kill + phase=-restart started=$(now) + start_server + restart_seconds=$(elapsed "$started") + scrape_retry recovery_seconds_count || { echo "$id: metrics scrape failed" >&2; exit 4; } + recovered=$(awk '/^greptime_logstore_object_store_wal_recovered_objects_total / {print $2}' "$dir/metrics-restart.prom") + python3 "$here/oswal-s3-bench-crash.py" write --dir "$dir/kept" --url "$sql_url" --table kept \ + --writers 1 --first-seq 0 --rows "$rows" --seconds 2 \ + || { echo "$id: the writer of kept stopped before 2 s" >&2; exit 10; } + sql "ADMIN FLUSH_TABLE('kept')" >/dev/null + # The flushed watermark becomes durable with the next object, so every sample writes one row. + deadline=$((SECONDS + ${OSWAL_COLLECT_SECONDS:-120})) + while :; do + sql "INSERT INTO kept VALUES (-1, $SECONDS)" >/dev/null + # The samples have their own file; the final scrape keeps the restart's name. + phase=-collected + scrape_retry indexed_objects || { phase=-restart; echo "$id: metrics scrape failed" >&2; exit 4; } + phase=-restart + indexed=$(awk '/^greptime_logstore_object_store_wal_indexed_objects / {print $2}' "$dir/metrics-collected.prom") + printf '%s %s\n' "$(date +%s)" "$indexed" >> "$dir/collected.tsv" + # A start object, the newest object and kept's own may stay; the dropped table's go. + awk -v i="$indexed" -v r="${recovered:-0}" 'BEGIN {exit !(i <= 3 || 2 * i < r)}' && break + ((SECONDS < deadline)) || { echo "$id: $indexed of ${recovered:-?} recovered objects still indexed" >&2; exit 9; } + sleep 1 + done } finish() { @@ -74,39 +331,62 @@ finish() { if [[ -n "$cid" ]]; then # A load that succeeded without its metrics is not a measurement. scrape_metrics || { echo "$id: metrics scrape failed" >&2; [[ $rc != 0 ]] || rc=4; } - docker stop -t 120 "$cid" >/dev/null - docker logs "$cid" > "$dir/server.log" 2>&1 - docker rm -fv "$cid" >/dev/null + stop_server fi [[ -z "$wal_dir" ]] || sudo -n rm -rf "$wal_dir" + [[ -z "$home_dir" ]] || sudo -n rm -rf "$home_dir" jq -n --argjson spec "$spec" --arg device "$wal_device" --arg version "$version" \ + --arg host_load_start "$host_load_start" --arg host_load_end "$(host_load)" \ --arg load_seconds "$load_seconds" --argjson exit_code "$rc" --argjson scraped "$metrics_scraped" \ + --argjson restart_scraped "$restart_scraped" --arg restart_seconds "$restart_seconds" \ + --arg restart_query_seconds "$restart_query_seconds" --arg restart_count "$restart_count" \ + --arg indexed_before_kill "$indexed_before_kill" \ + --arg server_cpu_seconds "$server_cpu_seconds" --arg host_cpu_busy "$host_cpu_busy" --arg host_cpus "$(nproc 2>/dev/null)" \ '$spec + {wal_device: $device, version: $version, exit_code: $exit_code, metrics_scraped: $scraped, - load_wall_seconds: ($load_seconds | tonumber? // null)}' > "$dir/cell.json" + load_wall_seconds: ($load_seconds | tonumber? // null), + server_cpu_seconds: ($server_cpu_seconds | tonumber? // null), + host_cpu_busy: ($host_cpu_busy | tonumber? // null), host_cpus: ($host_cpus | tonumber? // null), + restart_scraped: $restart_scraped, restart_seconds: ($restart_seconds | tonumber? // null), + restart_query_seconds: ($restart_query_seconds | tonumber? // null), + restart_count: ($restart_count | tonumber? // null), + indexed_before_kill: ($indexed_before_kill | tonumber? // null), + host_load_start: ($host_load_start | fromjson? // null), host_load_end: ($host_load_end | fromjson? // null)}' \ + > "$dir/cell.json" exit "$rc" } # One cell in its own process, so errexit holds inside it. The cell's state is # global because the EXIT trap runs after this function returns. cell() { - : "${PLAN:?}" "${RESULTS:?}" "${CORPUS_DIR:?}" "${OWNER:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_IMAGE:?}" + : "${PLAN:?}" "${RESULTS:?}" "${CORPUS_DIR:?}" "${OWNER:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_IMAGES:?}" : "${RUN_PREFIX:?}" "${WAL_ROOT:?}" "${S3_BUCKET:?}" "${S3_REGION:?}" "${AWS_ACCESS_KEY_ID:?}" "${AWS_SECRET_ACCESS_KEY:?}" - spec=$1 cid= wal_dir= version= load_seconds= metrics_scraped=false id= - local concurrency name setting started ended - id=$(jq -r .id <<< "$spec"); concurrency=$(jq -r .concurrency <<< "$spec") - dir="$RESULTS/cells/$id"; name="$OWNER-$id" + spec=$1 cid= wal_dir= home_dir= version= load_seconds= metrics_scraped=false id= phase= server_phase= + restart_scraped=false restart_seconds= restart_query_seconds= restart_count= server_cpu_seconds= host_cpu_busy= + indexed_before_kill= host_load_start=$(host_load) + local concurrency setting started cpu_before host_before host_after corpus partitions load script label + id=$(jq -r .id <<< "$spec"); concurrency=$(jq -r .concurrency <<< "$spec") label=$(jq -r .image <<< "$spec") + dir="$RESULTS/cells/$id" name="$OWNER-$id" + # The digest the image of the cell's label was pinned to before the run. + GREPTIMEDB_IMAGE=$(jq -er --arg label "$label" '.[] | select(.label == $label) | .image' <<< "$GREPTIMEDB_IMAGES") \ + || { echo "$id: no image labelled '$label'" >&2; exit 1; } mkdir -p "$dir" # The root holds data/ and the WAL's default prefix wal/ side by side. local settings=(STORAGE__TYPE=S3 "STORAGE__BUCKET=$S3_BUCKET" "STORAGE__REGION=$S3_REGION" "STORAGE__ROOT=$RUN_PREFIX/$id") [[ -z "${S3_CONTAINER_ENDPOINT:-}" ]] || settings+=("STORAGE__ENDPOINT=$S3_CONTAINER_ENDPOINT") + # A restarted server needs the metadata the loaded one kept in its data home. + jq -e '.restart or .mode != "load"' <<< "$spec" >/dev/null && settings+=(STORAGE__DATA_HOME=/oswal_home) while IFS= read -r setting; do settings+=("$setting"); done \ < <(jq -r '.wal | to_entries[] | "WAL__\(.key)=\(.value)"' <<< "$spec") - local docker_env=() mounts=() + docker_env=() mounts=() for setting in "${settings[@]}"; do docker_env+=(-e "GREPTIMEDB_STANDALONE__$setting") printf 'GREPTIMEDB_STANDALONE__%s\n' "$setting" >> "$dir/config.env" done + while IFS= read -r setting; do + docker_env+=(-e "$setting") + printf '%s\n' "$setting" >> "$dir/config.env" + done < <(jq -r '.env // {} | to_entries[] | "\(.key)=\(.value)"' <<< "$spec") # Pass the two keys by name: their values never reach argv, a file or the artifact. export GREPTIMEDB_STANDALONE__STORAGE__ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID export GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY @@ -119,33 +399,60 @@ cell() { else wal_device="s3://$S3_BUCKET/$RUN_PREFIX/$id/wal" fi + if jq -e '.restart or .mode != "load"' <<< "$spec" >/dev/null; then + home_dir="$WAL_ROOT/$id-home"; mkdir -p "$home_dir" + mounts+=(--mount "type=bind,src=$home_dir,dst=/oswal_home") + fi trap finish EXIT - cid=$(docker run -d --init --name "$name" --label "o11ybench.owner=$OWNER" \ - --add-host host.docker.internal:host-gateway \ - -p 127.0.0.1:4000:4000 -p 127.0.0.1:4001:4001 "${docker_env[@]}" ${mounts[@]+"${mounts[@]}"} \ - "$GREPTIMEDB_IMAGE" standalone start --http-addr=0.0.0.0:4000 --grpc-bind-addr=0.0.0.0:4001) - local deadline=$((SECONDS + 300)) - until curl --noproxy '*' -fsS --max-time 5 http://127.0.0.1:4000/health >/dev/null 2>&1; do - [[ $(docker inspect --format '{{.State.Running}}' "$cid") == true ]] || { echo "$id: server exited" >&2; exit 1; } - ((SECONDS < deadline)) || { echo "$id: server not healthy in 300s" >&2; exit 1; } - sleep 2 - done + start_server version=$(docker exec "$cid" greptime --version | sed -n 's/^version: //p') [[ -n "$version" ]] - started=$(date +%s.%N) - # Only the load phase, run exactly as run-agent-observability.sh runs it. + if jq -e '.mode == "crash"' <<< "$spec" >/dev/null; then crash_cycles; return; fi + if jq -e '.mode == "drop"' <<< "$spec" >/dev/null; then drop_restart; return; fi + corpus="$CORPUS_DIR/corpus-$(jq -r .rows <<< "$spec")" partitions=$(jq -r '.partitions // 1' <<< "$spec") + # Only the load phase, run exactly as run-agent-observability.sh runs it, or with the + # loader writing into a table split into regions that the cell created first. + load=(scripts/agent-observability/greptimedb.sh load) script=() + if ((partitions > 1)); then + sql "$(python3 "$here/oswal-s3-bench.py" table --summary "$corpus/summary.json" --partitions "$partitions")" >/dev/null + load=(/work/oswal-s3-bench-load.sh) + script=(--mount "type=bind,src=$here/oswal-s3-bench-load.sh,dst=/work/oswal-s3-bench-load.sh,readonly") + fi + cpu_before=$(server_cpu) host_before=$(host_cpu) + started=$(now) runtime --name "$name-load" --network host \ - --mount "type=bind,src=$CORPUS_DIR,dst=/work/corpus,readonly" --mount "type=bind,src=$dir,dst=/results" \ - -e GRPC_ENDPOINT=127.0.0.1:4001 -e "DATA_PROFILE=$(jq -r .data_profile "$PLAN")" -e RESULT_NAME=greptimedb-1 \ + --mount "type=bind,src=$corpus,dst=/work/corpus,readonly" ${script[@]+"${script[@]}"} \ + --mount "type=bind,src=$dir,dst=/results" \ + -e GRPC_ENDPOINT=127.0.0.1:4001 -e "DATA_PROFILE=$(jq -r .data_profile <<< "$spec")" -e RESULT_NAME=greptimedb-1 \ -e "ENDPOINT=http://127.0.0.1:4000/v1/sql?db=public&format=greptimedb_v1" \ -e "TARGET_IMAGE=$GREPTIMEDB_IMAGE" -e "TARGET_VERSION=$version" -e "TARGET_DIGEST=${GREPTIMEDB_IMAGE##*@}" \ -e "TARGET_ID=$name" -e TARGET_HOST=127.0.0.1 -e PREPARE_WORKERS= -e "MAX_IN_FLIGHT=$concurrency" \ -e "GREPTIMEDB_BATCH_SIZE=$(jq -r .batch_rows "$PLAN")" --entrypoint /bin/bash \ - "$RUNTIME_IMAGE" scripts/agent-observability/greptimedb.sh load > "$dir/load.log" 2>&1 - ended=$(date +%s.%N) - load_seconds=$(awk -v a="$started" -v b="$ended" 'BEGIN {printf "%.3f", b - a}') - jq -e --argjson rows "$(jq .rows "$PLAN")" '.count_before == 0 and .count_after == $rows' \ + "$RUNTIME_IMAGE" "${load[@]}" > "$dir/load.log" 2>&1 + load_seconds=$(elapsed "$started") + host_after=$(host_cpu) + server_cpu_seconds=$(awk -v a="$cpu_before" -v b="$(server_cpu)" 'BEGIN {if (a != "" && b != "") printf "%.3f", b - a}') + # Busy share of all host CPUs during the load, the loader included. + host_cpu_busy=$(awk -v a="$host_before" -v b="$host_after" \ + 'BEGIN {split(a, x, " "); split(b, y, " "); if (y[2] > x[2]) printf "%.4f", (y[1] - x[1]) / (y[2] - x[2])}') + jq -e --argjson rows "$(jq .rows <<< "$spec")" '.count_before == 0 and .count_after == $rows' \ "$dir/greptimedb-1/load-result.json" >/dev/null + if ((partitions > 1)); then + # The server's own account of the rows each region took. + sql "SELECT p.partition_name, s.region_id, s.region_rows, s.written_bytes_since_open + FROM information_schema.partitions p JOIN information_schema.region_statistics s + ON p.greptime_partition_id = s.region_id + WHERE p.table_schema = 'public' AND p.table_name = 'agent_observations' + ORDER BY s.region_id" > "$dir/regions.json" + jq -e --argjson n "$partitions" '[.output[0].records.rows[] | select(.[2] > 0)] | length == $n' \ + "$dir/regions.json" >/dev/null || { echo "$id: fewer than $partitions regions took rows" >&2; exit 8; } + fi + # An idle window: what the server sends with no writes, between two scrapes. + if jq -e '.idle_seconds > 0' <<< "$spec" >/dev/null; then + phase=-idle; scrape_metrics || { echo "$id: metrics scrape failed" >&2; exit 4; }; phase= + sleep "$(jq -r .idle_seconds <<< "$spec")" + fi + if jq -e '.restart' <<< "$spec" >/dev/null; then restart; fi } scrub() { @@ -160,6 +467,7 @@ scrub() { case "${1:-}" in check-runtime) check_runtime "${2:?}" ;; + check-images) check_images "${2:?}" ;; pull) pull_images ;; wal-disk) wal_disk "${2:?}" ;; generate) generate ;; @@ -175,5 +483,5 @@ case "${1:-}" in done < <(jq -c '.cells[]' "$PLAN") exit "$failed" ;; scrub) scrub ;; - *) echo "usage: $0 check-runtime IMAGE | pull | wal-disk DIR | generate | cells | scrub" >&2; exit 2 ;; + *) echo "usage: $0 check-runtime IMAGE | check-images JSON | pull | wal-disk DIR | generate | cells | scrub" >&2; exit 2 ;; esac diff --git a/.github/scripts/oswal-s3-bench-summary.py b/.github/scripts/oswal-s3-bench-summary.py index 67186d422da..e693b5f63db 100644 --- a/.github/scripts/oswal-s3-bench-summary.py +++ b/.github/scripts/oswal-s3-bench-summary.py @@ -26,14 +26,30 @@ from pathlib import Path PREFIX = "greptime_logstore_object_store_wal_" COUNTERS = ("created_objects_total", "stalled_appends_total", "poisoned_total", "create_failures_total", "create_conflicts_total") +# The Raft Engine controls every object store WAL cell is compared with: fsync per write and the default. +BASELINES = ("R", "R0") SAMPLE = re.compile(r'([a-zA-Z_:][a-zA-Z0-9_:]*)(?:\{([^}]*)\})?\s+(\S+)') +# Mito's line for every region it opened, with the WAL range it replayed. +REPLAY = re.compile(r"Replay WAL for region: (\d+)\(.*?provider: (\w+).*?rows recovered: (\d+), " + r"replay from entry id: (\d+), last entry id: (\d+),.*?elapsed: ([0-9.]+)(ns|µs|us|ms|s)\b") +UNITS = {"ns": 1e-9, "µs": 1e-6, "us": 1e-6, "ms": 1e-3, "s": 1.0} +# An object store WAL entry id is the object sequence above a 20-bit position. +POSITION_BITS = 20 def parse_metrics(text: str) -> dict: - """Map each unlabeled sample to its value and each histogram to its buckets.""" - samples, buckets = {}, {} + """Map each unlabeled sample to its value, each histogram to its buckets, and the S3 + operations (data files and WAL together, as OpenDAL counts its calls: a paginated list or a + multipart write is one operation over several HTTP requests) to their count per kind; the local cache + stores report through the same layer under another scheme and are left out.""" + samples, buckets, operations = {}, {}, {} for line in text.splitlines(): match = SAMPLE.fullmatch(line.strip()) + if match and match.group(1) == "opendal_operation_duration_seconds_count": + operation = re.search(r'operation="([^"]+)"', match.group(2) or "") + if operation and re.search(r'scheme="s3"', match.group(2) or ""): + operations[operation.group(1)] = operations.get(operation.group(1), 0) + float(match.group(3)) + continue if not match or not match.group(1).startswith(PREFIX): continue name, labels, value = match.group(1)[len(PREFIX):], match.group(2) or "", float(match.group(3)) @@ -42,7 +58,8 @@ def parse_metrics(text: str) -> dict: buckets.setdefault(name[:-len("_bucket")], []).append((float(le.group(1)), value)) else: samples[name] = value - return {"samples": samples, "buckets": {name: sorted(items) for name, items in buckets.items()}} + return {"samples": samples, "buckets": {name: sorted(items) for name, items in buckets.items()}, + "operations": operations} def quantile(q: float, buckets: list) -> float | None: @@ -64,25 +81,151 @@ def geomean(values: list) -> float | None: return math.exp(sum(map(math.log, values)) / len(values)) if values and all(v > 0 for v in values) else None +def read_metrics(path: Path) -> dict: + return parse_metrics(path.read_text() if path.exists() else "") + + +def ratio(a, b): + return a / b if a is not None and b else None + + +def replay(log: Path, metrics: dict) -> dict | None: + """What a restarted server replayed: every region it opened, the span of object sequences + its WAL range covers, and the log store's time serving region reads (`read_seconds`). The span + bounds the objects a region read from above: objects without its entries are skipped.""" + if not log.exists(): + return None + regions = [] + for match in REPLAY.finditer(log.read_text(errors="replace")): + region, provider, rows, first, last, elapsed, unit = match.groups() + # The range runs from the first entry to replay to the last one; it is empty when first > last. + span = None if provider != "ObjectStore" else 0 if int(first) > int(last) else \ + (int(last) >> POSITION_BITS) - (int(first) >> POSITION_BITS) + 1 + regions.append({"region": int(region), "rows": int(rows), "first_entry_id": int(first), + "last_entry_id": int(last), "seconds": float(elapsed) * UNITS[unit], "object_span": span}) + samples = metrics["samples"] + return {"regions": regions, "rows": sum(r["rows"] for r in regions), + "seconds": max((r["seconds"] for r in regions), default=None), + "object_span": sum(r["object_span"] for r in regions) + if regions and regions[0]["object_span"] is not None else None, + "read_seconds": samples.get("read_seconds_sum"), "reads": samples.get("read_seconds_count")} + + +def replayed(regions: list) -> str: + """Each region that replayed rows: its rows, the object sequences its range spans and its time.""" + return "; ".join(f"{r['region']}: {r['rows']:,} rows" + ( + f", object sequences {r['first_entry_id'] >> POSITION_BITS:,}–{r['last_entry_id'] >> POSITION_BITS:,}" + if r["object_span"] else "") + f", {r['seconds']:.2f} s" for r in regions if r["rows"]) or "—" + + +def client_latency(path: Path) -> tuple: + """The 50th and 99th percentile of the loader's own time per write request.""" + durations = [] + if path.exists(): + for line in path.read_text().splitlines(): + record = json.loads(line) if line.strip() else {} + if record.get("step") == "grpc_write_batch" and record.get("status") == "success": + durations.append(record["duration_seconds"]) + durations.sort() + return tuple(durations[max(0, math.ceil(q * len(durations)) - 1)] if durations else None for q in (0.5, 0.99)) + + +def region_rows(path: Path) -> list | None: + """Rows per region of a partitioned load cell, as the server reported them after the load.""" + if not path.exists(): + return None + return [int(row[2] or 0) for row in json.loads(path.read_text())["output"][0]["records"]["rows"]] + + +def interval_seconds(wal: dict) -> float: + value = wal.get("FLUSH_INTERVAL", "100ms") + return float(value[:-2]) / 1000 if value.endswith("ms") else float(value.rstrip("s")) + + +def traffic(end: dict, start: dict, seconds, wal: dict) -> dict: + """WAL traffic between two scrapes of one process over `seconds`; more objects than seal + ticks, or objects near the 8 MiB threshold, mean batches sealed on size rather than on the timer.""" + def delta(name): + value = end["samples"].get(name) + return None if value is None else value - start["samples"].get(name, 0) + created, written, objects = delta("created_objects_total"), delta("object_bytes_sum"), delta("object_bytes_count") + per_second = created / seconds if created and seconds else None + mean = written / objects if written and objects else None + return {"wal_mb_per_second": written / seconds / 1e6 if written and seconds else None, + "objects_per_second": per_second, "mean_object_bytes": mean, + "sealed_on_size": None if per_second is None else + per_second > 1.05 / interval_seconds(wal) or (mean or 0) >= 0.9 * 8 * 2**20, + "operations": {op: n - start["operations"].get(op, 0) for op, n in end["operations"].items() + if n != start["operations"].get(op, 0)}} + + +def crash_cycles(path: Path, wal: dict) -> list: + """Each cycle's check plus the loaded server's metrics before the kill and the restarted one's after.""" + cycles = [json.loads(line) for line in (path / "cycles.jsonl").read_text().splitlines() + if line.strip()] if (path / "cycles.jsonl").exists() else [] + for cycle in cycles: + start = read_metrics(path / f"metrics-start-{cycle['cycle']}.prom") + killed = read_metrics(path / f"metrics-cycle-{cycle['cycle']}.prom") + before = killed["samples"] + after = read_metrics(path / f"metrics-restart-{cycle['cycle']}.prom") + once = after["samples"].get("recovery_seconds_count") == 1 + cycle |= traffic(killed, start, cycle.get("load_seconds"), wal) + cycle |= {"created_before_kill": before.get("created_objects_total"), + "indexed_before_kill": before.get("indexed_objects"), + "recovery_seconds": after["samples"].get("recovery_seconds_sum") if once else None, + "recovered_objects": after["samples"].get("recovered_objects_total", 0 if once else None), + "restart_operations": after["operations"], + "replay": replay(path / f"server-restart-{cycle['cycle']}.log", after)} + return cycles + + def load_cell(path: Path, spec: dict | None = None) -> dict: """One cell's row; a planned cell that never recorded a result has no values, so no ratio uses its round.""" recorded = path / "cell.json" cell = json.loads(recorded.read_text()) if recorded.exists() else spec | {"exit_code": None, "missing": True} result_path = path / "greptimedb-1" / "load-result.json" result = json.loads(result_path.read_text()) if result_path.exists() else {} - metrics_path = path / "metrics.prom" - metrics = parse_metrics(metrics_path.read_text() if metrics_path.exists() else "") - ack = metrics["buckets"].get("append_ack_seconds", []) + final, restarted = read_metrics(path / "metrics.prom"), read_metrics(path / "metrics-restart.prom") + # With an idle window the scrape taken when the load ended describes the load. + metrics = read_metrics(path / "metrics-idle.prom") if cell.get("idle_seconds") else final load_seconds = result.get("load_seconds") + wall, cpu = cell.get("load_wall_seconds"), cell.get("server_cpu_seconds") row = { - "id": cell["id"], "round": cell["round"], "concurrency": cell["concurrency"], + "id": cell["id"], "group": cell.get("group", ""), "mode": cell.get("mode", "load"), + "image": cell.get("image", ""), "warmup": cell.get("warmup", False), + "round": cell["round"], "concurrency": cell["concurrency"], "config": cell["config"], "exit_code": cell["exit_code"], "wal_device": cell.get("wal_device"), - "missing": cell.get("missing", False), - "version": cell.get("version"), "load_wall_seconds": cell.get("load_wall_seconds"), + "missing": cell.get("missing", False), "env": cell.get("env", {}), + "version": cell.get("version"), "load_wall_seconds": wall, "rows": result.get("count_after"), "load_seconds": load_seconds, "rows_per_second": result.get("rows_per_second") if cell["exit_code"] == 0 else None, - "append_ack_p50_seconds": quantile(0.5, ack), "append_ack_p99_seconds": quantile(0.99, ack), + "server_cpu_cores": cpu / wall if cpu is not None and wall else None, + "host_cpu_busy": cell.get("host_cpu_busy"), "host_cpus": cell.get("host_cpus"), + "host_load_start": cell.get("host_load_start"), "host_load_end": cell.get("host_load_end"), } + for name, key in (("append_ack_seconds", "append_ack"), ("seal_to_durable_seconds", "seal_to_durable"), + ("object_bytes", "object_bytes")): + buckets = metrics["buckets"].get(name, []) + row[f"{key}_p50"], row[f"{key}_p99"] = quantile(0.5, buckets), quantile(0.99, buckets) + row["client_p50"], row["client_p99"] = client_latency(path / "greptimedb-1" / "load-transcript.jsonl") + # An entry is one region's write; an append, counted only in the durable mode, carries one or more. + samples = metrics["samples"] + row["entries_per_object"] = ratio(samples.get("object_entries_sum"), samples.get("object_entries_count")) + row["appends_per_object"] = ratio(samples.get("append_ack_seconds_count"), samples.get("object_bytes_count")) + row["partitions"] = cell.get("partitions", 1) + row["region_rows"] = region_rows(path / "regions.json") + if cell.get("restart"): + # A restarted process observed exactly one recovery, so the sum is its duration. + count = restarted["samples"].get("recovery_seconds_count") + row["restart"] = { + "wall_seconds": cell.get("restart_seconds"), "query_seconds": cell.get("restart_query_seconds"), + "rows": cell.get("restart_count"), + "recovery_seconds": restarted["samples"].get("recovery_seconds_sum") if count == 1 else None, + "recovered_objects": restarted["samples"].get("recovered_objects_total") + if cell.get("restart_scraped") else None, + "indexed_before_kill": cell.get("indexed_before_kill"), "min_objects": cell.get("min_objects", 0), + "replay": replay(path / "server-restart.log", restarted), + } # The counters register on their first event, so a scrape of an object store WAL # that created objects omits the ones that never moved: those are zero, not unknown. active = (cell.get("metrics_scraped") and cell.get("wal", {}).get("PROVIDER") == "experimental_object_store" @@ -91,49 +234,262 @@ def load_cell(path: Path, spec: dict | None = None) -> dict: row[counter] = metrics["samples"].get(counter, 0 if active else None) created = row["created_objects_total"] row["creates_per_minute"] = created * 60 / load_seconds if created is not None and load_seconds else None + row |= traffic(metrics, {"samples": {}, "operations": {}}, load_seconds, cell.get("wal", {})) + row["restart_operations"] = restarted["operations"] + if cell.get("idle_seconds"): + idle = traffic(final, metrics, cell["idle_seconds"], cell.get("wal", {})) + end = final["samples"].get("created_objects_total") + row["idle"] = {"seconds": cell["idle_seconds"], "operations": idle["operations"], + "creates": None if end is None else end - metrics["samples"].get("created_objects_total", 0)} + if row["mode"] == "drop": + written = read_metrics(path / "metrics-written.prom")["samples"] + restarted_samples = restarted["samples"] + samples = [line.split() for line in (path / "collected.tsv").read_text().splitlines() + if line.strip()] if (path / "collected.tsv").exists() else [] + row["drop"] = { + "acked_batches": sum(len((p.read_text()).split()) for p in sorted((path / "dropped").glob("acked-*.log"))), + "indexed_before_drop": written.get("indexed_objects"), + "restart_seconds": cell.get("restart_seconds"), + "recovered_objects": restarted_samples.get("recovered_objects_total"), + "indexed_at_end": int(float(samples[-1][1])) if samples else None, + "deleted_after_restart": restarted_samples.get("deleted_objects_total"), + "seconds_to_collect": int(samples[-1][0]) - int(samples[0][0]) if samples else None, + } + if row["mode"] == "crash": + row["cycles"] = crash_cycles(path, cell.get("wal", {})) + if (path / "objects.tsv").exists(): + # A sample is time, indexed, deleted, created; a failed scrape is time and "missing". + row["objects_samples"] = [[int(float(line.split()[0]))] + ( + [None] * 3 if line.split()[1:] == ["missing"] else [int(float(v)) for v in line.split()[1:]]) + for line in (path / "objects.tsv").read_text().splitlines() if line.strip()] return row def ratios(rows: list) -> list: - """D/R and E/R per round and their geometric mean over the rounds, per concurrency.""" - speed = {(r["round"], r["concurrency"], r["config"]): r["rows_per_second"] for r in rows} + """Each object store WAL cell over each Raft Engine control of the same image per round, and their + geometric mean over the rounds, per concurrency; warm-up cells are left out.""" + rows = [r for r in rows if r["mode"] == "load" and not r["warmup"]] + speed = {(r["group"], r["image"], r["round"], r["concurrency"], r["config"]): r["rows_per_second"] + for r in rows} result = [] - for concurrency in sorted({r["concurrency"] for r in rows}): - rounds = sorted({r["round"] for r in rows if r["concurrency"] == concurrency}) - for config in ("D", "E"): - per_round = [] - for round_number in rounds: - base, value = speed.get((round_number, concurrency, "R")), speed.get((round_number, concurrency, config)) - per_round.append(value / base if base and value else None) - complete = [value for value in per_round if value is not None] - result.append({"concurrency": concurrency, "ratio": f"{config}/R", "per_round": per_round, - "geomean": geomean(complete) if len(complete) == len(per_round) else None}) + for group, image in dict.fromkeys((r["group"], r["image"]) for r in rows): + members = [r for r in rows if (r["group"], r["image"]) == (group, image)] + series = sorted({r["config"] for r in members if r["config"] not in BASELINES}) + baselines = [b for b in BASELINES if any(r["config"] == b for r in members)] + for concurrency in sorted({r["concurrency"] for r in members}): + rounds = sorted({r["round"] for r in members if r["concurrency"] == concurrency}) + for config, baseline in ((s, b) for s in series for b in baselines): + per_round = [] + for round_number in rounds: + base = speed.get((group, image, round_number, concurrency, baseline)) + measured = speed.get((group, image, round_number, concurrency, config)) + per_round.append(measured / base if base and measured else None) + complete = [ratio for ratio in per_round if ratio is not None] + result.append({"group": group, "image": image, "concurrency": concurrency, + "ratio": f"{config}/{baseline}", "per_round": per_round, + "geomean": geomean(complete) if len(complete) == len(per_round) else None}) return result +def image_ratios(rows: list, images: list) -> list: + """Every load cell of every image per round, and each image's rows/s over the first image's in the + same round, with their geometric mean, minimum and maximum over the rounds; warm-up cells are left out.""" + rows = [r for r in rows if r["mode"] == "load" and not r["warmup"]] + cells = {(r["group"], r["image"], r["round"], r["concurrency"], r["config"]): r for r in rows} + result = [] + for group, concurrency, config in dict.fromkeys((r["group"], r["concurrency"], r["config"]) for r in rows): + rounds = sorted({r["round"] for r in rows if (r["group"], r["concurrency"], r["config"]) == (group, concurrency, config)}) + for image in images: + measured = [cells.get((group, image, n, concurrency, config)) or {} for n in rounds] + reference = [cells.get((group, images[0], n, concurrency, config)) or {} for n in rounds] + per_round = [ratio(m.get("rows_per_second"), r.get("rows_per_second")) for m, r in zip(measured, reference)] + complete = image != images[0] and None not in per_round + result.append({ + "group": group, "concurrency": concurrency, "config": config, "image": image, + "rounds": [{"round": n, "id": m.get("id"), "rows_per_second": m.get("rows_per_second"), + "client_p50": m.get("client_p50"), "client_p99": m.get("client_p99"), + "host_load_start": m.get("host_load_start"), "host_load_end": m.get("host_load_end")} + for n, m in zip(rounds, measured)], + "ratio_to": None if image == images[0] else images[0], + "per_round": None if image == images[0] else per_round, + "geomean": geomean(per_round) if complete else None, + "min": min(per_round) if complete else None, "max": max(per_round) if complete else None}) + return result + + +def operations(counts: dict) -> str: + return ", ".join(f"{op} {n:,.0f}" for op, n in sorted(counts.items())) or "—" + + +def overrides(env: dict) -> str: + return " ".join(f"{k}={v}" for k, v in sorted(env.items())) or "—" + + +def pair(a, b, digits=0, scale=1.0) -> str: + return f"{fmt(a, digits, scale)} / {fmt(b, digits, scale)}" + + def fmt(value, digits=0, scale=1.0) -> str: return "—" if value is None else f"{value * scale:,.{digits}f}" -def render(run: dict, rows: list, pairs: list) -> str: +def load_average(value) -> str: + """The 1-minute load average of a recorded [1, 5, 15] minute triple.""" + return "—" if not value else f"{value[0]:.2f}" + + +def render(run: dict, rows: list, pairs: list, compared: list = ()) -> str: lines = [f"## Object store WAL against S3 ({'smoke, MinIO' if run.get('smoke') else run.get('instance_type')})", ""] - lines += [f"GreptimeDB `{run.get('greptimedb_image')}` ({next((r['version'] for r in rows if r['version']), '—')}); " - f"{fmt(run.get('rows'))} rows in {run.get('batch_rows')}-row batches; bucket region {run.get('region')}.", ""] - lines += ["| Cell | Rows/s | Load s | Wall s | Rows | Ack p50 ms | Ack p99 ms | Creates | Creates/min " - "| Stalled | Poisoned | Create failures | Conflicts | WAL on |", - "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |"] - for r in rows: - cell = f"{r['id']} (not run)" if r["missing"] else r["id"] if r["exit_code"] == 0 \ - else f"{r['id']} (failed, exit {r['exit_code']})" + images = run.get("greptimedb_images") or [] + if len(images) > 1: + lines += [f"{fmt(run.get('rows'))} rows in {run.get('batch_rows')}-row batches; bucket region {run.get('region')}; " + f"each image is compared with the first, `{images[0]['label']}`.", "", + "| Image | Reference | Version | Commit |", "| --- | --- | --- | --- |"] + lines += [f"| {i['label']} | `{i['image']}` | {i['version']} | `{i['commit']}` |" for i in images] + [""] + else: + lines += [f"GreptimeDB `{run.get('greptimedb_image')}` ({next((r['version'] for r in rows if r['version']), '—')}); " + f"{fmt(run.get('rows'))} rows in {run.get('batch_rows')}-row batches; bucket region {run.get('region')}.", ""] + lines += ["| Cell | Overrides | Rows/s | Load s | Wall s | Ack p50 / p99 ms | Seal to durable p50 / p99 ms " + "| Object p50 / p99 MiB | Creates | Creates/min | Stalled | Poisoned | Create failures | Conflicts " + "| Server CPUs | Host busy | WAL on | Load average at start / end |", + "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: " + "| ---: | ---: | --- | ---: |"] + def label(r): + cell = f"{r['id']} (warm-up)" if r["warmup"] else r["id"] + return f"{cell} (not run)" if r["missing"] else cell if r["exit_code"] == 0 \ + else f"{cell} (failed, exit {r['exit_code']})" + loads = [r for r in rows if r["mode"] == "load"] + for r in loads: + cell = label(r) lines.append(" | ".join([ - f"| {cell}", fmt(r["rows_per_second"]), fmt(r["load_seconds"], 1), fmt(r["load_wall_seconds"], 1), - fmt(r["rows"]), fmt(r["append_ack_p50_seconds"], 1, 1000), fmt(r["append_ack_p99_seconds"], 1, 1000), + f"| {cell}", overrides(r["env"]), fmt(r["rows_per_second"]), fmt(r["load_seconds"], 1), + fmt(r["load_wall_seconds"], 1), pair(r["append_ack_p50"], r["append_ack_p99"], 1, 1000), + pair(r["seal_to_durable_p50"], r["seal_to_durable_p99"], 1, 1000), + pair(r["object_bytes_p50"], r["object_bytes_p99"], 1, 1 / 2**20), *(fmt(r[c]) for c in COUNTERS[:1]), fmt(r["creates_per_minute"], 1), - *(fmt(r[c]) for c in COUNTERS[1:]), f"`{r['wal_device'] or '—'}` |"])) - lines += ["", "| Concurrency | Ratio | Per round | Geometric mean |", "| ---: | --- | --- | ---: |"] + *(fmt(r[c]) for c in COUNTERS[1:]), fmt(r["server_cpu_cores"], 2), + "—" if r["host_cpu_busy"] is None else f"{r['host_cpu_busy']:.0%} of {r['host_cpus'] or '?'}", + f"`{r['wal_device'] or '—'}`", + f"{load_average(r['host_load_start'])} / {load_average(r['host_load_end'])} |"])) + lines += ["", "| Cell | WAL MB/s | Objects/s | Mean object MiB | Sealed on size | OpenDAL S3 operations during the load " + "| OpenDAL S3 operations of the restart | Idle s | Idle creates | Idle OpenDAL S3 operations |", + "| --- | ---: | ---: | ---: | --- | --- | --- | ---: | ---: | --- |"] + for r in loads: + idle = r.get("idle") or {} + lines.append(" | ".join([ + f"| {label(r)}", fmt(r["wal_mb_per_second"], 1), fmt(r["objects_per_second"], 1), + fmt(r["mean_object_bytes"], 2, 1 / 2**20), + "—" if r["sealed_on_size"] is None else ("yes" if r["sealed_on_size"] else "no"), + operations(r["operations"]), operations(r["restart_operations"]) if "restart" in r else "—", + fmt(idle.get("seconds")), fmt(idle.get("creates")), operations(idle["operations"]) if idle else "—"]) + " |") + crashes = [r for r in rows if r["mode"] == "crash"] + if crashes: + lines += ["", "| Cell | Cycle | Kill after s | Acked | Present | Lost | Lost rows | Extra | In flight " + "| All found | All acked | All lost | Indexed at kill | Restart to /health s | Recovery s " + "| Recovered objects | Pin rows | Verdict |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: " + "| ---: | ---: | ---: | --- |"] + for r in crashes: + for c in r["cycles"] or [{}]: + verdict = "—" if not c else "FAIL: " + "; ".join(c.get("kill_errors", []) + c.get("verify_errors", [])) \ + if not c.get("passed") else "pass" if c.get("metrics_complete", True) else "pass, scrape missing" + lines.append(" | ".join([ + f"| {label(r)}", fmt(c.get("cycle")), fmt(c.get("kill_after_seconds"), 1), + fmt(c.get("cycle_acked_batches")), fmt(c.get("cycle_present_batches")), + fmt(c.get("cycle_lost_batches")), fmt(c.get("cycle_lost_rows")), fmt(c.get("cycle_extra_batches")), + fmt(c.get("in_flight")), fmt(c.get("total_found_batches")), fmt(c.get("total_acked_batches")), + fmt(c.get("total_lost_batches")), fmt(c.get("indexed_before_kill")), + fmt(c.get("restart_seconds"), 2), fmt(c.get("recovery_seconds"), 3), + fmt(c.get("recovered_objects")), fmt(c.get("pin_rows")), verdict]) + " |") + lines += ["", "| Cell | Cycle | Load s | WAL MB/s | Objects/s | Mean object MiB | Sealed on size " + "| OpenDAL S3 operations until the kill | OpenDAL S3 operations of the restart |", + "| --- | ---: | ---: | ---: | ---: | ---: | --- | --- | --- |"] + for r in crashes: + for c in r["cycles"]: + lines.append(" | ".join([ + f"| {r['id']}", fmt(c["cycle"]), fmt(c.get("load_seconds"), 1), fmt(c["wal_mb_per_second"], 1), + fmt(c["objects_per_second"], 1), fmt(c["mean_object_bytes"], 2, 1 / 2**20), + "—" if c["sealed_on_size"] is None else ("yes" if c["sealed_on_size"] else "no"), + operations(c["operations"]), operations(c["restart_operations"])]) + " |") + pinned = [r for r in crashes if r.get("objects_samples")] + if pinned: + lines += ["", "| Cell | Indexed objects every 30 s | Missing samples | Deleted objects at the end " + "| Created objects at the end |", "| --- | --- | ---: | ---: | ---: |"] + for r in pinned: + samples = r["objects_samples"] + last = next((s for s in reversed(samples) if s[1] is not None), [None] * 4) + lines.append(f"| {r['id']} | {', '.join('missing' if s[1] is None else str(s[1]) for s in samples)} " + f"| {sum(s[1] is None for s in samples)} | {fmt(last[2])} | {fmt(last[3])} |") + details = [r for r in loads if any(r[k] is not None for k in ("client_p50", "entries_per_object", "region_rows"))] + if details: + lines += ["", "| Cell | Client p50 / p99 ms | Entries per object | Appends per object | Regions " + "| Rows per region min / max |", "| --- | ---: | ---: | ---: | ---: | ---: |"] + for r in details: + spread = r["region_rows"] or [] + lines.append(" | ".join([ + f"| {label(r)}", pair(r["client_p50"], r["client_p99"], 1, 1000), fmt(r["entries_per_object"], 1), + fmt(r["appends_per_object"], 2), fmt(len(spread) if spread else r["partitions"]), + pair(min(spread), max(spread)) if spread else "—"]) + " |") + replays = [(r["id"], "—", r["restart"]) for r in loads if r.get("restart", {}).get("replay")] + replays += [(r["id"], fmt(c["cycle"]), c | {"wall_seconds": c.get("restart_seconds")}) + for r in crashes for c in r["cycles"] if c.get("replay")] + if replays: + lines += ["", "| Cell | Cycle | Restart to /health s | Recovery s | Recovered objects | Regions opened " + "| Rows replayed | Object span | Slowest region replay s | Read seconds (reads) " + "| Regions with rows |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- | --- |"] + for cell, cycle, x in replays: + y = x["replay"] + lines.append(" | ".join([ + f"| {cell}", cycle, fmt(x.get("wall_seconds"), 2), fmt(x.get("recovery_seconds"), 3), + fmt(x.get("recovered_objects")), fmt(len(y["regions"])), fmt(y["rows"]), fmt(y["object_span"]), + fmt(y["seconds"], 2), f"{fmt(y['read_seconds'], 2)} ({fmt(y['reads'])})", + replayed(y["regions"])]) + " |") + drops = [r for r in rows if r.get("drop")] + if drops: + lines += ["", "| Cell | Acked batches before the drop | Indexed before the drop | Restart to /health s " + "| Recovered objects | Indexed at the end | Deleted after the restart | Seconds to collect | Verdict |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |"] + for r in drops: + x = r["drop"] + # Objects the release already collected before the kill leave nothing to show. + verdict = "—" if x["indexed_at_end"] is None else \ + "nothing left at the kill" if (x["recovered_objects"] or 0) <= 3 else \ + "collected" if r["exit_code"] == 0 else "still indexed" + lines.append(" | ".join([ + f"| {label(r)}", fmt(x["acked_batches"]), fmt(x["indexed_before_drop"]), fmt(x["restart_seconds"], 2), + fmt(x["recovered_objects"]), fmt(x["indexed_at_end"]), + fmt(x["deleted_after_restart"]), fmt(x["seconds_to_collect"]), verdict]) + " |") + restarts = [r for r in rows if "restart" in r] + if restarts: + lines += ["", "| Cell | Overrides | Objects created | Indexed at kill | Restart to /health s | Restart to count s " + "| Recovery s | Recovered objects | Rows after restart |", + "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |"] + for r in restarts: + x = r["restart"] + short = x["min_objects"] and (x["indexed_before_kill"] or 0) < x["min_objects"] + lines.append(" | ".join([ + f"| {label(r)}", overrides(r["env"]), fmt(r["created_objects_total"]), + fmt(x["indexed_before_kill"]) + (f" (under {x['min_objects']:,})" if short else ""), + fmt(x["wall_seconds"], 2), + fmt(x["query_seconds"], 2), fmt(x["recovery_seconds"], 3), fmt(x["recovered_objects"]), + f"{fmt(x['rows'])} |"])) + lines += ["", "| Group | Image | Concurrency | Ratio | Per round | Geometric mean |", + "| --- | --- | ---: | --- | --- | ---: |"] for p in pairs: - lines.append(f"| {p['concurrency']} | {p['ratio']} | {', '.join(fmt(v, 3) for v in p['per_round'])} " - f"| {fmt(p['geomean'], 3)} |") + lines.append(f"| {p['group'] or '—'} | {p['image'] or '—'} | {p['concurrency']} | {p['ratio']} " + f"| {', '.join(fmt(v, 3) for v in p['per_round'])} | {fmt(p['geomean'], 3)} |") + if compared: + lines += ["", f"| Group | Concurrency | Config | Image | Rows/s per round | Client p50 / p99 ms per round " + f"| Ratio to {compared[0]['image']} per round | Geometric mean | Min | Max |", + "| --- | ---: | --- | --- | --- | --- | --- | ---: | ---: | ---: |"] + for c in compared: + lines.append(" | ".join([ + f"| {c['group'] or '—'}", str(c["concurrency"]), c["config"], c["image"], + "; ".join(fmt(r["rows_per_second"]) for r in c["rounds"]), + "; ".join(pair(r["client_p50"], r["client_p99"], 1, 1000) for r in c["rounds"]), + "—" if c["per_round"] is None else ", ".join(fmt(v, 3) for v in c["per_round"]), + fmt(c["geomean"], 3), fmt(c["min"], 3), fmt(c["max"], 3)]) + " |") return "\n".join(lines) + "\n" @@ -145,17 +501,20 @@ def main() -> None: run_path = args.root / "run.json" run = json.loads(run_path.read_text()) if run_path.exists() else {} plan_path = args.root / "plan.json" - planned = json.loads(plan_path.read_text())["cells"] if plan_path.exists() else [] + plan = json.loads(plan_path.read_text()) if plan_path.exists() else {} + planned = plan.get("cells", []) # The plan, not what happened to run, defines the expected cells and rounds. rows = [load_cell(args.root / "cells" / spec["id"], spec) for spec in planned] rows += [load_cell(path) for path in sorted((args.root / "cells").glob("*")) if path.name not in {spec["id"] for spec in planned} and (path / "cell.json").exists()] pairs = ratios(rows) + images = plan.get("images", [""]) + compared = image_ratios(rows, images) if len(images) > 1 else [] missing = [r["id"] for r in rows if r["missing"]] (args.root / "summary.json").write_text(json.dumps( - {"run": run, "planned_cells": len(planned), "missing_cells": missing, "cells": rows, "ratios": pairs}, - indent=2) + "\n") - print(render(run, rows, pairs) if rows else "No benchmark cell completed.\n") + {"run": run, "planned_cells": len(planned), "missing_cells": missing, "cells": rows, "ratios": pairs, + "image_ratios": compared}, indent=2) + "\n") + print(render(run, rows, pairs, compared) if rows else "No benchmark cell completed.\n") if __name__ == "__main__": diff --git a/.github/scripts/oswal-s3-bench.py b/.github/scripts/oswal-s3-bench.py index faad8ace674..eedd906ffef 100644 --- a/.github/scripts/oswal-s3-bench.py +++ b/.github/scripts/oswal-s3-bench.py @@ -24,26 +24,84 @@ from __future__ import annotations import argparse import json +import math import os import re import sys import time import urllib.request +from datetime import datetime from pathlib import Path SMOKE_ROWS = 50_000 # The o11ybench data profiles; a matching row count generates the canonical dataset. DATA_PROFILES = {5_000: "S", 2_000_000: "D", 10_000_000: "P", 100_000_000: "M"} -# Identical [storage] for every configuration; only [wal] differs. +# Identical [storage] for every configuration; only [wal] differs. R0 is the standalone default. WAL_CONFIGS = { "R": {"PROVIDER": "raft_engine", "SYNC_WRITE": "true", "DIR": "/wal"}, + "R0": {"PROVIDER": "raft_engine", "SYNC_WRITE": "false", "DIR": "/wal"}, "D": {"PROVIDER": "experimental_object_store", "ACK_MODE": "durable"}, "E": {"PROVIDER": "experimental_object_store", "ACK_MODE": "enqueued"}, } -ORDERS = (("R", "D", "E"), ("E", "D", "R")) +BASELINES = ("R", "R0") +# Restart cells seal ten times as often as the default so recovery has thousands of objects. +RESTART_FLUSH_INTERVAL = "10ms" +# The loader's CREATE TABLE for the agent observability workload at the pinned runtime image. +AGENT_TABLE = """CREATE TABLE `agent_observations` ( + `event_time` TIMESTAMP(3) NOT NULL, + `biz_date` DATE NOT NULL, + `trace_id` STRING NOT NULL SKIPPING INDEX, + `session_id` STRING NOT NULL, + `observation_id` STRING NOT NULL SKIPPING INDEX, + `parent_observation_id` STRING NULL SKIPPING INDEX, + `seq_no` BIGINT NOT NULL, + `type` STRING NOT NULL INVERTED INDEX, + `status` STRING NOT NULL INVERTED INDEX, + `tenant` STRING NOT NULL, + `app` STRING NOT NULL, + `environment` STRING NOT NULL, + `task_category` STRING NULL, + `trace_archetype` STRING NULL, + `model` STRING NULL, + `tool_name` STRING NULL, + `input` STRING NULL FULLTEXT INDEX WITH(analyzer = 'English', backend = 'bloom', case_sensitive = 'false', granularity = '10240', false_positive_rate = '0.01'), + `output` STRING NULL FULLTEXT INDEX WITH(analyzer = 'English', backend = 'bloom', case_sensitive = 'false', granularity = '10240', false_positive_rate = '0.01'), + `input_tokens` BIGINT NULL, + `output_tokens` BIGINT NULL, + `total_cost` DOUBLE NULL, + `latency_ms` BIGINT NULL, + `payload` JSON2 NOT NULL, + TIME INDEX (`event_time`), + PRIMARY KEY (`tenant`, `app`, `environment`) +) +{partitions}ENGINE=mito +WITH( + append_mode = 'true', + 'compaction.type' = 'twcs', + 'compaction.twcs.time_window' = '2h' +)""" +GROUP_NAME = re.compile(r"[a-z][a-z0-9]{0,11}") +IMAGE_LABEL = re.compile(r"[a-z][a-z0-9]{0,7}") +RESERVED_ENV = ("GREPTIMEDB_STANDALONE__STORAGE__", "GREPTIMEDB_STANDALONE__WAL__") +ENV_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") +# Docker's reference grammar: path components separated by ".", "_", "__" or runs of "-". +IMAGE = re.compile(r"[a-z0-9]+(?:(?:[._]|__|-+)[a-z0-9]+)*(?::[0-9]+)?(?:/[a-z0-9]+(?:(?:[._]|__|-+)[a-z0-9]+)*)+" + r"(?::[A-Za-z0-9_][A-Za-z0-9_.-]{0,127})?(?:@sha256:[0-9a-f]{64})?") NIGHTLY = re.compile(r"v\d+\.\d+\.\d+(?:-[0-9A-Za-z.]+)?-nightly-(\d{8})") TAG = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}") DIGEST = re.compile(r".+@sha256:[0-9a-f]{64}") +# Rows/s of a load cell in 1000-row batches on c6id.2xlarge at the concurrency levels earlier runs +# measured, interpolated on a log scale between them; durable writes slow down with several regions. +# Smaller batches are taken to keep the request rate, so their rows/s shrink with the batch. +MEASURED_BATCH_ROWS = 1000 +LOAD_RATES = {"durable": [(1, 9_000), (8, 34_000), (32, 55_000)], + "durable, partitioned": [(1, 7_000), (8, 26_000), (32, 55_000)], + "other": [(1, 28_000), (8, 60_000)]} +# Starting, scraping and stopping a server, and the restart of a restart cell or a crash cycle. +CELL_SECONDS, RESTART_SECONDS, CYCLE_SECONDS = 5, 60, 15 +# Checkout, dependencies, image pulls and the bucket probe; corpus generation per million rows; and the +# deletion of the run's objects one at a time, which took 12-17 % of the cells' time in earlier runs. +JOB_SECONDS, GENERATE_SECONDS_PER_MILLION, CLEANUP_SHARE = 180, 30, 0.2 def append_github(variable: str, name: str, value: str) -> None: @@ -71,42 +129,228 @@ def check_regions(bucket: str, runner: str) -> None: "a cross-region bucket measures the WAN, not S3") -def build_manifest(plan: dict, greptimedb_inspect: list, runtime_image: str, +def build_manifest(plan: dict, images: list, runtime_image: str, instance_type: str, region: str, prefix: str) -> dict: - image = next((d for d in greptimedb_inspect[0].get("RepoDigests") or [] if DIGEST.fullmatch(d)), None) - if image is None: - raise ValueError("the GreptimeDB image has no sha256 repository digest") + """The run's record; `images` are the checked images, each pinned to its digest, the first the reference.""" + if [image["label"] for image in images] != plan["images"]: + raise ValueError(f"the checked images {[i['label'] for i in images]} are not the plan's {plan['images']}") + if not all(DIGEST.fullmatch(image["image"]) for image in images): + raise ValueError("every GreptimeDB image must be pinned to a sha256 digest") return {"smoke": plan["smoke"], "rows": plan["rows"], "rounds": plan["rounds"], "concurrency": plan["concurrency"], "batch_rows": plan["batch_rows"], - "greptimedb_image": image, "runtime_image": runtime_image, + "configs": plan["configs"], "extra_env": plan["extra_env"], + "restart": plan["restart"], "groups": plan["groups"], "warmup": plan["warmup"], + "greptimedb_image": images[0]["image"], "greptimedb_images": images, "runtime_image": runtime_image, "instance_type": instance_type, "region": region, "prefix": prefix} -def build_plan(concurrency: str, rows: str, batch_rows: str, rounds: str, smoke: bool) -> dict: +def parse_extra_env(text: str) -> dict: + """A JSON object or K=V lines; the storage and WAL settings that define a cell stay the benchmark's own.""" + if text.strip().startswith("{"): + parsed = json.loads(text) + if not all(isinstance(v, (str, int)) and not isinstance(v, bool) for v in parsed.values()): + raise ValueError("extra_env values must be strings or integers") + env = {k: str(v) for k, v in parsed.items()} + else: + env = {} + for line in filter(None, (line.strip() for line in text.splitlines())): + name, sep, value = line.partition("=") + if not sep or name.strip() in env: + raise ValueError(f"extra_env line is not a new K=V: {line!r}") + env[name.strip()] = value.strip() + for name, value in env.items(): + # GreptimeDB matches its environment keys without regard to case. + if not ENV_NAME.fullmatch(name) or name.upper().startswith(RESERVED_ENV) or "\n" in value: + raise ValueError(f"extra_env may not set {name!r}") + return env + + +def partitioned_table(summary: dict, partitions: int) -> str: + """The loader's table split into equal event_time ranges over the corpus window. The + generator draws event times uniformly from that window, so every range gets an equal + share of the rows and every batch writes to every region.""" + def millis(text): + return int(datetime.fromisoformat(text.replace("Z", "+00:00")).timestamp() * 1000) + start, end = millis(summary["start"]), millis(summary["end"]) + bounds = [start + (end - start) * i // partitions for i in range(1, partitions)] + if len(set(bounds)) != len(bounds) or not start < bounds[0]: + raise ValueError(f"the corpus window {summary['start']}..{summary['end']} is too short for {partitions} ranges") + ranges = [f"`event_time` < {bounds[0]}"] + ranges += [f"`event_time` >= {low} AND `event_time` < {high}" for low, high in zip(bounds, bounds[1:])] + ranges.append(f"`event_time` >= {bounds[-1]}") + clause = "PARTITION ON COLUMNS (`event_time`) (\n " + ",\n ".join(ranges) + "\n)\n" + return AGENT_TABLE.format(partitions=clause) + + +def parse_images(text: str) -> list: + """The images to compare: a JSON list of 2..4 {label, image}, the first the reference. Empty + means the one image of greptimedb_tag or greptimedb_image, with an empty label.""" + if not text.strip(): + return [{"label": "", "image": ""}] + images = json.loads(text) + if not isinstance(images, list) or not 2 <= len(images) <= 4 \ + or not all(isinstance(i, dict) and set(i) == {"label", "image"} for i in images): + raise ValueError("images must be a JSON list of 2..4 objects, each with only a label and an image") + labels = [image["label"] for image in images] + for image in images: + if not isinstance(image["label"], str) or not IMAGE_LABEL.fullmatch(image["label"]): + raise ValueError(f"an image label is up to 8 lowercase letters and digits, got {image['label']!r}") + full_reference("an image", image["image"]) + if len(set(labels)) != len(labels): + raise ValueError(f"image labels repeat: {labels}") + return images + + +def cell_id(group: str, round_number: int, level: int, config: str, label: str) -> str: + return (f"{group}-" if group else "") + f"{round_number}-{level}-{config}" + (f"-{label}" if label else "") + + +def group_cells(group: dict, smoke: bool, labels: list) -> list: + """The cells of one group, in run order; a group without a name is the top-level inputs. + Each cell runs once per image label, back to back.""" + name, mode = group.get("name", ""), group.get("mode", "load") + if mode not in ("load", "crash", "drop"): + raise ValueError(f"group {name!r}: mode must be load, crash or drop, got {mode!r}") + concurrency = str(group["concurrency"]) levels = [positive("concurrency", level.strip(), 256) for level in concurrency.split(",")] if len(set(levels)) != len(levels): raise ValueError(f"concurrency repeats a level: {concurrency!r}") - row_count = SMOKE_ROWS if smoke else positive("rows", rows, 100_000_000) - round_count = 1 if smoke else positive("rounds", rounds, 10) - profile = DATA_PROFILES.get(row_count) + chosen = [config.strip() for config in str(group["configs"]).split(",")] + allowed = set(WAL_CONFIGS) - (set(BASELINES) if mode != "load" else set()) + if not chosen or len(set(chosen)) != len(chosen) or not set(chosen) <= allowed: + raise ValueError(f"group {name!r}: configs must be distinct names among {', '.join(sorted(allowed))}, " + f"got {group['configs']!r}") + extra = group.get("extra_env", "") + env = parse_extra_env(json.dumps(extra) if isinstance(extra, dict) else str(extra)) + restart = group.get("restart", False) is True + rows = SMOKE_ROWS if smoke else positive("rows", str(group["rows"]), 100_000_000) + rounds = 1 if smoke else positive("rounds", str(group["rounds"]), 10) + idle = int(group.get("idle_seconds", 0)) + common = {"group": name, "mode": mode, "restart": restart, "idle_seconds": idle} + partitions = int(group.get("partitions", 1)) + if partitions != 1 and (mode != "load" or not 2 <= partitions <= 64): + raise ValueError(f"group {name!r}: partitions must be in 2..64 and needs mode load") + if mode == "load": + if not 0 <= idle <= 600: + raise ValueError(f"group {name!r}: idle_seconds must be in 0..600") + # The objects a restart cell must leave indexed at its kill; a smoke run is too small for any. + min_objects = int(group.get("min_objects", 0)) + if min_objects < 0 or (min_objects and not restart): + raise ValueError(f"group {name!r}: min_objects needs restart and must not be negative") + common |= {"rows": rows, "data_profile": DATA_PROFILES.get(rows, "custom"), + "min_objects": 0 if smoke else min_objects, "partitions": partitions} + elif mode == "drop": + # Writes one table for write_seconds, drops it, kills and restarts the server, then + # writes and flushes another table and watches the dropped table's objects go. + write_seconds = int(group.get("write_seconds", 30)) + if not 1 <= write_seconds <= 600 or restart or idle: + raise ValueError(f"group {name!r}: write_seconds must be in 1..600, without restart or idle") + common |= {"write_seconds": min(write_seconds, 10) if smoke else write_seconds} + else: + cycles = int(group.get("cycles", 5)) + kill_after = [int(v) for v in group.get("kill_after", [3, 10])] + if not 1 <= cycles <= 20 or len(kill_after) != 2 or not 1 <= kill_after[0] <= kill_after[1] <= 3600: + raise ValueError(f"group {name!r}: cycles must be in 1..20 and kill_after a [min, max] in 1..3600") + if restart or idle: + raise ValueError(f"group {name!r}: a crash group restarts after every kill and has no idle window") + if smoke: + cycles, kill_after = min(cycles, 2), [min(v, 60) for v in kill_after] + common |= {"cycles": cycles, "kill_after": kill_after, "pin": group.get("pin", False) is True} + series = [config for config in WAL_CONFIGS if config in chosen] cells = [] - for round_number in range(1, round_count + 1): - # Alternate the order between rounds so no configuration always runs first. + for round_number in range(1, rounds + 1): + # Reverse the order every other round so no configuration or image always runs first. + forward = round_number % 2 for level in levels: - for config in ORDERS[(round_number - 1) % 2]: - cells.append({ - "id": f"{round_number}-{level}-{config}", "round": round_number, - "concurrency": level, "config": config, "wal": WAL_CONFIGS[config], - }) + for config in series if forward else reversed(series): + wal = WAL_CONFIGS[config] | ({"FLUSH_INTERVAL": RESTART_FLUSH_INTERVAL} if restart and config not in BASELINES else {}) + for label in labels if forward else reversed(labels): + cells.append(common | { + "id": cell_id(name, round_number, level, config, label), "round": round_number, + "concurrency": level, "config": config, "image": label, "wal": wal, "env": env, + }) + return cells + + +def parse_groups(text: str, defaults: dict) -> list: + """A JSON list of groups; every key a group leaves out takes the top-level input's value.""" + if not text.strip(): + return [defaults] + groups = json.loads(text) + if not isinstance(groups, list) or not groups or not all(isinstance(g, dict) for g in groups): + raise ValueError("groups must be a non-empty JSON list of objects") + known = set(defaults) | {"name", "mode", "idle_seconds", "cycles", "kill_after", "pin", "min_objects", + "partitions", "write_seconds"} + names = [g.get("name") for g in groups] + for group, name in zip(groups, names): + if not isinstance(name, str) or not GROUP_NAME.fullmatch(name): + raise ValueError(f"every group needs a name of up to 12 lowercase letters and digits, got {name!r}") + if set(group) - known: + raise ValueError(f"group {name!r} has unknown keys {sorted(set(group) - known)}") + if len(set(names)) != len(names): + raise ValueError(f"group names repeat: {names}") + return [defaults | group for group in groups] + + +def build_plan(concurrency: str, rows: str, batch_rows: str, rounds: str, smoke: bool, + configs: str = "R,D,E", extra_env: str = "", restart: bool = False, groups: str = "", + images: str = "", warmup: bool = False) -> dict: + defaults = {"concurrency": concurrency, "rows": rows, "rounds": rounds, "configs": configs, + "extra_env": extra_env, "restart": restart} + resolved = parse_groups(groups, defaults) + labels = [image["label"] for image in parse_images(images)] + cells = [cell for group in resolved for cell in group_cells(group, smoke, labels)] + if warmup: + # Round 0 runs the first load cell once per image before anything else; no statistic uses it. + first = next((cell for cell in cells if cell["mode"] == "load"), None) + if first is None: + raise ValueError("warmup needs a load group") + cells = [first | {"id": cell_id(first["group"], 0, first["concurrency"], first["config"], label), + "round": 0, "image": label, "warmup": True} for label in labels] + cells + if len({cell["id"] for cell in cells}) != len(cells): + raise ValueError("two cells share an id") + corpora = [] + for count in dict.fromkeys(cell["rows"] for cell in cells if cell["mode"] == "load"): + profile = DATA_PROFILES.get(count) + corpora.append({"rows": count, "dir": f"corpus-{count}", "data_profile": profile or "custom", + "generator_args": ["--data-profile", profile] if profile else ["--line-count", str(count)]}) + first = cells[0] return { - "smoke": smoke, "rows": row_count, "rounds": round_count, - "batch_rows": positive("batch_rows", batch_rows, 100_000), "concurrency": levels, - "data_profile": profile or "custom", - "generator_args": ["--data-profile", profile] if profile else ["--line-count", str(row_count)], - "cells": cells, + "smoke": smoke, "rows": first.get("rows"), "rounds": max(cell["round"] for cell in cells), + "batch_rows": positive("batch_rows", batch_rows, 100_000), + "concurrency": sorted({cell["concurrency"] for cell in cells}), + "configs": list(dict.fromkeys(cell["config"] for cell in cells)), + "extra_env": parse_extra_env(extra_env), "groups": [g.get("name", "") for g in resolved], + "restart": restart, "images": labels, "warmup": warmup, "corpora": corpora, "cells": cells, } +def load_rate(cell: dict) -> float: + durable = cell["wal"].get("ACK_MODE") == "durable" + points = LOAD_RATES["other" if not durable else "durable, partitioned" if cell["partitions"] > 1 else "durable"] + level = cell["concurrency"] + for (low, low_rate), (high, high_rate) in zip(points, points[1:]): + if level < high: + return low_rate + (high_rate - low_rate) * math.log(max(level, low) / low) / math.log(high / low) + return points[-1][1] + + +def cell_seconds(cell: dict, batch_rows: int) -> float: + if cell["mode"] == "crash": + return CELL_SECONDS + cell["cycles"] * (sum(cell["kill_after"]) / 2 + CYCLE_SECONDS) + if cell["mode"] == "drop": + return CELL_SECONDS + cell["write_seconds"] + RESTART_SECONDS / 2 + rate = load_rate(cell) * min(1.0, batch_rows / MEASURED_BATCH_ROWS) + return CELL_SECONDS + cell["rows"] / rate + cell["idle_seconds"] + (RESTART_SECONDS if cell["restart"] else 0) + + +def estimate_minutes(plan: dict) -> int: + """The benchmark job's minutes on c6id.2xlarge, from the times of earlier runs.""" + generate = GENERATE_SECONDS_PER_MILLION * sum(corpus["rows"] for corpus in plan["corpora"]) / 1e6 + cells = sum(cell_seconds(cell, plan["batch_rows"]) for cell in plan["cells"]) + return math.ceil((JOB_SECONDS + generate + cells * (1 + CLEANUP_SHARE)) / 60) + + def newest_nightly(tags: list[str]) -> str: dated = [(match.group(1), tag) for tag in tags if (match := NIGHTLY.fullmatch(tag))] if not dated: @@ -114,6 +358,30 @@ def newest_nightly(tags: list[str]) -> str: return max(dated)[1] +def full_reference(name: str, image: str) -> str: + if not isinstance(image, str) or not IMAGE.fullmatch(image) \ + or not re.search(r"(:[^/]+|@sha256:[0-9a-f]{64})$", image): + raise ValueError(f"{name} must be a full reference with a tag or digest: {image!r}") + return image + + +def resolve_image(image: str, tag: str) -> str: + """The full image reference when given, else greptime/greptimedb at the resolved tag.""" + if image: + return full_reference("greptimedb_image", image) + return f"greptime/greptimedb:{resolve_tag(tag)}" + + +def resolve_images(images: str, image: str, tag: str) -> list: + """The images to pull and check: the list to compare, or the one image of the other two inputs.""" + entries = parse_images(images) + if entries[0]["label"]: + if image: + raise ValueError("set either images or greptimedb_image, not both") + return entries + return [{"label": "", "image": resolve_image(image, tag)}] + + def resolve_tag(requested: str) -> str: if requested != "nightly": if not TAG.fullmatch(requested): @@ -169,14 +437,25 @@ def main() -> int: for name in ("concurrency", "rows", "batch-rows", "rounds"): plan.add_argument(f"--{name}", required=True) plan.add_argument("--smoke", choices=("true", "false"), required=True) + plan.add_argument("--restart", choices=("true", "false"), default="false") + plan.add_argument("--configs", default="R,D,E") + plan.add_argument("--warmup", choices=("true", "false"), default="false") + for name in ("extra-env", "groups", "images"): + plan.add_argument(f"--{name}", default="") for name in ("timeout-minutes", "ttl-hours", "output"): plan.add_argument(f"--{name}") - commands.add_parser("tag").add_argument("requested") + image = commands.add_parser("image") + image.add_argument("--tag", required=True) + image.add_argument("--image", default="") + image.add_argument("--images", default="") regions = commands.add_parser("regions") for name in ("bucket", "runner"): regions.add_argument(f"--{name}", required=True) + table = commands.add_parser("table") + table.add_argument("--summary", type=Path, required=True) + table.add_argument("--partitions", type=int, required=True) manifest = commands.add_parser("manifest") - for name in ("results", "runtime-image", "instance-type", "region", "prefix"): + for name in ("results", "images", "runtime-image", "instance-type", "region", "prefix"): manifest.add_argument(f"--{name}", required=True) for name in ("probe", "cleanup"): bucket = commands.add_parser(name) @@ -187,26 +466,33 @@ def main() -> int: if args.command == "plan": if args.timeout_minutes or args.ttl_hours: validate_budgets(args.timeout_minutes or "", args.ttl_hours or "") - result = build_plan(args.concurrency, args.rows, args.batch_rows, args.rounds, args.smoke == "true") + result = build_plan(args.concurrency, args.rows, args.batch_rows, args.rounds, args.smoke == "true", + args.configs, args.extra_env, args.restart == "true", args.groups, + args.images, args.warmup == "true") + minutes = estimate_minutes(result) + if args.timeout_minutes and minutes > int(args.timeout_minutes): + raise ValueError(f"the plan needs about {minutes} minutes on c6id.2xlarge, " + f"more than the {args.timeout_minutes}-minute benchmark timeout") if args.output: with open(args.output, "w") as output: json.dump(result, output, indent=2) - print(f"{len(result['cells'])} cells, {result['rows']} rows, profile {result['data_profile']}") - elif args.command == "tag": - tag = resolve_tag(args.requested) - append_github("GITHUB_OUTPUT", "greptimedb_tag", tag) - print(f"GreptimeDB image tag: {tag}") + print(f"{len(result['cells'])} cells; corpora {[c['rows'] for c in result['corpora']]}; " + f"about {minutes} minutes on c6id.2xlarge") + elif args.command == "image": + images = resolve_images(args.images, args.image, args.tag) + append_github("GITHUB_OUTPUT", "requested_images", json.dumps(images)) + print("GreptimeDB images: " + ", ".join(f"{i['label']} {i['image']}".strip() for i in images)) elif args.command == "regions": check_regions(args.bucket, args.runner) print(f"bucket and runner both in {args.runner}") + elif args.command == "table": + print(partitioned_table(json.loads(args.summary.read_text()), args.partitions)) elif args.command == "manifest": results = Path(args.results) - manifest = build_manifest(json.loads((results / "plan.json").read_text()), - json.loads((results / "images" / "greptimedb.json").read_text()), + manifest = build_manifest(json.loads((results / "plan.json").read_text()), json.loads(args.images), args.runtime_image, args.instance_type, args.region, args.prefix) (results / "run.json").write_text(json.dumps(manifest, indent=2) + "\n") - append_github("GITHUB_ENV", "GREPTIMEDB_IMAGE", manifest["greptimedb_image"]) - print(f"GreptimeDB image: {manifest['greptimedb_image']}") + print("GreptimeDB images: " + ", ".join(i["image"] for i in manifest["greptimedb_images"])) else: if not re.fullmatch(r"oswal-bench/[0-9]+-[0-9]+", args.prefix): raise ValueError(f"refusing a prefix outside oswal-bench/: {args.prefix!r}") diff --git a/.github/scripts/previous-release-tag-test.py b/.github/scripts/previous-release-tag-test.py new file mode 100644 index 00000000000..5b1465eef93 --- /dev/null +++ b/.github/scripts/previous-release-tag-test.py @@ -0,0 +1,133 @@ +#!/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. + +import importlib.util +import subprocess +import sys +import unittest +from pathlib import Path + + +def load_module(): + path = Path(__file__).resolve().parent / "previous-release-tag.py" + spec = importlib.util.spec_from_file_location("previous_release_tag", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +TAGS = [ + "v1.1.0", + "v1.2.0", + "v1.2.1", + "v1.2.1^{}", + "v1.3.0-alpha.1", + "v1.3.0-beta.1", + "v1.3.0-beta.2", + "v1.2.0-nightly-20260706", +] + + +class PreviousReleaseTagTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.mod = load_module() + + def test_beta_uses_latest_older_stable_not_an_ancestor(self) -> None: + self.assertEqual(self.mod.previous_stable_tag("v1.3.0-beta.2", TAGS), "v1.2.1") + + def test_stable_release_skips_itself(self) -> None: + self.assertEqual(self.mod.previous_stable_tag("v1.2.1", TAGS), "v1.2.0") + + def test_same_line_stable_is_not_baseline_for_its_prerelease(self) -> None: + tags = TAGS + ["v1.3.0"] + self.assertEqual(self.mod.previous_stable_tag("v1.3.0-beta.2", tags), "v1.2.1") + self.assertEqual(self.mod.previous_stable_tag("v1.3.0", tags), "v1.2.1") + + def test_nightly_of_unreleased_version_uses_previous_stable(self) -> None: + self.assertEqual( + self.mod.previous_stable_tag("v1.3.0-nightly-20261013", TAGS), + "v1.2.1", + ) + + def test_no_older_stable_tag(self) -> None: + self.assertIsNone(self.mod.previous_stable_tag("v0.1.0", ["v0.1.0", "v1.2.1"])) + + def test_invalid_version(self) -> None: + with self.assertRaises(ValueError): + self.mod.previous_stable_tag("nightly-20261013", TAGS) + + def test_unordered_tags_use_numeric_order(self) -> None: + forward = ["v1.2.1", "v1.9.0", "v1.10.0", "v0.10.0", "v1.9.9"] + backward = list(reversed(forward)) + for tags in (forward, backward): + self.assertEqual(self.mod.previous_stable_tag("v1.10.0", tags), "v1.9.9") + self.assertEqual(self.mod.previous_stable_tag("v1.10.0-rc.1", tags), "v1.9.9") + + def test_multi_digit_components_compare_as_integers(self) -> None: + tags = ["v1.9.0", "v1.10.0", "v0.9.10", "v0.10.0", "v2.9.15", "v2.10.0", "v10.2.0"] + self.assertEqual(self.mod.previous_stable_tag("v1.10.0", tags), "v1.9.0") + self.assertEqual(self.mod.previous_stable_tag("v0.10.0", tags), "v0.9.10") + self.assertEqual(self.mod.previous_stable_tag("v2.10.0", tags), "v2.9.15") + self.assertEqual(self.mod.previous_stable_tag("v10.2.0", ["v10.0.9", "v9.20.0"]), "v10.0.9") + + +class PreviousReleaseTagCliTest(unittest.TestCase): + SCRIPT = Path(__file__).resolve().parent / "previous-release-tag.py" + + def run_cli( + self, version: str, tags: list[str] | None = None, stdin: str | None = None + ) -> subprocess.CompletedProcess[str]: + command = [sys.executable, str(self.SCRIPT), "--version", version] + if tags is not None: + command.extend(tags) + return subprocess.run( + command, + input=stdin, + text=True, + capture_output=True, + check=False, + ) + + def test_stdin_selects_baseline_and_exits_zero(self) -> None: + completed = self.run_cli( + "v1.10.0", + stdin="v1.2.1\nv1.9.0\n\nv1.10.0\nv1.9.0^{}\n", + ) + self.assertEqual(completed.returncode, 0) + self.assertEqual(completed.stdout, "v1.9.0\n") + self.assertEqual(completed.stderr, "") + + def test_positional_tags_ignore_stdin(self) -> None: + completed = self.run_cli("v1.10.0", tags=["v1.9.0"], stdin="v1.9.1\n") + self.assertEqual(completed.returncode, 0) + self.assertEqual(completed.stdout, "v1.9.0\n") + + def test_no_older_tag_exits_zero_with_empty_stdout(self) -> None: + completed = self.run_cli("v0.1.0", stdin="v0.1.0\nv1.2.1\n") + self.assertEqual(completed.returncode, 0) + self.assertEqual(completed.stdout, "") + self.assertEqual(completed.stderr, "") + + def test_invalid_version_exits_one(self) -> None: + completed = self.run_cli("nightly", stdin="v1.2.1\n") + self.assertEqual(completed.returncode, 1) + self.assertEqual(completed.stdout, "") + self.assertIn("not a release version: nightly", completed.stderr) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/previous-release-tag.py b/.github/scripts/previous-release-tag.py new file mode 100644 index 00000000000..b8a8f42ef72 --- /dev/null +++ b/.github/scripts/previous-release-tag.py @@ -0,0 +1,91 @@ +#!/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. + +"""Pick the query-regression baseline tag for a release. + +Compatibility tests do not use this tag. They read ``tests/compatibility/ci.toml``. + +Release branches diverge, so a tag on ``release/v1.3`` is not a descendant of +``v1.2.1``. Ancestry (``git tag --merged``) therefore skips the real previous +release. Compare version numbers instead: the baseline is the greatest +``vX.Y.Z`` tag whose version is strictly older than the release being built. +Prerelease and nightly tags are not baselines. +""" + +from __future__ import annotations + +import argparse +import re +import sys + +STABLE_TAG = re.compile(r"^v(\d+)\.(\d+)\.(\d+)$") +RELEASE_VERSION = re.compile(r"^v(\d+)\.(\d+)\.(\d+)(?:-.+)?$") + + +def _triple(match: re.Match[str]) -> tuple[int, int, int]: + return int(match.group(1)), int(match.group(2)), int(match.group(3)) + + +def release_version(version: str) -> tuple[int, int, int]: + match = RELEASE_VERSION.fullmatch(version.strip()) + if match is None: + raise ValueError(f"not a release version: {version}") + return _triple(match) + + +def previous_stable_tag(version: str, tags: list[str]) -> str | None: + """Return the newest stable tag older than ``version``, or None.""" + current = release_version(version) + best: tuple[int, int, int] | None = None + best_tag: str | None = None + for raw in tags: + tag = raw.strip() + if tag.endswith("^{}"): + tag = tag[: -len("^{}")] + match = STABLE_TAG.fullmatch(tag) + if match is None: + continue + triple = _triple(match) + if triple < current and (best is None or triple > best): + best = triple + best_tag = tag + return best_tag + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--version", required=True, help="Release version, e.g. v1.3.0-beta.2") + parser.add_argument( + "tags", + nargs="*", + help="Candidate tags. Read stdin when omitted.", + ) + args = parser.parse_args(argv) + tags = list(args.tags) + if not tags and not sys.stdin.isatty(): + tags = sys.stdin.read().splitlines() + try: + tag = previous_stable_tag(args.version, tags) + except ValueError as exc: + print(exc, file=sys.stderr) + return 1 + if tag is None: + return 0 + print(tag) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/release-linux-architectures-test.py b/.github/scripts/release-linux-architectures-test.py new file mode 100644 index 00000000000..95fcb6bc19f --- /dev/null +++ b/.github/scripts/release-linux-architectures-test.py @@ -0,0 +1,87 @@ +#!/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. + +import itertools +import os +import subprocess +import unittest +from pathlib import Path + + +SCRIPT = Path(__file__).with_name("release-linux-architectures.sh") +ARCHITECTURES = ("amd64", "arm64", "riscv64") + + +class ReleaseLinuxArchitecturesTest(unittest.TestCase): + def resolve(self, selected, event="workflow_dispatch", publish=False, images=False): + env = dict( + os.environ, + GITHUB_EVENT_NAME=event, + PUBLISH_GITHUB_RELEASE=str(publish).lower(), + RELEASE_IMAGES=str(images).lower(), + ) + for arch in ARCHITECTURES: + key = f"BUILD_LINUX_{arch.upper()}_ARTIFACTS" + env.pop(key, None) + if selected is not None: + env[key] = str(arch in selected).lower() + return subprocess.run( + ["bash", str(SCRIPT)], env=env, text=True, capture_output=True, check=False + ) + + def assert_targets(self, result, targets): + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(result.stderr, "") + self.assertEqual( + result.stdout, + "".join(f"{arch}={str(arch in targets).lower()}\n" for arch in ARCHITECTURES), + ) + + def test_checkbox_combinations(self): + for flags in itertools.product((False, True), repeat=3): + selected = {arch for arch, enabled in zip(ARCHITECTURES, flags) if enabled} + for event, publish, images in itertools.product( + ("workflow_dispatch", "push", "schedule"), (False, True), (False, True) + ): + with self.subTest(selected=selected, event=event, publish=publish, images=images): + expected = ( + set(ARCHITECTURES) + if event != "workflow_dispatch" or publish or images + else selected + ) + self.assert_targets(self.resolve(selected, event, publish, images), expected) + + def test_default_and_omitted_inputs_preserve_previous_behavior(self): + for selected in (None, set()): + for event, publish, images in itertools.product( + ("workflow_dispatch", "push", "schedule"), (False, True), (False, True) + ): + with self.subTest(selected=selected, event=event, publish=publish, images=images): + # The old build_linux_artifacts switch defaulted to false. + old_build_linux = publish or images or event in ("push", "schedule") + self.assert_targets( + self.resolve(selected, event, publish, images), + set(ARCHITECTURES) if old_build_linux else set(), + ) + + def test_unsupported_event(self): + result = self.resolve(set(), event="pull_request") + self.assertNotEqual(result.returncode, 0) + self.assertIn("Unsupported release event", result.stderr) + self.assertEqual(result.stdout, "") + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/release-linux-architectures.sh b/.github/scripts/release-linux-architectures.sh new file mode 100644 index 00000000000..cf6f819e84f --- /dev/null +++ b/.github/scripts/release-linux-architectures.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash + +set -euo pipefail + +# Publishing keeps the existing full-build behavior, regardless of checkboxes. +build_all=false +if [[ "${PUBLISH_GITHUB_RELEASE:-false}" == true || "${RELEASE_IMAGES:-false}" == true ]]; then + build_all=true +fi +case "${GITHUB_EVENT_NAME}" in + push|schedule) build_all=true ;; + workflow_dispatch) ;; + *) echo "Unsupported release event: ${GITHUB_EVENT_NAME}" >&2; exit 1 ;; +esac + +for arch in amd64 arm64 riscv64; do + input="BUILD_LINUX_${arch^^}_ARTIFACTS" + selected=false + if [[ "$build_all" == true || "${!input:-false}" == true ]]; then + selected=true + fi + printf '%s=%s\n' "$arch" "$selected" +done diff --git a/.github/scripts/tests/test_metric_snapshot_fixture.py b/.github/scripts/tests/test_metric_snapshot_fixture.py new file mode 100644 index 00000000000..f4c16a11af3 --- /dev/null +++ b/.github/scripts/tests/test_metric_snapshot_fixture.py @@ -0,0 +1,101 @@ +# 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. + +import copy +import importlib.util +from pathlib import Path +import unittest + +spec = importlib.util.spec_from_file_location( + "fixture", Path(__file__).resolve().parents[1] / "metric-snapshot-fixture.py" +) +fixture = importlib.util.module_from_spec(spec) +spec.loader.exec_module(fixture) + + +class ValidatorTests(unittest.TestCase): + def records(self): + return { + "schema": {"column_schemas": [ + {"name": "ts", "data_type": "TimestampMillisecond"}, + *[{"name": f"tag_{i:03d}", "data_type": "String"} for i in range(10)], + {"name": "val", "data_type": "Float64"}, + ]}, + "rows": [fixture.row(0, j, 50) for j in range(32)], + } + + def test_full_tag_coverage_and_heterogeneous_schemas(self): + for per_group in (50, 10000): + variants = {tuple(fixture.tags(i, per_group)) for i in range(per_group)} + self.assertEqual(set().union(*map(set, variants)), set(range(298))) + self.assertGreater(len(variants), 1) + self.assertTrue(all(len(tags) == 10 for tags in variants)) + + def test_accepts_engine_column_order_with_matching_values(self): + records = self.records() + order = [*range(1, 11), 0, 11] # The engine returns tags, timestamp, value. + records["schema"]["column_schemas"] = [ + records["schema"]["column_schemas"][index] for index in order + ] + records["rows"] = [[values[index] for index in order] for values in records["rows"]] + fixture.check_records(0, 50, records) + records["rows"][1][10] = 2 + with self.assertRaisesRegex(RuntimeError, "typed values"): + fixture.check_records(0, 50, records) + + def test_rejects_column_type_and_membership_changes(self): + fixture.check_records(0, 50, self.records()) + for mutation in ("type", "extra", "missing"): + records = self.records() + columns = records["schema"]["column_schemas"] + if mutation == "type": + columns[0]["data_type"] = "TimestampMicrosecond" + elif mutation == "extra": + columns.append({"name": "unexpected", "data_type": "String"}) + else: + del columns[1] + with self.subTest(mutation=mutation), self.assertRaisesRegex(RuntimeError, "schema"): + fixture.check_records(0, 50, records) + + def test_rejects_null_and_numeric_type_changes(self): + for value in ("", 0, False): + records = self.records() + records["rows"][0][-1] = value # Expected SQL NULL. + with self.subTest(value=value), self.assertRaisesRegex(RuntimeError, "typed values"): + fixture.check_records(0, 50, records) + records = self.records() + records["rows"][4][-1] = 1 # Numeric equality to 1.0 must not hide a type change. + with self.assertRaisesRegex(RuntimeError, "typed values"): + fixture.check_records(0, 50, records) + + def test_source_comparison_requires_same_ddl_and_remapped_ids(self): + source = {"fixture": {"tables_per_group": 50}, "schema_sha256": "schema", + "data_sha256": "data", "ddl_sha256": "ddl", + "table_ids": {"m000000": 1024, "phy_00": 1025}} + target = copy.deepcopy(source) + target["table_ids"] = {"m000000": 1041, "phy_00": 1042} + fixture.compare_source(target, source) + self.assertEqual(target["remapped_table_ids"], 2) + for field in ("schema_sha256", "data_sha256", "ddl_sha256"): + changed = copy.deepcopy(target) + changed[field] = "changed" + with self.subTest(field=field), self.assertRaisesRegex(RuntimeError, field): + fixture.compare_source(changed, source) + target["table_ids"]["phy_00"] = 1025 + with self.assertRaisesRegex(RuntimeError, "not all remapped"): + fixture.compare_source(target, source) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/tests/test_metric_snapshot_s3_bench.py b/.github/scripts/tests/test_metric_snapshot_s3_bench.py new file mode 100644 index 00000000000..c662ecbff76 --- /dev/null +++ b/.github/scripts/tests/test_metric_snapshot_s3_bench.py @@ -0,0 +1,74 @@ +#!/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. + +"""Failure accounting for the metric snapshot runner.""" +import importlib.util +import json +from pathlib import Path +import tempfile +import unittest +from unittest.mock import patch + +SPEC = importlib.util.spec_from_file_location( + "metric_snapshot_bench", Path(__file__).resolve().parents[1] / "metric-snapshot-s3-bench.py") +bench = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(bench) + + +class PhaseTests(unittest.TestCase): + def test_docker_nanosecond_timestamps(self): + start = bench.docker_timestamp("2026-10-14T06:06:44.72000818Z") + end = bench.docker_timestamp("2026-10-14T06:06:46.123456789Z") + self.assertAlmostEqual((end - start).total_seconds(), 1.403448) + + def test_resource_interference_invalidates_success_and_preserves_fixture(self): + with tempfile.TemporaryDirectory() as directory: + runner = bench.Benchmark.__new__(bench.Benchmark) + runner.round_dir = Path(directory) + runner.server_pid = 0 + runner.phases = [] + runner.metrics = lambda: "" + runner.summary = lambda: None + fixture = runner.round_dir / "source-validation.json" + fixture.write_text('{"data_sha256":"fixture"}') + before = {"available": True, "pswpout_pages": 0, "oom_kill": 0} + after = {"available": True, "pswpout_pages": 1, "oom_kill": 0} + with patch.object(bench, "host_memory", side_effect=[before, after]): + with self.assertRaisesRegex(RuntimeError, "Resource interference"): + runner.phase("source-validation", lambda: None) + self.assertEqual(json.loads(fixture.read_text()), {"data_sha256": "fixture"}) + evidence = json.loads((runner.round_dir / "source-validation-resources.json").read_text()) + self.assertEqual(evidence["outcome"], "resource-invalid") + self.assertEqual(evidence["new_swapout_pages"], 1) + + def test_action_failure_keeps_failed_resource_evidence(self): + with tempfile.TemporaryDirectory() as directory: + runner = bench.Benchmark.__new__(bench.Benchmark) + runner.round_dir = Path(directory) + runner.server_pid = 0 + runner.phases = [] + runner.metrics = lambda: "" + runner.summary = lambda: None + def fail(): + raise RuntimeError("export failed") + with patch.object(bench, "host_memory", return_value={"available": False}): + with self.assertRaisesRegex(RuntimeError, "export failed"): + runner.phase("export", fail) + evidence = json.loads((runner.round_dir / "export-resources.json").read_text()) + self.assertEqual(evidence["outcome"], "failed") + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/tests/test_metric_snapshot_s3_ops.py b/.github/scripts/tests/test_metric_snapshot_s3_ops.py new file mode 100644 index 00000000000..d22c09044b2 --- /dev/null +++ b/.github/scripts/tests/test_metric_snapshot_s3_ops.py @@ -0,0 +1,203 @@ +#!/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. + +"""Offline safety checks for the metric snapshot workflow's budget and cleanup.""" + +import importlib.util +import io +import json +from pathlib import Path +import tempfile +import unittest +from unittest import mock + +SPEC = importlib.util.spec_from_file_location("snapshot_ops", Path(__file__).parents[1] / "metric-snapshot-s3-ops.py") +ops = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(ops) +PREFIX = "metric-snapshot-bench/123-1/" + + +class FakeS3: + def __init__(self): + self.objects = [{"Key": PREFIX + "current", "Size": 11}, {"Key": PREFIX + "data", "Size": 7}] + self.deleted = [] + + def get_paginator(self, operation): + assert operation == "list_objects_v2" + parent = self + + class Paginator: + def paginate(self, **kwargs): + assert kwargs["Bucket"] == "bucket" and kwargs["Prefix"].startswith(PREFIX) + # Keys outside the run prefix model a misbehaving listing and are always returned. + return iter([{"Contents": [item for item in parent.objects if not item["Key"].startswith(PREFIX) + or item["Key"].startswith(kwargs["Prefix"])]}]) + + return Paginator() + + def put_object(self, **kwargs): + self.objects.append({"Key": kwargs["Key"], "Size": len(kwargs["Body"])}) + return {} + + def get_object(self, **kwargs): + return {"Body": io.BytesIO(b"metric snapshot benchmark")} + + def delete_object(self, **kwargs): + self.deleted.append(kwargs) + self.objects = [item for item in self.objects if item["Key"] != kwargs["Key"]] + + +class OperationsTest(unittest.TestCase): + def env(self, **changes): + return dict(GREPTIMEDB_IMAGE="greptime/greptimedb-nightly@sha256:" + "a" * 64, + EXPECTED_COMMIT="b" * 40, LOGICAL_TABLES_PER_GROUP="10000", + BENCHMARK_TIMEOUT_MINUTES="240", PRIOR_SPEND_USD="0", + S3_REGION="ap-southeast-1", RUNNER_REGION="ap-southeast-1", S3_BUCKET="bucket", + RUN_PREFIX=PREFIX) | changes + + def test_budget_includes_provision_cleanup_and_other_cost_reserve(self): + plan = ops.plan(self.env()) + self.assertEqual(plan["attempt_reserve_usd"], "6.4392") + ops.plan(self.env(PRIOR_SPEND_USD="13.5608")) + for changes in ({"PRIOR_SPEND_USD": "13.5609"}, {"PRIOR_SPEND_USD": "15"}, + {"PRIOR_SPEND_USD": "NaN"}, {"PRIOR_SPEND_USD": "-1"}, + {"BENCHMARK_TIMEOUT_MINUTES": "241"}, {"LOGICAL_TABLES_PER_GROUP": "49"}, + {"LOGICAL_TABLES_PER_GROUP": "10001"}, {"S3_REGION": "us-east-1"}, + {"GREPTIMEDB_IMAGE": "greptime/greptimedb:latest"}, + {"GREPTIMEDB_IMAGE": "docker.io/greptime/greptimedb@sha256:" + "a" * 64}, + {"GREPTIMEDB_IMAGE": "greptime/greptimedb:nightly@sha256:" + "a" * 64}, + {"EXPECTED_COMMIT": "abcdef"}): + with self.subTest(changes=changes), self.assertRaises(ValueError): + ops.plan(self.env(**changes)) + + def watcher_env(self): + return self.env(GITHUB_REPOSITORY="GreptimeTeam/greptimedb", GITHUB_RUN_ID="123", + GITHUB_RUN_ATTEMPT="2", GH_TOKEN="test-token", BENCHMARK_TIMEOUT_MINUTES="30") + + def test_watcher_completed_job_uses_original_attempt_and_does_not_wait(self): + for conclusion in ("success", "failure", "cancelled"): + response = io.BytesIO(json.dumps({"jobs": [{"name": "benchmark", "status": "completed", + "conclusion": conclusion}]}).encode()) + with self.subTest(conclusion=conclusion), \ + mock.patch.object(ops.urllib.request, "urlopen", return_value=response) as fetch, \ + mock.patch.object(ops.time, "time", return_value=1005), \ + mock.patch.object(ops.time, "sleep") as sleep: + self.assertEqual(ops.wait_for_benchmark(self.watcher_env(), 1000), + {"benchmark_conclusion": conclusion, "owned_attempt": "1"}) + self.assertIn("/runs/123/attempts/1/jobs?", fetch.call_args.args[0].full_url) + sleep.assert_not_called() + + def test_watcher_queue_reaches_durable_deadline(self): + deadline = 1000 + (45 + 30 + 10) * 60 + clock = [deadline - 31] + + def advance(seconds): + self.assertLessEqual(seconds, 30) + clock[0] += seconds + + def queued(*args, **kwargs): + return io.BytesIO(b'{"jobs": [{"name": "benchmark", "status": "queued"}]}') + + with mock.patch.object(ops.urllib.request, "urlopen", side_effect=queued) as fetch, \ + mock.patch.object(ops.time, "time", side_effect=lambda: clock[0]), \ + mock.patch.object(ops.time, "sleep", side_effect=advance), \ + self.assertRaisesRegex(TimeoutError, "proceeding to resource cleanup"): + ops.wait_for_benchmark(self.watcher_env(), 1000) + self.assertEqual(clock[0], deadline) + self.assertEqual(fetch.call_count, 2) + + def test_watcher_api_error_fails_immediately_for_cleanup(self): + with mock.patch.object(ops.urllib.request, "urlopen", side_effect=OSError("API unavailable")), \ + mock.patch.object(ops.time, "time", return_value=1005), \ + mock.patch.object(ops.time, "sleep") as sleep, \ + self.assertRaises(OSError): + ops.wait_for_benchmark(self.watcher_env(), 1000) + sleep.assert_not_called() + + def test_cleanup_deletes_every_listed_object(self): + s3 = FakeS3() + result = ops.cleanup(s3, "bucket", PREFIX) + self.assertEqual(result["objects"], 0) + self.assertEqual(s3.deleted, [{"Bucket": "bucket", "Key": PREFIX + "current"}, + {"Bucket": "bucket", "Key": PREFIX + "data"}]) + + def test_probe_refuses_occupied_prefix_and_removes_its_object(self): + with self.assertRaisesRegex(ValueError, "already occupied"): + ops.probe(FakeS3(), "bucket", PREFIX) + s3 = FakeS3() + s3.objects = [] + self.assertEqual(ops.probe(s3, "bucket", PREFIX)["objects"], 0) + self.assertEqual(s3.deleted, [{"Bucket": "bucket", "Key": PREFIX + "_probe"}]) + + def test_failed_probe_delete_still_has_hosted_cleanup_path(self): + s3 = FakeS3() + s3.objects = [] + with mock.patch.object(s3, "delete_object", side_effect=OSError("temporary delete failure")), \ + self.assertRaises(OSError): + ops.probe(s3, "bucket", PREFIX) + self.assertEqual([item["Key"] for item in s3.objects], [PREFIX + "_probe"]) + workflow = (Path(__file__).parents[2] / "workflows/metric-snapshot-s3-bench.yml").read_text() + teardown = workflow.split(" teardown:\n", 1)[1] + job_if = next(line.strip() for line in teardown.splitlines() if line.startswith(" if:")) + self.assertEqual(job_if, "if: ${{ always() && needs.validate.outputs.run_id != '' }}") + ec2 = teardown.split(" - name: Delete EC2, its temporary disks and runner registration\n", 1)[1] + self.assertEqual(ec2.splitlines()[0].strip(), "if: ${{ always() && needs.provision.result != 'skipped' }}") + self.assertIn(" RUN_PREFIX: ${{ needs.validate.outputs.run_prefix }}\n", teardown.split(" steps:", 1)[0]) + s3_cleanup = teardown.split(" - name: Remove and verify only this run's S3 prefix\n", 1)[1].split(" - ", 1)[0] + self.assertEqual(s3_cleanup.splitlines()[0].strip(), "if: always()") + self.assertIn("run: uv run greptimedb/.github/scripts/metric-snapshot-s3-ops.py cleanup\n", s3_cleanup) + self.assertEqual(ops.cleanup(s3, "bucket", PREFIX)["objects"], 0) + + def test_round_inventory_excludes_other_rounds_while_cleanup_covers_the_run(self): + s3 = FakeS3() + s3.objects = [{"Key": PREFIX + "smoke/pack-000000.bin", "Size": 3}, + {"Key": PREFIX + "full/pack-000000.bin", "Size": 5}, + {"Key": PREFIX + "full/index.json", "Size": 2}] + full = ops.inventory(s3, "bucket", PREFIX, "full") + self.assertEqual((full["prefix"], full["objects"], full["bytes"]), (PREFIX + "full/", 2, 7)) + self.assertEqual(ops.inventory(s3, "bucket", PREFIX, "smoke")["bytes"], 3) + self.assertEqual(ops.inventory(s3, "bucket", PREFIX)["objects"], 3) + for round in ("", "../full", "full/"): + with self.subTest(round=round), self.assertRaises(ValueError): + ops.inventory(s3, "bucket", PREFIX, round) + self.assertEqual(ops.cleanup(s3, "bucket", PREFIX)["objects"], 0) + self.assertEqual(len(s3.deleted), 3) + + def test_bad_prefix_and_out_of_scope_response_cannot_delete(self): + for prefix in ("", "metric-snapshot-bench/", "metric-snapshot-bench/123-1", PREFIX + "../", "oswal-bench/123-1/"): + s3 = FakeS3() + with self.subTest(prefix=prefix), self.assertRaises(ValueError): + ops.cleanup(s3, "bucket", prefix) + self.assertEqual(s3.deleted, []) + s3 = FakeS3() + s3.objects[0]["Key"] = "other-run/current" + with self.assertRaises(ValueError): + ops.inventory(s3, "bucket", PREFIX, "full") + with self.assertRaises(ValueError): + ops.cleanup(s3, "bucket", PREFIX) + self.assertEqual(s3.deleted, []) + + def test_credential_crossing_read_boundary_blocks_artifact_upload(self): + with tempfile.TemporaryDirectory() as directory: + logfile = Path(directory) / "server.log" + logfile.write_bytes(b"a" * (1024 * 1024 - 4) + b"private-key" + b"tail") + with self.assertRaisesRegex(ValueError, "upload is blocked"): + ops.scrub(directory, {"AWS_SECRET_ACCESS_KEY": "private-key"}) + logfile.write_text("safe log\n") + ops.scrub(directory, {"AWS_SECRET_ACCESS_KEY": "private-key"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/workflows/apidoc.yml b/.github/workflows/apidoc.yml index eb7a720a421..fb5601fdc5b 100644 --- a/.github/workflows/apidoc.yml +++ b/.github/workflows/apidoc.yml @@ -19,9 +19,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - run: cargo doc --workspace --no-deps --document-private-items - run: | diff --git a/.github/workflows/aws-ec2-runner-janitor.yml b/.github/workflows/aws-ec2-runner-janitor.yml index 1c4435022b4..0355a1159d9 100644 --- a/.github/workflows/aws-ec2-runner-janitor.yml +++ b/.github/workflows/aws-ec2-runner-janitor.yml @@ -10,6 +10,7 @@ permissions: jobs: sweep: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} runs-on: ubuntu-latest timeout-minutes: 30 steps: diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index deba5010c80..3bc54372d06 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -82,12 +82,18 @@ jobs: version: "latest" - name: Run check-version script tests run: uv run --no-project python .github/scripts/check-version-test.py + - name: Run release Linux architecture tests + run: python3 .github/scripts/release-linux-architectures-test.py + - name: Run previous-release-tag script tests + run: python3 .github/scripts/previous-release-tag-test.py - name: Run enterprise license check tests run: python3 scripts/check-enterprise-license-test.py - name: Run fuzz orchestration script tests run: .github/scripts/run-fuzz-targets-test.sh - name: Run CI comment-command tests run: python3 .github/scripts/tests/test_ci_slash.py + - name: Run metric snapshot benchmark tests + run: python3 -m unittest discover -s .github/scripts/tests -p 'test_metric_snapshot*.py' - name: Run query-regression Python tooling tests run: | python3 tests/perf/test_query_regression_runner_compaction_toctou.py @@ -113,9 +119,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Rust Cache uses: Swatinem/rust-cache@v2 diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index db4f4bef65b..afce4a8c796 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -54,9 +54,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Build Binaries Cache id: build-binaries-cache @@ -193,9 +193,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Rust Cache uses: Swatinem/rust-cache@v2 @@ -327,9 +327,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: cache: false @@ -364,13 +364,11 @@ jobs: env: RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} run: | - sudo apt-get install -y libfuzzer-14-dev cargo install cargo-fuzz --version "${CARGO_FUZZ_VERSION}" --locked - name: Build fuzz target binaries if: steps.fuzz-binaries.outputs.cache-hit != 'true' shell: bash env: - CUSTOM_LIBFUZZER_PATH: /usr/lib/llvm-14/lib/libFuzzer.a RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} FUZZ_TARGETS: | fuzz_create_database @@ -587,9 +585,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Build Profile-CI Cache id: build-greptime-ci-cache @@ -705,9 +703,9 @@ jobs: uses: ./.github/actions/setup-etcd-cluster # Only the Kafka WAL helper needs the Rust toolchain. - if: matrix.mode.kafkaWalHelper - uses: arduino/setup-protoc@v3 + uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - if: matrix.mode.kafkaWalHelper uses: actions-rust-lang/setup-rust-toolchain@v1 with: diff --git a/.github/workflows/metric-snapshot-s3-bench.yml b/.github/workflows/metric-snapshot-s3-bench.yml new file mode 100644 index 00000000000..e4a12dde69f --- /dev/null +++ b/.github/workflows/metric-snapshot-s3-bench.yml @@ -0,0 +1,230 @@ +name: Metric Snapshot S3 Benchmark +on: + workflow_dispatch: + inputs: + greptimedb_image: + description: 'Immutable greptime/greptimedb or greptimedb-nightly image (repository@sha256:digest)' + required: true + type: string + expected_commit: + description: 'Full source commit SHA reported by the image' + required: true + type: string + logical_tables_per_group: + description: 'Logical tables per physical table, 50..10000; 10 physical tables, 32 rows per logical table' + type: number + default: 10000 + benchmark_timeout_minutes: + description: 'Setup, 500-table smoke and full run timeout, 30..240 minutes' + type: number + default: 240 + prior_spend_usd: + description: 'Cumulative estimated cloud spend for all earlier attempts in this $20 campaign' + required: true + type: string + default: '0' +permissions: + contents: read +concurrency: + group: metric-snapshot-s3-bench + cancel-in-progress: false +env: + GREPTIMEDB_IMAGE: ${{ inputs.greptimedb_image }} + EXPECTED_COMMIT: ${{ inputs.expected_commit }} + LOGICAL_TABLES_PER_GROUP: ${{ inputs.logical_tables_per_group }} + BENCHMARK_TIMEOUT_MINUTES: ${{ inputs.benchmark_timeout_minutes }} + PRIOR_SPEND_USD: ${{ inputs.prior_spend_usd }} + RUN_PREFIX: metric-snapshot-bench/${{ github.run_id }}-${{ github.run_attempt }}/ + S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} + S3_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + RUNNER_REGION: ${{ vars.EC2_RUNNER_REGION }} +jobs: + validate: + runs-on: ubuntu-latest + timeout-minutes: 15 + outputs: + instance_type: ${{ steps.runner.outputs.instance_type }} + run_id: ${{ steps.identity.outputs.run_id }} + run_prefix: ${{ steps.identity.outputs.run_prefix }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Test benchmark boundaries + run: | + python3 greptimedb/.github/scripts/tests/test_metric_snapshot_s3_ops.py + python3 greptimedb/.github/scripts/tests/test_metric_snapshot_fixture.py + python3 greptimedb/.github/scripts/tests/test_metric_snapshot_s3_bench.py + - name: Validate pinned version, scale, region and cumulative budget + run: python3 greptimedb/.github/scripts/metric-snapshot-s3-ops.py plan + - name: Bind resource ownership to this validated attempt + id: identity + run: | + echo "run_id=$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT" >> "$GITHUB_OUTPUT" + echo "run_prefix=$RUN_PREFIX" >> "$GITHUB_OUTPUT" + - uses: astral-sh/setup-uv@v6 + - name: Check S3 access before allocating EC2 + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: uv run greptimedb/.github/scripts/metric-snapshot-s3-ops.py probe + - name: Resolve the fixed benchmark runner + id: runner + run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider AWS --instance-type c6id.4xlarge + provision: + needs: validate + runs-on: ubuntu-latest + timeout-minutes: 45 + outputs: + label: ${{ steps.aws.outputs.label }} + instance_id: ${{ steps.aws.outputs.instance_id }} + runner_name: ${{ steps.aws.outputs.runner_name }} + started_at: ${{ steps.start.outputs.started_at }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ env.RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + - name: Record allocation start + id: start + run: echo "started_at=$(date +%s)" >> "$GITHUB_OUTPUT" + - name: Create ephemeral EC2 runner + id: aws + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ env.RUNNER_REGION }} + AWS_EC2_IMAGE_ID: ${{ vars.BENCHMARK_EC2_IMAGE_ID || 'auto' }} + AWS_EC2_INSTANCE_TYPE: ${{ needs.validate.outputs.instance_type }} + AWS_EC2_SUBNET_ID: ${{ vars.EC2_RUNNER_SUBNET_ID }} + AWS_EC2_SECURITY_GROUP_ID: ${{ vars.EC2_RUNNER_SECURITY_GROUP_ID }} + AWS_EC2_SYSTEM_DISK_GIB: '60' + AWS_EC2_TTL_HOURS: '6' + AWS_EC2_RUN_ID: ${{ needs.validate.outputs.run_id }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-provision.py + benchmark: + needs: + - validate + - provision + runs-on: ${{ needs.provision.outputs.label }} + timeout-minutes: ${{ fromJSON(format('{0}', inputs.benchmark_timeout_minutes)) }} + env: + RUN_PREFIX: ${{ needs.validate.outputs.run_prefix }} + RESULTS: ${{ github.workspace }}/metric-snapshot-results + DISK_ROOT: /mnt/metric-snapshot-disk + DATA_ROOT: /mnt/metric-snapshot-disk/data + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Set up benchmark dependencies + uses: ./greptimedb/.github/actions/setup-benchmark + - uses: astral-sh/setup-uv@v6 + - name: Mount the temporary local File data disk + run: | + device=$(lsblk -dnpo NAME,MODEL | awk '/Instance Storage/ {print $1; exit}') + test -n "$device" + if [[ "$(lsblk -nro TYPE "$device" | wc -l)" -ne 1 ]] || findmnt --source "$device" >/dev/null; then + echo 'Refusing partitioned or mounted instance storage' >&2 + exit 1 + fi + probe_status=0 + sudo -n blkid -p "$device" || probe_status=$? + if ((probe_status != 2)); then + echo 'Refusing a disk with an existing signature or an unsuccessful signature probe' >&2 + exit 1 + fi + sudo -n mkfs.ext4 -q "$device" + sudo -n mkdir -p "$DISK_ROOT" + sudo -n mount -o noatime "$device" "$DISK_ROOT" + sudo -n chown "$(id -u):$(id -g)" "$DISK_ROOT" + findmnt -no SOURCE,FSTYPE -T "$DISK_ROOT" + - name: Probe the real same-region S3 bucket + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: uv run greptimedb/.github/scripts/metric-snapshot-s3-ops.py probe + - name: Run the 500-table smoke, then one full packed export and restore + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: python3 greptimedb/.github/scripts/metric-snapshot-s3-bench.py run + - name: Check results for credentials + id: scrub + if: always() + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: python3 greptimedb/.github/scripts/metric-snapshot-s3-ops.py scrub + - name: Publish the readable result + if: ${{ always() && steps.scrub.outcome == 'success' }} + run: | + if test -f "$RESULTS/summary.md"; then + cat "$RESULTS/summary.md" >> "$GITHUB_STEP_SUMMARY" + fi + - name: Upload bounded-retention evidence + if: ${{ always() && steps.scrub.outcome == 'success' }} + uses: actions/upload-artifact@v4 + with: + name: metric-snapshot-s3-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ env.RESULTS }} + if-no-files-found: warn + retention-days: 14 + teardown: + needs: + - validate + - provision + # Start alongside the benchmark so an offline runner cannot hold cleanup in its queue. + if: ${{ always() && needs.validate.outputs.run_id != '' }} + runs-on: ubuntu-latest + timeout-minutes: 330 + permissions: + contents: read + actions: read + env: + RUN_PREFIX: ${{ needs.validate.outputs.run_prefix }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - uses: astral-sh/setup-uv@v6 + - name: Wait for completion within the allocation deadline + if: ${{ !cancelled() && needs.provision.result == 'success' && needs.provision.outputs.started_at != '' }} + env: + GH_TOKEN: ${{ github.token }} + STARTED_AT: ${{ needs.provision.outputs.started_at }} + run: python3 greptimedb/.github/scripts/metric-snapshot-s3-ops.py wait-for-benchmark --started-at "$STARTED_AT" + - uses: aws-actions/configure-aws-credentials@v4 + if: ${{ always() && needs.provision.result != 'skipped' }} + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ env.RUNNER_REGION }} + - name: Delete EC2, its temporary disks and runner registration + if: ${{ always() && needs.provision.result != 'skipped' }} + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ env.RUNNER_REGION }} + AWS_EC2_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + AWS_EC2_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + AWS_EC2_RUN_ID: ${{ needs.validate.outputs.run_id }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-teardown.py + - name: Remove and verify only this run's S3 prefix + if: always() + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: uv run greptimedb/.github/scripts/metric-snapshot-s3-ops.py cleanup + - name: Estimate the cumulative cloud cost + if: ${{ always() && needs.provision.outputs.started_at != '' }} + env: + STARTED_AT: ${{ needs.provision.outputs.started_at }} + run: python3 greptimedb/.github/scripts/metric-snapshot-s3-ops.py cost --started-at "$STARTED_AT" >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index f7f955233a1..f6158a8d7ad 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -53,9 +53,9 @@ jobs: fetch-depth: 0 persist-credentials: false - uses: ./.github/actions/setup-cyborg - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Rust Cache uses: Swatinem/rust-cache@v2 @@ -82,9 +82,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: rui314/setup-mold@v1 - name: Install toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 @@ -134,9 +134,9 @@ jobs: fetch-depth: 0 persist-credentials: false - uses: ./.github/actions/setup-cyborg - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: KyleMayes/install-llvm-action@v1 with: version: "14.0" @@ -157,6 +157,7 @@ jobs: env: CARGO_BUILD_RUSTFLAGS: "-C linker=lld-link" RUST_BACKTRACE: 1 + RUST_MIN_STACK: 8388608 # 8 MiB, aligned with Rust CI and release tests. CARGO_INCREMENTAL: 0 GT_S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} diff --git a/.github/workflows/oswal-s3-bench.yml b/.github/workflows/oswal-s3-bench.yml index c31dfe8f1c3..e61754cb1f1 100644 --- a/.github/workflows/oswal-s3-bench.yml +++ b/.github/workflows/oswal-s3-bench.yml @@ -10,6 +10,34 @@ on: description: 'Docker Hub greptime/greptimedb tag; nightly = the newest v*-nightly-* tag (latest is a release without the object store WAL)' type: string default: nightly + greptimedb_image: + description: 'Full image reference with a tag or digest; used instead of greptimedb_tag when set' + type: string + default: '' + images: + description: 'JSON list of 2..4 {"label", "image"} to compare on one instance instead of greptimedb_tag and greptimedb_image: every cell runs once per image, in an order that reverses every other round; a label is up to 8 lowercase letters and digits, an image a full reference with a tag or digest, the first the reference of the ratios' + type: string + default: '' + warmup: + description: Before round 1, run the first load cell once per image and leave it out of every statistic + type: boolean + default: false + configs: + description: 'WAL configurations: R (Raft Engine, sync_write), R0 (Raft Engine, standalone default), D (durable), E (enqueued)' + type: string + default: R,D,E + extra_env: + description: 'Environment for every GreptimeDB container: a JSON object or K=V lines' + type: string + default: '' + restart: + description: 'After each load, SIGKILL the server, restart it on the same prefix and data home, and time recovery; object store cells flush every 10ms' + type: boolean + default: false + groups: + description: 'JSON list of cell groups, each overriding configs, extra_env, concurrency, rounds, rows, restart, plus idle_seconds and partitions (regions of the loaded table), or mode crash with cycles, kill_after [min, max] and pin, or mode drop with write_seconds; empty = one group from the inputs above' + type: string + default: '' instance_type: description: EC2 instance type in the bucket region; the d of c6id gives a local NVMe instance store for the Raft Engine WAL type: string @@ -27,7 +55,7 @@ on: type: number default: 1000 rounds: - description: Rounds; the configuration order alternates between rounds (R, D, E then E, D, R) + description: Rounds; the configuration and image order alternates between rounds (R, D, E then E, D, R) type: number default: 2 runtime_image: @@ -47,6 +75,13 @@ permissions: env: SMOKE: ${{ inputs.smoke }} GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} + GREPTIMEDB_IMAGE_INPUT: ${{ inputs.greptimedb_image }} + OSWAL_IMAGES: ${{ inputs.images }} + WARMUP: ${{ inputs.warmup }} + CONFIGS: ${{ inputs.configs }} + EXTRA_ENV: ${{ inputs.extra_env }} + RESTART: ${{ inputs.restart }} + OSWAL_GROUPS: ${{ inputs.groups }} REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} CONCURRENCY: ${{ inputs.concurrency }} ROWS: ${{ inputs.rows }} @@ -61,7 +96,7 @@ jobs: timeout-minutes: 15 outputs: instance_type: ${{ steps.runner.outputs.instance_type }} - greptimedb_tag: ${{ steps.plan.outputs.greptimedb_tag }} + greptimedb_images: ${{ steps.images.outputs.greptimedb_images }} runtime_image: ${{ steps.runtime.outputs.runtime_image }} steps: - uses: actions/checkout@v4 @@ -77,8 +112,16 @@ jobs: run: | python3 greptimedb/.github/scripts/oswal-s3-bench.py plan --smoke "$SMOKE" \ --concurrency "$CONCURRENCY" --rows "$ROWS" --batch-rows "$BATCH_ROWS" --rounds "$ROUNDS" \ + --configs "$CONFIGS" --extra-env "$EXTRA_ENV" --restart "$RESTART" --groups "$OSWAL_GROUPS" \ + --images "$OSWAL_IMAGES" --warmup "$WARMUP" \ --timeout-minutes "$BENCHMARK_TIMEOUT_MINUTES" --ttl-hours "$JANITOR_TTL_HOURS" - python3 greptimedb/.github/scripts/oswal-s3-bench.py tag "$GREPTIMEDB_TAG" + python3 greptimedb/.github/scripts/oswal-s3-bench.py image --tag "$GREPTIMEDB_TAG" --image "$GREPTIMEDB_IMAGE_INPUT" \ + --images "$OSWAL_IMAGES" + - name: Pull every GreptimeDB image and record its version + id: images + env: + REQUESTED_IMAGES: ${{ steps.plan.outputs.requested_images }} + run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh check-images "$REQUESTED_IMAGES" - name: Resolve CI runner id: runner run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider AWS --instance-type "$REQUESTED_INSTANCE_TYPE" @@ -87,7 +130,7 @@ jobs: BUCKET_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} RUNNER_REGION: ${{ vars.EC2_RUNNER_REGION }} run: python3 greptimedb/.github/scripts/oswal-s3-bench.py regions --bucket "$BUCKET_REGION" --runner "$RUNNER_REGION" - - name: Check the runtime image + - name: Pin the runtime image digest id: runtime run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh check-runtime "$RUNTIME_IMAGE" provision: @@ -132,7 +175,7 @@ jobs: timeout-minutes: ${{ fromJSON(format('{0}', inputs.benchmark_timeout_minutes)) }} env: INSTANCE_TYPE: ${{ inputs.smoke && 'ubuntu-latest' || needs.validate.outputs.instance_type }} - GREPTIMEDB_TAG: ${{ needs.validate.outputs.greptimedb_tag }} + GREPTIMEDB_IMAGES: ${{ needs.validate.outputs.greptimedb_images }} RUNTIME_IMAGE: ${{ needs.validate.outputs.runtime_image }} RUN_PREFIX: oswal-bench/${{ github.run_id }}-${{ github.run_attempt }} OWNER: o11ybench-oswal-${{ github.run_id }}-${{ github.run_attempt }} @@ -157,12 +200,14 @@ jobs: run: | bash greptimedb/.github/scripts/pull-benchmark-image.sh greptime/minio:2024 docker compose -f greptimedb/tests-integration/fixtures/docker-compose.yml up -d --wait minio - - name: Pull images and write the plan + - name: Pull and check images, write the plan run: | bash greptimedb/.github/scripts/oswal-s3-bench-run.sh pull python3 greptimedb/.github/scripts/oswal-s3-bench.py plan --smoke "$SMOKE" --output "$PLAN" \ - --concurrency "$CONCURRENCY" --rows "$ROWS" --batch-rows "$BATCH_ROWS" --rounds "$ROUNDS" - python3 greptimedb/.github/scripts/oswal-s3-bench.py manifest --results "$RESULTS" \ + --concurrency "$CONCURRENCY" --rows "$ROWS" --batch-rows "$BATCH_ROWS" --rounds "$ROUNDS" \ + --configs "$CONFIGS" --extra-env "$EXTRA_ENV" --restart "$RESTART" --groups "$OSWAL_GROUPS" \ + --images "$OSWAL_IMAGES" --warmup "$WARMUP" + python3 greptimedb/.github/scripts/oswal-s3-bench.py manifest --results "$RESULTS" --images "$GREPTIMEDB_IMAGES" \ --runtime-image "$RUNTIME_IMAGE" --instance-type "$INSTANCE_TYPE" --region "$S3_REGION" --prefix "$RUN_PREFIX" - name: Probe the bucket env: diff --git a/.github/workflows/query-regression-janitor.yml b/.github/workflows/query-regression-janitor.yml index ce277ce2f02..c615729206c 100644 --- a/.github/workflows/query-regression-janitor.yml +++ b/.github/workflows/query-regression-janitor.yml @@ -14,6 +14,7 @@ permissions: jobs: sweep: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} runs-on: ubuntu-latest timeout-minutes: 15 steps: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 0fb22f23e61..b292b586453 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -19,7 +19,6 @@ on: # NIGHTLY_RELEASE_MAX_AGE_DAYS days, see the 'Check nightly release necessity' step. - cron: '0 0 * * 1-5' workflow_dispatch: # Allows you to run this workflow manually. - # Notes: The GitHub Actions ONLY support 10 inputs; this final input is the tenth. inputs: linux_amd64_runner: type: choice @@ -57,9 +56,19 @@ on: description: Do not run integration tests during the build type: boolean default: true - build_linux_artifacts: + build_linux_amd64_artifacts: type: boolean - description: Build linux artifacts (amd64, arm64, and riscv64) + description: Build linux-amd64 artifacts (publishing always builds all Linux architectures) + required: false + default: false + build_linux_arm64_artifacts: + type: boolean + description: Build linux-arm64 artifacts (publishing always builds all Linux architectures) + required: false + default: false + build_linux_riscv64_artifacts: + type: boolean + description: Build linux-riscv64 artifacts (publishing always builds all Linux architectures) required: false default: false build_macos_artifacts: @@ -115,6 +124,9 @@ jobs: if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} runs-on: ubuntu-latest outputs: + build-linux-amd64: ${{ steps.linux-architectures.outputs.amd64 }} + build-linux-arm64: ${{ steps.linux-architectures.outputs.arm64 }} + build-linux-riscv64: ${{ steps.linux-architectures.outputs.riscv64 }} linux-amd64-runner: ${{ steps.start-linux-amd64-runner.outputs.label }} linux-arm64-runner: ${{ steps.start-linux-arm64-runner.outputs.label }} macos-runner: ${{ inputs.macos_runner || vars.DEFAULT_MACOS_RUNNER }} @@ -143,6 +155,17 @@ jobs: fetch-depth: 0 persist-credentials: false + # Publishing and automatic releases retain the full Linux build. + - name: Resolve Linux architectures + id: linux-architectures + env: + BUILD_LINUX_AMD64_ARTIFACTS: ${{ inputs.build_linux_amd64_artifacts }} + BUILD_LINUX_ARM64_ARTIFACTS: ${{ inputs.build_linux_arm64_artifacts }} + BUILD_LINUX_RISCV64_ARTIFACTS: ${{ inputs.build_linux_riscv64_artifacts }} + PUBLISH_GITHUB_RELEASE: ${{ inputs.publish_github_release }} + RELEASE_IMAGES: ${{ inputs.release_images }} + run: bash .github/scripts/release-linux-architectures.sh >> "$GITHUB_OUTPUT" + - name: Check Rust toolchain version shell: bash run: | @@ -182,7 +205,7 @@ jobs: ./.github/scripts/check-version.sh "${{ steps.create-version.outputs.version }}" - name: Allocate linux-amd64 runner - if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && steps.check-nightly.outputs.nightly-required == 'true' }} + if: ${{ (steps.linux-architectures.outputs.amd64 == 'true' || steps.linux-architectures.outputs.riscv64 == 'true') && steps.check-nightly.outputs.nightly-required == 'true' }} uses: ./.github/actions/start-runner id: start-linux-amd64-runner with: @@ -196,7 +219,7 @@ jobs: subnet-id: ${{ vars.EC2_RUNNER_SUBNET_ID }} - name: Allocate linux-arm64 runner - if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && steps.check-nightly.outputs.nightly-required == 'true' }} + if: ${{ steps.linux-architectures.outputs.arm64 == 'true' && steps.check-nightly.outputs.nightly-required == 'true' }} uses: ./.github/actions/start-runner id: start-linux-arm64-runner with: @@ -211,7 +234,7 @@ jobs: build-linux-amd64-artifacts: name: Build linux-amd64 artifacts - if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} + if: ${{ needs.allocate-runners.outputs.build-linux-amd64 == 'true' && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -237,7 +260,7 @@ jobs: build-linux-arm64-artifacts: name: Build linux-arm64 artifacts - if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} + if: ${{ needs.allocate-runners.outputs.build-linux-arm64 == 'true' && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -264,7 +287,7 @@ jobs: build-linux-riscv64-artifacts: name: Build linux-riscv64 artifacts - if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} + if: ${{ needs.allocate-runners.outputs.build-linux-riscv64 == 'true' && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -415,11 +438,15 @@ jobs: REQUESTED_VALIDATION_POLICY: ${{ inputs.release_validation }} run: | set -euo pipefail - base_commit="${GITHUB_SHA}" - if [[ "${GITHUB_REF_TYPE}" == "tag" ]]; then - base_commit="${GITHUB_SHA}^" - fi - previous_release_tag="$(git tag --merged "${base_commit}" --sort=-v:refname | grep -E '^v[0-9]+\.[0-9]+\.[0-9]+$' | head -n 1 || true)" + # Stable tags on older release branches are not ancestors of this + # commit. Pick the newest vX.Y.Z older than the version being released. + release_version="${{ needs.allocate-runners.outputs.version }}" + previous_release_tag="$( + git ls-remote --tags origin 'refs/tags/v*' \ + | awk '{print $2}' \ + | sed -e 's#^refs/tags/##' \ + | python3 .github/scripts/previous-release-tag.py --version "${release_version}" + )" if [[ -z "${previous_release_tag}" ]]; then previous_release_tag="main" fi @@ -489,9 +516,9 @@ jobs: fetch-depth: 0 persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: diff --git a/.github/workflows/run-multi-lang-tests.yml b/.github/workflows/run-multi-lang-tests.yml index 52fdca2117c..dad4fbaa080 100644 --- a/.github/workflows/run-multi-lang-tests.yml +++ b/.github/workflows/run-multi-lang-tests.yml @@ -83,10 +83,10 @@ jobs: java-version: '17' cache: 'maven' - - name: Setup Python 3.8 + - name: Setup Python 3.13 uses: actions/setup-python@v5 with: - python-version: '3.8' + python-version: '3.13' - name: Setup uv uses: astral-sh/setup-uv@v3 diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index e29b69678c1..d1b93a60e43 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -40,9 +40,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: components: rustfmt @@ -58,9 +58,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: components: clippy @@ -102,9 +102,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - name: Install riscv64 cross toolchain run: sudo apt-get update && sudo apt-get install -y gcc-riscv64-linux-gnu g++-riscv64-linux-gnu - uses: actions-rust-lang/setup-rust-toolchain@v1 @@ -130,9 +130,9 @@ jobs: with: persist-credentials: false - uses: ./.github/actions/setup-cyborg - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: KyleMayes/install-llvm-action@v1 with: version: "14.0" @@ -166,9 +166,9 @@ jobs: with: persist-credentials: false - uses: ./.github/actions/setup-cyborg - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: KyleMayes/install-llvm-action@v1 with: version: "14.0" @@ -210,9 +210,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: rui314/setup-mold@v1 - name: Install toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 @@ -265,9 +265,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: rui314/setup-mold@v1 - name: Install toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 @@ -294,6 +294,7 @@ jobs: env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 + RUST_MIN_STACK: 8388608 # 8MB, aligned with the test job CARGO_INCREMENTAL: 0 # Runs the object-store matrix against the local MinIO from # tests-integration/fixtures. Real AWS S3 is covered by the nightly diff --git a/Makefile b/Makefile index 7f92098262f..21c87cc3707 100644 --- a/Makefile +++ b/Makefile @@ -262,6 +262,7 @@ run-it-in-container: start-etcd ## Run integration tests in dev-builder. docker run --network=host \ -v ${PWD}:/greptimedb -v ${CARGO_REGISTRY_CACHE}:/root/.cargo/registry -v ${CARGO_GIT_CACHE}:/root/.cargo/git -v /tmp:/tmp \ -e GT_S3_BUCKET -e GT_S3_ACCESS_KEY_ID -e GT_S3_ACCESS_KEY -e GT_S3_REGION \ + -e RUST_MIN_STACK \ -w /greptimedb ${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-${BASE_IMAGE}:${DEV_BUILDER_IMAGE_TAG} \ make test sqlness-test BUILD_JOBS=${BUILD_JOBS} diff --git a/docs/how-to/metric-snapshot-s3-bench.md b/docs/how-to/metric-snapshot-s3-bench.md new file mode 100644 index 00000000000..86468bf3ffd --- /dev/null +++ b/docs/how-to/metric-snapshot-s3-bench.md @@ -0,0 +1,79 @@ +# Metric snapshot benchmark on AWS S3 + +The manual `Metric Snapshot S3 Benchmark` workflow tests packed Metric export and +restore against real S3 in `ap-southeast-1`. It must first exist on the repository's +default branch for GitHub to register `workflow_dispatch`. It was validated on AWS +by run [38544464731](https://github.com/GreptimeTeam/greptimedb/actions/runs/38544464731); +[#9573](https://github.com/GreptimeTeam/greptimedb/pull/9573) records its conditions and results. + +## Run + +Use the repository's existing EC2 runner configuration and CI S3 credentials. +The workflow passes credentials only through CI; no local AWS credentials are +needed. Source, target and CLI use the same immutable image. This image's source +SHA and required CLI capabilities have been checked: + +```bash +gh workflow run metric-snapshot-s3-bench.yml \ + --repo GreptimeTeam/greptimedb --ref main \ + -f greptimedb_image=greptime/greptimedb-nightly@sha256:20a1182bed0f452e6d573ca110b313e0b5cb76de7aedfa41a0c339963d4683d3 \ + -f expected_commit=f690bbb2040865c0d22b0b57d51612b60ae27cd0 \ + -f logical_tables_per_group=10000 \ + -f benchmark_timeout_minutes=240 \ + -f prior_spend_usd=0 +``` + +Start a new dispatch for each attempt and update `prior_spend_usd` to include all +earlier cloud attempts in the campaign, including failed setup and cleanup. +Do not reset the estimate to zero when retrying. The workflow accepts 50–10,000 +logical tables per group and a 30–240 minute benchmark job timeout. + +The runner is fixed at `c6id.4xlarge` (16 vCPU, 32 GiB RAM). At the recorded +Singapore Linux on-demand estimate of $0.9408/hour, the default attempt reserves +$6.4392: four benchmark hours, 45 minutes of provisioning, 10 minutes for queueing, +20 minutes of cleanup, and $1.50 for S3, disk and IPv4 costs. Preflight rejects prior spend at or above +$15, or an attempt whose reserve would exceed the cumulative $20 budget. +These estimates and checks are not an AWS hard spending cap or an AWS bill. + +## Workload and validation + +- A 10-physical-table × 50-logical-table smoke runs first. Its 500 logical tables + cover all 298 string tag columns before the larger workload starts. +- The default full run uses 10 physical tables and 100,000 logical tables, with + 32 rows per logical table: 3.2 million rows. Each logical schema has 10 + overlapping tags, a timestamp and a DOUBLE; schemas vary between tables. +- Source and target use local File storage on the temporary NVMe disk. Only + snapshots use S3. Source validation finishes before the source process stops; + restoration then starts a separate target process. +- Export uses `--pack-metric-data`, parallelism 4, chunk parallelism 1, and one + complete time chunk. Import uses task parallelism 1 and automatic batch DDL. +- Every source and restored logical table is checked for membership, schema, + DDL and typed values including NULLs. Target placeholders force table ID + remapping; successful batch DDL and removal of import resume state are checked. + +## Results and cleanup + +The job summary reports complete export, structural verify and import durations; +data generation and full data validation are timed separately. Evidence artifacts +are credential-scanned and retained for 14 days. Keep conclusions limited to +this version's absolute AWS results; this workflow does not compare old versions. + +Artifacts include sampled CPU/RSS, host memory, swap/OOM counters, server metrics, +CLI logs and S3 object/byte inventories. Configured swap is not RAM. A phase with +new swap-out or OOM fails and retains its evidence. OpenDAL operation counters +are not complete AWS HTTP request counts: multipart calls, retries and CLI traffic +are not fully covered. Object inventory bytes are not transferred bytes. + +A separate GitHub-hosted teardown job starts when provisioning finishes and polls +the benchmark job. Its deadline includes queue time: allocation start plus 45 +minutes, the benchmark timeout and 10 minutes. Completion, cancellation, an API +error or the deadline leads to cleanup. Failed provisioning skips the wait. +It terminates the EC2 instance and its disks, unregisters the runner, deletes every +object listed under the run's unique prefix and verifies the listing is empty. Like +the object store WAL benchmark, it needs only list, put, get and delete on the CI +bucket. It does not abort incomplete multipart uploads or remove noncurrent object +versions; those depend on the bucket's lifecycle configuration. Check teardown completion before another +attempt. A deadline may leave the GitHub benchmark job queued after its resources +are removed; cancel that workflow manually. If cleanup fails, rerun the teardown +job only; it keeps the original attempt's resource identity and deadline. A new benchmark needs a new dispatch because teardown +removes its runner. The daily EC2 janitor is only a fallback and does not clean S3. diff --git a/src/promql/src/aggregate.rs b/src/promql/src/aggregate.rs index 228f67ba4bc..0ab1e26a291 100644 --- a/src/promql/src/aggregate.rs +++ b/src/promql/src/aggregate.rs @@ -27,11 +27,14 @@ pub use state::{ use std::cmp::Reverse; use std::collections::hash_map::Entry; use std::collections::{BTreeMap, BinaryHeap, HashMap}; +use std::mem::size_of; use std::sync::Arc; use datafusion::error::{DataFusionError, Result}; +use datatypes::arrow::array::Float64Array; use datatypes::series::MetricName; -use datatypes::series::samples::{SampleRef, Series, SeriesSchema}; +use datatypes::series::samples::{SampleRef, Samples, Series, SeriesSchema}; +use datatypes::series::time::{GridSpan, TimeAxis}; use crate::go_format::format_float_f; use crate::grid_series::{GridSeriesBuilder, external}; @@ -97,6 +100,7 @@ pub struct FloatAggregator { labels: Vec>, group_index: HashMap, groups: Vec, + retained: usize, } impl FloatAggregator { @@ -104,25 +108,78 @@ impl FloatAggregator { schema: &Arc, op: AggregateOp, grouping: Grouping<'_>, + ) -> Result { + Self::try_new_budgeted(schema, op, grouping, |_| Ok(())) + } + + pub(crate) fn try_new_budgeted( + schema: &Arc, + op: AggregateOp, + grouping: Grouping<'_>, + mut budget: impl FnMut(usize) -> Result<()>, ) -> Result { validate_input(schema, &[], grouping)?; let (by, labels) = match grouping { Grouping::By(labels) => (true, labels), Grouping::Without(labels) => (false, labels), }; + let label_storage = + grouping_bytes(grouping).saturating_add(labels_capacity_bytes(grouping)); + budget(label_storage)?; + let mut owned_labels: Vec> = Vec::new(); + owned_labels + .try_reserve_exact(labels.len()) + .map_err(|error| { + DataFusionError::Execution(format!( + "cannot reserve aggregate grouping labels: {error}" + )) + })?; + owned_labels.extend(labels.iter().map(|label| Arc::from(*label))); + let retained = owned_labels + .capacity() + .saturating_mul(size_of::>()) + .saturating_add(owned_labels.iter().fold(0usize, |bytes, label| { + bytes.saturating_add(label.len() + 2 * size_of::()) + })); Ok(Self { schema: schema.clone(), op, by, - labels: labels.iter().map(|label| Arc::from(*label)).collect(), + labels: owned_labels, group_index: HashMap::new(), groups: Vec::new(), + retained, }) } /// Folds `series` into its group. A series that fails validation leaves the state unchanged. + /// + /// A budget refusal is not a validation failure: the fold may have opened a group or filled + /// blocks before the allocation it could not admit, so a failed [`Self::update_budgeted`] + /// leaves the state half folded and its retained figure stale. The runtime stops the execution + /// and drops the aggregator on that error, which is why no rollback is kept; a caller holding + /// one must drop it rather than fold into it again. pub fn update(&mut self, series: &Series) -> Result<()> { + self.update_budgeted(series, |_| Ok(())) + } + + pub(crate) fn update_budgeted( + &mut self, + series: &Series, + mut budget: impl FnMut(usize) -> Result<()>, + ) -> Result<()> { validate_float_series(&self.schema, series)?; + let existing_bytes = self.state_bytes(); + let grouping_workspace = grouping_workspace_bytes(&series.metric); + let mut workspace_bytes = + grouping_workspace.saturating_add(pending_bytes(self.schema.fields.len())); + budget(existing_bytes.saturating_add(workspace_bytes))?; + workspace_bytes = workspace_bytes.saturating_add( + self.labels + .len() + .saturating_mul(std::mem::size_of::<&str>()), + ); + budget(existing_bytes.saturating_add(workspace_bytes))?; let labels: Vec<&str> = self.labels.iter().map(|label| &**label).collect(); let grouping = if self.by { Grouping::By(&labels) @@ -130,39 +187,210 @@ impl FloatAggregator { Grouping::Without(&labels) }; let metric = grouping.apply(&series.metric)?; - let index = match self.group_index.entry(label_key(&metric)) { + let key_bytes = metric + .tags() + .len() + .saturating_add(1) + .saturating_mul(std::mem::size_of::<(Arc, Arc)>()); + workspace_bytes = workspace_bytes + .saturating_add(grouping_workspace_bytes(&metric)) + .saturating_add(key_bytes); + budget(existing_bytes.saturating_add(workspace_bytes))?; + let key = label_key(&metric); + let key_retained = key + .capacity() + .saturating_mul(size_of::<(Arc, Arc)>()); + let new_group = !self.group_index.contains_key(&key); + let group_capacity = self.groups.capacity(); + let map_capacity = self.group_index.capacity(); + let old_map_backing = hash_map_backing_bytes(map_capacity); + if new_group { + let next_group_capacity = if self.groups.len() == group_capacity { + group_capacity.saturating_mul(2).max(4) + } else { + group_capacity + }; + let group_realloc_bytes = if next_group_capacity > group_capacity { + next_group_capacity.saturating_mul(size_of::()) + } else { + 0 + }; + budget( + existing_bytes + .saturating_add(metric_bytes(&metric)) + .saturating_add(key_bytes) + // A realloc that cannot grow in place keeps both buffers alive while it + // copies, so a growing step admits the old buffer beside the new one -- + // the same reason the map below admits both tables. The old buffer alone + // is already in `existing_bytes`; what is added here is the new one. + .saturating_add(group_realloc_bytes) + .saturating_add(if self.group_index.len() == map_capacity { + hash_map_backing_bytes(next_hash_capacity(map_capacity)) + } else { + 0 + }) + .saturating_add(workspace_bytes), + )?; + } + let group_capacity_growth = if new_group && self.groups.len() == group_capacity { + group_capacity.saturating_mul(2).max(4) + } else { + group_capacity + }; + if new_group && group_capacity_growth > group_capacity { + self.groups.reserve(1); + } + if new_group && self.group_index.len() == self.group_index.capacity() { + self.group_index.reserve(1); + } + let index = match self.group_index.entry(key) { Entry::Occupied(entry) => *entry.get(), Entry::Vacant(entry) => { self.groups.push(Group { - metric, + metric: metric.clone(), blocks: BTreeMap::new(), }); - *entry.insert(self.groups.len() - 1) + let index = self.groups.len() - 1; + entry.insert(index); + index } }; - self.groups[index].accumulate(series, self.op); + if new_group { + self.retained = self + .retained + .saturating_add(metric_bytes(&metric)) + .saturating_add(key_retained) + .saturating_add( + self.groups + .capacity() + .saturating_sub(group_capacity) + .saturating_mul(size_of::()), + ) + .saturating_add( + hash_map_backing_bytes(self.group_index.capacity()) + .saturating_sub(old_map_backing), + ); + } + drop(labels); + drop(metric); + self.retained = self.groups[index].accumulate( + series, + self.op, + &mut budget, + self.retained, + self.schema.fields.len(), + )?; Ok(()) } + /// Number of groups currently retained by this fold. + pub(crate) fn groups(&self) -> usize { + self.groups.len() + } + + /// Bytes reserved for this fold, using bounded allowances for its maps and group payloads. + pub(crate) fn state_bytes(&self) -> usize { + self.retained + } + + #[cfg(test)] + pub(crate) fn recounted_state_bytes(&self) -> usize { + let groups = self.groups.iter().fold(0usize, |total, group| { + total + .saturating_add(metric_bytes(&group.metric)) + .saturating_add( + group + .blocks + .len() + .saturating_mul(BLOCK_POINTS) + .saturating_mul(self.schema.fields.len()) + .saturating_mul(std::mem::size_of::()), + ) + .saturating_add(btree_backing_bytes(group.blocks.len())) + }); + groups + .saturating_add( + self.labels + .capacity() + .saturating_mul(std::mem::size_of::>()), + ) + .saturating_add(hash_map_backing_bytes(self.group_index.capacity())) + .saturating_add(self.group_index.keys().fold(0usize, |bytes, key| { + bytes.saturating_add( + key.capacity() + .saturating_mul(size_of::<(Arc, Arc)>()), + ) + })) + .saturating_add( + self.labels + .iter() + .map(|label| label.len() + 2 * std::mem::size_of::()) + .sum::(), + ) + .saturating_add( + self.groups + .capacity() + .saturating_mul(std::mem::size_of::()), + ) + } + /// The result of every group, in the order the groups first appeared. pub fn finish(self) -> Result> { - finish_groups(self.schema, self.op, self.groups) + self.finish_budgeted(|_| Ok(())) + } + + pub(crate) fn finish_budgeted( + self, + budget: impl FnMut(usize) -> Result<()>, + ) -> Result> { + let state = self.state_bytes(); + finish_groups_budgeted(self.schema, self.op, self.groups, state, budget) } } -/// Turns the state of `groups` into result series, in group order. -fn finish_groups( +/// Turns the state of `groups` into result series, in group order, admitting as it builds. +fn finish_groups_budgeted( schema: Arc, op: AggregateOp, groups: Vec, + state: usize, + mut budget: impl FnMut(usize) -> Result<()>, ) -> Result> { let width = schema.fields.len(); + let group_count = groups.len(); + if group_count == 0 { + return Ok(Vec::new()); + } + let groups = groups.into_iter(); let mut builder = GridSeriesBuilder::try_new(schema)?; - let mut output = Vec::with_capacity(groups.len()); - let mut row = Vec::with_capacity(width); + let output_capacity_bytes = group_count.saturating_mul(std::mem::size_of::()); + let row_capacity_bytes = width.saturating_mul(std::mem::size_of::>()); + let base = state + .saturating_add(output_capacity_bytes) + .saturating_add(row_capacity_bytes); + budget(base)?; + let mut output = Vec::new(); + output.try_reserve_exact(group_count).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve aggregate outputs: {error}")) + })?; + let mut row = Vec::new(); + row.try_reserve_exact(width).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve aggregate row: {error}")) + })?; + let mut completed_output_bytes = 0usize; for group in groups { // One builder serves every group, so each result series starts explicitly. builder.begin_series(); + let (rows, runs) = group_rows_runs(&group, width); + let finish_bytes = finish_working_bytes(rows, runs, width); + budget( + state + .saturating_add(output_capacity_bytes) + .saturating_add(row_capacity_bytes) + .saturating_add(completed_output_bytes) + .saturating_add(finish_bytes), + )?; + builder.reserve_rows(rows)?; for (&block, states) in &group.blocks { // The last block can extend past the grid; only points with a sample are pushed. for (offset, point) in states.chunks_exact(width).enumerate() { @@ -174,11 +402,89 @@ fn finish_groups( builder.push(block * BLOCK_POINTS + offset, &row)?; } } - output.extend(builder.finish(group.metric)?); + let finished = builder.finish(group.metric)?; + completed_output_bytes = completed_output_bytes.saturating_add( + finished + .iter() + .map(crate::grid_series::series_bytes) + .sum::(), + ); + output.extend(finished); + budget( + state + .saturating_add(output_capacity_bytes) + .saturating_add(row_capacity_bytes) + .saturating_add(completed_output_bytes), + )?; } Ok(output) } +/// The result of every group, without a budget: the entry the partial-state merger takes, which +/// does not hold a query pool to admit against. +fn finish_groups( + schema: Arc, + op: AggregateOp, + groups: Vec, +) -> Result> { + finish_groups_budgeted(schema, op, groups, 0, |_| Ok(())) +} + +fn group_rows_runs(group: &Group, width: usize) -> (usize, usize) { + if width == 0 { + return (0, 0); + } + let mut rows = 0usize; + let mut runs = 0usize; + let mut previous: Option = None; + for (&block, states) in &group.blocks { + for (offset, point) in states.chunks_exact(width).enumerate() { + if point.iter().any(|state| state.seen) { + let index = block.saturating_mul(BLOCK_POINTS).saturating_add(offset); + rows = rows.saturating_add(1); + if previous.is_none_or(|previous| previous.checked_add(1) != Some(index)) { + runs = runs.saturating_add(1); + } + previous = Some(index); + } + } + } + (rows, runs) +} + +fn finish_working_bytes(rows: usize, runs: usize, width: usize) -> usize { + let values = rows + .saturating_mul(width) + .saturating_mul(std::mem::size_of::()); + let validity_bytes = rows + .saturating_add(7) + .checked_div(8) + .unwrap_or(0) + .div_ceil(64) + .saturating_mul(64); + let fields = crate::grid_series::field_builder_bytes(width).saturating_add( + width.saturating_mul( + std::mem::size_of::() + .saturating_add(std::mem::size_of::()) + .saturating_add(std::mem::size_of::()), + ), + ); + let indices = rows.saturating_mul(std::mem::size_of::()); + let spans = runs.saturating_mul(std::mem::size_of::()); + let axis_conversion = runs + .saturating_mul( + 2 * std::mem::size_of::() + std::mem::size_of::>(), + ) + .saturating_add(rows.saturating_mul(std::mem::size_of::())); + values + .saturating_add(width.saturating_mul(validity_bytes)) + .saturating_add(fields) + .saturating_add(indices) + .saturating_add(spans) + .saturating_add(axis_conversion) + .saturating_add(std::mem::size_of::() + std::mem::size_of::()) +} + /// `quantile(q, v)`, with `q` read at each grid point from an unlabelled single-field scalar on /// the same grid, as Prometheus evaluates the parameter at every step. A point where `q` has no /// value has no result; `q` outside `[0, 1]` or NaN gives what Prometheus' `quantile` gives. @@ -438,6 +744,105 @@ impl Grouping<'_> { type LabelKey = Vec<(Arc, Arc)>; +fn metric_bytes(metric: &MetricName) -> usize { + metric + .name() + .map_or(0, str::len) + .saturating_add( + metric + .tags() + .iter() + .map(|tag| tag.name.len() + tag.value.len()) + .sum::(), + ) + .saturating_add( + metric + .tags() + .len() + .saturating_mul(std::mem::size_of::<(Arc, Arc)>()), + ) + .saturating_add(std::mem::size_of::()) +} + +fn grouping_workspace_bytes(metric: &MetricName) -> usize { + metric_bytes(metric) + .saturating_add( + metric + .tags() + .len() + .saturating_add(1) + .saturating_mul(std::mem::size_of::<(Arc, Arc)>()), + ) + .saturating_add(std::mem::size_of::()) +} + +/// Conservative backing estimate for Rust 1.98.1's insert-only `HashMap` (hashbrown 0.17.1). +/// Bucket count follows its 7/8 load factor; entries, one control byte per bucket and the +/// 16-byte control-group tail are counted. This models requested table storage, not allocator slack. +/// An upper bound on the bytes the group index holds for `capacity` entries. +/// +/// The table keeps its entries in a power-of-two number of buckets below `16/7` of the usable +/// capacity, with one control byte beside each bucket. Three buckets per entry therefore covers the +/// table whatever exact growth the map uses, which is what keeps this a bound rather than a model of +/// one implementation: a toolchain that changes how the map grows cannot make it under-count. +fn hash_map_backing_bytes(capacity: usize) -> usize { + capacity + .saturating_mul(3) + .saturating_mul(size_of::<(LabelKey, usize)>() + 1) +} + +/// An upper bound on the capacity a map holding `capacity` entries grows to. +fn next_hash_capacity(capacity: usize) -> usize { + capacity.saturating_mul(2).max(7) +} + +/// An upper bound on the bytes one B-tree node holds. +/// +/// A node keeps at most eleven key/value pairs and twelve child pointers, plus a header of a few +/// words; the leaf form is smaller than that. The header allowance is generous on purpose, so the +/// bound does not depend on how a toolchain lays the node out. +fn btree_node_bytes() -> usize { + 11 * size_of::<(usize, Box<[FloatState]>)>() + 16 * size_of::() +} + +/// An upper bound on the bytes a B-tree of `entries` holds. +/// +/// Every node but the root keeps at least five entries, so `entries / 5` nodes is enough to hold +/// them; a tree of one entry is one node. +fn btree_backing_bytes(entries: usize) -> usize { + if entries == 0 { + return 0; + } + entries.div_ceil(5).saturating_mul(btree_node_bytes()) +} + +fn pending_bytes(fields: usize) -> usize { + working_set_bytes(fields, BLOCK_POINTS) +} + +fn working_set_bytes(fields: usize, pending_capacity: usize) -> usize { + pending_capacity + .saturating_mul(std::mem::size_of::<(usize, f64)>()) + .saturating_add(fields.saturating_mul(std::mem::size_of::>() + 1)) +} + +fn labels_capacity_bytes(grouping: Grouping<'_>) -> usize { + let labels = match grouping { + Grouping::By(labels) | Grouping::Without(labels) => labels, + }; + labels.len().saturating_mul(std::mem::size_of::>()) +} + +fn grouping_bytes(grouping: Grouping<'_>) -> usize { + let labels = match grouping { + Grouping::By(labels) | Grouping::Without(labels) => labels, + }; + labels + .iter() + .map(|label| label.len() + std::mem::size_of::>()) + .sum::() +} + fn label_key(metric: &MetricName) -> LabelKey { let mut labels = Vec::with_capacity(metric.tags().len() + 1); if let Some(name) = metric.name() { @@ -463,9 +868,29 @@ struct Group { blocks: BTreeMap>, } +/// Which field of which layout a fold is writing, and how it folds it. +/// +/// One value carried through the per-block work instead of three separate arguments: the three +/// never vary independently within one series, and the block loop is the only place that reads +/// them. +#[derive(Clone, Copy)] +struct FieldFold { + width: usize, + field_index: usize, + op: AggregateOp, +} + impl Group { - fn accumulate(&mut self, series: &Series, op: AggregateOp) { - let width = series.fields.len(); + fn accumulate( + &mut self, + series: &Series, + op: AggregateOp, + budget: &mut impl FnMut(usize) -> Result<()>, + mut retained: usize, + width: usize, + ) -> Result { + let pending_capacity_bytes = BLOCK_POINTS.saturating_mul(size_of::<(usize, f64)>()); + budget(retained.saturating_add(pending_capacity_bytes))?; let mut pending = Vec::with_capacity(BLOCK_POINTS); for (field_index, field) in series.fields.iter().enumerate() { // The axis says which absolute positions the field holds and the field iterator @@ -475,35 +900,81 @@ impl Group { let mut block = usize::MAX; for ((index, _), sample) in series.time.positions().zip(field.iter()) { if index / BLOCK_POINTS != block { - self.update_block(block, width, field_index, op, &mut pending); + retained = self.update_block( + block, + FieldFold { + width, + field_index, + op, + }, + &mut pending, + budget, + retained, + )?; block = index / BLOCK_POINTS; } if let SampleRef::Float(sample) = sample { pending.push((index % BLOCK_POINTS, sample)); } } - self.update_block(block, width, field_index, op, &mut pending); + retained = self.update_block( + block, + FieldFold { + width, + field_index, + op, + }, + &mut pending, + budget, + retained, + )?; } + Ok(retained) } /// Applies one field's samples within a block, so a block costs one lookup per field run. fn update_block( &mut self, block: usize, - width: usize, - field_index: usize, - op: AggregateOp, + fold: FieldFold, pending: &mut Vec<(usize, f64)>, - ) { + budget: &mut impl FnMut(usize) -> Result<()>, + retained: usize, + ) -> Result { if pending.is_empty() { - return; + return Ok(retained); + } + let new_block = !self.blocks.contains_key(&block); + let block_bytes = if new_block { + BLOCK_POINTS + .saturating_mul(fold.width) + .saturating_mul(std::mem::size_of::()) + } else { + 0 + }; + let backing_delta = if new_block { + btree_backing_bytes(self.blocks.len() + 1) + .saturating_sub(btree_backing_bytes(self.blocks.len())) + } else { + 0 + }; + if new_block { + budget( + retained + .saturating_add(backing_delta) + .saturating_add(block_bytes) + .saturating_add(pending.capacity().saturating_mul(size_of::<(usize, f64)>())), + )?; } let states = self.blocks.entry(block).or_insert_with(|| { - vec![FloatState::default(); BLOCK_POINTS * width].into_boxed_slice() + vec![FloatState::default(); BLOCK_POINTS * fold.width].into_boxed_slice() }); for (offset, sample) in pending.drain(..) { - states[offset * width + field_index].update(op, sample); + states[offset * fold.width + fold.field_index].update(fold.op, sample); } + Ok(retained + .saturating_add(backing_delta) + .saturating_add(block_bytes)) } } diff --git a/src/promql/src/aggregate/tests.rs b/src/promql/src/aggregate/tests.rs index 31082bd784a..0fd86561fdd 100644 --- a/src/promql/src/aggregate/tests.rs +++ b/src/promql/src/aggregate/tests.rs @@ -14,6 +14,7 @@ use std::collections::BTreeMap; +use datafusion::error::DataFusionError; use datatypes::arrow::array::{Float64Array, UInt32Array}; use datatypes::series::Label; use datatypes::series::histogram::HistogramBatchBuilder; @@ -345,6 +346,27 @@ fn avg_compensation_matches_prometheus_bits() { } } +#[test] +fn one_field_missing_and_another_nan_keep_row_nullability() { + let schema = schema(&["missing", "value"]); + let series = input( + &schema, + metric("m", &[]), + 0, + &[ + vec![None, None, None], + vec![Some(1.0), Some(f64::NAN), Some(2.0)], + ], + ); + let output = aggregate_float(&schema, &[series], AggregateOp::Sum, Grouping::By(&[])).unwrap(); + assert_eq!(output.len(), 1); + assert_eq!(values(&output[0], 0).1, vec![None, None, None]); + let values = values(&output[0], 1).1; + assert_eq!(values[0], Some(1.0)); + assert!(values[1].unwrap().is_nan()); + assert_eq!(values[2], Some(2.0)); +} + #[test] fn fields_are_aggregated_independently() { // A field without input values at a row is null there, even for count and group; the row @@ -476,6 +498,301 @@ fn fields_are_aggregated_independently() { } } +#[test] +fn incremental_retained_bytes_match_full_recount_after_updates() { + let schema = schema_with_grid_len(&["value", "other"], 256); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["zone", "rack"])) + .unwrap(); + assert_eq!( + aggregator.state_bytes(), + aggregator.recounted_state_bytes(), + "initial state" + ); + let updates = [ + input( + &schema, + metric("m", &[("zone", "a"), ("rack", "1")]), + 0, + &[vec![Some(1.0)], vec![Some(2.0)]], + ), + input( + &schema, + metric("m", &[("zone", "a"), ("rack", "1")]), + 64, + &[vec![Some(3.0)], vec![Some(4.0)]], + ), + input( + &schema, + metric("m", &[("zone", "a"), ("rack", "1")]), + 0, + &[vec![Some(5.0)], vec![Some(6.0)]], + ), + input( + &schema, + metric("m", &[("zone", "b"), ("rack", "2")]), + 128, + &[vec![Some(7.0)], vec![Some(8.0)]], + ), + ]; + for (index, series) in updates.iter().enumerate() { + let previous = aggregator.state_bytes(); + aggregator.update(series).unwrap(); + assert_eq!( + aggregator.state_bytes(), + aggregator.recounted_state_bytes(), + "update {index}" + ); + if index == 2 { + assert_eq!(aggregator.state_bytes(), previous); + } + } +} + +#[test] +fn incremental_retained_bytes_match_full_recount_for_many_groups() { + let schema = schema(&["value"]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["id"])).unwrap(); + for id in 0..512 { + let id = id.to_string(); + let series = input( + &schema, + metric("m", &[("id", &id)]), + id.parse::().unwrap() % 64, + &[vec![Some(1.0)]], + ); + aggregator.update(&series).unwrap(); + assert_eq!( + aggregator.state_bytes(), + aggregator.recounted_state_bytes(), + "group {id}" + ); + } +} + +/// The unbudgeted entry point is the budgeted one with a no-op budget, so this compares one path +/// with itself: it is a smoke check that the wrapper still runs, not evidence that the kernels are +/// right. The numbers are checked against Prometheus in `aggregates_match_prometheus` and +/// `stat_aggregates_match_prometheus`. +#[test] +fn every_supported_kind_runs_through_the_unbudgeted_entry_point() { + let schema = schema(&["value"]); + let series = input(&schema, metric("m", &[]), 0, &[vec![Some(1.0), Some(3.0)]]); + for op in [ + AggregateOp::Sum, + AggregateOp::Count, + AggregateOp::Avg, + AggregateOp::Min, + AggregateOp::Max, + AggregateOp::Group, + AggregateOp::Stddev, + AggregateOp::Stdvar, + ] { + let mut plain = FloatAggregator::try_new(&schema, op, Grouping::By(&[])).unwrap(); + plain.update(&series).unwrap(); + let plain = plain.finish().unwrap(); + let mut budgeted = + FloatAggregator::try_new_budgeted(&schema, op, Grouping::By(&[]), |_| Ok(())).unwrap(); + budgeted.update_budgeted(&series, |_| Ok(())).unwrap(); + let budgeted = budgeted.finish_budgeted(|_| Ok(())).unwrap(); + assert_eq!( + values_of(&plain, 0, 128), + values_of(&budgeted, 0, 128), + "{op:?}" + ); + } +} + +#[test] +fn repeated_updates_in_one_group_and_block_keep_one_block_charge() { + let schema = schema(&["value"]); + let first = input(&schema, metric("m", &[]), 0, &[vec![Some(1.0)]]); + let second = input(&schema, metric("m", &[]), 1, &[vec![Some(3.0)]]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&[])).unwrap(); + aggregator.update(&first).unwrap(); + let retained = aggregator.state_bytes(); + aggregator.update(&second).unwrap(); + assert_eq!(aggregator.state_bytes(), retained); +} + +#[test] +fn adjacent_blocks_are_admitted_and_accounted_separately() { + let schema = schema(&["value"]); + let series = input(&schema, metric("m", &[]), 63, &[vec![Some(1.0), Some(2.0)]]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&[])).unwrap(); + let mut admissions = Vec::new(); + aggregator + .update_budgeted(&series, |bytes| { + admissions.push(bytes); + Ok(()) + }) + .unwrap(); + assert_eq!(aggregator.groups(), 1); + assert_eq!(aggregator.groups[0].blocks.len(), 2); + assert_eq!( + admissions.last().copied(), + Some( + aggregator + .state_bytes() + .saturating_add(std::mem::size_of::<(usize, f64)>() * BLOCK_POINTS,) + ) + ); +} + +#[test] +fn a_rejected_block_allocation_does_not_insert_a_block() { + let schema = schema(&["value"]); + let series = input(&schema, metric("m", &[]), 0, &[vec![Some(1.0)]]); + // The figure a new block's admission adds: the block's states and its pending buffer, on + // top of whatever the fold already holds. Located by the value, not by which call reaches it, + // so reordering the earlier admissions cannot make this test pass without testing anything. + let block_states = BLOCK_POINTS + .saturating_mul(schema.fields.len()) + .saturating_mul(std::mem::size_of::()); + let pending = BLOCK_POINTS * std::mem::size_of::<(usize, f64)>(); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&[])).unwrap(); + let mut refused = None; + let result = aggregator.update_budgeted(&series, |bytes| { + // Everything up to the block's own admission is admitted; only the block is refused. + if bytes >= block_states.saturating_add(pending) { + refused = Some(bytes); + return Err(DataFusionError::Execution("block allocation budget".into())); + } + Ok(()) + }); + assert!( + result.is_err(), + "the block was admitted with no budget left" + ); + let refused = refused.expect("the refusal did not come from a block admission"); + assert!( + refused >= block_states.saturating_add(pending), + "the refusal came from an admission smaller than a block" + ); + assert_eq!(aggregator.groups(), 1); + assert!( + aggregator.groups[0].blocks.is_empty(), + "a refused block was inserted anyway" + ); +} + +#[test] +fn a_map_growth_rejection_does_not_insert_the_new_group() { + let schema = schema(&["value"]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["i"])).unwrap(); + let first = input(&schema, metric("m", &[("i", "0")]), 0, &[vec![Some(1.0)]]); + aggregator.update(&first).unwrap(); + let mut index = 1; + while aggregator.group_index.len() < aggregator.group_index.capacity() { + let label = index.to_string(); + let series = input( + &schema, + metric("m", &[("i", &label)]), + 0, + &[vec![Some(1.0)]], + ); + aggregator.update(&series).unwrap(); + index += 1; + } + let group_count = aggregator.groups(); + let old_state = aggregator.state_bytes(); + assert_eq!(aggregator.group_index.capacity(), group_count); + + // A resize keeps the old table while it fills the new one, so the figure the growth has to + // admit is the whole new table on top of the old state -- not the difference between them. + // That is the property asserted here: the refusal threshold is the state plus a whole new + // table, so the test fails if the admission ever drops to the delta. + let new_table = hash_map_backing_bytes(next_hash_capacity(group_count)); + let threshold = old_state.saturating_add(new_table); + let label = index.to_string(); + let series = input( + &schema, + metric("m", &[("i", &label)]), + 0, + &[vec![Some(1.0)]], + ); + let mut peak = 0; + let result = aggregator.update_budgeted(&series, |bytes| { + peak = peak.max(bytes); + if bytes >= threshold { + Err(DataFusionError::Execution("old+new map budget".into())) + } else { + Ok(()) + } + }); + assert!( + result.is_err(), + "the map grew with no budget for both tables" + ); + assert!( + peak >= threshold, + "the growth admitted {peak}, less than the old state plus a new table ({threshold})" + ); + assert_eq!(aggregator.group_index.capacity(), group_count); + assert_eq!(aggregator.groups(), group_count); + assert_eq!( + aggregator + .groups + .iter() + .map(|group| group.blocks.len()) + .sum::(), + group_count + ); +} + +#[test] +fn finish_rejection_after_first_group_is_budgeted_before_the_next() { + let schema = schema(&["value"]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["i"])).unwrap(); + for (label, value) in [("a", 1.0), ("b", 2.0)] { + let series = input( + &schema, + metric("m", &[("i", label)]), + 0, + &[vec![Some(value)]], + ); + aggregator.update(&series).unwrap(); + } + let mut calls = Vec::new(); + let result = aggregator.finish_budgeted(|bytes| { + calls.push(bytes); + if calls.len() == 4 { + Err(DataFusionError::Execution("budget".into())) + } else { + Ok(()) + } + }); + assert!(result.is_err()); + assert_eq!(calls.len(), 4); + assert!(calls[1] > calls[0], "first builder peak"); + assert!(calls[2] < calls[1], "first group settle"); + assert!(calls[3] > calls[2], "second builder peak"); + + // The base already includes the output Vec backing; later callbacks add only the + // builder workspace or completed Series allocation. + let builder_workspace = finish_working_bytes(1, 1, 1); + let expected_output = aggregate_float( + &schema, + &[ + input(&schema, metric("m", &[("i", "a")]), 0, &[vec![Some(1.0)]]), + input(&schema, metric("m", &[("i", "b")]), 0, &[vec![Some(2.0)]]), + ], + AggregateOp::Sum, + Grouping::By(&["i"]), + ) + .unwrap(); + let completed_series = crate::grid_series::series_bytes(&expected_output[0]); + assert_eq!(calls[1] - calls[0], builder_workspace); + assert_eq!(calls[2] - calls[0], completed_series); + assert_eq!(calls[3] - calls[2], builder_workspace); +} + #[test] fn state_blocks_keep_absolute_indices_across_boundaries_and_gaps() { // A 512-point grid spans eight state blocks; inputs straddle a block edge, leave whole diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index 750ee41b5e8..baf57a275c8 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -14,6 +14,7 @@ mod absent; mod empty_metric; +mod fused_series; mod histogram_fold; mod instant_manipulate; mod normalize; @@ -37,6 +38,7 @@ use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; use datatypes::data_type::DataType as _; pub use empty_metric::{EmptyMetric, EmptyMetricExec, EmptyMetricStream, build_special_time_expr}; +pub use fused_series::{FusedSeries, FusedSeriesBinding, FusedSeriesExec, FusedSeriesStream}; pub use histogram_fold::{ HistogramFold, HistogramFoldExec, HistogramFoldOperation, HistogramFoldStream, }; diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs new file mode 100644 index 00000000000..d957851902a --- /dev/null +++ b/src/promql/src/extension_plan/fused_series.rs @@ -0,0 +1,1555 @@ +// 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. + +//! One node that runs a compiled native tree over raw series. +//! +//! [`FusedSeries`] is the logical half: it holds the PromQL text the tree was compiled from, the +//! tree itself and the binding its input was compiled against. [`FusedSeriesExec`] is the physical +//! half: it owns one execution of that tree per partition and decodes its input frames into raw +//! series chunks and converts its evaluated chunks back into batches. +//! +//! The two halves show different things in `EXPLAIN`, which is the point of having both. The +//! logical node prints the query a reader wrote ([`FusedSeries`]'s `fmt_for_explain`), and the +//! physical node prints the tree it will run ([`FusedSeriesExec`]'s `fmt_as`, which embeds +//! [`NativeTree`]'s own display rather than a second renderer). A reader comparing the two plans +//! can therefore see whether a query was compiled the way they expected without decoding node +//! names. +//! +//! Scope of this layer: the input is a stream of native raw series frames +//! ([`raw_series_frame_schema`](datatypes::series::frame::raw_series_frame_schema)) — one row per +//! chunk, carrying the series' encoded key, its finality, its timestamps and its fields, plus the +//! complete labels when the frame is explicit — with each series' rows contiguous and its +//! timestamps ascending, which is what the caller's [`SeriesOrdering`] names. The output is one +//! batch per evaluated chunk, with the statement's grid timestamps in the declared time column. +//! Everything the tree cannot do is rejected when the execution plan is built, not when the first +//! batch arrives: a tree the native runtime does not execute, a frame whose schema the tree's input +//! was not compiled against, or an output schema the conversion cannot classify. + +use std::cmp::Ordering; +use std::collections::hash_map::Entry; +use std::collections::{HashMap, VecDeque}; +use std::fmt::{self, Formatter}; +use std::hash::{Hash, Hasher}; +use std::pin::Pin; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; + +use common_query::series_frame::{ensure_frame_matches_label_source, unpack_raw_series}; +use datafusion::arrow::array::{ + Array, ArrayRef, Float64Array, StringArray, TimestampMillisecondArray, +}; +use datafusion::arrow::buffer::ScalarBuffer; +use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::DFSchemaRef; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::context::TaskContext; +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool, +}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, + Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, +}; +use datatypes::series::SeriesOrdering; +use datatypes::series::raw::RawSeriesChunk; +use datatypes::series::samples::{DenseSamplesRef, SampleRef, Samples, Series, SeriesSchema}; +use datatypes::series::time::EvalGrid; +use futures::Stream; +use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream}; + +use crate::evaluated::EvaluatedChunk; +use crate::grid_series::{LabelSet, label_set}; +use crate::native::execute::{NativeExecutionState, NativeExecutor}; +use crate::native::metrics::{self, NodePath}; +use crate::native::plan::{NativeExpr, NativeTree}; +use crate::native::runs::{Run, absorb}; +use crate::source::materialize_raw_labels; + +#[cfg(test)] +mod tests; + +/// The label a metric name is stored under, which is not a tag column of the input. +const METRIC_NAME: &str = "__name__"; + +/// How many input steps one poll of [`FusedSeriesStream`] runs before it wakes itself. +/// +/// A step polls the input once and hands what it settled to the executor. Settling is not +/// guaranteed: raw samples before the first evaluation point settle nothing, so an input of many +/// such chunks would otherwise be read to its end inside one poll, holding the task's thread for +/// an unbounded time and delaying both other tasks on it and the cancellation the consumer asked +/// for. The budget makes that a poll that yields instead, and the waker makes sure it is polled +/// again, so nothing is lost by stopping early. +const POLL_BUDGET: usize = 64; + +/// What the tree's input was compiled against, and how its output columns are named. +/// +/// The binding travels with the tree: a tree is legal only against the schema it was bound to, and +/// a reader of the output needs to know which grid a chunk's absolute positions are on. Keeping it +/// in one value is what lets the logical node be built from the same description, and the physical +/// node reads the same facts back off the tree rather than trusting a second copy of them. +/// +/// What the logical node checks is that the schema the tree's input node declares is the one this +/// binding names. What it cannot check is whether the caller meant *some other* schema with the +/// same field names, or whether [`time_index`](Self::time_index) is the column the query reports +/// positions in: from here both are indistinguishable from a correct binding, so picking them is +/// the caller's obligation. The physical node reads neither the metric name nor a time column from +/// the input — a raw frame carries the labels of each series and its own timestamps. +#[derive(Clone, Debug)] +pub struct FusedSeriesBinding { + /// The raw series schema the tree's input node was bound to, without an evaluation grid. + pub series_schema: Arc, + /// The statement's grid, which is the grid every evaluated chunk is addressed on. + pub grid: Arc, + /// The metric name every input series is bound to, when the input reads one named metric. + /// + /// The logical node requires it when [`tags`](Self::tags) declares [`METRIC_NAME`]: a reader of + /// the output has to be able to tell that a range function dropped the name, which needs a name + /// to drop. The physical node takes the names from the frames, which carry each series' own. + pub metric_name: Option>, + /// The output column holding the grid timestamps. + pub time_index: String, + /// The logical tag names the output carries, read from each series' labels. + /// + /// [`METRIC_NAME`] is special: it is resolved to the series' metric name instead of a column, + /// which is how a reader sees that a range function dropped it. + pub tags: Vec, +} + +impl FusedSeriesBinding { + /// Whether `other` binds the same fields under the same names. + fn same_fields(&self, other: &SeriesSchema) -> bool { + self.series_schema.fields.len() == other.fields.len() + && self + .series_schema + .fields + .iter() + .zip(other.fields.iter()) + .all(|(mine, theirs)| mine.name == theirs.name) + } +} + +impl PartialEq for FusedSeriesBinding { + fn eq(&self, other: &Self) -> bool { + // The schemas are compared by the field names they declare, and the grids by their spec: + // neither type has structural equality, and neither identity nor allocation is meant here. + self.metric_name == other.metric_name + && self.time_index == other.time_index + && self.tags == other.tags + && self.grid.spec() == other.grid.spec() + && self.same_fields(&other.series_schema) + } +} + +impl Eq for FusedSeriesBinding {} + +impl Hash for FusedSeriesBinding { + fn hash(&self, state: &mut H) { + self.metric_name.hash(state); + self.time_index.hash(state); + self.tags.hash(state); + let spec = self.grid.spec(); + (spec.start_ms, spec.step_ms, spec.len).hash(state); + for field in self.series_schema.fields.iter() { + field.name.hash(state); + } + } +} + +impl PartialOrd for FusedSeriesBinding { + fn partial_cmp(&self, other: &Self) -> Option { + // A total order over the same fields equality reads, so equal bindings compare equal. + match self + .metric_name + .as_deref() + .partial_cmp(&other.metric_name.as_deref()) + { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.time_index.partial_cmp(&other.time_index) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.tags.partial_cmp(&other.tags) { + Some(Ordering::Equal) => {} + ord => return ord, + } + let (mine, theirs) = (self.grid.spec(), other.grid.spec()); + match (mine.start_ms, mine.step_ms, mine.len).partial_cmp(&( + theirs.start_ms, + theirs.step_ms, + theirs.len, + )) { + Some(Ordering::Equal) => {} + ord => return ord, + } + Some( + self.series_schema + .fields + .iter() + .map(|field| field.name.as_ref()) + .cmp( + other + .series_schema + .fields + .iter() + .map(|field| field.name.as_ref()), + ), + ) + } +} + +/// The logical plan of one node that runs a compiled native tree. +/// +/// The node adds nothing to the tree: it records the query the tree was compiled from, so +/// `EXPLAIN` shows the reader their own query rather than the shape it compiled to, and it carries +/// the binding the tree's input was compiled against so the execution plan can be built from it. +#[derive(Debug)] +pub struct FusedSeries { + /// The PromQL text this tree was compiled from. + promql: Arc, + /// The compiled tree to run. + tree: NativeTree, + /// What the tree's input was compiled against. + binding: FusedSeriesBinding, + /// The plan producing the input series. + input: LogicalPlan, + /// The columns this node produces. + output_schema: DFSchemaRef, +} + +impl FusedSeries { + /// The name this node is shown under. + pub const fn name() -> &'static str { + "FusedSeries" + } + + /// Creates the node, checking that `tree` is legal against `binding`. + pub fn try_new( + promql: Arc, + tree: NativeTree, + binding: FusedSeriesBinding, + input: LogicalPlan, + output_schema: DFSchemaRef, + ) -> DataFusionResult { + validate_binding(&tree, &binding)?; + Ok(Self { + promql, + tree, + binding, + input, + output_schema, + }) + } + + /// The PromQL text this node was compiled from. + pub fn promql(&self) -> &str { + &self.promql + } + + /// The compiled tree this node runs. + pub fn tree(&self) -> &NativeTree { + &self.tree + } + + /// What the tree's input was compiled against. + pub fn binding(&self) -> &FusedSeriesBinding { + &self.binding + } + + /// The physical node that runs `exec_input` as this node's input. + /// + /// `label_source` and `ordering` are the binding of the frames `exec_input` produces: the + /// frame carries the encoded key and, in explicit mode, the complete labels, but no schema can + /// say which of the two a producer meant or how it ordered the series, so the caller that built + /// the frame passes both on. Building the node rejects a tree the runtime cannot execute, a + /// frame schema the tree's input was not compiled against, and a schema the conversion cannot + /// classify, so a plan that is built is a plan that runs. + pub fn to_execution_plan( + &self, + exec_input: Arc, + label_source: RawLabelSource, + ordering: SeriesOrdering, + ) -> DataFusionResult> { + let output_schema: SchemaRef = self.output_schema.inner().clone(); + Ok(Arc::new(FusedSeriesExec::try_new( + Arc::clone(&self.promql), + exec_input, + label_source, + ordering, + Arc::new(self.tree.clone()), + output_schema, + self.binding.time_index.clone(), + self.binding.tags.clone(), + )?)) + } +} + +impl UserDefinedLogicalNodeCore for FusedSeries { + fn name(&self) -> &str { + Self::name() + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.output_schema + } + + fn expressions(&self) -> Vec { + // The tree's operands are bound inside it; there is no expression for an optimizer to + // rewrite separately from the tree that owns it. + Vec::new() + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "FusedSeries: {}", self.promql) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> DataFusionResult { + if inputs.len() != 1 { + return Err(DataFusionError::Internal( + "FusedSeries should have at exact one input".to_string(), + )); + } + + let input = inputs.pop().unwrap(); + Ok(Self { + promql: Arc::clone(&self.promql), + tree: self.tree.clone(), + binding: self.binding.clone(), + input, + output_schema: Arc::clone(&self.output_schema), + }) + } +} + +impl PartialEq for FusedSeries { + fn eq(&self, other: &Self) -> bool { + // `NativeTree` carries no structural equality, so the rendered tree stands in for it. The + // renderer names every node with its bound operation and grids, which is what makes two + // nodes the same computation; comparing the text alone would call two different trees + // equal whenever their PromQL happens to read the same. + self.promql == other.promql + && self.tree.to_string() == other.tree.to_string() + && self.binding == other.binding + && self.input == other.input + && self.output_schema == other.output_schema + } +} + +impl Eq for FusedSeries {} + +impl Hash for FusedSeries { + fn hash(&self, state: &mut H) { + self.promql.hash(state); + self.binding.hash(state); + self.input.hash(state); + self.output_schema.hash(state); + } +} + +impl PartialOrd for FusedSeries { + fn partial_cmp(&self, other: &Self) -> Option { + match self.promql.partial_cmp(&other.promql) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.binding.partial_cmp(&other.binding) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.input.partial_cmp(&other.input) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.tree.to_string().partial_cmp(&other.tree.to_string()) { + Some(Ordering::Equal) => {} + ord => return ord, + } + if self.output_schema == other.output_schema { + Some(Ordering::Equal) + } else { + None + } + } +} + +/// What the tree's input node reads, read off the tree itself. +/// +/// The tree is the one authority on the schema its input was compiled against: a caller's binding +/// can only agree or disagree with it, and the frame this node reads has to be decoded against the +/// same schema, so both halves read it here instead of each carrying a copy. +struct TreeInput { + /// The position the tree's raw chunks enter at: a position is how the runtime addresses a node, + /// so nothing on the node has to carry an identity that the tree itself does not. + path: NodePath, + /// The raw series schema the tree's input node was compiled against. + schema: Arc, +} + +/// Reads what the tree's input node was compiled against, rejecting a tree this node cannot feed. +/// +/// A node feeds one raw input, and its output is addressed on a grid; a tree that does not have +/// that shape is one this layer cannot hand chunks to or read positions from. +fn tree_input(tree: &NativeTree) -> DataFusionResult { + let mut sources = 0; + let mut input_path = None; + let mut schema = None; + metrics::walk(tree.root(), &mut |node, path| { + let NativeExpr::Input(input) = &node.expr else { + return; + }; + sources += 1; + input_path.get_or_insert_with(|| path.clone()); + // Raw, not evaluated: this node feeds the tree from a scan, so a tree whose input was + // compiled against an already-evaluated schema would read the wrong thing. + schema = input + .output + .samples() + .filter(|schema| schema.eval_grid.is_none()) + .map(Arc::clone); + }); + if sources != 1 { + return Err(DataFusionError::Plan(format!( + "FusedSeries runs a tree with {sources} inputs, and it needs exactly one" + ))); + } + let (Some(path), Some(schema)) = (input_path, schema) else { + return Err(DataFusionError::Plan( + "FusedSeries runs a tree whose input is not raw series".to_string(), + )); + }; + let evaluated = tree + .root() + .output + .samples() + .is_some_and(|schema| schema.eval_grid.is_some()); + if !evaluated { + return Err(DataFusionError::Plan( + "FusedSeries runs a tree whose output is not evaluated on a grid".to_string(), + )); + } + Ok(TreeInput { path, schema }) +} + +/// Refuses a tree whose logical series can be split across the partitions this node runs. +/// +/// Keeping the metric name is not what keeps two raw series apart. An empty tag value and an absent +/// tag are the same canonical label set, so `{host="h", zone=""}` and `{host="h"}` are one series to +/// Prometheus while the scan encodes them as two raw keys. A window that evaluated its input without +/// merging it therefore passes such a pair through as two series under one label set, exactly as the +/// functions that drop `__name__` do by merging them, and a map drops `__name__` outright. +/// +/// A stateful node above one of those needs both halves together, and the runtime does provide that +/// within one execution: a window whose input can hold two series under one label set collects the +/// fragments, merges each label set once, and only then evaluates. That merge is per execution, and +/// this node runs one execution per partition. Two raw series that become one label set can land in +/// different partitions - the scan keeps a *series* in one partition, which says nothing about two +/// different series that later become indistinguishable - so each execution merges only the half it +/// was given. The consumer above then answers from one half and reports nothing for the other, which +/// is an empty result rather than an error. +/// +/// Whether a window merged its input is runtime state this check does not have: it runs over a bare +/// tree, where a window's input is only something that produces samples. A window below another +/// window is therefore treated as one that did not merge, whatever function it runs, which refuses +/// `rate` under `rate` (which drops the name) and `last_over_time` under `rate` (which keeps it) +/// alike. +/// +/// Reading such a tree needs an exchange above this node that merges a logical series across +/// partitions, which does not exist yet. Until it does the tree is refused where it can hurt: over +/// one partition there is nothing to split, so only a tree run over more than one is checked. +/// +/// The compiler cannot build a nested window today - [`compile`](crate::promql::native_compile::compile) +/// requires a range function to apply directly to a selector, `([])`, +/// and a selector is raw samples - so no tree it produces has a window above a window, and this +/// check costs no executable query. It is checked because [`FusedSeriesExec::try_new`] accepts any +/// well-formed tree, and a tree it accepts is a tree it must answer correctly. +fn reject_a_split_logical_series( + tree: &NativeTree, + partition_count: usize, +) -> DataFusionResult<()> { + if partition_count == 1 { + return Ok(()); + } + // The walk is outermost first, so it meets nodes in the reverse of the order data flows: a + // consumer met before a rewriter is one that the rewriter feeds. + let mut consumer: Option = None; + let mut rejection: Option<(String, String)> = None; + metrics::walk(tree.root(), &mut |node, path| { + if rejection.is_some() { + return; + } + // A node can leave two raw series on one canonical label set. A map always can: it drops + // `__name__` and every empty-valued tag. An instant selection preserves its input labels, + // rather than dropping them, but can still leave two raw series on one canonical label set: + // `{zone=""}` and the absent `zone` are one label set and two raw keys. A window can + // whenever it did not merge its input, which a bare tree says nothing about: a check on the + // function alone accepts `last_over_time`, whose output keeps the name and can still hold + // both raw series. Every map, instant selection, and window is therefore taken to be a + // possible rewriter, and the refusal below is what narrows that to the shape that costs + // correctness - a window with a rewriter below it, a window being the only consumer here. + let rewrites = matches!( + node.expr, + NativeExpr::Map { .. } | NativeExpr::Instant { .. } | NativeExpr::Window { .. } + ); + if rewrites && let Some(above) = &consumer { + rejection = Some((path.as_str().to_string(), above.clone())); + return; + } + let consumes = matches!(node.expr, NativeExpr::Window { .. }); + if consumes && consumer.is_none() { + consumer = Some(path.as_str().to_string()); + } + }); + let Some((rewriter, above)) = rejection else { + return Ok(()); + }; + Err(DataFusionError::Plan(format!( + "FusedSeries runs a tree where `{rewriter}` can leave two raw series on one canonical label \ + set and `{above}` above it reads a whole series, over {partition_count} partitions: an \ + empty tag value and an absent tag are the same label set, the runtime merges such \ + fragments once per execution, and this node runs one execution per partition, so `{above}` \ + would read one half and answer an empty result. Reading this tree needs an exchange between \ + `{rewriter}` and `{above}` that merges a logical series across partitions, which is not \ + implemented" + ))) +} + +/// Checks that `tree` reads the schema `binding` declares, on the grid it declares. +/// +/// The output grid is compared by [`EvalGrid::spec`], not by allocation: that is the rule +/// [`NativeTree`]'s own validation and the kernels apply, and a chunk's timestamps already come +/// from its absolute grid positions, so a separately allocated grid with the same spec is the same +/// grid to everything downstream. +/// +/// Answers with what the tree's input reads, which is what a caller names when it hands a raw chunk +/// over. +fn validate_binding( + tree: &NativeTree, + binding: &FusedSeriesBinding, +) -> DataFusionResult { + if binding.series_schema.eval_grid.is_some() { + return Err(DataFusionError::Plan( + "FusedSeries binds a raw series schema, which cannot have an evaluation grid" + .to_string(), + )); + } + if binding.tags.iter().any(|tag| tag == METRIC_NAME) && binding.metric_name.is_none() { + return Err(DataFusionError::Plan(format!( + "FusedSeries declares a {METRIC_NAME:?} tag column but binds no metric name" + ))); + } + + // The tree's own input node says which schema it reads, so the binding is checked against the + // tree rather than trusted: a mismatch would evaluate one schema and label the result as + // another. + let input = tree_input(tree)?; + if !binding.same_fields(&input.schema) { + return Err(DataFusionError::Plan( + "FusedSeries binds a raw series schema the tree's input was not compiled against" + .to_string(), + )); + } + match tree.root().output.samples() { + Some(schema) + if schema + .eval_grid + .as_ref() + .is_some_and(|grid| grid.spec() == binding.grid.spec()) => + { + Ok(input) + } + _ => Err(DataFusionError::Plan( + "FusedSeries runs a tree whose output is not evaluated on the statement's grid" + .to_string(), + )), + } +} + +/// Where one output column of [`FusedSeriesExec`] comes from. +#[derive(Clone, Copy, Debug)] +enum OutputColumn { + /// The evaluation grid timestamps of the current chunk. + Time, + /// Series field by index in the bound series schema. + Field(usize), + /// Declared logical tag by index in the caller-supplied tag list. + Tag(usize), +} + +/// Classifies every output column, rejecting anything the conversion cannot produce. +fn classify_output( + output_schema: &SchemaRef, + series_schema: &SeriesSchema, + tags: &[String], + time_column: &str, +) -> DataFusionResult> { + for (index, tag) in tags.iter().enumerate() { + if tags[..index].contains(tag) { + return Err(DataFusionError::Plan(format!( + "duplicate tag column {tag:?} in FusedSeriesExec output schema" + ))); + } + if tag == time_column { + return Err(DataFusionError::Plan(format!( + "tag column {tag:?} of FusedSeriesExec cannot be the time column" + ))); + } + if series_schema + .fields + .iter() + .any(|field| field.name.as_ref() == tag) + { + return Err(DataFusionError::Plan(format!( + "tag column {tag:?} of FusedSeriesExec clashes with a series field" + ))); + } + } + + let mut time_seen = false; + let mut field_seen = vec![false; series_schema.fields.len()]; + let mut tag_seen = vec![false; tags.len()]; + let mut classification = Vec::with_capacity(output_schema.fields().len()); + for field in output_schema.fields() { + let name = field.name(); + if name == time_column { + if std::mem::replace(&mut time_seen, true) { + return Err(DataFusionError::Plan(format!( + "duplicate time column {name:?} in FusedSeriesExec output schema" + ))); + } + match field.data_type() { + DataType::Timestamp(TimeUnit::Millisecond, None) => {} + other => { + return Err(DataFusionError::Plan(format!( + "time column {name:?} of FusedSeriesExec must be Timestamp(Millisecond, \ + None), got {other:?}" + ))); + } + } + if field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "time column {name:?} of FusedSeriesExec must not be nullable" + ))); + } + classification.push(OutputColumn::Time); + continue; + } + if let Some(index) = series_schema + .fields + .iter() + .position(|field| field.name.as_ref() == name) + { + if std::mem::replace(&mut field_seen[index], true) { + return Err(DataFusionError::Plan(format!( + "duplicate field column {name:?} in FusedSeriesExec output schema" + ))); + } + if field.data_type() != &DataType::Float64 { + return Err(DataFusionError::Plan(format!( + "field column {name:?} of FusedSeriesExec must be Float64, got {:?}", + field.data_type() + ))); + } + // A grid point inside the statement's interval can have no value in this field. + if !field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "field column {name:?} of FusedSeriesExec must be nullable, because a grid \ + point can be missing" + ))); + } + classification.push(OutputColumn::Field(index)); + continue; + } + if let Some(index) = tags.iter().position(|tag| tag == name) { + if std::mem::replace(&mut tag_seen[index], true) { + return Err(DataFusionError::Plan(format!( + "duplicate tag column {name:?} in FusedSeriesExec output schema" + ))); + } + if field.data_type() != &DataType::Utf8 { + return Err(DataFusionError::Plan(format!( + "tag column {name:?} of FusedSeriesExec must be Utf8, got {:?}", + field.data_type() + ))); + } + // `rate` and the map kernels drop `__name__`, so a tag can be absent from a row. + if !field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "tag column {name:?} of FusedSeriesExec must be nullable, because range \ + functions and map kernels drop the metric name" + ))); + } + classification.push(OutputColumn::Tag(index)); + continue; + } + return Err(DataFusionError::Plan(format!( + "unknown column {name:?} in FusedSeriesExec output schema" + ))); + } + + if !time_seen { + return Err(DataFusionError::Plan(format!( + "missing time column {time_column:?} in FusedSeriesExec output schema" + ))); + } + for (index, seen) in field_seen.iter().enumerate() { + if !seen { + return Err(DataFusionError::Plan(format!( + "missing field column {:?} in FusedSeriesExec output schema", + series_schema.fields[index].name + ))); + } + } + for (index, seen) in tag_seen.iter().enumerate() { + if !seen { + return Err(DataFusionError::Plan(format!( + "missing tag column {:?} in FusedSeriesExec output schema", + tags[index] + ))); + } + } + Ok(classification) +} + +/// The physical node that runs a compiled native tree over its input's raw series. +pub struct FusedSeriesExec { + /// The PromQL text the tree was compiled from, which is what the plan line shows. + promql: Arc, + /// The tree this node runs, shared by every partition's execution. + tree: Arc, + /// The raw series schema the tree's input was compiled against, which every frame is decoded + /// against. + series_schema: Arc, + /// How the input frames name their series: labels encoded in the key, or carried along. + label_source: RawLabelSource, + /// The ordering the input frames guarantee. + ordering: SeriesOrdering, + /// The position of the tree's input node, where every raw chunk a partition reads enters. + input_path: NodePath, + /// Where each output column comes from. + classification: Arc<[OutputColumn]>, + /// The output column holding the grid timestamps, kept so this node can be rebuilt over a new + /// child without losing which column its time came from. + time_column: String, + /// The tag column names, indexed by [`OutputColumn::Tag`]. + tags: Arc<[String]>, + /// The plan producing the raw series frames. + input: Arc, + /// The columns this node produces. + output_schema: SchemaRef, + /// The grid positions every label set in this plan's output already occupies, and what they + /// keep charged on the memory pool. + /// + /// Plan-level rather than per-partition: see [`LabelSets`]. + labels: SharedLabelSets, + /// Per-node metrics, registered by each execution of the tree. + metric: ExecutionPlanMetricsSet, + /// This node's partitioning, emission and boundedness. + properties: Arc, +} + +impl FusedSeriesExec { + /// Creates the node, rejecting whatever this node cannot execute or convert. + /// + /// `input` produces native raw series frames: the series' encoded key, its finality, its + /// timestamps and its fields, plus the complete labels when `label_source` is + /// [`RawLabelSource::Explicit`]. `time_column` and `tags` name the columns of the output, which + /// is a row table rather than a frame. + /// + /// The parameters stay separate because each is a fact its own caller has: the query, the plan + /// below, the frame's label binding and ordering, the tree, and the three things the output + /// schema is classified by. A struct over them would be the same list under another name. + #[allow(clippy::too_many_arguments)] + pub fn try_new( + promql: Arc, + input: Arc, + label_source: RawLabelSource, + ordering: SeriesOrdering, + tree: Arc, + output_schema: SchemaRef, + time_column: String, + tags: Vec, + ) -> DataFusionResult { + let tree_input = tree_input(&tree)?; + // Building a state is what proves the runtime executes this tree: a node kind it does not + // run yet is refused here, so a plan that is built is a plan that starts. The throwaway + // state is dropped at once, returning its (unbounded) reservations. + let _ = NativeExecutionState::try_new( + &tree, + Arc::new(UnboundedMemoryPool::default()), + &ExecutionPlanMetricsSet::new(), + 0, + )?; + let classification = + classify_output(&output_schema, &tree_input.schema, &tags, &time_column)?; + // Read here rather than below, because whether the split below can happen decides whether + // the tree shape is one this node may run at all. + let input_properties = input.properties(); + let partition_count = input_properties.partitioning.partition_count(); + if partition_count == 0 { + return Err(DataFusionError::Plan( + "FusedSeriesExec input has no partitions".into(), + )); + } + reject_a_split_logical_series(&tree, partition_count)?; + // The frame states its own label layout and series ordering; this plan's binding is an + // expectation checked against that declaration, not an input to decoding. + ensure_frame_matches_label_source(input.schema().as_ref(), &label_source)?; + // The frame is checked against what the tree's input reads, so a frame this node cannot + // decode is reported when the plan is built. The shared decoder is the authority on that + // layout, so it is asked rather than copied: handing it an empty stream of this schema makes + // it validate the declared frame with no batch to decode, which is the check + // `unpack_raw_series` makes before its first poll too. + unpack_raw_series( + Box::pin(RecordBatchStreamAdapter::new( + input.schema(), + futures::stream::empty(), + )), + Some(Arc::clone(&tree_input.schema)), + )?; + // One evaluator per input partition, and the input's own partitioning kept as it is. + // + // A window holds the state of the series it is reading and closes it on the chunk carrying + // `is_last`, so it can answer a series only if the whole series arrives in the partition it + // reads. The input guarantees that: a raw series frame is produced by `SeriesScan`, whose + // contract is that the same primary key always reaches the same partition, and the frame + // decoder holds each series' chunks contiguous within that partition. + // + // What is *not* asked of the optimizer is a distribution. A `KeyPartitioned` or + // `HashPartitioned` requirement is classed as a hash requirement, and a hash requirement + // makes a hash exchange necessary whenever the child has more than one partition - even if + // the child already satisfies it, and even when the partition counts match. A hash exchange + // does not preserve order, so it would interleave the chunks of one series and the decoder + // would close a series it has not finished. `UnspecifiedDistribution` leaves the child's + // layout alone, and `benefits_from_input_partitioning` is false so no round-robin is added + // to raise parallelism either. Between them no exchange can appear between the scan and this + // node, and the scan's guarantee arrives intact. + // As many partitions as the input hands over: the node runs one evaluator per partition and + // answers each as it settles. + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(partition_count), + input_properties.emission_type, + input_properties.boundedness, + )); + // Shared by every partition's stream: a label set colliding is a property of the whole + // output vector, not of one partition. Building it needs no pool; the owner's reservation is + // attached in `execute`, which is the first point that has one. + let labels: SharedLabelSets = Arc::new(Mutex::new(LabelSets::default())); + Ok(Self { + promql, + tree, + series_schema: tree_input.schema, + label_source, + ordering, + input_path: tree_input.path, + classification: classification.into(), + time_column, + tags: tags.into(), + input, + output_schema, + labels, + metric: ExecutionPlanMetricsSet::new(), + properties, + }) + } + + /// The name this node is shown under. + pub const fn name() -> &'static str { + "FusedSeriesExec" + } +} + +impl std::fmt::Debug for FusedSeriesExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("FusedSeriesExec").finish() + } +} + +impl DisplayAs for FusedSeriesExec { + /// The query, then the tree it compiles to: the tree is rendered by [`NativeTree`]'s own + /// `Display`, so `EXPLAIN` and the node's metrics name the same nodes. + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "FusedSeriesExec: {}\n{}", self.promql, self.tree) + } +} + +impl ExecutionPlan for FusedSeriesExec { + fn name(&self) -> &str { + Self::name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + /// Nothing is asked of the input's distribution, and nothing may be inserted above it. + /// + /// A hash requirement would make a hash exchange necessary as soon as the child has more than + /// one partition - even where the child already satisfies the requirement - and a hash exchange + /// does not preserve order. The input's own layout is already the one this node needs, because + /// the scan it reads from guarantees a series stays in one partition, so the requirement says + /// only that the layout is left alone. + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + InputDistributionRequirements::new(vec![Distribution::UnspecifiedDistribution]) + } + + /// Splitting the input further would not help: a second partition is a second set of series + /// answered by a second tree, which this node already runs, and raising the count is what an + /// optimizer does by adding a round-robin exchange that reorders nothing it can rely on. + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + // The tree's operands are bound inside it; this node owns no expression of its own. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + let [child] = <[Arc; 1]>::try_from(children).map_err(|children| { + DataFusionError::Plan(format!( + "FusedSeriesExec requires exactly one child, got {}", + children.len() + )) + })?; + // Rebuilt rather than patched: the new child's schema is what the frame is validated + // against, so a child that changed shape is refused instead of decoded with a stale layout. + Ok(Arc::new(Self::try_new( + Arc::clone(&self.promql), + child, + self.label_source.clone(), + self.ordering, + Arc::clone(&self.tree), + self.output_schema.clone(), + self.time_column.clone(), + self.tags.to_vec(), + )?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + // The label-set state is the plan's, so the reservation that pays for it is attached exactly + // once, to the pool of the first execution. Every partition calls this and only the first + // takes effect: re-registering per stream would move a charge the other partitions keep + // growing, and would give it back when that stream ends while the entries stay. + self.labels + .lock() + .unwrap_or_else(|e| e.into_inner()) + .attach(&context.runtime_env().memory_pool); + let input = self.input.execute(partition, Arc::clone(&context))?; + // Batches become raw series chunks through the shared frame decoder, which is the inverse of + // what the producer packed: one chunk per frame row, each series closed exactly once by the + // row carrying `is_last`. + let input = unpack_raw_series(input, Some(Arc::clone(&self.series_schema)))?; + // A source bound in key mode carries its identity encoded and no decoded labels, and a chunk + // handed to the tree has to name the series it holds. + let input = + materialize_raw_labels(input).map_err(|e| DataFusionError::External(Box::new(e)))?; + let state = NativeExecutionState::try_new( + &self.tree, + Arc::clone(&context.runtime_env().memory_pool), + &self.metric, + partition, + )?; + // What the label sets hold is charged to the reservation attached above, so nothing is + // reserved per stream: the entries outlive every stream and are released with the plan. + Ok(Box::pin(FusedSeriesStream { + input: Some(input), + executor: Some(NativeExecutor::new(state)), + input_path: self.input_path.clone(), + classification: Arc::clone(&self.classification), + tags: Arc::clone(&self.tags), + output_schema: self.output_schema.clone(), + labels: Arc::clone(&self.labels), + ready: VecDeque::new(), + done: false, + })) + } + + fn metrics(&self) -> Option { + Some(self.metric.clone_inner()) + } +} + +/// Converts one evaluated chunk into one output batch. +/// +/// The time column is the absolute grid positions the chunk carries, missing samples are nulls and +/// ordinary NaN stays a value. +fn series_to_record_batch( + classification: &[OutputColumn], + tags: &[String], + output_schema: &SchemaRef, + chunk: &EvaluatedChunk, +) -> DataFusionResult { + let series = &chunk.series; + // A chunk that forms one run is read from its range, which lends the grid's buffers and so + // allocates nothing. A chunk whose rows form several runs has no single range to read + // ([`EvaluatedChunk::range`] answers `None` for it), and the positions between its runs are + // holes rather than rows: the axis is walked for its carried rows instead, which is the one + // traversal that visits exactly those. + let timestamps = match chunk.range() { + Some(range) => range.timestamps(), + None => series + .time + .timestamps() + .map(|timestamps| ScalarBuffer::from(timestamps.into_owned())), + } + .map_err(|error| DataFusionError::Execution(error.to_string()))?; + let rows = timestamps.len(); + let mut columns: Vec = Vec::with_capacity(classification.len()); + for (position, column) in classification.iter().enumerate() { + match column { + OutputColumn::Time => columns.push(Arc::new(TimestampMillisecondArray::new( + timestamps.clone(), + None, + ))), + OutputColumn::Field(index) => { + let samples = series.fields.get(*index).ok_or_else(|| { + DataFusionError::Execution(format!( + "evaluated series has no field at index {index}" + )) + })?; + if samples.len() != rows { + return Err(DataFusionError::Execution(format!( + "evaluated series field {index} has {} rows, time axis has {rows}", + samples.len() + ))); + } + columns.push(field_column(samples)?); + } + OutputColumn::Tag(index) => { + // `__name__` resolves to the metric name, every other tag through the label set. + let name = tags.get(*index).ok_or_else(|| { + DataFusionError::Execution(format!( + "output column {position} declares unavailable tag {index}" + )) + })?; + let value = series.metric.get(name); + columns.push(Arc::new(StringArray::from(vec![value; rows]))); + } + } + } + Ok(RecordBatch::try_new(output_schema.clone(), columns)?) +} + +/// Converts one evaluated series field into one `Float64` column. +/// +/// Dense float samples are reused as they are: the array already carries exactly the right nulls +/// and keeps NaN a value, so copying it through `Samples::iter()` would only cost an allocation +/// and a re-encode. Chunked samples go through the iterator, and histograms have no float column. +fn field_column(samples: &Samples) -> DataFusionResult { + match samples.dense() { + Some(DenseSamplesRef::Float(values)) => Ok(Arc::new(values.clone())), + Some(DenseSamplesRef::Histogram(_)) | Some(DenseSamplesRef::Mixed(_)) => { + Err(DataFusionError::NotImplemented( + "FusedSeriesExec cannot output histogram samples".into(), + )) + } + None => { + let values = samples + .iter() + .map(|sample| match sample { + SampleRef::Float(value) => Ok(Some(value)), + SampleRef::Missing => Ok(None), + SampleRef::Histogram(_) => Err(DataFusionError::NotImplemented( + "FusedSeriesExec cannot output histogram samples".into(), + )), + }) + .collect::>>()?; + Ok(Arc::new(Float64Array::from(values))) + } + } +} + +/// The bytes one label set holds, for the memory reservation. +/// +/// It counts each label's name and value with the pair that refers to them, which is the same +/// convention [`crate::native::execute`]'s own label-set table keeps. +fn labels_bytes(labels: &LabelSet) -> usize { + labels + .iter() + .map(|(name, value)| size_of::<(Arc, Arc)>() + name.len() + value.len()) + .sum() +} + +/// Renders the label set Prometheus compares the way it prints one, for the duplicate-label-set +/// error. +fn render_labels(labels: &LabelSet) -> String { + let parts = labels + .iter() + .map(|(name, value)| format!("{name}={value:?}")) + .collect::>(); + format!("{{{}}}", parts.join(", ")) +} + +/// The grid positions each final label set already occupies, per field and as merged, sorted runs. +/// +/// Shared by every partition's stream, because a label-set collision is a property of the whole +/// output vector and not of one partition. Two series that carry the same labels after the range +/// function and its maps drop a label land in different partitions, and the runtime's own per-node +/// `Emitted` table never sees both: it belongs to one execution. Checking only within a partition +/// would miss exactly those collisions. +/// +/// The key is [`crate::grid_series::label_set`]'s normalization of the final labels rather than +/// [`datatypes::series::MetricName`] itself: a label with an empty value is no label to +/// Prometheus, so `{host="h", zone=""}` and `{host="h"}` are one label set and have to be compared +/// as one. Encoding the tags as they are would file them apart and let one point hold two values. +/// +/// Occupancy is per field and per position, not per chunk: two series whose label sets became equal +/// while holding values in *different* fields merge cleanly in Prometheus, so only the same field +/// at the same position collides. That is the rule +/// [`crate::grid_series::merge_same_labelset`] applies. +/// +/// The map belongs to the plan, not to one stream: a partition that ends leaves what it claimed +/// for the partitions still running, so the entries are only released with the plan itself. The +/// charge has the same lifetime, which is what makes the bytes the map holds - one entry per label +/// set of the output, kept until the plan dies - visible to the memory pool. Charging each stream +/// for what it inserted instead returned those bytes when that stream ended, while the entries the +/// later partitions are checked against stayed: everything the map held after the first partition +/// was unaccounted for. +#[derive(Default)] +struct LabelSets { + /// The grid positions each final label set already occupies, one run list per field of the + /// series schema. Each list is sorted, disjoint and has the runs that touch joined, which is + /// what lets a claim look a position up instead of walking the history. + seen: HashMap>>, + /// What `seen` keeps charged on the task's memory pool. `None` until the first execution, + /// since construction has no `TaskContext` and therefore no pool to register against. + reservation: Option, +} + +impl LabelSets { + /// Attaches the reservation to the pool of the first execution. + /// + /// Later executions keep the first reservation: it already pays for what `seen` holds, and + /// moving it to another pool would drop the bytes other partitions are still growing. + fn attach(&mut self, pool: &Arc) { + if self.reservation.is_none() { + self.reservation = + Some(MemoryConsumer::new("FusedSeriesExec label sets").register(pool)); + } + } + + /// Claims the occupied grid positions of one chunk of a final series for its label set. + /// + /// A collision is an error only when both series hold a sample of *one field* at the same + /// evaluation timestamp, which is Prometheus' range-query rule. Series with disjoint points are + /// legal, and so are series that fill the same points field by field: Prometheus merges both + /// into one output series, while this node emits them as separate batches and defers that merge + /// to matrix assembly, so every sample survives. + fn claim(&mut self, series: &Series) -> DataFusionResult<()> { + let runs = occupied_runs(series)?; + // A chunk holding no value occupies nothing, so there is no collision to remember and no + // entry to write. Such a chunk is not rare: a selection over a series that staleness or the + // range left empty still closes it, and hands over a closing chunk with no rows. Recording + // one entry per such series would grow the table - and its charge, which the plan holds for + // its whole life - for a query that never answers a value, until an empty result exhausted + // the pool. The executor's own `Emitted::occupy` takes the same early return. + if runs.iter().all(Vec::is_empty) { + return Ok(()); + } + let key = label_set(&series.metric); + let Self { seen, reservation } = self; + // Nothing reaches a claim without an `execute` first, and an execution that ran attached the + // reservation. Claiming without one would leave the map growing uncharged, so it is refused + // rather than silently skipped. + let Some(reservation) = reservation.as_mut() else { + return Err(DataFusionError::Execution( + "FusedSeriesExec label-set state is not attached to a memory pool".to_string(), + )); + }; + match seen.entry(key) { + Entry::Occupied(entry) => { + if let Some(field) = colliding_field(entry.get(), &runs) { + return Err(DataFusionError::Execution(format!( + "vector cannot contain metrics with the same labelset: {} (field \ + {field})", + render_labels(entry.key()) + ))); + } + let stored = entry.into_mut(); + // An entry recorded from a chunk of fewer fields is padded, so every field of this + // chunk has a slot of its own: a slot is a field index of the schema, so both sides + // count the same field the same way. A chunk of *fewer* fields than the entry + // leaves the remaining slots to the chunks that fill them. + // Stage only the affected run lists, and reserve before committing. A forward + // append whose existing capacity suffices needs no copy of its stored history. + let mut growth = 0usize; + let mut shrink = 0usize; + let mut merged_runs = Vec::with_capacity(runs.len()); + for (index, incoming) in runs.iter().enumerate() { + let Some(old) = stored.get(index) else { + let copy = incoming.clone(); + growth += copy.capacity(); + merged_runs.push(Some(copy)); + continue; + }; + let append_fits = incoming.is_empty() + || (old.last().is_none_or(|last| incoming[0].start >= last.end) + && old.capacity() + >= old.len() + incoming.len() + - usize::from(old.last().is_some_and(|last| { + incoming + .first() + .is_some_and(|first| first.start == last.end) + }))); + if append_fits { + merged_runs.push(None); + } else { + let mut merged = old.clone(); + absorb(&mut merged, incoming); + if merged.capacity() >= old.capacity() { + growth += merged.capacity() - old.capacity(); + } else { + shrink += old.capacity() - merged.capacity(); + } + merged_runs.push(Some(merged)); + } + } + reservation.try_grow(growth * size_of::())?; + for (index, (incoming, merged)) in runs.iter().zip(merged_runs).enumerate() { + if let Some(merged) = merged { + if index < stored.len() { + stored[index] = merged; + } else { + stored.push(merged); + } + } else if index < stored.len() { + absorb(&mut stored[index], incoming); + } + } + reservation.try_shrink(shrink * size_of::())?; + } + Entry::Vacant(entry) => { + // Only the label-set state is this node's own cross-poll state that grows with + // the output, so it is the only thing reserved here. The runs are charged by the + // capacity they were gathered with rather than by their length, so that a later + // claim appending into that slack is charged nothing for it. + let bytes = labels_bytes(entry.key()) + + runs.iter().map(|runs| runs.capacity()).sum::() * size_of::(); + reservation.try_grow(bytes)?; + entry.insert(runs); + } + } + Ok(()) + } +} + +type SharedLabelSets = Arc>; + +/// The grid positions the rows of one chunk hold a value at, one run list per field. +/// +/// Each [`Run`] is a half-open range of absolute grid indices, and the runs of one field are sorted +/// and disjoint. Occupancy is kept per field because it is what a collision is: two series whose +/// label sets became equal merge field by field in Prometheus, so only two values of the same field +/// at one position collide. Marking the whole position as taken as soon as *any* field held a value +/// would refuse the complementary pair that merges cleanly. +fn occupied_runs(series: &Series) -> DataFusionResult>> { + // Rows are addressed on the statement's grid, so a row sits at the absolute position two + // series can collide on. An explicit axis has no such position, and a chunk never carries + // one: this is the refusal `EvaluatedChunk` already makes when the chunk is built. + if series.time.grid().is_none() { + return Err(DataFusionError::Execution( + "FusedSeriesExec requires evaluated series on a time grid".to_string(), + )); + } + let rows = series.time.len(); + let range = series.time.as_grid_range(); + let mut runs: Vec> = vec![Vec::new(); series.fields.len()]; + // A field that is one dense, null-free float run holds a value on every row the chunk carries, + // so it occupies exactly the range the chunk reports and no row of it has to be inspected. This + // is the shape the fused evaluator produces for a range function over ordinary float samples. + // Every other field is walked below, and a field filled here is skipped there. + let mut walked: Vec<(usize, &Samples)> = Vec::new(); + for (index, samples) in series.fields.iter().enumerate() { + match (range, samples.dense()) { + (Some(range), Some(DenseSamplesRef::Float(values))) + if rows > 0 && values.len() == rows && values.null_count() == 0 => + { + runs[index] = vec![Run { + start: range.start(), + end: range.start() + rows, + }]; + } + _ => walked.push((index, samples)), + } + } + if walked.is_empty() { + return Ok(runs); + } + // `positions` walks a chunk that forms one run and one whose rows form several the same way, + // in position order, so the holes between two runs are never read as rows. + let mut cursors: Vec<_> = walked.iter().map(|(_, samples)| samples.cursor()).collect(); + for (position, row) in series.time.positions() { + for ((index, _), cursor) in walked.iter().zip(&mut cursors) { + if !matches!(cursor.at(row), Some(sample) if !matches!(sample, SampleRef::Missing)) { + continue; + } + match runs[*index].last_mut() { + Some(run) if run.end == position => run.end += 1, + _ => runs[*index].push(Run { + start: position, + end: position + 1, + }), + } + } + } + Ok(runs) +} + +/// The field of `runs` that holds a value where `stored` already holds one, if any. +/// +/// A collision is two values of one field at one position, so each field is checked on its own. +/// Each incoming run is looked up rather than the stored history walked: the stored runs are sorted +/// and disjoint, so the only run that can share a position with `run` is the first one that does not +/// end before it, which is the look-up the runtime's own `Emitted::occupy` makes. Walking the +/// history instead costs a sparse series one pass over everything it stored per chunk - a query +/// whose chunks move forward one position at a time, a hole apart, is legal and would cost the +/// square of its chunk count under the lock every partition shares. +fn colliding_field(stored: &[Vec], runs: &[Vec]) -> Option { + for (field, (stored, new)) in stored.iter().zip(runs).enumerate() { + for run in new { + let at = stored.partition_point(|stored| stored.end <= run.start); + if stored.get(at).is_some_and(|stored| stored.overlaps(*run)) { + return Some(field); + } + } + } + None +} + +/// The batches one execution of [`FusedSeriesExec`] yields. +pub struct FusedSeriesStream { + /// The input frames, read as raw series chunks, or `None` once the stream ended in an error. + /// + /// The decoder behind it holds its current batch, the chunks it decoded from it and its own + /// source, so a terminated stream that kept it alive would hold upstream resources until the + /// consumer drops the stream. + input: Option, + /// The execution of this node's tree, or `None` once it has been closed. + executor: Option, + /// The position of the tree's input node, which is where every chunk read enters. + input_path: NodePath, + /// Where each output column comes from. + classification: Arc<[OutputColumn]>, + /// The tag column names, indexed by `OutputColumn::Tag`. + tags: Arc<[String]>, + /// The schema of every batch this stream yields. + output_schema: SchemaRef, + /// The label-set occupancy shared with the other partitions' streams. + /// + /// What it holds is charged to the plan's own reservation, attached by the first execution. + labels: SharedLabelSets, + /// The batches the executor settled since the last poll of the input. + ready: VecDeque, + /// Whether the input ended and everything it settled has been handed out. + done: bool, +} + +impl FusedSeriesStream { + /// Hands one raw chunk over and queues what it settled. + fn push(&mut self, chunk: RawSeriesChunk) -> DataFusionResult<()> { + let Some(executor) = self.executor.as_mut() else { + return Err(DataFusionError::Execution( + "FusedSeriesExec stream pushed a chunk after its execution closed".to_string(), + )); + }; + // The chunk enters at the tree's input node by position: a position is how the runtime + // addresses a node, and a tree that gains a second input names them the same way. + executor.push_raw_at(&self.input_path, chunk)?; + let mut settled = Vec::new(); + while let Some(chunk) = executor.next_chunk()? { + settled.push(chunk); + } + self.convert(settled) + } + + /// Ends the input, so the tree settles what it was holding. + /// + /// `finish()` would be the same thing plus a drain, but the drain collects: it holds the whole + /// answer in a `Vec` before anything is handed out, so a tree that answers many chunks after its + /// input ends would sit on all of them while the consumer waits for the first. Only the input is + /// ended here, and the answers are pulled one per call below, so each chunk leaves as the tree + /// settles it and what has not been pulled yet stays charged to the node holding it. + fn close_input(&mut self) -> DataFusionResult<()> { + match self.executor.as_mut() { + Some(executor) => executor.end_of_input(), + // Every answer was pulled and the execution closed itself. + None => Ok(()), + } + } + + /// Pulls the next answer the tree settles after its input ended, or `None` when it has none. + /// + /// One chunk per call is what keeps this a stream: collecting them first would hold the whole + /// answer while the consumer waits for the first, and would charge the pool for all of it at + /// once. `None` is also when the execution closes and returns its reservations, so it is + /// dropped here rather than staying alive until the stream is. + fn pull_one(&mut self) -> DataFusionResult> { + let Some(executor) = self.executor.as_mut() else { + return Ok(None); + }; + match executor.next_chunk() { + Ok(Some(chunk)) => Ok(Some(chunk)), + Ok(None) => { + self.executor = None; + Ok(None) + } + Err(error) => Err(error), + } + } + + /// Ends the stream in an error, releasing the execution and the input at once. + /// + /// An ordinary drop of the stream releases all of this as well, so this is not about the last + /// release: it is about a consumer that keeps the terminated stream, which would otherwise + /// hold the evaluator, the decoder's held batch and chunks, and the plan upstream of it. + fn fail(&mut self) { + self.done = true; + self.executor = None; + self.input = None; + // The label-set map and its charge are the plan's, so a stream stopping early gives back + // nothing: the entries another partition is still checking against have to stay charged, and + // are released when the plan itself is dropped. + // A batch queued before the error was settled by a run that has now stopped. + self.ready.clear(); + } + + /// Converts the chunks the executor settled into batches, in the order it produced them. + fn convert(&mut self, chunks: Vec) -> DataFusionResult<()> { + for chunk in &chunks { + // The label set is claimed before its batch is queued, so a collision surfaces at the + // chunk that reached it and nothing this plan would refuse is handed out after it. + self.labels + .lock() + .unwrap_or_else(|e| e.into_inner()) + .claim(&chunk.series)?; + let batch = series_to_record_batch( + &self.classification, + &self.tags, + &self.output_schema, + chunk, + )?; + self.ready.push_back(batch); + } + Ok(()) + } +} + +impl RecordBatchStream for FusedSeriesStream { + fn schema(&self) -> SchemaRef { + self.output_schema.clone() + } +} + +impl Stream for FusedSeriesStream { + type Item = DataFusionResult; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll>> { + let this = &mut *self; + // A batch that is already settled is handed over before anything is read, so a bounded + // poll never delays one that is ready behind more input. + for _ in 0..POLL_BUDGET { + if let Some(batch) = this.ready.pop_front() { + return Poll::Ready(Some(Ok(batch))); + } + if this.done { + return Poll::Ready(None); + } + // An error is the only thing that clears the input, and it ends the stream with it. + let Some(input) = this.input.as_mut() else { + // The input ended, so what the tree answers now comes from what it held. One chunk + // per iteration keeps it a stream, and `None` is both "nothing more" and the point + // where the execution releases its reservations. + match this.pull_one() { + Ok(Some(chunk)) => { + if let Err(error) = this.convert(vec![chunk]) { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + Ok(None) => { + this.done = true; + continue; + } + Err(error) => { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + continue; + }; + let polled = input.as_mut().poll_next(cx); + match polled { + Poll::Pending => return Poll::Pending, + Poll::Ready(None) => { + // The input ended, so every series it opened has to be closed. The tree settles + // what it was holding, and the answers are pulled below in this same poll + // rather than collected first. + this.input = None; + if let Err(error) = this.close_input() { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + Poll::Ready(Some(Err(error))) => { + // An input error stops the execution: a stopped evaluator must not be reused. + this.fail(); + return Poll::Ready(Some(Err(DataFusionError::External(Box::new(error))))); + } + Poll::Ready(Some(Ok(chunk))) => { + if let Err(error) = this.push(chunk) { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + } + } + // The budget ran out on input that settled nothing. Everything read so far is kept, so the + // next poll continues where this one stopped; waking first is what asks for that poll. + cx.waker().wake_by_ref(); + Poll::Pending + } +} diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs new file mode 100644 index 00000000000..017f4ef6628 --- /dev/null +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -0,0 +1,2882 @@ +// 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. + +//! What this node has to get right: the two `EXPLAIN` lines it shows, the batches it runs end to +//! end, and what it refuses to run. +//! +//! The trees here are built node by node with the same bindings the PromQL compiler produces for +//! `rate(x[5m])`: the compiler itself lives in the `query` crate, which depends on this one, so a +//! test in this crate cannot call it. +//! +//! The input is what a scan hands over: native raw series frames, one batch per chunk here, packed +//! by the same code the scan uses ([`pack_raw_series_chunk`]). A frame carries the series' encoded +//! key, the chunk's finality, its timestamps and its fields, plus the complete labels, so a test +//! says which series a chunk belongs to and whether it closes it instead of relying on where a row +//! boundary happens to fall. +//! +//! Kernel numerics are not the subject: `rate` is checked against a counter that advances one +//! second per second, so the value it has to produce is exactly 1 (or 2) at every grid point and +//! any position or finalization mistake shows up as a wrong, missing or duplicated row. + +use std::collections::VecDeque; +use std::fmt::{self, Display, Formatter}; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::task::{Context, Poll}; + +use datafusion::arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray}; +use datafusion::arrow::buffer::{Buffer, ScalarBuffer}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::ToDFSchema; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::context::TaskContext; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::logical_expr::{EmptyRelation, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_optimizer::ensure_requirements::EnsureRequirements; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricValue, MetricsSet}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, collect, displayable, +}; +use datafusion::prelude::SessionContext; +use datatypes::series::frame::{RawLabelMode, pack_raw_series_chunk, raw_series_frame_schema}; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{FieldSpec, Samples, Series, SeriesFields, SeriesSchema}; +use datatypes::series::time::{EvalGrid, GridSpan, GridSpec, TimeAxis}; +use datatypes::series::{Label, MetricName, SeriesOrdering}; +use futures::{Stream, StreamExt, TryStreamExt}; +use store_api::region_engine::RawLabelSource; + +use super::{ + FusedSeries, FusedSeriesBinding, LabelSets, OutputColumn, POLL_BUDGET, label_set, + occupied_runs, series_to_record_batch, +}; +use crate::evaluated::EvaluatedChunk; +use crate::native::execute::{NativeExecutionState, NativeExecutor}; +use crate::native::plan::{ + AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundWindow, Grouping, + NativeExpr, NativeNode, NativeTree, NativeType, +}; +use crate::native::runs::absorb_rebuilds; +use crate::range_eval::{RangeFunction, RangeWindow}; + +/// The query every tree here is compiled from. +const PROMQL: &str = "rate(x[5m])"; +/// The range the statement's window selects, which is the `5m` of [`PROMQL`]. +const RANGE_MS: i64 = 300_000; +/// The statement's grid: `rate(x[5m])` evaluated every minute from five minutes in. +const GRID_START: i64 = 300_000; +const GRID_STEP: i64 = 60_000; +const GRID_LEN: u64 = 5; +/// One raw sample per minute, from 0 to 9 minutes: enough to fill every window on the grid. +const SAMPLE_STEP: i64 = 60_000; +const SAMPLES: i64 = 10; +/// How many one-sample chunks the quiet series below arrives in. +/// +/// Two bounds hold at once: the series reaches the end of the last grid point's window (540s, so +/// indices `0..=540`), which is what lets it rate at 1 over the whole grid, and it arrives in more +/// batches than one poll's budget, which is what makes a poll that ignored the budget visible. +const QUIET_SAMPLES: i64 = 541; +/// One second between those samples, so the series rates at 1 once it is read to its end. +/// +/// The step is a second and not a minute on purpose: at a minute the series would arrive in ten +/// batches, fewer than [`POLL_BUDGET`], and the budget the test above exists to check would never +/// be reached. +const QUIET_STEP: i64 = 1_000; + +const TIME: &str = "time"; +const VALUE: &str = "value"; +const OTHER: &str = "other"; +const HOST: &str = "host"; +const ZONE: &str = "zone"; +const NAME: &str = "__name__"; +/// A metadata label, which the window drops along with the metric name. +const TYPE: &str = "__type__"; +const METRIC: &str = "http_requests_total"; + +fn grid() -> Arc { + Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, GRID_LEN).unwrap()).unwrap(), + ) +} + +/// The raw schema the tree's input is bound to: one float field called [`VALUE`]. +fn raw_schema() -> Arc { + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: VALUE.into() }]), + eval_grid: None, + }) +} + +/// A raw schema of two float fields, so one series can fill one of them and another the other. +fn two_field_schema() -> Arc { + Arc::new(SeriesSchema { + fields: Arc::from([ + FieldSpec { name: VALUE.into() }, + FieldSpec { name: OTHER.into() }, + ]), + eval_grid: None, + }) +} + +/// The tree `rate(x[5m])` compiles to over `schema`: raw samples folded by one window onto the grid. +fn rate_tree_of(schema: Arc) -> NativeTree { + let grid = grid(); + // One schema shared by the window's binding and its output, which is what `NativeTree`'s + // validation compares by identity. + let evaluated = NativeType::Samples(Arc::new(SeriesSchema { + fields: schema.fields.clone(), + eval_grid: Some(Arc::clone(&grid)), + })); + let raw = NativeType::Samples(schema); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: RANGE_MS, + offset_ms: 0, + anchor_ms: None, + }, + // The input is raw samples, so it has no grid of its own. + input_grid: None, + output_grid: grid, + }, + }, + output: evaluated, + }) + .unwrap() +} + +/// The tree `rate(x[5m])` compiles to over one field. +fn rate_tree() -> NativeTree { + rate_tree_of(raw_schema()) +} + +/// An instant selection feeding a range window. +fn instant_window_tree() -> NativeTree { + let instant = instant_tree().root().clone(); + let grid = instant + .output + .samples() + .expect("an instant selection produces samples") + .eval_grid + .clone() + .expect("an instant selection's output is evaluated on a grid"); + let evaluated = instant.output.clone(); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(instant), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: RANGE_MS, + offset_ms: 0, + anchor_ms: None, + }, + input_grid: Some(Arc::clone(&grid)), + output_grid: grid, + }, + }, + output: evaluated, + }) + .expect("the bound tree is legal") +} + +/// Two windows stacked, `rate(last_over_time(x[1s])[2s])`, which no compiler produces. +/// +/// A range function has to apply directly to a selector, and a selector is raw samples, so a tree +/// with a window above a window is only reachable by writing it down by hand - which is a tree +/// [`FusedSeriesExec`](super::FusedSeriesExec) accepts, and therefore one it has to answer correctly. +/// +/// The inner window is the interesting one: `last_over_time` keeps `__name__`, so the pair it can +/// split is not two series whose names it dropped but `{host="h", zone=""}` and `{host="h"}`, one +/// canonical label set and two raw keys. Its fragments are not merged unless the node above it +/// merges them, and that merge happens once per execution: with one execution per partition each +/// partition sees one fragment, so the outer window reads one half of the pair. +fn nested_window_tree() -> NativeTree { + let grid = grid(); + // The inner window's output, which is also the outer window's output: a window reports one value + // per input field on its own grid. + let evaluated = NativeType::Samples(Arc::new(SeriesSchema { + fields: raw_schema().fields.clone(), + eval_grid: Some(Arc::clone(&grid)), + })); + let raw = NativeType::Samples(raw_schema()); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }), + op: BoundWindow { + function: RangeFunction::LastOverTime, + window: RangeWindow { + range_ms: 1_000, + offset_ms: 0, + anchor_ms: None, + }, + // The inner window's input is raw samples, so it has no grid of its own. + input_grid: None, + output_grid: Arc::clone(&grid), + }, + }, + output: evaluated.clone(), + }), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: 2_000, + offset_ms: 0, + anchor_ms: None, + }, + // The outer window reads the inner one's evaluated output, on that output's grid, + // and reports one value per field there like the inner one does. + input_grid: Some(Arc::clone(&grid)), + output_grid: Arc::clone(&grid), + }, + }, + output: evaluated, + }) + .expect("the bound tree is legal") +} + +/// The same tree asking for `quantile`, whose state the runtime does not keep yet. +fn quantile_aggregate_tree() -> NativeTree { + let complete = aggregate_tree(); + let root = complete.root(); + let NativeExpr::Aggregate { input, op } = &root.expr else { + unreachable!("the fixture is an aggregate over a window"); + }; + NativeTree::try_new(NativeNode { + expr: NativeExpr::Aggregate { + input: input.clone(), + op: BoundAggregate { + kind: AggregateKind::Quantile { quantile: 0.5 }, + group: op.group.clone(), + stage: AggregateStage::Complete, + }, + }, + output: root.output.clone(), + }) + .expect("the bound tree is legal") +} + +/// The `rate(x[5m])` tree with a `sum by (job)` over it. +/// +/// A `Complete` aggregate: every contribution to each group lies inside one execution. +fn aggregate_tree() -> NativeTree { + let rate = rate_tree(); + let window = rate.root().clone(); + let grid = window + .output + .samples() + .expect("a window produces samples") + .eval_grid + .clone() + .expect("a window's output is evaluated on a grid"); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(window), + op: BoundAggregate { + kind: AggregateKind::Sum, + group: Grouping::By(Arc::new(["job".into()])), + stage: AggregateStage::Complete, + }, + }, + output: NativeType::Samples(Arc::new(SeriesSchema { + fields: raw_schema().fields.clone(), + eval_grid: Some(grid), + })), + }) + .expect("the bound tree is legal") +} + +/// The tree an instant selection compiles to: the last raw sample at or before each grid point. +/// +/// Unlike a window or a map, a selection rewrites no label: the series it answers carries the labels +/// its input did. +fn instant_tree() -> NativeTree { + let grid = grid(); + let evaluated = NativeType::Samples(Arc::new(SeriesSchema { + fields: raw_schema().fields.clone(), + eval_grid: Some(Arc::clone(&grid)), + })); + let raw = NativeType::Samples(raw_schema()); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Instant { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }), + op: BoundInstant { + lookback_ms: RANGE_MS, + offset_ms: 0, + output_grid: grid, + }, + }, + output: evaluated, + }) + .expect("the bound tree is legal") +} + +/// What the tree's input was compiled against, for the input schema and output tags given. +fn binding_of( + tree: &NativeTree, + series_schema: Arc, + tags: Vec, +) -> FusedSeriesBinding { + FusedSeriesBinding { + series_schema, + grid: tree + .root() + .output + .samples() + .expect("this tree produces samples") + .eval_grid + .clone() + .expect("this tree's output is evaluated on a grid"), + metric_name: Some(Arc::from(METRIC)), + time_index: TIME.to_string(), + tags, + } +} + +/// What the tree's input was compiled against, read off the tree itself. +fn binding(tree: &NativeTree) -> FusedSeriesBinding { + binding_of(tree, raw_schema(), vec![HOST.to_string(), NAME.to_string()]) +} + +/// The schema of a frame carrying series of `schema` with their labels explicit. +fn frame_schema_of(schema: &SeriesSchema) -> SchemaRef { + raw_series_frame_schema(schema, RawLabelMode::Explicit).expect("explicit labels are a frame") +} + +/// The schema of a frame carrying these series, which is what the node declares as its input. +fn frame_schema() -> SchemaRef { + frame_schema_of(&raw_schema()) +} + +/// The columns this node produces: the field, the tags and the grid timestamps. +fn output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +fn input_plan(schema: SchemaRef) -> LogicalPlan { + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.to_dfschema_ref().unwrap(), + }) +} + +/// A physical input producing `batches` in one partition, with `schema` as its schema. +fn memory_exec_with(schema: SchemaRef, batches: Vec) -> Arc { + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[batches], schema, None).unwrap(), + ))) +} + +/// A physical input producing `batches` in one partition, on the frame schema of these tests. +fn memory_exec(batches: Vec) -> Arc { + memory_exec_with(frame_schema(), batches) +} + +/// A physical input producing each of `partitions`' batches in a partition of its own, with +/// `schema` as its schema. +fn memory_exec_partitioned_with( + schema: SchemaRef, + partitions: Vec>, +) -> Arc { + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&partitions, schema, None) + .expect("every partition holds batches of the given frame schema"), + ))) +} + +/// A physical input producing each of `partitions`' batches in a partition of its own. +fn memory_exec_partitioned(partitions: Vec>) -> Arc { + memory_exec_partitioned_with(frame_schema(), partitions) +} + +/// A task context whose memory pool is `pool`, so what the node charges can be read afterwards. +fn task_context(pool: &Arc) -> Arc { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(pool) as Arc) + .build() + .expect("the runtime environment builds"); + Arc::new(TaskContext::default().with_runtime(Arc::new(runtime))) +} + +/// A physical input whose one partition's stream `make` builds. +fn staged( + schema: SchemaRef, + make: impl Fn() -> SendableRecordBatchStream + Send + Sync + 'static, +) -> Arc { + Arc::new(StagedPlan { + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )), + make: Box::new(make), + }) +} + +/// A physical input, so a test can stage what a memory source cannot: a pending poll, an error +/// after some rows, and a stream whose drop it can see. +struct StagedPlan { + make: Box SendableRecordBatchStream + Send + Sync>, + properties: Arc, +} + +impl std::fmt::Debug for StagedPlan { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("StagedPlan").finish() + } +} + +impl DisplayAs for StagedPlan { + fn fmt_as(&self, _t: DisplayFormatType, _f: &mut Formatter) -> fmt::Result { + Ok(()) + } +} + +impl ExecutionPlan for StagedPlan { + fn name(&self) -> &str { + "StagedPlan" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DataFusionResult> { + Ok(self) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + // A scan answers the partitions it declares and no others, which is what makes an index + // past this plan's count an error the node above has to pass on. + if partition >= self.properties.partitioning.partition_count() { + return Err(DataFusionError::Execution(format!( + "StagedPlan has no partition {partition}" + ))); + } + Ok((self.make)()) + } +} + +/// The stream a [`StagedPlan`] hands over: `batches`, then an optional error, with a pending poll +/// and a drop a test can observe. +struct StagedStream { + schema: SchemaRef, + batches: VecDeque, + handed: usize, + /// Hand this many batches, then answer [`Poll::Pending`] once. + pending_after: Option, + /// The error to report once the batches are exhausted. + error: Option<&'static str>, + /// Set when this stream is dropped, so a test can see the input being released. + dropped: Option>, + /// Counts the batches this stream has handed over, so a test can see how far one poll read. + handed_out: Option>, +} + +impl StagedStream { + fn new(schema: SchemaRef, batches: Vec) -> Self { + Self { + schema, + batches: batches.into(), + handed: 0, + pending_after: None, + error: None, + dropped: None, + handed_out: None, + } + } + + fn pending_after(mut self, batches: usize) -> Self { + self.pending_after = Some(batches); + self + } + + fn error(mut self, message: &'static str) -> Self { + self.error = Some(message); + self + } + + fn dropped(mut self, flag: &Arc) -> Self { + self.dropped = Some(Arc::clone(flag)); + self + } + + fn counted(mut self, counter: &Arc) -> Self { + self.handed_out = Some(Arc::clone(counter)); + self + } +} + +impl Stream for StagedStream { + type Item = DataFusionResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + if this.pending_after == Some(this.handed) { + this.pending_after = None; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + match this.batches.pop_front() { + Some(batch) => { + this.handed += 1; + if let Some(counter) = &this.handed_out { + counter.fetch_add(1, Ordering::SeqCst); + } + Poll::Ready(Some(Ok(batch))) + } + None => match this.error.take() { + Some(message) => { + Poll::Ready(Some(Err(DataFusionError::Execution(message.to_string())))) + } + None => Poll::Ready(None), + }, + } + } +} + +impl RecordBatchStream for StagedStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +impl Drop for StagedStream { + fn drop(&mut self) { + if let Some(dropped) = &self.dropped { + dropped.store(true, Ordering::SeqCst); + } + } +} + +fn fused(tree: NativeTree) -> FusedSeries { + let binding = binding(&tree); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The physical node of `node` over `input`, reading the complete labels every frame carries. +fn plan_of( + node: FusedSeries, + input: Arc, +) -> DataFusionResult> { + node.to_execution_plan(input, RawLabelSource::Explicit, SeriesOrdering::Contiguous) +} + +/// The physical node of `tree` over `input`. +fn plan( + tree: NativeTree, + input: Arc, +) -> DataFusionResult> { + plan_of(fused(tree), input) +} + +/// The timestamps of `rows`. +fn times_of(rows: &[(i64, f64)]) -> Vec { + rows.iter().map(|(time, _)| *time).collect() +} + +/// The values of `rows`, one sample per timestamp. +fn values_of(rows: &[(i64, f64)]) -> Vec> { + rows.iter().map(|(_, value)| Some(*value)).collect() +} + +/// One missing sample per timestamp of `rows`, which is a field the series holds no value in. +fn no_values(rows: &[(i64, f64)]) -> Vec> { + rows.iter().map(|_| None).collect() +} + +/// One chunk of `metric`, keyed `key`, holding `times` and one sample per field of `schema`. +fn chunk_of( + schema: &SeriesSchema, + key: &str, + metric: MetricName, + times: &[i64], + values: Vec>>, + last: bool, +) -> RawSeriesChunk { + assert_eq!( + schema.fields.len(), + values.len(), + "one value list per field of the schema" + ); + assert!( + values.iter().all(|values| values.len() == times.len()), + "one sample per timestamp per field" + ); + let fields = values + .into_iter() + .map(|values| Samples::from_float(Float64Array::from(values))) + .collect::>(); + RawSeriesChunk::try_new( + schema, + Buffer::from_slice_ref(key.as_bytes()), + Some(metric), + RawTimestamps::try_new(ScalarBuffer::from(times.to_vec())).expect("ascending timestamps"), + SeriesFields::try_from_vec(fields).expect("one sample per timestamp per field"), + last, + ) + .expect("one series' samples are a legal raw chunk") +} + +/// One chunk of the series `{__name__=METRIC, host=}`: `rows` as its samples, and the chunk +/// that closes the series when `last`. The host names the series, so it doubles as the key. +fn series_chunk(host: &str, rows: &[(i64, f64)], last: bool) -> RawSeriesChunk { + chunk_of( + &raw_schema(), + host, + metric(&[(HOST, host)], Some(METRIC)), + ×_of(rows), + vec![values_of(rows)], + last, + ) +} + +/// The metric these tests read, with `labels` on it and the name given. +fn metric(labels: &[(&str, &str)], name: Option<&str>) -> MetricName { + MetricName::try_new( + name.map(Arc::from), + labels + .iter() + .map(|(name, value)| Label { + name: Arc::from(*name), + value: Arc::from(*value), + }) + .collect(), + ) + .expect("the labels of these tests are legal") +} + +/// One frame batch carrying `chunk` of `schema`, packed by the scan's own encoder. +fn frame_of(schema: &SeriesSchema, chunk: &RawSeriesChunk) -> RecordBatch { + let mut batches = pack_raw_series_chunk(schema, RawLabelMode::Explicit, chunk) + .expect("a chunk of this schema packs into a frame"); + assert_eq!(1, batches.len(), "one small chunk is one frame batch"); + batches.remove(0) +} + +/// One frame batch carrying `chunk` of the one-field schema of these tests. +fn frame(chunk: &RawSeriesChunk) -> RecordBatch { + frame_of(&raw_schema(), chunk) +} + +/// The frames carrying `chunks` of `schema`, one batch each, so a test can say how much input one +/// poll reads. +fn frames_of(schema: &SeriesSchema, chunks: &[RawSeriesChunk]) -> Vec { + chunks.iter().map(|chunk| frame_of(schema, chunk)).collect() +} + +/// The frames carrying `chunks` of the one-field schema of these tests, one batch each. +fn frames(chunks: &[RawSeriesChunk]) -> Vec { + frames_of(&raw_schema(), chunks) +} + +/// One series' samples: one per minute from 0, `scale` seconds per second. +fn series_samples(scale: f64) -> Vec<(i64, f64)> { + (0..SAMPLES) + .map(|index| { + let time = index * SAMPLE_STEP; + (time, scale * time as f64 / 1000.0) + }) + .collect() +} + +/// Two series, each split over two chunks. A series' chunks stay adjacent, and the second of each +/// pair closes it, which is the input contract: a series' chunks arrive together. +fn two_series_frames() -> Vec { + let first = series_samples(1.0); + let second = series_samples(2.0); + frames(&[ + series_chunk("web-1", &first[..3], false), + series_chunk("web-1", &first[3..], true), + series_chunk("web-2", &second[..3], false), + series_chunk("web-2", &second[3..], true), + ]) +} + +/// One series split over three chunks, so its chunks have to be stitched. +fn one_series_frames() -> Vec { + let rows = series_samples(1.0); + frames(&[ + series_chunk("web-1", &rows[..4], false), + series_chunk("web-1", &rows[4..7], false), + series_chunk("web-1", &rows[7..], true), + ]) +} + +/// The rows of `batches`: timestamp, value, host and metric name. +type Row = (i64, Option, Option, Option); + +fn rows(batches: &[RecordBatch]) -> Vec { + let mut rows = Vec::new(); + for batch in batches { + let time = batch + .column_by_name(TIME) + .expect("every batch has a time column") + .as_any() + .downcast_ref::() + .expect("the time column is a millisecond timestamp"); + let value = batch + .column_by_name(VALUE) + .expect("every batch has the field column") + .as_any() + .downcast_ref::() + .expect("the field column is Float64"); + let host = batch + .column_by_name(HOST) + .expect("every batch has the tag column") + .as_any() + .downcast_ref::() + .expect("the tag column is Utf8"); + let name = batch + .column_by_name(NAME) + .expect("every batch has the metric name column") + .as_any() + .downcast_ref::() + .expect("the metric name column is Utf8"); + for row in 0..batch.num_rows() { + rows.push(( + time.value(row), + value.is_valid(row).then(|| value.value(row)), + host.is_valid(row).then(|| host.value(row).to_string()), + name.is_valid(row).then(|| name.value(row).to_string()), + )); + } + } + rows +} + +/// The output of [`zone_node`]: the field, `host`, `zone`, the dropped name and the grid times. +fn zone_output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(ZONE, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +/// `rate(x[5m])` over frames that carry `zone` as well as `host`, with both in the output. +fn zone_node() -> FusedSeries { + let tree = rate_tree(); + let binding = binding_of( + &tree, + raw_schema(), + vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], + ); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + zone_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The instant-selection node over frames that carry `zone` as well as `host`. +/// +/// Its output labels are the input's, so an empty-valued label reaches this node's own state as the +/// frame wrote it. +fn instant_node() -> FusedSeries { + let tree = instant_tree(); + let binding = binding_of( + &tree, + raw_schema(), + vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], + ); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + zone_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The two chunks of the collision: one series once with an empty `zone` and once with no `zone`. +/// +/// They are one identity to Prometheus — an empty label value is no label — but two different raw +/// series: the scan encodes the empty value and the absent one differently, so they reach the +/// runtime as two series that close on the same label set. +fn zone_chunks() -> Vec { + let rows = series_samples(1.0); + vec![ + chunk_of( + &raw_schema(), + "zone-empty", + metric(&[(HOST, "web-1"), (ZONE, "")], Some(METRIC)), + ×_of(&rows), + vec![values_of(&rows)], + true, + ), + chunk_of( + &raw_schema(), + "zone-absent", + metric(&[(HOST, "web-1")], Some(METRIC)), + ×_of(&rows), + vec![values_of(&rows)], + true, + ), + ] +} + +/// The canonical label set these tests collide on, as the runtime renders it in its refusal. +const DUPLICATE_LABEL_SET: &str = + "the node at root produces two series under the label set [(\"host\", \"web-1\")]"; + +/// Two raw series that are already one Prometheus identity are refused by the runtime that would +/// have to answer both. +/// +/// `{host="web-1", zone=""}` and `{host="web-1"}` are one identity to Prometheus: an empty label +/// value is no label, so the two chunks are the same series and `rate` creates no collision here. +/// What the two *raw* keys differ in is the empty `zone` against the absent one, which the scan +/// encodes differently and the decoder therefore opens two input series for. `rate` drops the +/// metric name and the empty value, so the second raw series arrives on the canonical label set the +/// first one already occupies; the refusal names that label set rather than answering it, because +/// answering would report one point twice. +#[tokio::test] +async fn differing_raw_keys_that_close_on_one_label_set_are_refused() { + let exec = plan_of(zone_node(), memory_exec(frames(&zone_chunks()))) + .expect("the runtime executes a rate tree"); + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("the second raw key cannot be answered under the first one's label set"); + assert!( + error.to_string().contains(DUPLICATE_LABEL_SET), + "refused for another reason: {error}" + ); +} + +/// A runtime error after buffered rows is reported, and the terminated stream releases its input. +/// +/// The decoder holds the batch it is reading, the chunks it decoded from it and its own source, so a +/// stream that kept its input after the error would hold all three until the consumer drops it. The +/// collision above is the error: it happens on the second series, after the first one's rows have +/// been handed out. +#[tokio::test] +async fn a_runtime_error_after_rows_releases_the_input() { + let dropped = Arc::new(AtomicBool::new(false)); + let schema = frame_schema(); + let batches = frames(&zone_chunks()); + let source = staged(schema.clone(), { + let dropped = Arc::clone(&dropped); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).dropped(&dropped)) + }); + let mut stream = plan_of(zone_node(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + let mut answered = Vec::new(); + let error = loop { + match stream.next().await { + Some(Ok(batch)) => answered.push(batch), + Some(Err(error)) => break error, + None => panic!( + "the collision was answered instead of refused: {:?}", + rows(&answered) + ), + } + }; + assert!( + error.to_string().contains(DUPLICATE_LABEL_SET), + "refused for another reason: {error}" + ); + // The first series' rows were handed out before the error, and the input is gone although the + // terminated stream is still alive here. + assert!( + !answered.is_empty(), + "the first series is answered before the error" + ); + assert!( + dropped.load(Ordering::SeqCst), + "the terminated stream still holds its input" + ); +} + +/// Two series with *different* names collapse once `rate` drops the name. +/// +/// This is the collision the canonical label set exists for, at the layer that sees it: +/// `http_requests_total{host="web-1"}` and `http_errors_total{host="web-1"}` reach the window as two +/// identities and leave it as `{host="web-1"}`. [`FusedSeries`] reads one metric per scan, so this +/// collision cannot be staged through it; the runtime is driven directly, with the names on the raw +/// chunks. +#[test] +fn two_differently_named_series_collide_once_rate_drops_the_name() { + let tree = rate_tree(); + let mut execution = NativeExecutor::new( + NativeExecutionState::try_new( + &tree, + Arc::new(UnboundedMemoryPool::default()), + &ExecutionPlanMetricsSet::new(), + 0, + ) + .expect("the runtime executes a rate tree"), + ); + let samples = series_samples(1.0); + execution + .push_raw(chunk_of( + &raw_schema(), + METRIC, + metric(&[(HOST, "web-1")], Some(METRIC)), + ×_of(&samples), + vec![values_of(&samples)], + true, + )) + .expect("the first series is a legal raw chunk"); + assert!( + execution + .next_chunk() + .expect("the first series runs") + .is_some(), + "the first series is answered" + ); + execution + .push_raw(chunk_of( + &raw_schema(), + "http_errors_total", + metric(&[(HOST, "web-1")], Some("http_errors_total")), + ×_of(&samples), + vec![values_of(&samples)], + true, + )) + .expect("the second series is a legal raw chunk"); + let error = execution + .next_chunk() + .expect_err("the second series puts a value where the first one already holds one"); + assert!( + error.to_string().contains(DUPLICATE_LABEL_SET), + "refused for another reason: {error}" + ); +} + +/// The grid timestamps of the statement, in order. +fn grid_timestamps() -> Vec { + (0..GRID_LEN) + .map(|index| GRID_START + index as i64 * GRID_STEP) + .collect() +} + +/// A waker a test can ask whether it was woken. +/// +/// A poll that stops on its own budget has to wake itself, and the default waker a `poll!` uses +/// swallows that, so the tests that check it poll with this one. +#[derive(Default)] +struct CountingWake(AtomicBool); + +impl futures::task::ArcWake for CountingWake { + fn wake_by_ref(woken: &Arc) { + woken.0.store(true, Ordering::SeqCst); + } +} + +impl CountingWake { + fn woken(&self) -> bool { + self.0.load(Ordering::SeqCst) + } +} + +/// One series as many chunks that settle nothing: one sample per chunk, all before the grid starts. +/// +/// The counter advances one second per second, so `rate` answers 1 at every grid point once the +/// series is read to its end; until a sample reaches the grid's first point, a push decides +/// nothing and settles no chunk. +fn quiet_series_frames() -> Vec { + frames( + &(0..QUIET_SAMPLES) + .map(|index| { + series_chunk( + "web-1", + &[(index * QUIET_STEP, index as f64)], + index == QUIET_SAMPLES - 1, + ) + }) + .collect::>(), + ) +} + +/// Asserts that `rows` are `rates`' series, each over the whole grid in order. +fn assert_series_rates(rows: &[Row], rates: &[(f64, &str)]) { + let expected = rates + .iter() + .flat_map(|(rate, host)| { + grid_timestamps() + .into_iter() + .map(move |time| (time, *rate, *host)) + }) + .collect::>(); + assert_eq!(expected.len(), rows.len(), "{rows:?}"); + for (row, (expected_time, expected_rate, expected_host)) in rows.iter().zip(expected) { + let (time, value, host, _) = row; + assert_eq!(&expected_time, time); + assert_eq!(Some(expected_host), host.as_deref()); + let value = value.expect("rate has a value at every grid point"); + assert!( + (expected_rate - value).abs() < 1e-9, + "{expected_host} at {expected_time}: {value} is not {expected_rate}" + ); + } +} + +fn count(metrics: &MetricsSet, name: &str) -> usize { + match metrics.sum_by_name(name) { + Some(MetricValue::Count { count, .. }) => count.value(), + other => panic!("no count named {name}: {other:?}"), + } +} + +/// Prints `node` the way the logical `EXPLAIN` prints it. +struct Explain<'a>(&'a FusedSeries); + +impl Display for Explain<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + self.0.fmt_for_explain(f) + } +} + +/// The logical node prints the query a reader wrote, not the tree it compiled to. +#[test] +fn fmt_for_explain_prints_the_promql() { + assert_eq!( + "FusedSeries: rate(x[5m])", + Explain(&fused(rate_tree())).to_string() + ); +} + +/// The physical node prints the tree it was built from, rendered by the tree's own display. +#[test] +fn fmt_as_prints_the_native_tree() { + let exec = + plan(rate_tree(), memory_exec(Vec::new())).expect("the runtime executes a rate tree"); + let text = displayable(exec.as_ref()).indent(false).to_string(); + + assert!( + text.starts_with("FusedSeriesExec: rate(x[5m])"), + "the node line names the query first: {text}" + ); + // A distinctive line of the compiled tree: the window's bound range, offset, anchor and grid. + assert!( + text.contains("root Window rate [300000ms] offset 0 anchor - on grid 300000/60000x5"), + "the node line carries the tree it runs: {text}" + ); + assert!( + text.contains("input Input raw samples"), + "the node line carries the tree's source: {text}" + ); +} + +/// `rate(x[5m])` runs end to end over two raw series: values, tags and the dropped metric name. +#[tokio::test] +async fn runs_rate_over_two_series() { + let exec = plan(rate_tree(), memory_exec(two_series_frames())) + .expect("the runtime executes a rate tree"); + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("the node runs"); + let rows = rows(&batches); + + // A counter advancing one second per second rates at 1, and at 2 for the series that advances + // two: the value at each grid point says the right window was read for the right series. + let expected = [(1.0, "web-1"), (2.0, "web-2")] + .into_iter() + .flat_map(|(rate, host)| { + grid_timestamps() + .into_iter() + .map(move |time| (time, rate, host)) + }) + .collect::>(); + assert_eq!(expected.len(), rows.len(), "{rows:?}"); + for (row, (expected_time, expected_rate, expected_host)) in rows.iter().zip(expected) { + let (time, value, host, _) = row; + assert_eq!(&expected_time, time); + assert_eq!(Some(expected_host), host.as_deref()); + let value = value.expect("rate has a value at every grid point"); + assert!( + (expected_rate - value).abs() < 1e-9, + "{expected_host} at {expected_time}: {value} is not {expected_rate}" + ); + } + // `rate` drops `__name__`, so the tag column that reports it is empty for every row. + assert!( + rows.iter().all(|(_, _, _, name)| name.is_none()), + "rate drops the metric name: {rows:?}" + ); +} + +/// A series that spans several input chunks keeps its absolute grid positions and closes once. +#[tokio::test] +async fn a_series_spanning_chunks_keeps_absolute_positions_and_closes_once() { + let exec = plan(rate_tree(), memory_exec(one_series_frames())) + .expect("the runtime executes a rate tree"); + let batches = collect(Arc::clone(&exec), SessionContext::default().task_ctx()) + .await + .expect("the node runs"); + let rows = rows(&batches); + + // Every grid position exactly once, in order: a chunk's range is never read as a row number. + let times = rows.iter().map(|row| row.0).collect::>(); + assert_eq!(grid_timestamps(), times); + assert!( + rows.iter().all(|(_, value, ..)| value.is_some()), + "every grid point of this series has a value: {rows:?}" + ); + + // The source closed the series on exactly one chunk, and the window saw exactly that one + // closing: one `is_last`, not one per input batch. + let metrics = exec + .metrics() + .expect("the node reports its runtime's metrics"); + assert_eq!(1, count(&metrics, "native.root.input.completed_series")); + assert_eq!(1, count(&metrics, "native.root.completed_series")); +} + +/// A tree the runtime does not execute is refused when the execution plan is built. +/// +/// The runtime folds the supported `Complete` kinds now, so the tree below asks for one it still +/// refuses on purpose: `quantile` needs the parameter and the member buffer that the other kinds +/// do not. What the test pins is the refusal happening while the plan is built, not the particular +/// kind of node that provokes it. +#[tokio::test] +async fn a_tree_the_runtime_cannot_run_is_rejected_when_the_exec_is_built() { + let node = fused(quantile_aggregate_tree()); + let error = plan_of(node, memory_exec(Vec::new())) + .expect_err("the runtime does not execute a quantile aggregate"); + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "the runtime refuses the tree it cannot run: {error}" + ); + assert!(error.to_string().contains("Quantile"), "{error}"); +} + +/// A tree evaluated on another grid than the node declares is refused before it runs. +#[test] +fn a_tree_bound_to_another_grid_is_rejected() { + let tree = rate_tree(); + let other_grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(0, GRID_STEP, GRID_LEN).unwrap()).unwrap()); + let binding = FusedSeriesBinding { + grid: other_grid, + ..binding(&tree) + }; + let error = FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect_err("the tree is not evaluated on the declared grid"); + assert!(error.to_string().contains("grid"), "{error}"); +} + +/// A grid allocated separately from the tree's, with the same spec, is the tree's grid. +/// +/// The comparison is by spec, which is the rule [`NativeTree`]'s own validation and the kernels +/// apply, and a chunk's timestamps come from its absolute grid positions rather than from the +/// allocation, so nothing downstream needs the `Arc` to be shared. This pairs with +/// [`a_tree_bound_to_another_grid_is_rejected`]: a different spec is still refused. +#[test] +fn a_separately_allocated_grid_with_an_equal_spec_is_accepted() { + let tree = rate_tree(); + let equal = Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, GRID_LEN).unwrap()).unwrap(), + ); + assert!( + !Arc::ptr_eq(&equal, &binding(&tree).grid), + "the two grids are separate allocations" + ); + let binding = FusedSeriesBinding { + grid: equal, + ..binding(&tree) + }; + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect("an equal grid spec is the same grid"); +} + +/// An input that ends without a series still runs the tree, and yields no rows. +#[tokio::test] +async fn an_empty_input_runs_the_tree_and_yields_no_rows() { + let exec = + plan(rate_tree(), memory_exec(Vec::new())).expect("the runtime executes a rate tree"); + let batches = collect(Arc::clone(&exec), SessionContext::default().task_ctx()) + .await + .expect("a tree with no series in it still closes"); + assert!(rows(&batches).is_empty(), "{:?}", rows(&batches)); + // The run closed on no series, so nothing was completed and nothing was decided. + let metrics = exec + .metrics() + .expect("the node reports its runtime's metrics"); + assert_eq!(0, count(&metrics, "native.root.completed_series")); + assert_eq!(0, count(&metrics, "native.root.input.completed_series")); +} + +/// An upstream error after some rows is reported, and the terminated stream releases that input. +#[tokio::test] +async fn an_upstream_error_after_rows_is_reported() { + let dropped = Arc::new(AtomicBool::new(false)); + let schema = frame_schema(); + let batches = one_series_frames(); + let source = staged(schema.clone(), { + let dropped = Arc::clone(&dropped); + move || { + Box::pin( + StagedStream::new(schema.clone(), batches.clone()) + .error("the source failed") + .dropped(&dropped), + ) + } + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The error is reported rather than swallowed by the rows the input already handed over. + let mut answered = Vec::new(); + let error = loop { + match stream.next().await { + Some(Ok(batch)) => answered.push(batch), + Some(Err(error)) => break error, + None => panic!("the source failure was swallowed: {:?}", rows(&answered)), + } + }; + assert!( + error.to_string().contains("the source failed"), + "reported another error: {error}" + ); + // Still alive here: the release is the stream's, not the drop of a temporary. + assert!( + dropped.load(Ordering::SeqCst), + "the terminated stream still holds its input" + ); +} + +/// A pending poll loses nothing when the next-item future is dropped and the stream polled again. +#[tokio::test] +async fn a_pending_boundary_loses_no_rows() { + let expected = rows( + &collect( + plan(rate_tree(), memory_exec(one_series_frames())) + .expect("the runtime executes a rate tree"), + SessionContext::default().task_ctx(), + ) + .await + .expect("the node runs"), + ); + assert!(!expected.is_empty(), "the fixture decides some grid points"); + + let schema = frame_schema(); + let batches = one_series_frames(); + let source = staged(schema.clone(), move || { + Box::pin( + StagedStream::new(schema.clone(), batches.clone()) + // One batch is in the decoder's hands when the stream reports pending. + .pending_after(1), + ) + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The future is dropped at the end of this block, as a consumer that gives up waiting does. + { + let next = stream.next(); + futures::pin_mut!(next); + assert!( + futures::poll!(next.as_mut()).is_pending(), + "the input's pending poll reaches the consumer" + ); + } + let batches = stream + .try_collect::>() + .await + .expect("the stream resumes where the pending poll left it"); + assert_eq!(expected, rows(&batches)); +} + +/// A frame whose labels are encoded in its key is refused when the plan is built. +/// +/// The tags this node declares are read off each chunk's labels, and a key-mode frame carries none: +/// decoding it with a binding that promises complete labels would produce chunks that name no +/// series. The frame's own declared layout says which of the two it is, so the mismatch is caught +/// before the first batch, not when the missing labels reach the output. +#[test] +fn a_frame_whose_labels_are_encoded_is_rejected_when_the_plan_is_built() { + let keyed = + raw_series_frame_schema(&raw_schema(), RawLabelMode::Key).expect("key labels are a frame"); + let error = plan(rate_tree(), memory_exec_with(keyed, Vec::new())) + .expect_err("the frames carry no labels to read tags from"); + assert!( + error.to_string().contains("raw frame"), + "refused for another reason: {error}" + ); +} + +/// A frame that carries an evaluation grid is refused when the plan is built. +/// +/// A frame holds raw series: its rows are samples, not grid positions, and the grid every output +/// row is addressed on is the tree's own. A schema that declares a grid is one this node would +/// read positions out of, so it is refused rather than decoded. The source below is one partition, +/// which is the shape the node accepts, so what refuses the plan here is the frame's own +/// declaration and not the node's partition count: the refusal has to name the grid. +#[test] +fn a_frame_with_an_evaluation_grid_is_rejected_when_the_plan_is_built() { + let mut metadata = frame_schema().metadata().clone(); + metadata.insert("greptime.series.grid".to_string(), "v1".to_string()); + let gridded = Arc::new(Schema::new_with_metadata( + frame_schema().fields().clone(), + metadata, + )); + // One empty partition: every check but the frame's passes. + let source = memory_exec_with(gridded, Vec::new()); + assert_eq!( + 1, + source.properties().output_partitioning().partition_count(), + "the node only looks at the frame when its input has a partition" + ); + let error = plan(rate_tree(), source).expect_err("the frame declares a grid"); + let message = error_chain(&error); + assert!( + message.contains("raw series schema cannot have an evaluation grid"), + "refused for another reason: {message}" + ); + assert!( + !message.contains("no partitions"), + "the partition count is not what refuses this plan: {message}" + ); +} + +/// A frame of another series schema is refused when the plan is built. +/// +/// The tree's input node was compiled against one schema, and a frame for a different one names +/// fields the tree does not read and omits ones it does: decoding it would hand the tree a column +/// it never bound. Nothing about the two schemas says which one was meant, so the frame is refused +/// instead of half-read. +#[test] +fn a_frame_of_another_series_schema_is_rejected_when_the_plan_is_built() { + let other = frame_schema_of(&two_field_schema()); + let error = plan(rate_tree(), memory_exec_with(other, Vec::new())) + .expect_err("the tree's input reads one field"); + assert!( + error.to_string().contains("raw frame"), + "refused for another reason: {error}" + ); +} + +/// An input with no partitions is refused when the plan is built, and no other count is. +/// +/// The node runs one evaluator per input partition and adds none of its own, so an input with no +/// partition is one that could answer nothing: the plan is refused rather than built as a node that +/// would report an empty result for data it never read. Any non-zero count is the input's to +/// choose: the node takes it and answers each partition's own series. +#[test] +fn an_input_with_no_partitions_is_rejected_when_the_plan_is_built() { + let schema = frame_schema(); + let empty = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[], Arc::clone(&schema), None) + .expect("no partitions is a legal source"), + ))); + let error = plan(rate_tree(), empty).expect_err("there is no partition to run"); + assert!( + error.to_string().contains("input has no partitions"), + "refused for another reason: {error}" + ); + // The only count refused is zero: two partitions build and answer, which is what the + // multi-partition tests below run. + let two = plan(rate_tree(), memory_exec_partitioned(vec![vec![], vec![]])) + .expect("a two-partition input is not refused"); + assert_eq!( + 2, + two.properties().output_partitioning().partition_count(), + "the node takes the input's count" + ); +} + +/// Two windows stacked are refused over more than one partition, because the inner one's fragments +/// need not have been merged. +/// +/// The counterexample is legal PromQL once nesting is expressible: the inner `last_over_time` runs +/// over two raw series that are one canonical label set, `{host="h", zone=""}` holding `(1s, 1)` and +/// `{host="h"}` holding `(2s, 2)`, and the outer `rate` needs both in one window. In a single +/// partition the outer window merges the two fragments and answers `(2s, 1)`; in two partitions each +/// executor is handed one fragment and has one sample to rate, so the outer window answers nothing. +/// That is an empty result rather than an error, which is why no check further down can see it. +/// +/// `last_over_time` keeps `__name__`, so a check that looked only for a dropped metric name accepts +/// this tree. What makes it unsafe is that a window's output can hold two series under one label set +/// whenever it did not merge its input, which a bare tree cannot say. +/// +/// One partition stays accepted: there is nothing to split there, which is what makes the refusal +/// about the split rather than about the shape of the tree. +#[test] +fn a_nested_window_over_two_partitions_is_refused() { + let error = plan( + nested_window_tree(), + memory_exec_partitioned(vec![vec![], vec![]]), + ) + .expect_err("each of two partitions would merge only the fragment it was given"); + let message = error.to_string(); + assert!( + message.contains("`root.input`"), + "the refusal must name the inner window: {message}" + ); + assert!( + message.contains("`root`"), + "the refusal must name the window above it: {message}" + ); + assert!( + message.contains("2 partitions"), + "the refusal must name the partition count: {message}" + ); + + // The same tree over one partition is accepted, because one execution sees every fragment and + // merges them into one series before the window above it evaluates. + plan(nested_window_tree(), memory_exec(Vec::new())) + .expect("a single partition cannot split a logical series"); +} + +/// An instant selection can leave distinct raw labels on one canonical label set for a window. +#[test] +fn an_instant_then_window_over_two_partitions_is_refused() { + let error = plan( + instant_window_tree(), + memory_exec_partitioned(vec![vec![], vec![]]), + ) + .expect_err("each partition could hand the window only one fragment"); + let message = error.to_string(); + assert!( + message.contains("`root.input`"), + "the refusal must name the instant selection: {message}" + ); + assert!( + message.contains("`root`"), + "the refusal must name the consuming window: {message}" + ); + assert!( + message.contains("2 partitions"), + "the refusal must name the partition count: {message}" + ); +} + +/// A single partition cannot split a logical series, even when an instant feeds a window. +#[test] +fn an_instant_then_window_over_one_partition_is_accepted() { + plan(instant_window_tree(), memory_exec(Vec::new())).expect("one partition sees all fragments"); +} + +/// A window directly over raw input remains valid across partitions without a preceding rewriter. +#[test] +fn a_window_over_raw_input_on_two_partitions_is_accepted() { + plan(rate_tree(), memory_exec_partitioned(vec![vec![], vec![]])) + .expect("ordinary input-to-window trees do not split a logical series"); +} + +/// One poll reads at most [`POLL_BUDGET`] chunks of input, and stopping there loses nothing. +/// +/// `rate(x[5m])` decides nothing until a sample reaches the grid's first point, so a series can +/// arrive as many chunks that settle no batch at all: a poll that kept reading them would hold the +/// task's thread until the series ended, which is what stops other tasks on that thread and delays +/// the cancellation a consumer asked for. The first series here is [`QUIET_SAMPLES`] such chunks; +/// the second one covers the grid, so a poll that resumes properly still answers both. +#[tokio::test] +async fn one_poll_reads_at_most_the_budget_and_loses_nothing() { + let mut batches = quiet_series_frames(); + batches.push(frame(&series_chunk("web-2", &series_samples(2.0), true))); + let total = batches.len(); + assert!( + total > POLL_BUDGET, + "the fixture has to outlast one poll: {total} batches" + ); + + let pulled = Arc::new(AtomicUsize::new(0)); + let woken = Arc::new(CountingWake::default()); + let schema = frame_schema(); + let source = staged(schema.clone(), { + let pulled = Arc::clone(&pulled); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).counted(&pulled)) + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The first poll: the quiet series settles nothing, so the budget is what ends this poll. + let waker = futures::task::waker(Arc::clone(&woken)); + let mut context = Context::from_waker(&waker); + assert!( + stream.as_mut().poll_next(&mut context).is_pending(), + "chunks that settle nothing cannot answer a poll" + ); + let read = pulled.load(Ordering::SeqCst); + assert!(read > 0, "the poll did read its input"); + // One step polls the decoder once, and the decoder reads at most one frame per poll: every + // frame here carries one chunk. The bound is therefore one frame per step, and the budget is + // what it is measured in. + assert!( + read <= POLL_BUDGET + 1, + "one poll read {read} of {total} batches, past its budget of {POLL_BUDGET} steps" + ); + assert!( + read < total, + "one poll read the whole input ({read} batches): the budget did not stop it" + ); + assert!( + woken.woken(), + "a poll that stopped on its budget has to wake itself" + ); + + // Nothing the budget stopped is lost: the stream resumes where the poll left it. + let batches = stream + .try_collect::>() + .await + .expect("the stream resumes where the bounded poll left it"); + assert_series_rates(&rows(&batches), &[(1.0, "web-1"), (2.0, "web-2")]); +} + +/// A batch that is ready is handed over before any batch behind it is read. +/// +/// Each series below is one chunk, and a chunk that carries `is_last` closes its series, so the +/// first answer needs only the first chunk. A poll that read on would answer the same batch only +/// after the whole input, which is exactly what the budget above exists to prevent. +#[tokio::test] +async fn a_ready_batch_is_returned_before_later_input_is_read() { + let batches = (0..20) + .map(|index| { + frame(&series_chunk( + &format!("web-{index}"), + &series_samples(1.0), + true, + )) + }) + .collect::>(); + let total = batches.len(); + let pulled = Arc::new(AtomicUsize::new(0)); + let schema = frame_schema(); + let source = staged(schema.clone(), { + let pulled = Arc::clone(&pulled); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).counted(&pulled)) + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + let first = stream + .next() + .await + .expect("the first series is answered") + .expect("the first series is answered without an error"); + let answered = rows(&[first]); + assert_series_rates(&answered, &[(1.0, "web-0")]); + // One batch: the series' own chunk closes it, so nothing behind it has to be read. + assert_eq!( + 1, + pulled.load(Ordering::SeqCst), + "the answer waited for input behind the series it reports" + ); + assert!(total > 1, "the fixture has to outlast the first answer"); +} + +/// A chunk whose rows form several runs is converted row by row from the rows it carries. +/// +/// A sparse result is one chunk holding two runs with a hole between them +/// ([`EvaluatedChunk::range`] answers `None` for it). Reading the hole as rows would invent grid +/// points this series holds no value at, and reading positions as compact rows would shift every +/// row after the hole onto the wrong timestamp: the batch has to hold exactly the carried rows, at +/// their absolute grid times, with nulls nowhere else. +#[test] +fn a_chunk_of_several_runs_converts_its_own_rows_only() { + let grid = grid(); + let axis = TimeAxis::from_grid_spans( + Arc::clone(&grid), + // Grid positions 0..2 and 4..5: position 2 and 3 hold no row at all. + vec![GridSpan::new(0, 2), GridSpan::new(4, 1)], + ) + .expect("the spans are increasing and inside the grid"); + let chunk = EvaluatedChunk::try_new( + Series { + metric: MetricName::try_new(Some(Arc::from(METRIC)), Vec::new()) + .expect("one metric name is a legal identity"), + time: axis, + fields: SeriesFields::try_from_vec(vec![Samples::from_float(Float64Array::from( + // One value per carried row, so the second run's value sits at grid position 4. + vec![7.0, 8.0, 11.0], + ))]) + .expect("one field is a legal field list"), + }, + true, + ) + .expect("a chunk on a grid"); + + // The columns of `output_schema`, in its order: the field, the two tags, then the time. + let batch = series_to_record_batch( + &[ + OutputColumn::Field(0), + OutputColumn::Tag(0), + OutputColumn::Tag(1), + OutputColumn::Time, + ], + &[HOST.to_string(), NAME.to_string()], + &output_schema(), + &chunk, + ) + .expect("the chunk converts"); + + let stamps = batch + .column_by_name(TIME) + .expect("the batch has its time column") + .as_any() + .downcast_ref::() + .expect("the time column is a millisecond timestamp"); + let values = batch + .column_by_name(VALUE) + .expect("the batch has its field column") + .as_any() + .downcast_ref::() + .expect("the field column is Float64"); + assert_eq!(3, batch.num_rows(), "only the carried rows are rows"); + let grid_stamps = grid.timestamps().expect("the grid timestamps build"); + assert_eq!( + vec![grid_stamps[0], grid_stamps[1], grid_stamps[4]], + stamps.values()[..3].to_vec(), + "the carried rows keep their absolute grid positions" + ); + assert_eq!(vec![7.0, 8.0, 11.0], values.values()[..3].to_vec()); +} + +/// Every message in `error`'s chain, so an error a layer above wrapped is still read. +fn error_chain(error: &DataFusionError) -> String { + let mut message = error.to_string(); + let mut source = std::error::Error::source(error); + while let Some(cause) = source { + message.push_str(": "); + message.push_str(&cause.to_string()); + source = cause.source(); + } + message +} + +/// Two partitions, one series each, and both series leave the window on the same label set. +/// +/// `rate` drops `__name__`, so `http_requests_total{host="web-1"}` reads as `{host="web-1"}` in +/// the output. The key column keeps each of the two whole inside its own partition, so a check that +/// lived inside one execution would see one series each and report nothing: only state shared by +/// the partitions sees that the output vector holds the label set twice. +#[tokio::test] +async fn a_duplicate_label_set_across_partitions_is_refused() { + let series = frame(&series_chunk("web-1", &series_samples(1.0), true)); + let exec = plan( + rate_tree(), + memory_exec_partitioned(vec![vec![series.clone()], vec![series]]), + ) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("the second partition lands on the label set the first one claimed"); + let message = error_chain(&error); + assert!( + message.contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {message}" + ); + assert!( + message.contains("host=\"web-1\""), + "the label set must be readable: {message}" + ); +} + +/// The same label set on points that do not overlap is legal, and every point survives. +/// +/// A rate needs two samples in its window, so the first series only has a value at the first grid +/// point and the second only at the last: the two share a final label set but no evaluation +/// timestamp. Prometheus merges such series into one; this node emits both batches and leaves the +/// merge to matrix assembly, so the check has to allow what it would otherwise refuse. +#[tokio::test] +async fn a_duplicate_label_set_on_disjoint_grid_points_is_emitted() { + let first = vec![frame(&series_chunk( + "web-1", + &[(60_000, 60.0), (GRID_START, GRID_START as f64 / 1_000.0)], + true, + ))]; + let last = GRID_START + (GRID_LEN as i64 - 1) * GRID_STEP; + let second = vec![frame(&series_chunk( + "web-1", + &[ + (GRID_START, GRID_START as f64 / 1_000.0), + (last, last as f64 / 1_000.0), + ], + true, + ))]; + let exec = plan(rate_tree(), memory_exec_partitioned(vec![first, second])) + .expect("the runtime executes a rate tree"); + + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("disjoint grid points are legal"); + assert_eq!(2, batches.len(), "each partition emits its own series"); + let mut rows = rows(&batches); + rows.sort_by_key(|row| row.0); + assert_eq!( + vec![GRID_START, last], + rows.iter().map(|row| row.0).collect::>(), + "both series must survive with the one point each holds" + ); + for ((_, value, host, _), rate) in rows.iter().zip([1.0, 1.0]) { + assert_eq!(Some("web-1"), host.as_deref()); + let value = value.expect("the point each series holds has a value"); + assert!((rate - value).abs() < 1e-9, "{value} is not {rate}"); + } +} + +/// What a label set holds stays charged after the stream that first claimed it ends, and what a +/// later partition adds to it is charged as well. +/// +/// The map belongs to the plan: it outlives every partition and is released with the plan. Charging +/// each stream for what it inserted returned the bytes of a label set when *that* stream ended, +/// while the entry the next partition is checked against stayed in the map, so everything the map +/// held after the first partition was invisible to the pool. What is charged here is therefore read +/// between the partitions: it has to survive the first stream, and the runs the second partition +/// merges into the entry have to cost what they hold. +#[tokio::test] +async fn a_label_set_stays_charged_after_its_first_stream_ends() { + // The first series holds a value only at the first grid point and the second only at the last, + // so the two share a final label set without sharing an evaluation timestamp: the second + // partition is legal and has to be answered, and it has to find the first one's entry there. + let first = vec![frame(&series_chunk( + "web-0", + &[(60_000, 60.0), (GRID_START, GRID_START as f64 / 1_000.0)], + true, + ))]; + let last = GRID_START + (GRID_LEN as i64 - 1) * GRID_STEP; + let second = vec![frame(&series_chunk( + "web-0", + &[ + (GRID_START, GRID_START as f64 / 1_000.0), + (last, last as f64 / 1_000.0), + ], + true, + ))]; + let exec = plan(rate_tree(), memory_exec_partitioned(vec![first, second])) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + + // The first partition claims the label set, answers it and ends. + let mut stream = exec + .execute(0, Arc::clone(&context)) + .expect("the node executes"); + let mut answered = Vec::new(); + while let Some(batch) = stream.next().await { + answered.push(batch.expect("the first partition runs")); + } + assert_eq!(1, answered.len(), "one series over one grid point"); + drop(stream); + let charged = pool.reserved(); + assert!( + charged > 0, + "the first partition's label set is still in the plan's map, but nothing is charged for it" + ); + + // The second partition lands on that same label set on points the first one does not hold: it + // is answered, and the run it merges into the entry is what its charge grows by. + let mut stream = exec + .execute(1, Arc::clone(&context)) + .expect("the node executes"); + let mut answered = Vec::new(); + while let Some(batch) = stream.next().await { + answered.push(batch.expect("the second partition runs")); + } + assert_eq!(1, answered.len(), "disjoint points are legal"); + drop(stream); + assert!( + pool.reserved() > charged, + "the second partition's run grew the entry, but nothing was charged for it: {} is not \ + above {charged}", + pool.reserved() + ); + + // The entries live as long as the plan, so dropping it is what gives their bytes back. + drop(exec); + assert_eq!( + 0, + pool.reserved(), + "the plan's drop releases what its label sets held" + ); +} + +/// Each partition runs its own series and answers it, with no label set claimed twice. +#[tokio::test] +async fn every_partition_evaluates_its_own_series() { + let partitions = [1.0, 2.0, 3.0] + .into_iter() + .enumerate() + .map(|(index, scale)| { + vec![frame(&series_chunk( + &format!("web-{index}"), + &series_samples(scale), + true, + ))] + }) + .collect::>(); + let exec = plan(rate_tree(), memory_exec_partitioned(partitions)) + .expect("the runtime executes a rate tree"); + assert_eq!( + 3, + exec.properties().output_partitioning().partition_count(), + "the node's partition count is its input's" + ); + + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("each partition runs its own tree"); + assert_eq!(3, batches.len(), "one batch per partition"); + let rows = rows(&batches); + assert_eq!( + 3 * GRID_LEN as usize, + rows.len(), + "every partition reports its own series over the whole grid" + ); + // Batches arrive in whatever order the partitions settle, so the series are read apart. + for (index, scale) in [1.0, 2.0, 3.0].into_iter().enumerate() { + let host = format!("web-{index}"); + let mut series = rows + .iter() + .filter(|row| row.2.as_deref() == Some(host.as_str())) + .cloned() + .collect::>(); + series.sort_by_key(|row| row.0); + assert_eq!(GRID_LEN as usize, series.len(), "{host}: {series:?}"); + assert_series_rates(&series, &[(scale, host.as_str())]); + } +} + +/// The optimizer inserts no exchange between a two-partition scan and this node. +/// +/// This is the claim the whole arrangement rests on, so it is checked against the real rule rather +/// than against the flags that are supposed to satisfy it: a plan that reports the right +/// requirements can still be repartitioned, and that is exactly how the previous `KeyPartitioned` +/// requirement failed - a hash requirement is classed as necessary for any child with more than one +/// partition, so a hash exchange was inserted even though the child already satisfied it. +/// +/// The input here is a two-partition source, and `target_partitions` is raised above it because +/// that is when an optimizer is most willing to repartition. Both are the conditions under which the +/// old requirement broke, so a plan that survives them is one where the guarantee the scan makes +/// arrives at this node intact. +#[test] +fn the_optimizer_inserts_no_exchange_between_the_input_and_this_node() { + let exec = plan( + rate_tree(), + memory_exec_partitioned(vec![ + vec![frame(&series_chunk("web-0", &series_samples(1.0), true))], + vec![frame(&series_chunk("web-1", &series_samples(2.0), true))], + ]), + ) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs a child with more than one partition" + ); + + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 8; + let optimized = EnsureRequirements::new() + .optimize(Arc::clone(&exec), &config) + .expect("the rule accepts the plan"); + + let exchanges = exchanges_between(&optimized, &exec); + assert!( + exchanges.is_empty(), + "the optimizer inserted {exchanges:?} between the input and the node, and an exchange \ + reorders a series' chunks: the node would close a series it has not finished" + ); + // No exchange is not the same as no change: a rule that raised the count some other way would + // still split a series. The count is the input's, which is what the node reports and what it + // runs one evaluator per. + assert_eq!( + 2, + optimized + .properties() + .output_partitioning() + .partition_count(), + "the optimizer kept no exchange but changed how many partitions the node runs" + ); + assert_eq!( + 2, + optimized.children()[0] + .properties() + .output_partitioning() + .partition_count(), + "the input's own partition count changed under the node" + ); +} + +/// The `RepartitionExec`s in `plan` that sit above `stop` and below no other `RepartitionExec`. +/// +/// Named by the kind the rule chose, so a failure says whether it was a hash or a round-robin +/// exchange: the two are inserted for different reasons and only the hash one is the trap. +fn exchanges_between(plan: &Arc, stop: &Arc) -> Vec { + let mut found = Vec::new(); + let mut stack = vec![Arc::clone(plan)]; + while let Some(node) = stack.pop() { + if Arc::ptr_eq(&node, stop) { + continue; + } + if let Some(exchange) = + node.downcast_ref::() + { + found.push(format!( + "{}({:?})", + exchange.name(), + exchange.partitioning() + )); + } + stack.extend(node.children().into_iter().cloned()); + } + found +} + +/// The node asks the optimizer for nothing, so no exchange can be inserted under it. +/// +/// A window holds the state of the series it is reading and closes it on the chunk carrying +/// `is_last`, so it can answer a series only if the whole series arrives in the partition it reads. +/// What makes that true is the input: a raw series frame comes from `SeriesScan`, whose contract is +/// that the same series (primary key) always reaches the same partition +/// (`src/mito2/src/read/series_scan.rs`), and the frame decoder keeps each series' chunks contiguous +/// within that partition. The node therefore needs the input's own layout kept, not a particular +/// layout of its own. +/// +/// The trap is the requirement writing that as `KeyPartitioned`, which reads like the honest +/// description of the guarantee. A `Hash|KeyPartitioned` requirement is classed as a hash +/// requirement, and a hash requirement forces a hash exchange whenever the child has more than one +/// partition - even where the child already satisfies it, and even when the partition counts match. +/// A hash exchange does not preserve order, so it would interleave one series' chunks across +/// partitions and the decoder would close a series it has not finished. Asking for +/// [`Distribution::UnspecifiedDistribution`] is what leaves the child alone, and declining extra +/// partitioning is the second half of the same thing: no round-robin is added to raise parallelism +/// either. The node is not pinned to one partition — it runs one tree per input partition and keeps +/// the input's count. +#[test] +fn the_input_is_left_alone_so_no_exchange_can_be_inserted() { + let exec = plan(rate_tree(), memory_exec_partitioned(vec![vec![], vec![]])) + .expect("the runtime executes a rate tree"); + let distributions = exec.input_distribution_requirements().into_per_child(); + // `Distribution` has no equality of its own, so the exact requirement is matched structurally: + // one child, and what it asks for is that the child's layout is left alone. + assert!( + matches!( + distributions.as_slice(), + [Distribution::UnspecifiedDistribution] + ), + "anything but leaving the child alone makes a hash exchange necessary: {distributions:?}" + ); + assert_eq!( + vec![false], + exec.benefits_from_input_partitioning(), + "raising the partition count would add a round-robin exchange" + ); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "the node's own count is its input's" + ); +} + +/// A partition index the input does not have is refused rather than answered. +/// +/// The node runs one evaluator per input partition and adds none of its own, so it forwards the +/// index to its input: the count it declares is the input's, and an index past that is the input's +/// refusal, not a series read from another partition. +#[tokio::test] +async fn a_partition_past_the_inputs_count_is_refused() { + let schema = frame_schema(); + let source = staged(schema.clone(), move || { + Box::pin(StagedStream::new(schema.clone(), Vec::new())) + }); + let exec = plan(rate_tree(), source).expect("the runtime executes a rate tree"); + assert_eq!( + 1, + exec.properties().output_partitioning().partition_count(), + "the node's partition count is its input's" + ); + exec.execute(0, SessionContext::default().task_ctx()) + .expect("the one partition the input has"); + let error = match exec.execute(1, SessionContext::default().task_ctx()) { + Ok(_) => panic!("a partition the input does not have is answered"), + Err(error) => error, + }; + assert!( + error.to_string().contains("partition"), + "refused for another reason: {error}" + ); +} + +/// The `rate(x[5m])` tree over two fields, with both fields and the tags in the output. +fn two_field_node() -> FusedSeries { + let tree = rate_tree_of(two_field_schema()); + let binding = binding_of( + &tree, + two_field_schema(), + vec![HOST.to_string(), NAME.to_string()], + ); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema_of(&two_field_schema())), + two_field_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The output of [`two_field_node`]: both fields, the tags and the grid times. +fn two_field_output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(OTHER, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +/// One chunk of `{__name__=METRIC, host=, __type__=}` whose `field` holds `rows` and +/// whose other field is missing everywhere. +fn one_field_chunk(field: usize, host: &str, kind: &str, last: bool) -> RawSeriesChunk { + let rows = series_samples(1.0); + let (value, other) = if field == 0 { + (values_of(&rows), no_values(&rows)) + } else { + (no_values(&rows), values_of(&rows)) + }; + chunk_of( + &two_field_schema(), + &format!("{host}-{kind}"), + metric(&[(HOST, host), (TYPE, kind)], Some(METRIC)), + ×_of(&rows), + vec![value, other], + last, + ) +} + +/// Two series that close on one label set may still fill it, one field each. +/// +/// `rate` drops `__name__` and `__type__`, so the two series below leave the window as +/// `{host="web-1"}` and read as one label set. They hold values in *different* fields, which is the +/// pair [`crate::grid_series::merge_same_labelset`] merges: one value of one field at each position, +/// never two. Occupancy is therefore per field, and refusing this pair would refuse a series +/// Prometheus answers. +#[tokio::test] +async fn a_duplicate_label_set_of_complementary_fields_is_emitted() { + let exec = plan_of( + two_field_node(), + memory_exec_partitioned_with( + frame_schema_of(&two_field_schema()), + vec![ + frames_of( + &two_field_schema(), + &[one_field_chunk(0, "web-1", "counter", true)], + ), + frames_of( + &two_field_schema(), + &[one_field_chunk(1, "web-1", "gauge", true)], + ), + ], + ), + ) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("complementary fields are legal"); + assert_eq!(2, batches.len(), "each partition emits its own series"); + // Each batch answers one field of the shared label set and leaves the other missing: the two + // together are the merged series, with no value of one field holding a position twice. + let answered = batches + .iter() + .map(|batch| { + let host = batch + .column_by_name(HOST) + .and_then(|host| host.as_any().downcast_ref::()) + .expect("every batch has its tag column"); + assert!( + (0..host.len()).all(|row| host.value(row) == "web-1"), + "both batches answer the one label set" + ); + let present = [VALUE, OTHER] + .into_iter() + .filter(|field| { + let column = batch + .column_by_name(field) + .unwrap_or_else(|| panic!("every batch has a {field} column")) + .as_any() + .downcast_ref::() + .expect("the field column is Float64"); + column.null_count() < column.len() + }) + .collect::>(); + assert_eq!( + 1, + present.len(), + "each series fills exactly one field: {present:?}" + ); + present[0] + }) + .collect::>(); + let mut answered = answered; + answered.sort_unstable(); + assert_eq!(vec![OTHER, VALUE], answered, "the two fields are filled"); +} + +/// Two series that close on one label set may not hold the same field at the same position. +/// +/// This is the pair no merge resolves: at one grid position one field of the label set would carry +/// two values. It is refused even though the two series name different `__type__` labels, because +/// the window drops that label and what reaches the output is one label set twice. +#[tokio::test] +async fn a_duplicate_label_set_of_one_field_is_refused() { + let exec = plan_of( + two_field_node(), + memory_exec_partitioned_with( + frame_schema_of(&two_field_schema()), + vec![ + frames_of( + &two_field_schema(), + &[one_field_chunk(0, "web-1", "counter", true)], + ), + frames_of( + &two_field_schema(), + &[one_field_chunk(0, "web-1", "gauge", true)], + ), + ], + ), + ) + .expect("the runtime executes a rate tree"); + + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("both series hold a value of one field at one position"); + let message = error_chain(&error); + assert!( + message.contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {message}" + ); + assert!( + message.contains("host=\"web-1\""), + "the label set must be readable: {message}" + ); +} + +/// Two partitions whose label sets differ only in an empty value hold one label set. +/// +/// `{host="web-1", zone=""}` and `{host="web-1"}` are one series to Prometheus: an empty label value +/// is no label. An instant selection rewrites nothing, so both label sets reach this node's own +/// state as their frames wrote them, and a state that compared the tags as they arrived would file +/// them apart. Neither partition sees the other's series, so nothing else in the plan can refuse +/// the pair: this node's state is what compares them, and what it has to compare is the label set +/// Prometheus compares. +#[tokio::test] +async fn partitions_whose_label_sets_differ_only_in_an_empty_value_are_refused() { + let [empty, absent] = <[_; 2]>::try_from(zone_chunks()).expect("the fixture has two chunks"); + let exec = plan_of( + instant_node(), + memory_exec_partitioned(vec![frames(&[empty]), frames(&[absent])]), + ) + .expect("the runtime executes an instant selection"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("both partitions hold one value of one field at the same position"); + let message = error_chain(&error); + assert!( + message.contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {message}" + ); + assert!( + message.contains("host=\"web-1\""), + "the label set must be readable: {message}" + ); +} + +/// A chunk holding no value records nothing, so a query that answers nothing holds nothing. +/// +/// A selection over a series that staleness or the range left empty still closes that series, and +/// hands over a closing chunk with no rows - [`crate::native::execute::no_row_chunk`] is where it is +/// built. Every chunk reaches [`LabelSets::claim`], so an implementation that wrote an entry per +/// chunk would record one per empty series and charge the plan's reservation for it, for a query +/// whose answer is empty: enough of them would exhaust a finite pool without a single row produced. +/// +/// The entries are checked rather than the charge, because a charge of zero is what a mistake here +/// would still show if the label bytes happened to be free: what must not grow is the table. +#[test] +fn a_chunk_holding_no_value_is_claimed_without_recording_it() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 0)]) + .expect("an empty run over the grid"); + let empty_series = |index: usize| Series { + metric: metric(&[(HOST, &format!("web-{index}"))], Some(METRIC)), + time: axis.clone(), + // Five positions were never filled, so no field occupies one. + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![f64::NAN; 0]))), + }; + + for index in 0..64 { + labels + .claim(&empty_series(index)) + .expect("an empty chunk cannot collide with anything"); + } + + assert_eq!( + 0, + labels.seen.len(), + "64 empty series were recorded, and the table outlives the query: {:?}", + labels.seen.keys().collect::>() + ); + assert_eq!( + 0, + pool.reserved(), + "nothing is held, so nothing may be charged" + ); +} + +/// An empty chunk does not stop a later chunk of the same label set from being recorded. +/// +/// Skipping the empty chunk is only safe while it also records nothing a later one needs: a series +/// that answers nothing holds no position, so a second series under one label set has to be free to +/// claim the positions it really holds. +#[test] +fn a_later_chunk_still_claims_the_label_set_an_empty_one_skipped() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let empty_axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 0)]) + .expect("an empty run over the grid"); + let full_axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 5)]) + .expect("one run over the whole grid"); + + labels + .claim(&Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: empty_axis, + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![f64::NAN; 0]))), + }) + .expect("an empty chunk cannot collide with anything"); + labels + .claim(&Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: full_axis, + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; 5]))), + }) + .expect("the first chunk to hold a value claims the positions it holds"); + + assert_eq!(1, labels.seen.len(), "the positions are recorded once"); + assert!(pool.reserved() > 0, "the recorded positions are charged"); +} + +/// The claim compares the label set Prometheus compares, so an empty value is no label. +/// +/// Driven directly at the layer that keeps the state: a window and a map both drop an empty-valued +/// label before a chunk reaches this node, so no tree this crate can build produces the pair below. +/// The state has to key on the same normalization they apply all the same, or the equality it +/// enforces would be one Prometheus does not have. +#[test] +fn a_claim_compares_label_sets_without_empty_values() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 5)]) + .expect("one run over the whole grid"); + let series_of = |metric: MetricName| Series { + metric, + time: axis.clone(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; 5]))), + }; + + labels + .claim(&series_of(metric(&[(HOST, "h"), (ZONE, "")], Some(METRIC)))) + .expect("the first label set is claimed"); + let error = labels + .claim(&series_of(metric(&[(HOST, "h")], Some(METRIC)))) + .expect_err("an empty label value is no label"); + assert!( + error + .to_string() + .contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {error}" + ); + assert_eq!( + vec![(Arc::from(HOST), Arc::from("h"))], + label_set(&metric(&[(HOST, "h"), (ZONE, "")], None)), + "the empty value is not part of the label set" + ); +} + +/// A series whose chunks move forward, a hole apart, appends to its stored runs rather than +/// re-merging them. +/// +/// One grid point per chunk with a hole after it (`0, 2, 4, ...`) is a legal stream and the shape a +/// short `count_over_time` over sub-sampled, chunked input produces: the chunks of one label set +/// hold their values at disjoint positions, so they merge into one series and no position ever +/// carries two values. What it used to cost was quadratic - `claim` walked the whole stored history +/// to look for a collision and then rebuilt all of it with a merge, for each chunk - and all of it +/// ran under the label-set lock every partition shares, so a sparse series of `K` chunks made the +/// other partitions wait for `K(K+1)/2` run moves. +/// +/// The counter is the runtime's own, so this test measures the append path the native table takes +/// as well: one fix, not two that can drift apart again. +#[test] +fn a_forward_chunk_does_not_rebuild_the_stored_runs() { + // How many one-position chunks the series arrives in: enough that the quadratic cost it used to + // pay would be unmistakable, since that cost was the square of this figure. + const CHUNKS: usize = 64; + + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + // Two positions per chunk, so a hole follows every point the series holds. + let grid = Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 2 * CHUNKS as u64).unwrap()) + .unwrap(), + ); + let chunk = |position: usize| Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(position, 1)]) + .expect("one position of the grid"), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + + let before = absorb_rebuilds(); + for index in 0..CHUNKS { + labels + .claim(&chunk(2 * index)) + .expect("the chunks of one label set hold values at disjoint positions"); + } + + assert_eq!( + 0, + absorb_rebuilds() - before, + "the {CHUNKS} forward chunks rebuilt the stored runs instead of appending to them" + ); + assert_eq!( + 1, + labels.seen.len(), + "every chunk belongs to one label set, so the history is one entry" + ); +} + +/// A backward disjoint claim after a spare-capacity run merge updates the exact memory charge. +#[test] +fn a_backward_disjoint_claim_after_spare_capacity_is_charged_correctly() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + let chunk = |start, len| Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, len)]) + .unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; len]))), + }; + labels.claim(&chunk(10, 2)).unwrap(); + let runs = labels + .seen + .values_mut() + .next() + .unwrap() + .first_mut() + .unwrap(); + let old_capacity = runs.capacity(); + runs.reserve(16); + let new_capacity = runs.capacity(); + labels + .reservation + .as_mut() + .unwrap() + .try_grow((new_capacity - old_capacity) * std::mem::size_of::()) + .unwrap(); + let before = pool.reserved(); + + labels + .claim(&chunk(0, 1)) + .expect("the position is disjoint and earlier than stored history"); + let runs = labels.seen.values().next().unwrap().first().unwrap(); + assert_eq!( + vec![(0, 1), (10, 12)], + runs.iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + assert_eq!( + before - (new_capacity - runs.capacity()) * std::mem::size_of::(), + pool.reserved(), + "the accounting follows the replacement vector's actual capacity" + ); +} + +/// A backward claim that joins two spare-capacity runs charges a replacement vector's shrinkage. +#[test] +fn a_backward_claim_merging_spare_capacity_runs_accounts_for_shrinkage() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 3).unwrap()).unwrap()); + let chunk = |values: Vec>| Series { + metric: metric(&[(HOST, "shrink")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 3)]) + .expect("the chunk covers the small grid"), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(values))), + }; + + labels + .claim(&chunk(vec![Some(1.0), None, Some(1.0)])) + .expect("the first claim stores runs at positions 0 and 2"); + let runs = labels.seen.values().next().unwrap().first().unwrap(); + assert_eq!( + vec![(0, 1), (2, 3)], + runs.iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + let before_capacity = runs.capacity(); + let before_reserved = pool.reserved(); + + labels + .claim(&chunk(vec![None, Some(1.0), None])) + .expect("the disjoint middle position joins both existing runs"); + + let runs = labels.seen.values().next().unwrap().first().unwrap(); + let after_capacity = runs.capacity(); + assert_eq!( + vec![(0, 3)], + runs.iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + assert!( + before_capacity > after_capacity, + "the regression requires replacement-vector shrinkage (before={before_capacity}, after={after_capacity})" + ); + let run_size = std::mem::size_of::(); + assert_eq!( + before_reserved - (before_capacity - after_capacity) * run_size, + pool.reserved(), + "the reservation tracks the replacement vector's actual capacity decrease" + ); +} + +/// A newly retained field is charged for the cloned run vector's capacity, not incoming slack. +#[test] +fn a_new_field_claim_is_charged_for_its_retained_capacity() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 4).unwrap()).unwrap()); + let first = Series { + metric: metric(&[(HOST, "new-field")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 1)]).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + labels.claim(&first).unwrap(); + + let incoming = Series { + metric: metric(&[(HOST, "new-field")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 4)]).unwrap(), + fields: SeriesFields::try_from_vec(vec![ + Samples::from_float(Float64Array::from(vec![None, None, None, None])), + Samples::from_float(Float64Array::from(vec![Some(1.0), None, Some(2.0), None])), + ]) + .unwrap(), + }; + let incoming_runs = occupied_runs(&incoming).unwrap(); + assert_eq!(2, incoming_runs[1].len()); + assert!( + incoming_runs[1].capacity() > incoming_runs[1].len(), + "the incoming run vector has spare capacity" + ); + let before = pool.reserved(); + labels.claim(&incoming).unwrap(); + + let retained = &labels.seen.values().next().unwrap()[1]; + assert_eq!(2, retained.len()); + assert_eq!( + retained.capacity() * std::mem::size_of::(), + pool.reserved() - before, + "the reservation charges the retained cloned vector capacity" + ); + assert!( + retained.capacity() < incoming_runs[1].capacity(), + "the incoming spare capacity was not retained or charged" + ); +} + +/// Chunked fields are traversed sequentially while occupancy is converted into grid runs. +#[test] +fn chunked_samples_occupancy_preserves_sparse_rows() { + let values = Samples::try_from_chunks(vec![ + Samples::from_float(Float64Array::from(vec![Some(1.0), None])), + Samples::from_float(Float64Array::from(vec![None, Some(2.0)])), + ]) + .unwrap(); + let series = Series { + metric: metric(&[(HOST, "chunked")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 4)]).unwrap(), + fields: SeriesFields::one(values), + }; + let runs = occupied_runs(&series).expect("chunked occupancy is read through a forward cursor"); + assert_eq!( + vec![(0, 1), (3, 4)], + runs[0] + .iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); +} + +/// Appending a run that touches stored history reuses its last slot instead of rebuilding it. +#[test] +fn a_touching_append_with_full_capacity_does_not_rebuild_history() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + let chunk = |start| Series { + metric: metric(&[(HOST, "touch")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, 1)]).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + labels.claim(&chunk(0)).unwrap(); + let stored = labels + .seen + .values_mut() + .next() + .unwrap() + .first_mut() + .unwrap(); + stored.shrink_to_fit(); + let cap = stored.capacity(); + let allocation = stored.as_ptr(); + assert_eq!( + stored.len(), + cap, + "the test requires a full-capacity run list" + ); + let before = absorb_rebuilds(); + labels + .claim(&chunk(1)) + .expect("the touching position extends the same run"); + let stored = labels.seen.values().next().unwrap().first().unwrap(); + assert_eq!( + vec![(0, 2)], + stored + .iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + assert_eq!( + cap, + stored.capacity(), + "extending a run did not allocate or copy its history" + ); + assert_eq!( + allocation, + stored.as_ptr(), + "extending a run retained the stored vector allocation" + ); + assert_eq!( + before, + absorb_rebuilds(), + "touching append uses the append path" + ); +} + +/// A failed claim does not mutate the shared entry, while its reservation owner remains alive. +#[test] +fn a_refused_growth_keeps_shared_label_entry_unchanged() { + use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; + let limit = 4096; + let pool = Arc::new(GreedyMemoryPool::new(limit)); + let erased: Arc = pool.clone(); + let mut labels = LabelSets::default(); + labels.attach(&erased); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + let chunk = |host: &str, start| Series { + metric: metric(&[(HOST, host)], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, 1)]).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + labels + .claim(&chunk("shared", 0)) + .expect("initial label fits"); + let key = label_set(&chunk("shared", 0).metric); + let before_runs = labels.seen[&key][0] + .iter() + .map(|run| (run.start, run.end)) + .collect::>(); + let before_reserved = pool.reserved(); + let blocker: MemoryReservation = MemoryConsumer::new("test blocker").register(&erased); + blocker + .try_grow(limit - before_reserved) + .expect("the blocker occupies remaining capacity"); + let error = labels + .claim(&chunk("shared", 2)) + .expect_err("growth exceeds the pool"); + assert!(!error.to_string().is_empty(), "pool refusal is reported"); + assert_eq!( + before_runs, + labels.seen[&key][0] + .iter() + .map(|run| (run.start, run.end)) + .collect::>(), + "the persistent entry is unchanged" + ); + assert_eq!(limit, pool.reserved(), "both reservations remain live"); + drop(blocker); + assert_eq!( + before_reserved, + pool.reserved(), + "the plan-owned entry remains charged" + ); +} + +/// Binding ordering includes field names, while schema ordering refuses unequal DF schemas. +#[test] +fn unequal_binding_and_output_schema_are_not_ordered_equal() { + let tree = rate_tree(); + let left = binding(&tree); + let other_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: "different".into(), + }]), + eval_grid: None, + }); + let right = binding_of(&tree, other_schema, left.tags.clone()); + assert_ne!(left, right); + assert_ne!(Some(std::cmp::Ordering::Equal), left.partial_cmp(&right)); + + let left_node = fused(tree.clone()); + let nullability_changed = Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, false), + Field::new(HOST, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])); + let right_node = FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding(&rate_tree()), + input_plan(frame_schema()), + nullability_changed.to_dfschema_ref().unwrap(), + ) + .unwrap(); + assert_ne!(left_node, right_node); + assert_ne!( + Some(std::cmp::Ordering::Equal), + left_node.partial_cmp(&right_node) + ); +} + +/// A chunk reaching back before the last stored run still merges into the history. +/// +/// The forward path only takes a chunk that starts at or after the end of the last stored run; a +/// chunk that starts before it can interleave with what is stored, so it has to go through the full +/// merge, and the result still has to be sorted, disjoint and with the runs that touch joined - +/// which is what the look-up in [`LabelSets::claim`] relies on. +/// +/// A chunk reaching back over a stored run is the other case, and it stays refused: `[0, 11)` shares +/// position 10 with the stored `[10, 12)`, so it is a collision and the stored history is left +/// as it was. +#[test] +fn a_chunk_reaching_back_still_merges() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + // One chunk holding `len` positions from `start`, every one of them with a value. + let chunk = |start: usize, len: usize| Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, len)]) + .expect("one run over the grid"), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; len]))), + }; + // The runs stored under the one label set, which one field of the output holds. + let stored = |labels: &LabelSets| -> Vec<(usize, usize)> { + labels + .seen + .values() + .next() + .expect("the chunk was recorded under one label set") + .first() + .expect("the series has one field") + .iter() + .map(|run| (run.start, run.end)) + .collect() + }; + + labels + .claim(&chunk(10, 2)) + .expect("the first chunk claims the positions it holds"); + assert_eq!( + vec![(10, 12)], + stored(&labels), + "the run `[10, 12)` is stored" + ); + + let error = labels + .claim(&chunk(0, 11)) + .expect_err("`[0, 11)` holds a value at position 10, where `[10, 12)` holds one"); + assert!( + error + .to_string() + .contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {error}" + ); + assert_eq!( + vec![(10, 12)], + stored(&labels), + "a refused claim left the stored history as it was" + ); + + let before = absorb_rebuilds(); + labels + .claim(&chunk(0, 10)) + .expect("`[0, 10)` reaches back before the last stored run without sharing a position"); + assert_eq!( + vec![(0, 12)], + stored(&labels), + "the run reaching back did not merge into the stored one and join it" + ); + assert_eq!( + 1, + absorb_rebuilds() - before, + "a chunk starting before the last stored run was appended rather than merged" + ); +} diff --git a/src/promql/src/grid_series.rs b/src/promql/src/grid_series.rs index c9d8bc19d1f..90cfa0f2e62 100644 --- a/src/promql/src/grid_series.rs +++ b/src/promql/src/grid_series.rs @@ -19,12 +19,12 @@ use std::collections::hash_map::Entry; use std::sync::Arc; use datafusion::arrow::array::builder::NullBufferBuilder; -use datafusion::arrow::array::{Float64Array, UInt32Array}; +use datafusion::arrow::array::{Array, Float64Array, UInt32Array}; use datafusion::error::{DataFusionError, Result}; use datatypes::series::MetricName; -use datatypes::series::histogram::HistogramBatchBuilder; +use datatypes::series::histogram::{HistogramBatch, HistogramBatchBuilder}; use datatypes::series::samples::{ - HistogramSamples, SampleRef, Samples, Series, SeriesFields, SeriesSchema, + DenseSamplesRef, HistogramSamples, SampleRef, Samples, Series, SeriesFields, SeriesSchema, }; use datatypes::series::time::{EvalGrid, GridSpan, TimeAxis}; @@ -35,6 +35,70 @@ pub(crate) mod reference; #[cfg(test)] mod tests; +pub(crate) fn series_bytes(series: &Series) -> usize { + let runs = series.time.runs().count(); + let segments: usize = series + .fields + .iter() + .map(|field| field.segments().len()) + .sum(); + std::mem::size_of::() + .saturating_add(std::mem::size_of::()) + .saturating_add(runs.saturating_mul(std::mem::size_of::<(usize, usize)>())) + .saturating_add( + segments.saturating_mul(std::mem::size_of::() + std::mem::size_of::()), + ) + .saturating_add(fields_bytes(&series.fields)) + .saturating_add(series.metric.name().map_or(0, str::len)) + .saturating_add(std::mem::size_of_val(series.metric.tags())) + .saturating_add( + series + .metric + .tags() + .iter() + .map(|tag| tag.name.len() + tag.value.len()) + .sum::(), + ) +} + +pub(crate) fn fields_bytes(fields: &SeriesFields) -> usize { + fields + .iter() + .map(|field| { + field.segments().iter().fold(0usize, |total, segment| { + total.saturating_add(segment_bytes(segment)) + }) + }) + .fold(0usize, usize::saturating_add) +} + +/// The bytes one dense segment of samples holds. +pub(crate) fn segment_bytes(segment: &Samples) -> usize { + match segment.dense() { + Some(DenseSamplesRef::Float(values)) => values.get_array_memory_size(), + Some(DenseSamplesRef::Histogram(histograms)) => histograms + .indices() + .get_array_memory_size() + .saturating_add(table_bytes(histograms.batch())), + Some(DenseSamplesRef::Mixed(mixed)) => mixed + .floats() + .get_array_memory_size() + .saturating_add(mixed.histograms().indices().get_array_memory_size()) + .saturating_add(table_bytes(mixed.histograms().batch())), + // `segments` hands out dense segments only, so a nested chunk never reaches here. + None => 0, + } +} + +/// The bytes one histogram table holds: the batch struct and the arrays its columns view. +pub(crate) fn table_bytes(table: &Arc) -> usize { + std::mem::size_of::().saturating_add(table.as_arrow().get_array_memory_size()) +} + +pub(crate) fn field_builder_bytes(width: usize) -> usize { + width.saturating_mul(std::mem::size_of::()) +} + /// Builds one grid series from rows pushed at absolute grid indices. /// /// A push names one position and the value each field holds there. Rows are accumulated @@ -114,6 +178,36 @@ impl GridSeriesBuilder { }) } + /// Reserves capacity for the exact number of valued rows in the next series. + /// + /// Call after `begin_series` and before its first push. Existing public callers may continue + /// to rely on normal growth; callers with a budget can reserve the concrete row capacity first. + pub(crate) fn reserve_rows(&mut self, rows: usize) -> Result<()> { + self.indices + .try_reserve_exact(rows.saturating_sub(self.indices.len())) + .map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid row indices: {error}")) + })?; + if rows == 0 { + return Ok(()); + } + if self.fields.is_empty() { + self.fields = (0..self.schema.fields.len()) + .map(|_| FieldBuilder::new(rows)) + .collect(); + } + for field in &mut self.fields { + field + .values + .try_reserve_exact(rows.saturating_sub(field.values.len())) + .map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid values: {error}")) + })?; + field.nulls = NullBufferBuilder::new(rows); + } + Ok(()) + } + /// Pushes the field values at absolute grid `index`; `None` is a missing value. /// /// Indices must be strictly increasing; skipped indices are positions that hold no value. @@ -176,7 +270,7 @@ impl GridSeriesBuilder { return Ok(false); } if self.fields.is_empty() { - self.fields = (0..len).map(|_| FieldBuilder::new()).collect(); + self.fields = (0..len).map(|_| FieldBuilder::new(0)).collect(); } self.indices.push(index); Ok(true) @@ -239,24 +333,28 @@ impl GridSeriesBuilder { } let fields = std::mem::take(&mut self.fields); let indices = std::mem::take(&mut self.indices); - let samples = fields - .into_iter() - .map(|mut field| { - let floats = Float64Array::new(field.values.into(), field.nulls.finish()); - let Some(column) = field.histograms else { - return Ok(Samples::from_float(floats)); - }; + let mut samples = Vec::new(); + samples.try_reserve_exact(fields.len()).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid field descriptions: {error}")) + })?; + for mut field in fields { + let floats = Float64Array::new(field.values.into(), field.nulls.finish()); + let sample = if let Some(column) = field.histograms { let batch = column.batch.finish().map_err(external)?; let histograms = HistogramSamples::try_new(UInt32Array::from(column.indices), batch) .map_err(external)?; - Samples::try_from_mixed(floats, histograms).map_err(external) - }) - .collect::>>()?; + Samples::try_from_mixed(floats, histograms).map_err(external)? + } else { + Samples::from_float(floats) + }; + samples.push(sample); + } let fields = SeriesFields::try_from_vec(samples).map_err(external)?; // The runs are all the axis has to remember about the gaps, and the rows are already // compact, so no position holding no value is visited or stored. - let spans = spans_of(&indices); + let run_count = count_runs(&indices); + let spans = spans_of(&indices, run_count)?; let time = TimeAxis::from_grid_spans(self.grid.clone(), spans).map_err(external)?; Series::try_new(&self.schema, metric, time, fields) .map(Some) @@ -265,8 +363,19 @@ impl GridSeriesBuilder { } /// Cuts the retained indices into maximal runs of consecutive positions. -fn spans_of(indices: &[usize]) -> Vec { +fn count_runs(indices: &[usize]) -> usize { + indices + .windows(2) + .filter(|pair| pair[1] != pair[0] + 1) + .count() + .saturating_add(usize::from(!indices.is_empty())) +} + +fn spans_of(indices: &[usize], runs: usize) -> Result> { let mut spans: Vec = Vec::new(); + spans.try_reserve_exact(runs).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid spans: {error}")) + })?; let mut start = 0; while let Some(&position) = indices.get(start) { let mut len = 1; @@ -276,14 +385,14 @@ fn spans_of(indices: &[usize]) -> Vec { spans.push(GridSpan::new(position, len)); start += len; } - spans + Ok(spans) } impl FieldBuilder { - fn new() -> Self { + fn new(rows: usize) -> Self { Self { values: Vec::new(), - nulls: NullBufferBuilder::new(0), + nulls: NullBufferBuilder::new(rows), histograms: None, } } @@ -454,7 +563,8 @@ pub fn merge_same_labelset(schema: &Arc, series: Vec) -> R // The runs are all the axis remembers about the holes, so a sparse merge stores one span // per run and no row for the positions between them. let time = - TimeAxis::from_grid_spans(grid.clone(), spans_of(&positions)).map_err(external)?; + TimeAxis::from_grid_spans(grid.clone(), spans_of(&positions, count_runs(&positions))?) + .map_err(external)?; let fields = SeriesFields::try_from_vec(fields).map_err(external)?; // The first fragment names the output, as it always has: a label set ignores empty values, // so the metric of the first member is the one the merged series keeps. diff --git a/src/promql/src/native/aggregate.rs b/src/promql/src/native/aggregate.rs new file mode 100644 index 00000000000..975dfb6c3c2 --- /dev/null +++ b/src/promql/src/native/aggregate.rs @@ -0,0 +1,243 @@ +// 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. + +//! The run-time state of an aggregate node. +//! +//! An aggregate folds a whole vector into one series per group, so its node is the one that cannot +//! answer as its input arrives: a group's value at a grid point is the fold of every value every +//! member of the group has there, and the last member to arrive is the one that decides it. What the +//! node keeps between chunks is therefore the fold's states and nothing else -- a chunk is folded as +//! it arrives and dropped -- and it produces once, when the input ends. +//! +//! That is also why this runtime cannot tell an incomplete group from a complete one: it folds what +//! it is handed. A [`AggregateStage::Complete`] aggregate is correct only if every contribution to +//! each group lies within the streams this execution reads, which is the physical plan's job -- the +//! caller that schedules scans and exchanges guarantees it, and a plan that ran a `Complete` +//! aggregate once per partition and then combined the answers would be combining values that were +//! each missing the other partitions' contributions. The stages that exist for that shape are +//! [`AggregateStage::Partial`] and [`AggregateStage::Final`], which this runtime does not execute +//! yet. + +use std::sync::Arc; + +use datafusion::error::{DataFusionError, Result}; +use datatypes::series::samples::{Series, SeriesSchema}; + +use crate::aggregate::{AggregateOp, FloatAggregator, Grouping}; +use crate::evaluated::EvaluatedChunk; +use crate::grid_series::series_bytes; +use crate::map::validate_float_series; +use crate::native::execute::{NodeBuffer, evaluated_bytes}; +use crate::native::metrics::NodePath; +use crate::native::plan::{ + AggregateKind, AggregateStage, BoundAggregate, Grouping as BoundGrouping, +}; + +/// One aggregate node's states, from the first chunk it is handed to the series it finishes with. +pub(crate) struct AggregateRuntime { + /// What the node computes, as the tree bound it: the kind it folds by and the stage it reports. + op: BoundAggregate, + /// The layout its input and output share: the fields every chunk has and the grid they are + /// addressed on. A `Complete` aggregate neither adds nor drops a field, and it evaluates on its + /// input's grid, so one schema describes both sides. + schema: Arc, + /// The fold. `None` once [`Self::compute`] has taken it, which is also what says the node has + /// nothing left to fold into. + accumulator: Option, + /// Whether the fold has been finished and its series handed over. A node computes once: a second + /// pull finds the series it queued rather than a fold that ran again. + /// + /// The driver reads this to decide whether the node settled: an aggregate answers only once its + /// fold has run, so a child with nothing left to hand over is not an answer. + pub(crate) computed: bool, +} + +impl AggregateRuntime { + /// What this node computes, for a rejection that has to say what it refused. + /// + /// It names the kind this node computes, which distinguishes it from another node in the tree. + fn describe(&self) -> String { + format!("{:?}", self.op.kind) + } + + /// The states of a `Complete` aggregate over `schema`. + /// + /// A stage other than `Complete` and a kind without a supported kernel are refused here, before a single + /// chunk is folded: a node this runtime cannot compute must not start, because a partial answer + /// would be one the caller never asked for. The rejection names what was found, so a reader can + /// tell a tree this runtime cannot run from one it cannot run correctly. + pub(crate) fn new( + op: BoundAggregate, + schema: Arc, + buffer: &mut NodeBuffer, + ) -> Result { + if op.stage != AggregateStage::Complete { + return Err(DataFusionError::NotImplemented(format!( + "the native runtime does not execute a `{:?}` aggregate yet: a `Partial` aggregate \ + reports the state a later stage merges and a `Final` aggregate merges that state, \ + and only `Complete` -- every contribution inside one execution -- is computed here", + op.stage + ))); + } + let kind = match op.kind { + AggregateKind::Sum => AggregateOp::Sum, + AggregateKind::Count => AggregateOp::Count, + AggregateKind::Avg => AggregateOp::Avg, + AggregateKind::Min => AggregateOp::Min, + AggregateKind::Max => AggregateOp::Max, + AggregateKind::Group => AggregateOp::Group, + AggregateKind::Stddev => AggregateOp::Stddev, + AggregateKind::Stdvar => AggregateOp::Stdvar, + other => { + return Err(DataFusionError::NotImplemented(format!( + "the native runtime does not execute the `{other:?}` aggregate yet: only the supported `Complete` aggregate kinds are computed here" + ))); + } + }; + let name_count = match &op.group { + BoundGrouping::By(labels) | BoundGrouping::Without(labels) => labels.len(), + }; + let name_bytes = name_count.saturating_mul(std::mem::size_of::<&str>()); + buffer.grow(name_bytes)?; + let names: Vec<&str> = match &op.group { + BoundGrouping::By(labels) | BoundGrouping::Without(labels) => { + labels.iter().map(|label| label.as_ref()).collect() + } + }; + let grouping = match &op.group { + BoundGrouping::By(_) => Grouping::By(&names), + BoundGrouping::Without(_) => Grouping::Without(&names), + }; + let name_bytes = names.len().saturating_mul(std::mem::size_of::<&str>()); + let initial = FloatAggregator::try_new_budgeted(&schema, kind, grouping, |bytes| { + buffer.resize(bytes.saturating_add(name_bytes)) + })?; + drop(names); + buffer.resize(initial.state_bytes())?; + Ok(Self { + op, + schema, + accumulator: Some(initial), + computed: false, + }) + } + + /// What folding `chunk` reserves while it is folded. + /// + /// The chunk is not retained: it is folded into the states and dropped at the end of the call, + /// so what this node holds afterwards is the fold and none of the input. The charge is therefore + /// the chunk's own bytes -- the figure [`evaluated_bytes`] reports, which is one pass over the + /// chunk and no state between chunks -- and not the `operand_chunk_bytes` figure a binary + /// operand keeps, which exists because such an operand holds a series for the whole execution + /// and so must count the slices of one buffer instead of the buffer once per chunk. + pub(crate) fn pending(&self, chunk: &EvaluatedChunk) -> usize { + evaluated_bytes(chunk) + } + + /// Folds one chunk into its group. + /// + /// A chunk that closes its series closes nothing here: a group is not a series, so the last + /// contribution of one member says nothing about the group. The fold keeps its states, which is + /// what makes a group's answer its whole group's answer once the input ends. + pub(crate) fn accumulate( + &mut self, + path: &NodePath, + chunk: EvaluatedChunk, + buffer: &mut NodeBuffer, + pending: usize, + ) -> Result<()> { + // `Series` keeps its fields public, so a chunk can be built by struct literal against + // another schema; the states are indexed by absolute position on this node's grid, so a + // chunk of another field count or another grid would be folded into the wrong places. + validate_float_series(&self.schema, &chunk.series) + .map_err(|error| folded_chunk(path, &self.describe(), error))?; + let accumulator = self.accumulator.as_mut().ok_or_else(|| no_fold(path))?; + // Keep the chunk covered while the fold decides whether it needs a new group or block. + accumulator.update_budgeted(&chunk.series, |fold_bytes| { + buffer.resize(pending.saturating_add(fold_bytes)) + })?; + drop(chunk); + buffer.resize(accumulator.state_bytes()) + } + + /// The number of retained groups, for budgeting the output queue backing storage. + pub(crate) fn groups(&self) -> usize { + self.accumulator.as_ref().map_or(0, FloatAggregator::groups) + } + + /// Finishes the fold and returns one series per group, or `None` if the fold already ran. + /// + /// `None` is what a node reports after it has queued its series: the caller reads them from the + /// queue, exactly as it reads the chunks a binary node queued, so a second pull does not finish + /// a fold whose states are gone. + pub(crate) fn compute( + &mut self, + path: &NodePath, + buffer: &mut NodeBuffer, + queue_bytes: usize, + ) -> Result>> { + if self.computed { + return Ok(None); + } + if self.accumulator.is_none() { + return Err(no_fold(path)); + } + self.computed = true; + let output = self + .accumulator + .take() + .ok_or_else(|| no_fold(path))? + .finish_budgeted(|total| buffer.resize(total.saturating_add(queue_bytes)))?; + let output_bytes = output.iter().fold(0usize, |bytes, series| { + bytes.saturating_add(series_bytes(series)) + }); + let output_vec_bytes = output + .capacity() + .saturating_mul(std::mem::size_of::()); + buffer.resize( + output_bytes + .saturating_add(output_vec_bytes) + .saturating_add(queue_bytes), + )?; + Ok(Some(output)) + } +} + +/// The rejection of a chunk that is not the layout this node folds. +/// +/// The fold indexes its states by absolute grid position and by field, so a chunk of another field +/// count or another grid would be folded into states that mean something else. `error` is what the +/// series check reported, which says which of the two disagreed. +/// Wraps the error of a chunk this node cannot fold. +/// +/// Only a chunk that does not fit the schema is the node's own refusal; anything else is the +/// kernel saying it does not compute this shape at all -- a native histogram in a float aggregate, +/// for instance -- and a caller that falls back to another planner on `NotImplemented` still can. +/// Wrapping that one in an `Execution` would deny it. +fn folded_chunk(path: &NodePath, aggregate: &str, error: DataFusionError) -> DataFusionError { + match error { + DataFusionError::NotImplemented(_) => error, + error => DataFusionError::Execution(format!( + "the `{aggregate}` aggregate at {path} folds chunks bound to its own schema, and this \ + one is not: {error}" + )), + } +} + +/// The error for a fold that has already been finished. +fn no_fold(path: &NodePath) -> DataFusionError { + DataFusionError::Execution(format!( + "the aggregate at {path} already finished its fold, so it takes no more input" + )) +} diff --git a/src/promql/src/native/display.rs b/src/promql/src/native/display.rs index 0191cdd9278..908b2179c53 100644 --- a/src/promql/src/native/display.rs +++ b/src/promql/src/native/display.rs @@ -127,6 +127,9 @@ fn describe(node: &NativeNode) -> String { BoundVectorOp::Histogram(HistogramFunction::Fraction { lower, upper }) => { format!("histogram_fraction {lower} {upper}") } + BoundVectorOp::CountValues { label, group } => { + format!("count_values {label:?} {}", grouping(group)) + } }, NativeExpr::Aggregate { op, .. } => format!( "{} {} [{}] {}", diff --git a/src/promql/src/native/display/tests.rs b/src/promql/src/native/display/tests.rs index a77a23fd0ed..c481f493455 100644 --- a/src/promql/src/native/display/tests.rs +++ b/src/promql/src/native/display/tests.rs @@ -446,6 +446,13 @@ fn a_vector_operation_names_its_binding() { }, r#" root Vector absent {job="x"}"#, ), + ( + BoundVectorOp::CountValues { + label: "value".into(), + group: Grouping::Without(Arc::new(["job".into()])), + }, + r#" root Vector count_values "value" without [job]"#, + ), ]; for (op, line) in lines { let tree = NativeTree::try_new(NativeNode { diff --git a/src/promql/src/native/execute.rs b/src/promql/src/native/execute.rs index 600a9b61d08..8813deff82a 100644 --- a/src/promql/src/native/execute.rs +++ b/src/promql/src/native/execute.rs @@ -91,9 +91,9 @@ use std::sync::Arc; use datafusion::error::{DataFusionError, Result}; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; -use datatypes::arrow::array::{Array, Float64Array}; +use datatypes::arrow::array::Float64Array; use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; -use datatypes::series::histogram::{HistogramBatch, HistogramView}; +use datatypes::series::histogram::HistogramView; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{ DenseSamplesRef, RowCursor, SampleRef, Samples, Series, SeriesFields, SeriesSchema, @@ -101,8 +101,10 @@ use datatypes::series::samples::{ use datatypes::series::time::{EvalGrid, GridRange, GridSpan, GridSpec, TimeAxis}; use datatypes::series::{MetricName, SeriesOrdering}; +use crate::aggregate; use crate::grid_series::{ - GridSample, GridSeriesBuilder, LabelSet, label_set, merge_same_labelset, owned, + GridSample, GridSeriesBuilder, LabelSet, fields_bytes, label_set, merge_same_labelset, owned, + series_bytes, table_bytes, }; use crate::histogram_eval::evaluate_histograms; use crate::instant_eval::select_instants; @@ -110,12 +112,16 @@ use crate::map::{ ArithmeticOp, Comparison, ComparisonOp, MathOp, ScalarOperand, drop_metric_name, map_arithmetic, map_compare, map_math, }; +use crate::native::aggregate::AggregateRuntime; use crate::native::batch::EvaluatedChunk; use crate::native::metrics::{NodeMetrics, NodePath}; use crate::native::plan::{ BoundInstant, BoundScalar, BoundVectorOp, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, }; +#[cfg(test)] +pub(crate) use crate::native::runs::absorb_rebuilds; +use crate::native::runs::{Run, absorb}; use crate::range_eval::RangeEvaluator; use crate::select::{self, SortOrder}; use crate::vector_match::{ @@ -189,34 +195,11 @@ fn raw_bytes(chunk: &RawSeriesChunk) -> usize { /// The bytes an evaluated chunk holds: its struct and the Arrow buffers behind its fields. /// /// The grid is shared by every chunk addressed on it, so it is not counted per chunk. -fn evaluated_bytes(chunk: &EvaluatedChunk) -> usize { - std::mem::size_of::().saturating_add(fields_bytes(&chunk.series.fields)) -} - -/// The bytes one evaluated series holds: its header, the axis that addresses it, the identity it -/// carries and the sample buffers behind its fields. /// -/// It is the figure a fragment of an affected window is charged with as it arrives. The axis's runs -/// and the segments of its fields are counted in one pass over the series, so the estimate costs -/// the fragment rather than a walk over everything the node holds. A payload a fragment shares with -/// another — a histogram table behind two slices of it, the grid behind every axis — is counted -/// once per fragment: that over-count is accepted here rather than a deduplication that would have -/// to keep state across fragments. -fn series_bytes(series: &Series) -> usize { - let runs = series.time.runs().count(); - let segments: usize = series - .fields - .iter() - .map(|field| field.segments().len()) - .sum(); - std::mem::size_of::() - .saturating_add(std::mem::size_of::()) - .saturating_add(runs.saturating_mul(std::mem::size_of::<(usize, usize)>())) - .saturating_add( - segments.saturating_mul(std::mem::size_of::() + std::mem::size_of::()), - ) - .saturating_add(fields_bytes(&series.fields)) - .saturating_add(metric_bytes(&series.metric)) +/// It is what an aggregate node charges for a chunk it is folding: that node hands the chunk over +/// to its fold and drops it, so the chunk alone is what it holds for as long as it holds anything. +pub(super) fn evaluated_bytes(chunk: &EvaluatedChunk) -> usize { + std::mem::size_of::().saturating_add(fields_bytes(&chunk.series.fields)) } /// The bytes one chunk a binary operand retains holds on its own. @@ -342,42 +325,6 @@ fn merge_workspace_bytes(rows: usize, fragments: usize, fields: usize) -> usize .saturating_add(output) } -/// The bytes the fields of a chunk hold: the sample buffers of every segment, and the histogram -/// tables those segments reference. -fn fields_bytes(fields: &SeriesFields) -> usize { - fields - .iter() - .map(|field| { - field.segments().iter().fold(0usize, |total, segment| { - total.saturating_add(segment_bytes(segment)) - }) - }) - .fold(0usize, usize::saturating_add) -} - -/// The bytes one dense segment of samples holds. -fn segment_bytes(segment: &Samples) -> usize { - match segment.dense() { - Some(DenseSamplesRef::Float(values)) => values.get_array_memory_size(), - Some(DenseSamplesRef::Histogram(histograms)) => histograms - .indices() - .get_array_memory_size() - .saturating_add(table_bytes(histograms.batch())), - Some(DenseSamplesRef::Mixed(mixed)) => mixed - .floats() - .get_array_memory_size() - .saturating_add(mixed.histograms().indices().get_array_memory_size()) - .saturating_add(table_bytes(mixed.histograms().batch())), - // `segments` hands out dense segments only, so a nested chunk never reaches here. - None => 0, - } -} - -/// The bytes one histogram table holds: the batch struct and the arrays its columns view. -fn table_bytes(table: &Arc) -> usize { - std::mem::size_of::().saturating_add(table.as_arrow().get_array_memory_size()) -} - /// The bytes the allocation `buffer` is a slice of. /// /// The capacity is the whole allocation, which is what a slice keeps alive: a key of a few bytes @@ -408,6 +355,8 @@ enum NodeTag { /// Matches two evaluated vectors and applies one operation to every matched pair of series, or /// broadcasts one constant over the chunks of one evaluated vector. Binary, + /// Folds every series of the input into one series per group, when the input ends. + Aggregate, /// Keeps some samples of its whole input vector at each grid point. Vector, } @@ -422,6 +371,7 @@ impl NodeTag { Self::Window => "window", Self::Map => "map", Self::Binary => "binary", + Self::Aggregate => "aggregate", Self::Vector => "vector operation", } } @@ -576,7 +526,9 @@ impl MergeInput { /// in the query; the metrics report the same figure so a slow node can be told apart from a wide /// one. The reservation is also the figure itself: a node holds nothing beside what it reserved, so /// a gauge reads its size rather than a copy that could drift from it. -struct NodeBuffer { +// Visible to the aggregate node so that its fold can reserve for what it is about to allocate; the +// reservation belongs to a node, not to a module. +pub(super) struct NodeBuffer { metrics: Arc, reservation: MemoryReservation, } @@ -592,7 +544,7 @@ impl NodeBuffer { } /// Reserves `bytes` beside what this node already holds. - fn grow(&mut self, bytes: usize) -> Result<()> { + pub(super) fn grow(&mut self, bytes: usize) -> Result<()> { self.reservation.try_grow(bytes)?; self.record(); Ok(()) @@ -603,7 +555,7 @@ impl NodeBuffer { /// What a node that consumes its input chunk by chunk holds afterwards is not the sum of what /// it consumed, so the reservation is set from the state that is left rather than grown per /// chunk. - fn resize(&mut self, bytes: usize) -> Result<()> { + pub(super) fn resize(&mut self, bytes: usize) -> Result<()> { self.reservation.try_resize(bytes)?; self.record(); Ok(()) @@ -628,20 +580,6 @@ impl NodeBuffer { } } -/// One run of grid positions, half-open, that hold a value in one field. -#[derive(Clone, Copy, Debug)] -struct Run { - start: usize, - end: usize, -} - -impl Run { - /// Whether two runs share a position. - fn overlaps(self, other: Self) -> bool { - self.start < other.end && other.start < self.end - } -} - /// The positions one node has handed out, per output label set and field. /// /// A node that rewrites labels — a `Map` dropping the metric name, a `Window` dropping the name @@ -876,109 +814,6 @@ fn misaligned_field( )) } -// How many times `absorb` rebuilt the stored runs instead of appending to them. -// -// Test-only, and per thread: the runs stored afterwards are the same either way, so a rebuild shows -// only in what it costs, and a test that reads its own thread's figure is not disturbed by another -// test's chunks. -#[cfg(test)] -thread_local! { - static ABSORB_REBUILDS: std::cell::Cell = const { std::cell::Cell::new(0) }; -} - -/// How many times this thread's [`absorb`] rebuilt the stored runs. -#[cfg(test)] -fn absorb_rebuilds() -> usize { - ABSORB_REBUILDS.with(std::cell::Cell::get) -} - -/// Merges `incoming` into `stored`, keeping the result in position order and joining the runs that -/// touch. -/// -/// Within a chunk one pass over both sides is what keeps this linear, since each side is sorted and -/// disjoint. Adding the runs one at a time and rebuilding the array after each would cost the square -/// of the runs a chunk yields: a chunk of one hundred thousand positions alternating between values -/// and nulls yields fifty thousand runs, and re-merging after each is over a billion element moves -/// for a single series. -/// -/// One pass per chunk is still a pass over everything the series has stored so far, so it is not -/// free once the *history* grows across chunks: a series split into `K` chunks of one position each, -/// a hole apart, stores one more run per chunk and moves `K(K+1)/2` of them, with no label collision -/// and no illegal input anywhere. A chunk that starts no earlier than the end of the last stored run -/// only ever adds to the end — the stored runs are sorted and disjoint, so every one of them ends -/// before anything this chunk holds — and it is appended instead, which is amortised constant per -/// run. That is the shape a scan hands over as its positions move forward chunk after chunk. A chunk -/// that starts before the last stored run can split a run or reach back past it, so it still costs -/// the full merge. -/// -/// Either way the result is sorted, disjoint, and has the runs that touch joined, which is what -/// [`Emitted::occupy`] looks a position up by. -fn absorb(stored: &mut Vec, incoming: &[Run]) { - if incoming.is_empty() { - return; - } - if stored - .last() - .is_none_or(|last| incoming[0].start >= last.end) - { - append_runs(stored, incoming); - } else { - merge_runs(stored, incoming); - } -} - -/// Appends `incoming` to `stored`, all of which lies past the last stored run. -/// -/// Nothing already stored has to be placed again: the two sides are each in position order and -/// nothing `incoming` holds can reach back into what is stored, so its runs are pushed on the end and -/// the vector's own growth is the only move — amortised over the runs appended. Where its first run -/// touches the end of the last stored run the two are one run, as they are after a merge, so it is -/// joined rather than left beside it. -fn append_runs(stored: &mut Vec, incoming: &[Run]) { - let mut incoming = incoming; - if let Some(last) = stored.last_mut() - && let Some(first) = incoming.first() - && first.start <= last.end - { - last.end = last.end.max(first.end); - incoming = &incoming[1..]; - } - stored.extend_from_slice(incoming); -} - -/// Merges `incoming` into `stored`, where the two can interleave. -/// -/// One pass is enough because each side is sorted and disjoint: whichever run starts first is -/// placed, and it joins the run before it where the two touch. -fn merge_runs(stored: &mut Vec, incoming: &[Run]) { - #[cfg(test)] - ABSORB_REBUILDS.with(|rebuilds| rebuilds.set(rebuilds.get() + 1)); - let mut merged: Vec = Vec::with_capacity(stored.len() + incoming.len()); - let (mut a, mut b) = (0, 0); - while a < stored.len() || b < incoming.len() { - let larger = match (stored.get(a), incoming.get(b)) { - (Some(left), Some(right)) if right.start < left.start => { - b += 1; - *right - } - (Some(left), _) => { - a += 1; - *left - } - (None, Some(right)) => { - b += 1; - *right - } - (None, None) => break, - }; - match merged.last_mut() { - Some(last) if larger.start <= last.end => last.end = last.end.max(larger.end), - _ => merged.push(larger), - } - } - *stored = merged; -} - /// What a map node applies. struct MapRuntime { ops: Arc<[MathOp]>, @@ -1915,6 +1750,19 @@ fn vector_series( BoundVectorOp::Histogram(function) => { evaluate_histograms(*function, schema, &merge_same_labelset(schema, series)?) } + BoundVectorOp::CountValues { label, group } => { + let (by, labels) = match group { + Grouping::By(labels) => (true, labels), + Grouping::Without(labels) => (false, labels), + }; + let labels: Vec<&str> = labels.iter().map(|label| &**label).collect(); + let grouping = if by { + aggregate::Grouping::By(&labels) + } else { + aggregate::Grouping::Without(&labels) + }; + aggregate::count_values(schema, &series, label, grouping) + } BoundVectorOp::Absent { labels } => Ok(select::absent(schema, &series, labels.clone())? .into_iter() .collect()), @@ -1933,6 +1781,10 @@ fn vector_series( /// - A sort: the value each series is keyed by, and the floats `sort` and `sort_desc` rebuild a /// series into when it holds histograms; the label sorts move the series they were given. /// - `absent`: a presence flag and a result value per grid point and field. +/// - `count_values`, per field sample, since every field's distinct values make groups of their +/// own and a gauge with distinct values makes about one group per sample: the group's label set +/// twice (its metric and its key) with the value written into it, its count per field, and the +/// result series it becomes. /// - `histogram_quantile` and `histogram_fraction`: the merge of label-set pieces before the kernel /// ([`merge_workspace_bytes`]); per series, its group entry and label pairs and a cursor per /// field; per row, the bucket it contributes and a result value per field. @@ -1979,6 +1831,21 @@ fn vector_workspace_bytes(op: &BoundVectorOp, schema: &SeriesSchema, series: &[S }, ); } + BoundVectorOp::CountValues { .. } => { + let label = std::mem::size_of::<(Arc, Arc)>(); + // The value label at the length of a 17-digit value with sign, point and exponent; + // the `'f'` form writes a large magnitude out in full, which this does not follow. + let value = 24; + let group = std::mem::size_of::>() + + fields * std::mem::size_of::() + + std::mem::size_of::() + + fields * std::mem::size_of::(); + return series.iter().fold(0usize, |total, item| { + let labels = 2 * (item.metric.tags().len() + 2) * label + value; + let samples = item.time.len().saturating_mul(fields); + total.saturating_add(samples.saturating_mul(group + labels)) + }); + } BoundVectorOp::TopK { .. } | BoundVectorOp::LimitK { .. } | BoundVectorOp::LimitRatio { .. } => {} @@ -2260,6 +2127,14 @@ struct NodeRuntime { /// `None` for every other node, and for a binary node that broadcasts a constant instead: that /// one holds no operand at all and is described by [`NodeRuntime::broadcast`]. binary: Option, + /// What an aggregate node holds: the fold over what it was handed. + /// + /// Unlike a binary node's operands, nothing of the input survives a chunk: the fold keeps its + /// states and the chunk it read is dropped, so this is the one node whose reservation does not + /// follow what it was handed but what it accumulated. That is also why this node cannot report + /// its input complete from what it holds -- there is nothing to hold -- and reads the driver's + /// [`NativeExecutor::inputs_finished`] and its child's [`NativeExecutionState::settled`] instead. + aggregate: Option, /// What a vector node holds. vector: Option, /// What a binary node holds when one operand is a constant: the operation and the layout of the @@ -2669,6 +2544,102 @@ impl NodeRuntime { self.buffer.resize(queued) } + /// Folds one chunk into this node's aggregate. + /// + /// The chunk is reserved for before it is folded, by what the node reports it is worth, and the + /// reservation settles afterwards on the fold that is left. Unlike [`Self::retain_binary`], what + /// stops being held here is the chunk itself: the fold keeps its states, so a chunk that closes + /// its series leaves nothing behind and the node is empty of input again. + fn retain_aggregate(&mut self, path: &NodePath, chunk: NativeChunk) -> Result<()> { + let NativeChunk::Evaluated(chunk) = chunk else { + return Err(DataFusionError::Execution(format!( + "the aggregate at {path} consumes evaluated chunks, not raw samples" + ))); + }; + // The figure is read before the chunk is folded, because reserving after the fact would be + // reserving for memory the fold already read. + let pending = self + .aggregate + .as_ref() + .ok_or_else(|| no_aggregate(path))? + .pending(&chunk); + let positions = chunk.series.time.len(); + self.buffer.grow(pending)?; + let metrics = Arc::clone(&self.buffer.metrics); + { + let _timer = metrics.timer(); + let (aggregate, buffer) = (&mut self.aggregate, &mut self.buffer); + aggregate + .as_mut() + .ok_or_else(|| no_aggregate(path))? + .accumulate(path, chunk, buffer, pending)?; + } + metrics.record_input(1, positions); + // A chunk that closed its series closed it for the node below: a group is not a series, so + // the fold completing nothing on it is why nothing is counted here. The series this node + // closes are the ones it hands out, counted when it queues them. + Ok(()) + } + + /// Finishes this node's fold and queues one chunk per group it produced. + /// + /// Each group is one series, decided in full by the fold, so the one chunk it is handed over as + /// closes it: a consumer above reads "nothing more here" from that chunk rather than from another + /// one arriving. + fn compute_aggregate(&mut self, path: &NodePath) -> Result<()> { + let metrics = Arc::clone(&self.buffer.metrics); + let _timer = metrics.timer(); + let groups = self + .aggregate + .as_ref() + .map_or(0, |aggregate| aggregate.groups()); + let queue_capacity = self.queued.capacity(); + let queue_target = self.queued.len().saturating_add(groups); + let queue_growth = queue_target + .saturating_sub(queue_capacity) + .saturating_mul(std::mem::size_of::<(NativeChunk, usize)>()); + self.buffer.grow(queue_growth)?; + if groups > 0 { + self.queued.try_reserve_exact(groups).map_err(|error| { + DataFusionError::Execution(format!( + "cannot reserve aggregate output queue: {error}" + )) + })?; + } + let queue_bytes = queue_target + .max(queue_capacity) + .saturating_mul(std::mem::size_of::<(NativeChunk, usize)>()); + let output = { + let aggregate = self.aggregate.as_mut().ok_or_else(|| no_aggregate(path))?; + aggregate.compute(path, &mut self.buffer, queue_bytes)? + }; + // The fold has returned its series; reserve their handover before appending them to the queue. + let Some(output) = output else { + return Ok(()); + }; + let mut closed = 0usize; + for series in output { + let bytes = series_bytes(&series); + let chunk = NativeChunk::Evaluated(EvaluatedChunk::try_new(series, true)?); + self.queued.push_back((chunk, bytes)); + closed += 1; + } + // A group is one series and this node closes each of them exactly once, on the one chunk it + // hands it over as. + self.buffer.metrics.record_completed_series(closed); + // The output Vec has been consumed; the node now holds only queue backing and chunks. + let queued_bytes = self + .queued + .iter() + .fold(0usize, |total, (_, bytes)| total.saturating_add(*bytes)); + let queue_backing = self + .queued + .capacity() + .saturating_mul(std::mem::size_of::<(NativeChunk, usize)>()); + self.buffer + .resize(queued_bytes.saturating_add(queue_backing)) + } + /// Feeds one chunk to this node's instant selection. /// /// `None` means the series this chunk belongs to has not ended: the selection reads one whole @@ -2871,6 +2842,8 @@ impl NodeRuntime { self.window = None; self.instant = None; self.binary = None; + // The fold is what an aggregate node holds, and dropping it is what gives its states back. + self.aggregate = None; self.vector = None; self.broadcast = None; self.source_schema = None; @@ -2963,15 +2936,14 @@ impl NativeExecutionState { // two vectors, or a constant broadcast over one -- is decided where the operands' // layouts are read. NativeExpr::Binary { .. } => NodeTag::Binary, + // An aggregate is built here and narrowed below, for the same reason: `validate` + // admits every kind and every stage as a legal tree, and the runtime computes one + // kind at one stage. + NativeExpr::Aggregate { .. } => NodeTag::Aggregate, NativeExpr::Vector { .. } => NodeTag::Vector, // A scalar is a value node: it produces the constant it was bound to, hands over no // chunk and holds no state, so it is built like any other node and pulled like none. NativeExpr::Scalar(_) => NodeTag::Scalar, - // An aggregate is legal in a bound tree, and `validate` stays the authority on that; - // this runtime simply does not execute it yet. Rejected before a single chunk is - // read: a tree the runtime cannot execute must not start, because a partial answer - // would be one the caller never asked for. - NativeExpr::Aggregate { .. } => return Err(not_implemented("Aggregate", &path)), }; let buffer = NodeBuffer::new(&path, handles, &pool); let mut runtime = NodeRuntime { @@ -2982,6 +2954,7 @@ impl NativeExecutionState { window: None, map: None, binary: None, + aggregate: None, vector: None, broadcast: None, source_schema: None, @@ -3081,6 +3054,29 @@ impl NativeExecutionState { // out. A broadcast that keeps the labels it reads is not where two series // become one either, so it keeps none; the one that drops them does, above. } + NativeExpr::Aggregate { input, op } => { + // A `Complete` aggregate neither adds nor drops a field and evaluates on its + // input's grid, so the node's own output is the layout it folds. What it does + // not execute -- another kind, or another stage -- is refused here, before the + // pass reaches the input that would feed it, and naming what it found is that + // rejection's job: a `Partial` aggregate's output is state and names no samples + // at all, so the input's layout stands in for it and the stage decides. + let schema = node + .output + .samples() + .or_else(|| input.output.samples()) + .cloned() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "the aggregate at {path} does not produce samples" + )) + })?; + runtime.aggregate = Some(AggregateRuntime::new( + op.clone(), + schema, + &mut runtime.buffer, + )?); + } NativeExpr::Vector { op, .. } => { let schema = node.output.samples().cloned().ok_or_else(|| { DataFusionError::Execution(format!( @@ -3103,10 +3099,9 @@ impl NativeExecutionState { // A vector operation keeps the labels of the series it reads, or answers under // label sets of its own, so it keeps no ledger of what it handed out. } - NativeExpr::Scalar(_) | NativeExpr::Aggregate { .. } => { + NativeExpr::Scalar(_) => { // A scalar holds nothing: what its consumer reads is the constant the node was - // bound to, which is on the tree rather than in this state. The aggregate is - // rejected above, where the pass has the position to name. + // bound to, which is on the tree rather than in this state. } } // The inputs this node reads, in the order the tree declares them: a unary node's one @@ -3114,12 +3109,11 @@ impl NativeExecutionState { // pushed back to front, so the node the spec names is the one the pass reaches next and // nothing but its own subtree lies between the two. let reads_input = match &node.expr { - NativeExpr::Input(_) | NativeExpr::Scalar(_) | NativeExpr::Aggregate { .. } => { - false - } + NativeExpr::Input(_) | NativeExpr::Scalar(_) => false, NativeExpr::Instant { input, .. } | NativeExpr::Window { input, .. } | NativeExpr::Map { input, .. } + | NativeExpr::Aggregate { input, .. } | NativeExpr::Vector { input, .. } => { pending.push((input.as_ref(), path.input())); true @@ -3189,6 +3183,14 @@ impl NativeExecutionState { holds_shared[index] = input_holds_shared; } NodeTag::Map => holds_shared[index] = true, + NodeTag::Aggregate => { + // An aggregate hands over one series per group, each under the label set its + // group was keyed by and each closed on its own chunk: two groups are two + // label sets, so nothing above has to join pieces of a series. Its own input + // is read whole -- the fold holds what it accumulated, not the chunks -- so + // nothing about its input's flag reaches it either. + holds_shared[index] = false; + } NodeTag::Scalar => { // A scalar produces one value and no labels, so it is not a stream of series // and nothing above it can read two of them under one label set. The flag @@ -3229,8 +3231,9 @@ impl NativeExecutionState { // what ranking the merged series would, and what is above treats them as // it treats a map's. `absent` answers with one series under the labels it // was bound to, the histogram functions merge the label sets they answer - // with, and `limitk` counts series in input order, so it merges the pieces - // first; none of them hands over pieces. + // with, `count_values` answers with one series per label set it makes, and + // `limitk` counts series in input order, so it merges the pieces first; + // none of them hands over pieces. let op = match expr { NativeExpr::Vector { op, .. } => Some(op), _ => None, @@ -3245,7 +3248,11 @@ impl NativeExecutionState { let whole = limitk || matches!( op, - Some(BoundVectorOp::Absent { .. } | BoundVectorOp::Histogram(_)) + Some( + BoundVectorOp::Absent { .. } + | BoundVectorOp::Histogram(_) + | BoundVectorOp::CountValues { .. } + ) ); holds_shared[index] = input_holds_shared && !whole; } @@ -3742,9 +3749,24 @@ impl NativeExecutor { // A scalar produces its value without reading anything, so it is settled from the // start: there is no chunk it owes its consumer. NodeTag::Scalar => true, + // These report as they are handed chunks, so a child with nothing left is all it takes + // for them to have nothing left either. NodeTag::Map | NodeTag::Window | NodeTag::Instant => self.state.nodes[index] .input .is_none_or(|child| self.settled(child)), + // An aggregate answers one series per group and only once its input ended, so its own + // fold has to have run as well: a child that has handed everything up is not enough, + // because every one of those chunks was folded and none of them answered yet. This is + // what keeps it out of the line above, which asks only about the child. + NodeTag::Aggregate => { + runtime + .aggregate + .as_ref() + .is_some_and(|aggregate| aggregate.computed) + && self.state.nodes[index] + .input + .is_none_or(|child| self.settled(child)) + } NodeTag::Vector => { runtime .vector @@ -3942,6 +3964,7 @@ impl NativeExecutor { // does, so its values are counted rather than read as its positions. runtime.pop_queued(evaluated_values) } + NodeTag::Aggregate => self.aggregate_node(index), NodeTag::Vector => { let child = input.ok_or_else(|| missing_input(tag))?; // Which samples are kept at a point is not decided until the whole input has @@ -3961,6 +3984,35 @@ impl NativeExecutor { } } + /// Folds everything the aggregate at `index` can read and hands over one group's series. + /// + /// What it reads is taken to the end of what its child has, because nothing the fold is handed + /// is retained: a chunk is folded as it arrives and dropped, so reading the child dry costs the + /// fold's states and not the input's bytes. Nothing is answered until the input has ended -- + /// before that, a group's value at a point is still missing the members that have not arrived. + /// + /// That is the difference from a binary node. A binary node retains both operands and so can ask + /// whether they hold anything still open; an aggregate keeps none of them and asks the driver + /// instead: [`Self::inputs_finished`] says the caller will hand over nothing more, and + /// [`NativeExecutionState::settled`] says the child has none left to produce. + fn aggregate_node(&mut self, index: usize) -> Result> { + let child = self.state.nodes[index] + .input + .ok_or_else(|| missing_input(NodeTag::Aggregate))?; + while let Some(chunk) = self.pull(child)? { + let path = &self.state.paths[index]; + self.state.runtimes[index].retain_aggregate(path, chunk)?; + } + if !self.inputs_finished || !self.settled(child) { + return Ok(None); + } + let path = &self.state.paths[index]; + self.state.runtimes[index].compute_aggregate(path)?; + // A group's series holds one value per field and point it decided, so its values are counted + // rather than read as its positions, as a binary result's are. + self.state.runtimes[index].pop_queued(evaluated_values) + } + /// Pulls the vector side of the broadcast at `index` until it hands over a chunk, and returns /// that chunk transformed. /// @@ -4188,11 +4240,7 @@ fn valid_values(chunk: &EvaluatedChunk) -> usize { .sum() } -/// The values a chunk a binary node hands over holds. -/// -/// A binary result is an evaluated chunk, so this is [`valid_values`]: one value per field sample -/// that holds one. A raw chunk cannot reach a binary node's queue, and is answered as its positions -/// so an impossible queue cannot under-count. +/// The values an evaluated chunk hands over, or the timestamps a raw source hands over. fn evaluated_values(chunk: &NativeChunk) -> usize { match chunk { NativeChunk::Evaluated(chunk) => valid_values(chunk), @@ -4200,18 +4248,6 @@ fn evaluated_values(chunk: &NativeChunk) -> usize { } } -/// The rejection of a node kind this runtime does not execute yet. -/// -/// It names the kind and the node's position: the position is what the metrics and `EXPLAIN` call -/// the same node, so a reader can find it in the tree. An aggregate is the one kind left here: a -/// scalar executes as a value node, and a binary node executes both as a vector match and as a -/// constant broadcast. -fn not_implemented(kind: &str, path: &NodePath) -> DataFusionError { - DataFusionError::NotImplemented(format!( - "the native runtime does not execute {kind} nodes yet: {kind} at {path}" - )) -} - /// The error for a window whose evaluator this execution already closed. fn closed(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!("the window at {path} is already closed")) @@ -4293,11 +4329,27 @@ fn truncated_input(path: &NodePath) -> DataFusionError { )) } +/// The rejection of a node kind this runtime does not execute yet. +/// +/// It names the kind and the node's position: the position is what the metrics and `EXPLAIN` call +/// the same node, so a reader can find it in the tree. +fn not_implemented(kind: &str, path: &NodePath) -> DataFusionError { + DataFusionError::NotImplemented(format!( + "the native runtime does not execute {kind} nodes yet: {kind} at {path}" + )) +} + /// The error for a node asked to compute without the runtime its kind holds. fn no_binary(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!("the node at {path} holds no binary operation")) } +/// The error for an aggregate node whose fold is missing. +fn no_aggregate(path: &NodePath) -> DataFusionError { + DataFusionError::Execution(format!("the node at {path} holds no aggregate")) +} + +/// The error for a node asked to run a vector operation without one. fn no_vector(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!("the node at {path} holds no vector operation")) } diff --git a/src/promql/src/native/execute/tests.rs b/src/promql/src/native/execute/tests.rs index ad83197cb3c..2e1bc5594f5 100644 --- a/src/promql/src/native/execute/tests.rs +++ b/src/promql/src/native/execute/tests.rs @@ -36,7 +36,7 @@ use datatypes::series::{Label, MetricName, SeriesOrdering}; use super::{ BinaryRuntime, Emitted, NativeChunk, NativeExecutionState, NativeExecutor, NodeBuffer, NodeRuntime, absorb_rebuilds, evaluated_bytes, grid_end_ms, map_chunk, operand_chunk_bytes, - raw_bytes, + raw_bytes, vector_workspace_bytes, }; use crate::evaluated::EvaluatedChunk; use crate::grid_series::merge_same_labelset; @@ -45,8 +45,9 @@ use crate::instant_eval::select_instants; use crate::map::MathOp; use crate::native::metrics::{NodeMetric, NodeMetrics, NodePath}; use crate::native::plan::{ - AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, + AggregateKind, AggregateStage, AggregateStateSchema, BoundAggregate, BoundInput, BoundInstant, + BoundMap, BoundScalar, BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, + NativeTree, NativeType, }; use crate::range_eval::{RangeEvaluator, RangeFunction, RangeWindow}; use crate::select::SortOrder; @@ -487,6 +488,7 @@ fn source_runtime(pool: &Arc) -> NodeRuntime { window: None, map: None, binary: None, + aggregate: None, broadcast: None, vector: None, source_schema: None, @@ -509,6 +511,7 @@ fn binary_runtime(pool: &Arc, schema: &Arc) -> Nod Arc::new(binary_matching(BinaryOp::Add)), [Arc::clone(schema), Arc::clone(schema)], )), + aggregate: None, broadcast: None, vector: None, source_schema: None, @@ -1088,30 +1091,17 @@ fn a_downstream_window_emits_before_its_input_series_ends() { } } -/// Every node kind the runtime does not execute is rejected, naming the kind and its position. +/// Every binary shape the runtime does not execute is rejected, naming the kind and its position. /// /// A scalar is not one of them any more: it is a value node, and a binary node over a constant is -/// executed as a broadcast. The binary shapes that stay refused are the ones no kernel here takes, -/// and the aggregate is still a kind this runtime does not run at all. +/// executed as a broadcast. The binary shapes that stay refused are the ones no kernel here takes; +/// the aggregate kinds and stages this runtime does not execute are refused by +/// `an_aggregate_kind_or_stage_the_runtime_does_not_compute_is_named`. #[test] fn an_unimplemented_node_is_reported_not_panicked() { // One grid shared by the nodes that validate by allocation identity. let statement = || grid(0, 120_000, 2); - let aggregate_tree = { - let out = statement(); - NativeNode { - expr: NativeExpr::Aggregate { - input: Box::new(window(input(), Arc::clone(&out))), - op: BoundAggregate { - kind: AggregateKind::Sum, - group: Grouping::By(Arc::new(["job".into()])), - stage: AggregateStage::Complete, - }, - }, - output: NativeType::Samples(evaluated_schema(out)), - } - }; // Two constants make a scalar, which is a legal node and a shape no binary here broadcasts // over: there is no vector to pair either of them with. let two_scalars_tree = { @@ -1143,12 +1133,6 @@ fn an_unimplemented_node_is_reported_not_panicked() { }; let cases = [ - ( - "Aggregate", - "root", - "does not execute Aggregate", - aggregate_tree, - ), ("Binary", "root", "two constants", two_scalars_tree), ("Binary", "root", "matching modifiers", matched_scalar_tree), ]; @@ -6563,6 +6547,158 @@ fn a_constant_on_the_right_multiplies_every_value() { assert_eq!(execution.state().reserved_bytes(), 0); } +/// A `Complete` aggregate over `inner`, carrying `inner`'s fields and evaluated on its grid. +/// +/// A `Complete` aggregate neither adds nor drops a field and lands on its input's grid, so the node's +/// own output is the layout its input produces. +fn aggregate(inner: NativeNode, op: BoundAggregate) -> NativeNode { + let fields = inner + .output + .samples() + .expect("an aggregate reads samples") + .fields + .clone(); + let grid = inner + .output + .grid() + .expect("a complete aggregate reads an evaluated vector") + .clone(); + let output = NativeType::Samples(Arc::new(SeriesSchema { + fields, + eval_grid: Some(grid), + })); + NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(inner), + op, + }, + output, + } +} + +/// A `Complete` `sum` grouped as `group`. +fn sum_grouping(group: Grouping) -> BoundAggregate { + BoundAggregate { + kind: AggregateKind::Sum, + group, + stage: AggregateStage::Complete, + } +} + +/// `name{job=...,instance=...}`, the two labels these tests group by. +fn job_instance_metric(name: &str, job: &str, instance: &str) -> MetricName { + MetricName::try_new( + Some(name.into()), + vec![ + Label { + name: "job".into(), + value: job.into(), + }, + Label { + name: "instance".into(), + value: instance.into(), + }, + ], + ) + .unwrap() +} + +/// What one aggregate result reports: its name, its tags and the value at each point it decided. +type GroupResult = ( + Option, + Vec<(String, String)>, + Vec<(usize, Option)>, +); + +/// The groups a run produced, each with the absolute positions it decided, sorted by identity. +/// +/// Sorting is the test's, not the runtime's: a group's place in the answer is the order its first +/// input appeared, which the tests that check it assert on the stream rather than here. +fn groups_of(chunks: &[EvaluatedChunk]) -> Vec { + let mut results: Vec = chunks + .iter() + .map(|chunk| { + let mut tags: Vec<(String, String)> = chunk + .series + .metric + .tags() + .iter() + .map(|tag| (tag.name.to_string(), tag.value.to_string())) + .collect(); + tags.sort(); + ( + chunk.series.metric.name().map(str::to_string), + tags, + decided(std::slice::from_ref(chunk)), + ) + }) + .collect(); + results.sort_by(|left, right| (&left.0, &left.1).cmp(&(&right.0, &right.1))); + results +} + +/// One group's answer under the name and tags `tags` give it, ordered by label name. +fn grouped(tags: &[(&str, &str)], points: Vec<(usize, Option)>) -> GroupResult { + let mut tags: Vec<(String, String)> = tags + .iter() + .map(|(name, value)| ((*name).to_string(), (*value).to_string())) + .collect(); + tags.sort(); + (None, tags, points) +} + +/// One group's answer when every label of its input was dropped. +fn unlabelled(points: Vec<(usize, Option)>) -> GroupResult { + grouped(&[], points) +} + +/// `sum` with no grouping drops every label, name and tags alike, and adds the members per point. +#[test] +fn an_aggregate_drops_every_label_and_adds_its_group_per_point() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let set = ExecutionPlanMetricsSet::new(); + let mut execution = executor(&tree, pool(), &set); + for (metric, values) in [ + ( + job_instance_metric("requests", "api", "a"), + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + ), + ( + job_instance_metric("requests", "db", "b"), + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + ), + ] { + execution + .push_evaluated(evaluated_chunk_of(&schema, metric, 0, values, true)) + .unwrap(); + } + + let output = execution.finish().unwrap(); + assert_eq!( + groups_of(&output), + vec![unlabelled(vec![ + (0, Some(11.0)), + (1, Some(22.0)), + (2, Some(33.0)), + (3, Some(44.0)), + ])], + "the two series were not folded into one unlabelled group" + ); + // One group is one series, so it closes exactly once and its chunk says so. + assert!(output.len() == 1 && output[0].is_last); + assert_eq!( + count(&set, &NodePath::root(), NodeMetric::COMPLETED_SERIES), + 1 + ); + assert_eq!(execution.state().reserved_bytes(), 0); +} + /// The side a constant is on decides which operand the vector's value is. /// /// Division is not commutative, so `x / 100` and `100 / x` answering different values is what says @@ -7036,6 +7172,208 @@ fn a_window_above_a_broadcast_merges_the_series_it_collapses() { assert_eq!(execution.state().reserved_bytes(), 0); } +/// Every supported `Complete` kind reaches its matching float kernel through the native runtime. +#[test] +fn complete_aggregate_kinds_fold_two_values() { + let cases = [ + (AggregateKind::Sum, 4.0), + (AggregateKind::Count, 2.0), + (AggregateKind::Avg, 2.0), + (AggregateKind::Min, 1.0), + (AggregateKind::Max, 3.0), + (AggregateKind::Group, 1.0), + (AggregateKind::Stddev, 1.0), + (AggregateKind::Stdvar, 1.0), + ]; + + for (kind, expected) in cases { + let grid = grid(0, 60_000, 1); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + BoundAggregate { + kind, + group: Grouping::By(Arc::new([])), + stage: AggregateStage::Complete, + }, + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + for value in [1.0, 3.0] { + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(value)], + true, + )) + .unwrap(); + } + + let output = execution.finish().unwrap(); + assert_eq!( + groups_of(&output), + vec![unlabelled(vec![(0, Some(expected))])], + "{kind:?} did not compute its native complete result" + ); + } +} + +/// The same series, delivered in two chunks, fold to what the whole series folds to. +#[test] +fn an_aggregate_folds_a_series_split_across_chunks_like_a_whole_one() { + let bound = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&bound)); + // A fresh tree per execution: the node's state is the execution's, not the tree's. + let tree = || { + NativeTree::try_new(aggregate( + evaluated_input(grid(0, 60_000, 4)), + sum_grouping(Grouping::By(Arc::new(["job".into()]))), + )) + .unwrap() + }; + let (first, second) = ( + job_instance_metric("requests", "api", "a"), + job_instance_metric("requests", "api", "b"), + ); + + // One chunk per series. + let mut whole = executor(&tree(), pool(), &ExecutionPlanMetricsSet::new()); + for (metric, values) in [ + ( + first.clone(), + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + ), + ( + second.clone(), + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + ), + ] { + whole + .push_evaluated(evaluated_chunk_of(&schema, metric, 0, values, true)) + .unwrap(); + } + let whole = whole.finish().unwrap(); + + // The first series split in two, the second whole: the group is the same group, folded from + // positions the chunks name absolutely. + let mut split = executor(&tree(), pool(), &ExecutionPlanMetricsSet::new()); + split + .push_evaluated(evaluated_chunk_of( + &schema, + first.clone(), + 0, + vec![Some(1.0), Some(2.0)], + false, + )) + .unwrap(); + assert!( + split.next_chunk().unwrap().is_none(), + "a chunk that does not close its series answered for the group" + ); + split + .push_evaluated(evaluated_chunk_of( + &schema, + first, + 2, + vec![Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + + split + .push_evaluated(evaluated_chunk_of( + &schema, + second, + 0, + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + true, + )) + .unwrap(); + + let split = split.finish().unwrap(); + assert_eq!(groups_of(&split), groups_of(&whole)); + assert_eq!( + groups_of(&split), + vec![grouped( + &[("job", "api")], + vec![ + (0, Some(11.0)), + (1, Some(22.0)), + (2, Some(33.0)), + (3, Some(44.0)), + ], + )], + "the split series was folded as two groups or at the wrong positions" + ); +} + +/// A point no member of a group has a value at is not a point of the group's result. +/// +/// Sparse input is what says the fold indexes its states by absolute position: a chunk holding +/// positions 0 and 1 and another holding 1 and 2 leave position 3 to nothing, and the hole between +/// the rows is never filled with a value. +#[test] +fn an_aggregate_reports_only_the_points_a_group_holds_a_value_at() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + // Positions 0 and 1 of the first series, 1 and 2 of the second, and nothing at 3. + execution + .push_evaluated(sparse_chunk( + &schema, + job_instance_metric("requests", "api", "a"), + &[(0, 2)], + vec![1.0, 2.0], + true, + )) + .unwrap(); + execution + .push_evaluated(sparse_chunk( + &schema, + job_instance_metric("requests", "db", "b"), + &[(1, 2)], + vec![20.0, 300.0], + true, + )) + .unwrap(); + + let output = execution.finish().unwrap(); + assert_eq!( + groups_of(&output), + vec![ + // Position 1 is the only one both members hold, and it is their sum. + unlabelled(vec![(0, Some(1.0)), (1, Some(22.0)), (2, Some(300.0))]), + ], + "the group was answered at a point no member holds, or summed the wrong members" + ); +} + +/// An aggregate handed no series produces nothing at all: no group, so no empty series either. +#[test] +fn an_aggregate_over_no_series_produces_nothing() { + let grid = grid(0, 60_000, 4); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new(["job".into()]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + let output = execution.finish().unwrap(); + assert!( + output.is_empty(), + "an aggregate with no input answered {} series", + output.len() + ); + assert_eq!(execution.state().reserved_bytes(), 0); +} + /// A broadcast hands one chunk over for every chunk it reads, and closes the series once. /// /// The metrics are what says the chunk-per-chunk behaviour held: one input batch per chunk read, one @@ -7251,6 +7589,411 @@ fn a_broadcast_refuses_two_series_that_share_a_label_set_and_releases_the_reserv ); } +/// A series that ended is one member of a group, not the group: nothing is answered until the input +/// ends. +/// +/// The observable state between the two is what this pins: the fold has read the closing chunk of +/// the first series and answered nothing, and `computed` still says so. +#[test] +fn an_aggregate_answers_nothing_until_its_whole_input_ends() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + let computed = |execution: &NativeExecutor| { + execution.state().runtimes[0] + .aggregate + .as_ref() + .expect("the root holds an aggregate") + .computed + }; + + execution + .push_evaluated(evaluated_chunk_of( + &schema, + job_instance_metric("requests", "api", "a"), + 0, + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + assert!(execution.next_chunk().unwrap().is_none()); + assert!(!computed(&execution), "the fold finished on one series"); + assert!( + execution.next_chunk().unwrap().is_none(), + "the closing chunk of one series answered for its group" + ); + // The second series closed too, and the input still has not ended. + execution + .push_evaluated(evaluated_chunk_of( + &schema, + job_instance_metric("requests", "db", "b"), + 0, + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + true, + )) + .unwrap(); + assert!( + !computed(&execution), + "the fold finished before the input ended" + ); + assert!( + execution.next_chunk().unwrap().is_none(), + "every series closed, and the group was answered before the input ended" + ); + + execution.end_of_input().unwrap(); + assert!( + !computed(&execution), + "ending the input finished the fold before anything pulled for it" + ); + let mut output = Vec::new(); + while let Some(chunk) = execution.next_chunk().unwrap() { + output.push(chunk); + } + assert_eq!( + groups_of(&output), + vec![unlabelled(vec![ + (0, Some(11.0)), + (1, Some(22.0)), + (2, Some(33.0)), + (3, Some(44.0)), + ])], + "the group was not answered once the input ended" + ); +} + +/// A chunk of another layout is refused, naming the aggregate and its position, and the execution +/// stops with nothing reserved. +/// +/// The fold indexes its states by absolute position and by field, so a chunk of another field count +/// would be folded into states that mean something else. It is refused rather than folded, and what +/// the refusal leaves behind is a stopped execution whose reservation went back to the pool. +#[test] +fn an_aggregate_refuses_a_chunk_of_another_layout_and_returns_its_reservation() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + // Two fields where the tree bound one: the source node accepts it, the fold refuses it. + let wider = Arc::new(SeriesSchema { + fields: field_count(2), + eval_grid: Some(Arc::clone(&grid)), + }); + execution + .push_evaluated(evaluated_columns( + &wider, + metric(), + 0, + vec![vec![Some(1.0); 4], vec![Some(2.0); 4]], + true, + )) + .unwrap(); + + let error = execution + .next_chunk() + .expect_err("a chunk of another field count was folded"); + assert!( + matches!(error, DataFusionError::Execution(_)), + "the refusal is not an execution error: {error}" + ); + let message = error.to_string(); + assert!( + message.contains("aggregate at root"), + "the refusal does not name the aggregate: {message}" + ); + assert!( + message.contains("field count"), + "the refusal does not say what disagreed: {message}" + ); + assert_eq!(execution.state().reserved_bytes(), 0); + // A stopped execution is not resumed: the fold must not be asked to continue after this. + assert!(execution.next_chunk().is_err()); +} + +/// A pool rejection while the first aggregate block is admitted releases the source and fold. +#[test] +fn aggregate_budget_rejects_first_block_and_releases_everything() { + let grid = grid(0, 1, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let pool = bounded_pool(2_000); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(2.0)], + true, + )) + .unwrap(); + assert!(execution.next_chunk().is_err()); + assert_eq!(pool.reserved(), 0); + assert_eq!(execution.state().reserved_bytes(), 0); +} + +/// The aggregate runtime charges its output queue, releases each output as it is pulled, and keeps +/// queue backing until the runtime itself is released. +#[test] +fn aggregate_runtime_queue_keeps_only_backing_after_pops_until_release() { + let grid = grid(0, 1, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new(["job".into()]))), + )) + .unwrap(); + let pool = bounded_pool(1 << 20); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + for (job, value) in [("api", 1.0), ("db", 2.0)] { + execution + .push_evaluated(evaluated_chunk_of( + &schema, + job_instance_metric("requests", job, "a"), + 0, + vec![Some(value)], + true, + )) + .unwrap(); + } + execution.end_of_input().unwrap(); + let first = execution.next_chunk().unwrap().expect("first group output"); + assert_eq!(first.series.time.positions().count(), 1); + let runtime = &execution.state().runtimes[0]; + let entry = std::mem::size_of::<(NativeChunk, usize)>(); + assert_eq!(runtime.queued.len(), 1); + let remaining = runtime.queued[0].1; + let expected = runtime.queued.capacity() * entry + remaining; + assert_eq!(runtime.buffer.reservation.size(), expected); + + let second = execution + .next_chunk() + .unwrap() + .expect("second group output"); + assert_eq!(second.series.time.positions().count(), 1); + let runtime = &execution.state().runtimes[0]; + assert!(runtime.queued.is_empty()); + let backing = runtime.queued.capacity() * entry; + assert!(backing > 0); + assert_eq!(runtime.buffer.reservation.size(), backing); + assert!(execution.next_chunk().unwrap().is_none()); + // Pulling the terminal None releases the queue backing; only explicit close checks the + // runtime's final cleanup below. + assert_eq!(execution.state().runtimes[0].buffer.reservation.size(), 0); + + execution.state.close().unwrap(); + assert_eq!(execution.state().runtimes[0].queued.capacity(), 0); + assert_eq!(execution.state().runtimes[0].buffer.reservation.size(), 0); + assert_eq!(pool.reserved(), 0); +} + +/// An aggregate whose sparse output has a few points on a large grid can finish in a small pool. +#[test] +fn aggregate_sparse_output_budget_uses_rows_not_grid_length() { + let grid = grid(0, 1, 100_000); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let pool = bounded_pool(64 * 1024); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 10, + vec![Some(2.0)], + true, + )) + .unwrap(); + let output = execution.finish().unwrap(); + assert_eq!(groups_of(&output), vec![unlabelled(vec![(10, Some(2.0))])]); + assert_eq!(pool.reserved(), 0); +} + +/// An aggregate retains its fold before finish and release returns its reservation. +#[test] +fn aggregate_fold_state_is_released_after_input() { + let grid = grid(0, 1, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let pool = bounded_pool(1 << 20); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(2.0)], + true, + )) + .unwrap(); + assert!(execution.next_chunk().unwrap().is_none()); + let runtime = &execution.state.runtimes[0]; + let aggregate = runtime.aggregate.as_ref().unwrap(); + assert_eq!(aggregate.groups(), 1); + assert!(!aggregate.computed); + assert!(pool.reserved() > 0); + execution.state.close().unwrap(); + assert_eq!(pool.reserved(), 0); +} + +/// An aggregate kind or stage this runtime does not compute is refused, naming what was found. +#[test] +fn an_aggregate_kind_or_stage_the_runtime_does_not_compute_is_named() { + let grid = || grid(0, 60_000, 4); + let kind_tree = NativeTree::try_new(aggregate( + evaluated_input(grid()), + BoundAggregate { + kind: AggregateKind::Quantile { quantile: 0.5 }, + group: Grouping::By(Arc::new([])), + stage: AggregateStage::Complete, + }, + )) + .unwrap(); + let partial_tree = { + let out = grid(); + let group = Grouping::By(Arc::new(["job".into()])); + let state = Arc::new(AggregateStateSchema { + group: group.clone(), + accumulators: fields(), + grid: Arc::clone(&out), + }); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(evaluated_input(Arc::clone(&out))), + op: BoundAggregate { + kind: AggregateKind::Sum, + group, + stage: AggregateStage::Partial, + }, + }, + output: NativeType::AggregateState(state), + }) + .unwrap() + }; + + for (found, tree) in [("Quantile", kind_tree), ("Partial", partial_tree)] { + let error = match NativeExecutionState::try_new( + &tree, + pool(), + &ExecutionPlanMetricsSet::new(), + 0, + ) { + Ok(_) => panic!("the runtime accepted a {found} aggregate"), + Err(error) => error, + }; + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "{found}: {error}" + ); + let message = error.to_string(); + assert!(message.contains(found), "{found}: {message}"); + // A `Partial` rejection has to say why it is not computed: the stages it would need. + if found == "Partial" { + assert!(message.contains("Final"), "{message}"); + } + } +} + +/// An aggregate above a window that collects and merges its input waits for that merge. +/// +/// The multiplication drops the metric name, so the two input series land on one label set and the +/// `last_over_time` above them has to collect and merge before it evaluates. The aggregate above the +/// window must wait for that merge: its guard is `inputs_finished && settled(child)`, and with only +/// the first half it finalizes on the pull that follows the caller's last chunk, before the window +/// has answered. +#[test] +fn an_aggregate_above_a_collecting_window_waits_for_the_merge() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = || { + NativeTree::try_new(aggregate( + window_with( + broadcast_right( + evaluated_input(Arc::clone(&grid)), + 2.0, + BinaryOp::Mul, + false, + ), + RangeFunction::LastOverTime, + 60_000, + Some(Arc::clone(&grid)), + Arc::clone(&grid), + ), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap() + }; + + let mut execution = executor(&tree(), pool(), &ExecutionPlanMetricsSet::new()); + // Two series that only differ by name, each holding half of the grid the other leaves empty: + // after the name is dropped they are one series only once merged. + execution + .push_evaluated(evaluated_chunk_of( + &schema, + named_metric("requests", "x"), + 0, + vec![Some(1.0), Some(2.0)], + true, + )) + .unwrap(); + assert!( + execution.next_chunk().unwrap().is_none(), + "the aggregate answered before the window above it had merged" + ); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + named_metric("requests_total", "x"), + 2, + vec![Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + + let out = run(&mut execution); + assert_eq!( + decided(&out), + vec![ + (0, Some(2.0)), + (1, Some(4.0)), + (2, Some(6.0)), + (3, Some(8.0)), + ], + "the aggregate did not fold the merged window result" + ); + assert_eq!(execution.state().reserved_bytes(), 0); +} + fn vector_op(inner: NativeNode, op: BoundVectorOp) -> NativeNode { let output = inner.output.clone(); NativeNode { @@ -7702,3 +8445,83 @@ fn histogram_quantile_reads_the_pieces_of_one_label_set_as_one_series() { .collect(); assert_eq!(positions, [0, 1]); } + +/// `count_values` counts the samples that share a value across the whole vector at a point, under +/// a label set per value; `a` arrives in two chunks and counts at both points. +#[test] +fn count_values_counts_the_samples_sharing_a_value() { + let grid = grid(0, 60_000, 2); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(vector_op( + evaluated_input(Arc::clone(&grid)), + BoundVectorOp::CountValues { + label: "v".into(), + group: Grouping::By(Arc::from([])), + }, + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + for (job, start, value, is_last) in [ + ("a", 0, 1.0, false), + ("a", 1, 1.0, true), + ("b", 0, 2.0, true), + ("c", 0, 1.0, true), + ] { + execution + .push_evaluated(evaluated_chunk_of( + &schema, + named_metric("m", job), + start, + vec![Some(value)], + is_last, + )) + .unwrap(); + } + let mut counts: Vec<(String, Decided)> = run(&mut execution) + .iter() + .map(|chunk| { + ( + chunk.series.metric.get("v").unwrap().to_string(), + decided(std::slice::from_ref(chunk)), + ) + }) + .collect(); + counts.sort_by(|a, b| a.0.cmp(&b.0)); + assert_eq!( + counts, + [ + ("1".to_string(), vec![(0, Some(2.0)), (1, Some(1.0))]), + ("2".to_string(), vec![(0, Some(1.0))]), + ] + ); +} + +/// Every field's distinct values make `count_values` groups of their own, so its estimate grows +/// with the field count, not only with the rows: a row of four fields reserves at least four +/// times what a row of one field does. +#[test] +fn count_values_reserves_a_group_per_field_sample() { + let grid = grid(0, 60_000, 1); + let op = BoundVectorOp::CountValues { + label: "v".into(), + group: Grouping::By(Arc::from([])), + }; + let estimate = |fields: usize| { + let schema = SeriesSchema { + fields: field_count(fields), + eval_grid: Some(Arc::clone(&grid)), + }; + let samples = (0..fields) + .map(|field| Samples::from_float(Float64Array::from(vec![field as f64]))) + .collect(); + let series = Series::try_new( + &schema, + named_metric("m", "a"), + TimeAxis::Grid(GridRange::try_new(Arc::clone(&grid), 0, 1).unwrap()), + SeriesFields::try_from_vec(samples).unwrap(), + ) + .unwrap(); + vector_workspace_bytes(&op, &schema, &[series]) + }; + assert!(estimate(4) >= 4 * estimate(1)); +} diff --git a/src/promql/src/native/mod.rs b/src/promql/src/native/mod.rs index c9545a9fe69..fd04dcfda32 100644 --- a/src/promql/src/native/mod.rs +++ b/src/promql/src/native/mod.rs @@ -27,8 +27,10 @@ //! - A series closes exactly once, on the chunk that says so, so a consumer can tell "nothing more //! here" from "no more input". +mod aggregate; pub mod batch; mod display; pub mod execute; pub mod metrics; pub mod plan; +pub(crate) mod runs; diff --git a/src/promql/src/native/plan.rs b/src/promql/src/native/plan.rs index 48ffdf52d3a..3e7b5fcf559 100644 --- a/src/promql/src/native/plan.rs +++ b/src/promql/src/native/plan.rs @@ -259,7 +259,7 @@ pub struct BoundAggregate { /// An operator over the whole input vector at each grid point: one that keeps some of the input's /// samples with their series' labels and values, orders the series, reports where none is present, -/// or reads a classic histogram from its bucket series. +/// reads a classic histogram from its bucket series, or counts the samples sharing a value. /// /// Its constants are bound here, as a window's are: a `k` or a ratio only known at run time has no /// field. @@ -282,6 +282,9 @@ pub enum BoundVectorOp { /// `histogram_quantile` or `histogram_fraction`: each classic histogram is read from the bucket /// series that share its labels, so this too reads the whole vector at a point. Histogram(HistogramFunction), + /// `count_values(label, v)`: how many samples of each group share a value, under a label set + /// per value. It writes a label per distinct value, which an aggregate's fixed groups cannot. + CountValues { label: Arc, group: Grouping }, } /// What a node computes. Every variant holds a bound operation, never a name to resolve. diff --git a/src/promql/src/native/runs.rs b/src/promql/src/native/runs.rs new file mode 100644 index 00000000000..8de325a5386 --- /dev/null +++ b/src/promql/src/native/runs.rs @@ -0,0 +1,146 @@ +// 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. + +//! The grid positions a field occupies, as sorted, disjoint runs. +//! +//! A node that rewrites labels — a `Map` dropping the metric name, a `Window` dropping the name and +//! the empty-valued labels — can turn two input series into one label set, and every layer above +//! then reads the two chunks as one series. What decides whether that pair is legal is which +//! positions of each field already hold a value: two values of one field at one position cannot be +//! merged away, while a chunk that fills positions the set left empty merges cleanly, as +//! [`crate::grid_series::merge_same_labelset`] does. +//! +//! The state is kept per label set, per field, as runs, so a chunk of many positions costs one run +//! rather than one entry per position. Keeping it in one place is what keeps the two layers that +//! check it — the native runtime's own [`crate::native::execute::Emitted`] table and the +//! cross-partition [`crate::extension_plan::fused_series::LabelSets`] — from drifting apart: both +//! ask the same functions for the same answers, so the fast path one of them earned is the fast path +//! the other gets. + +/// One run of grid positions, half-open, that hold a value in one field. +#[derive(Clone, Copy, Debug)] +pub(crate) struct Run { + pub(crate) start: usize, + pub(crate) end: usize, +} + +impl Run { + /// Whether two runs share a position. + pub(crate) fn overlaps(self, other: Self) -> bool { + self.start < other.end && other.start < self.end + } +} + +// How many times `absorb` rebuilt the stored runs instead of appending to them. +// +// Test-only, and per thread: the runs stored afterwards are the same either way, so a rebuild shows +// only in what it costs, and a test that reads its own thread's figure is not disturbed by another +// test's chunks. +#[cfg(test)] +thread_local! { + static ABSORB_REBUILDS: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +/// How many times this thread's [`absorb`] rebuilt the stored runs. +#[cfg(test)] +pub(crate) fn absorb_rebuilds() -> usize { + ABSORB_REBUILDS.with(std::cell::Cell::get) +} + +/// Merges `incoming` into `stored`, keeping the result in position order and joining the runs that +/// touch. +/// +/// Within a chunk one pass over both sides is what keeps this linear, since each side is sorted and +/// disjoint. Adding the runs one at a time and rebuilding the array after each would cost the square +/// of the runs a chunk yields: a chunk of one hundred thousand positions alternating between values +/// and nulls yields fifty thousand runs, and re-merging after each is over a billion element moves +/// for a single series. +/// +/// One pass per chunk is still a pass over everything the series has stored so far, so it is not +/// free once the *history* grows across chunks: a series split into `K` chunks of one position each, +/// a hole apart, stores one more run per chunk and moves `K(K+1)/2` of them, with no label collision +/// and no illegal input anywhere. A chunk that starts no earlier than the end of the last stored run +/// only ever adds to the end — the stored runs are sorted and disjoint, so every one of them ends +/// before anything this chunk holds — and it is appended instead, which is amortised constant per +/// run. That is the shape a scan hands over as its positions move forward chunk after chunk. A chunk +/// that starts before the last stored run can split a run or reach back past it, so it still costs +/// the full merge. +/// +/// Either way the result is sorted, disjoint, and has the runs that touch joined, which is what a +/// look-up by [`Run::overlaps`] on the runs stored here relies on. +pub(crate) fn absorb(stored: &mut Vec, incoming: &[Run]) { + if incoming.is_empty() { + return; + } + if stored + .last() + .is_none_or(|last| incoming[0].start >= last.end) + { + append_runs(stored, incoming); + } else { + merge_runs(stored, incoming); + } +} + +/// Appends `incoming` to `stored`, all of which lies past the last stored run. +/// +/// Nothing already stored has to be placed again: the two sides are each in position order and +/// nothing `incoming` holds can reach back into what is stored, so its runs are pushed on the end and +/// the vector's own growth is the only move — amortised over the runs appended. Where its first run +/// touches the end of the last stored run the two are one run, as they are after a merge, so it is +/// joined rather than left beside it. +fn append_runs(stored: &mut Vec, incoming: &[Run]) { + let mut incoming = incoming; + if let Some(last) = stored.last_mut() + && let Some(first) = incoming.first() + && first.start <= last.end + { + last.end = last.end.max(first.end); + incoming = &incoming[1..]; + } + stored.extend_from_slice(incoming); +} + +/// Merges `incoming` into `stored`, where the two can interleave. +/// +/// One pass is enough because each side is sorted and disjoint: whichever run starts first is +/// placed, and it joins the run before it where the two touch. +fn merge_runs(stored: &mut Vec, incoming: &[Run]) { + #[cfg(test)] + ABSORB_REBUILDS.with(|rebuilds| rebuilds.set(rebuilds.get() + 1)); + let mut merged: Vec = Vec::with_capacity(stored.len() + incoming.len()); + let (mut a, mut b) = (0, 0); + while a < stored.len() || b < incoming.len() { + let larger = match (stored.get(a), incoming.get(b)) { + (Some(left), Some(right)) if right.start < left.start => { + b += 1; + *right + } + (Some(left), _) => { + a += 1; + *left + } + (None, Some(right)) => { + b += 1; + *right + } + (None, None) => break, + }; + match merged.last_mut() { + Some(last) if larger.start <= last.end => last.end = last.end.max(larger.end), + _ => merged.push(larger), + } + } + *stored = merged; +} diff --git a/src/query/src/promql/native_compile/tests.rs b/src/query/src/promql/native_compile/tests.rs index d131af69ba9..d42ae4fa7b4 100644 --- a/src/query/src/promql/native_compile/tests.rs +++ b/src/query/src/promql/native_compile/tests.rs @@ -13,17 +13,17 @@ // limitations under the License. use std::sync::Arc; -use std::time::{Duration, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use datafusion::error::DataFusionError; use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datatypes::arrow::array::Float64Array; use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; -use datatypes::series::MetricName; use datatypes::series::SeriesOrdering; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{FieldSpec, SampleRef, Samples, SeriesFields, SeriesSchema}; +use datatypes::series::{Label, MetricName}; use promql::evaluated::EvaluatedChunk; use promql::map::MathOp; use promql::native::execute::{NativeExecutionState, NativeExecutor}; @@ -3138,11 +3138,25 @@ fn large_subquery_child_grids_compile_deferred() { assert_eq!(grid.2, range_ms); } +/// The largest millisecond endpoint the platform can represent as a `SystemTime`. +/// +/// The tests below need an endpoint close enough to `i64::MAX` milliseconds that a one-second +/// rounding or a one-millisecond offset crosses the representable range. Unix-like platforms can +/// hold that instant; Windows represents `SystemTime` as a 64-bit count of 100-nanosecond +/// intervals since 1601, so the same instant does not fit and `UNIX_EPOCH + duration` panics. The +/// behaviour under test is about arithmetic on the millisecond value, not about the platform's +/// clock range, so a platform that cannot represent the instant skips instead of failing. +fn extreme_millisecond_endpoint() -> Option { + UNIX_EPOCH.checked_add(Duration::from_millis(i64::MAX as u64)) +} + /// A sampled child endpoint can fit even when the unrounded read-window endpoint does not. #[test] fn a_subquery_accepts_an_in_range_sampled_endpoint() { + let Some(endpoint) = extreme_millisecond_endpoint() else { + return; + }; let mut stmt = eval_stmt("rate(x[1s:1s] offset -1ms)"); - let endpoint = UNIX_EPOCH + Duration::from_millis(i64::MAX as u64); stmt.start = endpoint; stmt.end = endpoint; stmt.interval = Duration::from_millis(1); @@ -3161,6 +3175,9 @@ fn a_subquery_accepts_an_in_range_sampled_endpoint() { /// An anchored fold offset that cannot be represented as i64 is a plan error. #[test] fn a_subquery_with_an_overflowing_fold_offset_is_refused() { + let Some(endpoint) = extreme_millisecond_endpoint() else { + return; + }; let mut stmt = eval_stmt("rate(x[5m:1m] @ 0)"); let PromExpr::Call(call) = &mut stmt.expr else { unreachable!("query parses to a call"); @@ -3168,9 +3185,7 @@ fn a_subquery_with_an_overflowing_fold_offset_is_refused() { let PromExpr::Subquery(subquery) = &mut *call.args.args[0] else { unreachable!("call argument parses to a subquery"); }; - subquery.at = Some(AtModifier::At( - UNIX_EPOCH + Duration::from_millis(i64::MAX as u64), - )); + subquery.at = Some(AtModifier::At(endpoint)); subquery.offset = Some(Offset::Neg(Duration::from_millis(1))); let error = compile(&stmt, raw_schema(), SeriesOrdering::KeyThenTime) @@ -3723,3 +3738,526 @@ fn whole_vector_functions_bind_a_vector_operation() { assert_eq!(root, format!(" root Vector {op}"), "`{query}`"); } } + +/// One raw chunk of `metric`, keyed by `key`, closing the series. +/// +/// `timestamps` are milliseconds and `values` one per timestamp, which is what a scan of the series +/// hands the tree. +fn agg_raw_chunk( + schema: &SeriesSchema, + key: &[u8], + metric: MetricName, + timestamps: &[i64], + values: &[f64], +) -> RawSeriesChunk { + agg_raw_chunk_with_last(schema, key, metric, timestamps, values, true) +} + +/// One raw chunk of `metric`, keyed by `key`, with an explicit series-ending flag. +fn agg_raw_chunk_with_last( + schema: &SeriesSchema, + key: &[u8], + metric: MetricName, + timestamps: &[i64], + values: &[f64], + is_last: bool, +) -> RawSeriesChunk { + RawSeriesChunk::try_new( + schema, + Buffer::from_slice_ref(key), + Some(metric), + RawTimestamps::try_new(ScalarBuffer::from(timestamps.to_vec())).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(values.to_vec()))), + is_last, + ) + .unwrap() +} + +/// `name` with the labels in `tags`, in the order given. +fn agg_metric(name: &str, tags: &[(&str, &str)]) -> MetricName { + MetricName::try_new( + Some(name.into()), + tags.iter() + .map(|(name, value)| Label { + name: (*name).into(), + value: (*value).into(), + }) + .collect(), + ) + .unwrap() +} + +/// What one result series of a native execution reports. +type AggSeries = ( + Option, + Vec<(String, String)>, + Vec<(usize, Option)>, +); + +/// Reads the outputs a native execution produced. +fn aggregate_results(output: &[EvaluatedChunk]) -> Vec { + let mut results: Vec = output + .iter() + .map(|chunk| { + let mut tags: Vec<(String, String)> = chunk + .series + .metric + .tags() + .iter() + .map(|tag| (tag.name.to_string(), tag.value.to_string())) + .collect(); + tags.sort(); + let points = chunk + .series + .time + .positions() + .map(|(position, row)| { + let sample = chunk.series.fields.as_slice()[0] + .iter() + .nth(row) + .expect("the axis and the field are aligned"); + ( + position, + match sample { + SampleRef::Float(value) => Some(value), + _ => None, + }, + ) + }) + .collect(); + (chunk.series.metric.name().map(str::to_string), tags, points) + }) + .collect(); + results.sort_by(|left, right| (&left.0, &left.1).cmp(&(&right.0, &right.1))); + results +} + +/// Runs `tree` over `chunks`, which enter at the tree's only input, and reads the answer. +/// +/// The answer is read the way a caller reads it: the whole input is pushed, the input is ended, and +/// what the root reports is collected. Each series is reported by its identity and the value it +/// decided at each absolute grid position. +fn run_aggregate_tree(tree: &NativeTree, chunks: Vec) -> Vec { + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 26)); + let state = + NativeExecutionState::try_new(tree, pool, &ExecutionPlanMetricsSet::new(), 0).unwrap(); + let mut execution = NativeExecutor::new(state); + for chunk in chunks { + execution + .push_raw(chunk) + .unwrap_or_else(|error| panic!("a source chunk was refused: {error}")); + } + let output = execution + .finish() + .unwrap_or_else(|error| panic!("the tree did not answer: {error}")); + aggregate_results(&output) +} + +/// The 21 grid points of a statement over 100 seconds in 5-second steps. +fn agg_points_0_to_100() -> Vec { + (0..21).collect() +} + +/// The timestamps of the same 21 points, in milliseconds. +fn agg_timestamps_0_to_100() -> Vec { + (0..21).map(|index| index * 5_000).collect() +} + +/// `sum(x)` over two series: one unlabelled group, and every label of both inputs dropped. +#[test] +fn a_compiled_sum_drops_every_label_and_adds_its_group() { + let tree = compile( + &eval_stmt("sum(x)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let schema = raw_schema(); + let timestamps = agg_timestamps_0_to_100(); + let chunks = vec![ + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a")]), + ×tamps, + &(1..=21).map(f64::from).collect::>(), + ), + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "db"), ("instance", "b")]), + ×tamps, + &(1..=21).map(|value| (2 * value) as f64).collect::>(), + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + // `3` at the first point, three more at each of the twenty after it: the two members summed. + let values: Vec<(usize, Option)> = agg_points_0_to_100() + .into_iter() + .map(|position| (position, Some(3.0 * (position as f64 + 1.0)))) + .collect(); + assert_eq!(output, vec![(None, Vec::new(), values)]); +} + +#[test] +fn a_compiled_sum_aggregates_series_chunks_after_input_ends() { + let schema = raw_schema(); + let timestamps = agg_timestamps_0_to_100(); + let values: Vec = (1..=21).map(f64::from).collect(); + let cases = [ + ("sum(x)", None, Vec::new()), + ( + "sum by(job)(x)", + None, + vec![("job".to_string(), "api".to_string())], + ), + ( + "sum without(instance)(x)", + None, + vec![("job".to_string(), "api".to_string())], + ), + ]; + + for (query, name, tags) in cases { + let tree = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + let metric = agg_metric("x", &[("job", "api"), ("instance", "a")]); + let chunks = vec![ + agg_raw_chunk_with_last( + &schema, + b"a", + metric.clone(), + ×tamps[..10], + &values[..10], + false, + ), + agg_raw_chunk_with_last( + &schema, + b"a", + metric, + ×tamps[10..], + &values[10..], + true, + ), + ]; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 26)); + let state = + NativeExecutionState::try_new(&tree, pool, &ExecutionPlanMetricsSet::new(), 0).unwrap(); + let mut execution = NativeExecutor::new(state); + for chunk in chunks { + execution.push_raw(chunk).unwrap(); + assert!( + execution.next_chunk().unwrap().is_none(), + "`{query}` produced output before the whole input ended" + ); + } + let output = execution.finish().unwrap(); + let results = aggregate_results(&output); + let expected_points = values + .iter() + .enumerate() + .map(|(position, value)| (position, Some(*value))) + .collect(); + assert_eq!( + results, + vec![(name.map(str::to_string), tags, expected_points)], + "`{query}`" + ); + } +} + +/// Empty source input is not a member of an aggregate's implicit group. +#[test] +fn a_compiled_aggregate_over_empty_input_produces_no_groups() { + for query in ["sum(x)", "sum by(job)(x)", "sum without(instance)(x)"] { + let tree = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + assert_eq!( + run_aggregate_tree(&tree, Vec::new()), + Vec::new(), + "`{query}` invented a group for empty input" + ); + } +} + +/// `sum by(job)(x)` over five series: two groups, each the sum of its own members. +/// +/// The members cover different points, so each group is a sum over a subset that changes from point +/// to point: one series holds the whole grid, one holds a single sample between two grid points, and +/// one stops a whole lookback before the grid does. No last sample sits exactly one lookback behind +/// a grid point, so which points each member covers does not depend on how the boundary is read. +#[test] +fn a_compiled_sum_by_groups_its_members_and_sums_each_group() { + let tree = compile( + &eval_stmt("sum by(job)(x)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let schema = raw_schema(); + let all = agg_timestamps_0_to_100(); + let chunks = vec![ + // The whole grid, holding one. + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a")]), + &all, + &[1.0; 21], + ), + // One sample halfway between grid points 0 and 1, so only grid point 1 selects it. + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "api"), ("instance", "b")]), + &[2_500], + &[2.0], + ), + // Four samples between grid points, covering grid points 1 to 4. + agg_raw_chunk( + &schema, + b"e", + agg_metric("x", &[("job", "api"), ("instance", "e")]), + &[2_400, 7_400, 12_400, 17_400], + &[5.0; 4], + ), + // The whole grid, holding ten. + agg_raw_chunk( + &schema, + b"c", + agg_metric("x", &[("job", "db"), ("instance", "c")]), + &all, + &[10.0; 21], + ), + // One sample halfway between grid points 0 and 1, holding twenty. + agg_raw_chunk( + &schema, + b"d", + agg_metric("x", &[("job", "db"), ("instance", "d")]), + &[2_500], + &[20.0], + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + // `api` is `a` alone at points 0 and 5 to 20, and every member at the points they cover: + // point 1 is `1 + 2 + 5`, points 2 to 4 are `1 + 5`. + let api: Vec<(usize, Option)> = agg_points_0_to_100() + .into_iter() + .map(|position| { + let value = match position { + 1 => 8.0, + 2..=4 => 6.0, + _ => 1.0, + }; + (position, Some(value)) + }) + .collect(); + // `db` is the same shape over its own members: point 1 is `10 + 20`, the rest is `c` alone. + let db: Vec<(usize, Option)> = agg_points_0_to_100() + .into_iter() + .map(|position| (position, Some(if position == 1 { 30.0 } else { 10.0 }))) + .collect(); + assert_eq!( + output, + vec![ + (None, vec![("job".to_string(), "api".to_string())], api), + (None, vec![("job".to_string(), "db".to_string())], db), + ] + ); +} + +/// `sum without(instance)(x)`: `instance` goes, `job` and `zone` stay, and so does `__name__` going. +#[test] +fn a_compiled_sum_without_drops_the_named_label_and_the_name() { + let tree = compile( + &eval_stmt("sum without(instance)(x)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let schema = raw_schema(); + let timestamps = agg_timestamps_0_to_100(); + let chunks = vec![ + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a"), ("zone", "west")]), + ×tamps, + &[1.0; 21], + ), + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "api"), ("instance", "b"), ("zone", "east")]), + ×tamps, + &[2.0; 21], + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + let values_of = |value: f64| -> Vec<(usize, Option)> { + agg_points_0_to_100() + .into_iter() + .map(|position| (position, Some(value))) + .collect() + }; + // Two groups, because `zone` is not the dropped label: the two series differ by more than the + // label the grouping removes. + assert_eq!( + output, + vec![ + ( + None, + vec![ + ("job".to_string(), "api".to_string()), + ("zone".to_string(), "east".to_string()), + ], + values_of(2.0), + ), + ( + None, + vec![ + ("job".to_string(), "api".to_string()), + ("zone".to_string(), "west".to_string()), + ], + values_of(1.0), + ), + ] + ); +} + +/// `sum by(job)(rate(x[5m]))`: an aggregate over a window, mid-chain. +/// +/// Every series is a counter rising by one per five seconds, and every grid point's window ends on a +/// sample and begins five seconds before one, so the extrapolated rate is the same at every point: +/// `59 / 295 = 0.2` per second, which the double arithmetic lands on the literals below. +#[test] +fn a_compiled_sum_over_a_rate_groups_the_window_it_reads() { + let statement = EvalStmt { + expr: promql_parser::parser::parse("sum by(job)(rate(x[5m]))").unwrap(), + start: UNIX_EPOCH + Duration::from_secs(300), + end: UNIX_EPOCH + Duration::from_secs(400), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(5), + }; + let tree = compile(&statement, raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + // The window reads 300 seconds, so the samples start a full window before the first grid point. + let schema = raw_schema(); + let timestamps: Vec = (0..=80).map(|index| index * 5_000).collect(); + let values: Vec = (0..=80).map(|index| f64::from(index) + 1.0).collect(); + let chunks = vec![ + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a")]), + ×tamps, + &values, + ), + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "api"), ("instance", "b")]), + ×tamps, + &values, + ), + agg_raw_chunk( + &schema, + b"c", + agg_metric("x", &[("job", "db"), ("instance", "c")]), + ×tamps, + &values, + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + let at_every_point = |value: f64| -> Vec<(usize, Option)> { + (0..21).map(|position| (position, Some(value))).collect() + }; + // One `rate` is `0.19999999999999998`, and the pair of them add to `0.39999999999999997`. + assert_eq!( + output, + vec![ + ( + None, + vec![("job".to_string(), "api".to_string())], + at_every_point(0.39999999999999997), + ), + ( + None, + vec![("job".to_string(), "db".to_string())], + at_every_point(0.19999999999999998), + ), + ] + ); +} + +/// Splitting one counter's samples does not change the rate aggregate or make it answer early. +#[test] +fn a_compiled_sum_over_a_rate_aggregates_series_chunks_after_input_ends() { + let statement = EvalStmt { + expr: promql_parser::parser::parse("sum by(job)(rate(x[5m]))").unwrap(), + start: UNIX_EPOCH + Duration::from_secs(300), + end: UNIX_EPOCH + Duration::from_secs(400), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(5), + }; + let tree = compile(&statement, raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + let schema = raw_schema(); + let timestamps: Vec = (0..=80).map(|index| index * 5_000).collect(); + let values: Vec = (0..=80).map(|index| f64::from(index) + 1.0).collect(); + let metric = agg_metric("x", &[("job", "api"), ("instance", "a")]); + let chunks = [ + agg_raw_chunk_with_last( + &schema, + b"a", + metric.clone(), + ×tamps[..41], + &values[..41], + false, + ), + agg_raw_chunk_with_last( + &schema, + b"a", + metric, + ×tamps[41..], + &values[41..], + true, + ), + ]; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 26)); + let state = + NativeExecutionState::try_new(&tree, pool, &ExecutionPlanMetricsSet::new(), 0).unwrap(); + let mut execution = NativeExecutor::new(state); + for chunk in chunks { + execution.push_raw(chunk).unwrap(); + assert!( + execution.next_chunk().unwrap().is_none(), + "the rate aggregate produced output before the whole input ended" + ); + } + let output = execution.finish().unwrap(); + assert_eq!( + aggregate_results(&output), + vec![( + None, + vec![("job".to_string(), "api".to_string())], + (0..21) + .map(|position| (position, Some(0.19999999999999998))) + .collect(), + )] + ); +} + +/// Empty source input is not a member of this aggregate's implicit group, with or without a window. +#[test] +fn a_compiled_sum_over_a_rate_with_empty_input_produces_no_groups() { + let statement = EvalStmt { + expr: promql_parser::parser::parse("sum by(job)(rate(x[5m]))").unwrap(), + start: UNIX_EPOCH + Duration::from_secs(300), + end: UNIX_EPOCH + Duration::from_secs(400), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(5), + }; + let tree = compile(&statement, raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + assert_eq!(run_aggregate_tree(&tree, Vec::new()), Vec::new()); +} diff --git a/tests/perf/test_oswal_s3_bench_scripts.py b/tests/perf/test_oswal_s3_bench_scripts.py index dc81240f1ae..d95b29f7a52 100644 --- a/tests/perf/test_oswal_s3_bench_scripts.py +++ b/tests/perf/test_oswal_s3_bench_scripts.py @@ -20,9 +20,12 @@ import json import os import subprocess import sys +import re import tempfile import unittest +import urllib.parse from pathlib import Path +from types import SimpleNamespace SCRIPTS = Path(__file__).parents[2] / ".github/scripts" @@ -35,6 +38,7 @@ def load(name, filename): bench = load("oswal_s3_bench", "oswal-s3-bench.py") +crash = load("oswal_s3_bench_crash", "oswal-s3-bench-crash.py") summary = load("oswal_s3_bench_summary", "oswal-s3-bench-summary.py") ACK_BUCKETS = """# TYPE greptime_logstore_object_store_wal_append_ack_seconds histogram @@ -57,8 +61,9 @@ class PlanTest(unittest.TestCase): self.assertEqual([c["id"] for c in plan["cells"]], [ "1-1-R", "1-1-D", "1-1-E", "1-8-R", "1-8-D", "1-8-E", "2-1-E", "2-1-D", "2-1-R", "2-8-E", "2-8-D", "2-8-R"]) - self.assertEqual(plan["generator_args"], ["--data-profile", "D"]) - self.assertEqual((plan["rows"], plan["batch_rows"], plan["data_profile"]), (2_000_000, 1000, "D")) + self.assertEqual(plan["corpora"], [{"rows": 2_000_000, "dir": "corpus-2000000", "data_profile": "D", + "generator_args": ["--data-profile", "D"]}]) + self.assertEqual((plan["rows"], plan["batch_rows"]), (2_000_000, 1000)) self.assertEqual({c["config"]: c["wal"] for c in plan["cells"]}, { "R": {"PROVIDER": "raft_engine", "SYNC_WRITE": "true", "DIR": "/wal"}, "D": {"PROVIDER": "experimental_object_store", "ACK_MODE": "durable"}, @@ -66,8 +71,8 @@ class PlanTest(unittest.TestCase): def test_smoke_overrides_rows_and_rounds(self): plan = bench.build_plan("8", "2000000", "1000", "2", smoke=True) - self.assertEqual((plan["rows"], plan["rounds"], plan["data_profile"]), (50_000, 1, "custom")) - self.assertEqual(plan["generator_args"], ["--line-count", "50000"]) + self.assertEqual((plan["rows"], plan["rounds"], plan["corpora"][0]["data_profile"]), (50_000, 1, "custom")) + self.assertEqual(plan["corpora"][0]["generator_args"], ["--line-count", "50000"]) self.assertEqual([c["id"] for c in plan["cells"]], ["1-8-R", "1-8-D", "1-8-E"]) def test_invalid_inputs_are_rejected(self): @@ -76,6 +81,129 @@ class PlanTest(unittest.TestCase): with self.subTest(args=args), self.assertRaises(ValueError): bench.build_plan(*args, smoke=False) + def test_the_baselines_and_the_order_reverse_every_other_round(self): + plan = bench.build_plan("8", "2000000", "1000", "2", smoke=False, configs="R,R0,D,E") + self.assertEqual([c["id"] for c in plan["cells"]], [ + "1-8-R", "1-8-R0", "1-8-D", "1-8-E", "2-8-E", "2-8-D", "2-8-R0", "2-8-R"]) + cells = {c["id"]: c for c in plan["cells"]} + self.assertEqual((cells["1-8-R0"]["env"], cells["1-8-R0"]["wal"]), + ({}, {"PROVIDER": "raft_engine", "SYNC_WRITE": "false", "DIR": "/wal"})) + + def test_extra_env_and_restart_reach_every_cell(self): + for extra in ('{"RUST_BACKTRACE": 1, "RUST_LOG": "info"}', "RUST_BACKTRACE=1\n\nRUST_LOG = info\n"): + plan = bench.build_plan("8", "2000000", "1000", "1", smoke=False, configs="R,D", extra_env=extra, + restart=True) + with self.subTest(extra=extra): + self.assertEqual([c["id"] for c in plan["cells"]], ["1-8-R", "1-8-D"]) + cells = {c["id"]: c for c in plan["cells"]} + self.assertEqual(cells["1-8-D"]["env"], {"RUST_BACKTRACE": "1", "RUST_LOG": "info"}) + self.assertEqual(cells["1-8-D"]["wal"]["FLUSH_INTERVAL"], "10ms") + self.assertNotIn("FLUSH_INTERVAL", cells["1-8-R"]["wal"]) + self.assertTrue(all(c["restart"] for c in plan["cells"])) + + def test_invalid_extra_env_and_configs_are_rejected(self): + for kwargs in ({"extra_env": "GREPTIMEDB_STANDALONE__STORAGE__BUCKET=other"}, {"extra_env": "NO_VALUE"}, + {"extra_env": "GREPTIMEDB_STANDALONE__WAL__FLUSH_INTERVAL=100ms", "restart": True}, + {"extra_env": "GREPTIMEDB_STANDALONE__WAL__SYNC_WRITE=false"}, + {"extra_env": "GREPTIMEDB_STANDALONE__wal__FLUSH_INTERVAL=100ms", "restart": True}, + {"extra_env": "greptimedb_standalone__Storage__bucket=other"}, + {"extra_env": "A=1\nA=2"}, {"extra_env": '{"A": true}'}, {"extra_env": "1A=x"}, + {"configs": "R,X"}, {"configs": "D,D"}): + with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): + bench.build_plan("8", "6000", "1000", "1", smoke=False, **kwargs) + + def test_the_image_is_the_full_reference_or_a_greptimedb_tag(self): + image = "docker.io/greptime/greptimedb-dev:dev-20261013-1791848986-a5a1a9f4" + for valid in (image, "ghcr.io/acme/greptimedb--bench:dev", "ghcr.io/acme/greptimedb__bench:dev", + "localhost:5000/a/b@sha256:" + "c" * 64): + self.assertEqual(bench.resolve_image(valid, "nightly"), valid) + self.assertEqual(bench.resolve_image("", "v1.4.0-nightly-20261012"), "greptime/greptimedb:v1.4.0-nightly-20261012") + for bad in ("greptimedb", "docker.io/greptime/greptimedb", "Docker.io/x/y:1", "x/y:bad tag", + "x/y___z:1", "x/-y:1"): + with self.subTest(bad=bad), self.assertRaises(ValueError): + bench.resolve_image(bad, "nightly") + + def test_groups_inherit_the_inputs_and_keep_their_own_settings(self): + groups = [{"name": "recovery", "configs": "D,E", "rows": 4000000, "rounds": 1, "restart": True}, + {"name": "crashe", "mode": "crash", "configs": "E", "rounds": 1, "cycles": 3, "kill_after": [5, 9]}, + {"name": "pin", "mode": "crash", "configs": "D", "rounds": 1, "cycles": 1, + "kill_after": [540, 540], "pin": True}, + {"name": "idle", "configs": "D", "rounds": 1, "idle_seconds": 180}] + plan = bench.build_plan("8", "2000000", "1000", "2", smoke=False, groups=json.dumps(groups)) + self.assertEqual([c["id"] for c in plan["cells"]], [ + "recovery-1-8-D", "recovery-1-8-E", "crashe-1-8-E", "pin-1-8-D", "idle-1-8-D"]) + self.assertEqual([c["rows"] for c in plan["corpora"]], [4_000_000, 2_000_000]) + cells = {c["id"]: c for c in plan["cells"]} + self.assertEqual(cells["recovery-1-8-D"]["wal"]["FLUSH_INTERVAL"], "10ms") + self.assertEqual((cells["crashe-1-8-E"]["cycles"], cells["crashe-1-8-E"]["kill_after"]), (3, [5, 9])) + self.assertTrue(cells["pin-1-8-D"]["pin"]) + self.assertNotIn("FLUSH_INTERVAL", cells["pin-1-8-D"]["wal"]) + self.assertEqual(cells["idle-1-8-D"]["idle_seconds"], 180) + smoke = bench.build_plan("8", "2000000", "1000", "2", smoke=True, groups=json.dumps(groups)) + pin = next(c for c in smoke["cells"] if c["group"] == "pin") + self.assertEqual((pin["cycles"], pin["kill_after"]), (1, [60, 60])) + self.assertEqual([c["rows"] for c in smoke["corpora"]], [50_000]) + recovery = [{"name": "rec", "configs": "D", "restart": True, "min_objects": 6000}] + self.assertEqual(bench.build_plan("8", "4000000", "1000", "1", smoke=False, + groups=json.dumps(recovery))["cells"][0]["min_objects"], 6000) + self.assertEqual(bench.build_plan("8", "4000000", "1000", "1", smoke=True, + groups=json.dumps(recovery))["cells"][0]["min_objects"], 0) + + def test_the_workflow_passes_groups_through_bash_unchanged(self): + workflow = (SCRIPTS.parent / "workflows" / "oswal-s3-bench.yml").read_text() + self.assertIn("OSWAL_GROUPS: ${{ inputs.groups }}", workflow) + self.assertEqual(workflow.count('--groups "$OSWAL_GROUPS"'), 2) + command = (f'python3 "{SCRIPTS}/oswal-s3-bench.py" plan --smoke false --concurrency 8 --rows 6000 ' + '--batch-rows 1000 --rounds 1 --groups "$OSWAL_GROUPS"') + for groups, expected in (("", "3 cells"), ('[{"name": "a", "configs": "D,E"}]', "2 cells")): + done = subprocess.run(["bash", "-c", command], env=os.environ | {"OSWAL_GROUPS": groups}, + capture_output=True, text=True, timeout=60) + self.assertEqual(done.returncode, 0, done.stderr) + self.assertTrue(done.stdout.startswith(expected), done.stdout) + + def test_a_partitioned_group_and_a_drop_group(self): + groups = [{"name": "multi", "configs": "R,D", "concurrency": "8,32", "partitions": 8}, + {"name": "drop", "mode": "drop", "configs": "D", "rounds": 1, "write_seconds": 45}] + plan = bench.build_plan("8", "2000000", "1000", "1", smoke=False, groups=json.dumps(groups)) + cells = {c["id"]: c for c in plan["cells"]} + self.assertEqual(list(cells), ["multi-1-8-R", "multi-1-8-D", "multi-1-32-R", "multi-1-32-D", "drop-1-8-D"]) + self.assertEqual({c["partitions"] for c in plan["cells"] if c["mode"] == "load"}, {8}) + self.assertEqual((cells["drop-1-8-D"]["write_seconds"], cells["drop-1-8-D"]["mode"]), (45, "drop")) + self.assertNotIn("rows", cells["drop-1-8-D"]) + self.assertEqual([c["rows"] for c in plan["corpora"]], [2_000_000]) + smoke = bench.build_plan("8", "2000000", "1000", "1", smoke=True, groups=json.dumps(groups)) + self.assertEqual(next(c for c in smoke["cells"] if c["mode"] == "drop")["write_seconds"], 10) + self.assertEqual(bench.build_plan("8", "6000", "1000", "1", smoke=False)["cells"][0]["partitions"], 1) + + def test_the_partitioned_table_splits_the_corpus_window_into_equal_ranges(self): + ddl = bench.partitioned_table({"start": "2026-07-10T00:00:00Z", "end": "2026-07-10T01:00:00Z"}, 8) + start = 1783641600000 + bounds = [int(b) for b in re.findall(r"`event_time` < (\d+)", ddl)] + self.assertEqual(bounds, [start + 450_000 * i for i in range(1, 8)]) + self.assertIn(f"`event_time` >= {start + 450_000 * 7}\n)\nENGINE=mito", ddl) + self.assertEqual(ddl.count("`event_time` >="), 7) + self.assertTrue(ddl.startswith("CREATE TABLE `agent_observations` (")) + self.assertIn("PRIMARY KEY (`tenant`, `app`, `environment`)\n)\nPARTITION ON COLUMNS (`event_time`) (", ddl) + self.assertNotIn("{partitions}", ddl) + with self.assertRaises(ValueError): + bench.partitioned_table({"start": "2026-07-10T00:00:00Z", "end": "2026-07-10T00:00:00.004Z"}, 8) + + def test_invalid_groups_are_rejected(self): + for groups in ([], [{"configs": "D"}], [{"name": "Bad"}], [{"name": "a"}, {"name": "a"}], + [{"name": "a", "colour": "red"}], [{"name": "a", "mode": "crash", "configs": "R"}], + [{"name": "a", "mode": "crash", "configs": "D", "restart": True}], + [{"name": "a", "mode": "crash", "configs": "D", "kill_after": [9, 5]}], + [{"name": "a", "mode": "soak"}], [{"name": "a", "idle_seconds": 601}], + [{"name": "a", "configs": "D", "min_objects": 6000}], + [{"name": "a", "configs": "D", "restart": True, "min_objects": -1}], + [{"name": "a", "configs": "D", "partitions": 65}], [{"name": "a", "configs": "D", "partitions": 0}], + [{"name": "a", "mode": "crash", "configs": "D", "partitions": 8}], + [{"name": "a", "mode": "drop", "configs": "R"}], + [{"name": "a", "mode": "drop", "configs": "D", "write_seconds": 0}], + [{"name": "a", "mode": "drop", "configs": "D", "restart": True}]): + with self.subTest(groups=groups), self.assertRaises(ValueError): + bench.build_plan("8", "6000", "1000", "1", smoke=False, groups=json.dumps(groups)) + def test_nightly_tag_is_the_newest_by_date(self): tags = ["latest", "v1.3.0-alpha.1-nightly-20260907", "v1.4.0-nightly-20261012", "v1.4.0-nightly-20261005", "nightly-arm64"] @@ -96,15 +224,285 @@ class PlanTest(unittest.TestCase): with self.subTest(bucket=bucket), self.assertRaisesRegex(ValueError, "cross-region"): bench.check_regions(bucket, runner) - def test_manifest_takes_the_sha256_digest(self): + def test_manifest_records_the_checked_images_pinned_to_their_digests(self): plan = bench.build_plan("1,8", "2000000", "1000", "2", smoke=False) - inspect = [{"RepoDigests": ["greptime/greptimedb@sha256:" + "a" * 64]}] - manifest = bench.build_manifest(plan, inspect, "runtime@sha256:" + "b" * 64, "c6id.2xlarge", + checked = [{"label": "", "image": "greptime/greptimedb@sha256:" + "a" * 64, + "requested": "greptime/greptimedb:v1.4.0-nightly-20261012", "version": "1.4.0", "commit": "abc"}] + manifest = bench.build_manifest(plan, checked, "runtime@sha256:" + "b" * 64, "c6id.2xlarge", "ap-southeast-1", "oswal-bench/1-1") self.assertEqual(manifest["greptimedb_image"], "greptime/greptimedb@sha256:" + "a" * 64) + self.assertEqual(manifest["greptimedb_images"], checked) self.assertEqual((manifest["rows"], manifest["rounds"], manifest["concurrency"]), (2_000_000, 2, [1, 8])) - with self.assertRaises(ValueError): - bench.build_manifest(plan, [{"RepoDigests": []}], "", "", "", "") + for images in ([checked[0] | {"image": "greptime/greptimedb:v1.4.0"}], [checked[0] | {"label": "base"}]): + with self.subTest(images=images), self.assertRaises(ValueError): + bench.build_manifest(plan, images, "", "", "", "") + + def test_every_cell_runs_once_per_image_in_an_order_that_reverses_every_other_round(self): + images = json.dumps([{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} + for label in ("base", "pr1", "pr2")]) + plan = bench.build_plan("8", "2000000", "1000", "3", smoke=False, configs="D", images=images) + self.assertEqual([c["id"] for c in plan["cells"]], [ + "1-8-D-base", "1-8-D-pr1", "1-8-D-pr2", "2-8-D-pr2", "2-8-D-pr1", "2-8-D-base", + "3-8-D-base", "3-8-D-pr1", "3-8-D-pr2"]) + self.assertEqual(plan["images"], ["base", "pr1", "pr2"]) + # The configurations reverse with the images, and a warm-up runs the first load cell per image. + groups = json.dumps([{"name": "crash", "mode": "crash", "configs": "E", "rounds": 1, "cycles": 1}, + {"name": "single", "configs": "D,E"}]) + plan = bench.build_plan("8", "2000000", "1000", "2", smoke=False, groups=groups, + images=json.dumps(json.loads(images)[:2]), warmup=True) + self.assertEqual([c["id"] for c in plan["cells"]], [ + "single-0-8-D-base", "single-0-8-D-pr1", "crash-1-8-E-base", "crash-1-8-E-pr1", + "single-1-8-D-base", "single-1-8-D-pr1", "single-1-8-E-base", "single-1-8-E-pr1", + "single-2-8-E-pr1", "single-2-8-E-base", "single-2-8-D-pr1", "single-2-8-D-base"]) + warmups = [c for c in plan["cells"] if c.get("warmup")] + self.assertEqual([(c["round"], c["image"], c["mode"]) for c in warmups], [(0, "base", "load"), (0, "pr1", "load")]) + self.assertEqual((plan["rounds"], plan["warmup"]), (2, True)) + # Without images every cell keeps its id and runs the one image. + single = bench.build_plan("8", "2000000", "1000", "1", smoke=False, configs="D") + self.assertEqual([(c["id"], c["image"]) for c in single["cells"]], [("1-8-D", "")]) + self.assertEqual((single["images"], single["warmup"]), ([""], False)) + + def test_invalid_images_and_warmups_are_rejected(self): + ref = "docker.io/greptime/greptimedb-dev:dev" + for images in ([{"label": "base", "image": ref}], [{"label": f"i{n}", "image": ref} for n in range(5)], + [{"label": "base", "image": ref}, {"label": "base", "image": ref}], + [{"label": "Base", "image": ref}, {"label": "pr1", "image": ref}], + [{"label": "ninechars", "image": ref}, {"label": "pr1", "image": ref}], + [{"label": 1, "image": ref}, {"label": "pr1", "image": ref}], + [{"label": "base", "image": "greptimedb"}, {"label": "pr1", "image": ref}], + [{"label": "base", "image": ref, "env": "A=1"}, {"label": "pr1", "image": ref}], + {"base": ref}): + with self.subTest(images=images), self.assertRaises(ValueError): + bench.build_plan("8", "6000", "1000", "1", smoke=False, images=json.dumps(images)) + crash_only = json.dumps([{"name": "c", "mode": "crash", "configs": "D"}]) + with self.assertRaisesRegex(ValueError, "warmup needs a load group"): + bench.build_plan("8", "6000", "1000", "1", smoke=False, groups=crash_only, warmup=True) + # One image twice compares it with itself. + pair = [{"label": "base", "image": ref}, {"label": "again", "image": ref}] + self.assertEqual(bench.resolve_images(json.dumps(pair), "", "nightly"), pair) + with self.assertRaisesRegex(ValueError, "not both"): + bench.resolve_images(json.dumps(pair), ref, "nightly") + self.assertEqual(bench.resolve_images("", ref, "nightly"), [{"label": "", "image": ref}]) + + def test_a_plan_longer_than_the_benchmark_timeout_is_refused(self): + groups = json.dumps([{"name": "single", "configs": "D,E"}, {"name": "multi", "configs": "D,E", "partitions": 8}]) + images = json.dumps([{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} + for label in ("base", "pr1", "pr2")]) + plan = bench.build_plan("8", "2000000", "1000", "5", smoke=False, groups=groups, images=images, warmup=True) + self.assertEqual(len(plan["cells"]), 63) + self.assertEqual(bench.estimate_minutes(plan), 75) + cells = {c["id"]: c for c in plan["cells"]} + # Durable at concurrency 8 writes 34,000 rows/s to one region and 26,000 to eight. + self.assertAlmostEqual(bench.cell_seconds(cells["single-1-8-D-base"], 1000), 5 + 2_000_000 / 34_000) + self.assertAlmostEqual(bench.cell_seconds(cells["multi-1-8-D-pr1"], 1000), 5 + 2_000_000 / 26_000) + self.assertAlmostEqual(bench.cell_seconds(cells["multi-1-8-E-pr2"], 4000), 5 + 2_000_000 / 60_000) + # A smaller batch keeps the request rate: one-row batches write a thousandth of the rows. + self.assertAlmostEqual(bench.cell_seconds(cells["single-1-8-D-base"], 1), 5 + 2_000_000 / 34) + durable = cells["single-1-8-D-base"] + self.assertAlmostEqual(bench.load_rate(durable | {"concurrency": 4}), 9_000 + 25_000 * 2 / 3) + self.assertEqual(bench.load_rate(durable | {"concurrency": 64}), 55_000) + crash = bench.build_plan("8", "6000", "1000", "1", smoke=False, groups=json.dumps( + [{"name": "c", "mode": "crash", "configs": "D", "cycles": 5, "kill_after": [5, 20]}]))["cells"][0] + self.assertEqual(bench.cell_seconds(crash, 1), 5 + 5 * (12.5 + 15)) + command = [sys.executable, str(SCRIPTS / "oswal-s3-bench.py"), "plan", "--smoke", "false", "--concurrency", + "8", "--rows", "2000000", "--rounds", "5", "--groups", groups, + "--images", images, "--warmup", "true", "--ttl-hours", "8"] + for batch_rows, timeout, expected in (("1000", "75", "about 75 minutes"), ("1000", "74", "more than the 74-minute"), + ("1", "360", "more than the 360-minute")): + done = subprocess.run(command + ["--batch-rows", batch_rows, "--timeout-minutes", timeout], + capture_output=True, text=True, timeout=60) + with self.subTest(batch_rows=batch_rows, timeout=timeout): + self.assertEqual(done.returncode, 0 if expected.startswith("about") else 1, done.stderr) + self.assertIn(expected, done.stdout + done.stderr) + + def test_the_workflow_checks_every_image_before_an_instance_is_allocated(self): + workflow = (SCRIPTS.parent / "workflows" / "oswal-s3-bench.yml").read_text() + self.assertIn("OSWAL_IMAGES: ${{ inputs.images }}", workflow) + self.assertEqual(workflow.count('--images "$OSWAL_IMAGES"'), 3) + self.assertEqual(workflow.count('--warmup "$WARMUP"'), 2) + validate = workflow[workflow.index(" validate:"):workflow.index(" provision:")] + self.assertIn("greptimedb_images: ${{ steps.images.outputs.greptimedb_images }}", validate) + self.assertIn("check-images \"$REQUESTED_IMAGES\"", validate) + self.assertIn("GREPTIMEDB_IMAGES: ${{ needs.validate.outputs.greptimedb_images }}", workflow) + + +def write_lines(path, lines): + path.write_text("".join(f"{line}\n" for line in lines)) + + +def crash_cycle(root, cycle, stripes, acked, requests, exits, found=None): + """One cycle of two writers: attempted stripes, acked batches, requests and exit times per + writer, and what the restart after its kill found, if given.""" + directory = root / f"cycle-{cycle}" + directory.mkdir(parents=True) + for index in range(2): + write_lines(directory / f"attempted-{index}.log", stripes[index]) + write_lines(directory / f"acked-{index}.log", acked[index]) + write_lines(directory / f"requests-{index}.log", requests[index]) + (directory / f"exit-{index}.log").write_text(f"{exits[index]}\n") + if found is not None: + (directory / "found.json").write_text(json.dumps( + {"output": [{"records": {"rows": [[seq, count] for seq, count in found.items()]}}]})) + return directory + + +class CrashOracleTest(unittest.TestCase): + # Writer 0 got batches 0 and 2 acknowledged and lost 4 in flight; writer 1 got 1 and lost 3. + STRIPES = ([0, 2, 4], [1, 3]) + ACKED = ([0, 2], [1]) + REQUESTS = (["0 200 10", "2 200 20", "4 000 105"], ["1 200 15", "3 000 101"]) + + def check(self, root, found, mode="durable", requests=REQUESTS, exits=(106, 102)): + crash_cycle(root, 1, self.STRIPES, self.ACKED, requests, exits) + result = crash.check_kill(root / "cycle-1", 2, 100, 104, 2000) + return result | crash.verify(root, found, 10, mode == "enqueued", 1) + + def test_acked_batches_present_once_and_an_in_flight_extra_pass(self): + with tempfile.TemporaryDirectory() as directory: + result = self.check(Path(directory), {0: 10, 1: 10, 2: 10, 4: 10}) + self.assertEqual((result["kill_errors"], result["verify_errors"]), ([], [])) + self.assertEqual((result["in_flight"], result["writers_at_kill"], result["cycle_extra_batches"], + result["cycle_lost_batches"], result["cycle_present_batches"]), (2, 2, 1, 0, 4)) + + def test_durable_violations_are_reported(self): + cases = {"acknowledged but absent": {0: 10, 1: 10}, "present more than once": {0: 20, 1: 10, 2: 10}, + "present in part": {0: 10, 1: 10, 2: 7}, "never attempted": {0: 10, 1: 10, 2: 10, 9: 10}} + for error, found in cases.items(): + with self.subTest(error=error), tempfile.TemporaryDirectory() as directory: + result = self.check(Path(directory), found) + self.assertTrue(any(error in e for e in result["verify_errors"]), result["verify_errors"]) + + def test_enqueued_loses_acked_batches_without_failing(self): + with tempfile.TemporaryDirectory() as directory: + result = self.check(Path(directory), {0: 10}, mode="enqueued") + self.assertEqual(result["verify_errors"], []) + self.assertEqual((result["cycle_lost_batches"], result["cycle_lost_rows"], result["total_lost_batches"]), + (2, 20, 2)) + + def test_a_later_cycle_reports_only_its_own_losses(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + crash_cycle(root, 1, self.STRIPES, self.ACKED, self.REQUESTS, (106, 102)) + crash_cycle(root, 2, ([10, 12], [11]), ([10], [11]), (["10 200 210", "12 000 305"], ["11 200 215"]), + (306, 302)) + # Cycle 1 lost batches 1 and 2; cycle 2 lost nothing. + result = crash.verify(root, {0: 10, 10: 10, 11: 10}, 10, True, 2) + self.assertEqual((result["cycle_acked_batches"], result["cycle_lost_batches"], + result["total_lost_batches"], result["total_acked_batches"]), (2, 0, 2, 5)) + + def later_cycle(self, root): + """Cycle 1's restart recovered batches 0, 1 and 2 and the in-flight 4; cycle 2 acknowledged + 10 and 11 and had 12 in flight.""" + crash_cycle(root, 1, self.STRIPES, self.ACKED, self.REQUESTS, (106, 102), found={0: 10, 1: 10, 2: 10, 4: 10}) + crash_cycle(root, 2, ([10, 12], [11]), ([10], [11]), (["10 200 210", "12 000 305"], ["11 200 215"]), + (306, 302)) + + def test_a_batch_an_earlier_restart_recovered_must_stay_present(self): + for mode in ("enqueued", "durable"): + with self.subTest(mode=mode), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.later_cycle(root) + # Acknowledged batch 1 and the in-flight 4 vanished after cycle 1's restart found them. + result = crash.verify(root, {0: 10, 2: 10, 10: 10, 11: 10}, 10, mode == "enqueued", 2) + self.assertIn("recovered by an earlier restart but absent: [1, 4]", result["verify_errors"]) + self.assertEqual(result["cycle_lost_batches"], 0) + + def test_only_the_current_cycle_may_lose_acknowledged_batches_in_enqueued_mode(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.later_cycle(root) + found = {0: 10, 1: 10, 2: 10, 4: 10, 10: 10} + result = crash.verify(root, found, 10, True, 2) + self.assertEqual(result["verify_errors"], []) + self.assertEqual((result["cycle_lost_batches"], result["total_lost_batches"]), (1, 1)) + durable = crash.verify(root, found, 10, False, 2) + self.assertEqual(durable["verify_errors"], ["acknowledged but absent: [11]"]) + duplicated = crash.verify(root, found | {10: 20}, 10, True, 2) + self.assertEqual(duplicated["verify_errors"], ["batches present more than once: [10]"]) + + def test_a_connection_lost_before_the_kill_fails(self): + with tempfile.TemporaryDirectory() as directory: + requests = (["0 200 10", "2 200 20", "4 000 50"], ["1 200 15", "3 500 101"]) + result = self.check(Path(directory), {0: 10, 1: 10, 2: 10}, requests=requests, exits=(51, 99)) + self.assertEqual(len(result["kill_errors"]), 3, result["kill_errors"]) + + def test_the_writer_stops_each_stripe_at_its_first_unacknowledged_batch(self): + import http.server + import threading + inserted, lock = [], threading.Lock() + + class Handler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + body = self.rfile.read(int(self.headers["Content-Length"])).decode() + with lock: + if len(inserted) >= 6: + self.close_connection = True + self.connection.close() + return + inserted.append(body) + self.send_response(200) + self.end_headers() + + def log_message(self, *args): + pass + + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + with tempfile.TemporaryDirectory() as directory: + run = Path(directory) + args = SimpleNamespace(dir=str(run / "cycle-1"), url=f"http://127.0.0.1:{server.server_port}/v1/sql", + table="crash", writers=2, first_seq=0, rows=3, seconds=0) + crash.write(args) + found = {} + for body in inserted: + for seq in set(re.findall(r"%28(\d+)%2C", body)): + found[int(seq)] = 3 + result = crash.verify(run, found, 3, False, 1) + self.assertEqual((result["verify_errors"], result["total_acked_batches"]), ([], 6)) + for index in range(2): + attempted = (run / "cycle-1" / f"attempted-{index}.log").read_text().split() + self.assertEqual([int(seq) % 2 for seq in attempted], [index] * len(attempted)) + finally: + server.shutdown() + server.server_close() + + +class WriterDeadlineTest(unittest.TestCase): + def test_writers_stop_at_the_deadline_with_every_batch_acknowledged(self): + import http.server + import threading + status = [200] + + class Handler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + self.rfile.read(int(self.headers["Content-Length"])) + self.send_response(status[0]) + self.end_headers() + + def log_message(self, *args): + pass + + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + with tempfile.TemporaryDirectory() as directory: + run = Path(directory) + args = SimpleNamespace(dir=str(run / "cycle-1"), url=f"http://127.0.0.1:{server.server_port}/", + table="dropped", writers=2, first_seq=0, rows=2, seconds=0.3) + self.assertEqual(crash.write(args), 0) + result = crash.verify(run, {}, 2, True, 1) + self.assertGreater(result["cycle_acked_batches"], 0) + self.assertEqual(result["verify_errors"], []) + for index in range(2): + self.assertTrue((run / "cycle-1" / f"exit-{index}.log").exists()) + # A rejected batch ends a timed writer early, which fails the command. + status[0] = 503 + args.dir, args.seconds = str(run / "cycle-2"), 30 + self.assertEqual(crash.write(args), 1) + finally: + server.shutdown() + server.server_close() class SummaryTest(unittest.TestCase): @@ -117,14 +515,15 @@ class SummaryTest(unittest.TestCase): self.assertEqual(summary.quantile(0.99, [(0.1, 1.0), (float("inf"), 4.0)]), 0.1) self.assertNotIn("kafka_client_produce_elapsed_count", summary.parse_metrics(ACK_BUCKETS)["samples"]) - def write_cell(self, root, cell_id, rows_per_second, exit_code=0, metrics=""): - round_number, concurrency, config = cell_id.split("-") + def write_cell(self, root, cell_id, rows_per_second, exit_code=0, metrics="", **extra): + round_number, concurrency, config, *image = cell_id.split("-") path = root / "cells" / cell_id (path / "greptimedb-1").mkdir(parents=True) (path / "cell.json").write_text(json.dumps({ "id": cell_id, "round": int(round_number), "concurrency": int(concurrency), "config": config, - "wal": bench.WAL_CONFIGS[config], "metrics_scraped": True, - "exit_code": exit_code, "wal_device": "/dev/nvme1n1", "version": "1.4.0", "load_wall_seconds": 61.0})) + "image": image[0] if image else "", + "wal": bench.WAL_CONFIGS[config], "metrics_scraped": True, "env": {}, + "exit_code": exit_code, "wal_device": "/dev/nvme1n1", "version": "1.4.0", "load_wall_seconds": 61.0} | extra)) if exit_code == 0: (path / "greptimedb-1" / "load-result.json").write_text(json.dumps({ "count_after": 6000, "load_seconds": 60.0, "rows_per_second": rows_per_second})) @@ -149,10 +548,172 @@ class SummaryTest(unittest.TestCase): d_cell = next(c for c in result["cells"] if c["id"] == "1-8-D") self.assertEqual(d_cell["creates_per_minute"], 120.0) self.assertEqual(d_cell["stalled_appends_total"], 3) - self.assertIn("| 1-8-D | 80 | 60.0 | 61.0 | 6,000 | 128.0 | 253.4 | 120 | 120.0 | 3 | 0 | 0 | 0 |", output) + self.assertIn("| 1-8-D | — | 80 | 60.0 | 61.0 | 128.0 / 253.4 | — / — | — / — | 120 | 120.0 | 3 | 0 | 0 | 0 | — | — |", output) r_cell = next(c for c in result["cells"] if c["id"] == "1-8-R") self.assertIsNone(r_cell["stalled_appends_total"]) + def test_ratios_cover_both_controls_and_restart_cells_report_recovery(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + restart = {"restart": True, "restart_scraped": True, "restart_seconds": 3.5, "indexed_before_kill": 117, + "min_objects": 6000, + "restart_query_seconds": 4.25, "restart_count": 6000, + "server_cpu_seconds": 122.0, "host_cpu_busy": 0.5, "host_cpus": 8} + self.write_cell(root, "1-8-R", 100.0) + self.write_cell(root, "1-8-R0", 125.0) + self.write_cell(root, "1-8-D", 50.0, metrics=ACK_BUCKETS, env={"RUST_LOG": "info"}, **restart) + self.write_cell(root, "1-8-E", 75.0) + (root / "cells" / "1-8-D" / "metrics-restart.prom").write_text( + "greptime_logstore_object_store_wal_recovery_seconds_sum 1.5\n" + "greptime_logstore_object_store_wal_recovery_seconds_count 1\n" + "greptime_logstore_object_store_wal_recovered_objects_total 118\n") + rows = [summary.load_cell(path) for path in sorted((root / "cells").iterdir())] + pairs = {p["ratio"]: p["geomean"] for p in summary.ratios(rows)} + self.assertEqual(pairs, {"D/R": 0.5, "D/R0": 0.4, "E/R": 0.75, "E/R0": 0.6}) + cell = next(r for r in rows if r["id"] == "1-8-D") + self.assertEqual(cell["restart"], {"wall_seconds": 3.5, "query_seconds": 4.25, "rows": 6000, + "recovery_seconds": 1.5, "recovered_objects": 118, + "indexed_before_kill": 117, "min_objects": 6000, "replay": None}) + self.assertEqual(cell["server_cpu_cores"], 2.0) + output = summary.render({}, rows, summary.ratios(rows)) + self.assertIn("| 1-8-D | RUST_LOG=info | 120 | 117 (under 6,000) | 3.50 | 4.25 | 1.500 | 118 " + "| 6,000 |", output) + self.assertIn("| 2.00 | 50% of 8 |", output) + self.assertIn("| 8 | E/R0 | 0.600 | 0.600 |", output) + + def test_traffic_idle_crash_and_pinned_cells_are_summarized(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + traffic = ("greptime_logstore_object_store_wal_created_objects_total 600\n" + "greptime_logstore_object_store_wal_object_bytes_sum 3000000000\n" + "greptime_logstore_object_store_wal_object_bytes_count 600\n" + 'opendal_operation_duration_seconds_count{scheme="s3",operation="write"} 650\n' + 'opendal_operation_duration_seconds_count{scheme="s3",operation="read"} 40\n') + # The final scrape follows a 180 s idle window that created 2 objects of 1 MB. + self.write_cell(root, "1-32-E", 90.0, idle_seconds=180, group="", metrics=traffic.replace( + " 600\n", " 602\n", 1).replace("3000000000", "3002000000").replace("count 600", "count 602") + .replace("} 650", "} 652")) + (root / "cells" / "1-32-E" / "metrics-idle.prom").write_text(traffic.replace("} 650", "} 648")) + cell = summary.load_cell(root / "cells" / "1-32-E") + self.assertEqual((cell["objects_per_second"], cell["mean_object_bytes"], cell["sealed_on_size"]), + (10.0, 5e6, False)) + self.assertAlmostEqual(cell["wal_mb_per_second"], 50.0) + self.assertEqual(cell["idle"], {"seconds": 180, "creates": 2, "operations": {"write": 4}}) + path = root / "cells" / "pin-1-8-D" + path.mkdir(parents=True) + (path / "cell.json").write_text(json.dumps({ + "id": "pin-1-8-D", "group": "pin", "mode": "crash", "round": 1, "concurrency": 8, "config": "D", + "wal": bench.WAL_CONFIGS["D"], "exit_code": 0, "metrics_scraped": True})) + (path / "cycles.jsonl").write_text(json.dumps({ + "cycle": 1, "kill_after_seconds": 540.0, "load_seconds": 540.0, "cycle_acked_batches": 4000, + "cycle_present_batches": 4003, "cycle_lost_batches": 0, "cycle_lost_rows": 0, "cycle_extra_batches": 3, + "in_flight": 8, "total_found_batches": 4003, "total_acked_batches": 4000, "total_lost_batches": 0, + "restart_seconds": 4.5, "pin_rows": 1, "passed": True, "metrics_complete": True, + "kill_errors": [], "verify_errors": []}) + "\n") + (path / "metrics-start-1.prom").write_text("greptime_logstore_object_store_wal_created_objects_total 1\n") + (path / "metrics-cycle-1.prom").write_text( + "greptime_logstore_object_store_wal_indexed_objects 5400\n" + "greptime_logstore_object_store_wal_created_objects_total 5401\n" + "greptime_logstore_object_store_wal_object_bytes_sum 10800000000\n" + "greptime_logstore_object_store_wal_object_bytes_count 5400\n" + 'opendal_operation_duration_seconds_count{scheme="s3",operation="write"} 9000\n' + 'opendal_operation_duration_seconds_count{scheme="fs",operation="write"} 70\n') + (path / "metrics-restart-1.prom").write_text( + "greptime_logstore_object_store_wal_recovery_seconds_sum 2.5\n" + "greptime_logstore_object_store_wal_recovery_seconds_count 1\n" + "greptime_logstore_object_store_wal_recovered_objects_total 5399\n") + (path / "objects.tsv").write_text("100 300 0 301\n130 600 0 601\n160 missing\n") + rows = [summary.load_cell(p) for p in sorted((root / "cells").iterdir())] + output = summary.render({}, rows, summary.ratios(rows)) + self.assertIn("| 1-32-E | 50.0 | 10.0 | 4.77 | no | read 40, write 648 | — | 180 | 2 | write 4 |", output) + self.assertIn("| pin-1-8-D | 1 | 540.0 | 4,000 | 4,003 | 0 | 0 | 3 | 8 | 4,003 | 4,000 | 0 | 5,400 " + "| 4.50 | 2.500 | 5,399 | 1 | pass |", output) + self.assertIn("| pin-1-8-D | 1 | 540.0 | 20.0 | 10.0 | 1.91 | no | write 9,000 | — |", output) + self.assertIn("| pin-1-8-D | 300, 600, missing | 1 | 0 | 601 |", output) + + def test_partitioned_load_cells_report_client_latency_entries_and_regions(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + metrics = (ACK_BUCKETS + "greptime_logstore_object_store_wal_object_entries_sum 4000\n" + "greptime_logstore_object_store_wal_object_entries_count 500\n" + "greptime_logstore_object_store_wal_object_bytes_count 500\n") + self.write_cell(root, "1-8-D", 80.0, metrics=metrics, partitions=8) + path = root / "cells" / "1-8-D" + transcript = [{"step": "create_table", "duration_seconds": 9.0}] + transcript += [{"step": "grpc_write_batch", "status": "success", "duration_seconds": d / 1000} + for d in range(1, 101)] + (path / "greptimedb-1" / "load-transcript.jsonl").write_text( + "".join(json.dumps(r) + "\n" for r in transcript)) + (path / "regions.json").write_text(json.dumps({"output": [{"records": {"rows": [ + [f"p{i}", 4398046511104 + i, 750 + i, 1000] for i in range(8)]}}]})) + cell = summary.load_cell(path) + self.assertEqual((cell["client_p50"], cell["client_p99"]), (0.05, 0.099)) + self.assertEqual((cell["entries_per_object"], cell["appends_per_object"]), (8.0, 0.2)) + self.assertEqual(cell["region_rows"], [750 + i for i in range(8)]) + output = summary.render({}, [cell], []) + self.assertIn("| 1-8-D | 50.0 / 99.0 | 8.0 | 0.20 | 8 | 750 / 757 |", output) + + def test_restarts_report_what_each_region_replayed(self): + log = ("INFO mito2::region::opener: Replay WAL for region: 4398046511104(1024, 0), provider: ObjectStore(" + 'ObjectStoreProvider { region_id: 4398046511104(1024, 0), prefix: "wal" }), rows recovered: 0, ' + "replay from entry id: 1048578, last entry id: 1048577, total timeseries replayed: 0, elapsed: 60.581µs\n" + "INFO mito2::region::opener: Replay WAL for region: 4402341478400(1025, 0), provider: ObjectStore(" + 'ObjectStoreProvider { region_id: 4402341478400(1025, 0), prefix: "wal" }), rows recovered: 12625000, ' + "replay from entry id: 4853858306, last entry id: 8164212737, total timeseries replayed: 82, " + "elapsed: 95.150155247s\n" + "INFO mito2::region::opener: Replay WAL for region: 4402341478401(1025, 1), provider: ObjectStore(" + 'ObjectStoreProvider { region_id: 4402341478401(1025, 1), prefix: "wal" }), rows recovered: 5, ' + "replay from entry id: 8164212736, last entry id: 8164212737, total timeseries replayed: 1, " + "elapsed: 12.5ms\n") + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) + (path / "server-restart.log").write_text(log) + metrics = summary.parse_metrics("greptime_logstore_object_store_wal_read_seconds_sum 95.15\n" + "greptime_logstore_object_store_wal_read_seconds_count 3\n") + result = summary.replay(path / "server-restart.log", metrics) + self.assertEqual((result["rows"], result["object_span"], result["read_seconds"], result["reads"]), + (12_625_005, 3158 + 1, 95.15, 3)) + self.assertAlmostEqual(result["seconds"], 95.150155247) + self.assertEqual([(r["region"], r["rows"], r["object_span"]) for r in result["regions"]], + [(4398046511104, 0, 0), (4402341478400, 12_625_000, 3158), (4402341478401, 5, 1)]) + self.assertEqual((result["regions"][1]["first_entry_id"], result["regions"][1]["last_entry_id"]), + (4853858306, 8164212737)) + self.assertAlmostEqual(result["regions"][0]["seconds"], 60.581e-6) + self.assertEqual(summary.replayed(result["regions"]), + "4402341478400: 12,625,000 rows, object sequences 4,629–7,786, 95.15 s; " + "4402341478401: 5 rows, object sequences 7,786–7,786, 0.01 s") + self.assertIsNone(summary.replay(path / "absent.log", metrics)) + raft = log.replace("ObjectStore(", "RaftEngine(") + (path / "server-restart.log").write_text(raft) + self.assertIsNone(summary.replay(path / "server-restart.log", metrics)["object_span"]) + # A region flushed in object 100 that next wrote in object 200 reads one object but spans 101. + gap = log.replace("replay from entry id: 4853858306", f"replay from entry id: {(100 << 20) + 2}") \ + .replace("last entry id: 8164212737", f"last entry id: {(200 << 20) + 1}") + (path / "server-restart.log").write_text(gap) + self.assertEqual(summary.replay(path / "server-restart.log", metrics)["regions"][1]["object_span"], 101) + + def test_a_drop_cell_reports_the_objects_collected_after_its_restart(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + path = root / "cells" / "drop-1-8-D" + (path / "dropped").mkdir(parents=True) + (path / "cell.json").write_text(json.dumps({ + "id": "drop-1-8-D", "group": "drop", "mode": "drop", "round": 1, "concurrency": 8, "config": "D", + "wal": bench.WAL_CONFIGS["D"], "exit_code": 0, "restart_seconds": 2.5, "restart_scraped": True})) + write_lines(path / "dropped" / "acked-0.log", [0, 8, 16]) + write_lines(path / "dropped" / "acked-1.log", [1, 9]) + (path / "metrics-written.prom").write_text("greptime_logstore_object_store_wal_indexed_objects 300\n") + (path / "metrics-restart.prom").write_text( + "greptime_logstore_object_store_wal_recovered_objects_total 280\n" + "greptime_logstore_object_store_wal_deleted_objects_total 279\n") + (path / "collected.tsv").write_text("100 281\n101 120\n104 2\n") + cell = summary.load_cell(path) + self.assertEqual(cell["drop"], {"acked_batches": 5, "indexed_before_drop": 300, "restart_seconds": 2.5, + "recovered_objects": 280, "indexed_at_end": 2, + "deleted_after_restart": 279, "seconds_to_collect": 4}) + output = summary.render({}, [cell], []) + self.assertIn("| drop-1-8-D | 5 | 300 | 2.50 | 280 | 2 | 279 | 4 | collected |", output) + def test_untouched_counters_are_zero_only_for_an_active_object_store_wal(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -195,29 +756,104 @@ class SummaryTest(unittest.TestCase): self.assertIn("1-1-D (failed, exit 1)", summary.render({}, rows, summary.ratios(rows))) + def test_images_are_compared_with_the_first_in_every_round(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + images = [{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} for label in ("base", "pr1")] + plan = bench.build_plan("8", "6000", "1000", "2", smoke=False, configs="R,D", images=json.dumps(images), + warmup=True) + (root / "plan.json").write_text(json.dumps(plan)) + (root / "run.json").write_text(json.dumps({"rows": 6000, "batch_rows": 1000, "greptimedb_images": [ + i | {"image": f"greptime/greptimedb-dev@sha256:{n * 64}", "version": "1.4.0", "commit": f"{n * 7}"} + for i, n in zip(images, "ab")]})) + # The warm-up cells are far off and must not move any ratio. + speeds = {"0-8-R-base": 10.0, "0-8-R-pr1": 999.0, + "1-8-R-base": 100.0, "1-8-R-pr1": 100.0, "1-8-D-base": 50.0, "1-8-D-pr1": 40.0, + "2-8-D-pr1": 66.0, "2-8-D-base": 55.0, "2-8-R-pr1": 100.0, "2-8-R-base": 110.0} + for cell_id, speed in speeds.items(): + self.write_cell(root, cell_id, speed, warmup=cell_id.startswith("0-"), + host_load_start=[0.5, 0.4, 0.3], host_load_end=[1.25, 0.5, 0.3]) + with open(root / "cells" / "2-8-D-pr1" / "greptimedb-1" / "load-transcript.jsonl", "w") as out: + for seconds in (0.2, 0.3): + out.write(json.dumps({"step": "grpc_write_batch", "status": "success", "duration_seconds": seconds}) + "\n") + output = subprocess.run([sys.executable, str(SCRIPTS / "oswal-s3-bench-summary.py"), "--root", str(root)], + check=True, capture_output=True, text=True).stdout + result = json.loads((root / "summary.json").read_text()) + compared = {(c["config"], c["image"]): c for c in result["image_ratios"]} + self.assertEqual(list(compared), [("R", "base"), ("R", "pr1"), ("D", "base"), ("D", "pr1")]) + durable = compared[("D", "pr1")] + self.assertEqual(durable["ratio_to"], "base") + self.assertEqual(durable["per_round"], [0.8, 1.2]) + self.assertAlmostEqual(durable["geomean"], 0.96 ** 0.5) + self.assertEqual((durable["min"], durable["max"]), (0.8, 1.2)) + self.assertEqual([r["rows_per_second"] for r in durable["rounds"]], [40.0, 66.0]) + self.assertEqual((durable["rounds"][1]["client_p50"], durable["rounds"][1]["client_p99"]), (0.2, 0.3)) + self.assertEqual(durable["rounds"][0]["host_load_end"], [1.25, 0.5, 0.3]) + self.assertEqual((compared[("D", "base")]["per_round"], compared[("D", "base")]["geomean"]), (None, None)) + # The ratios to the Raft Engine control stay within each image. + pairs = {(r["image"], r["ratio"]): r["per_round"] for r in result["ratios"]} + self.assertEqual(pairs, {("base", "D/R"): [0.5, 0.5], ("pr1", "D/R"): [0.4, 0.66]}) + self.assertIn("| base | `greptime/greptimedb-dev@sha256:" + "a" * 64 + "` | 1.4.0 | `aaaaaaa` |", output) + self.assertIn("| — | 8 | D | pr1 | 40; 66 | — / —; 200.0 / 300.0 | 0.800, 1.200 | 0.980 | 0.800 | 1.200 |", output) + self.assertIn("| — | 8 | D | base | 50; 55 | — / —; — / — | — | — | — | — |", output) + self.assertIn("| 0-8-R-pr1 (warm-up) |", output) + self.assertIn("| `/dev/nvme1n1` | 0.50 / 1.25 |", output) + self.assertIn("| — | pr1 | 8 | D/R | 0.400, 0.660 | 0.514 |", output) + + FAKE_DOCKER = r"""#!/usr/bin/env bash printf '%s\n' "docker $*" >> "$FAKE_LOG" case "$1" in run) - for arg in "$@"; do [[ "$arg" != -d ]] || { echo "keys=${GREPTIMEDB_STANDALONE__STORAGE__ACCESS_KEY_ID:+set}${GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY:+set}" >> "$FAKE_LOG"; echo fakecid; exit 0; }; done + for arg in "$@"; do [[ "$arg" != -d ]] || { echo $(($(cat "$FAKE_LOG.runs" 2>/dev/null || echo 0) + 1)) > "$FAKE_LOG.runs"; rm -f "${FAKE_KILL_FLAG:-/nonexistent}"; echo "keys=${GREPTIMEDB_STANDALONE__STORAGE__ACCESS_KEY_ID:+set}${GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY:+set}" >> "$FAKE_LOG"; echo fakecid; exit 0; }; done for arg in "$@"; do case "$arg" in type=bind,src=*,dst=/results) results=${arg#type=bind,src=}; results=${results%,dst=/results} ;; esac done [[ "$FAKE_LOAD_RC" == 0 ]] || exit "$FAKE_LOAD_RC" mkdir -p "$results/greptimedb-1" - echo "{\"count_before\": 0, \"count_after\": $FAKE_ROWS, \"load_seconds\": 2, \"rows_per_second\": 3000}" \ + echo "{\"count_before\": 0, \"count_after\": $FAKE_ROWS, \"load_seconds\": 2, \"rows_per_second\": 3000, \"table\": \"agent_observations\"}" \ > "$results/greptimedb-1/load-result.json" ;; + kill) touch "${FAKE_KILL_FLAG:-/dev/null}" ;; + rm) [[ -z "${FAKE_LOAD_AFTER:-}" ]] || echo "$FAKE_LOAD_AFTER" > "$OSWAL_LOADAVG" ;; inspect) echo true ;; - exec) echo 'version: 1.4.0-test' ;; - logs) echo 'server log' ;; + exec) [[ "$*" == *cpu.stat* ]] && echo 'usage_usec 5000000' || echo 'version: 1.4.0-test' ;; + logs) echo 'server log' + # The rows a server replayed are its start number, so each log can be told apart. + [[ -z "${FAKE_REPLAY:-}" ]] || echo "INFO Replay WAL for region: 4398046511104(1024, 0), provider: ObjectStore(x), rows recovered: $(cat "$FAKE_LOG.runs"), replay from entry id: 1048577, last entry id: 2097153, total timeseries replayed: 1, elapsed: 1.5s" ;; esac """ FAKE_CURL = r"""#!/usr/bin/env bash printf '%s\n' "curl $*" >> "$FAKE_LOG" +[[ "$*" != *"--max-time 10"* || "${FAKE_SAMPLE_RC:-0}" == 0 ]] || exit "$FAKE_SAMPLE_RC" for arg in "$@"; do + [[ "$arg" != *"GROUP BY seq"* ]] || { python3 -c 'import collections, json, sys +counts = collections.Counter(int(l) for l in open(sys.argv[1]) if l.strip()) +print(json.dumps({"output": [{"records": {"rows": [[k, v * int(sys.argv[2])] for k, v in counts.items()]}}]}))' \ + "$FAKE_FOUND" "$FAKE_BATCH_ROWS"; exit 0; } +done +for arg in "$@"; do + [[ "$arg" != *region_statistics* ]] || { python3 -c 'import json, sys +print(json.dumps({"output": [{"records": {"rows": [[f"p{i}", i, int(n), 1] for i, n in enumerate(sys.argv[1].split(","))]}}]}))' \ + "$FAKE_REGION_ROWS"; exit 0; } +done +for arg in "$@"; do + [[ "$arg" != */v1/sql* ]] || { echo "{\"output\": [{\"records\": {\"rows\": [[${FAKE_COUNT:-6000}]]}}]}"; exit 0; } [[ "$arg" != */metrics ]] || { [[ "$FAKE_METRICS_RC" == 0 ]] || exit "$FAKE_METRICS_RC" - printf '%s\n' 'greptime_logstore_object_store_wal_poisoned_total 1' 'other_total 2'; exit 0; + printf '%s\n' 'greptime_logstore_object_store_wal_poisoned_total 1' 'other_total 2' + [[ -z "${FAKE_FULL_METRICS:-}" ]] || printf '%s\n' \ + "greptime_logstore_object_store_wal_indexed_objects $FAKE_FULL_METRICS" \ + 'greptime_logstore_object_store_wal_recovery_seconds_count 1' \ + "greptime_logstore_object_store_wal_recovered_objects_total ${FAKE_RECOVERED:-0}" + # Every server counts the objects it collects from its start: with FAKE_COLLECTING one more on + # every scrape; with FAKE_FIRST_COLLECTS_AFTER the first server collects one object once it has + # been scraped that many times. + run=$(cat "$FAKE_LOG.runs") n=$(($(cat "$FAKE_LOG.scrapes.$run" 2>/dev/null || echo 0) + 1)) + echo "$n" > "$FAKE_LOG.scrapes.$run" + [[ -z "${FAKE_COLLECTING:-}" ]] || echo "greptime_logstore_object_store_wal_deleted_objects_total $((n - 1))" + [[ -z "${FAKE_FIRST_COLLECTS_AFTER:-}" || "$run" != 1 || "$n" -le "$FAKE_FIRST_COLLECTS_AFTER" ]] \ + || echo 'greptime_logstore_object_store_wal_deleted_objects_total 1' + exit 0; } done """ @@ -225,21 +861,24 @@ done class CellLifecycleTest(unittest.TestCase): # the run script against fake docker and curl - def run_cells(self, root, load_rc, metrics_rc=0): + def run_cells(self, root, load_rc, metrics_rc=0, restart=False, count=6000, groups="", images="", checked=None): bin_dir = root / "bin" bin_dir.mkdir() for name, body in (("docker", FAKE_DOCKER), ("curl", FAKE_CURL)): (bin_dir / name).write_text(body) (bin_dir / name).chmod(0o755) - plan = bench.build_plan("8", "6000", "1000", "1", smoke=False) - plan["cells"] = [cell for cell in plan["cells"] if cell["config"] != "R"] + plan = bench.build_plan("8", "6000", "1000", "1", smoke=False, configs="D,E", restart=restart, + extra_env="RUST_LOG=info", groups=groups, images=images) + if checked is None: + checked = [{"label": label, "image": f"greptimedb{label}@sha256:" + "a" * 64} for label in plan["images"]] (root / "results").mkdir() (root / "results" / "plan.json").write_text(json.dumps(plan)) env = os.environ | { "PATH": f"{bin_dir}:{os.environ['PATH']}", "FAKE_LOG": str(root / "fake.log"), - "FAKE_LOAD_RC": str(load_rc), "FAKE_METRICS_RC": str(metrics_rc), "FAKE_ROWS": "6000", "PLAN": str(root / "results" / "plan.json"), + "FAKE_LOAD_RC": str(load_rc), "FAKE_METRICS_RC": str(metrics_rc), "FAKE_ROWS": "6000", "FAKE_COUNT": str(count), + "PLAN": str(root / "results" / "plan.json"), "RESULTS": str(root / "results"), "CORPUS_DIR": str(root / "corpus"), "OWNER": "o11ybench-oswal-1-1", - "RUNTIME_IMAGE": "runtime@sha256:" + "b" * 64, "GREPTIMEDB_IMAGE": "greptimedb@sha256:" + "a" * 64, + "RUNTIME_IMAGE": "runtime@sha256:" + "b" * 64, "GREPTIMEDB_IMAGES": json.dumps(checked), "RUN_PREFIX": "oswal-bench/1-1", "WAL_ROOT": str(root / "wal"), "S3_BUCKET": "bucket", "S3_REGION": "ap-southeast-1", "AWS_ACCESS_KEY_ID": "AKIAFAKEKEY", "AWS_SECRET_ACCESS_KEY": "fake-secret-value", } @@ -302,6 +941,9 @@ class CellLifecycleTest(unittest.TestCase): # the run script against fake docke self.assertIn("GREPTIMEDB_STANDALONE__STORAGE__ROOT=oswal-bench/1-1/1-8-D", config) self.assertIn("GREPTIMEDB_STANDALONE__WAL__ACK_MODE=durable", config) self.assertIn("GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY=", config) + self.assertIn("RUST_LOG=info", config) + self.assertIn("-e RUST_LOG=info ", log) + self.assertNotIn("docker kill", log) recorded = json.loads((cell / "cell.json").read_text()) self.assertEqual((recorded["exit_code"], recorded["version"], recorded["metrics_scraped"]), (0, "1.4.0-test", True)) @@ -313,6 +955,410 @@ class CellLifecycleTest(unittest.TestCase): # the run script against fake docke self.assertNotEqual(leaked.returncode, 0) self.assertFalse((root / "results").exists()) + def test_a_restart_cell_kills_restarts_and_counts_the_rows(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, log, _ = self.run_cells(root, load_rc=0, restart=True) + self.assertEqual(cells.returncode, 0, cells.stdout + cells.stderr) + self.assertEqual(log.count("docker kill --signal KILL fakecid"), 2) + self.assertEqual(log.count("docker stop -t 120 fakecid"), 2) + self.assertIn("--name o11ybench-oswal-1-1-1-8-D-restart", log) + self.assertIn("sql=SELECT count(*) FROM agent_observations", log) + cell = root / "results" / "cells" / "1-8-D" + recorded = json.loads((cell / "cell.json").read_text()) + self.assertEqual((recorded["exit_code"], recorded["restart_count"], recorded["restart_scraped"]), (0, 6000, True)) + self.assertIsInstance(recorded["restart_seconds"], float) + config = (cell / "config.env").read_text().splitlines() + self.assertIn("GREPTIMEDB_STANDALONE__STORAGE__DATA_HOME=/oswal_home", config) + self.assertIn("GREPTIMEDB_STANDALONE__WAL__FLUSH_INTERVAL=10ms", config) + self.assertIn(f"type=bind,src={root}/wal/1-8-D-home,dst=/oswal_home", log) + for name in ("metrics.prom", "metrics-restart.prom", "server.log", "server-restart.log"): + self.assertTrue((cell / name).exists(), name) + + def test_an_idle_window_scrapes_before_and_after_it(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + groups = json.dumps([{"name": "idle", "configs": "D", "idle_seconds": 1}]) + cells, _, log, _ = self.run_cells(root, load_rc=0, groups=groups) + self.assertEqual(cells.returncode, 0, cells.stdout + cells.stderr) + cell = root / "results" / "cells" / "idle-1-8-D" + self.assertTrue((cell / "metrics-idle.prom").exists()) + self.assertTrue((cell / "metrics.prom").exists()) + self.assertEqual(log.count("/metrics"), 2) + self.assertIn(f"type=bind,src={root}/corpus/corpus-6000,dst=/work/corpus,readonly", log) + + def run_crash(self, root, metrics_rc=0, pin=False, sample_rc=0, full_metrics="5", groups=None, env=None, + status=200): + """A crash cell, or another group that writes through the crash writer, against fake docker + and a stub SQL endpoint that the fake kill takes down.""" + import http.server + import threading + flag, found = root / "killed", root / "found.log" + found.write_text("") + lock = threading.Lock() + + class Handler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + body = urllib.parse.unquote_plus(self.rfile.read(int(self.headers["Content-Length"])).decode()) + if flag.exists(): + self.close_connection = True + self.connection.close() + return + seq = re.search(r"VALUES [(](\d+),", body).group(1) + with lock, open(found, "a") as out: + out.write(seq + "\n") + self.send_response(status) + self.end_headers() + + def log_message(self, *args): + pass + + class Server(http.server.ThreadingHTTPServer): + request_queue_size = 64 # all eight writers connect at once + + server = Server(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + groups = groups or json.dumps([{"name": "crash", "configs": "D", "mode": "crash", "cycles": 2, + "kill_after": [1, 1], "pin": pin}]) + old = os.environ.copy() + os.environ.update({"OSWAL_SQL_URL": f"http://127.0.0.1:{server.server_port}/v1/sql", + "FAKE_KILL_FLAG": str(flag), "FAKE_FOUND": str(found), "FAKE_BATCH_ROWS": "1000", + "OSWAL_SAMPLE_SECONDS": "0.2", "FAKE_SAMPLE_RC": str(sample_rc), + "FAKE_FULL_METRICS": full_metrics} | (env or {})) + try: + return self.run_cells(root, load_rc=0, metrics_rc=metrics_rc, groups=groups, + count=(env or {}).get("FAKE_COUNT", 6000)) + finally: + os.environ.clear() + os.environ.update(old) + server.shutdown() + server.server_close() + + def test_a_crash_cell_kills_restarts_and_checks_every_cycle(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, log, _ = self.run_crash(root, env={"FAKE_REPLAY": "1"}) + self.assertEqual(cells.returncode, 0, cells.stdout + cells.stderr) + cell = root / "results" / "cells" / "crash-1-8-D" + # Server 1 is loaded and killed; servers 2 and 3 are the restarts of cycles 1 and 2. + replays = [c["replay"]["rows"] for c in summary.load_cell(cell)["cycles"]] + self.assertEqual(replays, [2, 3]) + cycles = [json.loads(line) for line in (cell / "cycles.jsonl").read_text().splitlines()] + self.assertEqual([c["cycle"] for c in cycles], [1, 2]) + for c in cycles: + self.assertTrue(c["passed"] and c["metrics_complete"], c) + self.assertEqual((c["cycle_lost_batches"], c["writers_at_kill"]), (0, 8)) + self.assertGreater(c["cycle_acked_batches"], 0) + self.assertEqual(log.count("docker kill --signal KILL fakecid"), 2) + self.assertEqual(log.count("docker wait fakecid"), 2) + self.assertIn("GREPTIMEDB_STANDALONE__STORAGE__DATA_HOME=/oswal_home", (cell / "config.env").read_text()) + for name in ("metrics-start-1.prom", "metrics-cycle-2.prom", "metrics-restart-2.prom", "server.log", + "server-restart-1.log", "server-restart-2.log"): + self.assertTrue((cell / name).exists(), name) + + def test_failed_pinned_samples_are_missing_not_zero(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, log, _ = self.run_crash(root, pin=True, sample_rc=28, env={"FAKE_COUNT": "1"}) + self.assertEqual(cells.returncode, 0, cells.stdout + cells.stderr) + cell = root / "results" / "cells" / "crash-1-8-D" + samples = (cell / "objects.tsv").read_text().splitlines() + self.assertTrue(samples and all(line.split()[1:] == ["missing"] for line in samples), samples) + self.assertIn("sql=INSERT INTO pin VALUES (1, 0)", log) + self.assertIn("WITH ('auto_flush_interval' = '1d')", urllib.parse.unquote_plus(log)) + self.assertEqual(len((cell / "cycles.jsonl").read_text().splitlines()), 2) + + def test_a_pinned_cell_fails_when_objects_the_pin_held_are_collected(self): + # Each server collects two objects before its kill and no periodic sample is taken. With one + # object below the pin the first server already exceeds it; with two, only both together do. + for below_pin, failed_cycle in (("1", 1), ("2", 2)): + with self.subTest(below_pin=below_pin), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, _, _ = self.run_crash(root, pin=True, full_metrics=below_pin, env={ + "FAKE_COUNT": "1", "FAKE_COLLECTING": "1", "OSWAL_SAMPLE_SECONDS": "30"}) + cell = root / "results" / "cells" / "crash-1-8-D" + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], 11, + cells.stdout + cells.stderr) + self.assertFalse((cell / "objects.tsv").exists()) + self.assertEqual(len((cell / "cycles.jsonl").read_text().splitlines()), failed_cycle) + self.assertIn(f"cycle {failed_cycle} collected WAL objects the pin held", cells.stdout + cells.stderr) + + def test_a_pinned_cell_may_collect_the_objects_below_its_pin_late(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + # The start object, the one object indexed before the pin, goes only after the first cycle's + # baseline: the pin scrape and that baseline are the first server's first two scrapes. + cells, _, _, _ = self.run_crash(root, pin=True, full_metrics="1", env={ + "FAKE_COUNT": "1", "FAKE_FIRST_COLLECTS_AFTER": "2", "OSWAL_SAMPLE_SECONDS": "30"}) + cell = root / "results" / "cells" / "crash-1-8-D" + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], 0, cells.stdout + cells.stderr) + self.assertNotIn("deleted_objects_total", (cell / "metrics-start-1.prom").read_text()) + self.assertIn("deleted_objects_total 1", (cell / "metrics-cycle-1.prom").read_text()) + self.assertEqual(len((cell / "cycles.jsonl").read_text().splitlines()), 2) + + def test_a_durable_pinned_cycle_fails_unless_its_pin_row_is_there_once(self): + for count in ("0", "2"): + with self.subTest(count=count), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.run_crash(root, pin=True, env={"FAKE_COUNT": count}) + cell = root / "results" / "cells" / "crash-1-8-D" + cycle = json.loads((cell / "cycles.jsonl").read_text().splitlines()[0]) + self.assertEqual((cycle["passed"], cycle["pin_rows"]), (False, int(count))) + self.assertIn(f"the pin table holds {count} rows, not 1", cycle["verify_errors"]) + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], 6) + + def test_a_scrape_without_the_wal_samples_is_incomplete(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, _, _ = self.run_crash(root, full_metrics="") + self.assertNotEqual(cells.returncode, 0) + cell = root / "results" / "cells" / "crash-1-8-D" + cycle = json.loads((cell / "cycles.jsonl").read_text().splitlines()[0]) + self.assertEqual((cycle["passed"], cycle["metrics_complete"]), (True, False)) + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], 4) + + def test_a_restart_cell_must_leave_its_minimum_objects_indexed(self): + for minimum, exit_code in ((5, 0), (6, 7)): + with self.subTest(minimum=minimum), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + groups = json.dumps([{"name": "rec", "configs": "D", "restart": True, "min_objects": minimum}]) + old = os.environ.copy() + os.environ["FAKE_FULL_METRICS"] = "5" + try: + cells, _, _, _ = self.run_cells(root, load_rc=0, groups=groups) + finally: + os.environ.clear() + os.environ.update(old) + recorded = json.loads((root / "results" / "cells" / "rec-1-8-D" / "cell.json").read_text()) + self.assertEqual((recorded["exit_code"], recorded["indexed_before_kill"], recorded["restart_count"]), + (exit_code, 5, 6000)) + + def test_a_partitioned_cell_creates_its_table_and_checks_every_region_took_rows(self): + for region_rows, exit_code in (("750,750,750,750,750,750,750,750", 0), ("1000,1000,1000,0,1000,1000,1000,0", 8)): + with self.subTest(region_rows=region_rows), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "corpus" / "corpus-6000").mkdir(parents=True) + (root / "corpus" / "corpus-6000" / "summary.json").write_text(json.dumps( + {"start": "2026-07-10T00:00:00Z", "end": "2026-07-10T01:00:00Z"})) + groups = json.dumps([{"name": "multi", "configs": "D", "partitions": 8}]) + old = os.environ.copy() + os.environ["FAKE_REGION_ROWS"] = region_rows + try: + cells, _, log, _ = self.run_cells(root, load_rc=0, groups=groups) + finally: + os.environ.clear() + os.environ.update(old) + cell = root / "results" / "cells" / "multi-1-8-D" + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], exit_code, + cells.stdout + cells.stderr) + self.assertIn("PARTITION ON COLUMNS (`event_time`)", urllib.parse.unquote_plus(log)) + self.assertIn(f"type=bind,src={SCRIPTS}/oswal-s3-bench-load.sh,dst=/work/oswal-s3-bench-load.sh,readonly", + log) + self.assertIn(" /work/oswal-s3-bench-load.sh", log) + self.assertNotIn("greptimedb.sh load", log) + self.assertEqual(len(json.loads((cell / "regions.json").read_text())["output"][0]["records"]["rows"]), 8) + + def test_a_drop_cell_restarts_and_waits_for_the_dropped_objects_to_go(self): + for indexed, recovered, exit_code in (("2", "0", 0), ("50", "200", 0), ("50", "60", 9)): + with self.subTest(indexed=indexed, recovered=recovered), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + groups = json.dumps([{"name": "drop", "mode": "drop", "configs": "D", "write_seconds": 1}]) + cells, _, log, _ = self.run_crash(root, full_metrics=indexed, groups=groups, + env={"FAKE_RECOVERED": recovered, "OSWAL_COLLECT_SECONDS": "1"}) + cell = root / "results" / "cells" / "drop-1-8-D" + recorded = json.loads((cell / "cell.json").read_text()) + self.assertEqual(recorded["exit_code"], exit_code, cells.stdout + cells.stderr) + self.assertIsInstance(recorded["restart_seconds"], float) + sql = urllib.parse.unquote_plus(log) + for statement in ("CREATE TABLE dropped", "CREATE TABLE kept", "DROP TABLE dropped", + "ADMIN FLUSH_TABLE('kept')", "INSERT INTO kept VALUES (-1, "): + self.assertIn(statement, sql) + self.assertLess(sql.index("DROP TABLE dropped"), log.index("docker kill --signal KILL fakecid")) + self.assertEqual(log.count("docker kill --signal KILL fakecid"), 1) + self.assertTrue(list((cell / "dropped").glob("acked-*.log"))) + self.assertTrue(list((cell / "kept").glob("acked-*.log"))) + for name in ("metrics-written.prom", "metrics-restart.prom", "server.log", + "server-restart.log", "collected.tsv"): + self.assertTrue((cell / name).exists(), name) + + def test_a_drop_cell_whose_writers_fail_stops_before_the_drop(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + groups = json.dumps([{"name": "drop", "mode": "drop", "configs": "D", "write_seconds": 1}]) + cells, _, log, _ = self.run_crash(root, full_metrics="2", groups=groups, status=503) + recorded = json.loads((root / "results" / "cells" / "drop-1-8-D" / "cell.json").read_text()) + self.assertEqual(recorded["exit_code"], 10, cells.stdout + cells.stderr) + self.assertNotIn("DROP TABLE", urllib.parse.unquote_plus(log)) + self.assertNotIn("docker kill", log) + + def test_a_crash_cycle_without_its_scrapes_fails_the_cell(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, _, _ = self.run_crash(root, metrics_rc=22) + self.assertNotEqual(cells.returncode, 0) + cell = root / "results" / "cells" / "crash-1-8-D" + cycle = json.loads((cell / "cycles.jsonl").read_text().splitlines()[0]) + self.assertEqual((cycle["passed"], cycle["metrics_complete"]), (True, False)) + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], 4) + + def test_a_restart_that_loses_rows_fails_the_cell(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, _, _ = self.run_cells(root, load_rc=0, restart=True, count=5999) + self.assertNotEqual(cells.returncode, 0) + recorded = json.loads((root / "results" / "cells" / "1-8-D" / "cell.json").read_text()) + self.assertEqual((recorded["exit_code"], recorded["restart_count"]), (5, 5999)) + + + def test_every_cell_starts_its_own_image_and_records_the_host_load(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "loadavg").write_text("0.50 0.40 0.30 1/200 99\n") + images = json.dumps([{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} + for label in ("base", "pr1")]) + old = os.environ.copy() + os.environ.update({"OSWAL_LOADAVG": str(root / "loadavg"), "FAKE_LOAD_AFTER": "1.25 1.00 0.75 1/200 99"}) + try: + cells, _, log, _ = self.run_cells(root, load_rc=0, images=images) + finally: + os.environ.clear() + os.environ.update(old) + self.assertEqual(cells.returncode, 0, cells.stdout + cells.stderr) + starts = [line for line in log.splitlines() if line.startswith("docker run -d")] + self.assertEqual([re.search(r"--name o11ybench-oswal-1-1-(\S+) ", line).group(1) for line in starts], + ["1-8-D-base", "1-8-D-pr1", "1-8-E-base", "1-8-E-pr1"]) + for line in starts: + label = line.split("--name ")[1].split()[0].rsplit("-", 1)[1] + self.assertIn(f" greptimedb{label}@sha256:{'a' * 64} standalone start", line) + self.assertIn(f"TARGET_IMAGE=greptimedbpr1@sha256:{'a' * 64}", log) + first, second = (json.loads((root / "results" / "cells" / i / "cell.json").read_text()) + for i in ("1-8-D-base", "1-8-D-pr1")) + self.assertEqual((first["image"], first["host_load_start"], first["host_load_end"]), + ("base", [0.5, 0.4, 0.3], [1.25, 1.0, 0.75])) + self.assertEqual(second["host_load_start"], [1.25, 1.0, 0.75]) + + def test_a_cell_without_a_checked_image_fails(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + images = json.dumps([{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} + for label in ("base", "pr1")]) + cells, _, log, _ = self.run_cells(root, load_rc=0, images=images, + checked=[{"label": "base", "image": "greptimedb@sha256:" + "a" * 64}]) + self.assertNotEqual(cells.returncode, 0) + self.assertIn("1-8-D-pr1: no image labelled 'pr1'", cells.stderr) + self.assertEqual(log.count("docker run -d"), 2) + + def test_check_images_pins_every_image_to_its_digest_and_records_its_version(self): + fake_docker = r"""#!/usr/bin/env bash +printf '%s\n' "docker $*" >> "$FAKE_LOG" +case "$1 $2" in + "image inspect") printf '[{"RepoDigests": ["%s@sha256:%064d"]}]\n' "${3%:*}" "$(printf '%s' "$3" | cksum | cut -d' ' -f1)" ;; + "run --rm") [[ "$*" == *nocommit* ]] || commit="commit: 0123456789abcdef$(printf '%s' "$4" | cksum | cut -c1-4)" + printf 'greptime\nbranch: main\n%s\nclean: true\nversion: 1.4.0\n' "${commit:-}" ;; +esac +""" + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + for name, body in (("docker", fake_docker), ("timeout", '#!/usr/bin/env bash\nshift 2\nexec "$@"\n')): + (root / name).write_text(body) + (root / name).chmod(0o755) + env = os.environ | {"PATH": f"{root}:{os.environ['PATH']}", "FAKE_LOG": str(root / "fake.log"), + "GITHUB_OUTPUT": str(root / "output")} + requested = [{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} for label in ("base", "pr1")] + done = subprocess.run(["bash", str(SCRIPTS / "oswal-s3-bench-run.sh"), "check-images", json.dumps(requested)], + env=env, capture_output=True, text=True, timeout=60) + self.assertEqual(done.returncode, 0, done.stderr) + name, _, value = (root / "output").read_text().strip().partition("=") + checked = json.loads(value) + self.assertEqual(name, "greptimedb_images") + self.assertEqual([(c["label"], c["requested"], c["version"]) for c in checked], + [(r["label"], r["image"], "1.4.0") for r in requested]) + for c in checked: + self.assertRegex(c["image"], r"^docker\.io/greptime/greptimedb-dev@sha256:[0-9]{64}$") + self.assertTrue(c["commit"].startswith("0123456789abcdef"), c) + self.assertIn(f"docker run --rm --entrypoint greptime {c['image']} --version", (root / "fake.log").read_text()) + self.assertNotEqual(checked[0]["image"], checked[1]["image"]) + requested[1]["image"] = "docker.io/greptime/nocommit:pr1" + done = subprocess.run(["bash", str(SCRIPTS / "oswal-s3-bench-run.sh"), "check-images", json.dumps(requested)], + env=env, capture_output=True, text=True, timeout=60) + self.assertNotEqual(done.returncode, 0) + self.assertIn("reports no version or commit", done.stderr) + + + def test_check_runtime_pins_a_named_digest_offline_and_resolves_only_a_tag(self): + digest = "sha256:" + "c" * 64 + fake_docker = f"""#!/usr/bin/env bash +printf '%s\\n' "docker $*" >> "$FAKE_LOG" +[[ "$1 $2 $3" == "buildx imagetools inspect" ]] || exit 9 +printf '{{"mediaType": "application/vnd.oci.image.index.v1+json", "digest": "{digest}", "size": 856}}\\n' +""" + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "docker").write_text(fake_docker) + (root / "docker").chmod(0o755) + env = os.environ | {"PATH": f"{root}:{os.environ['PATH']}", "FAKE_LOG": str(root / "fake.log"), + "GITHUB_OUTPUT": str(root / "output")} + repository, named = "public.ecr.aws/i8k6a5e1/greptime/o11ybenchruntime", "sha256:" + "d" * 64 + # Expected pin, or error, and whether the registry is asked. + for image, expected, asks in ((f"{repository}:77a84e1c8136", f"{repository}@{digest}", True), + (f"{repository}:77a84e1c8136@{named}", f"{repository}@{named}", False), + (f"{repository}@{named}", f"{repository}@{named}", False), + (f"{repository}:77a84e1c8136@sha256:abc", "not a full ECR Public", False), + (f"public.ecr.aws/@{named}", "not a full ECR Public", False), + (f"{repository}@@{named}", "not a full ECR Public", False), + (f"{repository}:bad tag", "not a full ECR Public", False), + ("docker.io/greptime/o11ybenchruntime:1", "must come from AWS ECR Public", False)): + with self.subTest(image=image): + (root / "output").write_text("") + (root / "fake.log").write_text("") + done = subprocess.run(["bash", str(SCRIPTS / "oswal-s3-bench-run.sh"), "check-runtime", image], + env=env, capture_output=True, text=True, timeout=60) + if expected.startswith(repository): + self.assertEqual(done.returncode, 0, done.stderr) + self.assertEqual((root / "output").read_text(), f"runtime_image={expected}\n") + else: + self.assertNotEqual(done.returncode, 0) + self.assertIn(expected, done.stderr) + self.assertEqual((root / "output").read_text(), "") + calls = (root / "fake.log").read_text().splitlines() + self.assertEqual(calls, [f"docker buildx imagetools inspect --format {{{{json .Manifest}}}} {image}"] + if asks else []) + + def test_the_benchmark_host_checks_the_runtime_contents_before_pulling_greptimedb(self): + # The fake runs the step's own check in a directory standing in for the image's working directory. + fake_docker = r"""#!/usr/bin/env bash +printf '%s\n' "docker $*" >> "$FAKE_LOG" +case "$1" in + run) cd "$FAKE_RUNTIME_ROOT" && exec sh -c "$7" ;; + image) echo '[{}]' ;; +esac +""" + loader, generator = "scripts/agent-observability/greptimedb.sh", "scripts/generate_agent_observations.py" + for present, passes in (((loader, generator), True), ((generator,), False), ((loader,), False)): + with self.subTest(present=present), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + for name, body in (("docker", fake_docker), ("timeout", '#!/usr/bin/env bash\nshift 2\nexec "$@"\n')): + (root / name).write_text(body) + (root / name).chmod(0o755) + for path in present: + (root / "image" / path).parent.mkdir(parents=True, exist_ok=True) + (root / "image" / path).write_text("") + runtime = "public.ecr.aws/i8k6a5e1/greptime/o11ybenchruntime@sha256:" + "c" * 64 + env = os.environ | {"PATH": f"{root}:{os.environ['PATH']}", "FAKE_LOG": str(root / "fake.log"), + "FAKE_RUNTIME_ROOT": str(root / "image"), "RESULTS": str(root / "results"), + "RUNTIME_IMAGE": runtime, + "GREPTIMEDB_IMAGES": json.dumps([{"label": "", "image": "greptimedb@sha256:" + "a" * 64}])} + done = subprocess.run(["bash", str(SCRIPTS / "oswal-s3-bench-run.sh"), "pull"], env=env, + capture_output=True, text=True, timeout=60) + log = (root / "fake.log").read_text().splitlines() + self.assertTrue(log[1].startswith(f"docker run --rm --entrypoint sh {runtime} -c "), log[1]) + if passes: + self.assertEqual(done.returncode, 0, done.stderr) + self.assertEqual([line.split()[1] for line in log], ["pull", "run", "pull", "image", "image"]) + else: + self.assertNotEqual(done.returncode, 0) + self.assertIn("lacks the agent observability loader or generator", done.stderr) + self.assertEqual([line.split()[1] for line in log], ["pull", "run"]) if __name__ == "__main__": unittest.main()