From 69ef5b0b0801f9997ed217ea7a9b96efc8dbff6a Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Tue, 18 Aug 2026 19:28:12 +0000 Subject: [PATCH 01/11] impl(bigquery): add query benchmark --- Cargo.lock | 26 ++ Cargo.toml | 1 + src/bigquery/benchmarks/queries/Cargo.toml | 49 +++ src/bigquery/benchmarks/queries/README.md | 126 ++++++++ src/bigquery/benchmarks/queries/src/args.rs | 184 +++++++++++ src/bigquery/benchmarks/queries/src/main.rs | 142 +++++++++ .../benchmarks/queries/src/metrics.rs | 267 ++++++++++++++++ .../benchmarks/queries/src/reporter.rs | 174 +++++++++++ src/bigquery/benchmarks/queries/src/runner.rs | 293 ++++++++++++++++++ src/bigquery/benchmarks/queries/src/sample.rs | 125 ++++++++ .../benchmarks/queries/src/scenarios.rs | 145 +++++++++ .../benchmarks/queries/src/telemetry.rs | 110 +++++++ 12 files changed, 1642 insertions(+) create mode 100644 src/bigquery/benchmarks/queries/Cargo.toml create mode 100644 src/bigquery/benchmarks/queries/README.md create mode 100644 src/bigquery/benchmarks/queries/src/args.rs create mode 100644 src/bigquery/benchmarks/queries/src/main.rs create mode 100644 src/bigquery/benchmarks/queries/src/metrics.rs create mode 100644 src/bigquery/benchmarks/queries/src/reporter.rs create mode 100644 src/bigquery/benchmarks/queries/src/runner.rs create mode 100644 src/bigquery/benchmarks/queries/src/sample.rs create mode 100644 src/bigquery/benchmarks/queries/src/scenarios.rs create mode 100644 src/bigquery/benchmarks/queries/src/telemetry.rs diff --git a/Cargo.lock b/Cargo.lock index 68b75b6de2..c2e456edec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -394,6 +394,32 @@ 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", + "futures", + "google-cloud-auth", + "google-cloud-bigquery", + "google-cloud-bigquery-v2", + "google-cloud-gax", + "humantime", + "integration-tests-o11y", + "opentelemetry", + "opentelemetry_sdk", + "rand 0.10.2", + "serde", + "serde_json", + "tokio", + "tokio-metrics", + "tracing", + "tracing-log", + "tracing-subscriber", + "uuid", +] + [[package]] name = "bigquery-samples" version = "0.0.0" diff --git a/Cargo.toml b/Cargo.toml index 47e48f545b..b3340e44f8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -68,6 +68,7 @@ members = [ "guide/samples", "src/auth", "src/bigquery", + "src/bigquery/benchmarks/queries", "src/bigquery-derive", "src/bigquery-write", "src/bigquery-write/grpc-mock", diff --git a/src/bigquery/benchmarks/queries/Cargo.toml b/src/bigquery/benchmarks/queries/Cargo.toml new file mode 100644 index 0000000000..1c50fe13a3 --- /dev/null +++ b/src/bigquery/benchmarks/queries/Cargo.toml @@ -0,0 +1,49 @@ +# 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"] } +futures.workspace = true +google-cloud-auth.workspace = true +google-cloud-bigquery = { workspace = true, features = ["default-rustls-provider"] } +google-cloud-bigquery-v2.workspace = true +google-cloud-gax.workspace = true +humantime.workspace = true +integration-tests-o11y.workspace = true +opentelemetry = { workspace = true, features = ["trace", "metrics"] } +opentelemetry_sdk = { workspace = true, features = ["rt-tokio", "trace", "metrics"] } +rand.workspace = true +serde = { workspace = true, features = ["derive"] } +serde_json.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time", "sync"] } +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..533f3b8efd --- /dev/null +++ b/src/bigquery/benchmarks/queries/README.md @@ -0,0 +1,126 @@ +# 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 + ``` + +2. **Project ID**: + Set the project ID environment variable: + ```shell + export GOOGLE_CLOUD_PROJECT="$(gcloud config get project)" + ``` + +--- + +## Running Benchmarks + +### 1. Zero-Setup Synthetic Benchmark (Default) + +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 +``` + +### 2. 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 +``` + +### 3. 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 +``` + +--- + +## 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` + +--- + +## 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..866689aa85 --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/args.rs @@ -0,0 +1,184 @@ +// 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. + /// + /// Defaults to 10 if `--duration` is not set. + #[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 (maps to `max_results`). + #[arg(long)] + pub max_results: 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, + + /// 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, + + /// 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, + + /// Whether to log debug details for retry decisions. + #[arg(long)] + pub debug_retry: bool, +} + +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." + ); + } + if let Some(iterations) = self.iterations + && iterations == 0 + && self.duration.is_none() + { + anyhow::bail!("--iterations must be greater than 0"); + } + Ok(()) + } + + /// Returns the effective iterations limit per worker (defaults to 10 if duration is not set). + pub fn effective_iterations(&self) -> Option { + match (self.iterations, self.duration) { + (Some(i), _) => Some(i), + (None, Some(_)) => None, + (None, None) => Some(10), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_default_args_validation() { + let args = Args::parse_from(["bigquery-benchmark-queries"]); + assert!(args.validate().is_ok()); + assert_eq!(args.task_count, 1); + assert_eq!(args.effective_iterations(), Some(10)); + } + + #[test] + fn test_custom_scenario_validation() { + let args = Args::parse_from(["bigquery-benchmark-queries", "--scenario", "custom"]); + assert!(args.validate().is_err()); + + let args = Args::parse_from([ + "bigquery-benchmark-queries", + "--scenario", + "custom", + "--sql", + "SELECT 1", + ]); + assert!(args.validate().is_ok()); + } + + #[test] + fn test_duration_mode() { + let args = Args::parse_from(["bigquery-benchmark-queries", "--duration", "5m"]); + assert!(args.validate().is_ok()); + assert_eq!(args.effective_iterations(), None); + assert_eq!(args.duration, Some(Duration::from_secs(300))); + } +} diff --git a/src/bigquery/benchmarks/queries/src/main.rs b/src/bigquery/benchmarks/queries/src/main.rs new file mode 100644 index 0000000000..aae9ac5886 --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/main.rs @@ -0,0 +1,142 @@ +// 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::collections::BTreeMap; +use std::time::Instant; +use tokio::task::JoinSet; + +#[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, + effective_iterations = ?args.effective_iterations(), + duration = ?args.duration, + "Starting BigQuery benchmark" + ); + + // Spawn periodic runtime monitor and counter logger + let handle = tokio::runtime::Handle::current(); + let runtime_monitor = tokio_metrics::RuntimeMonitor::new(&handle); + let monitor_freq = std::time::Duration::from_secs(5); + tokio::spawn(async move { + for metrics in runtime_monitor.intervals() { + let counters = BTreeMap::from_iter(metrics::get_counters()); + tracing::info!("Counters = {:?} RuntimeMetrics = {:?}", counters, metrics); + tokio::time::sleep(monitor_freq).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(); + + let channel_capacity = (1024 * args.task_count).max(64); + let (tx, rx) = tokio::sync::mpsc::channel(channel_capacity); + + // 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); + + 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:?}"), + } + + let final_counters = BTreeMap::from_iter(metrics::get_counters()); + tracing::info!("Final counters: {:?}", final_counters); + + 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..3fcaae85af --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/metrics.rs @@ -0,0 +1,267 @@ +// 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 opentelemetry::KeyValue; +use opentelemetry::metrics::{Counter, Histogram}; +use serde::{Deserialize, Serialize}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +/// Summary percentiles and metrics for execution latencies. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct LatencySummary { + pub min: Duration, + pub max: Duration, + pub mean: Duration, + pub p50: Duration, + pub p90: Duration, + 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, + }) +} + +// In-process global atomic counters for quick telemetry reporting. +static TOTAL_QUERIES: AtomicU64 = AtomicU64::new(0); +static SUCCESS_QUERIES: AtomicU64 = AtomicU64::new(0); +static ERROR_QUERIES: AtomicU64 = AtomicU64::new(0); +static RETRIED_QUERIES: AtomicU64 = AtomicU64::new(0); +static TOTAL_ROWS: AtomicU64 = AtomicU64::new(0); +static TOTAL_BYTES: AtomicU64 = AtomicU64::new(0); + +#[inline] +pub fn inc_total_queries() { + TOTAL_QUERIES.fetch_add(1, Ordering::SeqCst); +} + +#[inline] +pub fn inc_success_queries() { + SUCCESS_QUERIES.fetch_add(1, Ordering::SeqCst); +} + +#[inline] +pub fn inc_error_queries() { + ERROR_QUERIES.fetch_add(1, Ordering::SeqCst); +} + +#[inline] +pub fn inc_retried_queries() { + RETRIED_QUERIES.fetch_add(1, Ordering::SeqCst); +} + +#[inline] +pub fn add_rows_read(count: u64) { + TOTAL_ROWS.fetch_add(count, Ordering::SeqCst); +} + +#[inline] +pub fn add_bytes_processed(bytes: u64) { + TOTAL_BYTES.fetch_add(bytes, Ordering::SeqCst); +} + +/// Returns a snapshot of in-process counters. +pub fn get_counters() -> [(&'static str, u64); 6] { + [ + ("total_queries", TOTAL_QUERIES.load(Ordering::Relaxed)), + ("success_queries", SUCCESS_QUERIES.load(Ordering::Relaxed)), + ("error_queries", ERROR_QUERIES.load(Ordering::Relaxed)), + ("retried_queries", RETRIED_QUERIES.load(Ordering::Relaxed)), + ("total_rows_read", TOTAL_ROWS.load(Ordering::Relaxed)), + ("total_bytes_processed", TOTAL_BYTES.load(Ordering::Relaxed)), + ] +} + +/// OpenTelemetry metrics instruments. +#[derive(Clone)] +pub struct OtelMetrics { + pub queries_total: Counter, + pub queries_success: Counter, + pub queries_error: Counter, + pub queries_retried: Counter, + pub rows_read: Counter, + pub bytes_processed: Counter, + pub query_duration: Histogram, + pub send_duration: Histogram, + pub poll_duration: Histogram, + pub read_duration: Histogram, +} + +impl OtelMetrics { + pub fn new() -> Self { + let meter = opentelemetry::global::meter("bigquery-benchmark-queries"); + + let queries_total = meter + .u64_counter("bigquery.queries.total") + .with_description("Total number of BigQuery queries attempted") + .build(); + + let queries_success = meter + .u64_counter("bigquery.queries.success") + .with_description("Number of BigQuery queries completed successfully") + .build(); + + let queries_error = meter + .u64_counter("bigquery.queries.error") + .with_description("Number of BigQuery queries that failed") + .build(); + + let queries_retried = meter + .u64_counter("bigquery.queries.retries_detected") + .with_description("Number of queries where an under-the-hood job retry was detected") + .build(); + + let rows_read = meter + .u64_counter("bigquery.queries.rows_read") + .with_description("Total count of rows read from query results") + .build(); + + let bytes_processed = meter + .u64_counter("bigquery.queries.bytes_processed") + .with_description("Total estimated bytes processed by BigQuery jobs") + .build(); + + let query_duration = meter + .f64_histogram("bigquery.queries.duration_seconds") + .with_description("Total query end-to-end duration in seconds") + .build(); + + let send_duration = meter + .f64_histogram("bigquery.queries.send_duration_seconds") + .with_description("Duration for Query::send() execution") + .build(); + + let poll_duration = meter + .f64_histogram("bigquery.queries.poll_duration_seconds") + .with_description("Duration for Query::until_done() polling execution") + .build(); + + let read_duration = meter + .f64_histogram("bigquery.queries.read_duration_seconds") + .with_description("Duration for CompleteQuery::read() row streaming") + .build(); + + Self { + queries_total, + queries_success, + queries_error, + queries_retried, + rows_read, + bytes_processed, + query_duration, + send_duration, + poll_duration, + read_duration, + } + } + + pub fn record_sample(&self, scenario: &str, sample: &crate::sample::Sample) { + let is_ok = sample.status == crate::sample::SampleStatus::Ok; + let attrs = [ + KeyValue::new("scenario", scenario.to_string()), + KeyValue::new("status", if is_ok { "ok" } else { "error" }), + ]; + + 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); + if sample.bytes_processed > 0 { + self.bytes_processed + .add(sample.bytes_processed as u64, &attrs); + } + + self.send_duration.record( + Duration::from_micros(sample.send_duration_micros as u64).as_secs_f64(), + &attrs, + ); + self.poll_duration.record( + Duration::from_micros(sample.poll_duration_micros as u64).as_secs_f64(), + &attrs, + ); + self.read_duration.record( + Duration::from_micros(sample.read_duration_micros as u64).as_secs_f64(), + &attrs, + ); + } else { + self.queries_error.add(1, &attrs); + } + + self.query_duration.record( + Duration::from_micros(sample.total_duration_micros as u64).as_secs_f64(), + &attrs, + ); + } +} + +impl Default for OtelMetrics { + fn default() -> Self { + Self::new() + } +} + +#[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..624a477c12 --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/reporter.rs @@ -0,0 +1,174 @@ +// 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}; +use crate::scenarios::Scenario; +use serde::{Deserialize, Serialize}; +use std::fs::File; +use std::io::Write; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use tokio::sync::mpsc::Receiver; + +/// 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, +} + +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); + } + + if let Some(send) = &self.send_duration { + println!("\n--- Query::send() Latency ---"); + println!( + " P50: {:?} | P90: {:?} | P99: {:?}", + send.p50, send.p90, send.p99 + ); + } + + if let Some(poll) = &self.poll_duration { + println!("\n--- Query::until_done() Polling Latency ---"); + println!( + " P50: {:?} | P90: {:?} | P99: {:?}", + poll.p50, poll.p90, poll.p99 + ); + } + + if let Some(read) = &self.read_duration { + println!("\n--- CompleteQuery::read() Streaming Latency ---"); + println!( + " P50: {:?} | P90: {:?} | P99: {:?}", + read.p50, read.p90, read.p99 + ); + } + + println!("=======================================================\n"); + } +} + +/// 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 samples = Vec::new(); + let mut success_total_durations = Vec::new(); + let mut success_send_durations = Vec::new(); + let mut success_poll_durations = Vec::new(); + let mut success_read_durations = Vec::new(); + + let mut success_count = 0_usize; + let mut error_count = 0_usize; + let mut retries_detected_count = 0_usize; + let mut total_rows_read = 0_usize; + let mut total_bytes_processed = 0_i64; + + while let Some(sample) = rx.recv().await { + if sample.status == SampleStatus::Ok { + success_count += 1; + success_total_durations.push(sample.total_duration()); + success_send_durations.push(Duration::from_micros(sample.send_duration_micros as u64)); + success_poll_durations.push(Duration::from_micros(sample.poll_duration_micros as u64)); + success_read_durations.push(Duration::from_micros(sample.read_duration_micros as u64)); + total_rows_read += sample.rows_count; + total_bytes_processed += sample.bytes_processed; + } else { + error_count += 1; + } + + if sample.retry_detected { + retries_detected_count += 1; + } + + samples.push(sample); + } + + let report = BenchmarkReport { + scenario: scenario.name.clone(), + task_count: args.task_count, + total_samples: samples.len(), + success_count, + error_count, + retries_detected_count, + total_rows_read, + total_bytes_processed, + total_duration: metrics::compute_metrics(&success_total_durations), + send_duration: metrics::compute_metrics(&success_send_durations), + poll_duration: metrics::compute_metrics(&success_poll_durations), + read_duration: metrics::compute_metrics(&success_read_durations), + }; + + report.print_stdout(); + + if let Some(output_dir) = &args.output_dir { + std::fs::create_dir_all(output_dir)?; + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0); + + // Write raw samples CSV + let csv_path = output_dir.join(format!("samples-{}-{}.csv", scenario.name, timestamp)); + let mut csv_file = File::create(&csv_path)?; + writeln!(csv_file, "{}", Sample::HEADER)?; + for sample in &samples { + writeln!(csv_file, "{}", sample.to_csv_row())?; + } + println!("Raw samples written to: {}", csv_path.display()); + + // Write summary JSON + let json_path = output_dir.join(format!("summary-{}-{}.json", scenario.name, timestamp)); + let json_file = File::create(&json_path)?; + serde_json::to_writer_pretty(json_file, &report)?; + println!("Summary report written to: {}", json_path.display()); + } + + Ok(report) +} diff --git a/src/bigquery/benchmarks/queries/src/runner.rs b/src/bigquery/benchmarks/queries/src/runner.rs new file mode 100644 index 0000000000..90732a8b16 --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/runner.rs @@ -0,0 +1,293 @@ +// 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, OtelMetrics}; +use crate::sample::{Sample, SampleStatus}; +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 effective_iterations = self.args.effective_iterations(); + let mut iteration = 0_u64; + + loop { + if let Some(max_iter) = effective_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 = self.test_start.elapsed().as_micros(); + + let iteration_span = tracing::info_span!( + "bigquery.query_benchmark.iteration", + task_id = self.task_id, + iteration, + scenario = %self.scenario.name + ); + + let sample = self + .execute_iteration(iteration, start_offset_micros, iter_start) + .instrument(iteration_span) + .await; + + let _ = self.tx.send(sample).await; + iteration += 1; + } + + Ok(()) + } + + async fn execute_iteration( + &self, + iteration: u64, + start_offset_micros: u128, + iter_start: Instant, + ) -> Sample { + let mut query_builder = self + .client + .query(&self.scenario.sql) + .set_location(&self.args.location); + + if let Some(max_results) = self.args.max_results { + query_builder = query_builder.set_max_results(max_results); + } + 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; + let send_duration = send_start.elapsed(); + + let query_handle = match send_result { + Ok(handle) => handle, + Err(err) => { + let total_duration = iter_start.elapsed(); + metrics::inc_total_queries(); + metrics::inc_error_queries(); + tracing::error!(self.task_id, iteration, "Query::send failed: {err:?}"); + + let sample = Sample { + task_id: self.task_id, + iteration, + start_offset_micros, + send_duration_micros: send_duration.as_micros(), + poll_duration_micros: 0, + read_duration_micros: 0, + total_duration_micros: total_duration.as_micros(), + 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: err.to_string(), + }; + self.metrics.record_sample(&self.scenario.name, &sample); + return sample; + } + }; + + // Capture initial job_id if present + let 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 + ); + let done_result = query_handle.until_done().instrument(poll_span).await; + let poll_duration = poll_start.elapsed(); + + let complete_query = match done_result { + Ok(complete) => complete, + Err(err) => { + let total_duration = iter_start.elapsed(); + metrics::inc_total_queries(); + metrics::inc_error_queries(); + tracing::error!( + self.task_id, + iteration, + %initial_job_id, + "Query::until_done failed: {err:?}" + ); + + let sample = Sample { + task_id: self.task_id, + iteration, + start_offset_micros, + send_duration_micros: send_duration.as_micros(), + poll_duration_micros: poll_duration.as_micros(), + read_duration_micros: 0, + total_duration_micros: total_duration.as_micros(), + rows_count: 0, + bytes_processed: 0, + cache_hit: false, + initial_job_id, + final_job_id: String::new(), + retry_detected: false, + status: SampleStatus::Error, + error_message: err.to_string(), + }; + self.metrics.record_sample(&self.scenario.name, &sample); + return sample; + } + }; + + // Capture final job_id + let final_job_id = complete_query + .metadata() + .job_reference + .as_ref() + .map(|r| r.job_id.clone()) + .unwrap_or_default(); + + // Step 3: Detect if under-the-hood job retry occurred + let retry_detected = !initial_job_id.is_empty() + && !final_job_id.is_empty() + && initial_job_id != final_job_id; + + if retry_detected { + metrics::inc_retried_queries(); + tracing::warn!( + task_id = self.task_id, + iteration, + %initial_job_id, + %final_job_id, + "Query job retry detected under the hood (job_id mutated)!" + ); + } + + let bytes_processed = complete_query.metadata().total_bytes_processed.unwrap_or(0); + + let cache_hit = complete_query.metadata().cache_hit.unwrap_or(false); + + // Step 4: Stream and read result rows if enabled + let mut rows_count = 0_usize; + 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); + let _guard = read_span.enter(); + let mut rows = complete_query.read(); + while let Some(row_result) = rows.next().await { + match row_result { + Ok(_) => { + rows_count += 1; + } + Err(err) => { + tracing::error!(self.task_id, iteration, "Error streaming rows: {err:?}"); + read_error = Some(err.to_string()); + break; + } + } + } + } + let read_duration = read_start.elapsed(); + let total_duration = iter_start.elapsed(); + + metrics::inc_total_queries(); + + let sample = if let Some(err_msg) = read_error { + metrics::inc_error_queries(); + + Sample { + task_id: self.task_id, + iteration, + start_offset_micros, + send_duration_micros: send_duration.as_micros(), + poll_duration_micros: poll_duration.as_micros(), + read_duration_micros: read_duration.as_micros(), + total_duration_micros: total_duration.as_micros(), + rows_count, + bytes_processed, + cache_hit, + initial_job_id, + final_job_id, + retry_detected, + status: SampleStatus::Error, + error_message: err_msg, + } + } else { + metrics::inc_success_queries(); + metrics::add_rows_read(rows_count as u64); + if bytes_processed > 0 { + metrics::add_bytes_processed(bytes_processed as u64); + } + + Sample { + task_id: self.task_id, + iteration, + start_offset_micros, + send_duration_micros: send_duration.as_micros(), + poll_duration_micros: poll_duration.as_micros(), + read_duration_micros: read_duration.as_micros(), + total_duration_micros: total_duration.as_micros(), + rows_count, + bytes_processed, + cache_hit, + initial_job_id, + final_job_id, + retry_detected, + status: SampleStatus::Ok, + error_message: String::new(), + } + }; + + self.metrics.record_sample(&self.scenario.name, &sample); + 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..031f52115b --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/sample.rs @@ -0,0 +1,125 @@ +// 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: u128, + pub send_duration_micros: u128, + pub poll_duration_micros: u128, + pub read_duration_micros: u128, + pub total_duration_micros: u128, + 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" + ); + + pub fn to_csv_row(&self) -> String { + let clean_err = self.error_message.replace(',', ";").replace('\n', " "); + 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, + if self.initial_job_id.is_empty() { + "N/A" + } else { + &self.initial_job_id + }, + if self.final_job_id.is_empty() { + "N/A" + } else { + &self.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 as u64) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_sample_csv_serialization() { + let sample = Sample { + task_id: 1, + iteration: 42, + start_offset_micros: 100_000, + 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, + cache_hit: false, + initial_job_id: "job_init_123".to_string(), + final_job_id: "job_retry_456".to_string(), + retry_detected: true, + status: SampleStatus::Ok, + error_message: String::new(), + }; + + 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,")); + } +} diff --git a/src/bigquery/benchmarks/queries/src/scenarios.rs b/src/bigquery/benchmarks/queries/src/scenarios.rs new file mode 100644 index 0000000000..7d84905d1d --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/scenarios.rs @@ -0,0 +1,145 @@ +// 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: String, + 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 { + match args.scenario { + ScenarioName::Synthetic100k => Ok(Self { + name: "synthetic-100k".to_string(), + sql: concat!( + "SELECT ", + " x AS row_id, ", + " GENERATE_UUID() AS uuid, ", + " REPEAT('abcdefghij', 10) AS payload ", + "FROM UNNEST(GENERATE_ARRAY(1, 100000)) AS x" + ) + .to_string(), + description: "Generates 100,000 structured rows in-flight with no external table dependency.", + }), + ScenarioName::Synthetic10k => Ok(Self { + name: "synthetic-10k".to_string(), + sql: concat!( + "SELECT ", + " x AS row_id, ", + " GENERATE_UUID() AS uuid, ", + " REPEAT('abcdefghij', 10) AS payload ", + "FROM UNNEST(GENERATE_ARRAY(1, 10000)) AS x" + ) + .to_string(), + description: "Generates 10,000 structured rows in-flight with no external table dependency.", + }), + ScenarioName::UsaNamesScan => Ok(Self { + name: "usa-names-scan".to_string(), + sql: concat!( + "SELECT name, state, year, gender, number ", + "FROM `bigquery-public-data.usa_names.usa_1910_2013` ", + "WHERE year >= 2000 ", + "LIMIT 50000" + ) + .to_string(), + description: "Scans and retrieves 50,000 rows from the USA names public dataset.", + }), + ScenarioName::UsaNamesAgg => Ok(Self { + name: "usa-names-agg".to_string(), + sql: 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(), + description: "Aggregates 5.5M rows grouped by state and gender.", + }), + ScenarioName::WikipediaAgg => Ok(Self { + name: "wikipedia-agg".to_string(), + sql: 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(), + description: "Aggregates top 1000 article views from Wikipedia public samples.", + }), + ScenarioName::Custom => { + let sql = if let Some(sql) = &args.sql { + sql.clone() + } else if let Some(sql_file) = &args.sql_file { + fs::read_to_string(sql_file).map_err(|e| { + anyhow::anyhow!( + "Failed to read custom SQL file {}: {}", + sql_file.display(), + e + ) + })? + } else { + anyhow::bail!("Custom scenario requires --sql or --sql-file"); + }; + + Ok(Self { + name: "custom".to_string(), + sql, + description: "User-defined custom SQL query.", + }) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use clap::Parser; + + #[test] + fn test_synthetic_scenarios() { + let args = Args::parse_from(["bigquery-benchmark-queries", "--scenario", "synthetic-100k"]); + let s = Scenario::resolve(&args).unwrap(); + assert_eq!(s.name, "synthetic-100k"); + assert!(s.sql.contains("100000")); + + let args = Args::parse_from(["bigquery-benchmark-queries", "--scenario", "synthetic-10k"]); + let s = Scenario::resolve(&args).unwrap(); + assert_eq!(s.name, "synthetic-10k"); + assert!(s.sql.contains("10000")); + } + + #[test] + fn test_custom_scenario_from_string() { + let args = Args::parse_from([ + "bigquery-benchmark-queries", + "--scenario", + "custom", + "--sql", + "SELECT 42", + ]); + let s = Scenario::resolve(&args).unwrap(); + assert_eq!(s.name, "custom"); + assert_eq!(s.sql, "SELECT 42"); + } +} diff --git a/src/bigquery/benchmarks/queries/src/telemetry.rs b/src/bigquery/benchmarks/queries/src/telemetry.rs new file mode 100644 index 0000000000..b377c91407 --- /dev/null +++ b/src/bigquery/benchmarks/queries/src/telemetry.rs @@ -0,0 +1,110 @@ +// 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::otlp::Uri; +use opentelemetry_sdk::metrics::SdkMeterProvider; +use opentelemetry_sdk::trace::SdkTracerProvider; +use std::str::FromStr; +use tracing_subscriber::fmt::format::FmtSpan; +use tracing_subscriber::prelude::*; + +const SERVICE_NAME: &str = "bigquery-benchmark-queries"; + +/// Holds providers that need graceful flush and shutdown upon completion. +pub struct TelemetryGuard { + tracer_provider: Option, + meter_provider: Option, +} + +impl TelemetryGuard { + /// Flushes and shuts down telemetry providers. + pub fn shutdown(self) { + if let Some(tp) = self.tracer_provider + && let Err(e) = tp.shutdown() + { + eprintln!("Error shutting down trace provider: {e:?}"); + } + if let Some(mp) = self.meter_provider + && let Err(e) = mp.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::NEW | FmtSpan::CLOSE) + .with_writer(std::io::stderr) + .with_filter(env_filter); + + let registry = tracing_subscriber::Registry::default().with(fmt_layer); + + if let Some(project_id) = &args.project_id { + tracing::info!("Enabling OpenTelemetry Cloud Trace & Monitoring for project {project_id}"); + + let mut trace_builder = + integration_tests_o11y::otlp::trace::Builder::new(project_id, SERVICE_NAME) + .with_credentials(credentials.clone()); + + let mut meter_builder = + integration_tests_o11y::otlp::metrics::Builder::new(project_id, SERVICE_NAME) + .with_credentials(credentials.clone()); + + 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 { + tracer_provider: Some(tracer_provider), + meter_provider: Some(meter_provider), + }); + } + + tracing::subscriber::set_global_default(registry).expect("Setting global subscriber succeeds"); + + Ok(TelemetryGuard { + tracer_provider: None, + meter_provider: None, + }) +} From 12d7dfa5493673b0c1e62967939565a26adb1239 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Thu, 20 Aug 2026 15:23:04 +0000 Subject: [PATCH 02/11] impl: query cache and async reports --- src/bigquery/benchmarks/queries/Cargo.toml | 2 +- src/bigquery/benchmarks/queries/README.md | 70 ++++++++++ src/bigquery/benchmarks/queries/src/args.rs | 5 + src/bigquery/benchmarks/queries/src/main.rs | 33 +++-- .../benchmarks/queries/src/metrics.rs | 43 ++++++- .../benchmarks/queries/src/reporter.rs | 120 +++++++++++++----- src/bigquery/benchmarks/queries/src/runner.rs | 3 +- .../benchmarks/queries/src/telemetry.rs | 65 ++++++++-- 8 files changed, 284 insertions(+), 57 deletions(-) diff --git a/src/bigquery/benchmarks/queries/Cargo.toml b/src/bigquery/benchmarks/queries/Cargo.toml index 1c50fe13a3..3481187c5c 100644 --- a/src/bigquery/benchmarks/queries/Cargo.toml +++ b/src/bigquery/benchmarks/queries/Cargo.toml @@ -38,7 +38,7 @@ opentelemetry_sdk = { workspace = true, features = ["rt-tokio", " rand.workspace = true serde = { workspace = true, features = ["derive"] } serde_json.workspace = true -tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time", "sync"] } +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time", "sync", "signal"] } tokio-metrics = { workspace = true, features = ["rt"] } tracing.workspace = true tracing-log = { workspace = true, features = ["log-tracer", "std"] } diff --git a/src/bigquery/benchmarks/queries/README.md b/src/bigquery/benchmarks/queries/README.md index 533f3b8efd..f64ed7dad2 100644 --- a/src/bigquery/benchmarks/queries/README.md +++ b/src/bigquery/benchmarks/queries/README.md @@ -35,6 +35,9 @@ Benchmarks the Rust BigQuery client library (`google-cloud-bigquery`), measuring ## 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. + ### 1. Zero-Setup Synthetic Benchmark (Default) Runs 10 iterations per task with 4 concurrent tasks, generating and streaming 100,000 rows per query: @@ -87,6 +90,9 @@ cargo run --release -p bigquery-benchmark-queries -- \ --output-dir ./results ``` +> [!TIP] +> **Real-Time Reporting & Graceful Shutdown (`Ctrl+C`):** When `--output-dir` is provided, the benchmark writes raw sample CSV rows and updates the summary report JSON on disk in **real-time** as each query iteration completes. You can press `Ctrl+C` at any point during a long endurance test to immediately stop worker tasks, output the summary report to stdout, flush all OpenTelemetry metrics to Google Cloud, and preserve the recorded samples and summary report on disk. + --- ## OpenTelemetry & Cloud Observability @@ -112,6 +118,70 @@ When `--project-id` is specified, the benchmark automatically connects to `telem - `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()`): +```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 diff --git a/src/bigquery/benchmarks/queries/src/args.rs b/src/bigquery/benchmarks/queries/src/args.rs index 866689aa85..8c83b8e420 100644 --- a/src/bigquery/benchmarks/queries/src/args.rs +++ b/src/bigquery/benchmarks/queries/src/args.rs @@ -98,6 +98,10 @@ pub struct Args { #[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, @@ -157,6 +161,7 @@ mod tests { assert!(args.validate().is_ok()); assert_eq!(args.task_count, 1); assert_eq!(args.effective_iterations(), Some(10)); + assert!(!args.use_query_cache); } #[test] diff --git a/src/bigquery/benchmarks/queries/src/main.rs b/src/bigquery/benchmarks/queries/src/main.rs index aae9ac5886..a90ee272a1 100644 --- a/src/bigquery/benchmarks/queries/src/main.rs +++ b/src/bigquery/benchmarks/queries/src/main.rs @@ -48,6 +48,7 @@ async fn main() -> anyhow::Result<()> { task_count = args.task_count, effective_iterations = ?args.effective_iterations(), duration = ?args.duration, + use_query_cache = args.use_query_cache, "Starting BigQuery benchmark" ); @@ -73,6 +74,7 @@ async fn main() -> anyhow::Result<()> { let client = client_builder.build().await?; let otel_metrics = OtelMetrics::new(); + otel_metrics.init_scenario(&scenario.name); let channel_capacity = (1024 * args.task_count).max(64); let (tx, rx) = tokio::sync::mpsc::channel(channel_capacity); @@ -112,18 +114,27 @@ async fn main() -> anyhow::Result<()> { // Drop main sender so receiver terminates after all tasks complete drop(tx); - 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:?}"); - } + 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 diff --git a/src/bigquery/benchmarks/queries/src/metrics.rs b/src/bigquery/benchmarks/queries/src/metrics.rs index 3fcaae85af..868783f567 100644 --- a/src/bigquery/benchmarks/queries/src/metrics.rs +++ b/src/bigquery/benchmarks/queries/src/metrics.rs @@ -193,6 +193,27 @@ impl OtelMetrics { } } + /// Initializes all counter metrics with 0 so time series exist in Cloud Monitoring + /// even if no errors or retries occur during the benchmark run. + pub fn init_scenario(&self, scenario: &str) { + let ok_attrs = [ + KeyValue::new("scenario", scenario.to_string()), + KeyValue::new("status", "ok"), + ]; + let err_attrs = [ + KeyValue::new("scenario", scenario.to_string()), + KeyValue::new("status", "error"), + ]; + + self.queries_total.add(0, &ok_attrs); + self.queries_total.add(0, &err_attrs); + self.queries_success.add(0, &ok_attrs); + self.queries_error.add(0, &err_attrs); + self.queries_retried.add(0, &ok_attrs); + self.rows_read.add(0, &ok_attrs); + self.bytes_processed.add(0, &ok_attrs); + } + pub fn record_sample(&self, scenario: &str, sample: &crate::sample::Sample) { let is_ok = sample.status == crate::sample::SampleStatus::Ok; let attrs = [ @@ -203,14 +224,21 @@ impl OtelMetrics { self.queries_total.add(1, &attrs); if is_ok { self.queries_success.add(1, &attrs); + self.queries_error.add( + 0, + &[ + KeyValue::new("scenario", scenario.to_string()), + KeyValue::new("status", "error"), + ], + ); if sample.retry_detected { self.queries_retried.add(1, &attrs); + } else { + self.queries_retried.add(0, &attrs); } self.rows_read.add(sample.rows_count as u64, &attrs); - if sample.bytes_processed > 0 { - self.bytes_processed - .add(sample.bytes_processed as u64, &attrs); - } + self.bytes_processed + .add(sample.bytes_processed.max(0) as u64, &attrs); self.send_duration.record( Duration::from_micros(sample.send_duration_micros as u64).as_secs_f64(), @@ -226,6 +254,13 @@ impl OtelMetrics { ); } else { self.queries_error.add(1, &attrs); + self.queries_success.add( + 0, + &[ + KeyValue::new("scenario", scenario.to_string()), + KeyValue::new("status", "ok"), + ], + ); } self.query_duration.record( diff --git a/src/bigquery/benchmarks/queries/src/reporter.rs b/src/bigquery/benchmarks/queries/src/reporter.rs index 624a477c12..df689bf766 100644 --- a/src/bigquery/benchmarks/queries/src/reporter.rs +++ b/src/bigquery/benchmarks/queries/src/reporter.rs @@ -98,7 +98,7 @@ pub async fn collect_and_report( scenario: &Scenario, args: &Args, ) -> anyhow::Result { - let mut samples = Vec::new(); + let mut total_samples = 0_usize; let mut success_total_durations = Vec::new(); let mut success_send_durations = Vec::new(); let mut success_poll_durations = Vec::new(); @@ -110,7 +110,30 @@ pub async fn collect_and_report( let mut total_rows_read = 0_usize; let mut total_bytes_processed = 0_i64; + let mut realtime_files = if let Some(output_dir) = &args.output_dir { + 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-{}-{}.csv", scenario.name, timestamp)); + let json_path = output_dir.join(format!("summary-{}-{}.json", scenario.name, timestamp)); + + let mut csv_file = File::create(&csv_path)?; + writeln!(csv_file, "{}", Sample::HEADER)?; + csv_file.flush()?; + + println!("Writing real-time samples to: {}", csv_path.display()); + println!("Writing real-time summary to: {}", json_path.display()); + + Some((csv_file, csv_path, json_path)) + } else { + None + }; + while let Some(sample) = rx.recv().await { + total_samples += 1; if sample.status == SampleStatus::Ok { success_count += 1; success_total_durations.push(sample.total_duration()); @@ -127,48 +150,85 @@ pub async fn collect_and_report( retries_detected_count += 1; } - samples.push(sample); + if let Some((csv_file, _, json_path)) = &mut realtime_files { + let _ = writeln!(csv_file, "{}", sample.to_csv_row()); + let _ = csv_file.flush(); + + let stats = ReportStats { + scenario_name: &scenario.name, + task_count: args.task_count, + total_samples, + success_count, + error_count, + retries_detected_count, + total_rows_read, + total_bytes_processed, + success_total_durations: &success_total_durations, + success_send_durations: &success_send_durations, + success_poll_durations: &success_poll_durations, + success_read_durations: &success_read_durations, + }; + let current_report = build_report(&stats); + if let Ok(json_file) = File::create(json_path) { + let _ = serde_json::to_writer_pretty(json_file, ¤t_report); + } + } } - let report = BenchmarkReport { - scenario: scenario.name.clone(), + let stats = ReportStats { + scenario_name: &scenario.name, task_count: args.task_count, - total_samples: samples.len(), + total_samples, success_count, error_count, retries_detected_count, total_rows_read, total_bytes_processed, - total_duration: metrics::compute_metrics(&success_total_durations), - send_duration: metrics::compute_metrics(&success_send_durations), - poll_duration: metrics::compute_metrics(&success_poll_durations), - read_duration: metrics::compute_metrics(&success_read_durations), + success_total_durations: &success_total_durations, + success_send_durations: &success_send_durations, + success_poll_durations: &success_poll_durations, + success_read_durations: &success_read_durations, }; + let report = build_report(&stats); report.print_stdout(); - if let Some(output_dir) = &args.output_dir { - std::fs::create_dir_all(output_dir)?; - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0); - - // Write raw samples CSV - let csv_path = output_dir.join(format!("samples-{}-{}.csv", scenario.name, timestamp)); - let mut csv_file = File::create(&csv_path)?; - writeln!(csv_file, "{}", Sample::HEADER)?; - for sample in &samples { - writeln!(csv_file, "{}", sample.to_csv_row())?; - } - println!("Raw samples written to: {}", csv_path.display()); - - // Write summary JSON - let json_path = output_dir.join(format!("summary-{}-{}.json", scenario.name, timestamp)); - let json_file = File::create(&json_path)?; - serde_json::to_writer_pretty(json_file, &report)?; - println!("Summary report written to: {}", json_path.display()); + if let Some((_, csv_path, json_path)) = realtime_files { + println!("Final samples saved to: {}", csv_path.display()); + println!("Final summary saved to: {}", json_path.display()); } Ok(report) } + +struct ReportStats<'a> { + scenario_name: &'a str, + task_count: usize, + total_samples: usize, + success_count: usize, + error_count: usize, + retries_detected_count: usize, + total_rows_read: usize, + total_bytes_processed: i64, + success_total_durations: &'a [Duration], + success_send_durations: &'a [Duration], + success_poll_durations: &'a [Duration], + success_read_durations: &'a [Duration], +} + +fn build_report(stats: &ReportStats) -> BenchmarkReport { + BenchmarkReport { + scenario: stats.scenario_name.to_string(), + task_count: stats.task_count, + total_samples: stats.total_samples, + success_count: stats.success_count, + error_count: stats.error_count, + retries_detected_count: stats.retries_detected_count, + total_rows_read: stats.total_rows_read, + total_bytes_processed: stats.total_bytes_processed, + total_duration: metrics::compute_metrics(stats.success_total_durations), + send_duration: metrics::compute_metrics(stats.success_send_durations), + poll_duration: metrics::compute_metrics(stats.success_poll_durations), + read_duration: metrics::compute_metrics(stats.success_read_durations), + } +} diff --git a/src/bigquery/benchmarks/queries/src/runner.rs b/src/bigquery/benchmarks/queries/src/runner.rs index 90732a8b16..2dbc9bfa51 100644 --- a/src/bigquery/benchmarks/queries/src/runner.rs +++ b/src/bigquery/benchmarks/queries/src/runner.rs @@ -86,7 +86,8 @@ impl TaskRunner<'_> { let mut query_builder = self .client .query(&self.scenario.sql) - .set_location(&self.args.location); + .set_location(&self.args.location) + .set_use_query_cache(self.args.use_query_cache); if let Some(max_results) = self.args.max_results { query_builder = query_builder.set_max_results(max_results); diff --git a/src/bigquery/benchmarks/queries/src/telemetry.rs b/src/bigquery/benchmarks/queries/src/telemetry.rs index b377c91407..0c353cba64 100644 --- a/src/bigquery/benchmarks/queries/src/telemetry.rs +++ b/src/bigquery/benchmarks/queries/src/telemetry.rs @@ -14,15 +14,50 @@ 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"; +#[derive(Clone, Debug)] +struct GenericNodeDetector { + id: String, + location: String, + namespace: String, +} + +impl GenericNodeDetector { + pub fn new() -> Self { + let id = Uuid::new_v4().to_string(); + Self { + id, + location: "us-central1".to_string(), + namespace: "bigquery-benchmark-queries".to_string(), + } + } +} + +impl ResourceDetector for GenericNodeDetector { + fn detect(&self) -> Resource { + Resource::builder_empty() + .with_attributes([ + KeyValue::new("location", self.location.clone()), + KeyValue::new("namespace", self.namespace.clone()), + KeyValue::new("node_id", self.id.clone()), + ]) + .build() + } +} + /// Holds providers that need graceful flush and shutdown upon completion. pub struct TelemetryGuard { tracer_provider: Option, @@ -32,15 +67,17 @@ pub struct TelemetryGuard { impl TelemetryGuard { /// Flushes and shuts down telemetry providers. pub fn shutdown(self) { - if let Some(tp) = self.tracer_provider - && let Err(e) = tp.shutdown() - { - eprintln!("Error shutting down trace provider: {e:?}"); + if let Some(tp) = self.tracer_provider { + let _ = tp.force_flush(); + if let Err(e) = tp.shutdown() { + eprintln!("Error shutting down trace provider: {e:?}"); + } } - if let Some(mp) = self.meter_provider - && let Err(e) = mp.shutdown() - { - eprintln!("Error shutting down meter provider: {e:?}"); + if let Some(mp) = self.meter_provider { + let _ = mp.force_flush(); + if let Err(e) = mp.shutdown() { + eprintln!("Error shutting down meter provider: {e:?}"); + } } } } @@ -65,13 +102,21 @@ pub async fn enable_telemetry( 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_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_credentials(credentials.clone()) + .with_detector(node); if let Some(endpoint_str) = &args.otlp_endpoint { let uri = Uri::from_str(endpoint_str)?; From 9171eef076eaa7d2b996defab341e674d1419ff9 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Thu, 20 Aug 2026 17:31:50 +0000 Subject: [PATCH 03/11] impl: remove iteration limit when duration is not set --- src/bigquery/benchmarks/queries/README.md | 75 ++++++++++++++++++- src/bigquery/benchmarks/queries/src/args.rs | 15 +--- src/bigquery/benchmarks/queries/src/main.rs | 2 +- src/bigquery/benchmarks/queries/src/runner.rs | 3 +- 4 files changed, 79 insertions(+), 16 deletions(-) diff --git a/src/bigquery/benchmarks/queries/README.md b/src/bigquery/benchmarks/queries/README.md index f64ed7dad2..c824cbfa67 100644 --- a/src/bigquery/benchmarks/queries/README.md +++ b/src/bigquery/benchmarks/queries/README.md @@ -37,8 +37,10 @@ Benchmarks the Rust BigQuery client library (`google-cloud-bigquery`), measuring > [!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`. -### 1. Zero-Setup Synthetic Benchmark (Default) +### 1. Zero-Setup Synthetic Benchmark Runs 10 iterations per task with 4 concurrent tasks, generating and streaming 100,000 rows per query: @@ -90,6 +92,77 @@ cargo run --release -p bigquery-benchmark-queries -- \ --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. + > [!TIP] > **Real-Time Reporting & Graceful Shutdown (`Ctrl+C`):** When `--output-dir` is provided, the benchmark writes raw sample CSV rows and updates the summary report JSON on disk in **real-time** as each query iteration completes. You can press `Ctrl+C` at any point during a long endurance test to immediately stop worker tasks, output the summary report to stdout, flush all OpenTelemetry metrics to Google Cloud, and preserve the recorded samples and summary report on disk. diff --git a/src/bigquery/benchmarks/queries/src/args.rs b/src/bigquery/benchmarks/queries/src/args.rs index 8c83b8e420..9fd1e95b12 100644 --- a/src/bigquery/benchmarks/queries/src/args.rs +++ b/src/bigquery/benchmarks/queries/src/args.rs @@ -80,7 +80,7 @@ pub struct Args { /// Number of query iterations per worker task. /// - /// Defaults to 10 if `--duration` is not set. + /// If neither `--iterations` nor `--duration` is set, the benchmark runs indefinitely until interrupted (Ctrl+C). #[arg(long)] pub iterations: Option, @@ -140,15 +140,6 @@ impl Args { } Ok(()) } - - /// Returns the effective iterations limit per worker (defaults to 10 if duration is not set). - pub fn effective_iterations(&self) -> Option { - match (self.iterations, self.duration) { - (Some(i), _) => Some(i), - (None, Some(_)) => None, - (None, None) => Some(10), - } - } } #[cfg(test)] @@ -160,7 +151,7 @@ mod tests { let args = Args::parse_from(["bigquery-benchmark-queries"]); assert!(args.validate().is_ok()); assert_eq!(args.task_count, 1); - assert_eq!(args.effective_iterations(), Some(10)); + assert_eq!(args.iterations, None); assert!(!args.use_query_cache); } @@ -183,7 +174,7 @@ mod tests { fn test_duration_mode() { let args = Args::parse_from(["bigquery-benchmark-queries", "--duration", "5m"]); assert!(args.validate().is_ok()); - assert_eq!(args.effective_iterations(), None); + assert_eq!(args.iterations, None); assert_eq!(args.duration, Some(Duration::from_secs(300))); } } diff --git a/src/bigquery/benchmarks/queries/src/main.rs b/src/bigquery/benchmarks/queries/src/main.rs index a90ee272a1..80d9c72aea 100644 --- a/src/bigquery/benchmarks/queries/src/main.rs +++ b/src/bigquery/benchmarks/queries/src/main.rs @@ -46,7 +46,7 @@ async fn main() -> anyhow::Result<()> { scenario = %scenario.name, description = %scenario.description, task_count = args.task_count, - effective_iterations = ?args.effective_iterations(), + iterations = ?args.iterations, duration = ?args.duration, use_query_cache = args.use_query_cache, "Starting BigQuery benchmark" diff --git a/src/bigquery/benchmarks/queries/src/runner.rs b/src/bigquery/benchmarks/queries/src/runner.rs index 2dbc9bfa51..478b161e9f 100644 --- a/src/bigquery/benchmarks/queries/src/runner.rs +++ b/src/bigquery/benchmarks/queries/src/runner.rs @@ -39,11 +39,10 @@ impl TaskRunner<'_> { tokio::time::sleep(self.args.rampup_period * self.task_id as u32).await; } - let effective_iterations = self.args.effective_iterations(); let mut iteration = 0_u64; loop { - if let Some(max_iter) = effective_iterations + if let Some(max_iter) = self.args.iterations && iteration >= max_iter { break; From 7742f8a5f2aac7ea3cb2f016c366d65ef36d9425 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Fri, 21 Aug 2026 14:36:49 +0000 Subject: [PATCH 04/11] impl: improve error logs and add query timeout --- src/bigquery/benchmarks/queries/src/args.rs | 4 + .../benchmarks/queries/src/reporter.rs | 171 +++++++++++++++--- src/bigquery/benchmarks/queries/src/runner.rs | 48 ++++- .../benchmarks/queries/src/telemetry.rs | 30 ++- 4 files changed, 218 insertions(+), 35 deletions(-) diff --git a/src/bigquery/benchmarks/queries/src/args.rs b/src/bigquery/benchmarks/queries/src/args.rs index 9fd1e95b12..d59f172612 100644 --- a/src/bigquery/benchmarks/queries/src/args.rs +++ b/src/bigquery/benchmarks/queries/src/args.rs @@ -106,6 +106,10 @@ pub struct Args { #[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, diff --git a/src/bigquery/benchmarks/queries/src/reporter.rs b/src/bigquery/benchmarks/queries/src/reporter.rs index df689bf766..3b5ec8f445 100644 --- a/src/bigquery/benchmarks/queries/src/reporter.rs +++ b/src/bigquery/benchmarks/queries/src/reporter.rs @@ -18,10 +18,21 @@ use crate::sample::{Sample, SampleStatus}; use crate::scenarios::Scenario; use serde::{Deserialize, Serialize}; use std::fs::File; -use std::io::Write; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use std::io::{BufWriter, Write}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::sync::mpsc::Receiver; +/// 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 { @@ -37,6 +48,7 @@ pub struct BenchmarkReport { pub send_duration: Option, pub poll_duration: Option, pub read_duration: Option, + pub errors: Vec, } impl BenchmarkReport { @@ -88,6 +100,36 @@ impl BenchmarkReport { ); } + if !self.errors.is_empty() { + println!("\n-------------------------------------------------------"); + println!(" QUERY ERRORS ({} total)", self.errors.len()); + println!("-------------------------------------------------------"); + for (idx, err) in self.errors.iter().take(20).enumerate() { + let job_id = if !err.final_job_id.is_empty() && err.final_job_id != "N/A" { + &err.final_job_id + } else if !err.initial_job_id.is_empty() && err.initial_job_id != "N/A" { + &err.initial_job_id + } else { + "N/A" + }; + println!( + " {}. Task {:>2} | Iteration {:>6} | Offset: {:>7.1}s | Job: {}", + idx + 1, + err.task_id, + err.iteration, + err.offset_secs, + job_id + ); + println!(" Error: {}", err.error_message); + } + if self.errors.len() > 20 { + println!( + " ... and {} more error(s) recorded in full error log.", + self.errors.len() - 20 + ); + } + } + println!("=======================================================\n"); } } @@ -103,6 +145,7 @@ pub async fn collect_and_report( let mut success_send_durations = Vec::new(); let mut success_poll_durations = Vec::new(); let mut success_read_durations = Vec::new(); + let mut errors = Vec::new(); let mut success_count = 0_usize; let mut error_count = 0_usize; @@ -119,19 +162,33 @@ pub async fn collect_and_report( let csv_path = output_dir.join(format!("samples-{}-{}.csv", scenario.name, timestamp)); let json_path = output_dir.join(format!("summary-{}-{}.json", scenario.name, timestamp)); + let errors_path = output_dir.join(format!("errors-{}-{}.log", scenario.name, timestamp)); + + let csv_file = File::create(&csv_path)?; + let mut csv_writer = BufWriter::new(csv_file); + writeln!(csv_writer, "{}", Sample::HEADER)?; + csv_writer.flush()?; - let mut csv_file = File::create(&csv_path)?; - writeln!(csv_file, "{}", Sample::HEADER)?; - csv_file.flush()?; + let errors_file = File::create(&errors_path)?; + let mut errors_writer = BufWriter::new(errors_file); + writeln!( + errors_writer, + "# BigQuery Benchmark Error Log - Scenario: {}, Timestamp: {}", + scenario.name, timestamp + )?; + errors_writer.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()); - Some((csv_file, csv_path, json_path)) + Some((csv_writer, csv_path, json_path, errors_writer, errors_path)) } else { None }; + let mut last_json_write = Instant::now(); + while let Some(sample) = rx.recv().await { total_samples += 1; if sample.status == SampleStatus::Ok { @@ -144,33 +201,81 @@ pub async fn collect_and_report( total_bytes_processed += sample.bytes_processed; } else { error_count += 1; + let offset_secs = sample.start_offset_micros as f64 / 1_000_000.0; + let err_detail = ErrorDetail { + task_id: sample.task_id, + iteration: sample.iteration, + offset_secs, + initial_job_id: sample.initial_job_id.clone(), + final_job_id: sample.final_job_id.clone(), + error_message: sample.error_message.clone(), + }; + + let job_id = if !sample.final_job_id.is_empty() && sample.final_job_id != "N/A" { + &sample.final_job_id + } else if !sample.initial_job_id.is_empty() && sample.initial_job_id != "N/A" { + &sample.initial_job_id + } else { + "N/A" + }; + + // 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, offset_secs, job_id, sample.error_message + ); + + if let Some((_, _, _, errors_writer, _)) = &mut realtime_files { + let _ = writeln!( + errors_writer, + "Task: {}\nIteration: {}\nOffsetSecs: {:.3}\nInitialJobId: {}\nFinalJobId: {}\nError: {}\n--------------------------------------------------------------------------------", + sample.task_id, + sample.iteration, + offset_secs, + sample.initial_job_id, + sample.final_job_id, + sample.error_message + ); + let _ = errors_writer.flush(); + } + + errors.push(err_detail); } if sample.retry_detected { retries_detected_count += 1; } - if let Some((csv_file, _, json_path)) = &mut realtime_files { - let _ = writeln!(csv_file, "{}", sample.to_csv_row()); - let _ = csv_file.flush(); - - let stats = ReportStats { - scenario_name: &scenario.name, - task_count: args.task_count, - total_samples, - success_count, - error_count, - retries_detected_count, - total_rows_read, - total_bytes_processed, - success_total_durations: &success_total_durations, - success_send_durations: &success_send_durations, - success_poll_durations: &success_poll_durations, - success_read_durations: &success_read_durations, - }; - let current_report = build_report(&stats); - if let Ok(json_file) = File::create(json_path) { - let _ = serde_json::to_writer_pretty(json_file, ¤t_report); + if let Some((csv_writer, _, json_path, _, _)) = &mut realtime_files { + if let Err(err) = writeln!(csv_writer, "{}", sample.to_csv_row()) { + tracing::error!("Failed to write CSV sample row to disk: {err:?}"); + } + if let Err(err) = csv_writer.flush() { + tracing::error!("Failed to flush CSV sample file: {err:?}"); + } + + // Periodically update summary JSON (every 5 seconds) to avoid high disk I/O and sorting overhead + if last_json_write.elapsed() >= Duration::from_secs(5) { + last_json_write = Instant::now(); + let stats = ReportStats { + scenario_name: &scenario.name, + task_count: args.task_count, + total_samples, + success_count, + error_count, + retries_detected_count, + total_rows_read, + total_bytes_processed, + success_total_durations: &success_total_durations, + success_send_durations: &success_send_durations, + success_poll_durations: &success_poll_durations, + success_read_durations: &success_read_durations, + errors: &errors, + }; + let current_report = build_report(&stats); + if let Ok(json_file) = File::create(&*json_path) { + let _ = serde_json::to_writer_pretty(json_file, ¤t_report); + } } } } @@ -188,14 +293,22 @@ pub async fn collect_and_report( success_send_durations: &success_send_durations, success_poll_durations: &success_poll_durations, success_read_durations: &success_read_durations, + errors: &errors, }; let report = build_report(&stats); report.print_stdout(); - if let Some((_, csv_path, json_path)) = realtime_files { + if let Some((_, csv_path, json_path, _, errors_path)) = &realtime_files { + // Save final complete JSON summary + if let Ok(json_file) = File::create(json_path) { + let _ = serde_json::to_writer_pretty(json_file, &report); + } println!("Final samples saved to: {}", csv_path.display()); println!("Final summary saved to: {}", json_path.display()); + if error_count > 0 { + println!("Errors logged to: {}", errors_path.display()); + } } Ok(report) @@ -214,6 +327,7 @@ struct ReportStats<'a> { success_send_durations: &'a [Duration], success_poll_durations: &'a [Duration], success_read_durations: &'a [Duration], + errors: &'a [ErrorDetail], } fn build_report(stats: &ReportStats) -> BenchmarkReport { @@ -230,5 +344,6 @@ fn build_report(stats: &ReportStats) -> BenchmarkReport { send_duration: metrics::compute_metrics(stats.success_send_durations), poll_duration: metrics::compute_metrics(stats.success_poll_durations), read_duration: metrics::compute_metrics(stats.success_read_durations), + errors: stats.errors.to_vec(), } } diff --git a/src/bigquery/benchmarks/queries/src/runner.rs b/src/bigquery/benchmarks/queries/src/runner.rs index 478b161e9f..642117a2ef 100644 --- a/src/bigquery/benchmarks/queries/src/runner.rs +++ b/src/bigquery/benchmarks/queries/src/runner.rs @@ -64,10 +64,46 @@ impl TaskRunner<'_> { scenario = %self.scenario.name ); - let sample = self + let sample_fut = self .execute_iteration(iteration, start_offset_micros, iter_start) - .instrument(iteration_span) - .await; + .instrument(iteration_span); + + let sample = match tokio::time::timeout(self.args.query_timeout, sample_fut).await { + Ok(sample) => sample, + Err(_) => { + let total_duration = iter_start.elapsed(); + metrics::inc_total_queries(); + metrics::inc_error_queries(); + tracing::error!( + task_id = self.task_id, + iteration, + "Query iteration timed out after {:?}", + self.args.query_timeout + ); + let sample = Sample { + task_id: self.task_id, + iteration, + start_offset_micros, + send_duration_micros: 0, + poll_duration_micros: 0, + read_duration_micros: 0, + total_duration_micros: total_duration.as_micros(), + rows_count: 0, + bytes_processed: 0, + cache_hit: false, + initial_job_id: String::new(), + final_job_id: String::new(), + retry_detected: false, + status: SampleStatus::Timeout, + error_message: format!( + "Query iteration timed out after {:?}", + self.args.query_timeout + ), + }; + self.metrics.record_sample(&self.scenario.name, &sample); + sample + } + }; let _ = self.tx.send(sample).await; iteration += 1; @@ -124,7 +160,7 @@ impl TaskRunner<'_> { final_job_id: String::new(), retry_detected: false, status: SampleStatus::Error, - error_message: err.to_string(), + error_message: format!("Query::send: {err:#}"), }; self.metrics.record_sample(&self.scenario.name, &sample); return sample; @@ -178,7 +214,7 @@ impl TaskRunner<'_> { final_job_id: String::new(), retry_detected: false, status: SampleStatus::Error, - error_message: err.to_string(), + error_message: format!("Query::until_done: {err:#}"), }; self.metrics.record_sample(&self.scenario.name, &sample); return sample; @@ -230,7 +266,7 @@ impl TaskRunner<'_> { } Err(err) => { tracing::error!(self.task_id, iteration, "Error streaming rows: {err:?}"); - read_error = Some(err.to_string()); + read_error = Some(format!("CompleteQuery::read: {err:#}")); break; } } diff --git a/src/bigquery/benchmarks/queries/src/telemetry.rs b/src/bigquery/benchmarks/queries/src/telemetry.rs index 0c353cba64..550db77237 100644 --- a/src/bigquery/benchmarks/queries/src/telemetry.rs +++ b/src/bigquery/benchmarks/queries/src/telemetry.rs @@ -97,7 +97,35 @@ pub async fn enable_telemetry( .with_writer(std::io::stderr) .with_filter(env_filter); - let registry = tracing_subscriber::Registry::default().with(fmt_layer); + 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}"); From 2a406902bda44966b29ed662cf305c0accd3d958 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Tue, 25 Aug 2026 17:21:14 +0000 Subject: [PATCH 05/11] fix: reduce span logs --- src/bigquery/benchmarks/queries/src/telemetry.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/bigquery/benchmarks/queries/src/telemetry.rs b/src/bigquery/benchmarks/queries/src/telemetry.rs index 550db77237..07f6b2dfb3 100644 --- a/src/bigquery/benchmarks/queries/src/telemetry.rs +++ b/src/bigquery/benchmarks/queries/src/telemetry.rs @@ -93,7 +93,7 @@ pub async fn enable_telemetry( let fmt_layer = tracing_subscriber::fmt::layer() .with_level(true) .with_thread_ids(true) - .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE) + .with_span_events(FmtSpan::NONE) .with_writer(std::io::stderr) .with_filter(env_filter); From 950f2be202b63804d1ef3dcac5399db4c89af3b5 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Mon, 14 Sep 2026 15:07:47 +0000 Subject: [PATCH 06/11] docs: minor docs ajustments --- src/bigquery/benchmarks/queries/README.md | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/src/bigquery/benchmarks/queries/README.md b/src/bigquery/benchmarks/queries/README.md index c824cbfa67..ae8c8e2225 100644 --- a/src/bigquery/benchmarks/queries/README.md +++ b/src/bigquery/benchmarks/queries/README.md @@ -40,7 +40,7 @@ Benchmarks the Rust BigQuery client library (`google-cloud-bigquery`), measuring > > **Indefinite Execution by Default:** If neither `--iterations` nor `--duration` is specified, the benchmark runs indefinitely until interrupted with `Ctrl+C`. -### 1. Zero-Setup Synthetic Benchmark +### Synthetic Benchmark Runs 10 iterations per task with 4 concurrent tasks, generating and streaming 100,000 rows per query: @@ -53,7 +53,7 @@ cargo run --release -p bigquery-benchmark-queries -- \ --output-dir ./results ``` -### 2. Public Dataset Query Benchmark +### Public Dataset Query Benchmark Benchmark streaming 50,000 rows from the USA names public dataset: @@ -66,7 +66,7 @@ cargo run --release -p bigquery-benchmark-queries -- \ --output-dir ./results ``` -### 3. Custom SQL Query Benchmark +### Custom SQL Query Benchmark ```shell cargo run --release -p bigquery-benchmark-queries -- \ @@ -120,9 +120,9 @@ tail -f results/samples-synthetic-100k-*.csv To execute endurance tests on a GCE VM and stream all logs directly into **Google Cloud Console (Cloud Logging)**: -1. **Create the VM and Install Ops Agent**: +1. **Create the VM and install Ops Agent**: ```shell - # 1. Create a VM with full cloud-platform scope + # Create a VM with full cloud-platform scope gcloud compute instances create bq-benchmark-vm \ --zone=us-central1-a \ --machine-type=c2-standard-4 \ @@ -137,7 +137,7 @@ To execute endurance tests on a GCE VM and stream all logs directly into **Googl " ``` -2. **Run in the Background via `systemd-cat`**: +2. **Run in the background via `systemd-cat`**: Using `systemd-cat` tags the output in the system journal so the Ops Agent automatically forwards stdout and stderr to Cloud Logging: ```shell gcloud compute ssh bq-benchmark-vm --zone=us-central1-a @@ -163,9 +163,6 @@ To execute endurance tests on a GCE VM and stream all logs directly into **Googl ``` * Click **Stream logs** in the top right to watch real-time benchmark execution logs arrive in the console. -> [!TIP] -> **Real-Time Reporting & Graceful Shutdown (`Ctrl+C`):** When `--output-dir` is provided, the benchmark writes raw sample CSV rows and updates the summary report JSON on disk in **real-time** as each query iteration completes. You can press `Ctrl+C` at any point during a long endurance test to immediately stop worker tasks, output the summary report to stdout, flush all OpenTelemetry metrics to Google Cloud, and preserve the recorded samples and summary report on disk. - --- ## OpenTelemetry & Cloud Observability @@ -202,7 +199,7 @@ sum by (scenario, status) (rate(workload_googleapis_com:bigquery_queries_total[1 ``` #### 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()`): +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])) ``` From 95d17ae65ef222412ef0c7563f304568114e2f8b Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Mon, 14 Sep 2026 16:12:22 +0000 Subject: [PATCH 07/11] fix: address ai review comments --- .../benchmarks/queries/src/reporter.rs | 34 ++++++++++++------- src/bigquery/benchmarks/queries/src/runner.rs | 29 ++++++++++------ 2 files changed, 39 insertions(+), 24 deletions(-) diff --git a/src/bigquery/benchmarks/queries/src/reporter.rs b/src/bigquery/benchmarks/queries/src/reporter.rs index 3b5ec8f445..3748e59d93 100644 --- a/src/bigquery/benchmarks/queries/src/reporter.rs +++ b/src/bigquery/benchmarks/queries/src/reporter.rs @@ -246,15 +246,12 @@ pub async fn collect_and_report( retries_detected_count += 1; } - if let Some((csv_writer, _, json_path, _, _)) = &mut realtime_files { + if let Some((csv_writer, _, json_path, errors_writer, _)) = &mut realtime_files { if let Err(err) = writeln!(csv_writer, "{}", sample.to_csv_row()) { tracing::error!("Failed to write CSV sample row to disk: {err:?}"); } - if let Err(err) = csv_writer.flush() { - tracing::error!("Failed to flush CSV sample file: {err:?}"); - } - // Periodically update summary JSON (every 5 seconds) to avoid high disk I/O and sorting overhead + // Periodically update summary JSON and flush buffers (every 5 seconds) to avoid high disk I/O and sorting overhead if last_json_write.elapsed() >= Duration::from_secs(5) { last_json_write = Instant::now(); let stats = ReportStats { @@ -273,9 +270,14 @@ pub async fn collect_and_report( errors: &errors, }; let current_report = build_report(&stats); - if let Ok(json_file) = File::create(&*json_path) { - let _ = serde_json::to_writer_pretty(json_file, ¤t_report); - } + let json_path_clone = json_path.clone(); + tokio::task::block_in_place(|| { + let _ = csv_writer.flush(); + let _ = errors_writer.flush(); + if let Ok(json_file) = File::create(&json_path_clone) { + let _ = serde_json::to_writer_pretty(json_file, ¤t_report); + } + }); } } } @@ -299,11 +301,17 @@ pub async fn collect_and_report( report.print_stdout(); - if let Some((_, csv_path, json_path, _, errors_path)) = &realtime_files { - // Save final complete JSON summary - if let Ok(json_file) = File::create(json_path) { - let _ = serde_json::to_writer_pretty(json_file, &report); - } + if let Some((csv_writer, csv_path, json_path, errors_writer, errors_path)) = &mut realtime_files + { + let json_path_clone = json_path.clone(); + tokio::task::block_in_place(|| { + let _ = csv_writer.flush(); + let _ = errors_writer.flush(); + // Save final complete JSON summary + if let Ok(json_file) = File::create(&json_path_clone) { + let _ = serde_json::to_writer_pretty(json_file, &report); + } + }); println!("Final samples saved to: {}", csv_path.display()); println!("Final summary saved to: {}", json_path.display()); if error_count > 0 { diff --git a/src/bigquery/benchmarks/queries/src/runner.rs b/src/bigquery/benchmarks/queries/src/runner.rs index 642117a2ef..14a0a3b3a2 100644 --- a/src/bigquery/benchmarks/queries/src/runner.rs +++ b/src/bigquery/benchmarks/queries/src/runner.rs @@ -257,20 +257,27 @@ impl TaskRunner<'_> { if self.args.read_results { let read_span = tracing::info_span!("bigquery.read_rows", task_id = self.task_id, iteration); - let _guard = read_span.enter(); - let mut rows = complete_query.read(); - while let Some(row_result) = rows.next().await { - match row_result { - Ok(_) => { - rows_count += 1; - } - Err(err) => { - tracing::error!(self.task_id, iteration, "Error streaming rows: {err:?}"); - read_error = Some(format!("CompleteQuery::read: {err:#}")); - break; + async { + let mut rows = complete_query.read(); + while let Some(row_result) = rows.next().await { + match row_result { + Ok(_) => { + rows_count += 1; + } + Err(err) => { + tracing::error!( + self.task_id, + iteration, + "Error streaming rows: {err:?}" + ); + read_error = Some(format!("CompleteQuery::read: {err:#}")); + break; + } } } } + .instrument(read_span) + .await; } let read_duration = read_start.elapsed(); let total_duration = iter_start.elapsed(); From bcfde7152909dc5914d7635b98595fa09c1b9094 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Mon, 14 Sep 2026 16:28:33 +0000 Subject: [PATCH 08/11] fix: fmt all the things --- Cargo.toml | 2 +- src/bigquery/benchmarks/queries/Cargo.toml | 38 +++--- src/bigquery/benchmarks/queries/README.md | 133 +++++++++++++++------ 3 files changed, 114 insertions(+), 59 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 4afc31197e..72c175522b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -69,9 +69,9 @@ members = [ "guide/samples", "src/auth", "src/bigquery", - "src/bigquery/benchmarks/queries", "src/bigquery-derive", "src/bigquery-read", + "src/bigquery/benchmarks/queries", "src/bigquery/examples", "src/bigquery/grpc-mock", "src/bigtable", diff --git a/src/bigquery/benchmarks/queries/Cargo.toml b/src/bigquery/benchmarks/queries/Cargo.toml index 3481187c5c..86fd0b6425 100644 --- a/src/bigquery/benchmarks/queries/Cargo.toml +++ b/src/bigquery/benchmarks/queries/Cargo.toml @@ -24,26 +24,26 @@ keywords.workspace = true categories.workspace = true [dependencies] -anyhow.workspace = true -clap = { workspace = true, features = ["derive", "env", "help", "std", "usage"] } -futures.workspace = true -google-cloud-auth.workspace = true -google-cloud-bigquery = { workspace = true, features = ["default-rustls-provider"] } +anyhow.workspace = true +clap = { workspace = true, features = ["derive", "env", "help", "std", "usage"] } +futures.workspace = true +google-cloud-auth.workspace = true +google-cloud-bigquery = { workspace = true, features = ["default-rustls-provider"] } google-cloud-bigquery-v2.workspace = true -google-cloud-gax.workspace = true -humantime.workspace = true -integration-tests-o11y.workspace = true -opentelemetry = { workspace = true, features = ["trace", "metrics"] } -opentelemetry_sdk = { workspace = true, features = ["rt-tokio", "trace", "metrics"] } -rand.workspace = true -serde = { workspace = true, features = ["derive"] } -serde_json.workspace = true -tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time", "sync", "signal"] } -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 +google-cloud-gax.workspace = true +humantime.workspace = true +integration-tests-o11y.workspace = true +opentelemetry = { workspace = true, features = ["metrics", "trace"] } +opentelemetry_sdk = { workspace = true, features = ["metrics", "rt-tokio", "trace"] } +rand.workspace = true +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 index ae8c8e2225..3f2861fd25 100644 --- a/src/bigquery/benchmarks/queries/README.md +++ b/src/bigquery/benchmarks/queries/README.md @@ -1,48 +1,64 @@ # 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. +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. +- **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-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-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. + - `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: +1. **Authentication**: Ensure Application Default Credentials (ADC) are + configured: + ```shell gcloud auth application-default login ``` -2. **Project ID**: - Set the project ID environment variable: +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. +> [!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`. +> **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: +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 -- \ @@ -77,11 +93,12 @@ cargo run --release -p bigquery-benchmark-queries -- \ --iterations 25 ``` ---- +______________________________________________________________________ ## Running Endurance Tests -To test connection stability, memory leaks, and token refreshing over a prolonged period (e.g. 1 hour): +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 -- \ @@ -94,7 +111,9 @@ cargo run --release -p bigquery-benchmark-queries -- \ ### 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`: +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" \ @@ -108,6 +127,7 @@ TS=$(date +%s); RUSTFLAGS="-C target-cpu=native" \ ``` You can monitor the progress in real time with: + ```shell # Follow tracing & runtime logs tail -f bq-bm-${TS}.log @@ -118,9 +138,11 @@ tail -f results/samples-synthetic-100k-*.csv ### Running on a Google Compute Engine (GCE) VM & Viewing Cloud Console Logs -To execute endurance tests on a GCE VM and stream all logs directly into **Google Cloud Console (Cloud Logging)**: +To execute endurance tests on a GCE VM and stream all logs directly into +**Google Cloud Console (Cloud Logging)**: 1. **Create the VM and install Ops Agent**: + ```shell # Create a VM with full cloud-platform scope gcloud compute instances create bq-benchmark-vm \ @@ -137,8 +159,10 @@ To execute endurance tests on a GCE VM and stream all logs directly into **Googl " ``` -2. **Run in the background via `systemd-cat`**: - Using `systemd-cat` tags the output in the system journal so the Ops Agent automatically forwards stdout and stderr to Cloud Logging: +1. **Run in the background via `systemd-cat`**: Using `systemd-cat` tags the + output in the system journal so the Ops Agent automatically forwards stdout + and stderr to Cloud Logging: + ```shell gcloud compute ssh bq-benchmark-vm --zone=us-central1-a @@ -154,34 +178,44 @@ To execute endurance tests on a GCE VM and stream all logs directly into **Googl 2> >(tee bq-bm-${TS}.log | systemd-cat -t bq-benchmark) Logs Explorer** in Google Cloud Console. - * Filter by your identifier: +1. **View Logs in Google Cloud Console**: + + - Navigate to **Logging > 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. + - 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: +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.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. +- `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` @@ -190,34 +224,49 @@ When `--project-id` is specified, the benchmark automatically connects to `telem ### 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: +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): + +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: + +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]))) @@ -230,7 +279,9 @@ histogram_quantile(0.50, sum by (le, scenario) (rate(workload_googleapis_com:big ``` #### 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]))) @@ -243,7 +294,9 @@ histogram_quantile(0.95, sum by (le, scenario) (rate(workload_googleapis_com:big ``` #### 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])) @@ -252,11 +305,13 @@ sum by (scenario) (rate(workload_googleapis_com:bigquery_queries_rows_read[1m])) 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: +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 \ From a8819485710fbda162bdb63810be47091ad62cdc Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Mon, 14 Sep 2026 19:01:21 +0000 Subject: [PATCH 09/11] cleanup: benchmark code tidy up --- Cargo.lock | 4 - src/bigquery/benchmarks/queries/Cargo.toml | 36 +- src/bigquery/benchmarks/queries/src/args.rs | 43 +- src/bigquery/benchmarks/queries/src/main.rs | 32 +- .../benchmarks/queries/src/metrics.rs | 330 ++++++------- .../benchmarks/queries/src/reporter.rs | 447 +++++++++--------- src/bigquery/benchmarks/queries/src/runner.rs | 188 ++------ src/bigquery/benchmarks/queries/src/sample.rs | 109 ++++- .../benchmarks/queries/src/scenarios.rs | 156 +++--- .../benchmarks/queries/src/telemetry.rs | 44 +- 10 files changed, 633 insertions(+), 756 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2194d3af91..697204f985 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -425,16 +425,12 @@ version = "0.0.0" dependencies = [ "anyhow", "clap", - "futures", "google-cloud-auth", "google-cloud-bigquery", - "google-cloud-bigquery-v2", - "google-cloud-gax", "humantime", "integration-tests-o11y", "opentelemetry", "opentelemetry_sdk", - "rand 0.10.2", "serde", "serde_json", "tokio", diff --git a/src/bigquery/benchmarks/queries/Cargo.toml b/src/bigquery/benchmarks/queries/Cargo.toml index 86fd0b6425..390225f32b 100644 --- a/src/bigquery/benchmarks/queries/Cargo.toml +++ b/src/bigquery/benchmarks/queries/Cargo.toml @@ -24,26 +24,22 @@ keywords.workspace = true categories.workspace = true [dependencies] -anyhow.workspace = true -clap = { workspace = true, features = ["derive", "env", "help", "std", "usage"] } -futures.workspace = true -google-cloud-auth.workspace = true -google-cloud-bigquery = { workspace = true, features = ["default-rustls-provider"] } -google-cloud-bigquery-v2.workspace = true -google-cloud-gax.workspace = true -humantime.workspace = true -integration-tests-o11y.workspace = true -opentelemetry = { workspace = true, features = ["metrics", "trace"] } -opentelemetry_sdk = { workspace = true, features = ["metrics", "rt-tokio", "trace"] } -rand.workspace = true -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 +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/src/args.rs b/src/bigquery/benchmarks/queries/src/args.rs index d59f172612..83163e40f1 100644 --- a/src/bigquery/benchmarks/queries/src/args.rs +++ b/src/bigquery/benchmarks/queries/src/args.rs @@ -119,10 +119,6 @@ pub struct Args { /// Defaults to `https://telemetry.googleapis.com` if `--project-id` is provided. #[arg(long)] pub otlp_endpoint: Option, - - /// Whether to log debug details for retry decisions. - #[arg(long)] - pub debug_retry: bool, } impl Args { @@ -136,10 +132,9 @@ impl Args { "When using `--scenario custom`, either `--sql` or `--sql-file` must be provided." ); } - if let Some(iterations) = self.iterations - && iterations == 0 - && self.duration.is_none() - { + // 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(()) @@ -151,34 +146,18 @@ mod tests { use super::*; #[test] - fn test_default_args_validation() { - let args = Args::parse_from(["bigquery-benchmark-queries"]); - assert!(args.validate().is_ok()); - assert_eq!(args.task_count, 1); - assert_eq!(args.iterations, None); - assert!(!args.use_query_cache); - } - - #[test] - fn test_custom_scenario_validation() { - let args = Args::parse_from(["bigquery-benchmark-queries", "--scenario", "custom"]); + 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", - "--scenario", - "custom", - "--sql", - "SELECT 1", + "--iterations", + "0", + "--duration", + "5m", ]); - assert!(args.validate().is_ok()); - } - - #[test] - fn test_duration_mode() { - let args = Args::parse_from(["bigquery-benchmark-queries", "--duration", "5m"]); - assert!(args.validate().is_ok()); - assert_eq!(args.iterations, None); - assert_eq!(args.duration, Some(Duration::from_secs(300))); + assert!(args.validate().is_err()); } } diff --git a/src/bigquery/benchmarks/queries/src/main.rs b/src/bigquery/benchmarks/queries/src/main.rs index 80d9c72aea..bde5f1b6b7 100644 --- a/src/bigquery/benchmarks/queries/src/main.rs +++ b/src/bigquery/benchmarks/queries/src/main.rs @@ -28,10 +28,16 @@ use google_cloud_auth::credentials::Builder as CredentialsBuilder; use google_cloud_bigquery::client::BigQuery; use metrics::OtelMetrics; use scenarios::Scenario; -use std::collections::BTreeMap; -use std::time::Instant; +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()?; @@ -52,15 +58,14 @@ async fn main() -> anyhow::Result<()> { "Starting BigQuery benchmark" ); - // Spawn periodic runtime monitor and counter logger + // 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); - let monitor_freq = std::time::Duration::from_secs(5); tokio::spawn(async move { for metrics in runtime_monitor.intervals() { - let counters = BTreeMap::from_iter(metrics::get_counters()); - tracing::info!("Counters = {:?} RuntimeMetrics = {:?}", counters, metrics); - tokio::time::sleep(monitor_freq).await; + tracing::info!("RuntimeMetrics = {:?}", metrics); + tokio::time::sleep(RUNTIME_MONITOR_INTERVAL).await; } }); @@ -73,11 +78,13 @@ async fn main() -> anyhow::Result<()> { } let client = client_builder.build().await?; - let otel_metrics = OtelMetrics::new(); - otel_metrics.init_scenario(&scenario.name); + let otel_metrics = OtelMetrics::new(scenario.name); - let channel_capacity = (1024 * args.task_count).max(64); - let (tx, rx) = tokio::sync::mpsc::channel(channel_capacity); + 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(); @@ -144,9 +151,6 @@ async fn main() -> anyhow::Result<()> { Err(err) => tracing::error!("Reporter task panicked: {err:?}"), } - let final_counters = BTreeMap::from_iter(metrics::get_counters()); - tracing::info!("Final counters: {:?}", final_counters); - telemetry_guard.shutdown(); Ok(()) diff --git a/src/bigquery/benchmarks/queries/src/metrics.rs b/src/bigquery/benchmarks/queries/src/metrics.rs index 868783f567..5b679a008c 100644 --- a/src/bigquery/benchmarks/queries/src/metrics.rs +++ b/src/bigquery/benchmarks/queries/src/metrics.rs @@ -12,20 +12,46 @@ // 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::sync::atomic::{AtomicU64, Ordering}; 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, } @@ -60,219 +86,129 @@ pub fn compute_metrics(latencies: &[Duration]) -> Option { }) } -// In-process global atomic counters for quick telemetry reporting. -static TOTAL_QUERIES: AtomicU64 = AtomicU64::new(0); -static SUCCESS_QUERIES: AtomicU64 = AtomicU64::new(0); -static ERROR_QUERIES: AtomicU64 = AtomicU64::new(0); -static RETRIED_QUERIES: AtomicU64 = AtomicU64::new(0); -static TOTAL_ROWS: AtomicU64 = AtomicU64::new(0); -static TOTAL_BYTES: AtomicU64 = AtomicU64::new(0); - -#[inline] -pub fn inc_total_queries() { - TOTAL_QUERIES.fetch_add(1, Ordering::SeqCst); -} - -#[inline] -pub fn inc_success_queries() { - SUCCESS_QUERIES.fetch_add(1, Ordering::SeqCst); -} - -#[inline] -pub fn inc_error_queries() { - ERROR_QUERIES.fetch_add(1, Ordering::SeqCst); -} - -#[inline] -pub fn inc_retried_queries() { - RETRIED_QUERIES.fetch_add(1, Ordering::SeqCst); -} - -#[inline] -pub fn add_rows_read(count: u64) { - TOTAL_ROWS.fetch_add(count, Ordering::SeqCst); -} - -#[inline] -pub fn add_bytes_processed(bytes: u64) { - TOTAL_BYTES.fetch_add(bytes, Ordering::SeqCst); -} - -/// Returns a snapshot of in-process counters. -pub fn get_counters() -> [(&'static str, u64); 6] { - [ - ("total_queries", TOTAL_QUERIES.load(Ordering::Relaxed)), - ("success_queries", SUCCESS_QUERIES.load(Ordering::Relaxed)), - ("error_queries", ERROR_QUERIES.load(Ordering::Relaxed)), - ("retried_queries", RETRIED_QUERIES.load(Ordering::Relaxed)), - ("total_rows_read", TOTAL_ROWS.load(Ordering::Relaxed)), - ("total_bytes_processed", TOTAL_BYTES.load(Ordering::Relaxed)), - ] -} - -/// OpenTelemetry metrics instruments. +/// 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 { - pub queries_total: Counter, - pub queries_success: Counter, - pub queries_error: Counter, - pub queries_retried: Counter, - pub rows_read: Counter, - pub bytes_processed: Counter, - pub query_duration: Histogram, - pub send_duration: Histogram, - pub poll_duration: Histogram, - pub read_duration: Histogram, + 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 { - pub fn new() -> Self { + /// 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 queries_total = meter - .u64_counter("bigquery.queries.total") - .with_description("Total number of BigQuery queries attempted") - .build(); - - let queries_success = meter - .u64_counter("bigquery.queries.success") - .with_description("Number of BigQuery queries completed successfully") - .build(); - - let queries_error = meter - .u64_counter("bigquery.queries.error") - .with_description("Number of BigQuery queries that failed") - .build(); - - let queries_retried = meter - .u64_counter("bigquery.queries.retries_detected") - .with_description("Number of queries where an under-the-hood job retry was detected") - .build(); - - let rows_read = meter - .u64_counter("bigquery.queries.rows_read") - .with_description("Total count of rows read from query results") - .build(); - - let bytes_processed = meter - .u64_counter("bigquery.queries.bytes_processed") - .with_description("Total estimated bytes processed by BigQuery jobs") - .build(); - - let query_duration = meter - .f64_histogram("bigquery.queries.duration_seconds") - .with_description("Total query end-to-end duration in seconds") - .build(); - - let send_duration = meter - .f64_histogram("bigquery.queries.send_duration_seconds") - .with_description("Duration for Query::send() execution") - .build(); - - let poll_duration = meter - .f64_histogram("bigquery.queries.poll_duration_seconds") - .with_description("Duration for Query::until_done() polling execution") - .build(); - - let read_duration = meter - .f64_histogram("bigquery.queries.read_duration_seconds") - .with_description("Duration for CompleteQuery::read() row streaming") - .build(); - - Self { - queries_total, - queries_success, - queries_error, - queries_retried, - rows_read, - bytes_processed, - query_duration, - send_duration, - poll_duration, - read_duration, - } - } - - /// Initializes all counter metrics with 0 so time series exist in Cloud Monitoring - /// even if no errors or retries occur during the benchmark run. - pub fn init_scenario(&self, scenario: &str) { - let ok_attrs = [ - KeyValue::new("scenario", scenario.to_string()), - KeyValue::new("status", "ok"), - ]; - let err_attrs = [ - KeyValue::new("scenario", scenario.to_string()), - KeyValue::new("status", "error"), - ]; - - self.queries_total.add(0, &ok_attrs); - self.queries_total.add(0, &err_attrs); - self.queries_success.add(0, &ok_attrs); - self.queries_error.add(0, &err_attrs); - self.queries_retried.add(0, &ok_attrs); - self.rows_read.add(0, &ok_attrs); - self.bytes_processed.add(0, &ok_attrs); + 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 } - pub fn record_sample(&self, scenario: &str, sample: &crate::sample::Sample) { - let is_ok = sample.status == crate::sample::SampleStatus::Ok; - let attrs = [ - KeyValue::new("scenario", scenario.to_string()), - KeyValue::new("status", if is_ok { "ok" } else { "error" }), - ]; + /// 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); + self.queries_total.add(1, attrs); if is_ok { - self.queries_success.add(1, &attrs); - self.queries_error.add( - 0, - &[ - KeyValue::new("scenario", scenario.to_string()), - KeyValue::new("status", "error"), - ], - ); + self.queries_success.add(1, attrs); if sample.retry_detected { - self.queries_retried.add(1, &attrs); - } else { - self.queries_retried.add(0, &attrs); + self.queries_retried.add(1, attrs); } - self.rows_read.add(sample.rows_count as u64, &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( - Duration::from_micros(sample.send_duration_micros as u64).as_secs_f64(), - &attrs, - ); - self.poll_duration.record( - Duration::from_micros(sample.poll_duration_micros as u64).as_secs_f64(), - &attrs, - ); - self.read_duration.record( - Duration::from_micros(sample.read_duration_micros as u64).as_secs_f64(), - &attrs, - ); + .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.queries_success.add( - 0, - &[ - KeyValue::new("scenario", scenario.to_string()), - KeyValue::new("status", "ok"), - ], - ); + self.queries_error.add(1, attrs); } - self.query_duration.record( - Duration::from_micros(sample.total_duration_micros as u64).as_secs_f64(), - &attrs, - ); - } -} - -impl Default for OtelMetrics { - fn default() -> Self { - Self::new() + self.query_duration + .record(sample.total_duration().as_secs_f64(), attrs); } } diff --git a/src/bigquery/benchmarks/queries/src/reporter.rs b/src/bigquery/benchmarks/queries/src/reporter.rs index 3748e59d93..703731df0b 100644 --- a/src/bigquery/benchmarks/queries/src/reporter.rs +++ b/src/bigquery/benchmarks/queries/src/reporter.rs @@ -14,14 +14,31 @@ use crate::args::Args; use crate::metrics::{self, LatencySummary}; -use crate::sample::{Sample, SampleStatus}; +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 { @@ -48,6 +65,7 @@ pub struct BenchmarkReport { 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, } @@ -76,56 +94,37 @@ impl BenchmarkReport { println!(" Max: {:?}", total.max); } - if let Some(send) = &self.send_duration { - println!("\n--- Query::send() Latency ---"); - println!( - " P50: {:?} | P90: {:?} | P99: {:?}", - send.p50, send.p90, send.p99 - ); - } - - if let Some(poll) = &self.poll_duration { - println!("\n--- Query::until_done() Polling Latency ---"); - println!( - " P50: {:?} | P90: {:?} | P99: {:?}", - poll.p50, poll.p90, poll.p99 - ); - } - - if let Some(read) = &self.read_duration { - println!("\n--- CompleteQuery::read() Streaming Latency ---"); - println!( - " P50: {:?} | P90: {:?} | P99: {:?}", - read.p50, read.p90, read.p99 - ); + 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.errors.len()); + println!(" QUERY ERRORS ({} total)", self.error_count); println!("-------------------------------------------------------"); - for (idx, err) in self.errors.iter().take(20).enumerate() { - let job_id = if !err.final_job_id.is_empty() && err.final_job_id != "N/A" { - &err.final_job_id - } else if !err.initial_job_id.is_empty() && err.initial_job_id != "N/A" { - &err.initial_job_id - } else { - "N/A" - }; + 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, - job_id + display_job_id(&err.initial_job_id, &err.final_job_id) ); println!(" Error: {}", err.error_message); } - if self.errors.len() > 20 { + let shown = self.errors.len().min(ERRORS_PRINTED); + if self.error_count > shown { println!( " ... and {} more error(s) recorded in full error log.", - self.errors.len() - 20 + self.error_count - shown ); } } @@ -134,224 +133,244 @@ impl BenchmarkReport { } } -/// 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 total_samples = 0_usize; - let mut success_total_durations = Vec::new(); - let mut success_send_durations = Vec::new(); - let mut success_poll_durations = Vec::new(); - let mut success_read_durations = Vec::new(); - let mut errors = Vec::new(); - - let mut success_count = 0_usize; - let mut error_count = 0_usize; - let mut retries_detected_count = 0_usize; - let mut total_rows_read = 0_usize; - let mut total_bytes_processed = 0_i64; - - let mut realtime_files = if let Some(output_dir) = &args.output_dir { +/// 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-{}-{}.csv", scenario.name, timestamp)); - let json_path = output_dir.join(format!("summary-{}-{}.json", scenario.name, timestamp)); - let errors_path = output_dir.join(format!("errors-{}-{}.log", scenario.name, timestamp)); + 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 csv_file = File::create(&csv_path)?; - let mut csv_writer = BufWriter::new(csv_file); - writeln!(csv_writer, "{}", Sample::HEADER)?; - csv_writer.flush()?; + let mut csv = BufWriter::new(File::create(&csv_path)?); + writeln!(csv, "{}", Sample::HEADER)?; + csv.flush()?; - let errors_file = File::create(&errors_path)?; - let mut errors_writer = BufWriter::new(errors_file); + let mut errors = BufWriter::new(File::create(&errors_path)?); writeln!( - errors_writer, - "# BigQuery Benchmark Error Log - Scenario: {}, Timestamp: {}", - scenario.name, timestamp + errors, + "# BigQuery Benchmark Error Log - Scenario: {scenario}, Timestamp: {timestamp}" )?; - errors_writer.flush()?; + 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()); - Some((csv_writer, csv_path, json_path, errors_writer, errors_path)) - } else { - None - }; + 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:?}"); + } + } - let mut last_json_write = Instant::now(); + 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 { - total_samples += 1; - if sample.status == SampleStatus::Ok { - success_count += 1; - success_total_durations.push(sample.total_duration()); - success_send_durations.push(Duration::from_micros(sample.send_duration_micros as u64)); - success_poll_durations.push(Duration::from_micros(sample.poll_duration_micros as u64)); - success_read_durations.push(Duration::from_micros(sample.read_duration_micros as u64)); - total_rows_read += sample.rows_count; - total_bytes_processed += sample.bytes_processed; - } else { - error_count += 1; - let offset_secs = sample.start_offset_micros as f64 / 1_000_000.0; - let err_detail = ErrorDetail { - task_id: sample.task_id, - iteration: sample.iteration, - offset_secs, - initial_job_id: sample.initial_job_id.clone(), - final_job_id: sample.final_job_id.clone(), - error_message: sample.error_message.clone(), - }; - - let job_id = if !sample.final_job_id.is_empty() && sample.final_job_id != "N/A" { - &sample.final_job_id - } else if !sample.initial_job_id.is_empty() && sample.initial_job_id != "N/A" { - &sample.initial_job_id - } else { - "N/A" - }; + 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, offset_secs, job_id, sample.error_message + sample.task_id, + sample.iteration, + sample.start_offset_secs(), + sample.display_job_id(), + sample.error_message ); - - if let Some((_, _, _, errors_writer, _)) = &mut realtime_files { - let _ = writeln!( - errors_writer, - "Task: {}\nIteration: {}\nOffsetSecs: {:.3}\nInitialJobId: {}\nFinalJobId: {}\nError: {}\n--------------------------------------------------------------------------------", - sample.task_id, - sample.iteration, - offset_secs, - sample.initial_job_id, - sample.final_job_id, - sample.error_message - ); - let _ = errors_writer.flush(); - } - - errors.push(err_detail); - } - - if sample.retry_detected { - retries_detected_count += 1; } - if let Some((csv_writer, _, json_path, errors_writer, _)) = &mut realtime_files { - if let Err(err) = writeln!(csv_writer, "{}", sample.to_csv_row()) { - tracing::error!("Failed to write CSV sample row to disk: {err:?}"); + if let Some(files) = &mut files { + files.write_sample(&sample); + if sample.status != SampleStatus::Ok { + files.write_error(&sample); } - // Periodically update summary JSON and flush buffers (every 5 seconds) to avoid high disk I/O and sorting overhead - if last_json_write.elapsed() >= Duration::from_secs(5) { - last_json_write = Instant::now(); - let stats = ReportStats { - scenario_name: &scenario.name, - task_count: args.task_count, - total_samples, - success_count, - error_count, - retries_detected_count, - total_rows_read, - total_bytes_processed, - success_total_durations: &success_total_durations, - success_send_durations: &success_send_durations, - success_poll_durations: &success_poll_durations, - success_read_durations: &success_read_durations, - errors: &errors, - }; - let current_report = build_report(&stats); - let json_path_clone = json_path.clone(); - tokio::task::block_in_place(|| { - let _ = csv_writer.flush(); - let _ = errors_writer.flush(); - if let Ok(json_file) = File::create(&json_path_clone) { - let _ = serde_json::to_writer_pretty(json_file, ¤t_report); - } - }); + 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 stats = ReportStats { - scenario_name: &scenario.name, - task_count: args.task_count, - total_samples, - success_count, - error_count, - retries_detected_count, - total_rows_read, - total_bytes_processed, - success_total_durations: &success_total_durations, - success_send_durations: &success_send_durations, - success_poll_durations: &success_poll_durations, - success_read_durations: &success_read_durations, - errors: &errors, - }; - let report = build_report(&stats); - + let report = acc.build_report(scenario.name, args.task_count); report.print_stdout(); - if let Some((csv_writer, csv_path, json_path, errors_writer, errors_path)) = &mut realtime_files - { - let json_path_clone = json_path.clone(); - tokio::task::block_in_place(|| { - let _ = csv_writer.flush(); - let _ = errors_writer.flush(); - // Save final complete JSON summary - if let Ok(json_file) = File::create(&json_path_clone) { - let _ = serde_json::to_writer_pretty(json_file, &report); - } - }); - println!("Final samples saved to: {}", csv_path.display()); - println!("Final summary saved to: {}", json_path.display()); - if error_count > 0 { - println!("Errors logged to: {}", errors_path.display()); + 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) } -struct ReportStats<'a> { - scenario_name: &'a str, - task_count: usize, - total_samples: usize, - success_count: usize, - error_count: usize, - retries_detected_count: usize, - total_rows_read: usize, - total_bytes_processed: i64, - success_total_durations: &'a [Duration], - success_send_durations: &'a [Duration], - success_poll_durations: &'a [Duration], - success_read_durations: &'a [Duration], - errors: &'a [ErrorDetail], -} +#[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) + }); -fn build_report(stats: &ReportStats) -> BenchmarkReport { - BenchmarkReport { - scenario: stats.scenario_name.to_string(), - task_count: stats.task_count, - total_samples: stats.total_samples, - success_count: stats.success_count, - error_count: stats.error_count, - retries_detected_count: stats.retries_detected_count, - total_rows_read: stats.total_rows_read, - total_bytes_processed: stats.total_bytes_processed, - total_duration: metrics::compute_metrics(stats.success_total_durations), - send_duration: metrics::compute_metrics(stats.success_send_durations), - poll_duration: metrics::compute_metrics(stats.success_poll_durations), - read_duration: metrics::compute_metrics(stats.success_read_durations), - errors: stats.errors.to_vec(), + 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 index 14a0a3b3a2..e37da0c1f9 100644 --- a/src/bigquery/benchmarks/queries/src/runner.rs +++ b/src/bigquery/benchmarks/queries/src/runner.rs @@ -13,8 +13,8 @@ // limitations under the License. use crate::args::Args; -use crate::metrics::{self, OtelMetrics}; -use crate::sample::{Sample, SampleStatus}; +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}; @@ -55,7 +55,7 @@ impl TaskRunner<'_> { } let iter_start = Instant::now(); - let start_offset_micros = self.test_start.elapsed().as_micros(); + let start_offset_micros = as_micros_u64(self.test_start.elapsed()); let iteration_span = tracing::info_span!( "bigquery.query_benchmark.iteration", @@ -71,40 +71,27 @@ impl TaskRunner<'_> { let sample = match tokio::time::timeout(self.args.query_timeout, sample_fut).await { Ok(sample) => sample, Err(_) => { - let total_duration = iter_start.elapsed(); - metrics::inc_total_queries(); - metrics::inc_error_queries(); tracing::error!( task_id = self.task_id, iteration, "Query iteration timed out after {:?}", self.args.query_timeout ); - let sample = Sample { - task_id: self.task_id, - iteration, - start_offset_micros, - send_duration_micros: 0, - poll_duration_micros: 0, - read_duration_micros: 0, - total_duration_micros: total_duration.as_micros(), - rows_count: 0, - bytes_processed: 0, - cache_hit: false, - initial_job_id: String::new(), - final_job_id: String::new(), - retry_detected: false, + 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 ), - }; - self.metrics.record_sample(&self.scenario.name, &sample); - sample + ..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; } @@ -112,12 +99,18 @@ impl TaskRunner<'_> { 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: u128, + 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) @@ -135,40 +128,24 @@ impl TaskRunner<'_> { 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; - let send_duration = send_start.elapsed(); + sample.send_duration_micros = as_micros_u64(send_start.elapsed()); let query_handle = match send_result { Ok(handle) => handle, Err(err) => { - let total_duration = iter_start.elapsed(); - metrics::inc_total_queries(); - metrics::inc_error_queries(); - tracing::error!(self.task_id, iteration, "Query::send failed: {err:?}"); - - let sample = Sample { - task_id: self.task_id, + tracing::error!( + task_id = self.task_id, iteration, - start_offset_micros, - send_duration_micros: send_duration.as_micros(), - poll_duration_micros: 0, - read_duration_micros: 0, - total_duration_micros: total_duration.as_micros(), - 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: format!("Query::send: {err:#}"), - }; - self.metrics.record_sample(&self.scenario.name, &sample); + "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 - let initial_job_id = query_handle + sample.initial_job_id = query_handle .metadata() .job_reference .as_ref() @@ -181,76 +158,52 @@ impl TaskRunner<'_> { "bigquery.until_done", task_id = self.task_id, iteration, - %initial_job_id + initial_job_id = %sample.initial_job_id ); let done_result = query_handle.until_done().instrument(poll_span).await; - let poll_duration = poll_start.elapsed(); + sample.poll_duration_micros = as_micros_u64(poll_start.elapsed()); let complete_query = match done_result { Ok(complete) => complete, Err(err) => { - let total_duration = iter_start.elapsed(); - metrics::inc_total_queries(); - metrics::inc_error_queries(); tracing::error!( - self.task_id, + task_id = self.task_id, iteration, - %initial_job_id, + initial_job_id = %sample.initial_job_id, "Query::until_done failed: {err:?}" ); - - let sample = Sample { - task_id: self.task_id, - iteration, - start_offset_micros, - send_duration_micros: send_duration.as_micros(), - poll_duration_micros: poll_duration.as_micros(), - read_duration_micros: 0, - total_duration_micros: total_duration.as_micros(), - rows_count: 0, - bytes_processed: 0, - cache_hit: false, - initial_job_id, - final_job_id: String::new(), - retry_detected: false, - status: SampleStatus::Error, - error_message: format!("Query::until_done: {err:#}"), - }; - self.metrics.record_sample(&self.scenario.name, &sample); + 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 final_job_id = complete_query - .metadata() + 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 - let retry_detected = !initial_job_id.is_empty() - && !final_job_id.is_empty() - && initial_job_id != final_job_id; + sample.retry_detected = !sample.initial_job_id.is_empty() + && !sample.final_job_id.is_empty() + && sample.initial_job_id != sample.final_job_id; - if retry_detected { - metrics::inc_retried_queries(); + if sample.retry_detected { tracing::warn!( task_id = self.task_id, iteration, - %initial_job_id, - %final_job_id, + initial_job_id = %sample.initial_job_id, + final_job_id = %sample.final_job_id, "Query job retry detected under the hood (job_id mutated)!" ); } - let bytes_processed = complete_query.metadata().total_bytes_processed.unwrap_or(0); - - let cache_hit = complete_query.metadata().cache_hit.unwrap_or(false); - // Step 4: Stream and read result rows if enabled - let mut rows_count = 0_usize; let read_start = Instant::now(); let mut read_error = None; @@ -262,11 +215,11 @@ impl TaskRunner<'_> { while let Some(row_result) = rows.next().await { match row_result { Ok(_) => { - rows_count += 1; + sample.rows_count += 1; } Err(err) => { tracing::error!( - self.task_id, + task_id = self.task_id, iteration, "Error streaming rows: {err:?}" ); @@ -279,58 +232,15 @@ impl TaskRunner<'_> { .instrument(read_span) .await; } - let read_duration = read_start.elapsed(); - let total_duration = iter_start.elapsed(); - - metrics::inc_total_queries(); - - let sample = if let Some(err_msg) = read_error { - metrics::inc_error_queries(); - Sample { - task_id: self.task_id, - iteration, - start_offset_micros, - send_duration_micros: send_duration.as_micros(), - poll_duration_micros: poll_duration.as_micros(), - read_duration_micros: read_duration.as_micros(), - total_duration_micros: total_duration.as_micros(), - rows_count, - bytes_processed, - cache_hit, - initial_job_id, - final_job_id, - retry_detected, - status: SampleStatus::Error, - error_message: err_msg, - } - } else { - metrics::inc_success_queries(); - metrics::add_rows_read(rows_count as u64); - if bytes_processed > 0 { - metrics::add_bytes_processed(bytes_processed as u64); - } + sample.read_duration_micros = as_micros_u64(read_start.elapsed()); + sample.total_duration_micros = as_micros_u64(iter_start.elapsed()); - Sample { - task_id: self.task_id, - iteration, - start_offset_micros, - send_duration_micros: send_duration.as_micros(), - poll_duration_micros: poll_duration.as_micros(), - read_duration_micros: read_duration.as_micros(), - total_duration_micros: total_duration.as_micros(), - rows_count, - bytes_processed, - cache_hit, - initial_job_id, - final_job_id, - retry_detected, - status: SampleStatus::Ok, - error_message: String::new(), - } - }; + match read_error { + Some(err_msg) => sample.error_message = err_msg, + None => sample.status = SampleStatus::Ok, + } - self.metrics.record_sample(&self.scenario.name, &sample); sample } } diff --git a/src/bigquery/benchmarks/queries/src/sample.rs b/src/bigquery/benchmarks/queries/src/sample.rs index 031f52115b..908809b6df 100644 --- a/src/bigquery/benchmarks/queries/src/sample.rs +++ b/src/bigquery/benchmarks/queries/src/sample.rs @@ -38,11 +38,11 @@ impl SampleStatus { pub struct Sample { pub task_id: usize, pub iteration: u64, - pub start_offset_micros: u128, - pub send_duration_micros: u128, - pub poll_duration_micros: u128, - pub read_duration_micros: u128, - pub total_duration_micros: u128, + 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, @@ -60,8 +60,53 @@ impl Sample { "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, @@ -74,16 +119,8 @@ impl Sample { self.rows_count, self.bytes_processed, self.cache_hit, - if self.initial_job_id.is_empty() { - "N/A" - } else { - &self.initial_job_id - }, - if self.final_job_id.is_empty() { - "N/A" - } else { - &self.final_job_id - }, + initial_job_id, + final_job_id, self.retry_detected, self.status.as_str(), clean_err, @@ -91,10 +128,38 @@ impl Sample { } pub fn total_duration(&self) -> Duration { - Duration::from_micros(self.total_duration_micros as u64) + 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::*; @@ -102,24 +167,26 @@ mod tests { #[test] fn test_sample_csv_serialization() { let sample = Sample { - task_id: 1, - iteration: 42, - start_offset_micros: 100_000, 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, - cache_hit: false, initial_job_id: "job_init_123".to_string(), final_job_id: "job_retry_456".to_string(), retry_detected: true, status: SampleStatus::Ok, - error_message: String::new(), + ..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 index 7d84905d1d..45cc1e9ef9 100644 --- a/src/bigquery/benchmarks/queries/src/scenarios.rs +++ b/src/bigquery/benchmarks/queries/src/scenarios.rs @@ -18,7 +18,7 @@ use std::fs; /// Represents a configured query benchmark scenario. #[derive(Clone, Debug)] pub struct Scenario { - pub name: String, + pub name: &'static str, pub sql: String, pub description: &'static str, } @@ -26,56 +26,42 @@ pub struct Scenario { impl Scenario { /// Resolves the query scenario based on the provided CLI arguments. pub fn resolve(args: &Args) -> anyhow::Result { - match args.scenario { - ScenarioName::Synthetic100k => Ok(Self { - name: "synthetic-100k".to_string(), - sql: concat!( - "SELECT ", - " x AS row_id, ", - " GENERATE_UUID() AS uuid, ", - " REPEAT('abcdefghij', 10) AS payload ", - "FROM UNNEST(GENERATE_ARRAY(1, 100000)) AS x" - ) - .to_string(), - description: "Generates 100,000 structured rows in-flight with no external table dependency.", - }), - ScenarioName::Synthetic10k => Ok(Self { - name: "synthetic-10k".to_string(), - sql: concat!( - "SELECT ", - " x AS row_id, ", - " GENERATE_UUID() AS uuid, ", - " REPEAT('abcdefghij', 10) AS payload ", - "FROM UNNEST(GENERATE_ARRAY(1, 10000)) AS x" - ) - .to_string(), - description: "Generates 10,000 structured rows in-flight with no external table dependency.", - }), - ScenarioName::UsaNamesScan => Ok(Self { - name: "usa-names-scan".to_string(), - sql: concat!( + 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(), - description: "Scans and retrieves 50,000 rows from the USA names public dataset.", - }), - ScenarioName::UsaNamesAgg => Ok(Self { - name: "usa-names-agg".to_string(), - sql: concat!( + "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(), - description: "Aggregates 5.5M rows grouped by state and gender.", - }), - ScenarioName::WikipediaAgg => Ok(Self { - name: "wikipedia-agg".to_string(), - sql: concat!( + "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 ", @@ -84,62 +70,50 @@ impl Scenario { "LIMIT 1000" ) .to_string(), - description: "Aggregates top 1000 article views from Wikipedia public samples.", - }), - ScenarioName::Custom => { - let sql = if let Some(sql) = &args.sql { - sql.clone() - } else if let Some(sql_file) = &args.sql_file { - fs::read_to_string(sql_file).map_err(|e| { - anyhow::anyhow!( - "Failed to read custom SQL file {}: {}", - sql_file.display(), - e - ) - })? - } else { - anyhow::bail!("Custom scenario requires --sql or --sql-file"); - }; + "Aggregates top 1000 article views from Wikipedia public samples.", + ), + ScenarioName::Custom => ( + "custom", + custom_sql(args)?, + "User-defined custom SQL query.", + ), + }; - Ok(Self { - name: "custom".to_string(), - sql, - description: "User-defined custom SQL query.", - }) - } - } + Ok(Self { + name, + sql, + description, + }) } } -#[cfg(test)] -mod tests { - use super::*; - use clap::Parser; - - #[test] - fn test_synthetic_scenarios() { - let args = Args::parse_from(["bigquery-benchmark-queries", "--scenario", "synthetic-100k"]); - let s = Scenario::resolve(&args).unwrap(); - assert_eq!(s.name, "synthetic-100k"); - assert!(s.sql.contains("100000")); - - let args = Args::parse_from(["bigquery-benchmark-queries", "--scenario", "synthetic-10k"]); - let s = Scenario::resolve(&args).unwrap(); - assert_eq!(s.name, "synthetic-10k"); - assert!(s.sql.contains("10000")); - } +/// 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 + ) +} - #[test] - fn test_custom_scenario_from_string() { - let args = Args::parse_from([ - "bigquery-benchmark-queries", - "--scenario", - "custom", - "--sql", - "SELECT 42", - ]); - let s = Scenario::resolve(&args).unwrap(); - assert_eq!(s.name, "custom"); - assert_eq!(s.sql, "SELECT 42"); +/// 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 index 07f6b2dfb3..0e43923852 100644 --- a/src/bigquery/benchmarks/queries/src/telemetry.rs +++ b/src/bigquery/benchmarks/queries/src/telemetry.rs @@ -27,21 +27,21 @@ 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, - namespace: String, } impl GenericNodeDetector { pub fn new() -> Self { - let id = Uuid::new_v4().to_string(); Self { - id, - location: "us-central1".to_string(), - namespace: "bigquery-benchmark-queries".to_string(), + id: Uuid::new_v4().to_string(), + // Cloud Monitoring resolves `generic_node.location` to a GCP region + // or zone. + location: DEFAULT_MONITORING_REGION.to_string(), } } } @@ -51,7 +51,7 @@ impl ResourceDetector for GenericNodeDetector { Resource::builder_empty() .with_attributes([ KeyValue::new("location", self.location.clone()), - KeyValue::new("namespace", self.namespace.clone()), + KeyValue::new("namespace", SERVICE_NAME), KeyValue::new("node_id", self.id.clone()), ]) .build() @@ -59,25 +59,25 @@ impl ResourceDetector for GenericNodeDetector { } /// 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 { - tracer_provider: Option, - meter_provider: Option, + providers: Option<(SdkTracerProvider, SdkMeterProvider)>, } impl TelemetryGuard { /// Flushes and shuts down telemetry providers. pub fn shutdown(self) { - if let Some(tp) = self.tracer_provider { - let _ = tp.force_flush(); - if let Err(e) = tp.shutdown() { - eprintln!("Error shutting down trace provider: {e:?}"); - } + 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 Some(mp) = self.meter_provider { - let _ = mp.force_flush(); - if let Err(e) = mp.shutdown() { - eprintln!("Error shutting down meter provider: {e:?}"); - } + if let Err(e) = meter_provider.shutdown() { + eprintln!("Error shutting down meter provider: {e:?}"); } } } @@ -169,15 +169,11 @@ pub async fn enable_telemetry( .expect("Setting global subscriber succeeds"); return Ok(TelemetryGuard { - tracer_provider: Some(tracer_provider), - meter_provider: Some(meter_provider), + providers: Some((tracer_provider, meter_provider)), }); } tracing::subscriber::set_global_default(registry).expect("Setting global subscriber succeeds"); - Ok(TelemetryGuard { - tracer_provider: None, - meter_provider: None, - }) + Ok(TelemetryGuard { providers: None }) } From 95a3234584c6e0f20f9d7cf418b6f7afb1cae629 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Wed, 16 Sep 2026 19:27:37 +0000 Subject: [PATCH 10/11] fix: fmt all the things --- src/bigquery/benchmarks/queries/README.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/bigquery/benchmarks/queries/README.md b/src/bigquery/benchmarks/queries/README.md index 3f2861fd25..d7236896a9 100644 --- a/src/bigquery/benchmarks/queries/README.md +++ b/src/bigquery/benchmarks/queries/README.md @@ -46,10 +46,11 @@ ______________________________________________________________________ ## 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. +> [!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 From 6464ea6d76bd74e6e166d4bf344ec2e1199e4935 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Wed, 16 Sep 2026 19:54:32 +0000 Subject: [PATCH 11/11] fix: max results to page size --- src/bigquery/benchmarks/queries/src/args.rs | 4 ++-- src/bigquery/benchmarks/queries/src/runner.rs | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/bigquery/benchmarks/queries/src/args.rs b/src/bigquery/benchmarks/queries/src/args.rs index 83163e40f1..c85d061bb6 100644 --- a/src/bigquery/benchmarks/queries/src/args.rs +++ b/src/bigquery/benchmarks/queries/src/args.rs @@ -90,9 +90,9 @@ pub struct Args { #[arg(long, value_parser = parse_duration)] pub duration: Option, - /// The maximum number of rows per page returned from BigQuery (maps to `max_results`). + /// The maximum number of rows per page returned from BigQuery. #[arg(long)] - pub max_results: Option, + 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)] diff --git a/src/bigquery/benchmarks/queries/src/runner.rs b/src/bigquery/benchmarks/queries/src/runner.rs index e37da0c1f9..5eec9def61 100644 --- a/src/bigquery/benchmarks/queries/src/runner.rs +++ b/src/bigquery/benchmarks/queries/src/runner.rs @@ -117,8 +117,8 @@ impl TaskRunner<'_> { .set_location(&self.args.location) .set_use_query_cache(self.args.use_query_cache); - if let Some(max_results) = self.args.max_results { - query_builder = query_builder.set_max_results(max_results); + 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);