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::ops::{Deref, DerefMut};
26use std::sync::Arc;
27use std::sync::atomic::{AtomicU64, Ordering};
28use std::time::{Duration, Instant};
29
30use async_trait::async_trait;
31use deadpool::managed::{self, Hook, HookError, HookErrorCause, Object, Pool, RecycleResult};
32use deadpool_postgres::tokio_postgres::{self, Config};
33use deadpool_postgres::{
34    ClientWrapper as DeadpoolClient, Manager as PgManager, ManagerConfig, PoolError,
35    RecyclingMethod, Runtime, Status,
36};
37use mz_ore::cast::{CastFrom, CastLossy};
38use mz_ore::metrics::Counter;
39use mz_ore::now::SYSTEM_TIME;
40use mz_ore::url::SensitiveUrl;
41use tracing::debug;
42
43use crate::error::PostgresError;
44use crate::metrics::PostgresClientMetrics;
45
46/// Configuration knobs for [PostgresClient].
47pub trait PostgresClientKnobs: std::fmt::Debug + Send + Sync {
48    /// Maximum number of connections allowed in a pool.
49    fn connection_pool_max_size(&self) -> usize;
50    /// The maximum time to wait to obtain a connection, if any.
51    fn connection_pool_max_wait(&self) -> Option<Duration>;
52    /// Minimum TTL of a connection. It is expected that connections are
53    /// routinely culled to balance load to the backing store.
54    fn connection_pool_ttl(&self) -> Duration;
55    /// Minimum time between TTLing connections. Helps stagger reconnections
56    /// to avoid stampeding the backing store.
57    fn connection_pool_ttl_stagger(&self) -> Duration;
58    /// Time to wait for a connection to be made before retrying.
59    fn connect_timeout(&self) -> Duration;
60    /// TCP user timeout for connections.
61    fn tcp_user_timeout(&self) -> Duration;
62    /// Amount of idle time before a TCP keepalive packet is sent on a connection.
63    fn keepalives_idle(&self) -> Duration;
64    /// Time interval between TCP keepalive probes.
65    fn keepalives_interval(&self) -> Duration;
66    /// Maximum number of TCP keepalive probes that will be sent before dropping a connection.
67    fn keepalives_retries(&self) -> u32;
68    /// Server-side `statement_timeout` to set on each connection. A value of
69    /// zero is a sentinel that means "do not set a statement timeout".
70    fn statement_timeout(&self) -> Duration;
71}
72
73/// The transaction isolation level applied to new connections.
74#[derive(Clone, Copy, Debug, PartialEq, Eq)]
75pub enum IsolationLevel {
76    /// `SERIALIZABLE` — the strongest level; the historical default for consensus.
77    Serializable,
78    /// `READ COMMITTED` — for callers (e.g. consensus) whose queries are correct without
79    /// serializable isolation (relying instead on the `PRIMARY KEY` / `FOR UPDATE` / `ON CONFLICT`).
80    ReadCommitted,
81}
82
83impl IsolationLevel {
84    /// The `SET SESSION CHARACTERISTICS` statement that selects this isolation level.
85    fn set_characteristics_sql(self) -> &'static str {
86        match self {
87            IsolationLevel::Serializable => {
88                "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE"
89            }
90            IsolationLevel::ReadCommitted => {
91                "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL READ COMMITTED"
92            }
93        }
94    }
95}
96
97/// Resolves the isolation level to apply to a connection. It is invoked once per connection
98/// creation, so a dyncfg-backed resolver lets a change take effect as the pool cycles connections.
99pub type IsolationLevelFn = Arc<dyn Fn() -> IsolationLevel + Send + Sync>;
100
101/// A connection handed out by [`PostgresClient::get_connection`]. Dereferences to a [`Client`],
102/// which additionally records the [`IsolationLevel`] the connection was created under.
103pub type Connection = Object<Manager>;
104
105/// A pooled Postgres connection tagged with the [`IsolationLevel`] it was created under.
106///
107/// The isolation level is applied once at creation and is fixed for the life of the connection.
108#[derive(Debug)]
109pub struct Client {
110    inner: DeadpoolClient,
111    isolation: IsolationLevel,
112}
113
114impl Client {
115    /// The [`IsolationLevel`] this connection was configured with when it was created.
116    pub fn isolation_level(&self) -> IsolationLevel {
117        self.isolation
118    }
119}
120
121impl Deref for Client {
122    type Target = DeadpoolClient;
123
124    fn deref(&self) -> &DeadpoolClient {
125        &self.inner
126    }
127}
128
129impl DerefMut for Client {
130    fn deref_mut(&mut self) -> &mut DeadpoolClient {
131        &mut self.inner
132    }
133}
134
135/// A deadpool [`managed::Manager`] wrapping [`deadpool_postgres::Manager`]. It applies a
136/// per-connection isolation level at creation and records that level on every [`Client`] it hands
137/// out.
138pub struct Manager {
139    inner: PgManager,
140    /// Resolves the isolation level to apply. Invoked once per connection so a dyncfg-backed
141    /// resolver takes effect as the pool cycles connections.
142    isolation: IsolationLevelFn,
143    knobs: Arc<dyn PostgresClientKnobs>,
144    connections_created: Counter,
145}
146
147impl std::fmt::Debug for Manager {
148    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149        f.debug_struct("Manager")
150            .field("inner", &self.inner)
151            .finish_non_exhaustive()
152    }
153}
154
155#[async_trait]
156impl managed::Manager for Manager {
157    type Type = Client;
158    type Error = tokio_postgres::Error;
159
160    async fn create(&self) -> Result<Client, tokio_postgres::Error> {
161        let inner = self.inner.create().await?;
162        self.connections_created.inc();
163
164        // Resolved per connection so a dyncfg-backed isolation level takes effect as the pool
165        // cycles connections. Defaults to SERIALIZABLE (see `PostgresClientConfig::new`).
166        let isolation = (self.isolation)();
167        let mut setup = isolation.set_characteristics_sql().to_owned();
168        // A zero `statement_timeout` is our sentinel for "leave it unset". We only emit the `SET`
169        // when non-zero so we don't override a timeout configured out of band.
170        let statement_timeout = self.knobs.statement_timeout();
171        if !statement_timeout.is_zero() {
172            // A bare integer value for `statement_timeout` is interpreted as milliseconds.
173            write!(
174                setup,
175                "; SET statement_timeout = {}",
176                statement_timeout.as_millis()
177            )
178            .expect("writing to a String never fails");
179        }
180        debug!("opened new postgres connection");
181        // This must surface as `tokio_postgres::Error` (the pool's error type); using
182        // `mz_postgres_util` wrappers would change the error type.
183        #[allow(clippy::disallowed_methods)]
184        inner.batch_execute(&setup).await?;
185
186        Ok(Client { inner, isolation })
187    }
188
189    async fn recycle(&self, client: &mut Client) -> RecycleResult<tokio_postgres::Error> {
190        self.inner.recycle(&mut client.inner).await
191    }
192
193    fn detach(&self, client: &mut Client) {
194        self.inner.detach(&mut client.inner)
195    }
196}
197
198/// Configuration for creating a [PostgresClient].
199#[derive(Clone)]
200pub struct PostgresClientConfig {
201    url: SensitiveUrl,
202    knobs: Arc<dyn PostgresClientKnobs>,
203    metrics: PostgresClientMetrics,
204    isolation: IsolationLevelFn,
205}
206
207impl std::fmt::Debug for PostgresClientConfig {
208    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
209        f.debug_struct("PostgresClientConfig")
210            .field("url", &self.url)
211            .finish_non_exhaustive()
212    }
213}
214
215impl PostgresClientConfig {
216    /// Returns a new [PostgresClientConfig] for use in production. Connections default to
217    /// `SERIALIZABLE`; use [PostgresClientConfig::with_isolation] to override.
218    pub fn new(
219        url: SensitiveUrl,
220        knobs: Arc<dyn PostgresClientKnobs>,
221        metrics: PostgresClientMetrics,
222    ) -> Self {
223        PostgresClientConfig {
224            url,
225            knobs,
226            metrics,
227            isolation: Arc::new(|| IsolationLevel::Serializable),
228        }
229    }
230
231    /// Sets the resolver that picks the isolation level applied to each new connection.
232    pub fn with_isolation(mut self, isolation: IsolationLevelFn) -> Self {
233        self.isolation = isolation;
234        self
235    }
236}
237
238/// A Postgres client wrapper that uses deadpool as a connection pool.
239pub struct PostgresClient {
240    pool: Pool<Manager>,
241    metrics: PostgresClientMetrics,
242}
243
244impl std::fmt::Debug for PostgresClient {
245    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
246        f.debug_struct("PostgresClient").finish_non_exhaustive()
247    }
248}
249
250impl PostgresClient {
251    /// Open a [PostgresClient] using the given `config`.
252    pub fn open(config: PostgresClientConfig) -> Result<Self, PostgresError> {
253        let mut pg_config: Config = config.url.to_string_unredacted().parse()?;
254        pg_config.connect_timeout(config.knobs.connect_timeout());
255        pg_config.tcp_user_timeout(config.knobs.tcp_user_timeout());
256
257        // Configuring keepalives is important to ensure we can detect broken connections quickly.
258        // TCP_USER_TIMEOUT is not sufficient as it only enforces a timeout on ACKs for transmitted
259        // data, which only helps if we... transmit data.
260        pg_config.keepalives(true);
261        pg_config.keepalives_idle(config.knobs.keepalives_idle());
262        pg_config.keepalives_interval(config.knobs.keepalives_interval());
263        pg_config.keepalives_retries(config.knobs.keepalives_retries());
264
265        let tls = mz_tls_util::make_tls(&pg_config).map_err(|tls_err| match tls_err {
266            mz_tls_util::TlsError::Generic(e) => PostgresError::Indeterminate(e),
267            mz_tls_util::TlsError::OpenSsl(e) => PostgresError::Indeterminate(anyhow::anyhow!(e)),
268        })?;
269
270        let pg_manager = PgManager::from_config(
271            pg_config,
272            tls,
273            ManagerConfig {
274                recycling_method: RecyclingMethod::Fast,
275            },
276        );
277        // The isolation level and `statement_timeout` are applied inside `Manager::create` so the
278        // resolved level can be recorded on each connection it hands out.
279        let manager = Manager {
280            inner: pg_manager,
281            isolation: Arc::clone(&config.isolation),
282            knobs: Arc::clone(&config.knobs),
283            connections_created: config.metrics.connpool_connections_created.clone(),
284        };
285
286        let last_ttl_connection = AtomicU64::new(0);
287        let ttl_reconnections = config.metrics.connpool_ttl_reconnections.clone();
288        let builder = Pool::builder(manager);
289        let builder = match config.knobs.connection_pool_max_wait() {
290            None => builder,
291            Some(wait) => builder.wait_timeout(Some(wait)).runtime(Runtime::Tokio1),
292        };
293        let pool = builder
294            .max_size(config.knobs.connection_pool_max_size())
295            .pre_recycle(Hook::sync_fn(move |_client, conn_metrics| {
296                // proactively TTL connections to rebalance load to Postgres/CRDB. this helps
297                // fix skew when downstream DB operations (e.g. CRDB rolling restart) result
298                // in uneven load to each node, and works to reduce the # of connections
299                // maintained by the pool after bursty workloads.
300
301                // add a bias towards TTLing older connections first
302                if conn_metrics.age() < config.knobs.connection_pool_ttl() {
303                    return Ok(());
304                }
305
306                let last_ttl = last_ttl_connection.load(Ordering::SeqCst);
307                let now = (SYSTEM_TIME)();
308                let elapsed_since_last_ttl = Duration::from_millis(now.saturating_sub(last_ttl));
309
310                // stagger out reconnections to avoid stampeding the DB
311                if elapsed_since_last_ttl > config.knobs.connection_pool_ttl_stagger()
312                    && last_ttl_connection
313                        .compare_exchange_weak(last_ttl, now, Ordering::SeqCst, Ordering::SeqCst)
314                        .is_ok()
315                {
316                    ttl_reconnections.inc();
317                    return Err(HookError::Continue(Some(HookErrorCause::Message(
318                        "connection has been TTLed".to_string(),
319                    ))));
320                }
321
322                Ok(())
323            }))
324            .build()
325            .expect("postgres connection pool built with incorrect parameters");
326
327        Ok(PostgresClient {
328            pool,
329            metrics: config.metrics,
330        })
331    }
332
333    fn status_metrics(&self, status: Status) {
334        self.metrics
335            .connpool_available
336            .set(f64::cast_lossy(status.available));
337        self.metrics.connpool_size.set(u64::cast_from(status.size));
338        // Don't bother reporting the maximum size of the pool... we know that from config.
339    }
340
341    /// Gets connection from the pool or waits for one to become available.
342    pub async fn get_connection(&self) -> Result<Connection, PoolError> {
343        let start = Instant::now();
344        // note that getting the pool size here requires briefly locking the pool
345        self.status_metrics(self.pool.status());
346        let res = self.pool.get().await;
347        if let Err(PoolError::Backend(err)) = &res {
348            debug!("error establishing connection: {}", err);
349            self.metrics.connpool_connection_errors.inc();
350        }
351        self.metrics
352            .connpool_acquire_seconds
353            .inc_by(start.elapsed().as_secs_f64());
354        self.metrics.connpool_acquires.inc();
355        self.status_metrics(self.pool.status());
356        res
357    }
358}
mz_postgres_client/lib.rs

mz_postgres_client/
lib.rs
