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::Arc;
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
30/// Schema-change stamp shared by every pool a [`crate::Database`] opens.
31///
32/// A connection opened before the last recorded schema change is closed on
33/// its next acquire, so no prepared statement outlives the DDL under it.
34#[derive(Debug, Clone)]
35pub struct SchemaEpoch {
36    inner: Arc<EpochState>,
37}
38
39#[derive(Debug)]
40struct EpochState {
41    clock: Instant,
42    changed_at_nanos: AtomicU64,
43}
44
45impl Default for SchemaEpoch {
46    fn default() -> Self {
47        Self::new()
48    }
49}
50
51impl SchemaEpoch {
52    #[must_use]
53    pub fn new() -> Self {
54        Self {
55            inner: Arc::new(EpochState {
56                clock: Instant::now(),
57                changed_at_nanos: AtomicU64::new(0),
58            }),
59        }
60    }
61
62    pub fn mark_changed(&self) {
63        let now = nanos_since_clock(self.inner.clock.elapsed());
64        self.inner.changed_at_nanos.store(now, Ordering::Release);
65    }
66
67    fn opened_after_change(&self, age: Duration) -> bool {
68        let changed = self.inner.changed_at_nanos.load(Ordering::Acquire);
69        if changed == 0 {
70            return true;
71        }
72        let opened = self.inner.clock.elapsed().saturating_sub(age);
73        nanos_since_clock(opened) >= changed
74    }
75}
76
77fn nanos_since_clock(at: Duration) -> u64 {
78    u64::try_from(at.as_nanos()).unwrap_or(u64::MAX).max(1)
79}
80
81/// Operator-tunable connection-pool sizing for a `PostgresProvider`.
82///
83/// Only the sizing/lifetime knobs an operator needs to fit the pool to their
84/// Postgres `max_connections` and replica count are exposed; the connect, SSL
85/// and retry behaviour is fixed.
86#[derive(Debug, Clone, Copy)]
87pub struct PoolConfig {
88    pub max_connections: u32,
89    pub min_connections: u32,
90    pub acquire_timeout: Duration,
91    pub idle_timeout: Duration,
92    pub max_lifetime: Duration,
93    pub statement_cache_capacity: usize,
94}
95
96impl Default for PoolConfig {
97    fn default() -> Self {
98        Self {
99            max_connections: 50,
100            min_connections: 0,
101            acquire_timeout: Duration::from_secs(30),
102            idle_timeout: Duration::from_mins(5),
103            max_lifetime: Duration::from_mins(30),
104            statement_cache_capacity: DEFAULT_STATEMENT_CACHE_CAPACITY,
105        }
106    }
107}
108
109pub const DEFAULT_SESSION_MAX_CONNECTIONS: u32 = 10;
110
111impl PoolConfig {
112    #[must_use]
113    pub const fn for_session(self, max_connections: Option<u32>) -> Self {
114        Self {
115            max_connections: match max_connections {
116                Some(max) => max,
117                None => DEFAULT_SESSION_MAX_CONNECTIONS,
118            },
119            ..self
120        }
121    }
122}
123
124#[must_use]
125pub fn build_pool_options(cfg: &PoolConfig, epoch: &SchemaEpoch) -> PgPoolOptions {
126    let epoch = epoch.clone();
127    PgPoolOptions::new()
128        .max_connections(cfg.max_connections)
129        .min_connections(cfg.min_connections)
130        .max_lifetime(cfg.max_lifetime)
131        .acquire_timeout(cfg.acquire_timeout)
132        .idle_timeout(cfg.idle_timeout)
133        // Why: a cached prepared statement fails with SQLSTATE 0A000 ("cached
134        // plan must not change result type") once DDL changes the table under
135        // it. In-process migrations call `SchemaEpoch::mark_changed`, and every
136        // connection opened before that is closed on its next acquire instead
137        // of being handed out with stale plans.
138        .before_acquire(move |_conn, meta| {
139            let fresh = epoch.opened_after_change(meta.age);
140            Box::pin(async move { Ok(fresh) })
141        })
142}
143
144pub fn connect_options(database_url: &str) -> DatabaseResult<PgConnectOptions> {
145    let options = PgConnectOptions::from_str(database_url)?
146        .application_name("systemprompt")
147        // Why: sqlx 0.9 `PgConnection::get_or_prepare` Parses every persistent
148        // query as a NAMED statement and only sends Close when the cache is
149        // enabled and evicts, so capacity 0 never deallocates them: each
150        // backend grew without bound until Postgres OOMed. A bounded cache
151        // closes on eviction; stale plans after in-process DDL are handled by
152        // `mark_schema_changed` in `build_pool_options`.
153        .statement_cache_capacity(DEFAULT_STATEMENT_CACHE_CAPACITY)
154        // Why: a transaction left open by a crashed or stuck task holds its row
155        // locks and stops autovacuum reclaiming the hot tables until the
156        // connection is recycled (30 min). Postgres ends a session idle inside
157        // a transaction past this bound; no request or job legitimately sits
158        // idle mid-transaction that long.
159        .options([
160            ("client_min_messages", "warning"),
161            ("idle_in_transaction_session_timeout", "600000"),
162        ]);
163    Ok(options)
164}
165
166pub async fn connect_with_retry(
167    options: PgPoolOptions,
168    connect_options: PgConnectOptions,
169) -> DatabaseResult<PgPool> {
170    let connector = |opts: PgConnectOptions| {
171        let options = options.clone();
172        async move { options.connect_with(opts).await }
173    };
174    connect_with_retry_using(connect_options, MAX_ATTEMPTS, RETRY_DELAYS_MS, connector).await
175}
176
177pub async fn connect_with_retry_using<T, F, Fut>(
178    connect_options: PgConnectOptions,
179    max_attempts: u32,
180    delays_ms: &[u64],
181    connector: F,
182) -> DatabaseResult<T>
183where
184    T: Send,
185    F: Fn(PgConnectOptions) -> Fut + Send + Sync,
186    Fut: Future<Output = Result<T, sqlx::Error>> + Send,
187{
188    let cfg = RetryConfig {
189        max_attempts,
190        base_delay: Duration::from_millis(delays_ms.first().copied().unwrap_or(100)),
191        max_delay: Duration::from_millis(delays_ms.iter().copied().max().unwrap_or(1600)),
192        jitter: true,
193    };
194    let classify = |err: &sqlx::Error| {
195        if is_retryable(err) {
196            Outcome::Transient { retry_after: None }
197        } else {
198            Outcome::Permanent
199        }
200    };
201    retry_async(&cfg, "postgres-connect", classify, || {
202        connector(connect_options.clone())
203    })
204    .await
205    .map_err(Into::into)
206}
207
208fn is_retryable(err: &sqlx::Error) -> bool {
209    if let sqlx::Error::Io(io_err) = err
210        && io_err.kind() == std::io::ErrorKind::ConnectionRefused
211    {
212        return true;
213    }
214    let msg = err.to_string();
215    msg.contains("unexpected response from SSLRequest") || msg.contains("starting up")
216}
systemprompt_database/services/postgres/connection.rs

systemprompt_database/services/postgres/
connection.rs
