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    metrics: PostgresClientMetrics,
244}
245
246impl std::fmt::Debug for PostgresClient {
247    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
248        f.debug_struct("PostgresClient").finish_non_exhaustive()
249    }
250}
251
252impl PostgresClient {
253    /// Open a [PostgresClient] using the given `config`.
254    pub fn open(config: PostgresClientConfig) -> Result<Self, PostgresError> {
255        let mut pg_config: Config = config.url.to_string_unredacted().parse()?;
256        pg_config.connect_timeout(config.knobs.connect_timeout());
257        pg_config.tcp_user_timeout(config.knobs.tcp_user_timeout());
258
259        // Configuring keepalives is important to ensure we can detect broken connections quickly.
260        // TCP_USER_TIMEOUT is not sufficient as it only enforces a timeout on ACKs for transmitted
261        // data, which only helps if we... transmit data.
262        pg_config.keepalives(true);
263        pg_config.keepalives_idle(config.knobs.keepalives_idle());
264        pg_config.keepalives_interval(config.knobs.keepalives_interval());
265        pg_config.keepalives_retries(config.knobs.keepalives_retries());
266
267        let tls = mz_tls_util::make_tls(&pg_config).map_err(|tls_err| match tls_err {
268            mz_tls_util::TlsError::Generic(e) => PostgresError::Indeterminate(e),
269            mz_tls_util::TlsError::OpenSsl(e) => PostgresError::Indeterminate(anyhow::anyhow!(e)),
270        })?;
271
272        let pg_manager = PgManager::from_config(
273            pg_config,
274            tls,
275            ManagerConfig {
276                recycling_method: RecyclingMethod::Fast,
277            },
278        );
279        // The isolation level and `statement_timeout` are applied inside `Manager::create` so the
280        // resolved level can be recorded on each connection it hands out.
281        let manager = Manager {
282            inner: pg_manager,
283            isolation: Arc::clone(&config.isolation),
284            knobs: Arc::clone(&config.knobs),
285            connections_created: config.metrics.connpool_connections_created.clone(),
286        };
287
288        let last_ttl_connection = AtomicU64::new(0);
289        let ttl_reconnections = config.metrics.connpool_ttl_reconnections.clone();
290        let builder = Pool::builder(manager);
291        let builder = match config.knobs.connection_pool_max_wait() {
292            None => builder,
293            Some(wait) => builder.wait_timeout(Some(wait)).runtime(Runtime::Tokio1),
294        };
295        let pool = builder
296            .max_size(config.knobs.connection_pool_max_size())
297            .pre_recycle(Hook::sync_fn(move |_client, conn_metrics| {
298                // proactively TTL connections to rebalance load to Postgres/CRDB. this helps
299                // fix skew when downstream DB operations (e.g. CRDB rolling restart) result
300                // in uneven load to each node, and works to reduce the # of connections
301                // maintained by the pool after bursty workloads.
302
303                // add a bias towards TTLing older connections first
304                if conn_metrics.age() < config.knobs.connection_pool_ttl() {
305                    return Ok(());
306                }
307
308                let last_ttl = last_ttl_connection.load(Ordering::SeqCst);
309                let now = (SYSTEM_TIME)();
310                let elapsed_since_last_ttl = Duration::from_millis(now.saturating_sub(last_ttl));
311
312                // stagger out reconnections to avoid stampeding the DB
313                if elapsed_since_last_ttl > config.knobs.connection_pool_ttl_stagger()
314                    && last_ttl_connection
315                        .compare_exchange_weak(last_ttl, now, Ordering::SeqCst, Ordering::SeqCst)
316                        .is_ok()
317                {
318                    ttl_reconnections.inc();
319                    // A `pre_recycle` error discards the connection and the pool moves on to the
320                    // next one.
321                    return Err(HookError::message("connection has been TTLed"));
322                }
323
324                Ok(())
325            }))
326            .build()
327            .expect("postgres connection pool built with incorrect parameters");
328
329        Ok(PostgresClient {
330            pool,
331            metrics: config.metrics,
332        })
333    }
334
335    fn status_metrics(&self, status: Status) {
336        // Negative when tasks are waiting for a connection.
337        self.metrics
338            .connpool_available
339            .set(f64::cast_lossy(status.available) - f64::cast_lossy(status.waiting));
340        self.metrics.connpool_size.set(u64::cast_from(status.size));
341        // Don't bother reporting the maximum size of the pool... we know that from config.
342    }
343
344    /// Gets connection from the pool or waits for one to become available.
345    pub async fn get_connection(&self) -> Result<Connection, PoolError> {
346        let start = Instant::now();
347        // note that getting the pool size here requires briefly locking the pool
348        self.status_metrics(self.pool.status());
349        let res = self.pool.get().await;
350        if let Err(PoolError::Backend(err)) = &res {
351            debug!("error establishing connection: {}", err);
352            self.metrics.connpool_connection_errors.inc();
353        }
354        self.metrics
355            .connpool_acquire_seconds
356            .inc_by(start.elapsed().as_secs_f64());
357        self.metrics.connpool_acquires.inc();
358        self.status_metrics(self.pool.status());
359        res
360    }
361}
mz_postgres_client/lib.rs

mz_postgres_client/
lib.rs
