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/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..4e744a109 100644 --- a/src/cli/run/mod.rs +++ b/src/cli/run/mod.rs @@ -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,14 @@ pub struct RunArgs { #[arg(long)] pub mongo_uri_env_name: Option, + /// 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_env_name: Option, + #[arg(long, hide = true)] pub message_format: Option, @@ -87,6 +95,7 @@ impl RunArgs { }, instruments: vec![], mongo_uri_env_name: None, + postgres_dsn_env_name: None, message_format: None, command: vec![], } 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 c16a00de2..10167ba4f 100644 --- a/src/instruments/mod.rs +++ b/src/instruments/mod.rs @@ -7,20 +7,31 @@ 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 { + /// 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)] 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 +39,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 +50,10 @@ impl Instruments { names.push(InstrumentName::MongoDB); } + if self.is_postgres_enabled() { + names.push(InstrumentName::Postgres); + } + names } } @@ -47,6 +66,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 +84,23 @@ impl TryFrom<&RunArgs> for Instruments { None }; - Ok(Self { mongodb }) + let postgres = if validated_instrument_names.contains(&InstrumentName::Postgres) { + 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_env_name }) + } else if args.postgres_dsn_env_name.is_some() { + warn!( + "The Postgres instrument is disabled but a Postgres DSN env var name was provided, ignoring it" + ); + None + } else { + None + }; + + Ok(Self { mongodb, postgres }) } } @@ -76,6 +112,7 @@ impl Instruments { mongodb: Some(MongoDBConfig { uri_env_name: Some("MONGODB_URI".into()), }), + postgres: None, } } } @@ -133,4 +170,30 @@ mod tests { "Invalid instrument name: unknown" ); } + + #[test] + fn test_from_args_postgres() { + let args = RunArgs { + instruments: vec!["postgres".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_env_name, "PGTRACER_DSN"); + } + + #[test] + fn test_from_args_postgres_without_dsn_env_name() { + 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-dsn-env-name was not provided" + ); + } } diff --git a/src/instruments/postgres.rs b/src/instruments/postgres.rs new file mode 100644 index 000000000..d3dc537cd --- /dev/null +++ b/src/instruments/postgres.rs @@ -0,0 +1,409 @@ +use std::collections::{HashMap, HashSet}; +use std::path::Path; + +use runner_shared::artifacts::ExecutionTimestamps; +use serde::Serialize; +use serde_json::Value; +use tokio_postgres::{Client, NoTls, SimpleQueryMessage}; + +use crate::prelude::*; + +use super::PostgresConfig; + +/// 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 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 { + 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, +} + +impl PostgresInstrument { + /// Resolve the DSN from its env var, connect, and ensure the extension exists. + pub async fn connect(config: &PostgresConfig) -> Result { + 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 { + 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(), + }) + } + + /// 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(()) + } + + /// 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; + + // 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) + .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 })?; + 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 {count} benchmark(s) into {}", + dest.display() + ); + Ok(()) + } + + /// 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); + } + } + } +} + +/// 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], +) -> Option> { + if uri_by_ts.len() != snapshots.len() { + warn!( + "Postgres: {} benchmark URIs but {} snapshots; skipping artifact to avoid misattribution", + uri_by_ts.len(), + snapshots.len() + ); + return None; + } + 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. +fn is_self_query(sql: &str) -> bool { + let sql = sql.trim_start(); + sql.starts_with("EXPLAIN") || sql.contains("pg_stat_statements") +} + +/// 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"); +} + +/// 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()) +} + +#[cfg(test)] +mod tests { + use super::*; + + 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, + } + } + + #[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)]]; + + let out = zip_benchmarks(&uri_by_ts, &snapshots).unwrap(); + + 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)], + }, + ] + ); + } + + #[test] + 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)]]; + // A mismatch must yield None (skip the artifact) rather than misattribute. + assert!(zip_benchmarks(&uri_by_ts, &snapshots).is_none()); + } + + #[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")); + } + + /// 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; + } +}