runlight_sqlx/
postgres.rs

1//! Postgres, on the app's pool. Each statement runs on the pool on its own, so the store's writes never
2//! join a transaction the app has open; a transaction or a lock holds one connection until it ends.
3
4use std::sync::Arc;
5use std::time::Instant;
6
7use runlight::BoxFuture;
8use runlight::store::{
9    Db, DbError, Dialect, Hold, Param, Row, TablesBusyError, fill_placeholders, lock_wait, postgres_literal,
10};
11use sqlx::pool::PoolConnection;
12use sqlx::{AssertSqlSafe, PgPool, Postgres};
13use tokio::sync::Mutex;
14
15use crate::rows::{db_error, pg_row};
16
17/// Arbitrary but fixed, so every Runlight process takes the same lock to create tables.
18const MIGRATION_LOCK: i64 = 7_331_906;
19
20/// The store's driver over a Postgres pool. Values go in the statement's text as literals of unknown
21/// type, so Postgres infers each one's type as it does for node-postgres's untyped parameters.
22#[derive(Clone)]
23pub struct PostgresDb {
24    pool: PgPool,
25}
26
27impl PostgresDb {
28    /// The store's driver over a pool.
29    pub fn new(pool: PgPool) -> PostgresDb {
30        PostgresDb { pool }
31    }
32}
33
34fn text(sql: &str, params: &[Param]) -> Result<AssertSqlSafe<String>, DbError> {
35    Ok(AssertSqlSafe(fill_placeholders(sql, params, false, postgres_literal)?))
36}
37
38impl Db for PostgresDb {
39    fn dialect(&self) -> Dialect {
40        Dialect::Postgres
41    }
42
43    fn all<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<Vec<Row>, DbError>> {
44        Box::pin(async move {
45            let rows = sqlx::raw_sql(text(sql, &params)?).fetch_all(&self.pool).await.map_err(db_error)?;
46            rows.iter().map(pg_row).collect()
47        })
48    }
49
50    fn run<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<u64, DbError>> {
51        Box::pin(async move {
52            Ok(sqlx::raw_sql(text(sql, &params)?).execute(&self.pool).await.map_err(db_error)?.rows_affected())
53        })
54    }
55
56    fn hold(&self, kind: Hold) -> BoxFuture<'_, Result<Arc<dyn Db>, DbError>> {
57        Box::pin(async move {
58            // In its slot from the start, so a request cancelled part way closes the connection rather than
59            // give it back to the pool in a transaction or with the lock taken.
60            let mut slot = Slot { conn: Some(self.pool.acquire().await.map_err(db_error)?), nested: false };
61            let conn = slot.conn.as_mut().ok_or_else(ended)?;
62            match kind {
63                Hold::Transaction => {
64                    sqlx::raw_sql(AssertSqlSafe("BEGIN".to_string())).execute(&mut **conn).await.map_err(db_error)?;
65                }
66                Hold::Exclusive => {
67                    // Up to the lock wait, so a request does not wait on a long index build for good.
68                    let wait = lock_wait();
69                    let until = Instant::now() + wait;
70                    loop {
71                        // Asked for again and again rather than waited on: a waiting statement would hold up an
72                        // index being built CONCURRENTLY by whoever has the lock, and the two would wait on each
73                        // other for good.
74                        let rows = sqlx::raw_sql(AssertSqlSafe(format!(
75                            "SELECT pg_try_advisory_lock({MIGRATION_LOCK}) AS ok"
76                        )))
77                        .fetch_all(&mut **conn)
78                        .await
79                        .map_err(db_error)?;
80                        let ok = rows.first().map(pg_row).transpose()?.is_some_and(|r| r.num("ok") == 1.0);
81                        if ok {
82                            break;
83                        }
84                        if Instant::now() >= until {
85                            return Err(TablesBusyError { waited: wait }.into());
86                        }
87                        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
88                    }
89                }
90            }
91            Ok(Arc::new(Held { conn: Arc::new(Mutex::new(slot)), kind, role: Role::Owner }) as Arc<dyn Db>)
92        })
93    }
94}
95
96/// The held connection. Let go of with it still in (a request cancelled before `finish`), it closes the
97/// connection, so no transaction or lock goes back to the pool.
98struct Slot {
99    conn: Option<PoolConnection<Postgres>>,
100    /// A transaction begun inside the lock that has not ended.
101    nested: bool,
102}
103
104impl Slot {
105    fn close(&mut self) {
106        if let Some(conn) = self.conn.take() {
107            drop(conn.detach());
108        }
109    }
110}
111
112impl Drop for Slot {
113    fn drop(&mut self) {
114        self.close();
115    }
116}
117
118/// One connection held for a transaction or a lock.
119struct Held {
120    conn: Arc<Mutex<Slot>>,
121    kind: Hold,
122    role: Role,
123}
124
125impl Drop for Held {
126    fn drop(&mut self) {
127        // Its owner gone without `finish`: closed now rather than when the last joined handle goes.
128        if self.role == Role::Owner
129            && let Ok(mut slot) = self.conn.try_lock()
130        {
131            slot.close();
132        }
133    }
134}
135
136/// What finishing a held connection does.
137#[derive(Clone, Copy, PartialEq, Eq)]
138enum Role {
139    /// Ends the transaction or lets the lock go, and gives the connection back.
140    Owner,
141    /// A transaction begun inside a lock: ends the transaction, and keeps the connection.
142    Nested,
143    /// Joined what was already held: ends nothing.
144    Joined,
145}
146
147fn ended() -> DbError {
148    DbError("Runlight: the held connection has been let go".into())
149}
150
151impl Db for Held {
152    fn dialect(&self) -> Dialect {
153        Dialect::Postgres
154    }
155
156    fn all<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<Vec<Row>, DbError>> {
157        Box::pin(async move {
158            let mut slot = self.conn.lock().await;
159            let conn = slot.conn.as_mut().ok_or_else(ended)?;
160            let rows = sqlx::raw_sql(text(sql, &params)?).fetch_all(&mut **conn).await.map_err(db_error)?;
161            rows.iter().map(pg_row).collect()
162        })
163    }
164
165    fn run<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<u64, DbError>> {
166        Box::pin(async move {
167            let mut slot = self.conn.lock().await;
168            let conn = slot.conn.as_mut().ok_or_else(ended)?;
169            Ok(sqlx::raw_sql(text(sql, &params)?).execute(&mut **conn).await.map_err(db_error)?.rows_affected())
170        })
171    }
172
173    fn hold(&self, kind: Hold) -> BoxFuture<'_, Result<Arc<dyn Db>, DbError>> {
174        Box::pin(async move {
175            // A transaction inside the lock runs on the locked connection; anything else joins what is held.
176            if kind == Hold::Transaction && self.kind == Hold::Exclusive {
177                let mut slot = self.conn.lock().await;
178                let slot = &mut *slot;
179                let conn = slot.conn.as_mut().ok_or_else(ended)?;
180                // Marked before it is sent: a BEGIN cancelled part way may still have run.
181                slot.nested = true;
182                sqlx::raw_sql(AssertSqlSafe("BEGIN".to_string())).execute(&mut **conn).await.map_err(db_error)?;
183                return Ok(Arc::new(Held { conn: self.conn.clone(), kind, role: Role::Nested }) as Arc<dyn Db>);
184            }
185            Ok(Arc::new(Held { conn: self.conn.clone(), kind: self.kind, role: Role::Joined }) as Arc<dyn Db>)
186        })
187    }
188
189    fn finish(&self, ok: bool) -> BoxFuture<'_, Result<(), DbError>> {
190        Box::pin(async move {
191            let mut slot = self.conn.lock().await;
192            if self.role == Role::Joined {
193                return Ok(());
194            }
195            if self.role == Role::Nested {
196                let slot = &mut *slot;
197                let Some(conn) = slot.conn.as_mut() else { return Ok(()) };
198                let end = if ok { "COMMIT" } else { "ROLLBACK" };
199                sqlx::raw_sql(AssertSqlSafe(end.to_string())).execute(&mut **conn).await.map_err(db_error)?;
200                slot.nested = false;
201                return Ok(());
202            }
203            // A transaction inside the lock that never ended: closing the connection rolls it back and lets
204            // the lock go.
205            if slot.nested {
206                slot.close();
207                return Ok(());
208            }
209            let Some(mut conn) = slot.conn.take() else { return Ok(()) };
210            let result = match self.kind {
211                Hold::Transaction => {
212                    let end = if ok { "COMMIT" } else { "ROLLBACK" };
213                    sqlx::raw_sql(AssertSqlSafe(end.to_string()))
214                        .execute(&mut *conn)
215                        .await
216                        .map(|_| ())
217                        .map_err(db_error)
218                }
219                // A connection whose unlock failed may still have the lock, so it is closed below, and a
220                // lost connection ends its session, and the lock with it.
221                Hold::Exclusive => sqlx::raw_sql(AssertSqlSafe(format!("SELECT pg_advisory_unlock({MIGRATION_LOCK})")))
222                    .execute(&mut *conn)
223                    .await
224                    .map(|_| ())
225                    .map_err(db_error),
226            };
227            if result.is_err() {
228                // A connection still in its transaction or holding the lock must never go back to the pool.
229                drop(conn.detach());
230            }
231            if self.kind == Hold::Exclusive { Ok(()) } else { result }
232        })
233    }
234}
runlight_sqlx/postgres.rs

runlight_sqlx/
postgres.rs
