datumesh_store/
pool.rs

1//! A small connection pool: no global state, no background threads, and a per-connection
2//! cache of prepared statements (each statement is prepared once per connection, on first
3//! use).
4
5use std::collections::HashMap;
6use std::ops::{Deref, DerefMut};
7use std::sync::{Condvar, Mutex, MutexGuard};
8use std::time::{Duration, Instant};
9
10use postgres::{Client, Statement};
11
12use crate::error::{Error, Result};
13use crate::sql::Q;
14
15/// Opens a new connection.
16pub(crate) type Connector = Box<dyn Fn() -> core::result::Result<Client, postgres::Error> + Send + Sync>;
17
18/// How the pool behaves.
19#[derive(Clone, Debug, PartialEq, Eq)]
20pub struct PoolOptions {
21    max_size: usize,
22    checkout_timeout: Duration,
23}
24
25impl Default for PoolOptions {
26    fn default() -> PoolOptions {
27        PoolOptions::new()
28    }
29}
30
31impl PoolOptions {
32    /// At most 8 connections; wait at most 30 seconds for one.
33    pub fn new() -> PoolOptions {
34        PoolOptions { max_size: 8, checkout_timeout: Duration::from_secs(30) }
35    }
36
37    /// The most connections open at once (at least 1).
38    pub fn max_size(mut self, n: usize) -> PoolOptions {
39        self.max_size = n.max(1);
40        self
41    }
42
43    /// How long a call waits for a free connection before failing with
44    /// [`Error::PoolTimeout`] (reads: `None`).
45    pub fn checkout_timeout(mut self, d: Duration) -> PoolOptions {
46        self.checkout_timeout = d;
47        self
48    }
49}
50
51/// A connection and the statements prepared on it.
52pub(crate) struct Conn {
53    pub(crate) client: Client,
54    statements: HashMap<Q, Statement>,
55}
56
57impl Conn {
58    /// The prepared statement for `q`, preparing it on first use.
59    pub(crate) fn stmt(&mut self, q: Q) -> core::result::Result<Statement, postgres::Error> {
60        if let Some(s) = self.statements.get(&q) {
61            return Ok(s.clone());
62        }
63        let s = self.client.prepare_typed(q.sql(), q.types())?;
64        self.statements.insert(q, s.clone());
65        Ok(s)
66    }
67}
68
69struct State {
70    idle: Vec<Conn>,
71    open: usize,
72}
73
74pub(crate) struct Pool {
75    connect: Connector,
76    state: Mutex<State>,
77    freed: Condvar,
78    options: PoolOptions,
79}
80
81/// Settings every connection gets: the append-only triggers rely on READ COMMITTED.
82const SESSION: &str = "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL READ COMMITTED; \
83                       SET TIME ZONE 'UTC'; SET DateStyle = 'ISO, YMD'";
84
85impl Pool {
86    pub(crate) fn new(connect: Connector, options: PoolOptions) -> Pool {
87        Pool { connect, state: Mutex::new(State { idle: Vec::new(), open: 0 }), freed: Condvar::new(), options }
88    }
89
90    fn lock(&self) -> MutexGuard<'_, State> {
91        // A panic while holding the lock cannot leave State inconsistent (two fields, each
92        // updated in one step), so a poisoned lock is still usable.
93        self.state.lock().unwrap_or_else(|e| e.into_inner())
94    }
95
96    /// A free connection, a new one if fewer than `max_size` are open, or one freed within
97    /// the checkout timeout.
98    pub(crate) fn get(&self) -> Result<Pooled<'_>> {
99        let deadline = Instant::now() + self.options.checkout_timeout;
100        let mut state = self.lock();
101        loop {
102            while let Some(conn) = state.idle.pop() {
103                if conn.client.is_closed() {
104                    state.open -= 1;
105                    continue;
106                }
107                return Ok(Pooled { pool: self, conn: Some(conn) });
108            }
109            if state.open < self.options.max_size {
110                state.open += 1;
111                drop(state);
112                return match self.open() {
113                    Ok(conn) => Ok(Pooled { pool: self, conn: Some(conn) }),
114                    Err(e) => {
115                        self.lock().open -= 1;
116                        self.freed.notify_one();
117                        Err(e)
118                    }
119                };
120            }
121            let now = Instant::now();
122            if now >= deadline {
123                return Err(Error::PoolTimeout);
124            }
125            state = self.freed.wait_timeout(state, deadline - now).unwrap_or_else(|e| e.into_inner()).0;
126        }
127    }
128
129    fn open(&self) -> Result<Conn> {
130        let mut client = (self.connect)()?;
131        client.batch_execute(SESSION)?;
132        Ok(Conn { client, statements: HashMap::new() })
133    }
134
135    fn put_back(&self, conn: Conn) {
136        let mut state = self.lock();
137        if conn.client.is_closed() {
138            state.open -= 1;
139        } else {
140            state.idle.push(conn);
141        }
142        drop(state);
143        self.freed.notify_one();
144    }
145}
146
147/// A connection checked out of the pool; returned to it when dropped (or discarded if the
148/// server closed it).
149pub(crate) struct Pooled<'a> {
150    pool: &'a Pool,
151    conn: Option<Conn>,
152}
153
154impl Deref for Pooled<'_> {
155    type Target = Conn;
156    fn deref(&self) -> &Conn {
157        // Some from creation until drop.
158        self.conn.as_ref().expect("connection present until drop")
159    }
160}
161
162impl DerefMut for Pooled<'_> {
163    fn deref_mut(&mut self) -> &mut Conn {
164        self.conn.as_mut().expect("connection present until drop")
165    }
166}
167
168impl Drop for Pooled<'_> {
169    fn drop(&mut self) {
170        if let Some(conn) = self.conn.take() {
171            self.pool.put_back(conn);
172        }
173    }
174}
datumesh_store/pool.rs

datumesh_store/
pool.rs
