diff --git a/backend/tests/worker.rs b/backend/tests/worker.rs index 938763f21e..6966b40e1f 100644 --- a/backend/tests/worker.rs +++ b/backend/tests/worker.rs @@ -1758,8 +1758,9 @@ async fn test_postgresql_cached_connection_released_for_other_key( let server = ApiServer::start(db.clone()).await?; let port = server.addr.port(); - // Two cache keys for the same role: only the sslmode differs. - let run = |sslmode: &str| { + // Keys for the same role: one pool scope when only the sslmode differs, two + // when the database does, which the role's limit still counts together. + let run = |dbname: &str, sslmode: &str| { RunJob::from(JobPayload::Code(RawCode { hash: None, content: "SELECT 1 as n;".into(), @@ -1778,7 +1779,7 @@ async fn test_postgresql_cached_connection_released_for_other_key( })) .arg( "database", - json!({"host": "localhost", "port": 5432, "dbname": "windmill", + json!({"host": "localhost", "port": 5432, "dbname": dbname, "user": "wm_pg_cache_one_conn", "password": "changeme", "sslmode": sslmode}), ) .run_until_complete(&db, false, port) @@ -1787,16 +1788,21 @@ async fn test_postgresql_cached_connection_released_for_other_key( // The evicted connection's backend exits asynchronously, so a fresh // connection can briefly still count it. Retrying absorbs that; without the // eviction the cached connection stays open for 60s and every retry fails. - for sslmode in ["disable", "prefer", "disable"] { + for (dbname, sslmode) in [ + ("windmill", "disable"), + ("windmill", "prefer"), + ("postgres", "disable"), + ("windmill", "disable"), + ] { let mut result = json!(null); for _ in 0..5 { - result = run(sslmode).await.json_result().unwrap(); + result = run(dbname, sslmode).await.json_result().unwrap(); if result == json!([{"n": 1}]) { break; } tokio::time::sleep(std::time::Duration::from_millis(200)).await; } - assert_eq!(result, json!([{"n": 1}]), "sslmode={sslmode}"); + assert_eq!(result, json!([{"n": 1}]), "{dbname} sslmode={sslmode}"); } clear_pg_cache().await; @@ -1869,6 +1875,61 @@ async fn test_postgresql_open_transaction_not_reused(db: Pool) -> anyh Ok(()) } +/// A worker alternating between two databases keeps a connection for each. +#[sqlx::test(fixtures("base"))] +#[serial(pg_cache)] +async fn test_postgresql_cache_keeps_a_connection_per_database( + db: Pool, +) -> anyhow::Result<()> { + use std::sync::atomic::Ordering; + use windmill_worker::pg_executor::{clear_pg_cache, CACHE_HITS}; + + initialize_tracing().await; + clear_pg_cache().await; + + let server = ApiServer::start(db.clone()).await?; + let port = server.addr.port(); + + let run = |dbname: &str| { + RunJob::from(JobPayload::Code(RawCode { + hash: None, + content: "SELECT current_database() as d;".into(), + path: None, + lock: None, + language: ScriptLang::Postgresql, + cache_ttl: None, + cache_ignore_s3_path: None, + dedicated_worker: None, + concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default( + ) + .into(), + debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(), + modules: None, + tag: None, + })) + .arg( + "database", + json!({"host": "localhost", "port": 5432, "dbname": dbname, + "user": "postgres", "password": "changeme"}), + ) + .run_until_complete(&db, false, port) + }; + + let hits_before = CACHE_HITS.load(Ordering::Relaxed); + for dbname in ["windmill", "postgres", "windmill", "postgres"] { + let result = run(dbname).await.json_result().unwrap(); + assert_eq!(result, json!([{"d": dbname}])); + } + let hits = CACHE_HITS.load(Ordering::Relaxed) - hits_before; + assert_eq!( + hits, 2, + "the last two jobs should reuse a cached connection" + ); + + clear_pg_cache().await; + Ok(()) +} + /// Runs multiple PG jobs through a SINGLE worker (like production) to verify /// that SET ROLE / search_path changes do not leak across jobs. #[sqlx::test(fixtures("base"))] diff --git a/backend/windmill-worker/src/pg_executor.rs b/backend/windmill-worker/src/pg_executor.rs index 3222d7ca3e..b13cfa4535 100644 --- a/backend/windmill-worker/src/pg_executor.rs +++ b/backend/windmill-worker/src/pg_executor.rs @@ -1,7 +1,6 @@ use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; -use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use anyhow::Context; use base64::{engine, Engine as _}; @@ -13,7 +12,6 @@ use rust_decimal::{prelude::FromPrimitive, Decimal}; use serde_json::value::RawValue; use serde_json::Map; use serde_json::Value; -use tokio::sync::RwLock; use tokio_postgres::Client; use tokio_postgres::{types::ToSql, Row}; use tokio_postgres::{ @@ -54,11 +52,249 @@ use windmill_common::client::AuthedClient; use windmill_types::s3::S3Object; lazy_static! { - pub static ref CONNECTION_COUNTER: Arc>> = - Arc::new(RwLock::new(HashMap::new())); pub static ref CACHE_HITS: AtomicU64 = AtomicU64::new(0); } +/// Idle connections kept per cache key, checked out by a job for the duration +/// of its query and checked back in afterwards. +/// +/// Invariant: the process never keeps an idle connection in a pool scope +/// ([`pg_pool_scope`]) while one of its jobs runs on a connection of its own in +/// that scope. Behind a session-mode pooler (PgBouncer, RDS Proxy pinning, …) +/// an idle connection holds a server slot, and the job on its own connection +/// waits for that slot with no timeout, forever if other jobs keep reusing the +/// idle connection. +static PG_CONNECTION_CACHE: std::sync::Mutex = + std::sync::Mutex::new(PgConnectionCache { idle: Vec::new(), own: Vec::new(), next_own_id: 0 }); +static PG_CONNECTION_CACHE_SWEEPER: std::sync::Once = std::sync::Once::new(); +const PG_CONNECTION_IDLE_TIMEOUT: Duration = Duration::from_secs(60); +/// Across all keys. Native mode runs 8 jobs at once in one process. +const PG_MAX_IDLE_CONNECTIONS: usize = 8; +/// A pooler can also share slots across pool scopes: two host names for one +/// pooler, or a per-database or per-user limit. A job on its own connection for +/// longer than this may be waiting on such a slot, so no idle connection is +/// kept, in any scope, until it finishes. +const PG_OWN_CONNECTION_STALL: Duration = Duration::from_secs(5); + +struct PgConnectionCache { + /// Least recently used first. + idle: Vec, + /// Jobs of this process running on a connection of their own. + own: Vec, + next_own_id: u64, +} + +struct OwnConnection { + id: u64, + scope: String, + since: Instant, +} + +impl PgConnectionCache { + /// Registers a job on its own connection. Returns the idle connections of + /// its scope, to be closed before it connects. + fn register_own(&mut self, scope: &str) -> (u64, Vec) { + let id = self.next_own_id; + self.next_own_id += 1; + self.own + .push(OwnConnection { id, scope: scope.to_string(), since: Instant::now() }); + let (evicted, kept) = std::mem::take(&mut self.idle) + .into_iter() + .partition(|c| c.scope == scope); + self.idle = kept; + (id, evicted) + } + + fn stalled(&self) -> bool { + self.own + .iter() + .any(|o| o.since.elapsed() >= PG_OWN_CONNECTION_STALL) + } + + fn may_keep(&self, scope: &str) -> bool { + !self.stalled() && self.own.iter().all(|o| o.scope != scope) + } +} + +fn pg_connection_cache() -> std::sync::MutexGuard<'static, PgConnectionCache> { + PG_CONNECTION_CACHE + .lock() + .unwrap_or_else(|e| e.into_inner()) +} + +/// What a session pooler shares server slots across: PgBouncer pools by +/// database and user. Keys differing only in sslmode, TLS inputs, auth mode or +/// options share a scope. Hosts are compared by name, not address. +fn pg_pool_scope(database: &PgDatabase) -> String { + format!( + "{:?}", + ( + database.host.trim().to_ascii_lowercase(), + database.port.unwrap_or(5432), + &database.dbname, + database.login_name().trim(), + ) + ) +} + +struct PgConnection { + key: String, + scope: String, + client: Client, + connection_task: tokio::task::JoinHandle<()>, + last_used: Instant, +} + +impl Drop for PgConnection { + fn drop(&mut self) { + self.connection_task.abort(); + } +} + +/// A job's claim on a connection. Dropping it checks the connection back in, +/// or closes it when its query did not finish or keeping it would break the +/// cache invariant. +struct PgConnectionLease { + scope: String, + /// Set while the job is on a connection of its own. + own: Option, + conn: Option, + /// Check the connection back in on drop. Set once the job's query has read + /// all its results, so every other way out closes the connection. + keep: bool, +} + +impl PgConnectionLease { + /// Takes the most recently used idle connection for `key`. Otherwise + /// registers the job as on its own connection and closes the idle + /// connections of its scope, so they cannot hold the server slot its own + /// connection is about to need. + fn checkout(key: &str, scope: String) -> Self { + let (lease, evicted) = { + let mut cache = pg_connection_cache(); + match cache.idle.iter().rposition(|c| c.key == key) { + Some(i) => { + let conn = cache.idle.remove(i); + ( + Self { scope, own: None, conn: Some(conn), keep: false }, + vec![], + ) + } + None => { + let (id, evicted) = cache.register_own(&scope); + ( + Self { scope, own: Some(id), conn: None, keep: false }, + evicted, + ) + } + } + }; + drop(evicted); + lease + } + + fn uncached() -> Self { + Self { scope: String::new(), own: None, conn: None, keep: false } + } + + /// Gives up a cached connection that failed its probe; the job then + /// connects on its own. + fn discard(&mut self) { + let conn = self.conn.take(); + let evicted = if self.own.is_none() { + let (id, evicted) = pg_connection_cache().register_own(&self.scope); + self.own = Some(id); + evicted + } else { + vec![] + }; + drop(conn); + drop(evicted); + } + + fn client(&self) -> &Client { + &self.conn.as_ref().expect("lease holds a connection").client + } +} + +impl Drop for PgConnectionLease { + fn drop(&mut self) { + let conn = self.conn.take(); + let (kept, closed) = { + let mut cache = pg_connection_cache(); + if let Some(id) = self.own { + cache.own.retain(|o| o.id != id); + } + match conn { + Some(mut conn) if self.keep && cache.may_keep(&self.scope) => { + conn.last_used = Instant::now(); + cache.idle.push(conn); + let over = cache.idle.len().saturating_sub(PG_MAX_IDLE_CONNECTIONS); + (true, cache.idle.drain(..over).collect::>()) + } + conn => (false, conn.into_iter().collect()), + } + }; + drop(closed); + if kept { + PG_CONNECTION_CACHE_SWEEPER.call_once(|| { + tokio::spawn(sweep_idle_pg_connections()); + }); + } + } +} + +async fn sweep_idle_pg_connections() { + loop { + tokio::time::sleep(Duration::from_secs(1)).await; + let (stalled, closed) = { + let mut cache = pg_connection_cache(); + let stalled = cache.stalled(); + let (closed, kept): (Vec<_>, Vec<_>) = std::mem::take(&mut cache.idle) + .into_iter() + .partition(|c| stalled || c.last_used.elapsed() > PG_CONNECTION_IDLE_TIMEOUT); + cache.idle = kept; + (stalled, closed) + }; + if !closed.is_empty() { + if stalled { + tracing::info!( + "Closing {} cached pg executor connection(s): a job has been on its own connection for over {}s", + closed.len(), + PG_OWN_CONNECTION_STALL.as_secs() + ); + } else { + tracing::info!( + "Closing {} cached pg executor connection(s) due to inactivity", + closed.len() + ); + } + } + } +} + +/// PostgreSQL's role, database and server caps, by SQLSTATE 53300 since the +/// server translates the message under a non-English `lc_messages`, and +/// PgBouncer's client and user caps, whose messages are English only and come +/// with a generic code. +fn is_connection_cap_refusal(e: &Error) -> bool { + if let Error::Anyhow { error, .. } = e { + let postgres_cap = error + .chain() + .filter_map(|cause| cause.downcast_ref::()) + .any(|pg| pg.code() == Some(&tokio_postgres::error::SqlState::TOO_MANY_CONNECTIONS)); + if postgres_cap { + return true; + } + } + e.to_string().contains("no more connections allowed") +} + +fn close_idle_pg_connections() -> usize { + let idle = std::mem::take(&mut pg_connection_cache().idle); + idle.len() +} + /// How far a job's statements have got, read by its stall warning. struct PgProgress { start: std::time::Instant, @@ -149,173 +385,8 @@ async fn warn_on_stalled_statement( } } -/// One reusable connection per worker process, checked out by a job for the -/// duration of its query and checked back in afterwards. -/// -/// Invariant: the process never keeps an idle connection while one of its jobs -/// runs on a connection of its own. Behind a session-mode pooler (PgBouncer, -/// RDS Proxy pinning, …) the idle connection holds a server slot, and the job -/// on its own connection waits for that slot with no timeout, forever if the -/// idle connection keeps being reused by other jobs. -static PG_CONNECTION_CACHE: std::sync::Mutex = - std::sync::Mutex::new(PgConnectionCache { idle: None, checked_out: false, uncached: 0 }); -static PG_CONNECTION_CACHE_SWEEPER: std::sync::Once = std::sync::Once::new(); -const PG_CONNECTION_IDLE_TIMEOUT: Duration = Duration::from_secs(60); - -struct PgConnectionCache { - idle: Option, - /// The cached connection is out with a running job. - checked_out: bool, - /// Jobs of this process running on a connection that is not the cached one. - uncached: usize, -} - -fn pg_connection_cache() -> std::sync::MutexGuard<'static, PgConnectionCache> { - PG_CONNECTION_CACHE - .lock() - .unwrap_or_else(|e| e.into_inner()) -} - -struct PgConnection { - key: String, - client: Client, - connection_task: tokio::task::JoinHandle<()>, - last_used: std::time::Instant, -} - -impl Drop for PgConnection { - fn drop(&mut self) { - self.connection_task.abort(); - } -} - -#[derive(Clone, Copy, PartialEq)] -enum PgLeaseSlot { - CheckedOut, - Uncached, -} - -/// A job's claim on a connection. Dropping it checks the connection back in, -/// or closes it when keeping it would break the cache invariant. -struct PgConnectionLease { - slot: PgLeaseSlot, - conn: Option, - /// Cache this job's own connection on release. - cache_on_release: bool, -} - -impl PgConnectionLease { - /// Takes the cached connection when it matches `key`. Otherwise registers - /// the job as uncached and closes the idle connection, whatever its key, so - /// it cannot hold the server slot the job's own connection is about to need. - fn checkout(key: &str) -> Self { - let (lease, evicted) = { - let mut cache = pg_connection_cache(); - match cache.idle.take() { - Some(conn) if conn.key == key => { - cache.checked_out = true; - (Self::new(PgLeaseSlot::CheckedOut, Some(conn)), None) - } - other => { - cache.uncached += 1; - (Self::new(PgLeaseSlot::Uncached, None), other) - } - } - }; - drop(evicted); - lease - } - - fn uncached() -> Self { - pg_connection_cache().uncached += 1; - Self::new(PgLeaseSlot::Uncached, None) - } - - fn new(slot: PgLeaseSlot, conn: Option) -> Self { - Self { slot, conn, cache_on_release: false } - } - - /// Closes the lease's connection so it is never cached again: a cached one - /// that failed its probe (the job then connects on its own), or one a job - /// failed on or left inside a transaction. - fn discard(&mut self) { - let conn = self.conn.take(); - if self.slot == PgLeaseSlot::CheckedOut { - let mut cache = pg_connection_cache(); - cache.checked_out = false; - cache.uncached += 1; - self.slot = PgLeaseSlot::Uncached; - } - drop(conn); - } - - fn client(&self) -> &Client { - &self.conn.as_ref().expect("lease holds a connection").client - } -} - -impl Drop for PgConnectionLease { - fn drop(&mut self) { - let Some(mut conn) = self.conn.take() else { - if self.slot == PgLeaseSlot::Uncached { - pg_connection_cache().uncached -= 1; - } - return; - }; - let kept = { - let mut cache = pg_connection_cache(); - let wants_cache = match self.slot { - PgLeaseSlot::CheckedOut => { - cache.checked_out = false; - true - } - PgLeaseSlot::Uncached => { - cache.uncached -= 1; - self.cache_on_release - } - }; - if wants_cache && cache.uncached == 0 && !cache.checked_out && cache.idle.is_none() { - conn.last_used = std::time::Instant::now(); - cache.idle = Some(conn); - None - } else { - Some(conn) - } - }; - match kept { - None => PG_CONNECTION_CACHE_SWEEPER.call_once(|| { - tokio::spawn(sweep_idle_pg_connection()); - }), - Some(conn) => drop(conn), - } - } -} - -async fn sweep_idle_pg_connection() { - loop { - tokio::time::sleep(Duration::from_secs(5)).await; - let expired = { - let mut cache = pg_connection_cache(); - if cache - .idle - .as_ref() - .is_some_and(|c| c.last_used.elapsed() > PG_CONNECTION_IDLE_TIMEOUT) - { - cache.idle.take() - } else { - None - } - }; - if expired.is_some() { - tracing::info!("Closing cache pg executor connection due to inactivity"); - } - } -} - pub async fn clear_pg_cache() { - let idle = pg_connection_cache().idle.take(); - drop(idle); - CONNECTION_COUNTER.write().await.clear(); + close_idle_pg_connections(); } /// How the connection authenticates, which also keys the connection cache: a @@ -1125,11 +1196,11 @@ pub async fn do_postgresql( let connect_started = std::time::Instant::now(); let mut cached_connection_failed_reset = false; + let mut reused_cached_connection = false; let mut lease = if *CLOUD_HOSTED { PgConnectionLease::uncached() } else { - increment_connection_counter(&database_string).await; - PgConnectionLease::checkout(&database_string) + PgConnectionLease::checkout(&database_string, pg_pool_scope(&database)) }; if lease.conn.is_some() { @@ -1187,6 +1258,7 @@ pub async fn do_postgresql( { tracing::info!("Using cached connection"); CACHE_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + reused_cached_connection = true; } else { tracing::info!("Cached connection is stale, creating new one"); cached_connection_failed_reset = true; @@ -1195,16 +1267,29 @@ pub async fn do_postgresql( } if lease.conn.is_none() { let (client, connection_task) = - new_pg_connection(&database, auth_mode, conn.as_sql()).await?; + match new_pg_connection(&database, auth_mode, conn.as_sql()).await { + Ok(connected) => connected, + // A server or pooler can also refuse outright past a per-login or + // total connection cap, which pool scopes do not capture. Free the + // idle connections of every scope and try once more. + Err(e) => { + if !is_connection_cap_refusal(&e) || close_idle_pg_connections() == 0 { + return Err(e); + } + tokio::time::sleep(Duration::from_millis(250)).await; + new_pg_connection(&database, auth_mode, conn.as_sql()).await? + } + }; lease.conn = Some(PgConnection { key: database_string.clone(), + scope: lease.scope.clone(), client, connection_task, - last_used: std::time::Instant::now(), + last_used: Instant::now(), }); } - let fresh_connection = lease.slot == PgLeaseSlot::Uncached; + let fresh_connection = !reused_cached_connection; let log_progress = !run_inline && !annotations.prepare; let connect_time = connect_started.elapsed(); if log_progress && connect_time >= PG_SLOW_STEP { @@ -1381,17 +1466,7 @@ pub async fn do_postgresql( ) .await } - .map_err(|e| map_s3object_jsonb_overflow(e, had_s3object_input)); - // A failed job can leave the connection inside a transaction, or with its - // query still running after a timeout or cancel, and the status byte can - // lag behind an error. Never reuse it. - let result = match result { - Ok(result) => result, - Err(e) => { - lease.discard(); - return Err(e); - } - }; + .map_err(|e| map_s3object_jsonb_overflow(e, had_s3object_input))?; *mem_peak = size.load(Ordering::Relaxed) as i32; @@ -1400,7 +1475,6 @@ pub async fn do_postgresql( // back, as it always has without the cache. The status comes with the last // reply, so a script that ends cleanly pays nothing for this check. if lease.client().transaction_status() != tokio_postgres::TransactionStatus::Idle { - lease.discard(); if !run_inline { windmill_queue::append_logs( &job.id, @@ -1411,8 +1485,12 @@ pub async fn do_postgresql( ) .await; } - } else if !*CLOUD_HOSTED && lease.slot == PgLeaseSlot::Uncached { - lease.cache_on_release = is_most_used_conn(&database_string).await; + } else { + // Keep the connection only after a query that read its results to the end. + // An error may have stopped reading early (the result cap, a decode failure), + // as a timeout, a dropped future or a first-row collection do: the server + // keeps sending, and the next job's reset probe would wait on it. + lease.keep = !*CLOUD_HOSTED && !collection_strategy.collect_first_row_only(); } drop(lease); @@ -1421,18 +1499,6 @@ pub async fn do_postgresql( return Ok(result); } -async fn is_most_used_conn(database_string: &str) -> bool { - let counter_map = CONNECTION_COUNTER.read().await; - let current_count = counter_map.get(database_string).copied().unwrap_or(0); - let max_count = counter_map.values().copied().max().unwrap_or(0); - current_count >= max_count -} - -async fn increment_connection_counter(database_string: &str) { - let mut counter_map = CONNECTION_COUNTER.write().await; - *counter_map.entry(database_string.to_string()).or_insert(0) += 1; -} - /// For each `(s3object)` arg in `sig_args`: download the referenced file, decode it /// to JSON text, then rewrite the arg to bind as `jsonb`. Mutates `args_map` in place /// so the existing bind path picks up the materialized payload. Returns whether any