perf: keep a pg executor connection per database (#11492)

* fix: stop the pg executor cache from pinning a pooler slot

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

* test: retry while the evicted pg backend exits

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

* perf: keep a pg executor connection per database

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KHjj2iG8kKPCaFuK6bqXvi

* fix: close a pg connection whose query was cut off by a timeout or cancel

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KHjj2iG8kKPCaFuK6bqXvi

* fix: keep a pg connection only after its query read every result

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KHjj2iG8kKPCaFuK6bqXvi

* fix: free idle pg connections and retry when a connect is refused

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KHjj2iG8kKPCaFuK6bqXvi

* fix: retry a pg connect only when a connection cap refused it

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KHjj2iG8kKPCaFuK6bqXvi

* fix: detect pg connection caps by SQLSTATE and label cloud connections as new

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KHjj2iG8kKPCaFuK6bqXvi

---------

Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

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<Postgres>) -> 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<Postgres>,
+) -> 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"))]

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<RwLock<HashMap<String, u64>>> =
-        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<PgConnectionCache> =
+    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<PgConnection>,
+    /// Jobs of this process running on a connection of their own.
+    own: Vec<OwnConnection>,
+    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<PgConnection>) {
+        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<u64>,
+    conn: Option<PgConnection>,
+    /// 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::<Vec<_>>())
+                }
+                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::<tokio_postgres::Error>())
+            .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<PgConnectionCache> =
-    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<PgConnection>,
-    /// 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<PgConnection>,
-    /// 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<PgConnection>) -> 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