Skip to main content

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
10use std::future::Future;
11use std::time::Duration;
12
13use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions};
14
15use crate::error::DatabaseResult;
16use crate::resilience::classify::Outcome;
17use crate::resilience::config::RetryConfig;
18use crate::resilience::retry::retry_async;
19
20const RETRY_DELAYS_MS: &[u64] = &[100, 200, 400, 800, 1600];
21const MAX_ATTEMPTS: u32 = 5;
22
23/// Operator-tunable connection-pool sizing for a `PostgresProvider`.
24///
25/// [`PoolConfig::default`] reproduces the historical hardcoded values; callers
26/// that have profile config supply their own. The connect/SSL/retry behaviour
27/// is fixed and not exposed here — only the sizing/lifetime knobs an operator
28/// needs to fit the pool to their Postgres `max_connections` and replica count.
29#[derive(Debug, Clone, Copy)]
30pub struct PoolConfig {
31    pub max_connections: u32,
32    pub min_connections: u32,
33    pub acquire_timeout: Duration,
34    pub idle_timeout: Duration,
35    pub max_lifetime: Duration,
36}
37
38impl Default for PoolConfig {
39    fn default() -> Self {
40        Self {
41            max_connections: 50,
42            min_connections: 0,
43            acquire_timeout: Duration::from_secs(30),
44            idle_timeout: Duration::from_secs(300),
45            max_lifetime: Duration::from_secs(1800),
46        }
47    }
48}
49
50#[must_use]
51pub fn build_pool_options(cfg: &PoolConfig) -> PgPoolOptions {
52    PgPoolOptions::new()
53        .max_connections(cfg.max_connections)
54        .min_connections(cfg.min_connections)
55        .max_lifetime(cfg.max_lifetime)
56        .acquire_timeout(cfg.acquire_timeout)
57        .idle_timeout(cfg.idle_timeout)
58}
59
60pub async fn connect_with_retry(
61    options: PgPoolOptions,
62    connect_options: PgConnectOptions,
63) -> DatabaseResult<PgPool> {
64    let connector = |opts: PgConnectOptions| {
65        let options = options.clone();
66        async move { options.connect_with(opts).await }
67    };
68    connect_with_retry_using(connect_options, MAX_ATTEMPTS, RETRY_DELAYS_MS, connector).await
69}
70
71pub async fn connect_with_retry_using<T, F, Fut>(
72    connect_options: PgConnectOptions,
73    max_attempts: u32,
74    delays_ms: &[u64],
75    connector: F,
76) -> DatabaseResult<T>
77where
78    T: Send,
79    F: Fn(PgConnectOptions) -> Fut + Send + Sync,
80    Fut: Future<Output = Result<T, sqlx::Error>> + Send,
81{
82    let cfg = RetryConfig {
83        max_attempts,
84        base_delay: Duration::from_millis(delays_ms.first().copied().unwrap_or(100)),
85        max_delay: Duration::from_millis(delays_ms.iter().copied().max().unwrap_or(1600)),
86        jitter: false,
87    };
88    let classify = |err: &sqlx::Error| {
89        if is_retryable(err) {
90            Outcome::Transient { retry_after: None }
91        } else {
92            Outcome::Permanent
93        }
94    };
95    retry_async(&cfg, "postgres-connect", classify, || {
96        connector(connect_options.clone())
97    })
98    .await
99    .map_err(Into::into)
100}
101
102fn is_retryable(err: &sqlx::Error) -> bool {
103    if let sqlx::Error::Io(io_err) = err
104        && io_err.kind() == std::io::ErrorKind::ConnectionRefused
105    {
106        return true;
107    }
108    let msg = err.to_string();
109    msg.contains("unexpected response from SSLRequest") || msg.contains("starting up")
110}