miden_node_db/sqlite/
pool.rs

1//! Async connection pool over raw `rusqlite`.
2//!
3//! SQLite permits only a single writer at a time, so the pool is split into a **single** writer
4//! connection and a pool of read-only connections. Writes (`write`/`begin_write`) serialize on the
5//! one writer; reads (`read`/`begin_read`) run concurrently on the reader pool. This makes the
6//! single-writer model structural (rather than relying on lock contention) and lets a held write
7//! transaction stay open without starving readers.
8
9use std::num::NonZeroUsize;
10use std::path::{Path, PathBuf};
11
12use deadpool::Runtime;
13use deadpool::managed::{Manager, Metrics, Object, Pool, RecycleError, RecycleResult};
14use deadpool_sync::SyncWrapper;
15use rusqlite::{Connection, OpenFlags, TransactionBehavior};
16use tracing::Instrument;
17
18use crate::sqlite::tx::{ReadTx, WriteTx};
19use crate::{DatabaseError, default_connection_pool_size};
20
21/// Per-connection prepared-statement cache capacity. Raised well above rusqlite's default of 16
22/// because we keep a large set of distinct statements; the bounded connection pools cap total
23/// cached-statement memory.
24const STATEMENT_CACHE_CAPACITY: usize = 512;
25
26// CONNECTION MANAGER
27// =================================================================================================
28
29/// Errors raised while creating or recycling a pooled connection.
30///
31/// Internal to the pool: callers only ever observe a [`DatabaseError`] (pool failures are boxed into
32/// [`DatabaseError::ConnectionPoolObtainError`]), so this type is not part of the public API.
33#[derive(Debug, thiserror::Error)]
34pub(crate) enum SqliteManagerError {
35    /// Opening the database file failed.
36    #[error("failed to open the sqlite database")]
37    Open(#[source] rusqlite::Error),
38    /// Applying the per-connection PRAGMAs failed.
39    #[error("failed to configure the sqlite connection")]
40    Configure(#[source] rusqlite::Error),
41    /// The pooled connection's mutex was poisoned by a panic during a previous interaction.
42    #[error("the pooled sqlite connection is poisoned")]
43    Poisoned,
44}
45
46struct SqliteManager {
47    path: PathBuf,
48    /// When set, connections are configured `PRAGMA query_only = ON` and skip the writer-only
49    /// `journal_mode` setup — used for the reader pool.
50    read_only: bool,
51}
52
53impl Manager for SqliteManager {
54    type Type = SyncWrapper<Connection>;
55    type Error = SqliteManagerError;
56
57    async fn create(&self) -> Result<Self::Type, Self::Error> {
58        let path = self.path.clone();
59        let read_only = self.read_only;
60        SyncWrapper::new(Runtime::Tokio1, move || {
61            let conn = Connection::open_with_flags(&path, OpenFlags::SQLITE_OPEN_READ_WRITE)
62                .map_err(SqliteManagerError::Open)?;
63            configure_connection(&conn, read_only).map_err(SqliteManagerError::Configure)?;
64            Ok(conn)
65        })
66        .await
67    }
68
69    async fn recycle(
70        &self,
71        conn: &mut Self::Type,
72        _metrics: &Metrics,
73    ) -> RecycleResult<Self::Error> {
74        if conn.is_mutex_poisoned() {
75            return Err(RecycleError::Backend(SqliteManagerError::Poisoned));
76        }
77        // Safety net for a held transaction handle dropped without `commit`/`rollback`: roll back
78        // any still-open transaction so the next user gets a clean connection.
79        conn.interact(|conn| {
80            if !conn.is_autocommit() {
81                let _ = conn.execute_batch("ROLLBACK");
82            }
83        })
84        .await
85        .map_err(|_| RecycleError::Backend(SqliteManagerError::Poisoned))?;
86        Ok(())
87    }
88}
89
90/// Applies the per-connection PRAGMAs and statement-cache sizing.
91///
92/// Both pools open the file `READ_WRITE`; reader connections are made read-only at runtime with
93/// `PRAGMA query_only = ON` (which, unlike opening `READ_ONLY`, still lets them create the WAL
94/// `-shm` file and read a WAL database).
95fn configure_connection(conn: &Connection, read_only: bool) -> rusqlite::Result<()> {
96    // busy_timeout makes concurrent writers wait instead of failing immediately; foreign keys
97    // enforce referential integrity.
98    if read_only {
99        // A query_only connection cannot set `journal_mode` (it is a write); WAL is already
100        // persisted in the file header by the writer / migration path.
101        conn.execute_batch(
102            "PRAGMA busy_timeout = 5000;
103             PRAGMA foreign_keys = ON;
104             PRAGMA query_only = ON;",
105        )?;
106    } else {
107        // WAL allows concurrent readers while the writer holds the lock.
108        conn.execute_batch(
109            "PRAGMA busy_timeout = 5000;
110             PRAGMA journal_mode = WAL;
111             PRAGMA foreign_keys = ON;",
112        )?;
113    }
114    conn.set_prepared_statement_cache_capacity(STATEMENT_CACHE_CAPACITY);
115    // Register the `array` extension so the cacheable IN-list helpers can bind lists via
116    // `rarray(?)` (see `crate::sqlite::in_list`).
117    rusqlite::vtab::array::load_module(conn)?;
118    Ok(())
119}
120
121// DATABASE
122// =================================================================================================
123
124/// A rusqlite-backed connection pool. Cloning shares the underlying pools.
125///
126/// Holds a single writer connection and a pool of reader connections (see the module docs).
127#[derive(Clone)]
128pub struct Database {
129    writer: Pool<SqliteManager>,
130    readers: Pool<SqliteManager>,
131}
132
133impl Database {
134    /// Opens a database over `database_filepath` with the default reader-pool size.
135    pub fn new(database_filepath: &Path) -> Result<Self, DatabaseError> {
136        Self::new_with_pool_size(database_filepath, default_connection_pool_size())
137    }
138
139    /// Opens a database over `database_filepath` with the given reader-pool size. The writer is
140    /// always a single connection.
141    pub fn new_with_pool_size(
142        database_filepath: &Path,
143        connection_pool_size: NonZeroUsize,
144    ) -> Result<Self, DatabaseError> {
145        let writer = Pool::builder(SqliteManager {
146            path: database_filepath.to_path_buf(),
147            read_only: false,
148        })
149        .max_size(1)
150        .build()?;
151        let readers = Pool::builder(SqliteManager {
152            path: database_filepath.to_path_buf(),
153            read_only: true,
154        })
155        .max_size(connection_pool_size.get())
156        .build()?;
157        Ok(Self { writer, readers })
158    }
159
160    /// Checks the single writer connection out of the pool.
161    async fn checkout_writer(&self) -> Result<Object<SqliteManager>, DatabaseError> {
162        self.writer
163            .get()
164            .in_current_span()
165            .await
166            .map_err(|err| DatabaseError::ConnectionPoolObtainError(Box::new(err)))
167    }
168
169    /// Checks a reader connection out of the pool.
170    async fn checkout_reader(&self) -> Result<Object<SqliteManager>, DatabaseError> {
171        self.readers
172            .get()
173            .in_current_span()
174            .await
175            .map_err(|err| DatabaseError::ConnectionPoolObtainError(Box::new(err)))
176    }
177
178    /// Runs `query` inside a read-only (`DEFERRED`, never committed) transaction on a reader
179    /// connection.
180    pub async fn read<R, E, F>(&self, msg: impl ToString + Send, query: F) -> Result<R, E>
181    where
182        F: FnOnce(&ReadTx<'_>) -> Result<R, E> + Send + 'static,
183        R: Send + 'static,
184        E: From<DatabaseError> + Send + 'static,
185    {
186        let conn = self.checkout_reader().await.map_err(E::from)?;
187        let msg = msg.to_string();
188        let span = tracing::Span::current();
189        conn.interact(move |conn| {
190            let _guard = span.enter();
191            let tx = conn
192                .transaction_with_behavior(TransactionBehavior::Deferred)
193                .map_err(|err| E::from(DatabaseError::from(err)))?;
194            query(&ReadTx::new(&tx))
195            // `tx` is dropped here without a commit, rolling back any writes.
196        })
197        .await
198        .map_err(|err| E::from(DatabaseError::interact(&msg, &err)))?
199    }
200
201    /// Runs `query` inside a read-write (`IMMEDIATE`) transaction on the single writer connection,
202    /// committing on `Ok`.
203    pub async fn write<R, E, F>(&self, msg: impl ToString + Send, query: F) -> Result<R, E>
204    where
205        F: FnOnce(&WriteTx<'_>) -> Result<R, E> + Send + 'static,
206        R: Send + 'static,
207        E: From<DatabaseError> + Send + 'static,
208    {
209        let conn = self.checkout_writer().await.map_err(E::from)?;
210        let msg = msg.to_string();
211        let span = tracing::Span::current();
212        conn.interact(move |conn| {
213            let _guard = span.enter();
214            let tx = conn
215                .transaction_with_behavior(TransactionBehavior::Immediate)
216                .map_err(|err| E::from(DatabaseError::from(err)))?;
217            let result = query(&WriteTx::new(&tx))?;
218            tx.commit().map_err(|err| E::from(DatabaseError::from(err)))?;
219            Ok(result)
220        })
221        .await
222        .map_err(|err| E::from(DatabaseError::interact(&msg, &err)))?
223    }
224
225    /// Begins a read-only (`DEFERRED`) transaction on a reader connection and returns a handle held
226    /// across `.await` points. See [`ReadTransaction`].
227    pub async fn begin_read(&self) -> Result<ReadTransaction, DatabaseError> {
228        let conn = self.checkout_reader().await?;
229        run_tx_stmt(&conn, "BEGIN DEFERRED").await?;
230        Ok(ReadTransaction { conn })
231    }
232
233    /// Begins a read-write (`IMMEDIATE`) transaction on the single writer connection and returns a
234    /// handle held across `.await` points. The handle must be committed (or it rolls back). See
235    /// [`WriteTransaction`].
236    pub async fn begin_write(&self) -> Result<WriteTransaction, DatabaseError> {
237        let conn = self.checkout_writer().await?;
238        run_tx_stmt(&conn, "BEGIN IMMEDIATE").await?;
239        Ok(WriteTransaction { conn })
240    }
241}
242
243// HELD TRANSACTIONS
244// =================================================================================================
245
246/// Runs a transaction-control statement (`BEGIN`/`COMMIT`/`ROLLBACK`) on a checked-out connection.
247async fn run_tx_stmt(
248    conn: &Object<SqliteManager>,
249    stmt: &'static str,
250) -> Result<(), DatabaseError> {
251    conn.interact(move |conn| conn.execute_batch(stmt))
252        .await
253        .map_err(|err| DatabaseError::interact(stmt, &err))?
254        .map_err(DatabaseError::from)
255}
256
257/// A read transaction (`DEFERRED`) held across `.await` points, on a reader connection.
258///
259/// Run batches of synchronous queries with [`run`](Self::run); the transaction stays open between
260/// calls, so a request handler can interleave queries with async work on a single consistent
261/// snapshot. The transaction is read-only and ends (rolls back) when the handle is dropped, or
262/// explicitly via [`close`](Self::close).
263pub struct ReadTransaction {
264    conn: Object<SqliteManager>,
265}
266
267impl ReadTransaction {
268    /// Runs a batch of read queries against the open transaction.
269    pub async fn run<R, E, F>(&self, msg: impl ToString + Send, query: F) -> Result<R, E>
270    where
271        F: FnOnce(&ReadTx<'_>) -> Result<R, E> + Send + 'static,
272        R: Send + 'static,
273        E: From<DatabaseError> + Send + 'static,
274    {
275        let msg = msg.to_string();
276        let span = tracing::Span::current();
277        self.conn
278            .interact(move |conn| {
279                let _guard = span.enter();
280                query(&ReadTx::new(conn))
281            })
282            .await
283            .map_err(|err| E::from(DatabaseError::interact(&msg, &err)))?
284    }
285
286    /// Ends the transaction explicitly (rolls back; a read transaction has nothing to commit).
287    pub async fn close(self) -> Result<(), DatabaseError> {
288        run_tx_stmt(&self.conn, "ROLLBACK").await
289    }
290}
291
292/// A read-write transaction (`IMMEDIATE`) held across `.await` points, on the single writer
293/// connection.
294///
295/// Run batches of synchronous queries with [`run`](Self::run); the transaction stays open between
296/// calls, so a request handler can interleave reads and writes with async work atomically. Finish
297/// with [`commit`](Self::commit) to persist, or [`rollback`](Self::rollback) to discard; if the
298/// handle is dropped without either, the pool rolls the transaction back when the connection is
299/// recycled.
300///
301/// The handle holds the sole writer connection for its whole lifetime.
302pub struct WriteTransaction {
303    conn: Object<SqliteManager>,
304}
305
306impl WriteTransaction {
307    /// Runs a batch of read/write queries against the open transaction.
308    pub async fn run<R, E, F>(&self, msg: impl ToString + Send, query: F) -> Result<R, E>
309    where
310        F: FnOnce(&WriteTx<'_>) -> Result<R, E> + Send + 'static,
311        R: Send + 'static,
312        E: From<DatabaseError> + Send + 'static,
313    {
314        let msg = msg.to_string();
315        let span = tracing::Span::current();
316        self.conn
317            .interact(move |conn| {
318                let _guard = span.enter();
319                query(&WriteTx::new(conn))
320            })
321            .await
322            .map_err(|err| E::from(DatabaseError::interact(&msg, &err)))?
323    }
324
325    /// Commits the transaction, persisting all writes.
326    pub async fn commit(self) -> Result<(), DatabaseError> {
327        run_tx_stmt(&self.conn, "COMMIT").await
328    }
329
330    /// Rolls back the transaction, discarding all writes.
331    pub async fn rollback(self) -> Result<(), DatabaseError> {
332        run_tx_stmt(&self.conn, "ROLLBACK").await
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use std::num::NonZeroUsize;
339    use std::path::{Path, PathBuf};
340
341    use rusqlite::Connection;
342
343    use super::Database;
344    use crate::DatabaseError;
345
346    /// A throwaway file-backed database; the pools open existing files `READ_WRITE` only, so the
347    /// file and schema are created up front.
348    struct TempDb {
349        path: PathBuf,
350    }
351
352    impl TempDb {
353        fn new(name: &str) -> Self {
354            let path = std::env::temp_dir()
355                .join(format!("miden-node-db-pool-{name}-{}.sqlite3", std::process::id()));
356            let db = Self { path };
357            db.remove_files();
358            let conn = Connection::open(&db.path).expect("create db file");
359            conn.execute_batch("CREATE TABLE items (id INTEGER PRIMARY KEY);")
360                .expect("create table");
361            db
362        }
363
364        fn path(&self) -> &Path {
365            &self.path
366        }
367
368        fn remove_files(&self) {
369            let _ = fs_err::remove_file(&self.path);
370            let _ = fs_err::remove_file(self.path.with_extension("sqlite3-wal"));
371            let _ = fs_err::remove_file(self.path.with_extension("sqlite3-shm"));
372        }
373    }
374
375    impl Drop for TempDb {
376        fn drop(&mut self) {
377            self.remove_files();
378        }
379    }
380
381    fn open_db(temp: &TempDb) -> Database {
382        Database::new_with_pool_size(temp.path(), NonZeroUsize::new(4).unwrap()).unwrap()
383    }
384
385    async fn count_items(db: &Database) -> i64 {
386        db.read::<_, DatabaseError, _>("count", |r| {
387            Ok(r.query("SELECT COUNT(*) FROM items", &[], |row| row.get::<i64>(0))?
388                .into_iter()
389                .next()
390                .unwrap_or(0))
391        })
392        .await
393        .unwrap()
394    }
395
396    async fn insert_committed(db: &Database, id: i64) {
397        let tx = db.begin_write().await.unwrap();
398        tx.run::<_, DatabaseError, _>("insert", move |w| {
399            w.execute("INSERT INTO items (id) VALUES (?1)", &[&id])?;
400            Ok(())
401        })
402        .await
403        .unwrap();
404        tx.commit().await.unwrap();
405    }
406
407    #[tokio::test]
408    async fn held_write_transaction_commits_across_awaits() {
409        let temp = TempDb::new("commit");
410        let db = open_db(&temp);
411
412        let tx = db.begin_write().await.unwrap();
413        tx.run::<_, DatabaseError, _>("insert-1", |w| {
414            w.execute("INSERT INTO items (id) VALUES (?1)", &[&1i64])?;
415            Ok(())
416        })
417        .await
418        .unwrap();
419
420        // Interleave async work between statements on the same still-open transaction.
421        tokio::task::yield_now().await;
422
423        tx.run::<_, DatabaseError, _>("insert-2", |w| {
424            w.execute("INSERT INTO items (id) VALUES (?1)", &[&2i64])?;
425            Ok(())
426        })
427        .await
428        .unwrap();
429
430        tx.commit().await.unwrap();
431
432        assert_eq!(count_items(&db).await, 2);
433    }
434
435    #[tokio::test]
436    async fn dropped_write_transaction_rolls_back() {
437        let temp = TempDb::new("rollback");
438        let db = open_db(&temp);
439
440        {
441            let tx = db.begin_write().await.unwrap();
442            tx.run::<_, DatabaseError, _>("insert", |w| {
443                w.execute("INSERT INTO items (id) VALUES (?1)", &[&1i64])?;
444                Ok(())
445            })
446            .await
447            .unwrap();
448            // `tx` is dropped here without a commit.
449        }
450
451        // The sole writer connection is reused; `recycle` must have rolled back the orphaned
452        // transaction, otherwise this `BEGIN IMMEDIATE` would fail with "cannot start a transaction
453        // within a transaction". The first insert must not have persisted.
454        insert_committed(&db, 2).await;
455        assert_eq!(count_items(&db).await, 1);
456    }
457
458    #[tokio::test]
459    async fn reads_proceed_while_write_transaction_is_held() {
460        let temp = TempDb::new("concurrent");
461        let db = open_db(&temp);
462        insert_committed(&db, 1).await;
463
464        // Hold an open write transaction with an uncommitted insert.
465        let tx = db.begin_write().await.unwrap();
466        tx.run::<_, DatabaseError, _>("insert-uncommitted", |w| {
467            w.execute("INSERT INTO items (id) VALUES (?1)", &[&2i64])?;
468            Ok(())
469        })
470        .await
471        .unwrap();
472
473        // A read on the reader pool proceeds (does not block on the writer) and does not see the
474        // uncommitted row.
475        assert_eq!(count_items(&db).await, 1);
476
477        tx.commit().await.unwrap();
478        assert_eq!(count_items(&db).await, 2);
479    }
480
481    #[tokio::test]
482    async fn reader_connections_are_query_only() {
483        let temp = TempDb::new("query_only");
484        let db = open_db(&temp);
485
486        let query_only = db
487            .read::<_, DatabaseError, _>("pragma", |r| {
488                Ok(r.query("PRAGMA query_only", &[], |row| row.get::<i64>(0))?
489                    .into_iter()
490                    .next()
491                    .unwrap_or(0))
492            })
493            .await
494            .unwrap();
495        assert_eq!(query_only, 1, "reader connections must be query_only");
496
497        // A write attempted on a reader connection is rejected.
498        let result = db
499            .read::<(), DatabaseError, _>("rejected-write", |r| {
500                r.query("INSERT INTO items (id) VALUES (99)", &[], |_| Ok(()))?;
501                Ok(())
502            })
503            .await;
504        assert!(result.is_err(), "writes on a reader connection must fail");
505    }
506}
miden_node_db/sqlite/pool.rs

miden_node_db/sqlite/
pool.rs
