From 6c8dc15097c69b9c78a8483639a1023923e43951 Mon Sep 17 00:00:00 2001 From: fargito Date: Sat, 12 Sep 2026 03:35:39 +0200 Subject: [PATCH 1/3] feat: add postgres instrument to collect the pg_stat_statements dump Add a `postgres` instrument enabled via `--instruments postgres`. Unlike the MongoDB instrument it owns no process and does no proxying: the `codspeed/postgres` image runs a poller that writes a pg_stat_statements analytics dump to a file. This instrument copies that dump (path given by the required `--postgres-dump-path`) into `/instruments/postgres.json` after the run, so it rides along in the uploaded profile archive. collect() waits for the poller to flush a dump newer than the run end before copying, so the final queries aren't dropped by the poller's tick interval. --- src/cli/exec/mod.rs | 5 +- src/cli/run/mod.rs | 10 +- src/executor/mod.rs | 13 +++ src/instruments/mod.rs | 66 +++++++++++++- src/instruments/postgres.rs | 176 ++++++++++++++++++++++++++++++++++++ 5 files changed, 266 insertions(+), 4 deletions(-) create mode 100644 src/instruments/postgres.rs diff --git a/src/cli/exec/mod.rs b/src/cli/exec/mod.rs index 4a757f74b..171c9555e 100644 --- a/src/cli/exec/mod.rs +++ b/src/cli/exec/mod.rs @@ -75,7 +75,10 @@ fn build_orchestrator_config( working_directory: args.shared.working_directory, targets: vec![target], modes, - instruments: Instruments { mongodb: None }, // exec doesn't support MongoDB + instruments: Instruments { + mongodb: None, + postgres: None, + }, // exec doesn't support instruments perf_unwinding_mode: args.shared.profiler_run_args.perf.perf_unwinding_mode, enable_profiler: args.shared.profiler_run_args.resolve_enable_profiler(), walltime_profiler: args.shared.walltime_profiler, diff --git a/src/cli/run/mod.rs b/src/cli/run/mod.rs index a218155c7..312f4dfc6 100644 --- a/src/cli/run/mod.rs +++ b/src/cli/run/mod.rs @@ -9,7 +9,7 @@ use crate::run_environment::interfaces::RepositoryProvider; use crate::upload::poll_results::PollResultsOptions; use clap::{Args, ValueEnum}; use std::collections::HashMap; -use std::path::Path; +use std::path::{Path, PathBuf}; use url::Url; pub mod helpers; @@ -20,7 +20,7 @@ pub struct RunArgs { #[command(flatten)] pub shared: ExecAndRunSharedArgs, - /// Comma-separated list of instruments to enable. Possible values: mongodb. + /// Comma-separated list of instruments to enable. Possible values: mongodb, postgres. #[arg(long, value_delimiter = ',')] pub instruments: Vec, @@ -31,6 +31,11 @@ pub struct RunArgs { #[arg(long)] pub mongo_uri_env_name: Option, + /// Host path to the Postgres analytics dump written by the codspeed/postgres + /// image's poller. Required when the `postgres` instrument is enabled. + #[arg(long)] + pub postgres_dump_path: Option, + #[arg(long, hide = true)] pub message_format: Option, @@ -87,6 +92,7 @@ impl RunArgs { }, instruments: vec![], mongo_uri_env_name: None, + postgres_dump_path: None, message_format: None, command: vec![], } diff --git a/src/executor/mod.rs b/src/executor/mod.rs index 80d31ea87..6f9c1f85f 100644 --- a/src/executor/mod.rs +++ b/src/executor/mod.rs @@ -14,6 +14,7 @@ mod valgrind; mod wall_time; use crate::instruments::mongo_tracer::{MongoTracer, install_mongodb_tracer}; +use crate::instruments::postgres::PostgresInstrument; use crate::local_logger::rolling_buffer::{activate_rolling_buffer, deactivate_rolling_buffer}; use crate::prelude::*; use crate::runner_mode::RunnerMode; @@ -212,6 +213,18 @@ pub async fn run_executor( if let Some(mut mongo_tracer) = mongo_tracer { mongo_tracer.stop().await?; } + + // A benchmark run must not fail over analytics enrichment, so a collect + // error is logged and swallowed rather than propagated. + if let Some(postgres_config) = &execution_context.config.instruments.postgres { + if let Err(e) = + PostgresInstrument::new(&execution_context.profile_folder, postgres_config) + .collect() + .await + { + warn!("Failed to collect Postgres analytics: {e:#}"); + } + } debug!("Tearing down the executor"); executor.teardown(execution_context).await?; diff --git a/src/instruments/mod.rs b/src/instruments/mod.rs index c16a00de2..60f9ab1e8 100644 --- a/src/instruments/mod.rs +++ b/src/instruments/mod.rs @@ -7,20 +7,29 @@ use crate::cli::run::RunArgs; use crate::prelude::*; pub mod mongo_tracer; +pub mod postgres; #[derive(Debug, Clone, PartialEq, Eq)] pub struct MongoDBConfig { pub uri_env_name: Option, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PostgresConfig { + /// Host path to the analytics dump written by the codspeed/postgres poller. + pub dump_path: std::path::PathBuf, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct Instruments { pub mongodb: Option, + pub postgres: Option, } #[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq, Hash)] pub enum InstrumentName { MongoDB, + Postgres, } impl Instruments { @@ -28,6 +37,10 @@ impl Instruments { self.mongodb.is_some() } + pub fn is_postgres_enabled(&self) -> bool { + self.postgres.is_some() + } + pub fn get_active_instrument_names(&self) -> Vec { let mut names = vec![]; @@ -35,6 +48,10 @@ impl Instruments { names.push(InstrumentName::MongoDB); } + if self.is_postgres_enabled() { + names.push(InstrumentName::Postgres); + } + names } } @@ -47,6 +64,7 @@ impl TryFrom<&RunArgs> for Instruments { for instrument_name in &args.instruments { match instrument_name.as_str() { "mongodb" => validated_instrument_names.insert(InstrumentName::MongoDB), + "postgres" => validated_instrument_names.insert(InstrumentName::Postgres), _ => bail!("Invalid instrument name: {instrument_name}"), }; } @@ -64,7 +82,23 @@ impl TryFrom<&RunArgs> for Instruments { None }; - Ok(Self { mongodb }) + let postgres = if validated_instrument_names.contains(&InstrumentName::Postgres) { + let dump_path = args.postgres_dump_path.clone().ok_or_else(|| { + anyhow!( + "The Postgres instrument is enabled but --postgres-dump-path was not provided" + ) + })?; + Some(PostgresConfig { dump_path }) + } else if args.postgres_dump_path.is_some() { + warn!( + "The Postgres instrument is disabled but a Postgres dump path was provided, ignoring it" + ); + None + } else { + None + }; + + Ok(Self { mongodb, postgres }) } } @@ -76,6 +110,7 @@ impl Instruments { mongodb: Some(MongoDBConfig { uri_env_name: Some("MONGODB_URI".into()), }), + postgres: None, } } } @@ -133,4 +168,33 @@ mod tests { "Invalid instrument name: unknown" ); } + + #[test] + fn test_from_args_postgres() { + let args = RunArgs { + instruments: vec!["postgres".into()], + postgres_dump_path: Some("/tmp/x/dump.json".into()), + ..RunArgs::test() + }; + let instruments = Instruments::try_from(&args).unwrap(); + assert!(instruments.is_postgres_enabled()); + assert_eq!( + instruments.postgres.unwrap().dump_path, + std::path::PathBuf::from("/tmp/x/dump.json") + ); + } + + #[test] + fn test_from_args_postgres_without_dump_path() { + let args = RunArgs { + instruments: vec!["postgres".into()], + ..RunArgs::test() + }; + let instruments = Instruments::try_from(&args); + assert!(instruments.is_err()); + assert_eq!( + instruments.unwrap_err().to_string(), + "The Postgres instrument is enabled but --postgres-dump-path was not provided" + ); + } } diff --git a/src/instruments/postgres.rs b/src/instruments/postgres.rs new file mode 100644 index 000000000..71c5e6c50 --- /dev/null +++ b/src/instruments/postgres.rs @@ -0,0 +1,176 @@ +use std::path::{Path, PathBuf}; +use std::time::{Duration, SystemTime}; + +use tokio::fs; + +use crate::prelude::*; + +use super::PostgresConfig; + +/// Collects the analytics dump produced by the `codspeed/postgres` image's +/// poller and drops it into the profile folder so it is uploaded with the run. +/// +/// Unlike the MongoDB instrument, this owns no process and does no proxying: the +/// poller runs inside the database image and writes the dump to a file, so this +/// only copies that file into the profile folder. +#[derive(Debug)] +pub struct PostgresInstrument { + profile_folder: PathBuf, + dump_path: PathBuf, + max_flush_wait: Duration, + flush_poll_interval: Duration, +} + +/// Bound on how long to wait for the poller to flush a fresh dump after the +/// benchmark ends, and how often to re-check. The poller flushes on a fixed +/// interval (2s by default), so the bound is a few of those. +const MAX_FLUSH_WAIT: Duration = Duration::from_secs(6); +const FLUSH_POLL_INTERVAL: Duration = Duration::from_millis(200); + +impl PostgresInstrument { + pub fn new(profile_folder: &Path, config: &PostgresConfig) -> Self { + Self { + profile_folder: profile_folder.to_path_buf(), + dump_path: config.dump_path.clone(), + max_flush_wait: MAX_FLUSH_WAIT, + flush_poll_interval: FLUSH_POLL_INTERVAL, + } + } + + /// Copy the dump into `/instruments/postgres.json`, so it + /// rides along in the uploaded profile archive. A missing dump is a warning, + /// not a failure — a benchmark run must not fail over missing analytics. + pub async fn collect(&self) -> Result<()> { + // The benchmark has just finished. The poller flushes the dump on a fixed + // interval, so copying it immediately would drop up to one interval of the + // final queries. Wait for a flush that happened after now (the run end). + self.wait_for_fresh_dump(SystemTime::now()).await; + + if !self.dump_path.exists() { + warn!( + "Postgres instrument enabled but no dump found at {}; skipping", + self.dump_path.display() + ); + return Ok(()); + } + + let instruments_out_dir = self.profile_folder.join("instruments"); + fs::create_dir_all(&instruments_out_dir).await?; + let dest = instruments_out_dir.join("postgres.json"); + fs::copy(&self.dump_path, &dest).await.with_context(|| { + format!( + "Failed to copy Postgres dump from {}", + self.dump_path.display() + ) + })?; + debug!("Collected Postgres analytics into {}", dest.display()); + + Ok(()) + } + + /// Poll the dump's mtime until it advances past `after`, so the copied dump + /// reflects a poller flush that happened after the last benchmark query. The + /// poller writes atomically (temp file + rename), so a bumped mtime means a + /// complete document. Bounded: the poller may already be idle or gone. + async fn wait_for_fresh_dump(&self, after: SystemTime) { + let deadline = tokio::time::Instant::now() + self.max_flush_wait; + while tokio::time::Instant::now() < deadline { + if let Ok(meta) = fs::metadata(&self.dump_path).await { + if let Ok(mtime) = meta.modified() { + if mtime >= after { + return; + } + } + } + tokio::time::sleep(self.flush_poll_interval).await; + } + warn!( + "Postgres dump at {} did not refresh within {:?}; using the latest available snapshot", + self.dump_path.display(), + self.max_flush_wait + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::process::id; + + impl PostgresInstrument { + fn with_timing(mut self, max_flush_wait: Duration, flush_poll_interval: Duration) -> Self { + self.max_flush_wait = max_flush_wait; + self.flush_poll_interval = flush_poll_interval; + self + } + } + + fn scratch(tag: &str) -> PathBuf { + std::env::temp_dir().join(format!("pg-instr-{}-{tag}", id())) + } + + fn instrument(profile: &Path, dump_path: PathBuf, max: Duration) -> PostgresInstrument { + PostgresInstrument::new(profile, &PostgresConfig { dump_path }) + .with_timing(max, Duration::from_millis(20)) + } + + /// collect() must not copy the dump the poller last flushed before the run + /// ended: it waits for a flush newer than the run end, so the final queries + /// are captured instead of dropped by the poller's tick interval. + #[tokio::test] + async fn collect_waits_for_post_run_flush() { + let base = scratch("fresh"); + let _ = tokio::fs::remove_dir_all(&base).await; + tokio::fs::create_dir_all(&base).await.unwrap(); + let dump_path = base.join("dump.json"); + let profile = base.join("pf"); + tokio::fs::write(&dump_path, br#"{"queries":[]}"#) + .await + .unwrap(); + + let instr = instrument(&profile, dump_path.clone(), Duration::from_secs(6)); + + let dp = dump_path.clone(); + let flush = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(300)).await; + tokio::fs::write(&dp, br#"{"queries":[{"sql":"select 1"}]}"#) + .await + .unwrap(); + }); + + instr.collect().await.unwrap(); + flush.await.unwrap(); + + let copied = tokio::fs::read_to_string(profile.join("instruments/postgres.json")) + .await + .unwrap(); + assert!( + copied.contains("select 1"), + "collect copied a stale pre-flush dump: {copied}" + ); + let _ = tokio::fs::remove_dir_all(&base).await; + } + + /// A missing dump is not a failure — a benchmark run must not fail over + /// missing analytics — and nothing is written. + #[tokio::test] + async fn collect_missing_dump_is_ok_and_writes_nothing() { + let base = scratch("missing"); + let _ = tokio::fs::remove_dir_all(&base).await; + tokio::fs::create_dir_all(&base).await.unwrap(); + let profile = base.join("pf"); + + let instr = instrument( + &profile, + base.join("does-not-exist.json"), + Duration::from_millis(100), + ); + + instr.collect().await.unwrap(); + assert!( + !profile.join("instruments/postgres.json").exists(), + "no dump should be written when the source is missing" + ); + let _ = tokio::fs::remove_dir_all(&base).await; + } +} From 1c30274281a6a89d8daede7a1113a71617ba8c8e Mon Sep 17 00:00:00 2001 From: fargito Date: Sat, 12 Sep 2026 15:56:19 +0200 Subject: [PATCH 2/3] feat: capture postgres analytics per benchmark via the runner Replace the whole-run poller-copy instrument with per-benchmark capture driven by the runner. On its own connection (--postgres-dsn), reset pg_stat_statements at each benchmark's StartProfiler and snapshot + EXPLAIN at StopProfiler, in the pre-Ack window so the SQL stays outside the measured region. Key each snapshot to the benchmark URI (the same one its flamegraph uses) and write a per-benchmark artifact instruments/postgres.json = {benchmarks:[{uri, queries}]}. Drops --postgres-dump-path and the whole-run collect; the image now only needs pg_stat_statements preloaded, and the read+EXPLAIN logic lives in the runner. --- Cargo.lock | 275 ++++++++++++++++++- Cargo.toml | 1 + src/cli/run/mod.rs | 11 +- src/executor/mod.rs | 13 - src/executor/wall_time/executor.rs | 34 +++ src/instruments/mod.rs | 27 +- src/instruments/postgres.rs | 407 +++++++++++++++++++---------- 7 files changed, 590 insertions(+), 178 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9cf2031c7..bb585bdf2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9,7 +9,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "59317f77929f0e679d39364702289274de2f0f0b22cbf50b2b8cff2169a0b27a" dependencies = [ "cpp_demangle", - "fallible-iterator", + "fallible-iterator 0.3.0", "gimli", "memmap2", "object", @@ -377,6 +377,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + [[package]] name = "block2" version = "0.6.2" @@ -629,6 +638,12 @@ dependencies = [ "cc", ] +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + [[package]] name = "codspeed" version = "4.7.0" @@ -759,6 +774,7 @@ dependencies = [ "test-log", "test-with", "tokio", + "tokio-postgres", "tokio-util", "url", "uuid", @@ -828,6 +844,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "core-foundation" version = "0.9.4" @@ -949,6 +971,15 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + [[package]] name = "ctrlc" version = "3.5.2" @@ -960,6 +991,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "ctutils" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +dependencies = [ + "cmov", +] + [[package]] name = "debugid" version = "0.8.0" @@ -990,8 +1030,20 @@ version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "const-oid", + "crypto-common 0.2.2", + "ctutils", ] [[package]] @@ -1192,6 +1244,12 @@ dependencies = [ "tempfile", ] +[[package]] +name = "fallible-iterator" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" + [[package]] name = "fallible-iterator" version = "0.3.0" @@ -1279,7 +1337,7 @@ source = "git+https://github.com/CodSpeedHQ/framehop?rev=5852f775c2ffa55a24da469 dependencies = [ "arrayvec", "cfg-if", - "fallible-iterator", + "fallible-iterator 0.3.0", "gimli", "macho-unwind-info", "pe-unwind-info", @@ -1427,7 +1485,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "wasm-bindgen", ] @@ -1576,6 +1634,15 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hmac" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" +dependencies = [ + "digest 0.11.3", +] + [[package]] name = "home" version = "0.5.12" @@ -1648,6 +1715,15 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" +[[package]] +name = "hybrid-array" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27f864f10dfb56725ce5ce5472bc52252c8f93a4ab86327122cebf62c5f59a17" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.9.0" @@ -2163,9 +2239,9 @@ dependencies = [ [[package]] name = "libredox" -version = "0.1.16" +version = "0.1.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e02f3bb43d335493c96bf3fd3a321600bf6bd07ed34bc64118e9293bdffea46c" +checksum = "6480ccc157a1389bb2e4891b24751b0f798ba640d22386f23143fbcc89da195a" dependencies = [ "libc", ] @@ -2353,6 +2429,16 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4facc753ae494aeb6e3c22f839b158aebd4f9270f55cd3c79906c45476c47ab4" +[[package]] +name = "md-5" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" +dependencies = [ + "cfg-if", + "digest 0.11.3", +] + [[package]] name = "md5" version = "0.8.0" @@ -2466,7 +2552,7 @@ checksum = "50b7e5b27aa02a74bac8c3f23f448f8d87ff11f92d3aac1a6ed369ee08cc56c1" dependencies = [ "libc", "log", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] @@ -2645,6 +2731,15 @@ dependencies = [ "objc2-core-foundation", ] +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "objc2-core-foundation", +] + [[package]] name = "object" version = "0.39.1" @@ -2788,7 +2883,7 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "408d6fa13d943ee4b76ffda52cc28e817df9c2c4b2c46bd9aec8bff574377e1a" dependencies = [ - "fallible-iterator", + "fallible-iterator 0.3.0", "scroll", "uuid", ] @@ -2852,7 +2947,26 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "89815c69d36021a140146f26659a81d6c2afa33d216d736dd4be5381a7362220" dependencies = [ "pest", - "sha2", + "sha2 0.10.9", +] + +[[package]] +name = "phf" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf" +dependencies = [ + "phf_shared", + "serde", +] + +[[package]] +name = "phf_shared" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266" +dependencies = [ + "siphasher", ] [[package]] @@ -2910,6 +3024,35 @@ dependencies = [ "portable-atomic", ] +[[package]] +name = "postgres-protocol" +version = "0.6.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08808e3c483c46e999108051c78334f473d5adb59d78bb80a1268c7e6aa6c514" +dependencies = [ + "base64", + "byteorder", + "bytes", + "fallible-iterator 0.2.0", + "hmac", + "md-5", + "memchr", + "rand 0.10.1", + "sha2 0.11.0", + "stringprep", +] + +[[package]] +name = "postgres-types" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be" +dependencies = [ + "bytes", + "fallible-iterator 0.2.0", + "postgres-protocol", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -3626,7 +3769,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a22715a5d6deef63c637207afbe68d0c72c3f8d0022d7cf9714c442d6157606b" dependencies = [ "bitflags", - "fallible-iterator", + "fallible-iterator 0.3.0", "fallible-streaming-iterator", "hashlink", "libsqlite3-sys", @@ -4139,7 +4282,18 @@ checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", "cpufeatures 0.2.17", - "digest", + "digest 0.10.7", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", ] [[package]] @@ -4151,7 +4305,7 @@ dependencies = [ "async-trait", "bytes", "hex", - "sha2", + "sha2 0.10.9", "tokio", ] @@ -4233,6 +4387,12 @@ dependencies = [ "time", ] +[[package]] +name = "siphasher" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" + [[package]] name = "slab" version = "0.4.12" @@ -4287,6 +4447,17 @@ dependencies = [ "num-traits", ] +[[package]] +name = "stringprep" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b4df3d392d81bd458a8a621b8bffbd2302a12ffe288a9d931670948749463b1" +dependencies = [ + "unicode-bidi", + "unicode-normalization", + "unicode-properties", +] + [[package]] name = "strsim" version = "0.11.1" @@ -4704,6 +4875,32 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "tokio-postgres" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3" +dependencies = [ + "async-trait", + "byteorder", + "bytes", + "fallible-iterator 0.2.0", + "futures-channel", + "futures-util", + "log", + "parking_lot", + "percent-encoding", + "phf", + "pin-project-lite", + "postgres-protocol", + "postgres-types", + "rand 0.10.1", + "socket2", + "tokio", + "tokio-util", + "whoami", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -4918,12 +5115,33 @@ dependencies = [ "libc", ] +[[package]] +name = "unicode-bidi" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c1cb5db39152898a79168971543b1cb5020dff7fe43c8dc468b0885f5e29df5" + [[package]] name = "unicode-ident" version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-normalization" +version = "0.1.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fd4f6878c9cb28d874b009da9e8d183b5abc80117c40bbd187a1fde336be6e8" +dependencies = [ + "tinyvec", +] + +[[package]] +name = "unicode-properties" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" + [[package]] name = "unicode-width" version = "0.1.14" @@ -5072,6 +5290,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasi" +version = "0.14.7+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" +dependencies = [ + "wasip2", +] + [[package]] name = "wasip2" version = "1.0.3+wasi-0.2.9" @@ -5090,6 +5317,15 @@ dependencies = [ "wit-bindgen 0.51.0", ] +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + [[package]] name = "wasm-bindgen" version = "0.2.122" @@ -5278,6 +5514,19 @@ dependencies = [ "libc", ] +[[package]] +name = "whoami" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "626c4bac6755d76ffc12cb01b2eac751db1996b9e0041de9aa02c8c211ddc82c" +dependencies = [ + "libc", + "libredox", + "objc2-system-configuration", + "wasite", + "web-sys", +] + [[package]] name = "wholesym" version = "0.8.1" diff --git a/Cargo.toml b/Cargo.toml index 8bd28f039..fc5a75c0d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,7 @@ sha256 = "1.6" tokio = { version = "1", features = ["macros", "rt"] } tokio-tar = { package = "astral-tokio-tar", version = "0.6.2" } tokio-util = "0.7.18" +tokio-postgres = "0.7" md5 = "0.8" base64 = "0.22.1" async-compression = { version = "0.4.42", features = ["tokio", "gzip"] } diff --git a/src/cli/run/mod.rs b/src/cli/run/mod.rs index 312f4dfc6..543141037 100644 --- a/src/cli/run/mod.rs +++ b/src/cli/run/mod.rs @@ -9,7 +9,7 @@ use crate::run_environment::interfaces::RepositoryProvider; use crate::upload::poll_results::PollResultsOptions; use clap::{Args, ValueEnum}; use std::collections::HashMap; -use std::path::{Path, PathBuf}; +use std::path::Path; use url::Url; pub mod helpers; @@ -31,10 +31,11 @@ pub struct RunArgs { #[arg(long)] pub mongo_uri_env_name: Option, - /// Host path to the Postgres analytics dump written by the codspeed/postgres - /// image's poller. Required when the `postgres` instrument is enabled. + /// Connection string for the runner's own superuser connection to the + /// benchmarked database, used to reset/snapshot pg_stat_statements at + /// benchmark boundaries. Required when the `postgres` instrument is enabled. #[arg(long)] - pub postgres_dump_path: Option, + pub postgres_dsn: Option, #[arg(long, hide = true)] pub message_format: Option, @@ -92,7 +93,7 @@ impl RunArgs { }, instruments: vec![], mongo_uri_env_name: None, - postgres_dump_path: None, + postgres_dsn: None, message_format: None, command: vec![], } diff --git a/src/executor/mod.rs b/src/executor/mod.rs index 6f9c1f85f..80d31ea87 100644 --- a/src/executor/mod.rs +++ b/src/executor/mod.rs @@ -14,7 +14,6 @@ mod valgrind; mod wall_time; use crate::instruments::mongo_tracer::{MongoTracer, install_mongodb_tracer}; -use crate::instruments::postgres::PostgresInstrument; use crate::local_logger::rolling_buffer::{activate_rolling_buffer, deactivate_rolling_buffer}; use crate::prelude::*; use crate::runner_mode::RunnerMode; @@ -213,18 +212,6 @@ pub async fn run_executor( if let Some(mut mongo_tracer) = mongo_tracer { mongo_tracer.stop().await?; } - - // A benchmark run must not fail over analytics enrichment, so a collect - // error is logged and swallowed rather than propagated. - if let Some(postgres_config) = &execution_context.config.instruments.postgres { - if let Err(e) = - PostgresInstrument::new(&execution_context.profile_folder, postgres_config) - .collect() - .await - { - warn!("Failed to collect Postgres analytics: {e:#}"); - } - } debug!("Tearing down the executor"); executor.teardown(execution_context).await?; diff --git a/src/executor/wall_time/executor.rs b/src/executor/wall_time/executor.rs index a419f7bc3..18ec6831c 100644 --- a/src/executor/wall_time/executor.rs +++ b/src/executor/wall_time/executor.rs @@ -19,6 +19,7 @@ use crate::executor::shared::fifo::FifoBenchmarkData; use crate::executor::shared::fifo::RunnerFifo; use crate::executor::{ExecutionContext, ExecutorName, ExecutorSupport}; use crate::instruments::mongo_tracer::MongoTracer; +use crate::instruments::postgres::PostgresInstrument; use crate::prelude::*; use crate::runner_mode::RunnerMode; use crate::system::{SupportedOs, SystemInfo}; @@ -151,6 +152,18 @@ impl Executor for WallTimeExecutor { let status = match profiler.as_mut() { Some(profiler) if execution_context.config.enable_profiler => { + // Open the runner's own connection for per-benchmark pg_stat_statements + // capture; a failure disables the instrument but never fails the run. + let pg_observer = match &execution_context.config.instruments.postgres { + Some(cfg) => match PostgresInstrument::connect(cfg).await { + Ok(observer) => Some(observer), + Err(e) => { + warn!("Postgres instrument disabled: {e:#}"); + None + } + }, + None => None, + }; run_with_profiler( profiler.as_mut(), cmd_builder, @@ -158,6 +171,7 @@ impl Executor for WallTimeExecutor { &execution_context.profile_folder, requires_sudo, benchmark_state, + pg_observer, ) .await } @@ -213,6 +227,7 @@ async fn run_with_profiler( profile_folder: &Path, requires_sudo: bool, benchmark_state: &OnceCell<(FifoBenchmarkData, ExecutionTimestamps)>, + mut pg_observer: Option, ) -> Result { let wrapped = profiler .wrap_command(cmd_builder, config, profile_folder, requires_sudo) @@ -223,13 +238,26 @@ async fn run_with_profiler( let mut runner_fifo = RunnerFifo::new()?; run_command_with_log_pipe_and_callback(cmd, async move |mut child| { + // The Postgres instrument runs its reset/snapshot on the runner's own + // connection in the pre-Ack window, so the SQL is outside the measured + // region; failures are logged and never abort the run. let on_cmd = async |c: &FifoCommand| match c { FifoCommand::StartProfiler => { profiler.on_start_profiler().await?; + if let Some(pg) = pg_observer.as_ref() + && let Err(e) = pg.reset().await + { + warn!("Postgres reset at benchmark start failed: {e:#}"); + } Ok(None) } FifoCommand::StopProfiler => { profiler.on_stop_profiler().await?; + if let Some(pg) = pg_observer.as_mut() + && let Err(e) = pg.snapshot().await + { + warn!("Postgres snapshot at benchmark stop failed: {e:#}"); + } Ok(None) } #[allow(deprecated)] @@ -247,6 +275,12 @@ async fn run_with_profiler( let (timestamps, fifo_data, exit_status) = runner_fifo.handle_fifo_messages(&mut child, on_cmd).await?; + if let Some(mut pg) = pg_observer + && let Err(e) = pg.finalize(×tamps, profile_folder).await + { + warn!("Failed to write Postgres analytics: {e:#}"); + } + let _ = benchmark_state.set((fifo_data, timestamps)); Ok(exit_status) diff --git a/src/instruments/mod.rs b/src/instruments/mod.rs index 60f9ab1e8..027973797 100644 --- a/src/instruments/mod.rs +++ b/src/instruments/mod.rs @@ -16,8 +16,9 @@ pub struct MongoDBConfig { #[derive(Debug, Clone, PartialEq, Eq)] pub struct PostgresConfig { - /// Host path to the analytics dump written by the codspeed/postgres poller. - pub dump_path: std::path::PathBuf, + /// Connection string for the runner's own superuser connection, used to + /// reset and snapshot `pg_stat_statements` at benchmark boundaries. + pub dsn: String, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -83,15 +84,13 @@ impl TryFrom<&RunArgs> for Instruments { }; let postgres = if validated_instrument_names.contains(&InstrumentName::Postgres) { - let dump_path = args.postgres_dump_path.clone().ok_or_else(|| { - anyhow!( - "The Postgres instrument is enabled but --postgres-dump-path was not provided" - ) + let dsn = args.postgres_dsn.clone().ok_or_else(|| { + anyhow!("The Postgres instrument is enabled but --postgres-dsn was not provided") })?; - Some(PostgresConfig { dump_path }) - } else if args.postgres_dump_path.is_some() { + Some(PostgresConfig { dsn }) + } else if args.postgres_dsn.is_some() { warn!( - "The Postgres instrument is disabled but a Postgres dump path was provided, ignoring it" + "The Postgres instrument is disabled but a Postgres DSN was provided, ignoring it" ); None } else { @@ -173,19 +172,19 @@ mod tests { fn test_from_args_postgres() { let args = RunArgs { instruments: vec!["postgres".into()], - postgres_dump_path: Some("/tmp/x/dump.json".into()), + postgres_dsn: Some("postgresql://codspeed@localhost/codspeed_bench".into()), ..RunArgs::test() }; let instruments = Instruments::try_from(&args).unwrap(); assert!(instruments.is_postgres_enabled()); assert_eq!( - instruments.postgres.unwrap().dump_path, - std::path::PathBuf::from("/tmp/x/dump.json") + instruments.postgres.unwrap().dsn, + "postgresql://codspeed@localhost/codspeed_bench" ); } #[test] - fn test_from_args_postgres_without_dump_path() { + fn test_from_args_postgres_without_dsn() { let args = RunArgs { instruments: vec!["postgres".into()], ..RunArgs::test() @@ -194,7 +193,7 @@ mod tests { assert!(instruments.is_err()); assert_eq!( instruments.unwrap_err().to_string(), - "The Postgres instrument is enabled but --postgres-dump-path was not provided" + "The Postgres instrument is enabled but --postgres-dsn was not provided" ); } } diff --git a/src/instruments/postgres.rs b/src/instruments/postgres.rs index 71c5e6c50..dad6f5005 100644 --- a/src/instruments/postgres.rs +++ b/src/instruments/postgres.rs @@ -1,176 +1,317 @@ -use std::path::{Path, PathBuf}; -use std::time::{Duration, SystemTime}; +use std::collections::{HashMap, HashSet}; +use std::path::Path; -use tokio::fs; +use runner_shared::artifacts::ExecutionTimestamps; +use serde::Serialize; +use serde_json::Value; +use tokio_postgres::{Client, NoTls, SimpleQueryMessage}; use crate::prelude::*; use super::PostgresConfig; -/// Collects the analytics dump produced by the `codspeed/postgres` image's -/// poller and drops it into the profile folder so it is uploaded with the run. +/// One statement's `pg_stat_statements` counters plus its plan (attached at +/// finalize). This is the per-query shape written into the artifact. +#[derive(Debug, Clone, Serialize)] +pub struct PostgresQuery { + pub sql: String, + pub calls: i64, + pub rows: i64, + pub shared_blks_hit: i64, + pub shared_blks_read: i64, + /// EXPLAIN plan tree, or `null` when the statement can't be explained. + pub plan: Value, +} + +/// The queries a single benchmark issued, keyed to its URI (the same URI as its +/// flamegraph region). +#[derive(Debug, Clone, Serialize, PartialEq)] +pub struct BenchmarkQueries { + pub uri: String, + pub queries: Vec, +} + +#[derive(Debug, Serialize)] +struct PostgresDump { + benchmarks: Vec, +} + +// so BenchmarkQueries can derive PartialEq for the zip test +impl PartialEq for PostgresQuery { + fn eq(&self, other: &Self) -> bool { + self.sql == other.sql + && self.calls == other.calls + && self.rows == other.rows + && self.shared_blks_hit == other.shared_blks_hit + && self.shared_blks_read == other.shared_blks_read + && self.plan == other.plan + } +} + +/// Drives per-benchmark `pg_stat_statements` capture on the runner's own +/// (superuser) connection, keyed to benchmark boundaries. /// -/// Unlike the MongoDB instrument, this owns no process and does no proxying: the -/// poller runs inside the database image and writes the dump to a file, so this -/// only copies that file into the profile folder. -#[derive(Debug)] +/// Unlike the whole-run poller it replaces, this resets at each benchmark start +/// and snapshots at each benchmark stop — driven from the FIFO boundary hooks so +/// the SQL runs outside the measured region — then keys each snapshot to the +/// benchmark URI (the same one its flamegraph uses) at finalize. pub struct PostgresInstrument { - profile_folder: PathBuf, - dump_path: PathBuf, - max_flush_wait: Duration, - flush_poll_interval: Duration, + client: Client, + /// One counter snapshot per benchmark, in stop-boundary (execution) order. + snapshots: Vec>, + /// EXPLAIN plan cache keyed by normalized SQL — a plan does not change during + /// a run, so each statement is explained at most once. + plans: HashMap, } -/// Bound on how long to wait for the poller to flush a fresh dump after the -/// benchmark ends, and how often to re-check. The poller flushes on a fixed -/// interval (2s by default), so the bound is a few of those. -const MAX_FLUSH_WAIT: Duration = Duration::from_secs(6); -const FLUSH_POLL_INTERVAL: Duration = Duration::from_millis(200); - impl PostgresInstrument { - pub fn new(profile_folder: &Path, config: &PostgresConfig) -> Self { - Self { - profile_folder: profile_folder.to_path_buf(), - dump_path: config.dump_path.clone(), - max_flush_wait: MAX_FLUSH_WAIT, - flush_poll_interval: FLUSH_POLL_INTERVAL, - } + /// Connect on the given DSN and ensure the extension exists. + pub async fn connect(config: &PostgresConfig) -> Result { + let (client, connection) = tokio_postgres::connect(&config.dsn, NoTls) + .await + .context("connecting the Postgres instrument to the database")?; + tokio::spawn(async move { + if let Err(err) = connection.await { + warn!("Postgres instrument connection closed: {err}"); + } + }); + client + .simple_query("CREATE EXTENSION IF NOT EXISTS pg_stat_statements") + .await + .context("creating pg_stat_statements extension")?; + Ok(Self { + client, + snapshots: Vec::new(), + plans: HashMap::new(), + }) } - /// Copy the dump into `/instruments/postgres.json`, so it - /// rides along in the uploaded profile archive. A missing dump is a warning, - /// not a failure — a benchmark run must not fail over missing analytics. - pub async fn collect(&self) -> Result<()> { - // The benchmark has just finished. The poller flushes the dump on a fixed - // interval, so copying it immediately would drop up to one interval of the - // final queries. Wait for a flush that happened after now (the run end). - self.wait_for_fresh_dump(SystemTime::now()).await; - - if !self.dump_path.exists() { - warn!( - "Postgres instrument enabled but no dump found at {}; skipping", - self.dump_path.display() - ); - return Ok(()); - } + /// Reset `pg_stat_statements` at a benchmark start, so the next snapshot holds + /// only this benchmark's statements. + pub async fn reset(&self) -> Result<()> { + self.client + .simple_query("SELECT pg_stat_statements_reset()") + .await + .context("resetting pg_stat_statements")?; + Ok(()) + } - let instruments_out_dir = self.profile_folder.join("instruments"); - fs::create_dir_all(&instruments_out_dir).await?; - let dest = instruments_out_dir.join("postgres.json"); - fs::copy(&self.dump_path, &dest).await.with_context(|| { - format!( - "Failed to copy Postgres dump from {}", - self.dump_path.display() + /// Snapshot the current counters at a benchmark stop and buffer them (one + /// entry per benchmark, in execution order). + pub async fn snapshot(&mut self) -> Result<()> { + let rows = self + .client + .query( + "SELECT s.query, s.calls, s.rows, s.shared_blks_hit, s.shared_blks_read \ + FROM pg_stat_statements s JOIN pg_database d ON d.oid = s.dbid \ + WHERE d.datname = current_database() ORDER BY s.calls DESC, s.query", + &[], ) + .await + .context("reading pg_stat_statements")?; + + let queries = rows + .into_iter() + .filter_map(|row| { + let sql: String = row.get(0); + if is_self_query(&sql) { + return None; + } + Some(PostgresQuery { + sql, + calls: row.get(1), + rows: row.get(2), + shared_blks_hit: row.get(3), + shared_blks_read: row.get(4), + plan: Value::Null, + }) + }) + .collect(); + self.snapshots.push(queries); + Ok(()) + } + + /// EXPLAIN each distinct statement, key the buffered snapshots to benchmark + /// URIs (via `uri_by_ts`, the same zipping the flamegraph uses), and write the + /// per-benchmark artifact to `/instruments/postgres.json`. + pub async fn finalize( + &mut self, + timestamps: &ExecutionTimestamps, + profile_folder: &Path, + ) -> Result<()> { + self.attach_plans().await; + + let benchmarks = zip_benchmarks(×tamps.uri_by_ts, &self.snapshots); + + let out_dir = profile_folder.join("instruments"); + tokio::fs::create_dir_all(&out_dir) + .await + .with_context(|| format!("creating {}", out_dir.display()))?; + let dest = out_dir.join("postgres.json"); + let tmp = out_dir.join("postgres.json.tmp"); + let bytes = serde_json::to_vec_pretty(&PostgresDump { + benchmarks: benchmarks.clone(), })?; - debug!("Collected Postgres analytics into {}", dest.display()); + tokio::fs::write(&tmp, &bytes) + .await + .with_context(|| format!("writing {}", tmp.display()))?; + tokio::fs::rename(&tmp, &dest) + .await + .with_context(|| format!("renaming {} to {}", tmp.display(), dest.display()))?; + debug!( + "Collected Postgres analytics for {} benchmark(s) into {}", + benchmarks.len(), + dest.display() + ); Ok(()) } - /// Poll the dump's mtime until it advances past `after`, so the copied dump - /// reflects a poller flush that happened after the last benchmark query. The - /// poller writes atomically (temp file + rename), so a bumped mtime means a - /// complete document. Bounded: the poller may already be idle or gone. - async fn wait_for_fresh_dump(&self, after: SystemTime) { - let deadline = tokio::time::Instant::now() + self.max_flush_wait; - while tokio::time::Instant::now() < deadline { - if let Ok(meta) = fs::metadata(&self.dump_path).await { - if let Ok(mtime) = meta.modified() { - if mtime >= after { - return; - } + /// Fill each buffered query's `plan` from the EXPLAIN cache, explaining any + /// not-yet-seen statement once. + async fn attach_plans(&mut self) { + let distinct: HashSet = self + .snapshots + .iter() + .flat_map(|snap| snap.iter().map(|q| q.sql.clone())) + .collect(); + for sql in distinct { + if self.plans.contains_key(&sql) { + continue; + } + let plan = match explain(&self.client, &sql).await { + Ok(plan) => plan, + Err(err) => { + debug!("EXPLAIN skipped for `{sql}`: {err:#}"); + Value::Null } + }; + self.plans.insert(sql, plan); + } + for snap in &mut self.snapshots { + for query in snap.iter_mut() { + query.plan = self.plans.get(&query.sql).cloned().unwrap_or(Value::Null); } - tokio::time::sleep(self.flush_poll_interval).await; } + } +} + +/// Zip per-benchmark snapshots to their URIs in execution order — mirroring how +/// the flamegraph keys each `SampleStart..SampleEnd` region to `uri_by_ts`. +fn zip_benchmarks( + uri_by_ts: &[(u64, String)], + snapshots: &[Vec], +) -> Vec { + if uri_by_ts.len() != snapshots.len() { warn!( - "Postgres dump at {} did not refresh within {:?}; using the latest available snapshot", - self.dump_path.display(), - self.max_flush_wait + "Postgres: {} benchmark URIs but {} snapshots; zipping by index", + uri_by_ts.len(), + snapshots.len() ); } + uri_by_ts + .iter() + .zip(snapshots.iter()) + .map(|((_, uri), queries)| BenchmarkQueries { + uri: uri.clone(), + queries: queries.clone(), + }) + .collect() } -#[cfg(test)] -mod tests { - use super::*; - use std::process::id; +/// The instrument's own bookkeeping queries, which must not appear in the dump. +fn is_self_query(sql: &str) -> bool { + let sql = sql.trim_start(); + sql.starts_with("EXPLAIN") || sql.contains("pg_stat_statements") +} - impl PostgresInstrument { - fn with_timing(mut self, max_flush_wait: Duration, flush_poll_interval: Duration) -> Self { - self.max_flush_wait = max_flush_wait; - self.flush_poll_interval = flush_poll_interval; - self +/// Run `EXPLAIN` for `sql` and return the plan tree. Uses `GENERIC_PLAN` (PG16+) +/// when the statement is parameterized — `pg_stat_statements` normalizes literals +/// to `$1`, which can't be planned otherwise — over `simple_query` so the `$1` +/// stays part of the explained statement rather than becoming an EXPLAIN param. +async fn explain(client: &Client, sql: &str) -> Result { + let options = if has_placeholder(sql) { + "GENERIC_PLAN, FORMAT JSON" + } else { + "FORMAT JSON" + }; + let messages = client + .simple_query(&format!("EXPLAIN ({options}) {sql}")) + .await + .context("running EXPLAIN")?; + for msg in messages { + if let SimpleQueryMessage::Row(row) = msg { + let text = row.get(0).context("EXPLAIN row missing plan column")?; + return serde_json::from_str(text).context("parsing EXPLAIN JSON"); } } + bail!("EXPLAIN returned no rows"); +} - fn scratch(tag: &str) -> PathBuf { - std::env::temp_dir().join(format!("pg-instr-{}-{tag}", id())) - } - - fn instrument(profile: &Path, dump_path: PathBuf, max: Duration) -> PostgresInstrument { - PostgresInstrument::new(profile, &PostgresConfig { dump_path }) - .with_timing(max, Duration::from_millis(20)) - } +/// Whether `sql` contains a `$N` parameter placeholder. +fn has_placeholder(sql: &str) -> bool { + sql.as_bytes() + .windows(2) + .any(|w| w[0] == b'$' && w[1].is_ascii_digit()) +} - /// collect() must not copy the dump the poller last flushed before the run - /// ended: it waits for a flush newer than the run end, so the final queries - /// are captured instead of dropped by the poller's tick interval. - #[tokio::test] - async fn collect_waits_for_post_run_flush() { - let base = scratch("fresh"); - let _ = tokio::fs::remove_dir_all(&base).await; - tokio::fs::create_dir_all(&base).await.unwrap(); - let dump_path = base.join("dump.json"); - let profile = base.join("pf"); - tokio::fs::write(&dump_path, br#"{"queries":[]}"#) - .await - .unwrap(); +#[cfg(test)] +mod tests { + use super::*; - let instr = instrument(&profile, dump_path.clone(), Duration::from_secs(6)); + fn q(sql: &str, calls: i64) -> PostgresQuery { + PostgresQuery { + sql: sql.into(), + calls, + rows: 0, + shared_blks_hit: 0, + shared_blks_read: 0, + plan: Value::Null, + } + } - let dp = dump_path.clone(); - let flush = tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(300)).await; - tokio::fs::write(&dp, br#"{"queries":[{"sql":"select 1"}]}"#) - .await - .unwrap(); - }); + #[test] + fn zips_snapshots_to_uris_in_order() { + let uri_by_ts = vec![(10, "bench::a".to_string()), (20, "bench::b".to_string())]; + let snapshots = vec![vec![q("select 1", 3)], vec![q("select 2", 5)]]; - instr.collect().await.unwrap(); - flush.await.unwrap(); + let out = zip_benchmarks(&uri_by_ts, &snapshots); - let copied = tokio::fs::read_to_string(profile.join("instruments/postgres.json")) - .await - .unwrap(); - assert!( - copied.contains("select 1"), - "collect copied a stale pre-flush dump: {copied}" + assert_eq!( + out, + vec![ + BenchmarkQueries { + uri: "bench::a".into(), + queries: vec![q("select 1", 3)], + }, + BenchmarkQueries { + uri: "bench::b".into(), + queries: vec![q("select 2", 5)], + }, + ] ); - let _ = tokio::fs::remove_dir_all(&base).await; } - /// A missing dump is not a failure — a benchmark run must not fail over - /// missing analytics — and nothing is written. - #[tokio::test] - async fn collect_missing_dump_is_ok_and_writes_nothing() { - let base = scratch("missing"); - let _ = tokio::fs::remove_dir_all(&base).await; - tokio::fs::create_dir_all(&base).await.unwrap(); - let profile = base.join("pf"); - - let instr = instrument( - &profile, - base.join("does-not-exist.json"), - Duration::from_millis(100), - ); + #[test] + fn zip_tolerates_length_mismatch_by_truncating() { + let uri_by_ts = vec![(10, "bench::a".to_string())]; + let snapshots = vec![vec![q("select 1", 1)], vec![q("select 2", 2)]]; + let out = zip_benchmarks(&uri_by_ts, &snapshots); + assert_eq!(out.len(), 1); + assert_eq!(out[0].uri, "bench::a"); + } - instr.collect().await.unwrap(); - assert!( - !profile.join("instruments/postgres.json").exists(), - "no dump should be written when the source is missing" - ); - let _ = tokio::fs::remove_dir_all(&base).await; + #[test] + fn detects_parameter_placeholders() { + assert!(has_placeholder("select * from t where id = $1")); + assert!(!has_placeholder("select * from t")); + assert!(!has_placeholder("select '$x' from t")); + } + + #[test] + fn skips_self_queries() { + assert!(is_self_query("EXPLAIN (FORMAT JSON) select 1")); + assert!(is_self_query("SELECT pg_stat_statements_reset()")); + assert!(!is_self_query("select * from users")); } } From af05a497c7650ed37e83ce62f5b33da4e5896d13 Mon Sep 17 00:00:00 2001 From: fargito Date: Sat, 12 Sep 2026 16:05:26 +0200 Subject: [PATCH 3/3] refactor: pass the postgres DSN by env-var name, not value Take --postgres-dsn-env-name (the name of an env var holding the DSN) instead of the DSN itself, mirroring --mongo-uri-env-name. The DSN is resolved at connect time and never stored in the config, so it can't leak into the config debug dump, runner.log, or the uploaded archive, nor onto the command line. Also skip the artifact entirely (rather than zip by index) when the URI and snapshot counts disagree, so queries are never misattributed to the wrong benchmark, and keep a DSN-gated isolation test covering the reset-per-boundary invariant. --- src/cli/run/mod.rs | 12 ++-- src/instruments/mod.rs | 30 ++++---- src/instruments/postgres.rs | 138 ++++++++++++++++++++++++++++++------ 3 files changed, 137 insertions(+), 43 deletions(-) diff --git a/src/cli/run/mod.rs b/src/cli/run/mod.rs index 543141037..4e744a109 100644 --- a/src/cli/run/mod.rs +++ b/src/cli/run/mod.rs @@ -31,11 +31,13 @@ pub struct RunArgs { #[arg(long)] pub mongo_uri_env_name: Option, - /// Connection string for the runner's own superuser connection to the - /// benchmarked database, used to reset/snapshot pg_stat_statements at - /// benchmark boundaries. Required when the `postgres` instrument is enabled. + /// Name of the environment variable holding the connection string for the + /// runner's own superuser connection to the benchmarked database, used to + /// reset/snapshot pg_stat_statements at benchmark boundaries. Passed by name + /// (not value) so the DSN never lands on the command line or in logs. + /// Required when the `postgres` instrument is enabled. #[arg(long)] - pub postgres_dsn: Option, + pub postgres_dsn_env_name: Option, #[arg(long, hide = true)] pub message_format: Option, @@ -93,7 +95,7 @@ impl RunArgs { }, instruments: vec![], mongo_uri_env_name: None, - postgres_dsn: None, + postgres_dsn_env_name: None, message_format: None, command: vec![], } diff --git a/src/instruments/mod.rs b/src/instruments/mod.rs index 027973797..10167ba4f 100644 --- a/src/instruments/mod.rs +++ b/src/instruments/mod.rs @@ -16,9 +16,10 @@ pub struct MongoDBConfig { #[derive(Debug, Clone, PartialEq, Eq)] pub struct PostgresConfig { - /// Connection string for the runner's own superuser connection, used to - /// reset and snapshot `pg_stat_statements` at benchmark boundaries. - pub dsn: String, + /// Name of the environment variable holding the DSN for the runner's own + /// superuser connection. Resolved at connect time so the DSN is never stored + /// in the config (and thus never dumped to logs or the uploaded archive). + pub dsn_env_name: String, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -84,13 +85,15 @@ impl TryFrom<&RunArgs> for Instruments { }; let postgres = if validated_instrument_names.contains(&InstrumentName::Postgres) { - let dsn = args.postgres_dsn.clone().ok_or_else(|| { - anyhow!("The Postgres instrument is enabled but --postgres-dsn was not provided") + let dsn_env_name = args.postgres_dsn_env_name.clone().ok_or_else(|| { + anyhow!( + "The Postgres instrument is enabled but --postgres-dsn-env-name was not provided" + ) })?; - Some(PostgresConfig { dsn }) - } else if args.postgres_dsn.is_some() { + Some(PostgresConfig { dsn_env_name }) + } else if args.postgres_dsn_env_name.is_some() { warn!( - "The Postgres instrument is disabled but a Postgres DSN was provided, ignoring it" + "The Postgres instrument is disabled but a Postgres DSN env var name was provided, ignoring it" ); None } else { @@ -172,19 +175,16 @@ mod tests { fn test_from_args_postgres() { let args = RunArgs { instruments: vec!["postgres".into()], - postgres_dsn: Some("postgresql://codspeed@localhost/codspeed_bench".into()), + postgres_dsn_env_name: Some("PGTRACER_DSN".into()), ..RunArgs::test() }; let instruments = Instruments::try_from(&args).unwrap(); assert!(instruments.is_postgres_enabled()); - assert_eq!( - instruments.postgres.unwrap().dsn, - "postgresql://codspeed@localhost/codspeed_bench" - ); + assert_eq!(instruments.postgres.unwrap().dsn_env_name, "PGTRACER_DSN"); } #[test] - fn test_from_args_postgres_without_dsn() { + fn test_from_args_postgres_without_dsn_env_name() { let args = RunArgs { instruments: vec!["postgres".into()], ..RunArgs::test() @@ -193,7 +193,7 @@ mod tests { assert!(instruments.is_err()); assert_eq!( instruments.unwrap_err().to_string(), - "The Postgres instrument is enabled but --postgres-dsn was not provided" + "The Postgres instrument is enabled but --postgres-dsn-env-name was not provided" ); } } diff --git a/src/instruments/postgres.rs b/src/instruments/postgres.rs index dad6f5005..d3dc537cd 100644 --- a/src/instruments/postgres.rs +++ b/src/instruments/postgres.rs @@ -65,9 +65,15 @@ pub struct PostgresInstrument { } impl PostgresInstrument { - /// Connect on the given DSN and ensure the extension exists. + /// Resolve the DSN from its env var, connect, and ensure the extension exists. pub async fn connect(config: &PostgresConfig) -> Result { - let (client, connection) = tokio_postgres::connect(&config.dsn, NoTls) + let dsn = std::env::var(&config.dsn_env_name).with_context(|| { + format!( + "reading the Postgres DSN from ${} (--postgres-dsn-env-name)", + config.dsn_env_name + ) + })?; + let (client, connection) = tokio_postgres::connect(&dsn, NoTls) .await .context("connecting the Postgres instrument to the database")?; tokio::spawn(async move { @@ -141,7 +147,13 @@ impl PostgresInstrument { ) -> Result<()> { self.attach_plans().await; - let benchmarks = zip_benchmarks(×tamps.uri_by_ts, &self.snapshots); + // On a URI/snapshot count mismatch, skip the artifact entirely rather than + // zip by index — that would misattribute queries to the wrong benchmark, + // which is worse than emitting nothing (nothing downstream can detect it). + let Some(benchmarks) = zip_benchmarks(×tamps.uri_by_ts, &self.snapshots) else { + return Ok(()); + }; + let count = benchmarks.len(); let out_dir = profile_folder.join("instruments"); tokio::fs::create_dir_all(&out_dir) @@ -149,9 +161,7 @@ impl PostgresInstrument { .with_context(|| format!("creating {}", out_dir.display()))?; let dest = out_dir.join("postgres.json"); let tmp = out_dir.join("postgres.json.tmp"); - let bytes = serde_json::to_vec_pretty(&PostgresDump { - benchmarks: benchmarks.clone(), - })?; + let bytes = serde_json::to_vec_pretty(&PostgresDump { benchmarks })?; tokio::fs::write(&tmp, &bytes) .await .with_context(|| format!("writing {}", tmp.display()))?; @@ -160,8 +170,7 @@ impl PostgresInstrument { .with_context(|| format!("renaming {} to {}", tmp.display(), dest.display()))?; debug!( - "Collected Postgres analytics for {} benchmark(s) into {}", - benchmarks.len(), + "Collected Postgres analytics for {count} benchmark(s) into {}", dest.display() ); Ok(()) @@ -201,22 +210,25 @@ impl PostgresInstrument { fn zip_benchmarks( uri_by_ts: &[(u64, String)], snapshots: &[Vec], -) -> Vec { +) -> Option> { if uri_by_ts.len() != snapshots.len() { warn!( - "Postgres: {} benchmark URIs but {} snapshots; zipping by index", + "Postgres: {} benchmark URIs but {} snapshots; skipping artifact to avoid misattribution", uri_by_ts.len(), snapshots.len() ); + return None; } - uri_by_ts - .iter() - .zip(snapshots.iter()) - .map(|((_, uri), queries)| BenchmarkQueries { - uri: uri.clone(), - queries: queries.clone(), - }) - .collect() + Some( + uri_by_ts + .iter() + .zip(snapshots.iter()) + .map(|((_, uri), queries)| BenchmarkQueries { + uri: uri.clone(), + queries: queries.clone(), + }) + .collect(), + ) } /// The instrument's own bookkeeping queries, which must not appear in the dump. @@ -275,7 +287,7 @@ mod tests { let uri_by_ts = vec![(10, "bench::a".to_string()), (20, "bench::b".to_string())]; let snapshots = vec![vec![q("select 1", 3)], vec![q("select 2", 5)]]; - let out = zip_benchmarks(&uri_by_ts, &snapshots); + let out = zip_benchmarks(&uri_by_ts, &snapshots).unwrap(); assert_eq!( out, @@ -293,12 +305,11 @@ mod tests { } #[test] - fn zip_tolerates_length_mismatch_by_truncating() { + fn zip_skips_on_length_mismatch() { let uri_by_ts = vec![(10, "bench::a".to_string())]; let snapshots = vec![vec![q("select 1", 1)], vec![q("select 2", 2)]]; - let out = zip_benchmarks(&uri_by_ts, &snapshots); - assert_eq!(out.len(), 1); - assert_eq!(out[0].uri, "bench::a"); + // A mismatch must yield None (skip the artifact) rather than misattribute. + assert!(zip_benchmarks(&uri_by_ts, &snapshots).is_none()); } #[test] @@ -314,4 +325,85 @@ mod tests { assert!(is_self_query("SELECT pg_stat_statements_reset()")); assert!(!is_self_query("select * from users")); } + + /// Isolation invariant against a real database: `reset()` at each benchmark + /// boundary must prevent one benchmark's queries from leaking into the next, + /// and `EXPLAIN (GENERIC_PLAN)` must capture a plan for `$1`-parameterized SQL. + /// + /// A no-op unless `PG_SMOKE_DSN` is set. To run it: + /// docker run -d --name pg -e POSTGRES_USER=codspeed -e POSTGRES_PASSWORD=codspeed \ + /// -e POSTGRES_DB=codspeed_bench -p 5546:5432 fargito/test-pg-tracer:16 + /// psql "$DSN" -c "create table t(id int primary key, v text)" + /// PG_SMOKE_DSN="postgresql://codspeed:codspeed@127.0.0.1:5546/codspeed_bench?sslmode=disable" \ + /// cargo test --lib instruments::postgres::tests::capture_isolates_benchmarks_against_real_db + #[tokio::test] + async fn capture_isolates_benchmarks_against_real_db() { + let Ok(dsn) = std::env::var("PG_SMOKE_DSN") else { + return; + }; + let mut obs = PostgresInstrument::connect(&PostgresConfig { + dsn_env_name: "PG_SMOKE_DSN".into(), + }) + .await + .unwrap(); + let (app, conn) = tokio_postgres::connect(&dsn, NoTls).await.unwrap(); + tokio::spawn(async move { + let _ = conn.await; + }); + app.simple_query("create table if not exists t(id int primary key, v text)") + .await + .unwrap(); + + obs.reset().await.unwrap(); + for _ in 0..3 { + app.execute("select count(*) from t where id = $1", &[&7i32]) + .await + .unwrap(); + } + obs.snapshot().await.unwrap(); + + obs.reset().await.unwrap(); + for _ in 0..2 { + app.execute("select count(*) from t where v = $1", &[&"v9"]) + .await + .unwrap(); + } + obs.snapshot().await.unwrap(); + + let ts = ExecutionTimestamps { + uri_by_ts: vec![(1, "bench::a".into()), (2, "bench::b".into())], + markers: vec![], + }; + let dir = std::env::temp_dir().join(format!("pg-smoke-{}", std::process::id())); + let _ = tokio::fs::remove_dir_all(&dir).await; + obs.finalize(&ts, &dir).await.unwrap(); + + let out = tokio::fs::read_to_string(dir.join("instruments/postgres.json")) + .await + .unwrap(); + let v: Value = serde_json::from_str(&out).unwrap(); + let b = v["benchmarks"].as_array().unwrap(); + assert_eq!(b.len(), 2); + assert_eq!(b[0]["uri"], "bench::a"); + assert_eq!(b[1]["uri"], "bench::b"); + let a_has_id_with_plan = b[0]["queries"] + .as_array() + .unwrap() + .iter() + .any(|q| q["sql"].as_str().unwrap().contains("where id =") && !q["plan"].is_null()); + assert!( + a_has_id_with_plan, + "bench A must have the id query with a plan: {out}" + ); + let b_leaked = b[1]["queries"] + .as_array() + .unwrap() + .iter() + .any(|q| q["sql"].as_str().unwrap().contains("where id =")); + assert!( + !b_leaked, + "reset failed: bench B leaked bench A's queries: {out}" + ); + let _ = tokio::fs::remove_dir_all(&dir).await; + } }