systemprompt_database/services/postgres/
connection.rs

1//! Initial-connect retry policy for `PostgresProvider`.
2//!
3//! Wraps the first `PgPool` connect in a bounded exponential backoff so
4//! transient startup races (Postgres still booting, SSL handshake racing
5//! the TCP listener) recover without surfacing as user-visible failures.
6//! The retry loop intentionally targets a narrow set of error shapes so
7//! permanent failures (auth, missing database, bad URL) fail fast. The
8//! backoff itself runs on [`crate::resilience::retry::retry_async`].
9//!
10//! Copyright (c) systemprompt.io — Business Source License 1.1.
11//! See <https://systemprompt.io> for licensing details.
12
13use std::future::Future;
14use std::str::FromStr;
15use std::sync::OnceLock;
16use std::sync::atomic::{AtomicU64, Ordering};
17use std::time::{Duration, Instant};
18
19use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions};
20
21use crate::error::DatabaseResult;
22use crate::resilience::classify::Outcome;
23use crate::resilience::config::RetryConfig;
24use crate::resilience::retry::retry_async;
25
26const RETRY_DELAYS_MS: &[u64] = &[100, 200, 400, 800, 1600];
27const MAX_ATTEMPTS: u32 = 5;
28pub const DEFAULT_STATEMENT_CACHE_CAPACITY: usize = 100;
29
30static POOL_CLOCK: OnceLock<Instant> = OnceLock::new();
31static SCHEMA_CHANGED_AT_NANOS: AtomicU64 = AtomicU64::new(0);
32
33fn pool_clock() -> Instant {
34    *POOL_CLOCK.get_or_init(Instant::now)
35}
36
37fn nanos_since_clock(at: Duration) -> u64 {
38    u64::try_from(at.as_nanos()).unwrap_or(u64::MAX).max(1)
39}
40
41pub fn mark_schema_changed() {
42    let now = nanos_since_clock(pool_clock().elapsed());
43    SCHEMA_CHANGED_AT_NANOS.store(now, Ordering::Release);
44}
45
46fn opened_after_schema_change(age: Duration) -> bool {
47    let changed = SCHEMA_CHANGED_AT_NANOS.load(Ordering::Acquire);
48    if changed == 0 {
49        return true;
50    }
51    let opened = pool_clock().elapsed().saturating_sub(age);
52    nanos_since_clock(opened) >= changed
53}
54
55/// Operator-tunable connection-pool sizing for a `PostgresProvider`.
56///
57/// Only the sizing/lifetime knobs an operator needs to fit the pool to their
58/// Postgres `max_connections` and replica count are exposed; the connect, SSL
59/// and retry behaviour is fixed.
60#[derive(Debug, Clone, Copy)]
61pub struct PoolConfig {
62    pub max_connections: u32,
63    pub min_connections: u32,
64    pub acquire_timeout: Duration,
65    pub idle_timeout: Duration,
66    pub max_lifetime: Duration,
67    pub statement_cache_capacity: usize,
68}
69
70impl Default for PoolConfig {
71    fn default() -> Self {
72        Self {
73            max_connections: 50,
74            min_connections: 0,
75            acquire_timeout: Duration::from_secs(30),
76            idle_timeout: Duration::from_mins(5),
77            max_lifetime: Duration::from_mins(30),
78            statement_cache_capacity: DEFAULT_STATEMENT_CACHE_CAPACITY,
79        }
80    }
81}
82
83pub const DEFAULT_SESSION_MAX_CONNECTIONS: u32 = 10;
84
85impl PoolConfig {
86    #[must_use]
87    pub const fn for_session(self, max_connections: Option<u32>) -> Self {
88        Self {
89            max_connections: match max_connections {
90                Some(max) => max,
91                None => DEFAULT_SESSION_MAX_CONNECTIONS,
92            },
93            ..self
94        }
95    }
96}
97
98#[must_use]
99pub fn build_pool_options(cfg: &PoolConfig) -> PgPoolOptions {
100    pool_clock();
101    PgPoolOptions::new()
102        .max_connections(cfg.max_connections)
103        .min_connections(cfg.min_connections)
104        .max_lifetime(cfg.max_lifetime)
105        .acquire_timeout(cfg.acquire_timeout)
106        .idle_timeout(cfg.idle_timeout)
107        // Why: a cached prepared statement fails with SQLSTATE 0A000 ("cached
108        // plan must not change result type") once DDL changes the table under
109        // it. In-process migrations call `mark_schema_changed`, and every
110        // connection opened before that is closed on its next acquire instead
111        // of being handed out with stale plans.
112        .before_acquire(|_conn, meta| {
113            Box::pin(async move { Ok(opened_after_schema_change(meta.age)) })
114        })
115}
116
117pub fn connect_options(database_url: &str) -> DatabaseResult<PgConnectOptions> {
118    let options = PgConnectOptions::from_str(database_url)?
119        .application_name("systemprompt")
120        // Why: sqlx 0.9 `PgConnection::get_or_prepare` Parses every persistent
121        // query as a NAMED statement and only sends Close when the cache is
122        // enabled and evicts, so capacity 0 never deallocates them: each
123        // backend grew without bound until Postgres OOMed. A bounded cache
124        // closes on eviction; stale plans after in-process DDL are handled by
125        // `mark_schema_changed` in `build_pool_options`.
126        .statement_cache_capacity(DEFAULT_STATEMENT_CACHE_CAPACITY)
127        // Why: a transaction left open by a crashed or stuck task holds its row
128        // locks and stops autovacuum reclaiming the hot tables until the
129        // connection is recycled (30 min). Postgres ends a session idle inside
130        // a transaction past this bound; no request or job legitimately sits
131        // idle mid-transaction that long.
132        .options([
133            ("client_min_messages", "warning"),
134            ("idle_in_transaction_session_timeout", "600000"),
135        ]);
136    Ok(options)
137}
138
139pub async fn connect_with_retry(
140    options: PgPoolOptions,
141    connect_options: PgConnectOptions,
142) -> DatabaseResult<PgPool> {
143    let connector = |opts: PgConnectOptions| {
144        let options = options.clone();
145        async move { options.connect_with(opts).await }
146    };
147    connect_with_retry_using(connect_options, MAX_ATTEMPTS, RETRY_DELAYS_MS, connector).await
148}
149
150pub async fn connect_with_retry_using<T, F, Fut>(
151    connect_options: PgConnectOptions,
152    max_attempts: u32,
153    delays_ms: &[u64],
154    connector: F,
155) -> DatabaseResult<T>
156where
157    T: Send,
158    F: Fn(PgConnectOptions) -> Fut + Send + Sync,
159    Fut: Future<Output = Result<T, sqlx::Error>> + Send,
160{
161    let cfg = RetryConfig {
162        max_attempts,
163        base_delay: Duration::from_millis(delays_ms.first().copied().unwrap_or(100)),
164        max_delay: Duration::from_millis(delays_ms.iter().copied().max().unwrap_or(1600)),
165        jitter: false,
166    };
167    let classify = |err: &sqlx::Error| {
168        if is_retryable(err) {
169            Outcome::Transient { retry_after: None }
170        } else {
171            Outcome::Permanent
172        }
173    };
174    retry_async(&cfg, "postgres-connect", classify, || {
175        connector(connect_options.clone())
176    })
177    .await
178    .map_err(Into::into)
179}
180
181fn is_retryable(err: &sqlx::Error) -> bool {
182    if let sqlx::Error::Io(io_err) = err
183        && io_err.kind() == std::io::ErrorKind::ConnectionRefused
184    {
185        return true;
186    }
187    let msg = err.to_string();
188    msg.contains("unexpected response from SSLRequest") || msg.contains("starting up")
189}
systemprompt_database/services/postgres/connection.rs

systemprompt_database/services/postgres/
connection.rs
