mz_postgres_client/
lib.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! A Postgres client that uses deadpool as a connection pool and comes with
11//! common/default configuration options.
12
13#![warn(missing_docs, missing_debug_implementations)]
14#![warn(
15    clippy::cast_possible_truncation,
16    clippy::cast_precision_loss,
17    clippy::cast_sign_loss,
18    clippy::clone_on_ref_ptr
19)]
20
21pub mod error;
22pub mod metrics;
23
24use std::fmt::Write;
25use std::sync::Arc;
26use std::sync::atomic::{AtomicU64, Ordering};
27use std::time::{Duration, Instant};
28
29use deadpool_postgres::tokio_postgres::Config;
30use deadpool_postgres::{
31    Hook, HookError, HookErrorCause, Manager, ManagerConfig, Object, Pool, PoolError,
32    RecyclingMethod, Runtime, Status,
33};
34use mz_ore::cast::{CastFrom, CastLossy};
35use mz_ore::now::SYSTEM_TIME;
36use mz_ore::url::SensitiveUrl;
37use tracing::debug;
38
39use crate::error::PostgresError;
40use crate::metrics::PostgresClientMetrics;
41
42/// Configuration knobs for [PostgresClient].
43pub trait PostgresClientKnobs: std::fmt::Debug + Send + Sync {
44    /// Maximum number of connections allowed in a pool.
45    fn connection_pool_max_size(&self) -> usize;
46    /// The maximum time to wait to obtain a connection, if any.
47    fn connection_pool_max_wait(&self) -> Option<Duration>;
48    /// Minimum TTL of a connection. It is expected that connections are
49    /// routinely culled to balance load to the backing store.
50    fn connection_pool_ttl(&self) -> Duration;
51    /// Minimum time between TTLing connections. Helps stagger reconnections
52    /// to avoid stampeding the backing store.
53    fn connection_pool_ttl_stagger(&self) -> Duration;
54    /// Time to wait for a connection to be made before retrying.
55    fn connect_timeout(&self) -> Duration;
56    /// TCP user timeout for connections.
57    fn tcp_user_timeout(&self) -> Duration;
58    /// Amount of idle time before a TCP keepalive packet is sent on a connection.
59    fn keepalives_idle(&self) -> Duration;
60    /// Time interval between TCP keepalive probes.
61    fn keepalives_interval(&self) -> Duration;
62    /// Maximum number of TCP keepalive probes that will be sent before dropping a connection.
63    fn keepalives_retries(&self) -> u32;
64    /// Server-side `statement_timeout` to set on each connection. A value of
65    /// zero is a sentinel that means "do not set a statement timeout".
66    fn statement_timeout(&self) -> Duration;
67}
68
69/// The transaction isolation level applied to new connections.
70#[derive(Clone, Copy, Debug, PartialEq, Eq)]
71pub enum IsolationLevel {
72    /// `SERIALIZABLE` — the strongest level; the historical default for consensus.
73    Serializable,
74    /// `READ COMMITTED` — for callers (e.g. consensus) whose queries are correct without
75    /// serializable isolation (relying instead on the `PRIMARY KEY` / `FOR UPDATE` / `ON CONFLICT`).
76    ReadCommitted,
77}
78
79impl IsolationLevel {
80    /// The `SET SESSION CHARACTERISTICS` statement that selects this isolation level.
81    fn set_characteristics_sql(self) -> &'static str {
82        match self {
83            IsolationLevel::Serializable => {
84                "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE"
85            }
86            IsolationLevel::ReadCommitted => {
87                "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL READ COMMITTED"
88            }
89        }
90    }
91}
92
93/// Resolves the isolation level to apply to a connection. It is invoked once per connection
94/// creation, so a dyncfg-backed resolver lets a change take effect as the pool cycles connections.
95pub type IsolationLevelFn = Arc<dyn Fn() -> IsolationLevel + Send + Sync>;
96
97/// Configuration for creating a [PostgresClient].
98#[derive(Clone)]
99pub struct PostgresClientConfig {
100    url: SensitiveUrl,
101    knobs: Arc<dyn PostgresClientKnobs>,
102    metrics: PostgresClientMetrics,
103    isolation: IsolationLevelFn,
104}
105
106impl std::fmt::Debug for PostgresClientConfig {
107    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
108        f.debug_struct("PostgresClientConfig")
109            .field("url", &self.url)
110            .finish_non_exhaustive()
111    }
112}
113
114impl PostgresClientConfig {
115    /// Returns a new [PostgresClientConfig] for use in production. Connections default to
116    /// `SERIALIZABLE`; use [PostgresClientConfig::with_isolation] to override.
117    pub fn new(
118        url: SensitiveUrl,
119        knobs: Arc<dyn PostgresClientKnobs>,
120        metrics: PostgresClientMetrics,
121    ) -> Self {
122        PostgresClientConfig {
123            url,
124            knobs,
125            metrics,
126            isolation: Arc::new(|| IsolationLevel::Serializable),
127        }
128    }
129
130    /// Sets the resolver that picks the isolation level applied to each new connection.
131    pub fn with_isolation(mut self, isolation: IsolationLevelFn) -> Self {
132        self.isolation = isolation;
133        self
134    }
135}
136
137/// A Postgres client wrapper that uses deadpool as a connection pool.
138pub struct PostgresClient {
139    pool: Pool,
140    metrics: PostgresClientMetrics,
141}
142
143impl std::fmt::Debug for PostgresClient {
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        f.debug_struct("PostgresClient").finish_non_exhaustive()
146    }
147}
148
149impl PostgresClient {
150    /// Open a [PostgresClient] using the given `config`.
151    pub fn open(config: PostgresClientConfig) -> Result<Self, PostgresError> {
152        let mut pg_config: Config = config.url.to_string_unredacted().parse()?;
153        pg_config.connect_timeout(config.knobs.connect_timeout());
154        pg_config.tcp_user_timeout(config.knobs.tcp_user_timeout());
155
156        // Configuring keepalives is important to ensure we can detect broken connections quickly.
157        // TCP_USER_TIMEOUT is not sufficient as it only enforces a timeout on ACKs for transmitted
158        // data, which only helps if we... transmit data.
159        pg_config.keepalives(true);
160        pg_config.keepalives_idle(config.knobs.keepalives_idle());
161        pg_config.keepalives_interval(config.knobs.keepalives_interval());
162        pg_config.keepalives_retries(config.knobs.keepalives_retries());
163
164        let tls = mz_tls_util::make_tls(&pg_config).map_err(|tls_err| match tls_err {
165            mz_tls_util::TlsError::Generic(e) => PostgresError::Indeterminate(e),
166            mz_tls_util::TlsError::OpenSsl(e) => PostgresError::Indeterminate(anyhow::anyhow!(e)),
167        })?;
168
169        let manager = Manager::from_config(
170            pg_config,
171            tls,
172            ManagerConfig {
173                recycling_method: RecyclingMethod::Fast,
174            },
175        );
176
177        let last_ttl_connection = AtomicU64::new(0);
178        let connections_created = config.metrics.connpool_connections_created.clone();
179        let ttl_reconnections = config.metrics.connpool_ttl_reconnections.clone();
180        let knobs = Arc::clone(&config.knobs);
181        let isolation = Arc::clone(&config.isolation);
182        let builder = Pool::builder(manager);
183        let builder = match config.knobs.connection_pool_max_wait() {
184            None => builder,
185            Some(wait) => builder.wait_timeout(Some(wait)).runtime(Runtime::Tokio1),
186        };
187        let pool = builder
188            .max_size(config.knobs.connection_pool_max_size())
189            .post_create(Hook::async_fn(move |client, _| {
190                connections_created.inc();
191                let knobs = Arc::clone(&knobs);
192                // Resolved per connection so a dyncfg-backed isolation level takes effect as the
193                // pool cycles connections. Defaults to SERIALIZABLE (see `new`).
194                let isolation = Arc::clone(&isolation);
195                Box::pin(async move {
196                    debug!("opened new consensus postgres connection");
197                    let mut setup = isolation().set_characteristics_sql().to_owned();
198                    // A zero `statement_timeout` is our sentinel for "leave it
199                    // unset". We only emit the `SET` when non-zero so we don't
200                    // override a timeout configured out of band.
201                    let statement_timeout = knobs.statement_timeout();
202                    if !statement_timeout.is_zero() {
203                        // A bare integer value for `statement_timeout` is
204                        // interpreted as milliseconds.
205                        write!(
206                            setup,
207                            "; SET statement_timeout = {}",
208                            statement_timeout.as_millis()
209                        )
210                        .expect("writing to a String never fails");
211                    }
212                    // This hook must return `tokio_postgres::Error`; using
213                    // `mz_postgres_util` wrappers would change the error type.
214                    #[allow(clippy::disallowed_methods)]
215                    client
216                        .batch_execute(&setup)
217                        .await
218                        .map_err(|e| HookError::Abort(HookErrorCause::Backend(e)))
219                })
220            }))
221            .pre_recycle(Hook::sync_fn(move |_client, conn_metrics| {
222                // proactively TTL connections to rebalance load to Postgres/CRDB. this helps
223                // fix skew when downstream DB operations (e.g. CRDB rolling restart) result
224                // in uneven load to each node, and works to reduce the # of connections
225                // maintained by the pool after bursty workloads.
226
227                // add a bias towards TTLing older connections first
228                if conn_metrics.age() < config.knobs.connection_pool_ttl() {
229                    return Ok(());
230                }
231
232                let last_ttl = last_ttl_connection.load(Ordering::SeqCst);
233                let now = (SYSTEM_TIME)();
234                let elapsed_since_last_ttl = Duration::from_millis(now.saturating_sub(last_ttl));
235
236                // stagger out reconnections to avoid stampeding the DB
237                if elapsed_since_last_ttl > config.knobs.connection_pool_ttl_stagger()
238                    && last_ttl_connection
239                        .compare_exchange_weak(last_ttl, now, Ordering::SeqCst, Ordering::SeqCst)
240                        .is_ok()
241                {
242                    ttl_reconnections.inc();
243                    return Err(HookError::Continue(Some(HookErrorCause::Message(
244                        "connection has been TTLed".to_string(),
245                    ))));
246                }
247
248                Ok(())
249            }))
250            .build()
251            .expect("postgres connection pool built with incorrect parameters");
252
253        Ok(PostgresClient {
254            pool,
255            metrics: config.metrics,
256        })
257    }
258
259    fn status_metrics(&self, status: Status) {
260        self.metrics
261            .connpool_available
262            .set(f64::cast_lossy(status.available));
263        self.metrics.connpool_size.set(u64::cast_from(status.size));
264        // Don't bother reporting the maximum size of the pool... we know that from config.
265    }
266
267    /// Gets connection from the pool or waits for one to become available.
268    pub async fn get_connection(&self) -> Result<Object, PoolError> {
269        let start = Instant::now();
270        // note that getting the pool size here requires briefly locking the pool
271        self.status_metrics(self.pool.status());
272        let res = self.pool.get().await;
273        if let Err(PoolError::Backend(err)) = &res {
274            debug!("error establishing connection: {}", err);
275            self.metrics.connpool_connection_errors.inc();
276        }
277        self.metrics
278            .connpool_acquire_seconds
279            .inc_by(start.elapsed().as_secs_f64());
280        self.metrics.connpool_acquires.inc();
281        self.status_metrics(self.pool.status());
282        res
283    }
284}
mz_postgres_client/lib.rs

mz_postgres_client/
lib.rs
