From 44b312c0c5c820c48bfa2ae98081cc7ac7493edc Mon Sep 17 00:00:00 2001
From: Bryce Wilkinson <bwilkinson@alpha-sense.com>
Date: Tue, 11 Aug 2026 14:04:22 +0300
Subject: [PATCH] feat(server): make the Postgres pool ceiling configurable
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The Postgres store hardcoded max_connections(10), which makes the pool
the throughput limit under concurrent sandbox creation. Once all ten are
checked out, callers queue on acquire and sqlx starts logging "time to
acquire exceeded slow threshold"; sandbox creates then time out and get
retried, which adds load rather than shedding it.

The right ceiling is deployment-specific — it depends on how many gateway
replicas share the server and what max_connections the server itself
allows — so it cannot be a single number baked into the binary.

OPENSHELL_DB_MAX_CONNECTIONS now sets it, defaulting to 10 so existing
deployments are unchanged. Unset, unparseable and zero all fall back to
the default: a typo should degrade to today's behaviour rather than
crash-loop the gateway, and zero would deadlock every acquire.

Refs: #2561
Signed-off-by: Bryce Wilkinson <bwilkinson@alpha-sense.com>
---
 .../src/persistence/postgres.rs               | 100 +++++++++++++++++-
 1 file changed, 99 insertions(+), 1 deletion(-)

diff --git a/crates/openshell-server/src/persistence/postgres.rs b/crates/openshell-server/src/persistence/postgres.rs
index e3f9e5a7b3..491d1bb3f3 100644
--- a/crates/openshell-server/src/persistence/postgres.rs
+++ b/crates/openshell-server/src/persistence/postgres.rs
@@ -25,10 +25,43 @@ pub struct PostgresStore {
     pool: PgPool,
 }
 
+/// Default pool ceiling, unchanged from the previously hardcoded value.
+const DEFAULT_MAX_CONNECTIONS: u32 = 10;
+
+/// Pool ceiling for the Postgres store.
+///
+/// A fixed ceiling makes the pool the throughput limit under concurrent
+/// sandbox creation: once every connection is checked out, callers queue on
+/// `acquire` and sqlx logs `time to acquire exceeded slow threshold`. The
+/// right ceiling depends on the deployment — how many gateway replicas share
+/// the server, and what `max_connections` the server itself allows — so it
+/// cannot be one number baked into the binary. See NVIDIA/OpenShell#2561.
+///
+/// Falls back to the default when unset, unparseable, or zero, so a typo
+/// degrades to today's behaviour instead of crash-looping the gateway.
+fn max_connections() -> u32 {
+    let Ok(raw) = std::env::var("OPENSHELL_DB_MAX_CONNECTIONS") else {
+        return DEFAULT_MAX_CONNECTIONS;
+    };
+    match raw.trim().parse::<u32>() {
+        Ok(n) if n > 0 => n,
+        _ => {
+            tracing::warn!(
+                value = %raw,
+                default = DEFAULT_MAX_CONNECTIONS,
+                "OPENSHELL_DB_MAX_CONNECTIONS is not a positive integer; using the default"
+            );
+            DEFAULT_MAX_CONNECTIONS
+        }
+    }
+}
+
 impl PostgresStore {
     pub async fn connect(url: &str) -> PersistenceResult<Self> {
+        let max_connections = max_connections();
+        tracing::info!(max_connections, "connecting Postgres store");
         let pool = PgPoolOptions::new()
-            .max_connections(10)
+            .max_connections(max_connections)
             .connect(url)
             .await
             .map_err(|e| map_db_error(&e))?;
@@ -1024,3 +1057,68 @@ fn row_to_draft_chunk_record(row: sqlx::postgres::PgRow) -> PersistenceResult<Dr
         updated_at_ms,
     )
 }
+
+#[cfg(test)]
+mod max_connections_tests {
+    use super::{DEFAULT_MAX_CONNECTIONS, max_connections};
+
+    const KEY: &str = "OPENSHELL_DB_MAX_CONNECTIONS";
+
+    struct EnvVarGuard(Option<String>);
+
+    impl EnvVarGuard {
+        #[allow(unsafe_code)]
+        fn set(value: Option<&str>) -> Self {
+            let original = std::env::var(KEY).ok();
+            // SAFETY: tests serialize environment mutation with TEST_ENV_LOCK.
+            match value {
+                Some(v) => unsafe { std::env::set_var(KEY, v) },
+                None => unsafe { std::env::remove_var(KEY) },
+            }
+            Self(original)
+        }
+    }
+
+    impl Drop for EnvVarGuard {
+        #[allow(unsafe_code)]
+        fn drop(&mut self) {
+            // SAFETY: tests serialize environment mutation with TEST_ENV_LOCK.
+            match self.0.as_deref() {
+                Some(v) => unsafe { std::env::set_var(KEY, v) },
+                None => unsafe { std::env::remove_var(KEY) },
+            }
+        }
+    }
+
+    fn with_env(value: Option<&str>) -> u32 {
+        let _lock = crate::TEST_ENV_LOCK
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let _env = EnvVarGuard::set(value);
+        max_connections()
+    }
+
+    #[test]
+    fn unset_keeps_the_previous_hardcoded_ceiling() {
+        assert_eq!(with_env(None), DEFAULT_MAX_CONNECTIONS);
+    }
+
+    #[test]
+    fn a_positive_integer_overrides_the_ceiling() {
+        assert_eq!(with_env(Some("64")), 64);
+        assert_eq!(with_env(Some("  32  ")), 32);
+    }
+
+    #[test]
+    fn unusable_values_fall_back_rather_than_crash_looping() {
+        // A typo must not take the gateway down on startup, and zero would
+        // deadlock every acquire.
+        for bad in ["", "0", "-5", "many", "10.5"] {
+            assert_eq!(
+                with_env(Some(bad)),
+                DEFAULT_MAX_CONNECTIONS,
+                "input {bad:?}"
+            );
+        }
+    }
+}