diff --git a/backend/src/db_connect.rs b/backend/src/db_connect.rs index d4dee19f82..b18f675e49 100644 --- a/backend/src/db_connect.rs +++ b/backend/src/db_connect.rs @@ -153,6 +153,15 @@ pub async fn connect( pool_options = pool_options.idle_timeout(Duration::from_secs(60)); } pool_options + // Clears transaction state sqlx is not tracking, so a connection whose session was + // left inside a transaction is not handed to the next borrower. See + // `windmill_common::db::connection_reset` for how such a session comes about and why + // this only runs once one has been observed. + .after_release(|conn, _| { + Box::pin(windmill_common::db::connection_reset::reset_on_release( + conn, + )) + }) .after_connect(move |conn, _| { if worker_mode { Box::pin(async move { diff --git a/backend/windmill-common/src/db.rs b/backend/windmill-common/src/db.rs index 0048abf438..6f62038c1c 100644 --- a/backend/windmill-common/src/db.rs +++ b/backend/windmill-common/src/db.rs @@ -290,3 +290,76 @@ impl<'b> DbExecutor<'b> for &'b mut PgConnection { &mut **self } } + +/// Guards the pool against a Postgres session left inside a transaction sqlx is not +/// tracking. +/// +/// sqlx only queues its rollback-on-drop once `begin()` has returned, so a future +/// cancelled while `BEGIN` is still in flight — a disconnecting API client, a +/// `tokio::time::timeout`, an aborted task — hands the connection back with the +/// transaction still open server-side. sqlx's on-release `ping` is a bare +/// `wait_until_ready` that reports such a connection as healthy, so it is reused: the +/// next borrower's statements silently run inside that transaction and hold their locks +/// for as long as it lives, and the first error turns the session into +/// `idle in transaction (aborted)`, after which every unrelated query on it fails with +/// `25P02` until `max_lifetime` recycles it half an hour later. +/// +/// Rolling back on every release would add a round trip to every query, which measured at +/// about a third of the throughput on small ones, so the reset is armed only once Postgres +/// has reported a state that proves a connection is carrying such leftover state. +pub mod connection_reset { + use std::sync::atomic::{AtomicU64, Ordering}; + use std::sync::LazyLock; + use std::time::{Duration, Instant}; + + static PROCESS_START: LazyLock = LazyLock::new(Instant::now); + /// Milliseconds since `PROCESS_START` until which releases roll back; 0 = never armed. + static RESET_UNTIL_MS: AtomicU64 = AtomicU64::new(0); + + /// Long enough to outlast the connections that were already checked out when the + /// first poisoned one surfaced; one release cycle clears the pool. + const RESET_WINDOW: Duration = Duration::from_secs(60); + + fn now_ms() -> u64 { + PROCESS_START.elapsed().as_millis() as u64 + } + + fn arm() { + RESET_UNTIL_MS.fetch_max( + now_ms() + RESET_WINDOW.as_millis() as u64, + Ordering::Relaxed, + ); + } + + /// Reading the clock is skipped entirely while the pool has never been armed, which + /// is the steady state. + fn armed() -> bool { + let until = RESET_UNTIL_MS.load(Ordering::Relaxed); + until != 0 && until > now_ms() + } + + /// Body of the pool's `after_release` hook. `ROLLBACK` ends both a leaked-open and an + /// aborted transaction, and is a no-op warning on a session that has neither. Reporting + /// the failure makes sqlx discard the connection, which is also what a session we could + /// not clean deserves. + pub async fn reset_on_release(conn: &mut sqlx::PgConnection) -> Result { + if armed() { + use sqlx::Executor; + conn.execute("ROLLBACK").await?; + } + Ok(true) + } + + /// Arms the reset when an error proves a session is stuck in an aborted transaction. + /// Called from `Error`'s `From`, so it sees every query that reports one. + pub(crate) fn note_sqlx_error(err: &sqlx::Error) { + let sqlx::Error::Database(db_err) = err else { + return; + }; + // 25P02 `in_failed_sql_transaction`: the connection this query ran on is sitting + // in a failed transaction block that nothing in sqlx will ever roll back. + if db_err.code().as_deref() == Some("25P02") { + arm(); + } + } +} diff --git a/backend/windmill-common/src/error.rs b/backend/windmill-common/src/error.rs index 0d5ebdc2c3..7ee8ee21cf 100644 --- a/backend/windmill-common/src/error.rs +++ b/backend/windmill-common/src/error.rs @@ -153,6 +153,7 @@ impl From for Error { impl From for Error { #[track_caller] fn from(e: sqlx::Error) -> Self { + crate::db::connection_reset::note_sqlx_error(&e); Self::SqlErr { error: e, location: prettify_location(std::panic::Location::caller()) } } } diff --git a/backend/windmill-common/tests/connection_reset.rs b/backend/windmill-common/tests/connection_reset.rs new file mode 100644 index 0000000000..0e5d8b6e95 --- /dev/null +++ b/backend/windmill-common/tests/connection_reset.rs @@ -0,0 +1,45 @@ +//! A pooled connection whose session is left inside a transaction sqlx is not tracking stays +//! broken for every later borrower — `Rollback::drop` only fires while sqlx's own transaction +//! depth is non-zero, and the on-release `ping` is a bare `wait_until_ready` that reports such +//! a session as healthy. Before the reset hook this cost half an hour of unrelated `25P02` +//! failures across the whole process, until `max_lifetime` recycled the connection. + +use sqlx::{Executor, Pool, Postgres}; +use windmill_common::db::connection_reset; +use windmill_common::error::Error; + +#[sqlx::test(migrations = "../migrations")] +async fn poisoned_connection_is_reset_before_being_handed_out_again(db: Pool) { + // One connection, so every query below lands on the same session. + let pool = sqlx::postgres::PgPoolOptions::new() + .max_connections(1) + .min_connections(0) + .after_release(|conn, _| Box::pin(connection_reset::reset_on_release(conn))) + .connect_with((*db.connect_options()).clone()) + .await + .expect("failed to build pool"); + + // A batch that fails between BEGIN and COMMIT never reaches the COMMIT, leaving the + // session `idle in transaction (aborted)` while sqlx still believes it is not in a + // transaction — the same state a future cancelled during `begin()` leaves behind. + pool.execute("BEGIN; SELECT 1/0; COMMIT;") + .await + .expect_err("the batch must fail"); + + let err = sqlx::query_scalar::<_, i32>("SELECT 1") + .fetch_one(&pool) + .await + .expect_err("the next borrower inherits the aborted transaction"); + assert_eq!( + err.as_database_error().and_then(|e| e.code()).as_deref(), + Some("25P02"), + ); + // Converting the error is what arms the reset in production, where every `?` on a + // sqlx result goes through this same `From`. + let _ = Error::from(err); + + sqlx::query_scalar::<_, i32>("SELECT 1") + .fetch_one(&pool) + .await + .expect("pool must hand out a usable connection again"); +}