diff --git a/Cargo.lock b/Cargo.lock index cc47268b5d..3befeef586 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -411,6 +411,28 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bigquery-benchmark-queries" +version = "0.0.0" +dependencies = [ + "anyhow", + "clap", + "google-cloud-auth", + "google-cloud-bigquery", + "humantime", + "integration-tests-o11y", + "opentelemetry", + "opentelemetry_sdk", + "serde", + "serde_json", + "tokio", + "tokio-metrics", + "tracing", + "tracing-log", + "tracing-subscriber", + "uuid", +] + [[package]] name = "bigquery-grpc-mock" version = "0.0.0" diff --git a/Cargo.toml b/Cargo.toml index 55da18e7b4..1cdd04b5da 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -69,6 +69,7 @@ members = [ "src/auth", "src/bigquery", "src/bigquery-derive", + "src/bigquery/benchmarks/queries", "src/bigquery/benchmarks/write-throughput", "src/bigquery/examples", "src/bigquery/grpc-mock", diff --git a/src/bigquery/benchmarks/queries/Cargo.toml b/src/bigquery/benchmarks/queries/Cargo.toml new file mode 100644 index 0000000000..390225f32b --- /dev/null +++ b/src/bigquery/benchmarks/queries/Cargo.toml @@ -0,0 +1,45 @@ +# Copyright 2026 Google LLC +# +# 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 +# +# https://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. + +[package] +name = "bigquery-benchmark-queries" +version = "0.0.0" +publish = false +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +keywords.workspace = true +categories.workspace = true + +[dependencies] +anyhow.workspace = true +clap = { workspace = true, features = ["derive", "env", "help", "std", "usage"] } +google-cloud-auth.workspace = true +google-cloud-bigquery = { workspace = true, features = ["default-rustls-provider"] } +humantime.workspace = true +integration-tests-o11y.workspace = true +opentelemetry = { workspace = true, features = ["metrics", "trace"] } +opentelemetry_sdk = { workspace = true, features = ["metrics", "rt-tokio", "trace"] } +serde = { workspace = true, features = ["derive"] } +serde_json.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "signal", "sync", "time"] } +tokio-metrics = { workspace = true, features = ["rt"] } +tracing.workspace = true +tracing-log = { workspace = true, features = ["log-tracer", "std"] } +tracing-subscriber = { workspace = true, features = ["env-filter", "fmt", "std"] } +uuid.workspace = true + +[lints] +workspace = true diff --git a/src/bigquery/benchmarks/queries/README.md b/src/bigquery/benchmarks/queries/README.md new file mode 100644 index 0000000000..d7236896a9 --- /dev/null +++ b/src/bigquery/benchmarks/queries/README.md @@ -0,0 +1,322 @@ +# BigQuery SDK Benchmark & Endurance Test Suite + +Benchmarks the Rust BigQuery client library (`google-cloud-bigquery`), measuring +latency distributions (P50/P90/P99), query throughput, and endurance under +long-running workloads with OpenTelemetry tracking and under-the-hood job retry +detection. + +## Features + +- **Benchmark & Endurance Modes**: Run fixed iterations or continuous + duration-based tests (e.g. 10m, 2h, 24h). +- **Under-the-Hood Retry Detection**: Detects BigQuery job retries by checking + if the `job_id` changed between `Query::send()` and `Query::until_done()`. +- **OpenTelemetry & Cloud Observability**: Automatically exports distributed + traces to **Google Cloud Trace** and metric instruments to **Google Cloud + Monitoring** when `--project-id` is provided. +- **Configurable Scenarios**: + - `synthetic-100k` (default): Zero-dependency query generating 100,000 + structured rows in-flight with `UNNEST(GENERATE_ARRAY(1, 100000))`. + - `synthetic-10k`: Zero-dependency query generating 10,000 rows. + - `usa-names-scan`: Scans and retrieves 50,000 rows from + `bigquery-public-data.usa_names.usa_1910_2013`. + - `usa-names-agg`: Aggregates 5.5M rows grouped by state and gender. + - `wikipedia-agg`: Aggregates top 1000 page views from Wikipedia public + dataset. + - `custom`: Executes user-provided queries via `--sql` or `--sql-file`. + +______________________________________________________________________ + +## Pre-requisites + +1. **Authentication**: Ensure Application Default Credentials (ADC) are + configured: + + ```shell + gcloud auth application-default login + ``` + +1. **Project ID**: Set the project ID environment variable: + + ```shell + export GOOGLE_CLOUD_PROJECT="$(gcloud config get project)" + ``` + +______________________________________________________________________ + +## Running Benchmarks + +> [!NOTE] +> **Query Caching is disabled by default (`--use-query-cache false`)** to force +> queries to always execute against storage and provide accurate, repeatable +> latency measurements. You can pass `--use-query-cache true` if you wish to +> benchmark cache hit performance. +> +> **Indefinite Execution by Default:** If neither `--iterations` nor +> `--duration` is specified, the benchmark runs indefinitely until interrupted +> with `Ctrl+C`. + +### Synthetic Benchmark + +Runs 10 iterations per task with 4 concurrent tasks, generating and streaming +100,000 rows per query: + +```shell +cargo run --release -p bigquery-benchmark-queries -- \ + --project-id ${GOOGLE_CLOUD_PROJECT} \ + --scenario synthetic-100k \ + --task-count 4 \ + --iterations 10 \ + --output-dir ./results +``` + +### Public Dataset Query Benchmark + +Benchmark streaming 50,000 rows from the USA names public dataset: + +```shell +cargo run --release -p bigquery-benchmark-queries -- \ + --project-id ${GOOGLE_CLOUD_PROJECT} \ + --scenario usa-names-scan \ + --task-count 2 \ + --iterations 20 \ + --output-dir ./results +``` + +### Custom SQL Query Benchmark + +```shell +cargo run --release -p bigquery-benchmark-queries -- \ + --project-id ${GOOGLE_CLOUD_PROJECT} \ + --scenario custom \ + --sql "SELECT word, SUM(word_count) as total FROM \`bigquery-public-data.samples.shakespeare\` GROUP BY word ORDER BY total DESC LIMIT 100" \ + --task-count 4 \ + --iterations 25 +``` + +______________________________________________________________________ + +## Running Endurance Tests + +To test connection stability, memory leaks, and token refreshing over a +prolonged period (e.g. 1 hour): + +```shell +cargo run --release -p bigquery-benchmark-queries -- \ + --project-id ${GOOGLE_CLOUD_PROJECT} \ + --scenario synthetic-100k \ + --task-count 4 \ + --duration 1h \ + --output-dir ./results +``` + +### Running as a Background Job + +To run a long endurance test in the background (similar to the Pub/Sub benchmark +pattern) while capturing stdout reports to `.txt`, stderr tracing logs to +`.log`, and real-time CSV/JSON samples to `./results`: + +```shell +TS=$(date +%s); RUSTFLAGS="-C target-cpu=native" \ + cargo run --release -p bigquery-benchmark-queries -- \ + --project-id "${GOOGLE_CLOUD_PROJECT:-$(gcloud config get project)}" \ + --scenario synthetic-100k \ + --task-count 4 \ + --duration 1h \ + --output-dir ./results \ + >bq-bm-${TS}.txt 2>bq-bm-${TS}.log >(tee bq-bm-${TS}.txt | systemd-cat -t bq-benchmark) \ + 2> >(tee bq-bm-${TS}.log | systemd-cat -t bq-benchmark) Logs Explorer** in Google Cloud Console. + - Filter by your identifier: + ```log + resource.type="gce_instance" + jsonPayload.SYSLOG_IDENTIFIER="bq-benchmark" + ``` + - Click **Stream logs** in the top right to watch real-time benchmark + execution logs arrive in the console. + +______________________________________________________________________ + +## OpenTelemetry & Cloud Observability + +When `--project-id` is specified, the benchmark automatically connects to +`telemetry.googleapis.com` and records: + +### Cloud Monitoring Metrics + +- `bigquery.queries.total`: Total count of executed queries. +- `bigquery.queries.success`: Successful query executions. +- `bigquery.queries.error`: Query failures. +- `bigquery.queries.retries_detected`: Queries where an under-the-hood job retry + was detected. +- `bigquery.queries.rows_read`: Cumulative rows streamed. +- `bigquery.queries.bytes_processed`: Total bytes processed. +- `bigquery.queries.duration_seconds`: Histogram of total end-to-end query + latency. +- `bigquery.queries.send_duration_seconds`: Histogram of `Query::send()` + latency. +- `bigquery.queries.poll_duration_seconds`: Histogram of `Query::until_done()` + polling latency. +- `bigquery.queries.read_duration_seconds`: Histogram of `CompleteQuery::read()` + row streaming latency. + +### Cloud Trace Spans + +- Root span: `bigquery.query_benchmark.iteration` +- Child spans: + - `bigquery.send` + - `bigquery.until_done` + - `bigquery.read_rows` + +### Monitoring with PromQL in Google Cloud Monitoring + +In the Google Cloud Console, navigate to **Monitoring** > **Metrics Explorer** +and select the **PromQL** tab. You can use the following PromQL queries to +visualize the benchmark metrics: + +#### 1. Query Throughput (QPS by Status and Scenario) + +Tracks the rate of queries executed per second: + +```promql +sum by (scenario, status) (rate(workload_googleapis_com:bigquery_queries_total[1m])) +``` + +#### 2. Under-the-Hood Job Retries (Retry Rate & Count) + +Tracks how often queries triggered a backend job retry (where the job ID mutated +between `Query::send()` and `Query::until_done()` - needs to be improved to +detect when queries are retried on creation): + +```promql +sum by (scenario) (rate(workload_googleapis_com:bigquery_queries_retries_detected[1m])) +``` + +To calculate the **Percentage of Queries Retried**: + +```promql +100 * sum(rate(workload_googleapis_com:bigquery_queries_retries_detected[1m])) + / sum(rate(workload_googleapis_com:bigquery_queries_total[1m])) +``` + +#### 3. Error Rate (%) + +Tracks the percentage of query executions that failed: + +```promql +100 * sum(rate(workload_googleapis_com:bigquery_queries_error[1m])) + / sum(rate(workload_googleapis_com:bigquery_queries_total[1m])) +``` + +#### 4. End-to-End Query Latency Percentiles (P50, P90, P99) + +Calculates the 50th, 90th, and 99th percentile query latencies from histogram +buckets: + +```promql +# P99 Latency (seconds) +histogram_quantile(0.99, sum by (le, scenario) (rate(workload_googleapis_com:bigquery_queries_duration_seconds_bucket[1m]))) + +# P90 Latency (seconds) +histogram_quantile(0.90, sum by (le, scenario) (rate(workload_googleapis_com:bigquery_queries_duration_seconds_bucket[1m]))) + +# P50 (Median) Latency (seconds) +histogram_quantile(0.50, sum by (le, scenario) (rate(workload_googleapis_com:bigquery_queries_duration_seconds_bucket[1m]))) +``` + +#### 5. Query Phase Latency Breakdown (P95 comparison) + +Compare where time is spent across `send()`, `until_done()`, and `read()`: + +```promql +# Query::send() P95 Latency +histogram_quantile(0.95, sum by (le, scenario) (rate(workload_googleapis_com:bigquery_queries_send_duration_seconds_bucket[1m]))) + +# Query::until_done() P95 Polling Latency +histogram_quantile(0.95, sum by (le, scenario) (rate(workload_googleapis_com:bigquery_queries_poll_duration_seconds_bucket[1m]))) + +# CompleteQuery::read() P95 Row Streaming Latency +histogram_quantile(0.95, sum by (le, scenario) (rate(workload_googleapis_com:bigquery_queries_read_duration_seconds_bucket[1m]))) +``` + +#### 6. Row and Byte Streaming Throughput + +Tracks data processing throughput (rows/sec and MiB/sec): + +```promql +# Rows streamed per second +sum by (scenario) (rate(workload_googleapis_com:bigquery_queries_rows_read[1m])) + +# MiB processed per second +sum by (scenario) (rate(workload_googleapis_com:bigquery_queries_bytes_processed[1m])) / (1024 * 1024) +``` + +______________________________________________________________________ + +## Uploading Results to BigQuery + +If `--output-dir` is specified, the suite saves per-iteration raw samples to a +CSV file (e.g. `results/samples-synthetic-100k-*.csv`). You can upload the +samples to BigQuery for SQL analysis: + +```shell +bq load --source_format=CSV --skip_leading_rows=1 \ + ${GOOGLE_CLOUD_PROJECT}:benchmark_dataset.bigquery_samples \ + ./results/samples-*.csv \ + Task:int64,Iteration:int64,StartOffsetMicros:int64,SendDurationMicros:int64,PollDurationMicros:int64,ReadDurationMicros:int64,TotalDurationMicros:int64,RowsCount:int64,BytesProcessed:int64,CacheHit:bool,InitialJobId:string,FinalJobId:string,RetryDetected:bool,Status:string,ErrorMessage:string +``` diff --git a/src/bigquery/benchmarks/queries/src/args.rs b/src/bigquery/benchmarks/queries/src/args.rs new file mode 100644 index 0000000000..c85d061bb6 --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/args.rs @@ -0,0 +1,163 @@ +// Copyright 2026 Google LLC +// +// 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 +// +// https://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. + +use clap::{Parser, ValueEnum}; +use humantime::parse_duration; +use std::path::PathBuf; +use std::time::Duration; + +const DESCRIPTION: &str = concat!( + "A benchmark and endurance test runner for the Rust BigQuery client library.\n\n", + "Supports fixed-iteration benchmarking (measuring P50/P90/P99 latencies) and\n", + "long-running endurance tests with OpenTelemetry trace and metric export." +); + +/// Preset query scenarios for benchmarking. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, ValueEnum)] +pub enum ScenarioName { + /// Zero-dependency synthetic query generating 100,000 structured rows via UNNEST(GENERATE_ARRAY(1, 100000)). + #[value(name = "synthetic-100k")] + #[default] + Synthetic100k, + /// Zero-dependency synthetic query generating 10,000 structured rows. + #[value(name = "synthetic-10k")] + Synthetic10k, + /// Scans up to 50,000 rows from `bigquery-public-data.usa_names.usa_1910_2013`. + #[value(name = "usa-names-scan")] + UsaNamesScan, + /// Aggregation query grouping 5.5M rows by state and gender from `bigquery-public-data.usa_names.usa_1910_2013`. + #[value(name = "usa-names-agg")] + UsaNamesAgg, + /// Aggregation query on `bigquery-public-data.samples.wikipedia`. + #[value(name = "wikipedia-agg")] + WikipediaAgg, + /// Custom SQL query provided via `--sql` or `--sql-file`. + #[value(name = "custom")] + Custom, +} + +/// Runs the BigQuery benchmark and endurance test suite. +#[derive(Clone, Debug, Parser)] +#[command(version, about, long_about = DESCRIPTION)] +pub struct Args { + /// The Google Cloud Project ID used for BigQuery billing and OpenTelemetry export. + /// + /// If not provided, the default project from Application Default Credentials or + /// the `GOOGLE_CLOUD_PROJECT` environment variable will be used. + #[arg(long, env = "GOOGLE_CLOUD_PROJECT")] + pub project_id: Option, + + /// The geographic location for BigQuery datasets and job execution (e.g. `US`, `EU`). + #[arg(long, default_value = "US")] + pub location: String, + + /// The query workload scenario to run. + #[arg(long, value_enum, default_value_t = ScenarioName::Synthetic100k)] + pub scenario: ScenarioName, + + /// Custom SQL query string (required if `--scenario custom` and `--sql-file` is not set). + #[arg(long)] + pub sql: Option, + + /// Path to a file containing custom SQL query. + #[arg(long)] + pub sql_file: Option, + + /// Number of concurrent worker tasks running query loops. + #[arg(long, default_value_t = 1)] + pub task_count: usize, + + /// Number of query iterations per worker task. + /// + /// If neither `--iterations` nor `--duration` is set, the benchmark runs indefinitely until interrupted (Ctrl+C). + #[arg(long)] + pub iterations: Option, + + /// Total runtime duration for the benchmark or endurance test (e.g. "30s", "10m", "2h"). + /// + /// When set, tasks will continue executing queries until the duration expires. + #[arg(long, value_parser = parse_duration)] + pub duration: Option, + + /// The maximum number of rows per page returned from BigQuery. + #[arg(long)] + pub page_size: Option, + + /// Whether to consume all returned rows by iterating over the stream (`read().next().await`). + #[arg(long, default_value_t = true, action = clap::ArgAction::Set)] + pub read_results: bool, + + /// Whether to enable BigQuery query results cache. Defaults to false so queries always execute against storage. + #[arg(long, default_value_t = false, action = clap::ArgAction::Set)] + pub use_query_cache: bool, + + /// Ramp-up delay between spawning subsequent worker tasks to prevent thundering herd. + #[arg(long, value_parser = parse_duration, default_value = "250ms")] + pub rampup_period: Duration, + + /// Timeout for an individual query iteration (including send, until_done, and row streaming). + #[arg(long, value_parser = parse_duration, default_value = "120s")] + pub query_timeout: Duration, + + /// Directory where raw CSV samples and summary JSON metrics will be written. + #[arg(long)] + pub output_dir: Option, + + /// Custom OpenTelemetry collector endpoint for traces and metrics. + /// + /// Defaults to `https://telemetry.googleapis.com` if `--project-id` is provided. + #[arg(long)] + pub otlp_endpoint: Option, +} + +impl Args { + /// Validates the command line arguments. + pub fn validate(&self) -> anyhow::Result<()> { + if self.task_count == 0 { + anyhow::bail!("--task-count must be at least 1"); + } + if self.scenario == ScenarioName::Custom && self.sql.is_none() && self.sql_file.is_none() { + anyhow::bail!( + "When using `--scenario custom`, either `--sql` or `--sql-file` must be provided." + ); + } + // Rejected even when `--duration` is set: the task loop breaks on + // `iteration >= 0` before running anything, producing an empty report. + if self.iterations == Some(0) { + anyhow::bail!("--iterations must be greater than 0"); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_zero_iterations_rejected() { + let args = Args::parse_from(["bigquery-benchmark-queries", "--iterations", "0"]); + assert!(args.validate().is_err()); + + // Also rejected alongside --duration: the task loop would exit immediately. + let args = Args::parse_from([ + "bigquery-benchmark-queries", + "--iterations", + "0", + "--duration", + "5m", + ]); + assert!(args.validate().is_err()); + } +} diff --git a/src/bigquery/benchmarks/queries/src/main.rs b/src/bigquery/benchmarks/queries/src/main.rs new file mode 100644 index 0000000000..bde5f1b6b7 --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/main.rs @@ -0,0 +1,157 @@ +// Copyright 2026 Google LLC +// +// 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 +// +// https://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. + +//! BigQuery query benchmark and endurance test runner. + +mod args; +mod metrics; +mod reporter; +mod runner; +mod sample; +mod scenarios; +mod telemetry; + +use args::Args; +use clap::Parser; +use google_cloud_auth::credentials::Builder as CredentialsBuilder; +use google_cloud_bigquery::client::BigQuery; +use metrics::OtelMetrics; +use scenarios::Scenario; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::task::JoinSet; + +/// How often the Tokio runtime metrics are logged. +const RUNTIME_MONITOR_INTERVAL: Duration = Duration::from_secs(5); + +/// Samples buffered per worker task before senders start blocking. +const CHANNEL_CAPACITY_PER_TASK: usize = 1024; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + tracing_log::LogTracer::init()?; + let args = Args::parse(); + args.validate()?; + + let scenario = Scenario::resolve(&args)?; + let credentials = CredentialsBuilder::default().build()?; + let telemetry_guard = telemetry::enable_telemetry(&args, &credentials).await?; + + tracing::info!( + scenario = %scenario.name, + description = %scenario.description, + task_count = args.task_count, + iterations = ?args.iterations, + duration = ?args.duration, + use_query_cache = args.use_query_cache, + "Starting BigQuery benchmark" + ); + + // Spawn periodic runtime monitor. Query counts are reported by the reporter + // and exported as OpenTelemetry metrics. + let handle = tokio::runtime::Handle::current(); + let runtime_monitor = tokio_metrics::RuntimeMonitor::new(&handle); + tokio::spawn(async move { + for metrics in runtime_monitor.intervals() { + tracing::info!("RuntimeMetrics = {:?}", metrics); + tokio::time::sleep(RUNTIME_MONITOR_INTERVAL).await; + } + }); + + let mut client_builder = BigQuery::builder() + .with_credentials(credentials.clone()) + .with_tracing(); + + if let Some(project_id) = &args.project_id { + client_builder = client_builder.with_project_id(project_id); + } + + let client = client_builder.build().await?; + let otel_metrics = OtelMetrics::new(scenario.name); + + let (tx, rx) = tokio::sync::mpsc::channel(CHANNEL_CAPACITY_PER_TASK * args.task_count); + + // Shared by every worker task; the SQL text and CLI options are read-only. + let args = Arc::new(args); + let scenario = Arc::new(scenario); + + // Spawn reporter in background to process samples as they arrive + let reporter_scenario = scenario.clone(); + let reporter_args = args.clone(); + let reporter_handle = tokio::spawn(async move { + reporter::collect_and_report(rx, &reporter_scenario, &reporter_args).await + }); + + let test_start = Instant::now(); + let mut tasks = JoinSet::new(); + + for task_id in 0..args.task_count { + let task_client = client.clone(); + let task_scenario = scenario.clone(); + let task_args = args.clone(); + let task_tx = tx.clone(); + let task_metrics = otel_metrics.clone(); + + tasks.spawn(async move { + let runner = runner::TaskRunner { + task_id, + test_start, + client: &task_client, + scenario: &task_scenario, + args: &task_args, + tx: &task_tx, + metrics: &task_metrics, + }; + let result = runner.run().await; + (task_id, result) + }); + } + + // Drop main sender so receiver terminates after all tasks complete + drop(tx); + + tokio::select! { + _ = tokio::signal::ctrl_c() => { + tracing::warn!("Ctrl+C received, stopping benchmark tasks and generating report..."); + tasks.abort_all(); + while (tasks.join_next().await).is_some() {} + } + _ = async { + while let Some(res) = tasks.join_next().await { + match res { + Ok((task_id, Ok(_))) => { + tracing::debug!(task_id, "Task worker completed successfully"); + } + Ok((task_id, Err(err))) => { + tracing::error!(task_id, "Task worker encountered error: {err:?}"); + } + Err(err) => { + tracing::error!("Failed to join task: {err:?}"); + } + } + } + } => {} + } + + // Wait for reporter to finish outputting summary and files + match reporter_handle.await { + Ok(Ok(_)) => {} + Ok(Err(err)) => tracing::error!("Reporter failed: {err:?}"), + Err(err) => tracing::error!("Reporter task panicked: {err:?}"), + } + + telemetry_guard.shutdown(); + + Ok(()) +} diff --git a/src/bigquery/benchmarks/queries/src/metrics.rs b/src/bigquery/benchmarks/queries/src/metrics.rs new file mode 100644 index 0000000000..5b679a008c --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/metrics.rs @@ -0,0 +1,238 @@ +// Copyright 2026 Google LLC +// +// 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 +// +// https://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. + +use crate::sample::{Sample, SampleStatus}; +use opentelemetry::KeyValue; +use opentelemetry::metrics::{Counter, Histogram}; +use serde::{Deserialize, Serialize}; +use std::time::Duration; + +/// Serializes a [`Duration`] as fractional milliseconds. +/// +/// The serde default for [`Duration`] emits `{"secs": 1, "nanos": 234000000}`, +/// which is awkward to chart or load into BigQuery. +mod duration_millis { + use serde::{Deserialize, Deserializer, Serializer}; + use std::time::Duration; + + pub fn serialize(duration: &Duration, ser: S) -> Result { + ser.serialize_f64(duration.as_secs_f64() * 1_000.0) + } + + pub fn deserialize<'de, D: Deserializer<'de>>(de: D) -> Result { + let millis = f64::deserialize(de)?; + Ok(Duration::from_secs_f64(millis / 1_000.0)) + } +} + +/// Summary percentiles and metrics for execution latencies. +/// +/// Durations are serialized as fractional milliseconds. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct LatencySummary { + #[serde(with = "duration_millis", rename = "min_millis")] + pub min: Duration, + #[serde(with = "duration_millis", rename = "max_millis")] + pub max: Duration, + #[serde(with = "duration_millis", rename = "mean_millis")] + pub mean: Duration, + #[serde(with = "duration_millis", rename = "p50_millis")] + pub p50: Duration, + #[serde(with = "duration_millis", rename = "p90_millis")] + pub p90: Duration, + #[serde(with = "duration_millis", rename = "p99_millis")] + pub p99: Duration, + pub count: usize, +} + +/// Computes statistical metrics (min, max, mean, p50, p90, p99) from a slice of latencies. +pub fn compute_metrics(latencies: &[Duration]) -> Option { + if latencies.is_empty() { + return None; + } + + let mut sorted = latencies.to_vec(); + sorted.sort(); + + let count = sorted.len(); + let min = sorted[0]; + let max = sorted[count - 1]; + let sum: Duration = sorted.iter().sum(); + let mean = sum / count as u32; + + let p50 = sorted[((count - 1) as f64 * 0.50).round() as usize]; + let p90 = sorted[((count - 1) as f64 * 0.90).round() as usize]; + let p99 = sorted[((count - 1) as f64 * 0.99).round() as usize]; + + Some(LatencySummary { + min, + max, + mean, + p50, + p90, + p99, + count, + }) +} + +/// OpenTelemetry metrics instruments for a single scenario. +/// +/// The scenario/status attribute sets are built once and reused, so recording a +/// sample allocates nothing. +#[derive(Clone)] +pub struct OtelMetrics { + queries_total: Counter, + queries_success: Counter, + queries_error: Counter, + queries_retried: Counter, + rows_read: Counter, + bytes_processed: Counter, + query_duration: Histogram, + send_duration: Histogram, + poll_duration: Histogram, + read_duration: Histogram, + ok_attrs: [KeyValue; 2], + error_attrs: [KeyValue; 2], +} + +impl OtelMetrics { + /// Creates the instruments for `scenario`. + /// + /// Every counter is seeded with 0 so the time series exist in Cloud + /// Monitoring even if no errors or retries occur during the run. + pub fn new(scenario: &str) -> Self { + let meter = opentelemetry::global::meter("bigquery-benchmark-queries"); + + let metrics = Self { + queries_total: meter + .u64_counter("bigquery.queries.total") + .with_description("Total number of BigQuery queries attempted") + .build(), + queries_success: meter + .u64_counter("bigquery.queries.success") + .with_description("Number of BigQuery queries completed successfully") + .build(), + queries_error: meter + .u64_counter("bigquery.queries.error") + .with_description("Number of BigQuery queries that failed") + .build(), + queries_retried: meter + .u64_counter("bigquery.queries.retries_detected") + .with_description( + "Number of queries where an under-the-hood job retry was detected", + ) + .build(), + rows_read: meter + .u64_counter("bigquery.queries.rows_read") + .with_description("Total count of rows read from query results") + .build(), + bytes_processed: meter + .u64_counter("bigquery.queries.bytes_processed") + .with_description("Total estimated bytes processed by BigQuery jobs") + .build(), + query_duration: meter + .f64_histogram("bigquery.queries.duration_seconds") + .with_description("Total query end-to-end duration in seconds") + .build(), + send_duration: meter + .f64_histogram("bigquery.queries.send_duration_seconds") + .with_description("Duration for Query::send() execution") + .build(), + poll_duration: meter + .f64_histogram("bigquery.queries.poll_duration_seconds") + .with_description("Duration for Query::until_done() polling execution") + .build(), + read_duration: meter + .f64_histogram("bigquery.queries.read_duration_seconds") + .with_description("Duration for CompleteQuery::read() row streaming") + .build(), + ok_attrs: [ + KeyValue::new("scenario", scenario.to_string()), + KeyValue::new("status", "ok"), + ], + error_attrs: [ + KeyValue::new("scenario", scenario.to_string()), + KeyValue::new("status", "error"), + ], + }; + + metrics.queries_total.add(0, &metrics.ok_attrs); + metrics.queries_total.add(0, &metrics.error_attrs); + metrics.queries_success.add(0, &metrics.ok_attrs); + metrics.queries_error.add(0, &metrics.error_attrs); + metrics.queries_retried.add(0, &metrics.ok_attrs); + metrics.rows_read.add(0, &metrics.ok_attrs); + metrics.bytes_processed.add(0, &metrics.ok_attrs); + + metrics + } + + /// Records a single completed query execution. + pub fn record_sample(&self, sample: &Sample) { + let is_ok = sample.status == SampleStatus::Ok; + let attrs = if is_ok { + &self.ok_attrs + } else { + &self.error_attrs + }; + + self.queries_total.add(1, attrs); + if is_ok { + self.queries_success.add(1, attrs); + if sample.retry_detected { + self.queries_retried.add(1, attrs); + } + self.rows_read.add(sample.rows_count as u64, attrs); + self.bytes_processed + .add(sample.bytes_processed.max(0) as u64, attrs); + + self.send_duration + .record(sample.send_duration().as_secs_f64(), attrs); + self.poll_duration + .record(sample.poll_duration().as_secs_f64(), attrs); + self.read_duration + .record(sample.read_duration().as_secs_f64(), attrs); + } else { + self.queries_error.add(1, attrs); + } + + self.query_duration + .record(sample.total_duration().as_secs_f64(), attrs); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_compute_metrics() { + assert!(compute_metrics(&[]).is_none()); + + let single = compute_metrics(&[Duration::from_millis(100)]).unwrap(); + assert_eq!(single.min, Duration::from_millis(100)); + assert_eq!(single.max, Duration::from_millis(100)); + assert_eq!(single.p50, Duration::from_millis(100)); + assert_eq!(single.p90, Duration::from_millis(100)); + assert_eq!(single.p99, Duration::from_millis(100)); + + let hundred: Vec = (1..=100).map(Duration::from_millis).collect(); + let summary = compute_metrics(&hundred).unwrap(); + assert_eq!(summary.min, Duration::from_millis(1)); + assert_eq!(summary.max, Duration::from_millis(100)); + assert_eq!(summary.p50, Duration::from_millis(51)); + assert_eq!(summary.p90, Duration::from_millis(90)); + assert_eq!(summary.p99, Duration::from_millis(99)); + } +} diff --git a/src/bigquery/benchmarks/queries/src/reporter.rs b/src/bigquery/benchmarks/queries/src/reporter.rs new file mode 100644 index 0000000000..703731df0b --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/reporter.rs @@ -0,0 +1,376 @@ +// Copyright 2026 Google LLC +// +// 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 +// +// https://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. + +use crate::args::Args; +use crate::metrics::{self, LatencySummary}; +use crate::sample::{Sample, SampleStatus, display_job_id}; +use crate::scenarios::Scenario; +use serde::{Deserialize, Serialize}; +use std::fs::File; +use std::io::{BufWriter, Write}; +use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use tokio::sync::mpsc::Receiver; + +/// Maximum number of error details retained in memory for the final report. +/// +/// Every error is still written to the error log; this only bounds the +/// in-memory copy so that long endurance runs cannot exhaust memory. +const MAX_RETAINED_ERRORS: usize = 1_000; + +/// Number of errors printed in the stdout summary. +const ERRORS_PRINTED: usize = 20; + +/// How often the interim summary JSON is rebuilt, and the ceiling it backs off to. +/// +/// Rebuilding sorts every retained latency, so the interval doubles as the run +/// grows to keep that cost from crowding out sample collection. +const MIN_SUMMARY_INTERVAL: Duration = Duration::from_secs(5); +const MAX_SUMMARY_INTERVAL: Duration = Duration::from_secs(60); + +/// Details of a single query error for diagnostics. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ErrorDetail { + pub task_id: usize, + pub iteration: u64, + pub offset_secs: f64, + pub initial_job_id: String, + pub final_job_id: String, + pub error_message: String, +} + +/// Structured benchmark summary report. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct BenchmarkReport { + pub scenario: String, + pub task_count: usize, + pub total_samples: usize, + pub success_count: usize, + pub error_count: usize, + pub retries_detected_count: usize, + pub total_rows_read: usize, + pub total_bytes_processed: i64, + pub total_duration: Option, + pub send_duration: Option, + pub poll_duration: Option, + pub read_duration: Option, + /// The first [`MAX_RETAINED_ERRORS`] errors; see `error_count` for the total. + pub errors: Vec, +} + +impl BenchmarkReport { + /// Prints the benchmark summary report to stdout. + pub fn print_stdout(&self) { + println!("\n======================================================="); + println!(" BigQuery Benchmark Report "); + println!("======================================================="); + println!("Scenario: {}", self.scenario); + println!("Task Count: {}", self.task_count); + println!("Total Queries Executed: {}", self.total_samples); + println!("Successful Queries: {}", self.success_count); + println!("Errors: {}", self.error_count); + println!("Job Retries Detected: {}", self.retries_detected_count); + println!("Total Rows Read: {}", self.total_rows_read); + println!("Total Bytes Processed: {}", self.total_bytes_processed); + + if let Some(total) = &self.total_duration { + println!("\n--- End-to-End Query Latency ---"); + println!(" Min: {:?}", total.min); + println!(" Mean: {:?}", total.mean); + println!(" P50: {:?}", total.p50); + println!(" P90: {:?}", total.p90); + println!(" P99: {:?}", total.p99); + println!(" Max: {:?}", total.max); + } + + for (label, summary) in [ + ("Query::send()", &self.send_duration), + ("Query::until_done() Polling", &self.poll_duration), + ("CompleteQuery::read() Streaming", &self.read_duration), + ] { + if let Some(s) = summary { + println!("\n--- {label} Latency ---"); + println!(" P50: {:?} | P90: {:?} | P99: {:?}", s.p50, s.p90, s.p99); + } + } + + if !self.errors.is_empty() { + println!("\n-------------------------------------------------------"); + println!(" QUERY ERRORS ({} total)", self.error_count); + println!("-------------------------------------------------------"); + for (idx, err) in self.errors.iter().take(ERRORS_PRINTED).enumerate() { + println!( + " {}. Task {:>2} | Iteration {:>6} | Offset: {:>7.1}s | Job: {}", + idx + 1, + err.task_id, + err.iteration, + err.offset_secs, + display_job_id(&err.initial_job_id, &err.final_job_id) + ); + println!(" Error: {}", err.error_message); + } + let shown = self.errors.len().min(ERRORS_PRINTED); + if self.error_count > shown { + println!( + " ... and {} more error(s) recorded in full error log.", + self.error_count - shown + ); + } + } + + println!("=======================================================\n"); + } +} + +/// Running totals over the samples received so far. +#[derive(Default)] +struct Accumulator { + total_samples: usize, + success_count: usize, + error_count: usize, + retries_detected_count: usize, + total_rows_read: usize, + total_bytes_processed: i64, + total_durations: Vec, + send_durations: Vec, + poll_durations: Vec, + read_durations: Vec, + errors: Vec, +} + +impl Accumulator { + fn push(&mut self, sample: &Sample) { + self.total_samples += 1; + + if sample.status == SampleStatus::Ok { + self.success_count += 1; + self.total_durations.push(sample.total_duration()); + self.send_durations.push(sample.send_duration()); + self.poll_durations.push(sample.poll_duration()); + self.read_durations.push(sample.read_duration()); + self.total_rows_read += sample.rows_count; + self.total_bytes_processed += sample.bytes_processed; + } else { + self.error_count += 1; + if self.errors.len() < MAX_RETAINED_ERRORS { + self.errors.push(ErrorDetail { + task_id: sample.task_id, + iteration: sample.iteration, + offset_secs: sample.start_offset_secs(), + initial_job_id: sample.initial_job_id.clone(), + final_job_id: sample.final_job_id.clone(), + error_message: sample.error_message.clone(), + }); + } + } + + if sample.retry_detected { + self.retries_detected_count += 1; + } + } + + fn build_report(&self, scenario: &str, task_count: usize) -> BenchmarkReport { + BenchmarkReport { + scenario: scenario.to_string(), + task_count, + total_samples: self.total_samples, + success_count: self.success_count, + error_count: self.error_count, + retries_detected_count: self.retries_detected_count, + total_rows_read: self.total_rows_read, + total_bytes_processed: self.total_bytes_processed, + total_duration: metrics::compute_metrics(&self.total_durations), + send_duration: metrics::compute_metrics(&self.send_durations), + poll_duration: metrics::compute_metrics(&self.poll_durations), + read_duration: metrics::compute_metrics(&self.read_durations), + errors: self.errors.clone(), + } + } +} + +/// The set of files written when `--output-dir` is provided. +struct OutputFiles { + csv: BufWriter, + csv_path: PathBuf, + json_path: PathBuf, + errors: BufWriter, + errors_path: PathBuf, +} + +impl OutputFiles { + fn create(output_dir: &Path, scenario: &str) -> anyhow::Result { + std::fs::create_dir_all(output_dir)?; + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0); + + let csv_path = output_dir.join(format!("samples-{scenario}-{timestamp}.csv")); + let json_path = output_dir.join(format!("summary-{scenario}-{timestamp}.json")); + let errors_path = output_dir.join(format!("errors-{scenario}-{timestamp}.log")); + + let mut csv = BufWriter::new(File::create(&csv_path)?); + writeln!(csv, "{}", Sample::HEADER)?; + csv.flush()?; + + let mut errors = BufWriter::new(File::create(&errors_path)?); + writeln!( + errors, + "# BigQuery Benchmark Error Log - Scenario: {scenario}, Timestamp: {timestamp}" + )?; + errors.flush()?; + + println!("Writing real-time samples to: {}", csv_path.display()); + println!("Writing real-time summary to: {}", json_path.display()); + println!("Writing error details to: {}", errors_path.display()); + + Ok(Self { + csv, + csv_path, + json_path, + errors, + errors_path, + }) + } + + fn write_sample(&mut self, sample: &Sample) { + if let Err(err) = writeln!(self.csv, "{}", sample.to_csv_row()) { + tracing::error!("Failed to write CSV sample row to disk: {err:?}"); + } + } + + fn write_error(&mut self, sample: &Sample) { + let _ = writeln!( + self.errors, + "Task: {}\nIteration: {}\nOffsetSecs: {:.3}\nInitialJobId: {}\nFinalJobId: {}\nError: {}\n{}", + sample.task_id, + sample.iteration, + sample.start_offset_secs(), + sample.initial_job_id, + sample.final_job_id, + sample.error_message, + "-".repeat(80) + ); + let _ = self.errors.flush(); + } + + fn write_summary(&mut self, report: &BenchmarkReport) { + let _ = self.csv.flush(); + let _ = self.errors.flush(); + match File::create(&self.json_path) { + Ok(file) => { + if let Err(err) = serde_json::to_writer_pretty(file, report) { + tracing::error!("Failed to write summary JSON: {err:?}"); + } + } + Err(err) => tracing::error!("Failed to create summary JSON: {err:?}"), + } + } +} + +/// Receives sample results, logs real-time output, and generates the final report. +pub async fn collect_and_report( + mut rx: Receiver, + scenario: &Scenario, + args: &Args, +) -> anyhow::Result { + let mut acc = Accumulator::default(); + let mut files = args + .output_dir + .as_deref() + .map(|dir| OutputFiles::create(dir, scenario.name)) + .transpose()?; + + let mut last_summary = Instant::now(); + let mut summary_interval = MIN_SUMMARY_INTERVAL; + + while let Some(sample) = rx.recv().await { + acc.push(&sample); + + if sample.status != SampleStatus::Ok { + // Loud alert to console immediately + eprintln!( + "\n🚨 [QUERY FAILURE] Task {:>2} | Iteration {:>6} | Offset: {:>7.1}s | Job: {} | Error: {}\n", + sample.task_id, + sample.iteration, + sample.start_offset_secs(), + sample.display_job_id(), + sample.error_message + ); + } + + if let Some(files) = &mut files { + files.write_sample(&sample); + if sample.status != SampleStatus::Ok { + files.write_error(&sample); + } + + if last_summary.elapsed() >= summary_interval { + last_summary = Instant::now(); + summary_interval = (summary_interval * 2).min(MAX_SUMMARY_INTERVAL); + files.write_summary(&acc.build_report(scenario.name, args.task_count)); + } + } + } + + let report = acc.build_report(scenario.name, args.task_count); + report.print_stdout(); + + if let Some(files) = &mut files { + files.write_summary(&report); + println!("Final samples saved to: {}", files.csv_path.display()); + println!("Final summary saved to: {}", files.json_path.display()); + if acc.error_count > 0 { + println!("Errors logged to: {}", files.errors_path.display()); + } + } + + Ok(report) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ok_sample(total_micros: u64) -> Sample { + Sample { + total_duration_micros: total_micros, + rows_count: 10, + bytes_processed: 100, + status: SampleStatus::Ok, + ..Sample::new(0, 0, 0) + } + } + + #[test] + fn test_accumulator_tallies_successes_and_errors() { + let mut acc = Accumulator::default(); + acc.push(&ok_sample(1_000)); + acc.push(&ok_sample(3_000)); + acc.push(&Sample { + error_message: "boom".to_string(), + ..Sample::new(1, 0, 0) + }); + + let report = acc.build_report("test", 2); + assert_eq!(report.total_samples, 3); + assert_eq!(report.success_count, 2); + assert_eq!(report.error_count, 1); + assert_eq!(report.total_rows_read, 20); + assert_eq!(report.total_bytes_processed, 200); + assert_eq!(report.errors.len(), 1); + // Latency summaries only cover successful samples. + assert_eq!(report.total_duration.unwrap().count, 2); + } +} diff --git a/src/bigquery/benchmarks/queries/src/runner.rs b/src/bigquery/benchmarks/queries/src/runner.rs new file mode 100644 index 0000000000..5eec9def61 --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/runner.rs @@ -0,0 +1,246 @@ +// Copyright 2026 Google LLC +// +// 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 +// +// https://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. + +use crate::args::Args; +use crate::metrics::OtelMetrics; +use crate::sample::{Sample, SampleStatus, as_micros_u64}; +use crate::scenarios::Scenario; +use google_cloud_bigquery::client::BigQuery; +use std::time::{Duration, Instant}; +use tokio::sync::mpsc::Sender; +use tracing::Instrument; + +/// Manages and executes benchmark queries for a single worker task. +pub struct TaskRunner<'a> { + pub task_id: usize, + pub test_start: Instant, + pub client: &'a BigQuery, + pub scenario: &'a Scenario, + pub args: &'a Args, + pub tx: &'a Sender, + pub metrics: &'a OtelMetrics, +} + +impl TaskRunner<'_> { + /// Executes the task loop until iterations or duration limit is reached. + pub async fn run(&self) -> anyhow::Result<()> { + if self.args.rampup_period > Duration::ZERO { + tokio::time::sleep(self.args.rampup_period * self.task_id as u32).await; + } + + let mut iteration = 0_u64; + + loop { + if let Some(max_iter) = self.args.iterations + && iteration >= max_iter + { + break; + } + + if let Some(max_duration) = self.args.duration + && self.test_start.elapsed() >= max_duration + { + break; + } + + let iter_start = Instant::now(); + let start_offset_micros = as_micros_u64(self.test_start.elapsed()); + + let iteration_span = tracing::info_span!( + "bigquery.query_benchmark.iteration", + task_id = self.task_id, + iteration, + scenario = %self.scenario.name + ); + + let sample_fut = self + .execute_iteration(iteration, start_offset_micros, iter_start) + .instrument(iteration_span); + + let sample = match tokio::time::timeout(self.args.query_timeout, sample_fut).await { + Ok(sample) => sample, + Err(_) => { + tracing::error!( + task_id = self.task_id, + iteration, + "Query iteration timed out after {:?}", + self.args.query_timeout + ); + Sample { + total_duration_micros: as_micros_u64(iter_start.elapsed()), + status: SampleStatus::Timeout, + error_message: format!( + "Query iteration timed out after {:?}", + self.args.query_timeout + ), + ..Sample::new(self.task_id, iteration, start_offset_micros) + } + } + }; + + // Every execution path funnels through here, so no outcome can go + // unrecorded. + self.metrics.record_sample(&sample); + let _ = self.tx.send(sample).await; + iteration += 1; + } + + Ok(()) + } + + /// Runs one query end to end, returning a sample describing the outcome. + /// + /// The sample starts out marked as an error and is upgraded to + /// [`SampleStatus::Ok`] only once every phase has succeeded. + async fn execute_iteration( + &self, + iteration: u64, + start_offset_micros: u64, + iter_start: Instant, + ) -> Sample { + let mut sample = Sample::new(self.task_id, iteration, start_offset_micros); + + let mut query_builder = self + .client + .query(&self.scenario.sql) + .set_location(&self.args.location) + .set_use_query_cache(self.args.use_query_cache); + + if let Some(page_size) = self.args.page_size { + query_builder = query_builder.set_page_size(page_size); + } + if let Some(project_id) = &self.args.project_id { + query_builder = query_builder.with_project_id(project_id); + } + + // Step 1: Execute Query::send() + let send_start = Instant::now(); + let send_span = tracing::info_span!("bigquery.send", task_id = self.task_id, iteration); + let send_result = query_builder.send().instrument(send_span).await; + sample.send_duration_micros = as_micros_u64(send_start.elapsed()); + + let query_handle = match send_result { + Ok(handle) => handle, + Err(err) => { + tracing::error!( + task_id = self.task_id, + iteration, + "Query::send failed: {err:?}" + ); + sample.total_duration_micros = as_micros_u64(iter_start.elapsed()); + sample.error_message = format!("Query::send: {err:#}"); + return sample; + } + }; + + // Capture initial job_id if present + sample.initial_job_id = query_handle + .metadata() + .job_reference + .as_ref() + .map(|r| r.job_id.clone()) + .unwrap_or_default(); + + // Step 2: Execute Query::until_done() + let poll_start = Instant::now(); + let poll_span = tracing::info_span!( + "bigquery.until_done", + task_id = self.task_id, + iteration, + initial_job_id = %sample.initial_job_id + ); + let done_result = query_handle.until_done().instrument(poll_span).await; + sample.poll_duration_micros = as_micros_u64(poll_start.elapsed()); + + let complete_query = match done_result { + Ok(complete) => complete, + Err(err) => { + tracing::error!( + task_id = self.task_id, + iteration, + initial_job_id = %sample.initial_job_id, + "Query::until_done failed: {err:?}" + ); + sample.total_duration_micros = as_micros_u64(iter_start.elapsed()); + sample.error_message = format!("Query::until_done: {err:#}"); + return sample; + } + }; + + // Capture final job_id + let metadata = complete_query.metadata(); + sample.final_job_id = metadata + .job_reference + .as_ref() + .map(|r| r.job_id.clone()) + .unwrap_or_default(); + sample.bytes_processed = metadata.total_bytes_processed.unwrap_or(0); + sample.cache_hit = metadata.cache_hit.unwrap_or(false); + + // Step 3: Detect if under-the-hood job retry occurred + sample.retry_detected = !sample.initial_job_id.is_empty() + && !sample.final_job_id.is_empty() + && sample.initial_job_id != sample.final_job_id; + + if sample.retry_detected { + tracing::warn!( + task_id = self.task_id, + iteration, + initial_job_id = %sample.initial_job_id, + final_job_id = %sample.final_job_id, + "Query job retry detected under the hood (job_id mutated)!" + ); + } + + // Step 4: Stream and read result rows if enabled + let read_start = Instant::now(); + let mut read_error = None; + + if self.args.read_results { + let read_span = + tracing::info_span!("bigquery.read_rows", task_id = self.task_id, iteration); + async { + let mut rows = complete_query.read(); + while let Some(row_result) = rows.next().await { + match row_result { + Ok(_) => { + sample.rows_count += 1; + } + Err(err) => { + tracing::error!( + task_id = self.task_id, + iteration, + "Error streaming rows: {err:?}" + ); + read_error = Some(format!("CompleteQuery::read: {err:#}")); + break; + } + } + } + } + .instrument(read_span) + .await; + } + + sample.read_duration_micros = as_micros_u64(read_start.elapsed()); + sample.total_duration_micros = as_micros_u64(iter_start.elapsed()); + + match read_error { + Some(err_msg) => sample.error_message = err_msg, + None => sample.status = SampleStatus::Ok, + } + + sample + } +} diff --git a/src/bigquery/benchmarks/queries/src/sample.rs b/src/bigquery/benchmarks/queries/src/sample.rs new file mode 100644 index 0000000000..908809b6df --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/sample.rs @@ -0,0 +1,192 @@ +// Copyright 2026 Google LLC +// +// 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 +// +// https://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. + +use serde::{Deserialize, Serialize}; +use std::time::Duration; + +/// Result status of an individual query execution sample. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum SampleStatus { + Ok, + Error, + Timeout, +} + +impl SampleStatus { + pub fn as_str(&self) -> &'static str { + match self { + Self::Ok => "OK", + Self::Error => "ERR", + Self::Timeout => "TIMEOUT", + } + } +} + +/// A recorded sample of a single query execution. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct Sample { + pub task_id: usize, + pub iteration: u64, + pub start_offset_micros: u64, + pub send_duration_micros: u64, + pub poll_duration_micros: u64, + pub read_duration_micros: u64, + pub total_duration_micros: u64, + pub rows_count: usize, + pub bytes_processed: i64, + pub cache_hit: bool, + pub initial_job_id: String, + pub final_job_id: String, + pub retry_detected: bool, + pub status: SampleStatus, + pub error_message: String, +} + +impl Sample { + pub const HEADER: &'static str = concat!( + "Task,Iteration,StartOffsetMicros,SendDurationMicros,PollDurationMicros,", + "ReadDurationMicros,TotalDurationMicros,RowsCount,BytesProcessed,CacheHit,", + "InitialJobId,FinalJobId,RetryDetected,Status,ErrorMessage" + ); + + /// Creates a sample for an iteration that has just started. + /// + /// The status defaults to [`SampleStatus::Error`] so that an execution path + /// which returns early without recording an outcome is never mistaken for a + /// success. Each phase fills in its own fields as it completes. + pub fn new(task_id: usize, iteration: u64, start_offset_micros: u64) -> Self { + Self { + task_id, + iteration, + start_offset_micros, + send_duration_micros: 0, + poll_duration_micros: 0, + read_duration_micros: 0, + total_duration_micros: 0, + rows_count: 0, + bytes_processed: 0, + cache_hit: false, + initial_job_id: String::new(), + final_job_id: String::new(), + retry_detected: false, + status: SampleStatus::Error, + error_message: String::new(), + } + } + + /// Returns the job ID to display in reports, preferring the final one. + pub fn display_job_id(&self) -> &str { + display_job_id(&self.initial_job_id, &self.final_job_id) + } + + /// Returns the offset from the start of the test run, in seconds. + pub fn start_offset_secs(&self) -> f64 { + self.start_offset_micros as f64 / 1_000_000.0 + } + + pub fn to_csv_row(&self) -> String { + let clean_err = self.error_message.replace(',', ";").replace('\n', " "); + let initial_job_id = if self.initial_job_id.is_empty() { + UNKNOWN_JOB_ID + } else { + &self.initial_job_id + }; + let final_job_id = if self.final_job_id.is_empty() { + UNKNOWN_JOB_ID + } else { + &self.final_job_id + }; + format!( + "{},{},{},{},{},{},{},{},{},{},{},{},{},{},{}", + self.task_id, + self.iteration, + self.start_offset_micros, + self.send_duration_micros, + self.poll_duration_micros, + self.read_duration_micros, + self.total_duration_micros, + self.rows_count, + self.bytes_processed, + self.cache_hit, + initial_job_id, + final_job_id, + self.retry_detected, + self.status.as_str(), + clean_err, + ) + } + + pub fn total_duration(&self) -> Duration { + Duration::from_micros(self.total_duration_micros) + } + + pub fn send_duration(&self) -> Duration { + Duration::from_micros(self.send_duration_micros) + } + + pub fn poll_duration(&self) -> Duration { + Duration::from_micros(self.poll_duration_micros) + } + + pub fn read_duration(&self) -> Duration { + Duration::from_micros(self.read_duration_micros) + } +} + +/// Placeholder used in reports when a job ID was never assigned. +pub const UNKNOWN_JOB_ID: &str = "N/A"; + +/// Converts a duration to whole microseconds, saturating at [`u64::MAX`]. +pub fn as_micros_u64(duration: Duration) -> u64 { + u64::try_from(duration.as_micros()).unwrap_or(u64::MAX) +} + +/// Returns the most specific job ID available, preferring the final one. +pub fn display_job_id<'a>(initial: &'a str, final_id: &'a str) -> &'a str { + [final_id, initial] + .into_iter() + .find(|id| !id.is_empty()) + .unwrap_or(UNKNOWN_JOB_ID) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_sample_csv_serialization() { + let sample = Sample { + send_duration_micros: 20_000, + poll_duration_micros: 30_000, + read_duration_micros: 50_000, + total_duration_micros: 100_000, + rows_count: 500, + bytes_processed: 1024, + initial_job_id: "job_init_123".to_string(), + final_job_id: "job_retry_456".to_string(), + retry_detected: true, + status: SampleStatus::Ok, + ..Sample::new(1, 42, 100_000) + }; + + let row = sample.to_csv_row(); + assert!(row.contains("1,42,100000,20000,30000,50000,100000,500,1024,false,job_init_123,job_retry_456,true,OK,")); + } + + #[test] + fn test_new_sample_defaults_to_error() { + // Guards against an early return being reported as a success. + assert_eq!(Sample::new(0, 0, 0).status, SampleStatus::Error); + } +} diff --git a/src/bigquery/benchmarks/queries/src/scenarios.rs b/src/bigquery/benchmarks/queries/src/scenarios.rs new file mode 100644 index 0000000000..45cc1e9ef9 --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/scenarios.rs @@ -0,0 +1,119 @@ +// Copyright 2026 Google LLC +// +// 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 +// +// https://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. + +use crate::args::{Args, ScenarioName}; +use std::fs; + +/// Represents a configured query benchmark scenario. +#[derive(Clone, Debug)] +pub struct Scenario { + pub name: &'static str, + pub sql: String, + pub description: &'static str, +} + +impl Scenario { + /// Resolves the query scenario based on the provided CLI arguments. + pub fn resolve(args: &Args) -> anyhow::Result { + let (name, sql, description) = match args.scenario { + ScenarioName::Synthetic100k => ( + "synthetic-100k", + synthetic_sql(100_000), + "Generates 100,000 structured rows in-flight with no external table dependency.", + ), + ScenarioName::Synthetic10k => ( + "synthetic-10k", + synthetic_sql(10_000), + "Generates 10,000 structured rows in-flight with no external table dependency.", + ), + ScenarioName::UsaNamesScan => ( + "usa-names-scan", + concat!( + "SELECT name, state, year, gender, number ", + "FROM `bigquery-public-data.usa_names.usa_1910_2013` ", + "WHERE year >= 2000 ", + "LIMIT 50000" + ) + .to_string(), + "Scans and retrieves 50,000 rows from the USA names public dataset.", + ), + ScenarioName::UsaNamesAgg => ( + "usa-names-agg", + concat!( + "SELECT state, gender, SUM(number) AS total_count ", + "FROM `bigquery-public-data.usa_names.usa_1910_2013` ", + "GROUP BY state, gender ", + "ORDER BY total_count DESC" + ) + .to_string(), + "Aggregates 5.5M rows grouped by state and gender.", + ), + ScenarioName::WikipediaAgg => ( + "wikipedia-agg", + concat!( + "SELECT title, SUM(views) AS total_views ", + "FROM `bigquery-public-data.samples.wikipedia` ", + "WHERE wp_namespace = 0 ", + "GROUP BY title ", + "ORDER BY total_views DESC ", + "LIMIT 1000" + ) + .to_string(), + "Aggregates top 1000 article views from Wikipedia public samples.", + ), + ScenarioName::Custom => ( + "custom", + custom_sql(args)?, + "User-defined custom SQL query.", + ), + }; + + Ok(Self { + name, + sql, + description, + }) + } +} + +/// Builds a zero-dependency query that generates `rows` structured rows in-flight. +fn synthetic_sql(rows: u64) -> String { + format!( + concat!( + "SELECT ", + " x AS row_id, ", + " GENERATE_UUID() AS uuid, ", + " REPEAT('abcdefghij', 10) AS payload ", + "FROM UNNEST(GENERATE_ARRAY(1, {})) AS x" + ), + rows + ) +} + +/// Reads the user-supplied SQL from `--sql` or `--sql-file`. +fn custom_sql(args: &Args) -> anyhow::Result { + if let Some(sql) = &args.sql { + return Ok(sql.clone()); + } + let Some(sql_file) = &args.sql_file else { + anyhow::bail!("Custom scenario requires --sql or --sql-file"); + }; + fs::read_to_string(sql_file).map_err(|e| { + anyhow::anyhow!( + "Failed to read custom SQL file {}: {}", + sql_file.display(), + e + ) + }) +} diff --git a/src/bigquery/benchmarks/queries/src/telemetry.rs b/src/bigquery/benchmarks/queries/src/telemetry.rs new file mode 100644 index 0000000000..0e43923852 --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/telemetry.rs @@ -0,0 +1,179 @@ +// Copyright 2026 Google LLC +// +// 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 +// +// https://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. + +use crate::args::Args; +use google_cloud_auth::credentials::Credentials; +use integration_tests_o11y::detector::GoogleCloudResourceDetector; +use integration_tests_o11y::otlp::Uri; +use opentelemetry::KeyValue; +use opentelemetry_sdk::Resource; +use opentelemetry_sdk::metrics::SdkMeterProvider; +use opentelemetry_sdk::resource::ResourceDetector; +use opentelemetry_sdk::trace::SdkTracerProvider; +use std::str::FromStr; +use tracing_subscriber::fmt::format::FmtSpan; +use tracing_subscriber::prelude::*; +use uuid::Uuid; + +const SERVICE_NAME: &str = "bigquery-benchmark-queries"; +const DEFAULT_MONITORING_REGION: &str = "us-central1"; + +#[derive(Clone, Debug)] +struct GenericNodeDetector { + id: String, + location: String, +} + +impl GenericNodeDetector { + pub fn new() -> Self { + Self { + id: Uuid::new_v4().to_string(), + // Cloud Monitoring resolves `generic_node.location` to a GCP region + // or zone. + location: DEFAULT_MONITORING_REGION.to_string(), + } + } +} + +impl ResourceDetector for GenericNodeDetector { + fn detect(&self) -> Resource { + Resource::builder_empty() + .with_attributes([ + KeyValue::new("location", self.location.clone()), + KeyValue::new("namespace", SERVICE_NAME), + KeyValue::new("node_id", self.id.clone()), + ]) + .build() + } +} + +/// Holds providers that need graceful flush and shutdown upon completion. +/// +/// Both providers are installed together, or not at all when no project ID is +/// configured. +pub struct TelemetryGuard { + providers: Option<(SdkTracerProvider, SdkMeterProvider)>, +} + +impl TelemetryGuard { + /// Flushes and shuts down telemetry providers. + pub fn shutdown(self) { + let Some((tracer_provider, meter_provider)) = self.providers else { + return; + }; + + if let Err(e) = tracer_provider.shutdown() { + eprintln!("Error shutting down trace provider: {e:?}"); + } + if let Err(e) = meter_provider.shutdown() { + eprintln!("Error shutting down meter provider: {e:?}"); + } + } +} + +/// Initializes tracing subscriber, OpenTelemetry distributed tracing, and metrics export. +pub async fn enable_telemetry( + args: &Args, + credentials: &Credentials, +) -> anyhow::Result { + let env_filter = tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")); + + let fmt_layer = tracing_subscriber::fmt::layer() + .with_level(true) + .with_thread_ids(true) + .with_span_events(FmtSpan::NONE) + .with_writer(std::io::stderr) + .with_filter(env_filter); + + let error_file_layer = if let Some(output_dir) = &args.output_dir { + let _ = std::fs::create_dir_all(output_dir); + let error_log_path = output_dir.join("tracing-errors.log"); + match std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&error_log_path) + { + Ok(file) => { + let layer = tracing_subscriber::fmt::layer() + .with_level(true) + .with_thread_ids(true) + .with_ansi(false) + .with_writer(std::sync::Mutex::new(file)) + .with_filter(tracing_subscriber::filter::LevelFilter::ERROR); + Some(layer) + } + Err(e) => { + eprintln!("Could not open tracing error log file: {e:?}"); + None + } + } + } else { + None + }; + + let registry = tracing_subscriber::Registry::default() + .with(fmt_layer) + .with(error_file_layer); + + if let Some(project_id) = &args.project_id { + tracing::info!("Enabling OpenTelemetry Cloud Trace & Monitoring for project {project_id}"); + + let node = GenericNodeDetector::new(); + let detector = GoogleCloudResourceDetector::builder() + .with_fallback(node.detect()) + .build() + .await?; + + let mut trace_builder = + integration_tests_o11y::otlp::trace::Builder::new(project_id, SERVICE_NAME) + .with_credentials(credentials.clone()) + .with_detector(detector); + + let mut meter_builder = + integration_tests_o11y::otlp::metrics::Builder::new(project_id, SERVICE_NAME) + .with_credentials(credentials.clone()) + .with_detector(node); + + if let Some(endpoint_str) = &args.otlp_endpoint { + let uri = Uri::from_str(endpoint_str)?; + trace_builder = trace_builder.with_endpoint(endpoint_str.clone()); + meter_builder = meter_builder.with_endpoint(uri); + } + + let tracer_provider = trace_builder + .build() + .await + .inspect_err(|e| eprintln!("Failed to create tracer provider: {e:?}"))?; + + let meter_provider = meter_builder + .build() + .await + .inspect_err(|e| eprintln!("Failed to create meter provider: {e:?}"))?; + + opentelemetry::global::set_meter_provider(meter_provider.clone()); + + let otel_layer = integration_tests_o11y::tracing::trace_layer(tracer_provider.clone()); + tracing::subscriber::set_global_default(registry.with(otel_layer)) + .expect("Setting global subscriber succeeds"); + + return Ok(TelemetryGuard { + providers: Some((tracer_provider, meter_provider)), + }); + } + + tracing::subscriber::set_global_default(registry).expect("Setting global subscriber succeeds"); + + Ok(TelemetryGuard { providers: None }) +}