diff --git a/CHANGELOG.md b/CHANGELOG.md index ce819d4..0c9264f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ## [Unreleased] +### Added + +- `sqlx_otel.connection.abandoned_wait.duration` records the time already waited, once, when a polled `acquire()` is dropped before SQLx answers. `db.client.connection.wait_time` and `db.client.connection.timeouts` keep their meaning, and a never-polled future records nothing. +- `PoolAttribution` splits `db.client.connection.wait_time`, `use_time` and `timeouts`, plus the new instruments, by one bounded `&'static str` attribute captured at acquisition. It uses a poll-scoped thread-local with no new dependency. **Queries that read those series without aggregating over attributes will see them split once a scope is used**; totals are unchanged. `pending_requests`, the pool gauges and the operation metrics are unaffected. Spawned tasks are not covered unless passed `PoolAttribution::inherit`. +- `sqlx_otel.transaction.begin.duration` (`ok`, `error`, `abandoned`) and `sqlx_otel.transaction.duration` (`commit_ok`, `commit_error`, `commit_abandoned`, `rollback_ok`, `rollback_error`, `rollback_abandoned`, `dropped`) report transaction boundaries and terminal client observations, distinct from `use_time`. They add no await, round trip or change to SQLx's ownership and drop behaviour. + ### Fixed - Implicit pool queries and `Pool::begin()` now route through the instrumented `Pool::acquire()`, so `db.client.connection.wait_time`, `use_time`, `pending_requests`, and `timeouts` account for every acquisition rather than only explicit `acquire()` calls. A workload written as `query(..).fetch_one(&pool)` previously borrowed its connection straight from `sqlx::Pool` and could saturate the pool without any of its waits or timeouts reaching telemetry. Queries issued inside a transaction still acquire once, and connection use time now follows the lease through transaction commit, rollback, or drop ([#40](https://github.com/chmodas/sqlx-otel/pull/40)). diff --git a/Cargo.toml b/Cargo.toml index 0da2108..e85e68f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,7 +55,7 @@ proptest = "1" serial_test = "3" sqlx = { version = "0.9", features = ["macros", "sqlite", "postgres", "mysql", "mysql-rsa", "runtime-tokio"] } testcontainers = "0.28" -tokio = { version = "1", features = ["macros", "rt-multi-thread"] } +tokio = { version = "1", features = ["macros", "rt-multi-thread", "io-util", "net", "sync", "time"] } [[example]] name = "sqlite" diff --git a/README.md b/README.md index 4d2182f..7ff8783 100644 --- a/README.md +++ b/README.md @@ -206,20 +206,71 @@ These mirror the bounded portion of the span attribute set: connection-level att | `db.client.connection.max` | Gauge | | Maximum number of connections allowed | | `db.client.connection.idle.max` | Gauge | | Maximum idle connections (equals `max` in SQLx) | | `db.client.connection.idle.min` | Gauge | | Configured minimum connections | +| `sqlx_otel.connection.abandoned_wait.duration` | Histogram | `s` | Time waited by a caller that dropped its `acquire()` | -Explicit `acquire()`, implicit pool queries (including annotated and streaming queries), and `begin()` share one instrumented acquisition; queries issued inside a transaction do not acquire again. Wait time covers completed acquisition attempts – semaphore wait, connection creation, and validation – but excludes `BEGIN` and query execution. Dropping a waiting future decrements pending requests without recording a completed wait or a timeout, since the pool never failed to deliver: the caller left. Only `sqlx::Error::PoolTimedOut` increments the timeout counter, so `PoolClosed` on shutdown does not inflate your timeout rate. Pending-request and timeout series are initialised at zero so dashboards start from a known baseline rather than waiting for the first event. Connection use time follows the lease through transaction commit, rollback, or drop. +Explicit `acquire()`, implicit pool queries (including annotated and streaming queries), and `begin()` share one instrumented acquisition; queries issued inside a transaction do not acquire again. Wait time covers completed acquisition attempts – semaphore wait, connection creation, and validation – but excludes `BEGIN` and query execution. Dropping a waiting future decrements pending requests without recording a completed wait or a timeout, since the pool never failed to deliver: the caller left. It records the time already waited once on `sqlx_otel.connection.abandoned_wait.duration` instead, a custom instrument because semantic conventions have no equivalent; a future dropped before its first poll never waited and records nothing. Only `sqlx::Error::PoolTimedOut` increments the timeout counter, so `PoolClosed` on shutdown does not inflate your timeout rate. Pending-request and timeout series are initialised at zero so dashboards start from a known baseline rather than waiting for the first event. Connection use time follows the lease through transaction commit, rollback, or drop, and ends when the lease owner is dropped. SQLx returns the connection to the pool in the background after that, so a sample does not prove that cleanup finished. + +`db.client.operation.duration` and the query span start when the executor method is called and end when the client observes the result. On `&Pool` that interval includes acquiring the connection, plus client and network time. It is not database execution time and does not prove a lock wait; compare it with `db.client.connection.wait_time` before reading it that way. `db.client.connection.count` is polled by a background task and requires both a runtime feature (`runtime-tokio` or `runtime-async-std`) and a pool name set via `PoolBuilder::with_pool_name`. Polling and export intervals can miss brief utilisation or pending-request peaks. The remaining three are static gauges recorded once at construction. +### Workload attribution + +Wrap a workload in `PoolAttribution::scope` to split the connection-lease instruments by one bounded attribute: + +```rust,ignore +use sqlx_otel::PoolAttribution; + +PoolAttribution::new("myapp.job.kind", "nightly_rollup") + .scope(run_job(&pool)) + .await; +``` + +- **Carries it:** `db.client.connection.wait_time`, `use_time` and `timeouts`, plus `sqlx_otel.connection.abandoned_wait.duration` and the two transaction histograms. It splits series without adding samples, so summing across it reproduces the pool-wide totals. A workload's timeout series is initialized at zero at its first acquisition; on backends that compute rates from cumulative samples, the first timeout is visible only if that zero was exported first. +- **Never carries it:** `pending_requests` and the pool gauges, which stay pool-wide, and `db.client.operation.duration` and the row-count histograms, whose `db.query.summary` dimension would multiply with it. +- **Reserved keys:** attribution using a connection attribute or either transaction-outcome key is ignored, even when that optional pool attribute is absent. This preserves pool identity and outcomes. +- **Bounded:** key and value are `&'static str`, so runtime identifiers cannot reach a label by accident. The added series are at most `pools × (values + 1)`, times the outcome count on the transaction histograms. +- **Immutable per lease:** it is captured when an acquisition first runs and stored with the attempt and its lease, so dropping a connection or transaction elsewhere cannot move the sample to another workload. +- **Poll-scoped, not task-scoped:** futures polled inside the scope inherit it, including `join!`, `select!` and `FuturesUnordered` children, and concurrent scopes on one task stay separate. `tokio::spawn`, `spawn_blocking` and other threads do **not** inherit it; use `PoolAttribution::inherit(future)` to pass it to a spawned task. A stream acquires when first polled, so the scope at that poll wins. It has no runtime dependency. + +Absence of the attribute means "acquired without valid attribution". + +### Transaction metrics + +| Instrument | Type | Unit | Bounded attribute | +|----------------------------------------|-----------|------|--------------------------------------------------------------| +| `sqlx_otel.transaction.begin.duration` | Histogram | `s` | `sqlx_otel.transaction.begin.outcome`: `ok`, `error`, `abandoned` | +| `sqlx_otel.transaction.duration` | Histogram | `s` | `sqlx_otel.transaction.outcome`, below | + +Semantic conventions define no transaction metrics, so both are custom instruments. Each carries the lease's connection attributes and attribution. `begin.duration` runs from holding a connection to observing the result of `BEGIN`, so it excludes acquisition. Only `ok` starts `transaction.duration`, which runs from there to the terminal client observation. That makes it distinct from `use_time`, which also spans `BEGIN` and the connection's handover, and which also ends for a failed `BEGIN`. + +Outcomes report what the client observed, never what the server did: + +| Outcome | The client observed | What it does not prove | +|----------------------|---------------------------------------------------------|--------------------------------------------------------------------------| +| `commit_ok` | `commit()` returned `Ok` | That data committed: PostgreSQL answers `COMMIT` on a failed transaction with a `ROLLBACK` tag and no error | +| `commit_error` | `commit()` returned `Err` | A rejected commit: after a transport error the server outcome is unknown | +| `commit_abandoned` | The `commit()` future was dropped mid-flight | Anything: `COMMIT` may have reached the server, so this is neither a rollback nor a failed commit | +| `rollback_ok` | `rollback()` returned `Ok` | | +| `rollback_error` | `rollback()` returned `Err` | That the transaction is still open | +| `rollback_abandoned` | The `rollback()` future was dropped mid-flight | A completed rollback: SQLx queued another for the connection's next use | +| `dropped` | Dropped with no terminal call started, or a `commit()` / `rollback()` future dropped before its first poll | A server rollback: SQLx only queued one for the connection's next use | + +`dropped` is a client lifecycle observation, not an error. It can be expected when a guard releases a transaction or an early return drops one; SQLx has only queued rollback, so it does not confirm server cleanup. + +Telemetry adds no await, round trip, or `Drop` behaviour to the transaction itself. `begin.duration` uses the boundaries below; `transaction.duration` extends them with `120, 300, 600, 1800, 3600`, because no client timeout caps a transaction and long batch transactions must not all collapse into the overflow bucket. + ### Histogram buckets -The three latency histograms – `db.client.operation.duration`, `db.client.connection.wait_time`, and `db.client.connection.use_time` – are recorded in seconds, as the semantic conventions require. The OpenTelemetry SDK's default bucket boundaries (`0, 5, 10, 25, … 10000`) are shaped for milliseconds, so on a seconds-valued histogram every observation below five seconds falls into a single bucket and any quantile derived from it is an interpolation artefact rather than a measurement. The crate therefore ships explicit boundaries spanning the range a database client actually occupies, from sub-millisecond local round trips up to the 30-second acquire timeout SQLx applies by default: +The three semantic-convention latency histograms – `db.client.operation.duration`, `db.client.connection.wait_time`, and `db.client.connection.use_time` – are recorded in seconds, as the semantic conventions require. The OpenTelemetry SDK's default bucket boundaries (`0, 5, 10, 25, … 10000`) are shaped for milliseconds, so on a seconds-valued histogram every observation below five seconds falls into a single bucket and any quantile derived from it is an interpolation artefact rather than a measurement. The crate therefore ships explicit boundaries spanning the range a database client actually occupies, from sub-millisecond local round trips up to the 30-second acquire timeout SQLx applies by default: ```text 0.0001, 0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0 ``` +`sqlx_otel.connection.abandoned_wait.duration` and `sqlx_otel.transaction.begin.duration` use the same boundaries. `sqlx_otel.transaction.duration` extends them as described above. + These are instrument *advice*, not a fixed configuration. The SDK applies advice only when no matching view has already set an aggregation, so an application that registers a `View` with its own explicit bucket boundaries overrides what the crate ships – no API is needed from this crate to retune them. (A view whose aggregation the SDK rejects as incompatible with the instrument is discarded, and the crate's advice applies after all, so match histogram names only.) ```rust,ignore diff --git a/src/attribution.rs b/src/attribution.rs new file mode 100644 index 0000000..d45000f --- /dev/null +++ b/src/attribution.rs @@ -0,0 +1,385 @@ +//! Poll-scoped attribution of connection-lifecycle metrics to a bounded workload. + +use std::cell::Cell; +use std::future::{Future, poll_fn}; +use std::pin::pin; + +use opentelemetry::KeyValue; + +thread_local! { + /// The attribution of the future this thread is polling right now. + /// + /// It is set only for the duration of one synchronous `poll`, so another task the runtime + /// schedules on this thread between polls can never observe it. The value has no destructor, + /// so reading it stays valid during thread teardown. + static CURRENT: Cell> = const { Cell::new(None) }; +} + +/// A single bounded attribute identifying the workload that uses a connection pool. +/// +/// While a future runs inside [`scope`](Self::scope), every connection it acquires from any +/// [`Pool`](crate::Pool) carries this attribute on `db.client.connection.wait_time`, +/// `db.client.connection.use_time`, `db.client.connection.timeouts`, +/// `sqlx_otel.connection.abandoned_wait.duration`, `sqlx_otel.transaction.begin.duration` and +/// `sqlx_otel.transaction.duration`. It splits existing series rather than adding samples, so +/// summing across it reproduces the pool-wide totals. +/// +/// Pool-wide state stays pool-wide: `db.client.connection.pending_requests` and the +/// `db.client.connection.count` / `max` / `idle.*` gauges never carry it. Neither do +/// `db.client.operation.duration` and the row-count histograms, whose `db.query.summary` dimension +/// would multiply with it. A workload's timeout series is initialized at zero at its first +/// acquisition. On backends that compute rates from cumulative samples, the first timeout is +/// visible only if that zero was exported first. +/// +/// # Cardinality +/// +/// Both strings are `&'static str`, so a value derived at runtime from an identifier, a query or +/// user input cannot reach a metric label without deliberately leaking memory. One attribution is +/// exactly one key and one value. Choose a key in your own namespace. Attribution using any of +/// these reserved keys is ignored, even when the corresponding optional pool attribute is absent: +/// +/// - `db.system.name`, `db.namespace`, `server.address`, `server.port` +/// - `network.peer.address`, `network.peer.port`, `network.protocol.name`, `network.transport` +/// - `db.client.connection.pool.name` +/// - `sqlx_otel.transaction.begin.outcome`, `sqlx_otel.transaction.outcome` +/// +/// The added series are bounded by `pools × (values + 1)`, multiplied by the outcome count +/// on the two transaction instruments; the extra value is work without emitted attribution. +/// +/// # What is and is not covered +/// +/// Attribution is captured once, when an acquisition first runs, and stored with the attempt and +/// the resulting lease. Everything recorded later uses that captured value, wherever and whenever +/// the connection or transaction is finally dropped, so it cannot drift to another workload. +/// +/// The scope itself follows polling, not "the task" or the call stack: +/// +/// - Futures polled inside the scope inherit it, including `join!`, `select!` and +/// `FuturesUnordered` children. Concurrent scopes on one task stay separate. +/// - **`tokio::spawn`, `spawn_blocking` and other threads do not inherit it.** Their acquisitions +/// record without this attribute unless you pass it on explicitly with +/// [`inherit`](Self::inherit). +/// - A stream or query future acquires when first polled, so the scope active at that first poll +/// wins, not the one where it was constructed. +/// - Anything polled re-entrantly inside the scope on the same thread, such as tasks driven by a +/// nested `LocalSet` or `block_on`, is inside the scope too. +/// +/// Absence of the attribute therefore means "acquired without valid attribution", not +/// "pool-internal". +/// +/// # Example +/// +/// ```no_run +/// # #[cfg(feature = "sqlite")] +/// # async fn _doc() -> Result<(), sqlx::Error> { +/// use sqlx::Executor as _; +/// use sqlx_otel::{PoolAttribution, PoolBuilder}; +/// # let pool = PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await?).build(); +/// +/// PoolAttribution::new("myapp.job.kind", "nightly_rollup") +/// .scope(async { +/// pool.execute("SELECT 1").await?; +/// +/// // A spawned task starts outside the scope unless it is passed on explicitly. +/// let spawned_pool = pool.clone(); +/// tokio::spawn(PoolAttribution::inherit(async move { +/// spawned_pool.execute("SELECT 2").await +/// })); +/// Ok::<_, sqlx::Error>(()) +/// }) +/// .await +/// # } +/// ``` +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PoolAttribution { + key: &'static str, + value: &'static str, +} + +impl PoolAttribution { + /// Create an attribution from one compile-time key and one compile-time value. + /// + /// Reserved connection and transaction-outcome keys are ignored when recording metrics; see + /// [cardinality](Self#cardinality) for the exact list. + #[must_use] + pub const fn new(key: &'static str, value: &'static str) -> Self { + Self { key, value } + } + + /// The attribution of the future being polled on this thread, if any. + #[must_use] + pub fn current() -> Option { + CURRENT.with(Cell::get) + } + + /// Run `future` with this attribution active during each of its polls. + /// + /// A nested scope replaces this one only while the inner future is being polled. Dropping the + /// returned future, even mid-way, leaves no state behind. + pub fn scope(self, future: F) -> impl Future { + poll_scoped(Some(self), future) + } + + /// Capture the attribution active *now* and run `future` under it, for handing work to a + /// spawned task. + /// + /// The capture happens when this function is called, not when the returned future is first + /// polled, so call it in the scope you want to propagate. When no scope is active, the future + /// also runs without one, even if it is later polled inside another scope. + pub fn inherit(future: F) -> impl Future { + poll_scoped(Self::current(), future) + } + + fn key_value(self) -> KeyValue { + KeyValue::new(self.key, self.value) + } +} + +async fn poll_scoped(attribution: Option, future: F) -> F::Output { + let mut future = pin!(future); + poll_fn(|cx| { + let _active = ActiveAttribution::enter(attribution); + future.as_mut().poll(cx) + }) + .await +} + +/// Restores the enclosing attribution on every exit from a poll, including a panic. +struct ActiveAttribution { + enclosing: Option, +} + +impl ActiveAttribution { + fn enter(attribution: Option) -> Self { + Self { + enclosing: CURRENT.with(|current| current.replace(attribution)), + } + } +} + +impl Drop for ActiveAttribution { + fn drop(&mut self) { + CURRENT.with(|current| current.set(self.enclosing)); + } +} + +// Reserve optional connection keys even when their pool values are absent, as well as the +// outcome keys appended later by transaction recording. Telemetry must not overwrite identity or +// outcomes, and invalid attribution must not panic in the acquisition path. +const RESERVED_KEYS: &[&str] = &[ + "db.system.name", + "db.namespace", + "server.address", + "server.port", + "network.peer.address", + "network.peer.port", + "network.protocol.name", + "network.transport", + "db.client.connection.pool.name", + "sqlx_otel.transaction.begin.outcome", + "sqlx_otel.transaction.outcome", +]; + +/// The pool-wide attribute set plus any valid attribution active right now. +/// +/// Called once per acquisition. The result is owned by the attempt and lease, so nothing recorded +/// later reads the thread-local again. +pub(crate) fn lease_attributes(pool_attrs: &[KeyValue]) -> Vec { + let mut attrs = Vec::with_capacity(pool_attrs.len() + 2); + attrs.extend_from_slice(pool_attrs); + if let Some(attribution) = PoolAttribution::current() + && !RESERVED_KEYS.contains(&attribution.key) + { + attrs.push(attribution.key_value()); + } + attrs +} + +#[cfg(test)] +mod tests { + use super::{PoolAttribution, lease_attributes}; + use futures::executor::block_on; + use futures::future::{FutureExt, lazy, poll_fn}; + use opentelemetry::KeyValue; + use std::panic::{AssertUnwindSafe, catch_unwind}; + use std::pin::pin; + use std::task::Poll; + + const A: PoolAttribution = PoolAttribution::new("test.kind", "a"); + const B: PoolAttribution = PoolAttribution::new("test.kind", "b"); + + /// Returns `Pending` once and wakes itself, so a sibling future gets polled in between. + async fn yield_once() { + let mut yielded = false; + poll_fn(|cx| { + if yielded { + Poll::Ready(()) + } else { + yielded = true; + cx.waker().wake_by_ref(); + Poll::Pending + } + }) + .await; + } + + async fn unreachable_after_yield() { + yield_once().await; + unreachable!("dropped while suspended"); + } + + fn assert_send(value: T) -> T { + value + } + + #[test] + fn no_attribution_outside_a_scope() { + assert_eq!(PoolAttribution::current(), None); + assert_eq!(lease_attributes(&[]), Vec::::new()); + } + + /// The scope is visible on every poll and gone between polls: this is what stops another task + /// sharing the worker thread from seeing it. + #[test] + fn scope_is_active_only_while_polling() { + block_on(async { + let mut scoped = pin!(A.scope(async { + assert_eq!(PoolAttribution::current(), Some(A)); + yield_once().await; + assert_eq!(PoolAttribution::current(), Some(A)); + })); + assert_eq!(PoolAttribution::current(), None); + assert!(futures::poll!(&mut scoped).is_pending()); + assert_eq!( + PoolAttribution::current(), + None, + "must not outlive the poll" + ); + assert!(futures::poll!(&mut scoped).is_ready()); + assert_eq!(PoolAttribution::current(), None); + }); + } + + #[test] + fn nested_scope_replaces_and_restores() { + block_on(A.scope(async { + B.scope(async { + assert_eq!(PoolAttribution::current(), Some(B)); + yield_once().await; + assert_eq!(PoolAttribution::current(), Some(B)); + }) + .await; + assert_eq!(PoolAttribution::current(), Some(A)); + })); + } + + /// Two scopes interleaved on one task, each resuming after the other ran. + #[test] + fn concurrent_scopes_on_one_task_never_see_each_other_across_awaits() { + async fn observe(expected: Option, rounds: usize) -> usize { + for _ in 0..rounds { + assert_eq!(PoolAttribution::current(), expected); + yield_once().await; + } + assert_eq!(PoolAttribution::current(), expected); + rounds + } + block_on(async { + let (a, b, unscoped) = futures::join!( + A.scope(observe(Some(A), 25)), + B.scope(observe(Some(B), 25)), + // The enclosing unscoped body must also stay unscoped between its siblings. + observe(None, 25), + ); + assert_eq!((a, b, unscoped), (25, 25, 25)); + }); + } + + #[test] + fn dropping_a_suspended_scope_leaves_nothing_behind() { + block_on(async { + let mut scoped = Box::pin(A.scope(unreachable_after_yield())); + assert!(futures::poll!(&mut scoped).is_pending()); + drop(scoped); + assert_eq!(PoolAttribution::current(), None); + }); + } + + #[test] + fn a_never_polled_scope_has_no_effect() { + drop(A.scope(lazy(|_| -> () { unreachable!("never polled") }))); + assert_eq!(PoolAttribution::current(), None); + } + + #[test] + fn a_panicking_poll_restores_the_enclosing_attribution() { + let panicked = catch_unwind(AssertUnwindSafe(|| { + A.scope(lazy(|_| -> () { panic!("inside scope") })) + .now_or_never(); + })); + assert!(panicked.is_err()); + assert_eq!(PoolAttribution::current(), None); + } + + #[test] + fn inherit_captures_at_the_call_not_at_the_first_poll() { + let inherited = block_on(A.scope(lazy(|_| { + PoolAttribution::inherit(async { PoolAttribution::current() }) + }))); + assert_eq!(PoolAttribution::current(), None); + assert_eq!(block_on(inherited), Some(A), "polled outside any scope"); + + let inherited = block_on(A.scope(lazy(|_| { + PoolAttribution::inherit(async { PoolAttribution::current() }) + }))); + assert_eq!(block_on(B.scope(inherited)), Some(A), "polled inside B"); + } + + #[test] + fn inheriting_no_attribution_masks_a_later_scope() { + let unscoped = PoolAttribution::inherit(async { PoolAttribution::current() }); + assert_eq!(block_on(A.scope(unscoped)), None); + } + + #[test] + fn scoped_futures_stay_send_for_spawning() { + drop(assert_send(A.scope(yield_once()))); + drop(assert_send(PoolAttribution::inherit(yield_once()))); + } + + #[test] + fn lease_attributes_append_exactly_the_active_attribution() { + let base = [KeyValue::new("db.system.name", "sqlite")]; + block_on(A.scope(async { + assert_eq!( + lease_attributes(&base), + vec![base[0].clone(), KeyValue::new("test.kind", "a")] + ); + })); + assert_eq!(lease_attributes(&base), base); + } + + #[test] + fn reserved_attribution_never_changes_pool_attributes_even_when_optional_keys_are_absent() { + for key in [ + "db.system.name", + "db.namespace", + "server.address", + "server.port", + "network.peer.address", + "network.peer.port", + "network.protocol.name", + "network.transport", + "db.client.connection.pool.name", + "sqlx_otel.transaction.begin.outcome", + "sqlx_otel.transaction.outcome", + ] { + let attribution = PoolAttribution::new(key, "spoofed"); + let present = [KeyValue::new(key, "canonical")]; + block_on(attribution.scope(async { + assert_eq!(lease_attributes(&present), present, "{key} already present"); + assert!(lease_attributes(&[]).is_empty(), "{key} absent"); + })); + } + } +} diff --git a/src/connection.rs b/src/connection.rs index 7b53e1c..ae4e665 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -11,8 +11,9 @@ use crate::pool::SharedState; /// /// Acquired via [`Pool::acquire`](crate::Pool::acquire). `&mut PoolConnection` implements /// [`sqlx::Executor`], so the value plugs straight into the standard `SQLx` query builders. When -/// the value is dropped, the connection is returned to the pool and `db.client.connection.use_time` -/// is recorded with the elapsed time between `acquire` and `drop`. +/// the value is dropped, `db.client.connection.use_time` is recorded with the elapsed time between +/// `acquire` and `drop`, and the connection starts its way back to the pool. `SQLx` finishes that +/// return in the background, so the sample is not proof that the connection is available again. /// /// Use [`with_annotations`](Self::with_annotations) / [`with_operation`](Self::with_operation) to /// attach per-query semantic-convention attributes – the methods are inherited via the same macro @@ -42,13 +43,15 @@ pub struct PoolConnection { pub(crate) struct ConnectionUsage { pub(crate) use_time: Arc>, pub(crate) acquired_at: Instant, - pub(crate) base_attrs: Vec, + /// Pool attributes plus the [`PoolAttribution`](crate::PoolAttribution) captured when the + /// acquisition started. Never re-read, so the lease's end cannot drift to another workload. + pub(crate) lease_attrs: Vec, } impl Drop for ConnectionUsage { fn drop(&mut self) { self.use_time - .record(self.acquired_at.elapsed().as_secs_f64(), &self.base_attrs); + .record(self.acquired_at.elapsed().as_secs_f64(), &self.lease_attrs); } } diff --git a/src/lib.rs b/src/lib.rs index d9a0f53..d329b93 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -98,6 +98,11 @@ //! These carry the connection-level attributes (`db.system.name`, `db.namespace`, `server.address`, //! `server.port`). //! +//! `db.client.operation.duration` and the query span both start when the executor method is called +//! and end when the client observes the result. On `&Pool` that interval includes acquiring the +//! connection, plus client and network time. It is not server SQL execution time and does not +//! prove a lock wait; compare it with `db.client.connection.wait_time` before reading it that way. +//! //! ## Connection-pool metrics //! //! | Instrument | Type | Unit | Description | @@ -110,20 +115,47 @@ //! | `db.client.connection.max` | Gauge | | Maximum number of connections allowed | //! | `db.client.connection.idle.max` | Gauge | | Maximum idle connections (equals `max` in `SQLx`) | //! | `db.client.connection.idle.min` | Gauge | | Configured minimum connections | +//! | `sqlx_otel.connection.abandoned_wait.duration` | Histogram | `s` | Time waited by a caller that dropped its `acquire()` | //! //! Explicit acquisition, implicit pool queries (including annotated and streaming queries), and //! `begin()` share one instrumented acquisition; queries inside a transaction do not acquire //! again. Wait time covers completed acquisition attempts and excludes `BEGIN` and query //! execution. Caller cancellation clears pending without recording a completed wait or a pool -//! timeout, since the pool never failed to deliver – the caller left. Only -//! `sqlx::Error::PoolTimedOut` increments the timeout counter, so `PoolClosed` on shutdown does +//! timeout, since the pool never failed to deliver – the caller left; instead it records the time +//! already waited once on `sqlx_otel.connection.abandoned_wait.duration`, which has no semantic +//! convention equivalent. A future dropped before its first poll never waited and records nothing. +//! Only `sqlx::Error::PoolTimedOut` increments the timeout counter, so `PoolClosed` on shutdown does //! not inflate the timeout rate. Connection use time follows the lease through transaction -//! completion. Pending and timeout series start at zero. Polling/export can miss brief gauge -//! peaks. `connection.count` is polled by a background task and requires both +//! completion and ends when the lease owner is dropped; `SQLx` returns the connection to the pool +//! in the background after that, so a sample does not prove that cleanup finished. Pending and +//! timeout series start at zero. Polling/export can miss brief gauge peaks. `connection.count` is +//! polled by a background task and requires both //! [`PoolBuilder::with_pool_name`] and a runtime feature (`runtime-tokio` or //! `runtime-async-std`); without either, the gauge is silent. The remaining three are static //! gauges recorded once at [`PoolBuilder::build`]. //! +//! ## Workload attribution +//! +//! [`PoolAttribution::scope`] splits the connection-lease instruments by one bounded `&'static str` +//! attribute for connections acquired inside the scope, captured once at acquisition. It splits +//! series without adding samples; `pending_requests`, the pool gauges and the operation metrics +//! stay untouched. Scopes follow polling, so spawned tasks need +//! [`PoolAttribution::inherit`]. See [`PoolAttribution`] for the exact coverage. +//! +//! ## Transaction metrics +//! +//! | Instrument | Type | Unit | Bounded attribute | +//! |----------------------------------------|-----------|------|---------------------------------------| +//! | `sqlx_otel.transaction.begin.duration` | Histogram | `s` | `sqlx_otel.transaction.begin.outcome` | +//! | `sqlx_otel.transaction.duration` | Histogram | `s` | `sqlx_otel.transaction.outcome` | +//! +//! Semantic conventions define no transaction metrics, so both live in this crate's namespace and +//! carry the lease attributes. The first covers `BEGIN` once a connection is held (`ok`, `error`, +//! `abandoned`); only `ok` starts the second, whose interval ends at the terminal client +//! observation and so excludes acquisition and `BEGIN`. See [`Transaction`] for its outcomes, none +//! of which claims a server result: a dropped transaction only queued a rollback, and an abandoned +//! `COMMIT` has an unknown outcome. +//! //! ## Histogram buckets //! //! `db.client.operation.duration`, `db.client.connection.wait_time`, and @@ -141,6 +173,10 @@ //! retuning them. A view whose aggregation the SDK rejects as incompatible falls back to the //! implicit default view, which applies this advice after all. See the README for an example. //! +//! `sqlx_otel.connection.abandoned_wait.duration` and `sqlx_otel.transaction.begin.duration` use +//! the same boundaries. `sqlx_otel.transaction.duration` extends them to one hour, because no +//! client timeout caps a transaction. +//! //! ## A note on `db.response.affected_rows` //! //! `db.response.affected_rows` is **not part of the OpenTelemetry semantic conventions** at @@ -251,6 +287,7 @@ #[macro_use] mod annotations; pub(crate) mod attributes; +mod attribution; mod compact; mod connection; mod database; @@ -264,9 +301,11 @@ mod query_ext; mod rebuilt_query; mod runtime; mod transaction; +mod transaction_metrics; pub use annotations::{Annotated, AnnotatedMut, QueryAnnotations}; pub use attributes::QueryTextMode; +pub use attribution::PoolAttribution; pub use connection::PoolConnection; pub use database::Database; pub use pool::{Pool, PoolBuilder}; diff --git a/src/metrics.rs b/src/metrics.rs index d6ed1a5..1607d76 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -4,7 +4,8 @@ use opentelemetry::KeyValue; use opentelemetry::metrics::{Histogram, Meter}; use opentelemetry_semantic_conventions::metric; -/// Explicit bucket boundaries, in seconds, for every latency histogram the crate records. +/// Shared bucket boundaries, in seconds, for operation, connection and `BEGIN` latency. +/// Transaction duration uses the extended set below. /// /// The `OTel` SDK's default boundaries (`0, 5, 10, 25, ... 10000`) are shaped for milliseconds, /// but semconv requires these histograms be recorded in seconds. Left on the defaults, every @@ -23,6 +24,17 @@ pub(crate) const LATENCY_BUCKETS_SECONDS: &[f64] = &[ 30.0, 60.0, ]; +/// Explicit bucket boundaries, in seconds, for `sqlx_otel.transaction.duration`. +/// +/// The shared latency set, extended past 60 seconds. Unlike an acquisition, a transaction has no +/// client-side timeout that caps it, and the batch transactions that run for minutes are exactly +/// the ones holding a connection for longest, so they must not all collapse into the overflow +/// bucket. +pub(crate) const TRANSACTION_BUCKETS_SECONDS: &[f64] = &[ + 0.0001, 0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, + 30.0, 60.0, 120.0, 300.0, 600.0, 1800.0, 3600.0, +]; + /// Holds the OpenTelemetry metric instruments for per-operation recording. /// /// Created once per [`Pool`](crate::Pool) and shared (via `Arc`) across all wrapper types @@ -91,7 +103,19 @@ impl Metrics { #[cfg(test)] mod tests { - use super::LATENCY_BUCKETS_SECONDS; + use super::{LATENCY_BUCKETS_SECONDS, TRANSACTION_BUCKETS_SECONDS}; + + /// The transaction set must stay comparable with `use_time` wherever the two ranges overlap, + /// and well-formed for an explicit-bucket histogram. + #[test] + fn transaction_buckets_extend_the_latency_set() { + assert!(TRANSACTION_BUCKETS_SECONDS.starts_with(LATENCY_BUCKETS_SECONDS)); + assert!( + TRANSACTION_BUCKETS_SECONDS + .windows(2) + .all(|pair| pair[0] < pair[1]) + ); + } /// The boundary set must be well-formed for an explicit-bucket histogram. /// diff --git a/src/pool.rs b/src/pool.rs index a57d4c3..d74f593 100644 --- a/src/pool.rs +++ b/src/pool.rs @@ -1,5 +1,5 @@ use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use opentelemetry_semantic_conventions::metric as semconv_metric; @@ -9,6 +9,7 @@ use crate::connection::PoolConnection; use crate::database::Database; use crate::metrics::Metrics; use crate::transaction::Transaction; +use crate::transaction_metrics::TransactionMetrics; /// Shared state propagated to every wrapper type derived from a pool. #[derive(Debug, Clone)] @@ -183,7 +184,8 @@ impl PoolBuilder { /// `db.client.connection.idle.max`, `db.client.connection.idle.min`) are recorded /// once with the connection-level attributes – they do not change over the pool's /// lifetime. The wait-time / use-time / timeout / pending-request instruments are - /// created here and updated inline on every `acquire()` and connection drop. + /// created here and updated inline on every `acquire()` and connection drop. So are the + /// abandoned-wait and transaction instruments, which have no series until first used. #[must_use] pub fn build(self) -> Pool { let metrics_shutdown = self.spawn_pool_metrics_task(); @@ -262,6 +264,17 @@ impl PoolBuilder { .with_description("The number of pending requests for an open connection.") .build(), ), + abandoned_wait: Arc::new( + meter + .f64_histogram("sqlx_otel.connection.abandoned_wait.duration") + .with_unit("s") + .with_description( + "The time a caller waited for a connection before dropping the attempt.", + ) + .with_boundaries(crate::metrics::LATENCY_BUCKETS_SECONDS.to_vec()) + .build(), + ), + transactions: Arc::new(TransactionMetrics::new(&meter)), }; pool.pending_requests.add(0, &base_attrs); pool.timeouts.add(0, &base_attrs); @@ -354,6 +367,11 @@ pub struct Pool { /// Up/down counter for `db.client.connection.pending_requests`, tracks callers /// currently waiting in `acquire()`. pending_requests: Arc>, + /// Histogram for `sqlx_otel.connection.abandoned_wait.duration`, recorded when an `acquire()` + /// already waiting is dropped. + abandoned_wait: Arc>, + /// Instruments for the `sqlx_otel.transaction.*` family. + transactions: Arc, } impl Clone for Pool { @@ -366,6 +384,8 @@ impl Clone for Pool { use_time: self.use_time.clone(), timeouts: self.timeouts.clone(), pending_requests: self.pending_requests.clone(), + abandoned_wait: self.abandoned_wait.clone(), + transactions: self.transactions.clone(), } } } @@ -379,22 +399,35 @@ impl Pool { /// returned [`PoolConnection`] records `db.client.connection.use_time` when dropped /// and is itself an [`sqlx::Executor`] via `&mut conn`. /// + /// Dropping this future while it waits records the time already waited once on + /// `sqlx_otel.connection.abandoned_wait.duration`, not as a completed wait or a timeout. A + /// future that was never polled never waited and records nothing. The active + /// [`PoolAttribution`](crate::PoolAttribution) is captured on first poll and used for every + /// measurement of this attempt and its lease. + /// /// # Errors /// /// Returns `sqlx::Error` if a connection cannot be obtained from the pool – typically /// `PoolTimedOut` when the configured acquire timeout elapses, or `PoolClosed` after /// [`close`](Self::close). pub async fn acquire(&self) -> Result, sqlx::Error> { - let attrs = self.state.attrs.base_key_values(); - self.pending_requests.add(1, &attrs); + let pool_attrs = self.state.attrs.base_key_values(); + let attrs = crate::attribution::lease_attributes(&pool_attrs); + // Initialize the workload series like the pool-wide series. A rate can include the + // first timeout in a cumulative-sample rate only if the zero was exported first. + self.timeouts.add(0, &attrs); + self.pending_requests.add(1, &pool_attrs); let pending = PendingAcquisition { counter: &self.pending_requests, - attrs: &attrs, + pool_attrs: &pool_attrs, + abandoned_wait: &self.abandoned_wait, + lease_attrs: &attrs, + started_at: Instant::now(), + completed: false, }; - let start = std::time::Instant::now(); let result = self.inner.acquire().await; - drop(pending); - self.wait_time.record(start.elapsed().as_secs_f64(), &attrs); + let waited = pending.complete(); + self.wait_time.record(waited.as_secs_f64(), &attrs); if let Err(sqlx::Error::PoolTimedOut) = &result { self.timeouts.add(1, &attrs); @@ -405,8 +438,8 @@ impl Pool { state: self.state.clone(), usage: crate::connection::ConnectionUsage { use_time: self.use_time.clone(), - acquired_at: std::time::Instant::now(), - base_attrs: attrs, + acquired_at: Instant::now(), + lease_attrs: attrs, }, }) } @@ -419,6 +452,10 @@ impl Pool { /// it; dropping the value without doing either rolls back implicitly (per `SQLx`'s /// usual behaviour). /// + /// Records one `sqlx_otel.transaction.begin.duration` sample once a connection is held, + /// including when `BEGIN` fails or this future is dropped during it. Only a successful `BEGIN` + /// starts `sqlx_otel.transaction.duration`; see [`Transaction`] for its outcomes. + /// /// # Errors /// /// Returns `sqlx::Error` if `BEGIN` fails – typically due to a connection problem or @@ -429,11 +466,15 @@ impl Pool { state, usage, } = self.acquire().await?; - let inner = sqlx::Transaction::begin(inner, None).await?; + let (inner, lifecycle) = self + .transactions + .observe_begin(&usage.lease_attrs, sqlx::Transaction::begin(inner, None)) + .await?; Ok(Transaction { inner, state, usage, + lifecycle, }) } @@ -504,20 +545,37 @@ impl Pool { } /// Decrements `db.client.connection.pending_requests` on every exit path from -/// [`Pool::acquire`]. +/// [`Pool::acquire`], and measures the attempts that never got an answer. /// /// A caller timeout or task abort can drop the `acquire()` future while it is still awaiting /// `SQLx`, which would otherwise leave the counter permanently incremented. Tying the decrement -/// to a `Drop` guard covers success, error, and cancellation alike. Cancellation deliberately -/// stops there: it records no completed wait and is not classified as `PoolTimedOut`, because -/// the pool never failed to deliver – the caller left. +/// to a `Drop` guard covers success, error, and cancellation alike. Cancellation records no +/// completed wait and is not classified as `PoolTimedOut`, because the pool never failed to +/// deliver – the caller left. It records the elapsed wait on its own instrument instead, exactly +/// once, because the guard exists only after a first poll and is dropped only once. struct PendingAcquisition<'a> { counter: &'a opentelemetry::metrics::UpDownCounter, - attrs: &'a [opentelemetry::KeyValue], + pool_attrs: &'a [opentelemetry::KeyValue], + abandoned_wait: &'a opentelemetry::metrics::Histogram, + lease_attrs: &'a [opentelemetry::KeyValue], + started_at: Instant, + completed: bool, +} + +impl PendingAcquisition<'_> { + /// `SQLx` answered, so this attempt was not abandoned. Returns how long it took. + fn complete(mut self) -> Duration { + self.completed = true; + self.started_at.elapsed() + } } impl Drop for PendingAcquisition<'_> { fn drop(&mut self) { - self.counter.add(-1, self.attrs); + self.counter.add(-1, self.pool_attrs); + if !self.completed { + self.abandoned_wait + .record(self.started_at.elapsed().as_secs_f64(), self.lease_attrs); + } } } diff --git a/src/transaction.rs b/src/transaction.rs index 9fefbd9..c4f6ec6 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -1,5 +1,6 @@ use crate::database::Database; use crate::pool::SharedState; +use crate::transaction_metrics::{TransactionLifecycle, TransactionTermination}; /// An in-progress database transaction instrumented for OpenTelemetry. /// @@ -14,6 +15,25 @@ use crate::pool::SharedState; /// [`with_operation`](Self::with_operation) to attach per-query semantic-convention /// attributes. /// +/// `sqlx_otel.transaction.duration` runs from the client observing a successful `BEGIN` to the +/// terminal client observation, and its `sqlx_otel.transaction.outcome` attribute reports what the +/// client saw, not what the server did: +/// +/// | Outcome | The client observed | +/// |----------------------|----------------------------------------------------------------------| +/// | `commit_ok` | `commit()` returned `Ok` | +/// | `commit_error` | `commit()` returned `Err` | +/// | `commit_abandoned` | the `commit()` future was dropped mid-flight: server outcome unknown | +/// | `rollback_ok` | `rollback()` returned `Ok` | +/// | `rollback_error` | `rollback()` returned `Err` | +/// | `rollback_abandoned` | the `rollback()` future was dropped mid-flight | +/// | `dropped` | dropped with no terminal call started: a rollback was only queued | +/// +/// `dropped` is an observation, not an error: guard release or an early return can produce it +/// normally. It confirms only that `SQLx` queued rollback, not that server cleanup completed. +/// +/// Telemetry adds no await, round trip or `Drop` behaviour to the transaction itself. +/// /// # Example /// /// ```no_run @@ -34,6 +54,7 @@ pub struct Transaction<'c, DB: sqlx::Database> { pub(crate) inner: sqlx::Transaction<'c, DB>, pub(crate) state: SharedState, pub(crate) usage: crate::connection::ConnectionUsage, + pub(crate) lifecycle: TransactionLifecycle, } impl Transaction<'_, DB> @@ -46,7 +67,10 @@ where /// /// Returns `sqlx::Error` if the commit fails. pub async fn commit(self) -> Result<(), sqlx::Error> { - let result = self.inner.commit().await; + let result = self + .lifecycle + .finish(TransactionTermination::Commit, self.inner.commit()) + .await; drop(self.usage); result } @@ -57,7 +81,10 @@ where /// /// Returns `sqlx::Error` if the rollback fails. pub async fn rollback(self) -> Result<(), sqlx::Error> { - let result = self.inner.rollback().await; + let result = self + .lifecycle + .finish(TransactionTermination::Rollback, self.inner.rollback()) + .await; drop(self.usage); result } diff --git a/src/transaction_metrics.rs b/src/transaction_metrics.rs new file mode 100644 index 0000000..e084042 --- /dev/null +++ b/src/transaction_metrics.rs @@ -0,0 +1,474 @@ +//! Transaction lifecycle telemetry, distinct from the connection lease behind it. +//! +//! Every outcome names what the *client* observed. None of them is a statement about the server: +//! see the outcome variants for what each one does and does not prove. + +use std::future::Future; +use std::time::Instant; + +use opentelemetry::KeyValue; +use opentelemetry::metrics::{Histogram, Meter}; + +use crate::metrics::{LATENCY_BUCKETS_SECONDS, TRANSACTION_BUCKETS_SECONDS}; + +const BEGIN_DURATION: &str = "sqlx_otel.transaction.begin.duration"; +const BEGIN_OUTCOME: &str = "sqlx_otel.transaction.begin.outcome"; +const TRANSACTION_DURATION: &str = "sqlx_otel.transaction.duration"; +const TRANSACTION_OUTCOME: &str = "sqlx_otel.transaction.outcome"; + +/// How a `BEGIN` attempt on an already leased connection ended, as seen by the client. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BeginOutcome { + /// `BEGIN` returned `Ok`. The only outcome that starts a [`TransactionLifecycle`]. + Ok, + /// `BEGIN` returned `Err`. + Error, + /// The future was dropped after `BEGIN` started and before its result was observed. + Abandoned, +} + +impl BeginOutcome { + const fn as_str(self) -> &'static str { + match self { + Self::Ok => "ok", + Self::Error => "error", + Self::Abandoned => "abandoned", + } + } +} + +/// How a started transaction ended, as seen by the client. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TransactionOutcome { + /// `COMMIT` returned `Ok`. The server acknowledged the statement, which is not proof that data + /// committed: `PostgreSQL` answers `COMMIT` on an already failed transaction with a `ROLLBACK` + /// tag and no error. + CommitOk, + /// `COMMIT` returned `Err`. A database error means it did not commit; a transport error after + /// the statement was sent leaves the server outcome unknown. + CommitError, + /// The `commit()` future was dropped mid-flight. `COMMIT` may already have reached the server, + /// so the outcome is **unknown**: neither a rollback nor a failed commit. + CommitAbandoned, + /// `ROLLBACK` returned `Ok`. + RollbackOk, + /// `ROLLBACK` returned `Err`. + RollbackError, + /// The `rollback()` future was dropped mid-flight. `SQLx` queues another rollback for the + /// connection's next use; completion is unobserved. + RollbackAbandoned, + /// The transaction was dropped without a polled `commit()` or `rollback()`. `SQLx` only queued + /// a rollback for the connection's next use, so this never confirms a server rollback. + Dropped, +} + +impl TransactionOutcome { + const fn as_str(self) -> &'static str { + match self { + Self::CommitOk => "commit_ok", + Self::CommitError => "commit_error", + Self::CommitAbandoned => "commit_abandoned", + Self::RollbackOk => "rollback_ok", + Self::RollbackError => "rollback_error", + Self::RollbackAbandoned => "rollback_abandoned", + Self::Dropped => "dropped", + } + } +} + +/// The client request that ends a transaction. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum TransactionTermination { + Commit, + Rollback, +} + +impl TransactionTermination { + const fn abandoned(self) -> TransactionOutcome { + match self { + Self::Commit => TransactionOutcome::CommitAbandoned, + Self::Rollback => TransactionOutcome::RollbackAbandoned, + } + } + + const fn observed(self, succeeded: bool) -> TransactionOutcome { + match (self, succeeded) { + (Self::Commit, true) => TransactionOutcome::CommitOk, + (Self::Commit, false) => TransactionOutcome::CommitError, + (Self::Rollback, true) => TransactionOutcome::RollbackOk, + (Self::Rollback, false) => TransactionOutcome::RollbackError, + } + } +} + +/// Instruments shared by every transaction started from one [`Pool`](crate::Pool). +#[derive(Debug)] +pub(crate) struct TransactionMetrics { + begin_duration: Histogram, + duration: Histogram, +} + +impl TransactionMetrics { + pub(crate) fn new(meter: &Meter) -> Self { + Self { + begin_duration: meter + .f64_histogram(BEGIN_DURATION) + .with_unit("s") + .with_description( + "Time from holding a connection to observing the result of BEGIN, by outcome.", + ) + .with_boundaries(LATENCY_BUCKETS_SECONDS.to_vec()) + .build(), + duration: meter + .f64_histogram(TRANSACTION_DURATION) + .with_unit("s") + .with_description( + "Time from a successful BEGIN to the terminal client observation, by outcome.", + ) + .with_boundaries(TRANSACTION_BUCKETS_SECONDS.to_vec()) + .build(), + } + } + + /// Await `begin` on an already leased connection, recording exactly one `BEGIN` sample, and + /// start the lifecycle only if it succeeded. + /// + /// A failed or cancelled `BEGIN` never started a transaction, so it records no + /// `sqlx_otel.transaction.duration` sample. `attrs` are the lease's captured attributes. + pub(crate) async fn observe_begin( + &self, + attrs: &[KeyValue], + begin: F, + ) -> Result<(T, TransactionLifecycle), sqlx::Error> + where + F: Future>, + { + let mut attempt = BeginAttempt { + histogram: &self.begin_duration, + started_at: Instant::now(), + attrs, + outcome: BeginOutcome::Abandoned, + }; + let result = begin.await; + attempt.outcome = if result.is_ok() { + BeginOutcome::Ok + } else { + BeginOutcome::Error + }; + drop(attempt); + result.map(|inner| (inner, TransactionLifecycle::start(&self.duration, attrs))) + } +} + +/// `attrs` plus one bounded outcome attribute. +fn with_outcome(attrs: &[KeyValue], key: &'static str, outcome: &'static str) -> Vec { + let mut with_outcome = Vec::with_capacity(attrs.len() + 1); + with_outcome.extend_from_slice(attrs); + with_outcome.push(KeyValue::new(key, outcome)); + with_outcome +} + +/// Records one `BEGIN` sample when dropped, so cancellation is counted by the same path. +struct BeginAttempt<'a> { + histogram: &'a Histogram, + started_at: Instant, + attrs: &'a [KeyValue], + outcome: BeginOutcome, +} + +impl Drop for BeginAttempt<'_> { + fn drop(&mut self) { + self.histogram.record( + self.started_at.elapsed().as_secs_f64(), + &with_outcome(self.attrs, BEGIN_OUTCOME, self.outcome.as_str()), + ); + } +} + +/// Records one `sqlx_otel.transaction.duration` sample when dropped. +/// +/// It is owned by the [`Transaction`](crate::Transaction), so whichever of `commit`, `rollback` or +/// plain drop consumes the transaction also consumes this, exactly once. Dropping it does nothing +/// but record: `SQLx` keeps sole ownership of ending the transaction. +#[derive(Debug)] +pub(crate) struct TransactionLifecycle { + histogram: Histogram, + started_at: Instant, + attrs: Vec, + outcome: TransactionOutcome, +} + +impl TransactionLifecycle { + fn start(histogram: &Histogram, attrs: &[KeyValue]) -> Self { + Self { + histogram: histogram.clone(), + started_at: Instant::now(), + attrs: attrs.to_vec(), + outcome: TransactionOutcome::Dropped, + } + } + + /// Await the terminal `operation`, classifying it from the moment it first runs. + /// + /// Because the classification happens inside this future's body, a `commit()` that is created + /// but never polled stays [`Dropped`](TransactionOutcome::Dropped): `COMMIT` was never sent. + pub(crate) async fn finish( + mut self, + termination: TransactionTermination, + operation: F, + ) -> Result<(), sqlx::Error> + where + F: Future>, + { + self.outcome = termination.abandoned(); + let result = operation.await; + self.outcome = termination.observed(result.is_ok()); + result + } +} + +impl Drop for TransactionLifecycle { + fn drop(&mut self) { + self.histogram.record( + self.started_at.elapsed().as_secs_f64(), + &with_outcome(&self.attrs, TRANSACTION_OUTCOME, self.outcome.as_str()), + ); + } +} + +#[cfg(test)] +mod tests { + use super::{ + BEGIN_DURATION, BEGIN_OUTCOME, TRANSACTION_DURATION, TRANSACTION_OUTCOME, + TransactionLifecycle, TransactionMetrics, TransactionTermination, + }; + use futures::executor::block_on; + use futures::future::{pending, ready}; + use opentelemetry::KeyValue; + use opentelemetry::metrics::MeterProvider as _; + use opentelemetry_sdk::metrics::data::{ + AggregatedMetrics, HistogramDataPoint, MetricData, ResourceMetrics, ScopeMetrics, + }; + use opentelemetry_sdk::metrics::{InMemoryMetricExporter, PeriodicReader, SdkMeterProvider}; + use std::time::Duration; + + /// A meter of its own for each test: nothing here touches the global provider. + struct Harness { + provider: SdkMeterProvider, + exporter: InMemoryMetricExporter, + metrics: TransactionMetrics, + attrs: Vec, + } + + impl Harness { + fn new() -> Self { + let exporter = InMemoryMetricExporter::default(); + let provider = SdkMeterProvider::builder() + .with_reader(PeriodicReader::builder(exporter.clone()).build()) + .build(); + let metrics = TransactionMetrics::new(&provider.meter("transaction-metrics-test")); + Self { + provider, + exporter, + metrics, + attrs: vec![KeyValue::new("test.kind", "a")], + } + } + + /// `(outcome, count, sum)` per data point of `name`, sorted by outcome. + fn samples(&self, name: &str, outcome_key: &str) -> Vec<(String, u64, f64)> { + self.provider.force_flush().unwrap(); + let snapshots = self.exporter.get_finished_metrics().unwrap(); + let Some(metric) = snapshots + .iter() + .rev() + .flat_map(ResourceMetrics::scope_metrics) + .flat_map(ScopeMetrics::metrics) + .find(|metric| metric.name() == name) + else { + return Vec::new(); + }; + assert_eq!(metric.unit(), "s"); + let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = metric.data() else { + panic!("{name} must be an f64 histogram"); + }; + let mut samples: Vec<_> = histogram + .data_points() + .map(|point: &HistogramDataPoint| { + assert!( + point + .attributes() + .any(|kv| kv == &KeyValue::new("test.kind", "a")), + "the captured attributes must be kept" + ); + let outcome = point + .attributes() + .find(|kv| kv.key.as_str() == outcome_key) + .expect("every sample carries its outcome") + .value + .to_string(); + (outcome, point.count(), point.sum()) + }) + .collect(); + samples.sort_by(|left, right| left.0.cmp(&right.0)); + samples + } + + fn begins(&self) -> Vec<(String, u64)> { + self.samples(BEGIN_DURATION, BEGIN_OUTCOME) + .into_iter() + .map(|(outcome, count, _)| (outcome, count)) + .collect() + } + + fn transactions(&self) -> Vec<(String, u64)> { + self.samples(TRANSACTION_DURATION, TRANSACTION_OUTCOME) + .into_iter() + .map(|(outcome, count, _)| (outcome, count)) + .collect() + } + + fn started(&self) -> TransactionLifecycle { + let ((), lifecycle) = + block_on(self.metrics.observe_begin(&self.attrs, ready(Ok(())))).unwrap(); + lifecycle + } + } + + fn one(outcome: &str) -> Vec<(String, u64)> { + vec![(outcome.to_owned(), 1)] + } + + #[test] + fn successful_begin_starts_a_lifecycle_that_records_only_when_it_ends() { + let harness = Harness::new(); + let lifecycle = harness.started(); + assert_eq!(harness.begins(), one("ok")); + assert_eq!(harness.transactions(), Vec::new(), "still open"); + drop(lifecycle); + assert_eq!(harness.transactions(), one("dropped")); + } + + #[test] + fn failed_begin_records_one_error_and_no_transaction() { + let harness = Harness::new(); + let result = block_on( + harness + .metrics + .observe_begin::<(), _>(&harness.attrs, ready(Err(sqlx::Error::PoolClosed))), + ); + assert!(matches!(result, Err(sqlx::Error::PoolClosed))); + assert_eq!(harness.begins(), one("error")); + assert_eq!(harness.transactions(), Vec::new()); + } + + #[test] + fn dropped_begin_records_one_abandonment_and_no_transaction() { + let harness = Harness::new(); + block_on(async { + let mut begin = Box::pin( + harness + .metrics + .observe_begin::<(), _>(&harness.attrs, pending()), + ); + assert!(futures::poll!(&mut begin).is_pending()); + assert_eq!(harness.begins(), Vec::new(), "the attempt is still running"); + drop(begin); + }); + assert_eq!(harness.begins(), one("abandoned")); + assert_eq!(harness.transactions(), Vec::new()); + } + + #[test] + fn never_polled_begin_records_nothing() { + let harness = Harness::new(); + drop( + harness + .metrics + .observe_begin::<(), _>(&harness.attrs, pending()), + ); + assert_eq!(harness.begins(), Vec::new()); + assert_eq!(harness.transactions(), Vec::new()); + } + + #[test] + fn every_terminal_result_maps_to_one_outcome() { + for (termination, result, expected) in [ + (TransactionTermination::Commit, Ok(()), "commit_ok"), + ( + TransactionTermination::Commit, + Err(sqlx::Error::PoolClosed), + "commit_error", + ), + (TransactionTermination::Rollback, Ok(()), "rollback_ok"), + ( + TransactionTermination::Rollback, + Err(sqlx::Error::PoolClosed), + "rollback_error", + ), + ] { + let harness = Harness::new(); + let is_ok = result.is_ok(); + let finished = block_on(harness.started().finish(termination, ready(result))); + assert_eq!(finished.is_ok(), is_ok, "the result must pass through"); + assert_eq!(harness.transactions(), one(expected)); + } + } + + /// A cancelled terminal call must say "abandoned", never rollback or a failed commit. + #[test] + fn dropping_a_started_terminal_call_records_its_abandonment_once() { + for (termination, expected) in [ + (TransactionTermination::Commit, "commit_abandoned"), + (TransactionTermination::Rollback, "rollback_abandoned"), + ] { + let harness = Harness::new(); + let lifecycle = harness.started(); + block_on(async { + let mut finish = Box::pin(lifecycle.finish(termination, pending())); + assert!(futures::poll!(&mut finish).is_pending()); + assert_eq!(harness.transactions(), Vec::new(), "still in flight"); + drop(finish); + }); + assert_eq!(harness.transactions(), one(expected)); + } + } + + /// `COMMIT` was never sent, so this must not read as an abandoned commit. + #[test] + fn a_never_polled_terminal_call_counts_as_a_plain_drop() { + for termination in [ + TransactionTermination::Commit, + TransactionTermination::Rollback, + ] { + let harness = Harness::new(); + drop(harness.started().finish(termination, pending())); + assert_eq!(harness.transactions(), one("dropped")); + } + } + + #[tokio::test] + async fn begin_time_and_transaction_time_do_not_overlap() { + let harness = Harness::new(); + let ((), lifecycle) = harness + .metrics + .observe_begin(&harness.attrs, async { + tokio::time::sleep(Duration::from_millis(200)).await; + Ok(()) + }) + .await + .unwrap(); + lifecycle + .finish(TransactionTermination::Commit, ready(Ok(()))) + .await + .unwrap(); + let begin = harness.samples(BEGIN_DURATION, BEGIN_OUTCOME); + let transaction = harness.samples(TRANSACTION_DURATION, TRANSACTION_OUTCOME); + assert!(begin[0].2 >= 0.2, "BEGIN time: {}", begin[0].2); + assert!( + transaction[0].2 < 0.2, + "transaction time starts after BEGIN: {}", + transaction[0].2 + ); + } +} diff --git a/tests/lifecycle_metrics.rs b/tests/lifecycle_metrics.rs new file mode 100644 index 0000000..dd6be27 --- /dev/null +++ b/tests/lifecycle_metrics.rs @@ -0,0 +1,770 @@ +#![cfg(feature = "sqlite")] + +//! Connection-lifecycle contracts on the real wrappers: abandoned waits, workload attribution and +//! transaction outcomes. Cancellation of a terminal call at a known protocol point needs a real +//! server, so that half lives in `tests/postgres.rs`; the classification itself is covered +//! deterministically by the unit tests in `src/transaction_metrics.rs`. + +mod common; + +use std::collections::BTreeMap; +use std::task::Poll; +use std::time::Duration; + +use common::TestTelemetry; +use futures::TryStreamExt; +use opentelemetry_sdk::metrics::data::{AggregatedMetrics, Metric, MetricData, ResourceMetrics}; +use serial_test::serial; +use sqlx::Executor as _; +use sqlx_otel::{Pool, PoolAttribution, PoolBuilder, QueryAnnotateExt}; + +const KEY: &str = "test.workload.kind"; +const A: PoolAttribution = PoolAttribution::new(KEY, "a"); +const B: PoolAttribution = PoolAttribution::new(KEY, "b"); +const C: PoolAttribution = PoolAttribution::new(KEY, "c"); +/// Series with no emitted attribution: outside a scope or using a reserved key. +const UNSCOPED: &str = ""; + +const WAIT: &str = "db.client.connection.wait_time"; +const USE: &str = "db.client.connection.use_time"; +const TIMEOUTS: &str = "db.client.connection.timeouts"; +const PENDING: &str = "db.client.connection.pending_requests"; +const ABANDONED: &str = "sqlx_otel.connection.abandoned_wait.duration"; +const BEGIN: &str = "sqlx_otel.transaction.begin.duration"; +const BEGIN_OUTCOME: &str = "sqlx_otel.transaction.begin.outcome"; +const TRANSACTION: &str = "sqlx_otel.transaction.duration"; +const OUTCOME: &str = "sqlx_otel.transaction.outcome"; + +fn snapshot(tel: &TestTelemetry) -> Vec { + tel.reset(); + tel.metrics() +} + +fn find<'a>(metrics: &'a [ResourceMetrics], name: &str) -> Option<&'a Metric> { + metrics + .iter() + .flat_map(ResourceMetrics::scope_metrics) + .flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics) + .find(|metric| metric.name() == name) +} + +fn label<'a>( + mut attributes: impl Iterator, + key: &str, +) -> String { + attributes + .find(|kv| kv.key.as_str() == key) + .map_or_else(|| UNSCOPED.to_owned(), |kv| kv.value.to_string()) +} + +/// Sample count of a histogram per value of `key`. Empty when no sample was ever recorded. +fn histogram_by(tel: &TestTelemetry, name: &str, key: &str) -> BTreeMap { + let metrics = snapshot(tel); + let Some(metric) = find(&metrics, name) else { + return BTreeMap::new(); + }; + let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = metric.data() else { + panic!("{name} must be an f64 histogram"); + }; + let mut counts = BTreeMap::new(); + for point in histogram.data_points() { + *counts.entry(label(point.attributes(), key)).or_insert(0) += point.count(); + } + counts +} + +fn histogram_sum(tel: &TestTelemetry, name: &str) -> f64 { + let metrics = snapshot(tel); + let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = + find(&metrics, name).expect("histogram recorded").data() + else { + panic!("{name} must be an f64 histogram"); + }; + histogram + .data_points() + .map(opentelemetry_sdk::metrics::data::HistogramDataPoint::sum) + .sum() +} + +fn timeouts_by(tel: &TestTelemetry) -> BTreeMap { + let metrics = snapshot(tel); + let AggregatedMetrics::U64(MetricData::Sum(sum)) = find(&metrics, TIMEOUTS) + .expect("timeouts start at zero") + .data() + else { + panic!("timeouts must be a u64 counter"); + }; + sum.data_points() + .map(|point| (label(point.attributes(), KEY), point.value())) + .collect() +} + +fn pending(tel: &TestTelemetry) -> i64 { + let metrics = snapshot(tel); + let AggregatedMetrics::I64(MetricData::Sum(sum)) = find(&metrics, PENDING) + .expect("pending starts at zero") + .data() + else { + panic!("pending must be an i64 sum"); + }; + assert_eq!( + sum.data_points().count(), + 1, + "pending stays one pool-wide series" + ); + let point = sum.data_points().next().unwrap(); + assert_eq!( + label(point.attributes(), KEY), + UNSCOPED, + "pending must never carry attribution" + ); + point.value() +} + +fn expected(entries: [(&str, u64); N]) -> BTreeMap { + entries + .into_iter() + .map(|(key, count)| (key.to_owned(), count)) + .collect() +} + +async fn memory_pool() -> Pool { + PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await.unwrap()).build() +} + +/// A pool whose only connection the caller can hold, to make every later acquisition wait. +async fn single_connection_pool(acquire_timeout: Duration) -> Pool { + let raw = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .acquire_timeout(acquire_timeout) + .connect(":memory:") + .await + .unwrap(); + PoolBuilder::from(raw).build() +} + +const LONG: Duration = Duration::from_secs(30); + +// =========================================================================== +// Abandoned waits +// =========================================================================== + +#[tokio::test] +#[serial] +async fn a_never_polled_acquire_records_nothing() { + let tel = TestTelemetry::install(); + let pool = single_connection_pool(LONG).await; + let held = pool.acquire().await.unwrap(); + + drop(pool.acquire()); + drop(pool.begin()); + drop(pool.fetch_optional("SELECT 1")); + drop(pool.fetch("SELECT 1")); + + assert_eq!(histogram_by(&tel, ABANDONED, KEY), expected([])); + assert_eq!(pending(&tel), 0); + drop(held); +} + +#[tokio::test] +#[serial] +async fn a_dropped_wait_records_its_elapsed_time_once() { + let tel = TestTelemetry::install(); + let pool = single_connection_pool(LONG).await; + let held = pool.acquire().await.unwrap(); + let waits_before = histogram_by(&tel, WAIT, KEY); + + let mut waiting = Box::pin(pool.acquire()); + assert!(futures::poll!(&mut waiting).is_pending()); + tokio::time::sleep(Duration::from_millis(30)).await; + assert!(futures::poll!(&mut waiting).is_pending()); + assert_eq!( + histogram_by(&tel, ABANDONED, KEY), + expected([]), + "still waiting, whatever the number of polls" + ); + drop(waiting); + + assert_eq!( + histogram_by(&tel, ABANDONED, KEY), + expected([(UNSCOPED, 1)]) + ); + let waited = histogram_sum(&tel, ABANDONED); + assert!(waited >= 0.03, "abandoned wait must be measured: {waited}"); + assert_eq!(histogram_by(&tel, WAIT, KEY), waits_before); + assert_eq!(timeouts_by(&tel), expected([(UNSCOPED, 0)])); + assert_eq!(pending(&tel), 0); + drop(held); +} + +#[tokio::test] +#[serial] +async fn a_wait_that_gets_an_answer_is_not_abandoned() { + let tel = TestTelemetry::install(); + let pool = single_connection_pool(Duration::from_millis(50)).await; + drop(pool.acquire().await.unwrap()); + let held = pool.acquire().await.unwrap(); + assert!(matches!( + pool.acquire().await, + Err(sqlx::Error::PoolTimedOut) + )); + drop(held); + pool.close().await; + assert!(matches!(pool.acquire().await, Err(sqlx::Error::PoolClosed))); + + assert_eq!(histogram_by(&tel, ABANDONED, KEY), expected([])); + assert_eq!(histogram_by(&tel, WAIT, KEY), expected([(UNSCOPED, 4)])); + assert_eq!(timeouts_by(&tel), expected([(UNSCOPED, 1)])); +} + +// =========================================================================== +// Attribution +// =========================================================================== + +/// Every acquisition path, inside one scope. +#[tokio::test] +#[serial] +async fn attribution_covers_every_acquisition_path_without_changing_totals() { + let tel = TestTelemetry::install(); + let pool = memory_pool().await; + + A.scope(async { + drop(pool.acquire().await.unwrap()); + pool.fetch_optional("SELECT 1").await.unwrap(); + pool.fetch("SELECT 1") + .try_collect::>() + .await + .unwrap(); + pool.with_operation("SELECT", "test") + .fetch_optional("SELECT 1") + .await + .unwrap(); + sqlx::query("SELECT 1") + .with_operation("SELECT", "test") + .fetch_optional(&pool) + .await + .unwrap(); + pool.prepare_with(sqlx::SqlStr::from_static("SELECT 1"), &[]) + .await + .unwrap(); + pool.describe(sqlx::SqlStr::from_static("SELECT 1")) + .await + .unwrap(); + + let mut transaction = pool.begin().await.unwrap(); + sqlx::query("SELECT 1") + .execute(&mut transaction) + .await + .unwrap(); + sqlx::query("SELECT 2") + .execute(&mut transaction) + .await + .unwrap(); + transaction.commit().await.unwrap(); + }) + .await; + pool.fetch_optional("SELECT 1").await.unwrap(); + + // Eight acquisitions in the scope, the two transaction queries among no extra ones. + assert_eq!( + histogram_by(&tel, WAIT, KEY), + expected([("a", 8), (UNSCOPED, 1)]) + ); + assert_eq!( + histogram_by(&tel, USE, KEY), + expected([("a", 8), (UNSCOPED, 1)]) + ); + assert_eq!(histogram_by(&tel, BEGIN, KEY), expected([("a", 1)])); + assert_eq!(histogram_by(&tel, TRANSACTION, KEY), expected([("a", 1)])); + + // Query instruments and pool-wide state stay exactly as they were. + let metrics = snapshot(&tel); + for name in [ + "db.client.operation.duration", + "db.client.connection.max", + "db.client.connection.idle.max", + "db.client.connection.idle.min", + PENDING, + ] { + let metric = find(&metrics, name).unwrap_or_else(|| panic!("{name} not recorded")); + let mut keys = Vec::new(); + match metric.data() { + AggregatedMetrics::F64(MetricData::Histogram(data)) => { + for point in data.data_points() { + keys.extend(point.attributes().map(|kv| kv.key.to_string())); + } + } + AggregatedMetrics::I64(MetricData::Gauge(data)) => { + for point in data.data_points() { + keys.extend(point.attributes().map(|kv| kv.key.to_string())); + } + } + AggregatedMetrics::I64(MetricData::Sum(data)) => { + for point in data.data_points() { + keys.extend(point.attributes().map(|kv| kv.key.to_string())); + } + } + _ => panic!("unexpected aggregation for {name}"), + } + assert!( + !keys.iter().any(|key| key == KEY), + "{name} must not carry attribution" + ); + } +} + +/// A lease keeps what was captured at acquisition, wherever it ends. +#[tokio::test] +#[serial] +async fn the_lease_ignores_the_scope_it_ends_in() { + let tel = TestTelemetry::install(); + let pool = memory_pool().await; + + // Connection: acquired under A, dropped under B. + let connection = A.scope(pool.acquire()).await.unwrap(); + B.scope(async { drop(connection) }).await; + assert_eq!(histogram_by(&tel, USE, KEY), expected([("a", 1)])); + + // Connection: acquired under A, dropped by an unscoped task. + let connection = A.scope(pool.acquire()).await.unwrap(); + tokio::spawn(async move { drop(connection) }).await.unwrap(); + assert_eq!(histogram_by(&tel, USE, KEY), expected([("a", 2)])); + + // Transaction: begun under A, finished under B, and dropped under B. + let transaction = A.scope(pool.begin()).await.unwrap(); + B.scope(transaction.commit()).await.unwrap(); + let transaction = A.scope(pool.begin()).await.unwrap(); + B.scope(async { drop(transaction) }).await; + assert_eq!(histogram_by(&tel, TRANSACTION, KEY), expected([("a", 2)])); + assert_eq!(histogram_by(&tel, USE, KEY), expected([("a", 4)])); + assert_eq!(histogram_by(&tel, BEGIN, KEY), expected([("a", 2)])); + assert_eq!(histogram_by(&tel, WAIT, KEY), expected([("a", 4)])); +} + +/// A stream acquires on first poll, so that scope wins for the whole lease. +#[tokio::test] +#[serial] +async fn a_stream_keeps_the_scope_of_its_first_poll() { + let tel = TestTelemetry::install(); + let pool = memory_pool().await; + + // Constructed under B, which does not count: nothing was acquired yet. + let mut rows = B + .scope(async { pool.fetch("SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3") }) + .await; + A.scope(rows.try_next()).await.unwrap().unwrap(); + B.scope(rows.try_next()).await.unwrap().unwrap(); + rows.try_next().await.unwrap().unwrap(); + assert_eq!( + histogram_by(&tel, USE, KEY), + expected([]), + "lease still open" + ); + B.scope(async { drop(rows) }).await; + + assert_eq!(histogram_by(&tel, WAIT, KEY), expected([("a", 1)])); + assert_eq!(histogram_by(&tel, USE, KEY), expected([("a", 1)])); +} + +/// One acquisition of every kind per round, yielding across each await while holding the lease. +async fn workload(pool: Pool, rounds: u64) { + for _ in 0..rounds { + pool.fetch_optional("SELECT 1").await.unwrap(); + tokio::task::yield_now().await; + + let connection = pool.acquire().await.unwrap(); + tokio::task::yield_now().await; + drop(connection); + + let mut transaction = pool.begin().await.unwrap(); + tokio::task::yield_now().await; + sqlx::query("SELECT 1") + .execute(&mut transaction) + .await + .unwrap(); + tokio::task::yield_now().await; + transaction.commit().await.unwrap(); + } +} + +/// Reserved labels must not replace pool identity or transaction outcomes in exported telemetry. +#[tokio::test] +#[serial] +async fn reserved_attribution_cannot_spoof_pool_identity_or_transaction_outcomes() { + let tel = TestTelemetry::install(); + let raw = sqlx::SqlitePool::connect(":memory:").await.unwrap(); + let pool = PoolBuilder::from(raw).with_pool_name("canonical").build(); + + for key in [ + "db.client.connection.pool.name", + "network.transport", // Optional and absent on this pool: still reserved. + BEGIN_OUTCOME, + OUTCOME, + ] { + PoolAttribution::new(key, "spoofed") + .scope(async { pool.begin().await.unwrap().commit().await.unwrap() }) + .await; + } + + for instrument in [WAIT, USE, BEGIN, TRANSACTION] { + assert_eq!( + histogram_by(&tel, instrument, "db.client.connection.pool.name"), + expected([("canonical", 4)]) + ); + assert_eq!( + histogram_by(&tel, instrument, "network.transport"), + expected([(UNSCOPED, 4)]) + ); + } + assert_eq!( + histogram_by(&tel, BEGIN, BEGIN_OUTCOME), + expected([("ok", 4)]) + ); + assert_eq!( + histogram_by(&tel, TRANSACTION, OUTCOME), + expected([("commit_ok", 4)]) + ); +} + +/// Scopes interleaved on one task and spread over worker threads never see each other. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[serial] +async fn concurrent_scopes_never_leak_across_awaits_or_worker_threads() { + const ROUNDS: u64 = 40; + let tel = TestTelemetry::install(); + let raw = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(3) + .connect(":memory:") + .await + .unwrap(); + let pool = PoolBuilder::from(raw).build(); + + // A and an unscoped sibling interleave on one spawned task. + let same_task = { + let (first, second) = (pool.clone(), pool.clone()); + tokio::spawn(async move { + futures::join!(A.scope(workload(first, ROUNDS)), workload(second, ROUNDS)); + }) + }; + // B and C run as two tasks each, so workers keep migrating between the two scopes. + let mut tasks = vec![same_task]; + for attribution in [B, B, C, C] { + let pool = pool.clone(); + tasks.push(tokio::spawn(attribution.scope(workload(pool, ROUNDS)))); + } + for task in tasks { + task.await.unwrap(); + } + + let acquisitions = expected([ + ("a", 3 * ROUNDS), + ("b", 6 * ROUNDS), + ("c", 6 * ROUNDS), + (UNSCOPED, 3 * ROUNDS), + ]); + assert_eq!(histogram_by(&tel, WAIT, KEY), acquisitions); + assert_eq!(histogram_by(&tel, USE, KEY), acquisitions); + let transactions = expected([ + ("a", ROUNDS), + ("b", 2 * ROUNDS), + ("c", 2 * ROUNDS), + (UNSCOPED, ROUNDS), + ]); + assert_eq!(histogram_by(&tel, BEGIN, KEY), transactions); + assert_eq!(histogram_by(&tel, TRANSACTION, KEY), transactions); + assert_eq!(pending(&tel), 0); +} + +/// Cancellation and timeouts belong to the attempt, not to whoever dropped it or ran alongside it. +#[tokio::test] +#[serial] +async fn waits_are_cancelled_and_timed_out_under_their_own_scope() { + let tel = TestTelemetry::install(); + let pool = single_connection_pool(Duration::from_millis(50)).await; + let held = C.scope(pool.acquire()).await.unwrap(); + assert_eq!( + timeouts_by(&tel), + expected([(UNSCOPED, 0), ("c", 0)]), + "a workload's series exists at zero before its first timeout" + ); + + let mut cancelled = Box::pin(A.scope(pool.acquire())); + let mut timed_out = Box::pin(B.scope(pool.acquire())); + assert!(futures::poll!(&mut cancelled).is_pending()); + assert!(futures::poll!(&mut timed_out).is_pending()); + assert_eq!(pending(&tel), 2); + + // Dropped from inside the other attempt's scope. + B.scope(async { drop(cancelled) }).await; + assert_eq!(histogram_by(&tel, ABANDONED, KEY), expected([("a", 1)])); + assert_eq!(pending(&tel), 1); + + assert!(matches!(timed_out.await, Err(sqlx::Error::PoolTimedOut))); + assert_eq!( + timeouts_by(&tel), + expected([(UNSCOPED, 0), ("a", 0), ("b", 1), ("c", 0)]) + ); + assert_eq!( + histogram_by(&tel, WAIT, KEY), + expected([("b", 1), ("c", 1)]) + ); + assert_eq!(histogram_by(&tel, ABANDONED, KEY), expected([("a", 1)])); + assert_eq!(pending(&tel), 0); + drop(held); +} + +/// The documented limit: a spawned task starts outside the scope until told otherwise. +#[tokio::test] +#[serial] +async fn spawned_work_is_unscoped_unless_explicitly_inherited() { + let tel = TestTelemetry::install(); + let pool = memory_pool().await; + + let (plain, inherited) = A + .scope(async { + let plain = { + let pool = pool.clone(); + tokio::spawn(async move { drop(pool.acquire().await.unwrap()) }) + }; + let inherited = { + let pool = pool.clone(); + tokio::spawn(PoolAttribution::inherit(async move { + drop(pool.acquire().await.unwrap()); + })) + }; + (plain, inherited) + }) + .await; + plain.await.unwrap(); + inherited.await.unwrap(); + + assert_eq!( + histogram_by(&tel, WAIT, KEY), + expected([("a", 1), (UNSCOPED, 1)]) + ); + assert_eq!( + histogram_by(&tel, USE, KEY), + expected([("a", 1), (UNSCOPED, 1)]) + ); +} + +// =========================================================================== +// Transaction outcomes +// =========================================================================== + +#[tokio::test] +#[serial] +async fn each_way_of_ending_a_transaction_records_one_outcome_on_one_acquisition() { + let tel = TestTelemetry::install(); + let pool = memory_pool().await; + + let mut transaction = pool.begin().await.unwrap(); + assert_eq!( + histogram_by(&tel, TRANSACTION, OUTCOME), + expected([]), + "still open" + ); + for query in ["SELECT 1", "SELECT 2", "SELECT 3"] { + sqlx::query(query).execute(&mut transaction).await.unwrap(); + } + assert_eq!( + histogram_by(&tel, WAIT, KEY), + expected([(UNSCOPED, 1)]), + "queries inside a transaction must not acquire" + ); + transaction.commit().await.unwrap(); + assert_eq!( + histogram_by(&tel, TRANSACTION, OUTCOME), + expected([("commit_ok", 1)]) + ); + + pool.begin().await.unwrap().rollback().await.unwrap(); + assert_eq!( + histogram_by(&tel, TRANSACTION, OUTCOME), + expected([("commit_ok", 1), ("rollback_ok", 1)]) + ); + + drop(pool.begin().await.unwrap()); + assert_eq!( + histogram_by(&tel, TRANSACTION, OUTCOME), + expected([("commit_ok", 1), ("dropped", 1), ("rollback_ok", 1)]) + ); + + // Created but never polled: nothing was sent, so these are plain drops. + drop(pool.begin().await.unwrap().commit()); + drop(pool.begin().await.unwrap().rollback()); + assert_eq!( + histogram_by(&tel, TRANSACTION, OUTCOME), + expected([("commit_ok", 1), ("dropped", 3), ("rollback_ok", 1)]) + ); + + assert_eq!( + histogram_by(&tel, BEGIN, BEGIN_OUTCOME), + expected([("ok", 5)]) + ); + assert_eq!(histogram_by(&tel, WAIT, KEY), expected([(UNSCOPED, 5)])); + assert_eq!(histogram_by(&tel, USE, KEY), expected([(UNSCOPED, 5)])); +} + +/// A deferred constraint fails at `COMMIT` itself, so this is a real server rejection. +#[tokio::test] +#[serial] +async fn a_rejected_commit_is_an_error_outcome() { + let tel = TestTelemetry::install(); + let pool = single_connection_pool(LONG).await; + pool.execute( + "CREATE TABLE parent (id INTEGER PRIMARY KEY); + CREATE TABLE child ( + parent_id INTEGER REFERENCES parent (id) DEFERRABLE INITIALLY DEFERRED + );", + ) + .await + .unwrap(); + + let mut transaction = pool.begin().await.unwrap(); + sqlx::query("INSERT INTO child VALUES (1)") + .execute(&mut transaction) + .await + .unwrap(); + assert!(transaction.commit().await.is_err()); + + assert_eq!( + histogram_by(&tel, TRANSACTION, OUTCOME), + expected([("commit_error", 1)]) + ); + assert_eq!( + histogram_by(&tel, USE, KEY)[UNSCOPED], + 2, + "setup plus this lease" + ); +} + +/// `BEGIN` fails on a connection the pool believes idle but which is already in a transaction. +#[tokio::test] +#[serial] +async fn a_failed_begin_keeps_its_lease_accounting_and_starts_no_transaction() { + let tel = TestTelemetry::install(); + let raw = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .after_connect(|connection, _| { + Box::pin(async move { + connection.execute("BEGIN").await?; + Ok(()) + }) + }) + .connect(":memory:") + .await + .unwrap(); + let pool = PoolBuilder::from(raw).build(); + + assert!(pool.begin().await.is_err()); + + assert_eq!( + histogram_by(&tel, BEGIN, BEGIN_OUTCOME), + expected([("error", 1)]) + ); + assert_eq!(histogram_by(&tel, TRANSACTION, OUTCOME), expected([])); + assert_eq!(histogram_by(&tel, WAIT, KEY), expected([(UNSCOPED, 1)])); + assert_eq!(histogram_by(&tel, USE, KEY), expected([(UNSCOPED, 1)])); +} + +/// A wait abandoned before `BEGIN` never reached the transaction instruments. +#[tokio::test] +#[serial] +async fn a_begin_dropped_while_waiting_for_a_connection_is_only_an_abandoned_wait() { + let tel = TestTelemetry::install(); + let pool = single_connection_pool(LONG).await; + let held = pool.acquire().await.unwrap(); + + let mut begin = Box::pin(pool.begin()); + assert!(futures::poll!(&mut begin).is_pending()); + drop(begin); + + assert_eq!( + histogram_by(&tel, ABANDONED, KEY), + expected([(UNSCOPED, 1)]) + ); + assert_eq!(histogram_by(&tel, BEGIN, BEGIN_OUTCOME), expected([])); + assert_eq!(histogram_by(&tel, TRANSACTION, OUTCOME), expected([])); + drop(held); +} + +/// Through the real wrappers. The `SQLite` worker thread could answer before the first poll +/// returns, so either result is accepted as long as it matches what the caller saw; the +/// deterministic version of this lives in the `src/transaction_metrics.rs` unit tests. +#[tokio::test] +#[serial] +async fn a_dropped_terminal_call_reports_what_the_caller_could_see_exactly_once() { + let tel = TestTelemetry::install(); + let pool = memory_pool().await; + + let mut commit = Box::pin(pool.begin().await.unwrap().commit()); + let commit_outcome = match futures::poll!(&mut commit) { + Poll::Pending => "commit_abandoned", + Poll::Ready(result) => { + result.unwrap(); + "commit_ok" + } + }; + drop(commit); + assert_eq!( + histogram_by(&tel, TRANSACTION, OUTCOME), + expected([(commit_outcome, 1)]) + ); + + let mut rollback = Box::pin(pool.begin().await.unwrap().rollback()); + let rollback_outcome = match futures::poll!(&mut rollback) { + Poll::Pending => "rollback_abandoned", + Poll::Ready(result) => { + result.unwrap(); + "rollback_ok" + } + }; + drop(rollback); + assert_eq!( + histogram_by(&tel, TRANSACTION, OUTCOME), + expected([(commit_outcome, 1), (rollback_outcome, 1)]) + ); + assert_eq!(histogram_by(&tel, USE, KEY), expected([(UNSCOPED, 2)])); +} + +// =========================================================================== +// Shape +// =========================================================================== + +/// Hardcoded on purpose: a shared constant would make the assertion tautological. +#[tokio::test] +#[serial] +async fn new_instruments_export_their_unit_and_boundaries() { + const LATENCY: &[f64] = &[ + 0.0001, 0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, + 10.0, 30.0, 60.0, + ]; + const LONG_RUNNING: &[f64] = &[ + 0.0001, 0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, + 10.0, 30.0, 60.0, 120.0, 300.0, 600.0, 1800.0, 3600.0, + ]; + let tel = TestTelemetry::install(); + let pool = single_connection_pool(LONG).await; + pool.begin().await.unwrap().commit().await.unwrap(); + let held = pool.acquire().await.unwrap(); + let mut waiting = Box::pin(pool.acquire()); + assert!(futures::poll!(&mut waiting).is_pending()); + drop(waiting); + drop(held); + + let metrics = snapshot(&tel); + for (name, bounds) in [ + (ABANDONED, LATENCY), + (BEGIN, LATENCY), + (TRANSACTION, LONG_RUNNING), + ] { + let metric = find(&metrics, name).expect("instrument recorded"); + assert_eq!(metric.unit(), "s", "{name}"); + let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = metric.data() else { + panic!("{name} must be an f64 histogram"); + }; + let point = histogram.data_points().next().expect("one data point"); + assert_eq!(point.bounds().collect::>(), bounds, "{name}"); + } +} diff --git a/tests/pool_metrics.rs b/tests/pool_metrics.rs index 663ef19..9678de3 100644 --- a/tests/pool_metrics.rs +++ b/tests/pool_metrics.rs @@ -147,8 +147,17 @@ async fn transactions_keep_one_acquisition_and_measure_the_complete_lease() { histogram_count(&tel, "db.client.connection.use_time"), finish + 1 ); + assert_eq!( + histogram_count(&tel, "sqlx_otel.transaction.duration"), + finish + 1, + "transaction time is distinct from the lease, and recorded once" + ); } assert_eq!(pending_count(&tel), 0); + assert_eq!( + histogram_count(&tel, "sqlx_otel.transaction.begin.duration"), + 3 + ); } async fn waiting_operation( @@ -205,6 +214,16 @@ async fn all_acquisition_paths_clear_pending_on_cancel_and_report_real_timeouts( u64::from(path), "caller cancellation is not a pool timeout" ); + assert_eq!( + histogram_count(&tel, "sqlx_otel.connection.abandoned_wait.duration"), + u64::from(path) + 1, + "each abandoned wait is measured exactly once" + ); + assert_eq!( + histogram_count(&tel, "db.client.connection.wait_time"), + u64::from(path) + 1, + "an abandoned wait is not a completed wait" + ); assert!(matches!( waiting_operation(&pool, path).await, Err(sqlx::Error::PoolTimedOut) @@ -215,6 +234,11 @@ async fn all_acquisition_paths_clear_pending_on_cancel_and_report_real_timeouts( histogram_count(&tel, "db.client.connection.wait_time"), u64::from(path) + 2 ); + assert_eq!( + histogram_count(&tel, "sqlx_otel.connection.abandoned_wait.duration"), + u64::from(path) + 1, + "a timed-out wait got an answer, so it was not abandoned" + ); } drop(held); } diff --git a/tests/postgres.rs b/tests/postgres.rs index d848008..448ca3d 100644 --- a/tests/postgres.rs +++ b/tests/postgres.rs @@ -1191,3 +1191,397 @@ async fn query_scalar_macro_fetch_all_with_annotations_via_pool() { async fn executor_side_query_side_parity_via_pool() { test_executor_side_query_side_parity_via_pool!(test_pool().await, common::POSTGRES_DIALECT); } + +// =========================================================================== +// Connection lifecycle against a real server +// =========================================================================== + +mod connection_lifecycle { + use std::collections::BTreeMap; + + use opentelemetry_sdk::metrics::data::{AggregatedMetrics, MetricData}; + use sqlx::Executor as _; + use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; + use sqlx_otel::PoolAttribution; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + use tokio::sync::watch; + use tokio::task::JoinHandle; + use tokio::time::{sleep, timeout}; + + use super::{Duration, PoolBuilder, common, serial, shared_container}; + + const KEY: &str = "test.workload.kind"; + const REFRESH: PoolAttribution = PoolAttribution::new(KEY, "refresh"); + const EVALUATE: PoolAttribution = PoolAttribution::new(KEY, "evaluate"); + const OUTCOME: &str = "sqlx_otel.transaction.outcome"; + const BEGIN_OUTCOME: &str = "sqlx_otel.transaction.begin.outcome"; + const WAIT: &str = "db.client.connection.wait_time"; + const USE: &str = "db.client.connection.use_time"; + const ABANDONED: &str = "sqlx_otel.connection.abandoned_wait.duration"; + const BEGIN: &str = "sqlx_otel.transaction.begin.duration"; + const TRANSACTION: &str = "sqlx_otel.transaction.duration"; + /// Long enough for a statement to reach the server through the gate, with its answer withheld. + const HELD_STATEMENT: Duration = Duration::from_millis(200); + + fn counts(entries: [(&str, u64); N]) -> BTreeMap { + entries + .into_iter() + .map(|(value, count)| (value.to_owned(), count)) + .collect() + } + + /// Sample count of a histogram per value of `key`, with `` for series without it. + fn series(tel: &common::TestTelemetry, name: &str, key: &str) -> BTreeMap { + tel.reset(); + let metrics = tel.metrics(); + let mut series = BTreeMap::new(); + let Some(metric) = metrics + .iter() + .flat_map(opentelemetry_sdk::metrics::data::ResourceMetrics::scope_metrics) + .flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics) + .find(|metric| metric.name() == name) + else { + return series; + }; + let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = metric.data() else { + panic!("{name} must be an f64 histogram"); + }; + for point in histogram.data_points() { + let value = point + .attributes() + .find(|attribute| attribute.key.as_str() == key) + .map_or_else( + || "".to_owned(), + |attribute| attribute.value.to_string(), + ); + *series.entry(value).or_insert(0) += point.count(); + } + series + } + + /// Forwards `PostgreSQL` traffic and can withhold the server's answers. A statement then sits + /// at a known point, sent to the server with its result unseen by the client, independent of + /// timing. + struct ResponseGate { + port: u16, + held: watch::Sender, + accept: JoinHandle<()>, + } + + impl ResponseGate { + async fn start(options: &PgConnectOptions) -> Self { + let upstream = (options.get_host().to_owned(), options.get_port()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let (held, gate) = watch::channel(false); + let accept = tokio::spawn(async move { + while let Ok((client, _)) = listener.accept().await { + tokio::spawn(forward(client, upstream.clone(), gate.clone())); + } + }); + Self { port, held, accept } + } + + fn hold(&self) { + self.held.send_replace(true); + } + + fn release(&self) { + self.held.send_replace(false); + } + } + + impl Drop for ResponseGate { + fn drop(&mut self) { + self.accept.abort(); + } + } + + async fn forward(client: TcpStream, upstream: (String, u16), mut held: watch::Receiver) { + let Ok(server) = TcpStream::connect(upstream).await else { + return; + }; + let (mut client_read, mut client_write) = client.into_split(); + let (mut server_read, mut server_write) = server.into_split(); + let requests = tokio::io::copy(&mut client_read, &mut server_write); + let responses = async { + let mut chunk = [0_u8; 8192]; + loop { + let length = server_read.read(&mut chunk).await?; + if length == 0 || held.wait_for(|held| !*held).await.is_err() { + return Ok::<_, std::io::Error>(()); + } + client_write.write_all(&chunk[..length]).await?; + } + }; + tokio::select! { + _ = requests => {} + _ = responses => {} + } + } + + async fn wait_until_idle(pool: &sqlx::PgPool) { + timeout(Duration::from_secs(5), async { + while pool.num_idle() != 1 { + sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("the connection must return to the pool"); + } + + async fn terminate_backend(admin: &sqlx::PgPool, backend: i32) { + let terminated: bool = sqlx::query_scalar("SELECT pg_terminate_backend($1)") + .bind(backend) + .fetch_one(admin) + .await + .unwrap(); + assert!(terminated); + timeout(Duration::from_secs(5), async { + loop { + let alive: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE pid = $1)", + ) + .bind(backend) + .fetch_one(admin) + .await + .unwrap(); + if !alive { + return; + } + sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("the backend must exit"); + } + + /// Two workloads share one connection on one task: the second waits while the first holds a + /// transaction, across awaits, and the first ends inside the second's scope. + #[tokio::test] + #[serial] + async fn attribution_stays_with_each_lease() { + let tel = common::TestTelemetry::install(); + let raw = PgPoolOptions::new() + .max_connections(1) + .acquire_timeout(Duration::from_secs(5)) + .connect(&shared_container().await.url) + .await + .unwrap(); + let pool = PoolBuilder::from(raw).build(); + + let mut transaction = REFRESH.scope(pool.begin()).await.unwrap(); + let mut waiting = Box::pin(EVALUATE.scope(pool.fetch_one("SELECT 1"))); + assert!(futures::poll!(&mut waiting).is_pending()); + let mut abandoned = Box::pin(EVALUATE.scope(pool.acquire())); + assert!(futures::poll!(&mut abandoned).is_pending()); + REFRESH.scope(async { drop(abandoned) }).await; + assert_eq!(series(&tel, ABANDONED, KEY), counts([("evaluate", 1)])); + + sqlx::query("SELECT 2") + .execute(&mut transaction) + .await + .unwrap(); + EVALUATE.scope(transaction.commit()).await.unwrap(); + waiting.await.unwrap(); + + let leases = counts([("evaluate", 1), ("refresh", 1)]); + assert_eq!(series(&tel, WAIT, KEY), leases); + assert_eq!(series(&tel, USE, KEY), leases); + assert_eq!(series(&tel, BEGIN, KEY), counts([("refresh", 1)])); + assert_eq!(series(&tel, TRANSACTION, KEY), counts([("refresh", 1)])); + assert_eq!( + series(&tel, TRANSACTION, OUTCOME), + counts([("commit_ok", 1)]) + ); + pool.close().await; + } + + /// Each terminal path against a real server, including the ones where the client cannot know + /// the server's result. The row checks show the labels must not be read as server facts. + #[tokio::test] + #[serial] + #[allow(clippy::too_many_lines)] + async fn transaction_outcomes_report_only_what_the_client_saw() { + let tel = common::TestTelemetry::install(); + let url = &shared_container().await.url; + let options: PgConnectOptions = url.parse().unwrap(); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect_with(options.clone()) + .await + .unwrap(); + admin + .execute( + "DROP TABLE IF EXISTS lifecycle_row, lifecycle_child, lifecycle_parent; + CREATE TABLE lifecycle_row (id integer PRIMARY KEY); + CREATE TABLE lifecycle_parent (id integer PRIMARY KEY); + CREATE TABLE lifecycle_child ( + parent_id integer REFERENCES lifecycle_parent (id) + DEFERRABLE INITIALLY DEFERRED + );", + ) + .await + .unwrap(); + let row_exists = async |id: i32| -> bool { + sqlx::query_scalar("SELECT EXISTS (SELECT 1 FROM lifecycle_row WHERE id = $1)") + .bind(id) + .fetch_one(&admin) + .await + .unwrap() + }; + + let gate = ResponseGate::start(&options).await; + // No ping on acquire: the next statement on an idle connection is the one under test. + let raw = PgPoolOptions::new() + .max_connections(1) + .test_before_acquire(false) + .acquire_timeout(Duration::from_secs(5)) + .connect_with(options.clone().host("127.0.0.1").port(gate.port)) + .await + .unwrap(); + let pool = PoolBuilder::from(raw.clone()).build(); + let insert = |id: i32| sqlx::query("INSERT INTO lifecycle_row VALUES ($1)").bind(id); + + // Acknowledged commit, with several queries on one acquisition. + let mut transaction = pool.begin().await.unwrap(); + insert(1).execute(&mut transaction).await.unwrap(); + sqlx::query("SELECT 1") + .execute(&mut transaction) + .await + .unwrap(); + assert_eq!(series(&tel, WAIT, KEY), counts([("", 1)])); + transaction.commit().await.unwrap(); + assert!(row_exists(1).await); + assert_eq!( + series(&tel, TRANSACTION, OUTCOME), + counts([("commit_ok", 1)]) + ); + + // COMMIT of an already failed transaction: PostgreSQL answers with a ROLLBACK tag and no + // error, so the client saw success and nothing committed. + let mut transaction = pool.begin().await.unwrap(); + insert(2).execute(&mut transaction).await.unwrap(); + assert!( + sqlx::query("SELECT 1 / 0") + .execute(&mut transaction) + .await + .is_err() + ); + transaction.commit().await.unwrap(); + assert!(!row_exists(2).await); + assert_eq!( + series(&tel, TRANSACTION, OUTCOME), + counts([("commit_ok", 2)]) + ); + + // A deferred constraint rejects the COMMIT statement itself. + let mut transaction = pool.begin().await.unwrap(); + sqlx::query("INSERT INTO lifecycle_child VALUES (1)") + .execute(&mut transaction) + .await + .unwrap(); + assert!(transaction.commit().await.is_err()); + assert_eq!( + series(&tel, TRANSACTION, OUTCOME), + counts([("commit_error", 1), ("commit_ok", 2)]) + ); + + // Acknowledged rollback. + let mut transaction = pool.begin().await.unwrap(); + insert(3).execute(&mut transaction).await.unwrap(); + transaction.rollback().await.unwrap(); + assert!(!row_exists(3).await); + + // Plain drop: recorded at once, while the rollback is only queued for the return. + let mut transaction = pool.begin().await.unwrap(); + insert(4).execute(&mut transaction).await.unwrap(); + gate.hold(); + drop(transaction); + assert_eq!( + series(&tel, TRANSACTION, OUTCOME), + counts([ + ("commit_error", 1), + ("commit_ok", 2), + ("dropped", 1), + ("rollback_ok", 1), + ]) + ); + gate.release(); + wait_until_idle(&raw).await; + assert!(!row_exists(4).await); + + // COMMIT reaches the server and its answer is withheld until the caller gives up. The + // client recorded an abandonment, and the row committed anyway. + let mut transaction = pool.begin().await.unwrap(); + insert(5).execute(&mut transaction).await.unwrap(); + gate.hold(); + assert!(timeout(HELD_STATEMENT, transaction.commit()).await.is_err()); + assert_eq!(series(&tel, TRANSACTION, OUTCOME)["commit_abandoned"], 1); + gate.release(); + wait_until_idle(&raw).await; + assert!(row_exists(5).await); + + // The same for ROLLBACK. + let mut transaction = pool.begin().await.unwrap(); + insert(6).execute(&mut transaction).await.unwrap(); + gate.hold(); + assert!( + timeout(HELD_STATEMENT, transaction.rollback()) + .await + .is_err() + ); + assert_eq!(series(&tel, TRANSACTION, OUTCOME)["rollback_abandoned"], 1); + gate.release(); + wait_until_idle(&raw).await; + assert!(!row_exists(6).await); + + // BEGIN abandoned on a held connection: a completed acquisition and no transaction. + let transactions_before = series(&tel, TRANSACTION, OUTCOME); + gate.hold(); + assert!(timeout(HELD_STATEMENT, pool.begin()).await.is_err()); + assert_eq!(series(&tel, BEGIN, BEGIN_OUTCOME)["abandoned"], 1); + gate.release(); + wait_until_idle(&raw).await; + assert_eq!(series(&tel, TRANSACTION, OUTCOME), transactions_before); + + // BEGIN failed: the backend is gone, and with no ping the pool hands out its dead connection. + let backend: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&pool) + .await + .unwrap(); + wait_until_idle(&raw).await; + terminate_backend(&admin, backend).await; + assert!(pool.begin().await.is_err()); + assert_eq!(series(&tel, BEGIN, BEGIN_OUTCOME)["error"], 1); + assert_eq!(series(&tel, TRANSACTION, OUTCOME), transactions_before); + + // Exactly once across everything above, on exactly one acquisition each. + assert_eq!( + series(&tel, BEGIN, BEGIN_OUTCOME), + counts([("abandoned", 1), ("error", 1), ("ok", 7)]) + ); + assert_eq!( + series(&tel, TRANSACTION, OUTCOME), + counts([ + ("commit_abandoned", 1), + ("commit_error", 1), + ("commit_ok", 2), + ("dropped", 1), + ("rollback_abandoned", 1), + ("rollback_ok", 1), + ]) + ); + // Seven transactions, the failed and abandoned BEGINs, and the backend-id query. + assert_eq!(series(&tel, WAIT, KEY), counts([("", 10)])); + assert_eq!(series(&tel, USE, KEY), counts([("", 10)])); + + pool.close().await; + admin + .execute("DROP TABLE lifecycle_row, lifecycle_child, lifecycle_parent") + .await + .unwrap(); + admin.close().await; + } +}