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 deadpool::managed::{self, Hook, HookError, Metrics, Object, Pool, RecycleResult};
31use deadpool_postgres::tokio_postgres::{self, Config};
32use deadpool_postgres::{
33    ClientWrapper as DeadpoolClient, Manager as PgManager, ManagerConfig, PoolError,
34    RecyclingMethod, Runtime, Status,
35};
36use mz_ore::cast::{CastFrom, CastLossy};
37use mz_ore::metrics::Counter;
38use mz_ore::now::SYSTEM_TIME;
39use mz_ore::url::SensitiveUrl;
40use tracing::debug;
41
42use crate::error::PostgresError;
43use crate::metrics::PostgresClientMetrics;
44
45/// Configuration knobs for [PostgresClient].
46pub trait PostgresClientKnobs: std::fmt::Debug + Send + Sync {
47    /// Maximum number of connections allowed in a pool.
48    fn connection_pool_max_size(&self) -> usize;
49    /// The maximum time to wait to obtain a connection, if any.
50    fn connection_pool_max_wait(&self) -> Option<Duration>;
51    /// Minimum TTL of a connection. It is expected that connections are
52    /// routinely culled to balance load to the backing store.
53    fn connection_pool_ttl(&self) -> Duration;
54    /// Minimum time between TTLing connections. Helps stagger reconnections
55    /// to avoid stampeding the backing store.
56    fn connection_pool_ttl_stagger(&self) -> Duration;
57    /// Time to wait for a connection to be made before retrying.
58    fn connect_timeout(&self) -> Duration;
59    /// TCP user timeout for connections.
60    fn tcp_user_timeout(&self) -> Duration;
61    /// Amount of idle time before a TCP keepalive packet is sent on a connection.
62    fn keepalives_idle(&self) -> Duration;
63    /// Time interval between TCP keepalive probes.
64    fn keepalives_interval(&self) -> Duration;
65    /// Maximum number of TCP keepalive probes that will be sent before dropping a connection.
66    fn keepalives_retries(&self) -> u32;
67    /// Server-side `statement_timeout` to set on each connection. A value of
68    /// zero is a sentinel that means "do not set a statement timeout".
69    fn statement_timeout(&self) -> Duration;
70}
71
72/// The transaction isolation level applied to new connections.
73#[derive(Clone, Copy, Debug, PartialEq, Eq)]
74pub enum IsolationLevel {
75    /// `SERIALIZABLE` — the strongest level; the historical default for consensus.
76    Serializable,
77    /// `READ COMMITTED` — for callers (e.g. consensus) whose queries are correct without
78    /// serializable isolation (relying instead on the `PRIMARY KEY` / `FOR UPDATE` / `ON CONFLICT`).
79    ReadCommitted,
80}
81
82impl IsolationLevel {
83    /// The `SET SESSION CHARACTERISTICS` statement that selects this isolation level.
84    fn set_characteristics_sql(self) -> &'static str {
85        match self {
86            IsolationLevel::Serializable => {
87                "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE"
88            }
89            IsolationLevel::ReadCommitted => {
90                "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL READ COMMITTED"
91            }
92        }
93    }
94}
95
96/// Resolves the isolation level to apply to a connection. It is invoked once per connection
97/// creation, so a dyncfg-backed resolver lets a change take effect as the pool cycles connections.
98pub type IsolationLevelFn = Arc<dyn Fn() -> IsolationLevel + Send + Sync>;
99
100/// A connection handed out by [`PostgresClient::get_connection`]. Dereferences to a [`Client`],
101/// which additionally records the [`IsolationLevel`] the connection was created under.
102pub type Connection = Object<Manager>;
103
104/// A pooled Postgres connection tagged with the [`IsolationLevel`] it was created under.
105///
106/// The isolation level is applied once at creation and is fixed for the life of the connection.
107#[derive(Debug)]
108pub struct Client {
109    inner: DeadpoolClient,
110    isolation: IsolationLevel,
111}
112
113impl Client {
114    /// The [`IsolationLevel`] this connection was configured with when it was created.
115    pub fn isolation_level(&self) -> IsolationLevel {
116        self.isolation
117    }
118}
119
120impl Deref for Client {
121    type Target = DeadpoolClient;
122
123    fn deref(&self) -> &DeadpoolClient {
124        &self.inner
125    }
126}
127
128impl DerefMut for Client {
129    fn deref_mut(&mut self) -> &mut DeadpoolClient {
130        &mut self.inner
131    }
132}
133
134/// A deadpool [`managed::Manager`] wrapping [`deadpool_postgres::Manager`]. It applies a
135/// per-connection isolation level at creation and records that level on every [`Client`] it hands
136/// out.
137pub struct Manager {
138    inner: PgManager,
139    /// Resolves the isolation level to apply. Invoked once per connection so a dyncfg-backed
140    /// resolver takes effect as the pool cycles connections.
141    isolation: IsolationLevelFn,
142    knobs: Arc<dyn PostgresClientKnobs>,
143    connections_created: Counter,
144}
145
146impl std::fmt::Debug for Manager {
147    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
148        f.debug_struct("Manager")
149            .field("inner", &self.inner)
150            .finish_non_exhaustive()
151    }
152}
153
154impl managed::Manager for Manager {
155    type Type = Client;
156    type Error = tokio_postgres::Error;
157
158    async fn create(&self) -> Result<Client, tokio_postgres::Error> {
159        let inner = self.inner.create().await?;
160        self.connections_created.inc();
161
162        // Resolved per connection so a dyncfg-backed isolation level takes effect as the pool
163        // cycles connections. Defaults to SERIALIZABLE (see `PostgresClientConfig::new`).
164        let isolation = (self.isolation)();
165        let mut setup = isolation.set_characteristics_sql().to_owned();
166        // A zero `statement_timeout` is our sentinel for "leave it unset". We only emit the `SET`
167        // when non-zero so we don't override a timeout configured out of band.
168        let statement_timeout = self.knobs.statement_timeout();
169        if !statement_timeout.is_zero() {
170            // A bare integer value for `statement_timeout` is interpreted as milliseconds.
171            write!(
172                setup,
173                "; SET statement_timeout = {}",
174                statement_timeout.as_millis()
175            )
176            .expect("writing to a String never fails");
177        }
178        debug!("opened new postgres connection");
179        // This must surface as `tokio_postgres::Error` (the pool's error type); using
180        // `mz_postgres_util` wrappers would change the error type.
181        #[allow(clippy::disallowed_methods)]
182        inner.batch_execute(&setup).await?;
183
184        Ok(Client { inner, isolation })
185    }
186
187    async fn recycle(
188        &self,
189        client: &mut Client,
190        metrics: &Metrics,
191    ) -> RecycleResult<tokio_postgres::Error> {
192        self.inner.recycle(&mut client.inner, metrics).await
193    }
194
195    fn detach(&self, client: &mut Client) {
196        self.inner.detach(&mut client.inner)
197    }
198}
199
200/// Configuration for creating a [PostgresClient].
201#[derive(Clone)]
202pub struct PostgresClientConfig {
203    url: SensitiveUrl,
204    knobs: Arc<dyn PostgresClientKnobs>,
205    metrics: PostgresClientMetrics,
206    isolation: IsolationLevelFn,
207}
208
209impl std::fmt::Debug for PostgresClientConfig {
210    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
211        f.debug_struct("PostgresClientConfig")
212            .field("url", &self.url)
213            .finish_non_exhaustive()
214    }
215}
216
217impl PostgresClientConfig {
218    /// Returns a new [PostgresClientConfig] for use in production. Connections default to
219    /// `SERIALIZABLE`; use [PostgresClientConfig::with_isolation] to override.
220    pub fn new(
221        url: SensitiveUrl,
222        knobs: Arc<dyn PostgresClientKnobs>,
223        metrics: PostgresClientMetrics,
224    ) -> Self {
225        PostgresClientConfig {
226            url,
227            knobs,
228            metrics,
229            isolation: Arc::new(|| IsolationLevel::Serializable),
230        }
231    }
232
233    /// Sets the resolver that picks the isolation level applied to each new connection.
234    pub fn with_isolation(mut self, isolation: IsolationLevelFn) -> Self {
235        self.isolation = isolation;
236        self
237    }
238}
239
240/// A Postgres client wrapper that uses deadpool as a connection pool.
241pub struct PostgresClient {
242    pool: Pool<Manager>,
243    knobs: Arc<dyn PostgresClientKnobs>,
244    metrics: PostgresClientMetrics,
245}
246
247impl std::fmt::Debug for PostgresClient {
248    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
249        f.debug_struct("PostgresClient").finish_non_exhaustive()
250    }
251}
252
253impl PostgresClient {
254    /// Open a [PostgresClient] using the given `config`.
255    pub fn open(config: PostgresClientConfig) -> Result<Self, PostgresError> {
256        let mut pg_config: Config = config.url.to_string_unredacted().parse()?;
257        pg_config.connect_timeout(config.knobs.connect_timeout());
258        pg_config.tcp_user_timeout(config.knobs.tcp_user_timeout());
259
260        // Configuring keepalives is important to ensure we can detect broken connections quickly.
261        // TCP_USER_TIMEOUT is not sufficient as it only enforces a timeout on ACKs for transmitted
262        // data, which only helps if we... transmit data.
263        pg_config.keepalives(true);
264        pg_config.keepalives_idle(config.knobs.keepalives_idle());
265        pg_config.keepalives_interval(config.knobs.keepalives_interval());
266        pg_config.keepalives_retries(config.knobs.keepalives_retries());
267
268        let tls = mz_tls_util::make_tls(&pg_config).map_err(|tls_err| match tls_err {
269            mz_tls_util::TlsError::Generic(e) => PostgresError::Indeterminate(e),
270            mz_tls_util::TlsError::OpenSsl(e) => PostgresError::Indeterminate(anyhow::anyhow!(e)),
271        })?;
272
273        let pg_manager = PgManager::from_config(
274            pg_config,
275            tls,
276            ManagerConfig {
277                recycling_method: RecyclingMethod::Fast,
278            },
279        );
280        // The isolation level and `statement_timeout` are applied inside `Manager::create` so the
281        // resolved level can be recorded on each connection it hands out.
282        let manager = Manager {
283            inner: pg_manager,
284            isolation: Arc::clone(&config.isolation),
285            knobs: Arc::clone(&config.knobs),
286            connections_created: config.metrics.connpool_connections_created.clone(),
287        };
288
289        let last_ttl_connection = AtomicU64::new(0);
290        let ttl_reconnections = config.metrics.connpool_ttl_reconnections.clone();
291        let knobs = Arc::clone(&config.knobs);
292        let builder = Pool::builder(manager);
293        let builder = match config.knobs.connection_pool_max_wait() {
294            None => builder,
295            Some(wait) => builder.wait_timeout(Some(wait)).runtime(Runtime::Tokio1),
296        };
297        let pool = builder
298            .max_size(config.knobs.connection_pool_max_size())
299            .pre_recycle(Hook::sync_fn(move |_client, conn_metrics| {
300                // proactively TTL connections to rebalance load to Postgres/CRDB. this helps
301                // fix skew when downstream DB operations (e.g. CRDB rolling restart) result
302                // in uneven load to each node, and works to reduce the # of connections
303                // maintained by the pool after bursty workloads.
304
305                // add a bias towards TTLing older connections first
306                if conn_metrics.age() < config.knobs.connection_pool_ttl() {
307                    return Ok(());
308                }
309
310                let last_ttl = last_ttl_connection.load(Ordering::SeqCst);
311                let now = (SYSTEM_TIME)();
312                let elapsed_since_last_ttl = Duration::from_millis(now.saturating_sub(last_ttl));
313
314                // stagger out reconnections to avoid stampeding the DB
315                if elapsed_since_last_ttl > config.knobs.connection_pool_ttl_stagger()
316                    && last_ttl_connection
317                        .compare_exchange_weak(last_ttl, now, Ordering::SeqCst, Ordering::SeqCst)
318                        .is_ok()
319                {
320                    ttl_reconnections.inc();
321                    // A `pre_recycle` error discards the connection and the pool moves on to the
322                    // next one.
323                    return Err(HookError::message("connection has been TTLed"));
324                }
325
326                Ok(())
327            }))
328            .build()
329            .expect("postgres connection pool built with incorrect parameters");
330
331        Ok(PostgresClient {
332            pool,
333            knobs,
334            metrics: config.metrics,
335        })
336    }
337
338    fn status_metrics(&self, status: Status) {
339        // Negative when tasks are waiting for a connection.
340        self.metrics
341            .connpool_available
342            .set(f64::cast_lossy(status.available) - f64::cast_lossy(status.waiting));
343        self.metrics.connpool_size.set(u64::cast_from(status.size));
344        // Don't bother reporting the maximum size of the pool... we know that from config.
345    }
346
347    /// The current [`Status`] of the connection pool.
348    ///
349    /// NOTE: this briefly locks the pool.
350    pub fn status(&self) -> Status {
351        self.pool.status()
352    }
353
354    /// Gets connection from the pool or waits for one to become available.
355    pub async fn get_connection(&self) -> Result<Connection, PoolError> {
356        let start = Instant::now();
357        // note that getting the pool status here requires briefly locking the pool
358        let status = self.pool.status();
359        // Apply knob changes to the pool cap without a restart, so an operator
360        // can grow (or shrink) the pool during an incident or ahead of planned
361        // CRDB maintenance. Shrinking is graceful: deadpool drops surplus
362        // connections as they are returned.
363        let max_size = self.knobs.connection_pool_max_size();
364        if status.max_size != max_size {
365            self.pool.resize(max_size);
366        }
367        self.status_metrics(status);
368        let res = self.pool.get().await;
369        if let Err(PoolError::Backend(err)) = &res {
370            debug!("error establishing connection: {}", err);
371            self.metrics.connpool_connection_errors.inc();
372        }
373        self.metrics
374            .connpool_acquire_seconds
375            .inc_by(start.elapsed().as_secs_f64());
376        self.metrics.connpool_acquires.inc();
377        self.status_metrics(self.pool.status());
378        res
379    }
380}
381
382#[cfg(test)]
383mod tests {
384    use std::str::FromStr;
385    use std::sync::atomic::AtomicUsize;
386
387    use mz_ore::metrics::MetricsRegistry;
388
389    use super::*;
390
391    #[derive(Debug)]
392    struct TestKnobs {
393        max_size: AtomicUsize,
394    }
395
396    impl PostgresClientKnobs for TestKnobs {
397        fn connection_pool_max_size(&self) -> usize {
398            self.max_size.load(Ordering::SeqCst)
399        }
400        fn connection_pool_max_wait(&self) -> Option<Duration> {
401            Some(Duration::from_secs(1))
402        }
403        fn connection_pool_ttl(&self) -> Duration {
404            Duration::MAX
405        }
406        fn connection_pool_ttl_stagger(&self) -> Duration {
407            Duration::MAX
408        }
409        fn connect_timeout(&self) -> Duration {
410            Duration::from_secs(5)
411        }
412        fn tcp_user_timeout(&self) -> Duration {
413            Duration::ZERO
414        }
415        fn keepalives_idle(&self) -> Duration {
416            Duration::from_secs(10)
417        }
418        fn keepalives_interval(&self) -> Duration {
419            Duration::from_secs(5)
420        }
421        fn keepalives_retries(&self) -> u32 {
422            5
423        }
424        fn statement_timeout(&self) -> Duration {
425            Duration::ZERO
426        }
427    }
428
429    /// Verifies that a changed `connection_pool_max_size` knob is applied to a
430    /// live pool on the next acquire, without reopening the client.
431    ///
432    /// Requires a running Postgres/CRDB, opted into via the same environment
433    /// variable as the persist external-storage tests. No-op otherwise so
434    /// `cargo test` works on unconfigured environments.
435    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
436    async fn pool_resize_applies_on_acquire() {
437        let url = match std::env::var("MZ_PERSIST_EXTERNAL_STORAGE_TEST_POSTGRES_URL") {
438            Ok(url) => SensitiveUrl::from_str(&url).expect("valid url"),
439            Err(_) => return,
440        };
441        let knobs = Arc::new(TestKnobs {
442            max_size: AtomicUsize::new(2),
443        });
444        let dyn_knobs: Arc<dyn PostgresClientKnobs> = Arc::<TestKnobs>::clone(&knobs);
445        let config = PostgresClientConfig::new(
446            url,
447            dyn_knobs,
448            PostgresClientMetrics::new(&MetricsRegistry::new(), "mz_postgres_client_test"),
449        );
450        let client = PostgresClient::open(config).expect("open client");
451
452        let conn = client.get_connection().await.expect("connection");
453        assert_eq!(client.status().max_size, 2);
454
455        // Growing applies on the next acquire.
456        knobs.max_size.store(5, Ordering::SeqCst);
457        let conn2 = client.get_connection().await.expect("connection");
458        assert_eq!(client.status().max_size, 5);
459
460        // Shrinking applies too. Surplus connections are dropped as they are
461        // returned to the pool.
462        knobs.max_size.store(1, Ordering::SeqCst);
463        drop(conn);
464        drop(conn2);
465        let _conn3 = client.get_connection().await.expect("connection");
466        assert_eq!(client.status().max_size, 1);
467    }
468}
mz_postgres_client/lib.rs

mz_postgres_client/
lib.rs
