diff --git a/.github/scripts/query-regression-run.py b/.github/scripts/query-regression-run.py index 268cb0bae6c..e86ef707734 100644 --- a/.github/scripts/query-regression-run.py +++ b/.github/scripts/query-regression-run.py @@ -24,7 +24,9 @@ import re import signal import socket import subprocess +import threading import time +import urllib.parse import urllib.request from dataclasses import dataclass from pathlib import Path @@ -295,6 +297,257 @@ def load_plan(fixture_generator: Path, case_path: Path) -> dict[str, Any]: return json.loads(result.stdout) + +CPU_PROFILE_SECONDS = 10 +CPU_PROFILE_FREQUENCY = 49 +CPU_PROFILE_TIMEOUT_SECONDS = 180 +CPU_PROFILE_MANIFEST = Path("logs/cpu-profile-manifest.json") + + +def select_cpu_profile_query(plan: dict[str, Any]) -> tuple[dict[str, str] | None, str | None]: + """Select the first timed SQL or TQL query from a supported query plan.""" + scenario = plan.get("scenario") + if not isinstance(scenario, dict): + return None, "fixture plan has no scenario" + kind = scenario.get("kind") + if kind == "direct_readable_sst": + tables = scenario.get("tables") + database = ( + tables[0].get("database") + if isinstance(tables, list) and tables and isinstance(tables[0], dict) + else None + ) + elif kind == "prom_remote_write_then_query": + remote_write = scenario.get("remote_write") + database = remote_write.get("database") if isinstance(remote_write, dict) else None + else: + return None, f"CPU profiling does not cover scenario kind {kind!r}" + if not isinstance(database, str) or not database: + return None, "fixture plan has no query database" + queries = scenario.get("queries") + if not isinstance(queries, list): + return None, "fixture plan has no configured queries" + for query in queries: + if not isinstance(query, dict): + continue + query_kind = query.get("kind") + statement = query.get("query") + iterations = query.get("iterations", 1) + if ( + query_kind in {"sql", "tql"} + and isinstance(statement, str) + and statement + and isinstance(iterations, int) + and iterations > 0 + ): + name = query.get("name") + return { + "name": name if isinstance(name, str) else statement, + "kind": query_kind, + "query": statement, + "database": database, + }, None + return None, "no timed SQL or TQL query is configured" + + +def binary_version(binary: Path) -> dict[str, str]: + result: dict[str, str] = {"path": str(binary)} + try: + completed = subprocess.run( + [str(binary), "--version"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=30, + check=False, + ) + output = completed.stdout.strip() or completed.stderr.strip() + if completed.returncode == 0: + result["version"] = output + else: + result["error"] = f"--version exited {completed.returncode}: {output}" + except Exception as err: # noqa: BLE001 - version metadata must not hide profiling artifacts + result["error"] = repr(err) + return result + + +def query_response_error(payload: Any) -> str | None: + """Return a `/v1/sql` failure using the runner's response contract.""" + if not isinstance(payload, dict): + return "query response was not a JSON object" + for key in ("error", "err_msg", "error_msg"): + value = payload.get(key) + if value is not None and value != "" and value is not False and value != 0: + return f"query response {key}: {value}" + error_code = payload.get("error_code") + if error_code is not None and str(error_code).lower() not in {"", "0", "success"}: + return f"query response error_code {error_code}: {payload.get('message')}" + code = payload.get("code") + if "output" not in payload and code is not None and str(code).lower() not in {"", "0", "success"}: + return f"query response code {code}: {payload.get('message')}" + return None + + +def post_profile_query( + target: RunTarget, + query: dict[str, str], + timeout_s: float, +) -> tuple[bool, str | None]: + body = urllib.parse.urlencode( + {"sql": query["query"], "db": query["database"], "format": "json"}, + ).encode() + request = urllib.request.Request( + f"http://127.0.0.1:{target.http_port}/v1/sql", + data=body, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=timeout_s) as response: + payload = response.read() + if response.status >= 400: + return False, f"HTTP {response.status}" + error = query_response_error(json.loads(payload)) + return error is None, error + except Exception as err: # noqa: BLE001 - request diagnostics belong in the profile manifest + return False, repr(err) + + +def download_cpu_profile(port: int, output: Path) -> None: + request = urllib.request.Request( + f"http://127.0.0.1:{port}/debug/prof/cpu?seconds={CPU_PROFILE_SECONDS}&frequency={CPU_PROFILE_FREQUENCY}&output=proto", + data=b"", + method="POST", + ) + with urllib.request.urlopen(request, timeout=CPU_PROFILE_TIMEOUT_SECONDS) as response: + if response.status >= 400: + raise RuntimeError(f"CPU profile request returned HTTP {response.status}") + payload = response.read() + output.parent.mkdir(parents=True, exist_ok=True) + output.write_bytes(payload) + if not payload: + raise RuntimeError("CPU profile response was empty") + + +def collect_target_cpu_profiles( + target: RunTarget, + query: dict[str, str], + work_dir: Path, + manifest_target: dict[str, Any], + http_timeout: float, +) -> None: + """Profile one target while executing only during the sampling interval.""" + profiles = manifest_target["profiles"] + workers: list[threading.Thread] = [] + try: + for component, port in (("frontend", target.http_port), ("datanode", target.datanode_http_port)): + artifact = work_dir / "logs" / "cpu-profiles" / target.name / f"{component}.pb" + profile: dict[str, Any] = { + "component": component, + "endpoint": f"http://127.0.0.1:{port}/debug/prof/cpu", + "path": str(artifact.relative_to(work_dir)), + "status": "running", + } + profiles.append(profile) + + def run_profile(record: dict[str, Any] = profile, profile_port: int = port, profile_path: Path = artifact) -> None: + try: + download_cpu_profile(profile_port, profile_path) + record["status"] = "success" + except Exception as err: # noqa: BLE001 - retain each component failure in the manifest + record["status"] = "failure" + record["error"] = repr(err) + + worker = threading.Thread( + target=run_profile, + name=f"cpu-profile-{target.name}-{component}", + ) + worker.start() + workers.append(worker) + + deadline = time.monotonic() + CPU_PROFILE_SECONDS + query_result = manifest_target["query_execution"] + while time.monotonic() < deadline: + # The deadline controls launching new work. A query already launched may use + # its normal timeout and finish after sampling ends. + ok, error = post_profile_query(target, query, http_timeout) + query_result["attempts"] += 1 + if ok: + query_result["successes"] += 1 + else: + query_result["failures"] += 1 + if len(query_result["errors"]) < 5: + query_result["errors"].append(error) + finally: + # Reports can still be symbolizing if startup or query execution raises. + # Join before the caller reaches stop_all and terminates their components. + for worker in workers: + worker.join() + + +def collect_cpu_profiles(args: argparse.Namespace, plan: dict[str, Any], work_dir: Path, targets: list[RunTarget]) -> str: + """Collect default CPU profiles for direct-SST and remote-write query cases.""" + query, unsupported_reason = select_cpu_profile_query(plan) + manifest: dict[str, Any] = { + "status": "unsupported", + "sampling": { + "seconds": CPU_PROFILE_SECONDS, + "frequency": CPU_PROFILE_FREQUENCY, + "output": "proto", + "request_timeout_seconds": CPU_PROFILE_TIMEOUT_SECONDS, + }, + "binaries": {target.name: binary_version(target.binary) for target in targets}, + "query": { + "selection": "first timed SQL or TQL query", + "selected": query, + "unsupported_reason": unsupported_reason, + }, + "targets": [], + "errors": [], + } + if query is not None: + for target in targets: # Base and candidate intentionally do not profile concurrently. + target_manifest: dict[str, Any] = { + "role": target.name, + "query_execution": {"attempts": 0, "successes": 0, "failures": 0, "errors": []}, + "profiles": [], + } + manifest["targets"].append(target_manifest) + try: + collect_target_cpu_profiles( + target, + query, + work_dir, + target_manifest, + float(args.http_timeout), + ) + except Exception as err: # noqa: BLE001 - save the partial manifest before failing the case + target_manifest["error"] = repr(err) + manifest["errors"].append(f"{target.name}: {err!r}") + failed_profiles = [ + profile for target in manifest["targets"] for profile in target["profiles"] + if profile["status"] != "success" + ] + failed_queries = [ + target + for target in manifest["targets"] + if ( + target["query_execution"]["successes"] == 0 + or target["query_execution"]["failures"] + or target.get("error") + ) + ] + manifest["status"] = "failure" if failed_profiles or failed_queries else "success" + manifest_path = work_dir / CPU_PROFILE_MANIFEST + manifest_path.parent.mkdir(parents=True, exist_ok=True) + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8") + if manifest["status"] != "success": + if manifest["status"] == "unsupported": + reason = unsupported_reason + else: + reason = manifest["errors"] or "profile or workload query failed" + print(f"CPU profile {manifest['status']}: {reason}; manifest: {manifest_path}", flush=True) + return manifest["status"] + def run_direct_case( args: argparse.Namespace, case_path: Path, @@ -303,6 +556,7 @@ def run_direct_case( candidate_bin: Path, fixture_generator: Path, runner: Path, + plan: dict[str, Any], ) -> int: ports = allocate_ports(16) targets = [ @@ -378,7 +632,9 @@ def run_direct_case( for target in targets: restart_component(target, "frontend", procs) status = subprocess.run(measure, check=False).returncode - return status + if status != 0: + return status + return 0 if collect_cpu_profiles(args, plan, work_dir, targets) == "success" else 1 finally: stop_all(targets, procs) @@ -391,6 +647,7 @@ def run_remote_case( candidate_bin: Path, fixture_generator: Path, runner: Path, + plan: dict[str, Any], ) -> int: ports = allocate_ports(16) targets = [ @@ -454,6 +711,9 @@ def run_remote_case( measure_status = subprocess.run(measure, check=False).returncode if measure_status != 0 and not report.exists(): return measure_status + profile_status = "unsupported" + if measure_status == 0: + profile_status = collect_cpu_profiles(args, plan, work_dir, targets) for target in targets: stop_component(target, "datanode", procs) @@ -470,7 +730,11 @@ def run_remote_case( if finalize_status != 0: return finalize_status final_report = json.loads(report.read_text(encoding="utf-8")) - return 1 if measure_status != 0 or final_report.get("status") == "failed" else 0 + return 1 if ( + measure_status != 0 + or profile_status != "success" + or final_report.get("status") == "failed" + ) else 0 finally: stop_all(targets, procs) @@ -556,6 +820,7 @@ def run_case(args: argparse.Namespace, case_path: Path, work_dir: Path) -> int: candidate_bin, fixture_generator, runner, + plan, ) if kind == "prom_remote_write_then_query": return run_remote_case( @@ -566,6 +831,7 @@ def run_case(args: argparse.Namespace, case_path: Path, work_dir: Path) -> int: candidate_bin, fixture_generator, runner, + plan, ) if kind == "otlp_trace_load": return run_otlp_case( diff --git a/tests/perf/README.md b/tests/perf/README.md index d518fec996c..a9e7575ae99 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -366,6 +366,30 @@ uv run --no-project python .github/scripts/query-regression-run.py \ --work-dir /tmp/query-regression-work ``` +### Default CPU profiles + +After a successful `direct_readable_sst` or `prom_remote_write_then_query` +measure and validation, the outer runner automatically records CPU protobufs +for the base target and then the candidate target. For each role, frontend and datanode profile concurrently via +`POST /debug/prof/cpu?seconds=10&frequency=49&output=proto`, while the first +configured timed SQL or TQL query is launched only during the ten-second +sampling window. A query launched before that deadline retains its normal HTTP +timeout and may finish after sampling ends. The ordinary measurement samples and +report remain unchanged. + +Each case writes `logs/cpu-profile-manifest.json` and +`logs/cpu-profiles//{frontend,datanode}.pb`. The manifest records the +selected query name and text, role, binary version, artifact paths, query and +profile errors, and `success`, `failure`, or `unsupported` status. A profile or +profile-workload failure fails the direct-SST case after retaining those files; +an unsupported direct-SST query selection is recorded explicitly. The existing +workflow artifact upload includes the case `logs/**` tree, so these outputs are +retained with the report. + +This default profiling stage covers all seven routine default cases: three +direct-SST and four remote-write query cases. `otlp_trace_load` is an ingest +scenario and is explicitly not covered. + The Rust runner subcommands are also useful for focused diagnostics: ```bash diff --git a/tests/perf/test_query_regression_case_selection.py b/tests/perf/test_query_regression_case_selection.py index 84b3cd44fa7..706d0f67909 100644 --- a/tests/perf/test_query_regression_case_selection.py +++ b/tests/perf/test_query_regression_case_selection.py @@ -16,6 +16,7 @@ """Coverage for query-regression case group selection.""" import importlib.util +import json import sys import unittest from pathlib import Path @@ -55,5 +56,432 @@ class QueryRegressionCaseSelectionTest(unittest.TestCase): self.assertEqual(runner.split_cases([case]), [case]) +class QueryRegressionCpuProfileTest(unittest.TestCase): + def make_target(self, name: str, root: Path, port_offset: int) -> object: + return runner.make_target( + name, + Path(f"/{name}/greptime"), + root, + list(range(port_offset, port_offset + 8)), + ) + + def test_direct_profiles_run_after_measurement_and_before_cleanup(self) -> None: + from tempfile import TemporaryDirectory + from types import SimpleNamespace + from unittest.mock import patch + + with TemporaryDirectory() as directory: + root = Path(directory) + events: list[str] = [] + + def run(command: list[str], **_kwargs: object) -> object: + if "prepare-direct" in command: + Path(command[command.index("--output") + 1]).write_text(json.dumps({"fixtures": []})) + if "measure" in command: + events.append("measure") + Path(command[command.index("--output") + 1]).write_text( + json.dumps({"targets": [{"measurements": [{"samples": ["ordinary-sample"]}]}]}), + ) + return SimpleNamespace(returncode=0) + + def profile(*_args: object, **_kwargs: object) -> str: + events.append("profile") + return "success" + + def stop(_target: object, _name: str, _procs: object) -> None: + events.append("cleanup" if "measure" in events else "setup-stop") + + args = SimpleNamespace(http_timeout="1", allow_large_fixture="false") + plan = { + "scenario": { + "tables": [{"database": "public"}], + "queries": [{"name": "selected", "kind": "sql", "query": "SELECT 1", "iterations": 1}], + # Measurements belong to the ordinary runner report, not CPU profiling. + "measurements": [{"samples": ["ordinary-sample"]}], + }, + } + with ( + patch.object(runner, "allocate_ports", return_value=list(range(1000, 1016))), + patch.object(runner, "start_component", side_effect=lambda *_args: events.append("start")), + patch.object(runner, "stop_component", side_effect=stop), + patch.object(runner.subprocess, "run", side_effect=run), + patch.object(runner, "collect_cpu_profiles", side_effect=profile) as collect, + ): + status = runner.run_direct_case( + args, root / "case.toml", root, Path("base"), Path("candidate"), Path("fixture"), Path("runner"), plan, + ) + + self.assertEqual(0, status) + self.assertLess(events.index("measure"), events.index("profile")) + self.assertLess(events.index("profile"), events.index("cleanup")) + self.assertEqual(plan["scenario"]["measurements"][0]["samples"], ["ordinary-sample"]) + report = json.loads((root / "query-regression-report.json").read_text()) + self.assertEqual(report["targets"][0]["measurements"][0]["samples"], ["ordinary-sample"]) + self.assertEqual([target.name for target in collect.call_args.args[3]], ["base", "candidate"]) + + def test_cpu_profiles_write_base_and_candidate_manifest_without_network(self) -> None: + from tempfile import TemporaryDirectory + from types import SimpleNamespace + from unittest.mock import patch + + with TemporaryDirectory() as directory: + root = Path(directory) + targets = [self.make_target("base", root, 1000), self.make_target("candidate", root, 2000)] + seen: list[str] = [] + + def collect(target: object, _query: object, work_dir: Path, target_manifest: dict[str, object], _timeout: float) -> None: + seen.append(target.name) + target_manifest["query_execution"] = {"attempts": 1, "successes": 1, "failures": 0, "errors": []} + target_manifest["profiles"] = [{"component": "frontend", "path": str(work_dir / target.name), "status": "success"}] + + plan = {"scenario": {"kind": "direct_readable_sst", "tables": [{"database": "public"}], "queries": [ + {"name": "not-supported", "kind": "prom_http", "query": "metric", "iterations": 1}, + {"name": "chosen", "kind": "tql", "query": "TQL ANALYZE x", "iterations": 1}, + ]}} + with ( + patch.object(runner, "binary_version", side_effect=lambda binary: {"path": str(binary), "version": "test"}), + patch.object(runner, "collect_target_cpu_profiles", side_effect=collect), + ): + status = runner.collect_cpu_profiles(SimpleNamespace(http_timeout="1"), plan, root, targets) + + manifest = json.loads((root / runner.CPU_PROFILE_MANIFEST).read_text()) + self.assertEqual("success", status) + self.assertEqual(["base", "candidate"], seen) + self.assertEqual("chosen", manifest["query"]["selected"]["name"]) + self.assertEqual(["base", "candidate"], [target["role"] for target in manifest["targets"]]) + + def test_profile_query_response_errors_are_failures(self) -> None: + self.assertIsNone(runner.query_response_error({"code": 0})) + self.assertEqual("query response error: denied", runner.query_response_error({"error": "denied"})) + self.assertEqual("query response error_code 7: bad request", runner.query_response_error({"error_code": 7, "message": "bad request"})) + self.assertEqual("query response was not a JSON object", runner.query_response_error([])) + + def test_target_cpu_profiles_use_mocked_profile_and_query_http_helpers(self) -> None: + from tempfile import TemporaryDirectory + from unittest.mock import patch + + with TemporaryDirectory() as directory: + root = Path(directory) + target = self.make_target("base", root, 1000) + manifest_target: dict[str, object] = { + "query_execution": {"attempts": 0, "successes": 0, "failures": 0, "errors": []}, + "profiles": [], + } + ticks = iter([0.0, 0.0, float(runner.CPU_PROFILE_SECONDS)]) + query = {"name": "chosen", "kind": "sql", "query": "SELECT 1", "database": "public"} + with ( + patch.object(runner, "download_cpu_profile") as download, + patch.object(runner, "post_profile_query", return_value=(True, None)) as post, + patch.object(runner.time, "monotonic", side_effect=lambda: next(ticks)), + ): + runner.collect_target_cpu_profiles(target, query, root, manifest_target, 1.0) + + self.assertEqual(2, download.call_count) + self.assertEqual({target.http_port, target.datanode_http_port}, {call.args[0] for call in download.call_args_list}) + post.assert_called_once_with(target, query, 1.0) + self.assertEqual(1, manifest_target["query_execution"]["successes"]) + self.assertTrue(all(profile["status"] == "success" for profile in manifest_target["profiles"])) + + def test_unsupported_direct_profile_prints_reason_and_manifest_and_fails(self) -> None: + from contextlib import redirect_stdout + from io import StringIO + from tempfile import TemporaryDirectory + from types import SimpleNamespace + from unittest.mock import patch + + with TemporaryDirectory() as directory: + root = Path(directory) + targets = [self.make_target("base", root, 1000), self.make_target("candidate", root, 2000)] + plan = {"scenario": {"kind": "direct_readable_sst", "tables": [{"database": "public"}], "queries": [{"kind": "prom_http", "query": "metric", "iterations": 1}]}} + output = StringIO() + with patch.object(runner, "binary_version", return_value={"version": "test"}), redirect_stdout(output): + status = runner.collect_cpu_profiles(SimpleNamespace(http_timeout="300"), plan, root, targets) + + manifest_path = root / runner.CPU_PROFILE_MANIFEST + manifest = json.loads(manifest_path.read_text()) + self.assertEqual("unsupported", status) + self.assertEqual("unsupported", manifest["status"]) + self.assertIn("no timed SQL or TQL query is configured", output.getvalue()) + self.assertIn(str(manifest_path), output.getvalue()) + self.assertNotEqual("success", status) + + def test_unsupported_direct_case_preserves_measurement_and_returns_nonzero(self) -> None: + from contextlib import redirect_stdout + from io import StringIO + from tempfile import TemporaryDirectory + from types import SimpleNamespace + from unittest.mock import patch + + with TemporaryDirectory() as directory: + root = Path(directory) + + def run(command: list[str], **_kwargs: object) -> object: + if "prepare-direct" in command: + Path(command[command.index("--output") + 1]).write_text(json.dumps({"fixtures": []})) + if "measure" in command: + Path(command[command.index("--output") + 1]).write_text(json.dumps({"status": "ok", "samples": ["ordinary"]})) + return SimpleNamespace(returncode=0) + + args = SimpleNamespace(http_timeout="1", allow_large_fixture="false") + plan = {"scenario": {"kind": "direct_readable_sst", "tables": [{"database": "public"}], "queries": [{"kind": "prom_http", "query": "metric", "iterations": 1}]}} + output = StringIO() + with ( + patch.object(runner, "allocate_ports", return_value=list(range(1000, 1016))), + patch.object(runner, "start_component"), + patch.object(runner, "stop_component"), + patch.object(runner.subprocess, "run", side_effect=run), + patch.object(runner, "binary_version", return_value={"version": "test"}), + redirect_stdout(output), + ): + status = runner.run_direct_case( + args, root / "case.toml", root, Path("base"), Path("candidate"), Path("fixture"), Path("runner"), plan, + ) + + manifest_path = root / runner.CPU_PROFILE_MANIFEST + self.assertEqual(1, status) + self.assertEqual(["ordinary"], json.loads((root / "query-regression-report.json").read_text())["samples"]) + self.assertTrue(manifest_path.exists()) + self.assertIn("no timed SQL or TQL query is configured", output.getvalue()) + self.assertIn(str(manifest_path), output.getvalue()) + + def test_profile_query_near_deadline_keeps_normal_timeout_and_stops_launching(self) -> None: + from tempfile import TemporaryDirectory + from unittest.mock import patch + + with TemporaryDirectory() as directory: + root = Path(directory) + target = self.make_target("base", root, 1000) + manifest_target: dict[str, object] = { + "query_execution": {"attempts": 0, "successes": 0, "failures": 0, "errors": []}, + "profiles": [], + } + # The first query starts just before the deadline; after it completes, + # the next loop check must not launch another query. + ticks = iter([0.0, 9.99, 10.01]) + query = {"name": "chosen", "kind": "sql", "query": "SELECT 1", "database": "public"} + with ( + patch.object(runner, "download_cpu_profile"), + patch.object(runner, "post_profile_query", return_value=(True, None)) as post, + patch.object(runner.time, "monotonic", side_effect=lambda: next(ticks)), + ): + runner.collect_target_cpu_profiles(target, query, root, manifest_target, 300.0) + + post.assert_called_once_with(target, query, 300.0) + self.assertEqual(1, manifest_target["query_execution"]["attempts"]) + + def test_profile_workers_join_when_query_loop_raises(self) -> None: + from tempfile import TemporaryDirectory + from unittest.mock import patch + + class Thread: + instances: list["Thread"] = [] + + def __init__(self, **_kwargs: object) -> None: + self.started = False + self.joined = False + self.instances.append(self) + + def start(self) -> None: + self.started = True + + def join(self) -> None: + self.joined = True + + with TemporaryDirectory() as directory: + root = Path(directory) + target = self.make_target("base", root, 1000) + manifest_target: dict[str, object] = { + "query_execution": {"attempts": 0, "successes": 0, "failures": 0, "errors": []}, + "profiles": [], + } + query = {"name": "chosen", "kind": "sql", "query": "SELECT 1", "database": "public"} + with ( + patch.object(runner.threading, "Thread", Thread), + patch.object(runner, "post_profile_query", side_effect=RuntimeError("query loop failed")), + patch.object(runner.time, "monotonic", side_effect=[0.0, 0.0]), + ): + with self.assertRaisesRegex(RuntimeError, "query loop failed"): + runner.collect_target_cpu_profiles(target, query, root, manifest_target, 300.0) + + self.assertEqual(2, len(Thread.instances)) + self.assertTrue(all(thread.started and thread.joined for thread in Thread.instances)) + + def test_second_profile_worker_start_failure_joins_first_and_is_captured(self) -> None: + from tempfile import TemporaryDirectory + from types import SimpleNamespace + from unittest.mock import patch + + class Thread: + instances: list["Thread"] = [] + + def __init__(self, **_kwargs: object) -> None: + self.number = len(self.instances) + self.started = False + self.joined = False + self.instances.append(self) + + def start(self) -> None: + self.started = True + if self.number == 1: + raise RuntimeError("second start failed") + + def join(self) -> None: + self.joined = True + + with TemporaryDirectory() as directory: + root = Path(directory) + target = self.make_target("base", root, 1000) + plan = {"scenario": {"kind": "direct_readable_sst", "tables": [{"database": "public"}], "queries": [{"kind": "sql", "query": "SELECT 1", "iterations": 1}]}} + with ( + patch.object(runner.threading, "Thread", Thread), + patch.object(runner, "binary_version", return_value={"version": "test"}), + ): + status = runner.collect_cpu_profiles(SimpleNamespace(http_timeout="300"), plan, root, [target]) + + manifest = json.loads((root / runner.CPU_PROFILE_MANIFEST).read_text()) + self.assertEqual("failure", status) + self.assertTrue(Thread.instances[0].started) + self.assertTrue(Thread.instances[0].joined) + self.assertTrue(Thread.instances[1].started) + self.assertFalse(Thread.instances[1].joined) + self.assertIn("second start failed", manifest["targets"][0]["error"]) + + def test_zero_profile_query_successes_fail_manifest(self) -> None: + from tempfile import TemporaryDirectory + from types import SimpleNamespace + from unittest.mock import patch + + with TemporaryDirectory() as directory: + root = Path(directory) + targets = [self.make_target("base", root, 1000), self.make_target("candidate", root, 2000)] + + def collect(_target: object, _query: object, _work_dir: Path, target_manifest: dict[str, object], _timeout: float) -> None: + target_manifest["query_execution"] = {"attempts": 0, "successes": 0, "failures": 0, "errors": []} + target_manifest["profiles"] = [{"component": "frontend", "status": "success"}] + + plan = {"scenario": {"kind": "direct_readable_sst", "tables": [{"database": "public"}], "queries": [{"kind": "sql", "query": "SELECT 1", "iterations": 1}]}} + with ( + patch.object(runner, "binary_version", return_value={"version": "test"}), + patch.object(runner, "collect_target_cpu_profiles", side_effect=collect), + ): + status = runner.collect_cpu_profiles(SimpleNamespace(http_timeout="300"), plan, root, targets) + + manifest = json.loads((root / runner.CPU_PROFILE_MANIFEST).read_text()) + self.assertEqual("failure", status) + self.assertEqual("failure", manifest["status"]) + + def test_cpu_profile_failure_writes_manifest_without_network(self) -> None: + from tempfile import TemporaryDirectory + from types import SimpleNamespace + from unittest.mock import patch + + with TemporaryDirectory() as directory: + root = Path(directory) + targets = [self.make_target("base", root, 1000), self.make_target("candidate", root, 2000)] + + def collect(_target: object, _query: object, _work_dir: Path, target_manifest: dict[str, object], _timeout: float) -> None: + target_manifest["query_execution"] = {"attempts": 1, "successes": 0, "failures": 1, "errors": ["query failed"]} + target_manifest["profiles"] = [{"component": "frontend", "status": "failure", "error": "profile failed"}] + + plan = {"scenario": {"kind": "direct_readable_sst", "tables": [{"database": "public"}], "queries": [{"kind": "sql", "query": "SELECT 1", "iterations": 1}]}} + with ( + patch.object(runner, "binary_version", return_value={"version": "test"}), + patch.object(runner, "collect_target_cpu_profiles", side_effect=collect), + ): + status = runner.collect_cpu_profiles(SimpleNamespace(http_timeout="1"), plan, root, targets) + + manifest = json.loads((root / runner.CPU_PROFILE_MANIFEST).read_text()) + self.assertEqual("failure", status) + self.assertEqual("failure", manifest["status"]) + self.assertEqual("profile failed", manifest["targets"][0]["profiles"][0]["error"]) + self.assertEqual(["base", "candidate"], [target["role"] for target in manifest["targets"]]) + + def test_remote_profiles_before_finalize_and_failure_keeps_finalize(self) -> None: + from tempfile import TemporaryDirectory + from types import SimpleNamespace + from unittest.mock import patch + + for profile_status in ("success", "failure"): + with self.subTest(profile_status=profile_status), TemporaryDirectory() as directory: + root = Path(directory) + events: list[str] = [] + + def run(command: list[str], **_kwargs: object) -> object: + if "measure" in command: + events.append("measure") + Path(command[command.index("--output") + 1]).write_text(json.dumps({"status": "ok"})) + if "finalize-remote" in command: + events.append("finalize") + return SimpleNamespace(returncode=0) + + def profile(_args: object, plan: dict[str, object], _work_dir: Path, _targets: object) -> str: + events.append("profile") + self.assertEqual("remote_db", plan["scenario"]["remote_write"]["database"]) + return profile_status + + args = SimpleNamespace(http_timeout="1") + plan = {"scenario": {"kind": "prom_remote_write_then_query", "remote_write": {"database": "remote_db"}, "queries": [{"kind": "tql", "query": "TQL ANALYZE x", "iterations": 1}]}} + with ( + patch.object(runner, "allocate_ports", return_value=list(range(1000, 1016))), + patch.object(runner, "start_component"), + patch.object(runner, "stop_component", side_effect=lambda _target, name, _procs: events.append(f"stop-{name}")), + patch.object(runner.subprocess, "run", side_effect=run), + patch.object(runner, "collect_cpu_profiles", side_effect=profile), + ): + status = runner.run_remote_case( + args, root / "case.toml", root, Path("base"), Path("candidate"), Path("fixture"), Path("runner"), plan, + ) + + self.assertLess(events.index("measure"), events.index("profile")) + self.assertLess(events.index("profile"), events.index("stop-datanode")) + self.assertLess(events.index("stop-datanode"), events.index("finalize")) + self.assertEqual(0 if profile_status == "success" else 1, status) + + def test_remote_profile_selection_uses_remote_write_database(self) -> None: + plan = {"scenario": {"kind": "prom_remote_write_then_query", "remote_write": {"database": "remote_db"}, "queries": [{"name": "chosen", "kind": "sql", "query": "SELECT 1", "iterations": 1}]}} + query, reason = runner.select_cpu_profile_query(plan) + self.assertIsNone(reason) + self.assertEqual("remote_db", query["database"]) + + def test_cpu_profile_failure_manifest_is_saved_and_direct_case_cleans_up(self) -> None: + from tempfile import TemporaryDirectory + from types import SimpleNamespace + from unittest.mock import patch + + with TemporaryDirectory() as directory: + root = Path(directory) + events: list[str] = [] + + def run(command: list[str], **_kwargs: object) -> object: + if "prepare-direct" in command: + Path(command[command.index("--output") + 1]).write_text(json.dumps({"fixtures": []})) + if "measure" in command: + events.append("measure") + return SimpleNamespace(returncode=0) + + def failure(_args: object, _plan: object, work_dir: Path, _targets: object) -> str: + manifest = work_dir / runner.CPU_PROFILE_MANIFEST + manifest.parent.mkdir(parents=True, exist_ok=True) + manifest.write_text(json.dumps({"status": "failure", "errors": ["profile unavailable"]})) + return "failure" + + args = SimpleNamespace(http_timeout="1", allow_large_fixture="false") + plan = {"scenario": {"kind": "direct_readable_sst", "tables": [{"database": "public"}], "queries": [{"kind": "sql", "query": "SELECT 1", "iterations": 1}]}} + with ( + patch.object(runner, "allocate_ports", return_value=list(range(1000, 1016))), + patch.object(runner, "start_component"), + patch.object(runner, "stop_component", side_effect=lambda *_args: events.append("cleanup")), + patch.object(runner.subprocess, "run", side_effect=run), + patch.object(runner, "collect_cpu_profiles", side_effect=failure), + ): + status = runner.run_direct_case( + args, root / "case.toml", root, Path("base"), Path("candidate"), Path("fixture"), Path("runner"), plan, + ) + + self.assertEqual(1, status) + self.assertIn("cleanup", events) + self.assertEqual("failure", json.loads((root / runner.CPU_PROFILE_MANIFEST).read_text())["status"]) + + if __name__ == "__main__": unittest.main() diff --git a/tests/perf/test_query_regression_runner_compaction_toctou.py b/tests/perf/test_query_regression_runner_compaction_toctou.py index 58f7f31b7fa..2f3a0e06d86 100644 --- a/tests/perf/test_query_regression_runner_compaction_toctou.py +++ b/tests/perf/test_query_regression_runner_compaction_toctou.py @@ -54,16 +54,18 @@ class RemoteWriteLifecycleTest(unittest.TestCase): def stop(target, component, _procs): events.append(f"stop:{target.name}:{component}") + plan = {"scenario": {"kind": "prom_remote_write_then_query", "remote_write": {"database": "public"}, "queries": [{"kind": "sql", "query": "SELECT 1", "iterations": 1}]}} with ( tempfile.TemporaryDirectory() as tmpdir, patch.object(runner, "allocate_ports", return_value=self.ports), patch.object(runner, "start_component", side_effect=start), patch.object(runner, "stop_component", side_effect=stop), patch.object(runner.subprocess, "run", side_effect=run), + patch.object(runner, "collect_cpu_profiles", return_value="success"), ): self.assertEqual( runner.run_remote_case( - self.args, Path("case.toml"), Path(tmpdir), Path("base"), Path("candidate"), Path("fixture"), Path("runner") + self.args, Path("case.toml"), Path(tmpdir), Path("base"), Path("candidate"), Path("fixture"), Path("runner"), plan ), 0, ) @@ -93,6 +95,7 @@ class RemoteWriteLifecycleTest(unittest.TestCase): def stop(target, component, _procs): events.append(f"stop:{target.name}:{component}") + plan = {"scenario": {"kind": "prom_remote_write_then_query", "remote_write": {"database": "public"}, "queries": [{"kind": "sql", "query": "SELECT 1", "iterations": 1}]}} with ( tempfile.TemporaryDirectory() as tmpdir, patch.object(runner, "allocate_ports", return_value=self.ports), @@ -102,7 +105,7 @@ class RemoteWriteLifecycleTest(unittest.TestCase): ): with self.assertRaisesRegex(RuntimeError, "prepare exploded"): runner.run_remote_case( - self.args, Path("case.toml"), Path(tmpdir), Path("base"), Path("candidate"), Path("fixture"), Path("runner") + self.args, Path("case.toml"), Path(tmpdir), Path("base"), Path("candidate"), Path("fixture"), Path("runner"), plan ) self.assertIn("prepare-remote", events)