Skip to content

Commit 196efc1

Browse files
committed
feat(server): make the database pool ceiling configurable
Both stores hardcoded their pool ceiling: 10 connections for Postgres, 5 for on-disk SQLite. One pool is shared by every database-backed RPC, so that number is also the gateway's ceiling on concurrent database work. Once every connection is checked out, callers queue on acquire and sqlx logs "time to acquire exceeded slow threshold"; sandbox creates then time out and are retried, which adds load rather than shedding it. Scaling the database up does not help, because the pool never opens more connections. The right ceiling is deployment-specific — it depends on how many gateway replicas share the database and what max_connections the server itself allows — so it cannot be a single number baked into the binary. Expose it on the same config surface as the rest of the gateway's settings: --db-max-connections, OPENSHELL_DB_MAX_CONNECTIONS, and the TOML key database_max_connections, resolved in that precedence order and rendered by the Helm chart from server.dbMaxConnections. Omitting it keeps each backend's previous value, so existing deployments do not move. Values below 1 are rejected rather than silently replaced by the default: a zero pool would block every acquire, and quietly ignoring a typo would reproduce the ceiling the operator is trying to lift. An in-memory SQLite database stays pinned to one connection, since the database lives inside that connection. Refs: #2561 Signed-off-by: Bryce Wilkinson <bwilkinson@alpha-sense.com>
1 parent 3a16012 commit 196efc1

15 files changed

Lines changed: 328 additions & 9 deletions

File tree

‎crates/openshell-core/src/config.rs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -461,6 +461,16 @@ pub struct Config {
461461
/// Database URL for persistence.
462462
pub database_url: String,
463463

464+
/// Connection ceiling for the persistence pool.
465+
///
466+
/// `None` leaves the backend's built-in default in place. The right
467+
/// ceiling depends on the deployment — how many gateway replicas share
468+
/// the database, and what `max_connections` a `Postgres` server itself
469+
/// allows — so it cannot be one number baked into the binary. An
470+
/// in-memory `SQLite` database ignores it: the database lives in its single
471+
/// connection.
472+
pub database_max_connections: Option<u32>,
473+
464474
/// Compute drivers configured for the gateway.
465475
///
466476
/// The config shape allows multiple drivers so the gateway can evolve
@@ -835,6 +845,7 @@ impl Config {
835845
mtls_auth: MtlsAuthConfig::default(),
836846
gateway_jwt: None,
837847
database_url: String::new(),
848+
database_max_connections: None,
838849
compute_drivers: vec![],
839850
compute_driver_endpoints: BTreeMap::new(),
840851
credential_drivers: Vec::new(),
@@ -879,6 +890,15 @@ impl Config {
879890
self
880891
}
881892

893+
/// Create a new configuration with a database pool connection ceiling.
894+
///
895+
/// `None` keeps the persistence backend's default.
896+
#[must_use]
897+
pub const fn with_database_max_connections(mut self, max_connections: Option<u32>) -> Self {
898+
self.database_max_connections = max_connections;
899+
self
900+
}
901+
882902
/// Create a new configuration with the configured compute drivers.
883903
#[must_use]
884904
pub fn with_compute_drivers<I, D>(mut self, drivers: I) -> Self

‎crates/openshell-server/src/cli.rs‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,21 @@ struct RunArgs {
9494
#[arg(long, env = "OPENSHELL_DB_URL")]
9595
db_url: Option<String>,
9696

97+
/// Connection ceiling for the database persistence pool.
98+
///
99+
/// One shared pool serves every database-backed RPC, so this also caps
100+
/// the gateway's concurrent database work. On `Postgres`, total server
101+
/// connections are this value times the replica count, which must stay
102+
/// under the server's own `max_connections`. When unset, the gateway uses
103+
/// its backend default: 10 for `Postgres`, 5 for on-disk `SQLite`. An
104+
/// in-memory `SQLite` database is always a single connection.
105+
#[arg(
106+
long,
107+
env = "OPENSHELL_DB_MAX_CONNECTIONS",
108+
value_parser = clap::value_parser!(u32).range(1..)
109+
)]
110+
db_max_connections: Option<u32>,
111+
97112
/// Compute drivers configured for this gateway.
98113
///
99114
/// Accepts a comma-delimited list such as `kubernetes` or
@@ -383,6 +398,7 @@ fn prepare_server_config(args: &mut RunArgs, matches: &ArgMatches) -> Result<Ser
383398

384399
config = config
385400
.with_database_url(db_url)
401+
.with_database_max_connections(args.db_max_connections)
386402
.with_compute_drivers(args.drivers.clone())
387403
.with_grpc_rate_limit(
388404
args.grpc_rate_limit_requests,
@@ -723,6 +739,12 @@ fn merge_file_into_args(args: &mut RunArgs, file: &GatewayFileSection, matches:
723739
{
724740
args.grpc_rate_limit_window_seconds = Some(window);
725741
}
742+
if let Some(max_connections) = file.database_max_connections
743+
&& args.db_max_connections.is_none()
744+
&& arg_defaulted(matches, "db_max_connections")
745+
{
746+
args.db_max_connections = Some(max_connections);
747+
}
726748
}
727749

728750
fn validate_grpc_rate_limit_args(requests: Option<u64>, window_seconds: Option<u64>) -> Result<()> {
@@ -1489,6 +1511,75 @@ grpc_rate_limit_window_seconds = 30
14891511
assert_eq!(args.grpc_rate_limit_window_seconds, Some(30));
14901512
}
14911513

1514+
#[test]
1515+
fn file_db_max_connections_populates_args_when_cli_omits() {
1516+
let _lock = ENV_LOCK
1517+
.lock()
1518+
.unwrap_or_else(std::sync::PoisonError::into_inner);
1519+
let _g = EnvVarGuard::remove("OPENSHELL_DB_MAX_CONNECTIONS");
1520+
1521+
let (mut args, matches) =
1522+
parse_with_args(&["openshell-gateway", "--db-url", "sqlite::memory:"]);
1523+
assert_eq!(
1524+
args.db_max_connections, None,
1525+
"default leaves the pool alone"
1526+
);
1527+
let file = config_file_from_toml(
1528+
r"
1529+
[openshell.gateway]
1530+
database_max_connections = 64
1531+
",
1532+
);
1533+
merge_file_into_args(&mut args, &file.openshell.gateway, &matches);
1534+
1535+
assert_eq!(args.db_max_connections, Some(64));
1536+
}
1537+
1538+
#[test]
1539+
fn env_db_max_connections_overrides_file_value() {
1540+
let _lock = ENV_LOCK
1541+
.lock()
1542+
.unwrap_or_else(std::sync::PoisonError::into_inner);
1543+
let _g = EnvVarGuard::set("OPENSHELL_DB_MAX_CONNECTIONS", "128");
1544+
1545+
let (mut args, matches) =
1546+
parse_with_args(&["openshell-gateway", "--db-url", "sqlite::memory:"]);
1547+
let file = config_file_from_toml(
1548+
r"
1549+
[openshell.gateway]
1550+
database_max_connections = 64
1551+
",
1552+
);
1553+
merge_file_into_args(&mut args, &file.openshell.gateway, &matches);
1554+
1555+
assert_eq!(args.db_max_connections, Some(128));
1556+
}
1557+
1558+
#[test]
1559+
fn db_max_connections_rejects_a_zero_or_unparseable_ceiling() {
1560+
let _lock = ENV_LOCK
1561+
.lock()
1562+
.unwrap_or_else(std::sync::PoisonError::into_inner);
1563+
let _g = EnvVarGuard::remove("OPENSHELL_DB_MAX_CONNECTIONS");
1564+
1565+
// Zero would deadlock every `acquire`, and a silent fallback to the
1566+
// default would reproduce the ceiling the operator is trying to lift.
1567+
for bad in ["0", "-5", "many"] {
1568+
assert!(
1569+
command()
1570+
.try_get_matches_from([
1571+
"openshell-gateway",
1572+
"--db-url",
1573+
"sqlite::memory:",
1574+
"--db-max-connections",
1575+
bad,
1576+
])
1577+
.is_err(),
1578+
"input {bad:?} must be rejected"
1579+
);
1580+
}
1581+
}
1582+
14921583
#[test]
14931584
fn aux_listener_preserves_file_ip_against_public_bind() {
14941585
use std::net::SocketAddr;

‎crates/openshell-server/src/config_file.rs‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,17 @@ pub struct GatewayFileSection {
9999
#[serde(default)]
100100
pub log_level: Option<String>,
101101

102+
// ── Database ─────────────────────────────────────────────────────────
103+
/// Connection ceiling for the persistence pool. One shared pool serves
104+
/// every database-backed RPC, so this is also the gateway's ceiling on
105+
/// concurrent database work. On `Postgres`, total server connections are
106+
/// `database_max_connections × replica_count`, which must stay under the
107+
/// server's own `max_connections`. Omit to keep the backend default (10
108+
/// for `Postgres`, 5 for on-disk `SQLite`); an in-memory `SQLite` database
109+
/// is always a single connection.
110+
#[serde(default)]
111+
pub database_max_connections: Option<u32>,
112+
102113
// ── Drivers ──────────────────────────────────────────────────────────
103114
#[serde(default)]
104115
pub compute_drivers: Option<Vec<String>>,
@@ -399,6 +410,14 @@ pub fn load(path: &Path) -> Result<ConfigFile, ConfigFileError> {
399410
cli: "--db-url",
400411
});
401412
}
413+
// A zero-sized pool would deadlock every `acquire` rather than degrade,
414+
// so reject it at load instead of starting a gateway that cannot serve.
415+
if file.openshell.gateway.database_max_connections == Some(0) {
416+
return Err(ConfigFileError::InvalidValue {
417+
field: "openshell.gateway.database_max_connections",
418+
message: "must be greater than zero; omit the key to use the built-in default",
419+
});
420+
}
402421
if file
403422
.openshell
404423
.gateway
@@ -912,6 +931,34 @@ database_url = "sqlite::memory:"
912931
));
913932
}
914933

934+
#[test]
935+
fn parses_database_max_connections() {
936+
let toml = r"
937+
[openshell.gateway]
938+
database_max_connections = 64
939+
";
940+
let tmp = write_tmp(toml);
941+
let file = load(tmp.path()).expect("a positive pool ceiling parses");
942+
assert_eq!(file.openshell.gateway.database_max_connections, Some(64));
943+
}
944+
945+
#[test]
946+
fn rejects_zero_database_max_connections() {
947+
let toml = r"
948+
[openshell.gateway]
949+
database_max_connections = 0
950+
";
951+
let tmp = write_tmp(toml);
952+
let err = load(tmp.path()).expect_err("a zero pool ceiling must be rejected");
953+
assert!(matches!(
954+
err,
955+
ConfigFileError::InvalidValue {
956+
field: "openshell.gateway.database_max_connections",
957+
..
958+
}
959+
));
960+
}
961+
915962
#[test]
916963
fn rejects_unknown_gateway_field() {
917964
let toml = r"

‎crates/openshell-server/src/lib.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -554,7 +554,9 @@ pub(crate) async fn run_server(
554554
.map_err(|error| Error::config(format!("middleware registration failed: {error}")))?,
555555
);
556556

557-
let store = Arc::new(Store::connect(database_url).await?);
557+
let store = Arc::new(
558+
Store::connect_with_pool_size(database_url, config.database_max_connections).await?,
559+
);
558560
let credentials = credentials::CredentialRuntime::from_config_file_with_store(
559561
&config,
560562
config_file.as_ref(),

‎crates/openshell-server/src/persistence/mod.rs‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -201,10 +201,24 @@ impl Store {
201201
matches!(self, Self::Sqlite(_))
202202
}
203203

204-
/// Connect to a persistence store based on the database URL.
204+
/// Connect to a persistence store based on the database URL, leaving pool
205+
/// sizing at the backend default.
205206
pub async fn connect(url: &str) -> CoreResult<Self> {
207+
Self::connect_with_pool_size(url, None).await
208+
}
209+
210+
/// Connect to a persistence store based on the database URL.
211+
///
212+
/// `max_connections` overrides the backend's pool ceiling, which stays at
213+
/// its built-in default when `None`. The two backends have different
214+
/// defaults, and an in-memory `SQLite` database ignores the override
215+
/// entirely — it must be held by exactly one connection.
216+
pub async fn connect_with_pool_size(
217+
url: &str,
218+
max_connections: Option<u32>,
219+
) -> CoreResult<Self> {
206220
if url.starts_with("postgres://") || url.starts_with("postgresql://") {
207-
let store = PostgresStore::connect(url)
221+
let store = PostgresStore::connect(url, max_connections)
208222
.await
209223
.map_err(|e| CoreError::execution(e.to_string()))?;
210224
store
@@ -213,7 +227,7 @@ impl Store {
213227
.map_err(|e| CoreError::execution(e.to_string()))?;
214228
Ok(Self::Postgres(store))
215229
} else if url.starts_with("sqlite:") {
216-
let store = SqliteStore::connect(url)
230+
let store = SqliteStore::connect(url, max_connections)
217231
.await
218232
.map_err(|e| CoreError::execution(e.to_string()))?;
219233
store

‎crates/openshell-server/src/persistence/postgres.rs‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,21 @@ pub struct PostgresStore {
2525
pool: PgPool,
2626
}
2727

28+
/// Pool ceiling used when the operator has not configured one.
29+
///
30+
/// A single shared pool serves every database-backed RPC, so this value is
31+
/// also the gateway's ceiling on concurrent database work. It is high enough
32+
/// for a single-replica gateway against a default-sized server; fleets that
33+
/// outgrow it raise it through `[openshell.gateway] database_max_connections`.
34+
pub const DEFAULT_MAX_CONNECTIONS: u32 = 10;
35+
2836
impl PostgresStore {
29-
pub async fn connect(url: &str) -> PersistenceResult<Self> {
37+
/// Connect and size the pool, falling back to [`DEFAULT_MAX_CONNECTIONS`].
38+
pub async fn connect(url: &str, max_connections: Option<u32>) -> PersistenceResult<Self> {
39+
let max_connections = max_connections.unwrap_or(DEFAULT_MAX_CONNECTIONS);
40+
tracing::info!(max_connections, "sizing Postgres connection pool");
3041
let pool = PgPoolOptions::new()
31-
.max_connections(10)
42+
.max_connections(max_connections)
3243
.connect(url)
3344
.await
3445
.map_err(|e| map_db_error(&e))?;

‎crates/openshell-server/src/persistence/sqlite.rs‎

Lines changed: 37 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,14 @@ static SQLITE_MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations/s
2323

2424
use super::{DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE};
2525

26+
/// Pool ceiling for an on-disk database when the operator has not configured
27+
/// one.
28+
///
29+
/// `SQLite` serializes writers, so raising this mainly buys read concurrency;
30+
/// a gateway that needs more write throughput wants `Postgres`, not a larger
31+
/// `SQLite` pool.
32+
pub const DEFAULT_MAX_CONNECTIONS: u32 = 5;
33+
2634
#[derive(Debug, Clone)]
2735
pub struct SqliteStore {
2836
pool: SqlitePool,
@@ -35,17 +43,43 @@ impl SqliteStore {
3543
self.pool.close().await;
3644
}
3745

38-
pub async fn connect(url: &str) -> PersistenceResult<Self> {
46+
/// Test support only: the pool's effective connection ceiling.
47+
#[cfg(test)]
48+
pub(crate) fn max_connections_for_test(&self) -> u32 {
49+
self.pool.options().get_max_connections()
50+
}
51+
52+
/// Connect and size the pool, falling back to [`DEFAULT_MAX_CONNECTIONS`]
53+
/// for an on-disk database.
54+
///
55+
/// An in-memory database is always held by exactly one connection — the
56+
/// database lives in that connection, so a second one would see a
57+
/// different (empty) database. A configured ceiling is ignored there.
58+
pub async fn connect(url: &str, max_connections: Option<u32>) -> PersistenceResult<Self> {
3959
let is_in_memory = url.contains(":memory:") || url.contains("mode=memory");
40-
let max_connections = if is_in_memory { 1 } else { 5 };
60+
let max_connections = if is_in_memory {
61+
if max_connections.is_some_and(|configured| configured != 1) {
62+
tracing::warn!(
63+
"ignoring the configured database pool ceiling: an in-memory SQLite database is held by a single connection"
64+
);
65+
}
66+
1
67+
} else {
68+
max_connections.unwrap_or(DEFAULT_MAX_CONNECTIONS)
69+
};
70+
// Keep the pool eagerly warm up to the built-in size, but let a raised
71+
// ceiling grow on demand rather than pinning that many file handles
72+
// open for the life of the process.
73+
let min_connections = max_connections.min(DEFAULT_MAX_CONNECTIONS);
74+
tracing::info!(max_connections, "sizing SQLite connection pool");
4175

4276
let options = SqliteConnectOptions::from_str(url)
4377
.map_err(|e| map_db_error(&e))?
4478
.create_if_missing(true);
4579

4680
let mut pool_options = SqlitePoolOptions::new()
4781
.max_connections(max_connections)
48-
.min_connections(max_connections);
82+
.min_connections(min_connections);
4983

5084
if is_in_memory {
5185
pool_options = pool_options.idle_timeout(None).max_lifetime(None);

‎crates/openshell-server/src/persistence/tests.rs‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,37 @@ async fn sqlite_connect_runs_embedded_migrations() {
132132
assert!(records.is_empty());
133133
}
134134

135+
#[tokio::test]
136+
async fn sqlite_pool_ceiling_defaults_and_honours_an_override() {
137+
use super::sqlite::{DEFAULT_MAX_CONNECTIONS, SqliteStore};
138+
139+
let tmp = tempfile::tempdir().expect("tempdir");
140+
let db_path = tmp.path().join("openshell.db");
141+
let url = format!("sqlite:{}?mode=rwc", db_path.display());
142+
143+
let unset = SqliteStore::connect(&url, None)
144+
.await
145+
.expect("connect to sqlite");
146+
assert_eq!(unset.max_connections_for_test(), DEFAULT_MAX_CONNECTIONS);
147+
148+
let raised = SqliteStore::connect(&url, Some(32))
149+
.await
150+
.expect("connect to sqlite");
151+
assert_eq!(raised.max_connections_for_test(), 32);
152+
}
153+
154+
/// An in-memory database lives inside its single connection: a second one
155+
/// would open a different, empty database, so the ceiling cannot be raised.
156+
#[tokio::test]
157+
async fn in_memory_sqlite_pins_the_pool_to_one_connection() {
158+
use super::sqlite::SqliteStore;
159+
160+
let store = SqliteStore::connect("sqlite::memory:", Some(32))
161+
.await
162+
.expect("connect to in-memory sqlite");
163+
assert_eq!(store.max_connections_for_test(), 1);
164+
}
165+
135166
#[cfg(unix)]
136167
#[tokio::test]
137168
async fn sqlite_connect_restricts_db_file_permissions() {

0 commit comments

Comments
 (0)