siderite_backends/sqlite/
group.rs

1//! Opt-in group commit for autocommit SQLite writes.
2//!
3//! Every autocommit write is its own transaction, so each pays SQLite's
4//! commit sync (several `fsync`s in rollback-journal mode). With group commit
5//! one task performs the writes: whatever queued while the previous commit
6//! was running shares the next transaction, and therefore one commit sync.
7//!
8//! - Each write runs inside its own savepoint, so a failing statement fails
9//!   only its own caller; the rest of the batch still commits.
10//! - A caller is answered only after the shared `COMMIT` succeeds, so an
11//!   acknowledged write is as durable as an ungrouped one.
12//! - If `COMMIT` fails, every caller in the batch gets the error. Nothing is
13//!   retried: the outcome of a failed commit is not always known.
14//! - A lone queued write runs as a plain autocommit statement, so light load
15//!   pays no transaction or savepoint overhead.
16//!
17//! What changes: writes in one batch become visible together, a commit
18//! failure fails every write in the batch, and a statement that ends the
19//! surrounding transaction itself (`INSERT OR ROLLBACK`, `RAISE(ROLLBACK)`)
20//! fails the writes queued with it.
21
22use super::gate::WriteGate;
23use super::run_write;
24use crate::shared::map_error;
25use crate::sql::CompiledQuery;
26use siderite_orm::{BackendError, ExecResult, OrmError};
27use sqlx::TransactionManager;
28use sqlx::sqlite::{SqliteConnection, SqlitePool, SqliteTransactionManager};
29use tokio::sync::{mpsc, oneshot};
30
31/// Group commit settings for
32/// [`SqliteBackend::group_commit`](super::SqliteBackend::group_commit).
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub struct GroupCommit {
35    /// Most writes sharing one transaction.
36    pub max_batch: usize,
37    /// Writes waiting for the committer before callers wait to enqueue.
38    pub queue: usize,
39}
40
41impl Default for GroupCommit {
42    fn default() -> Self {
43        Self {
44            max_batch: 64,
45            queue: 1024,
46        }
47    }
48}
49
50type Reply = Result<ExecResult, OrmError>;
51
52struct Job {
53    compiled: CompiledQuery,
54    returning: bool,
55    reply: oneshot::Sender<Reply>,
56}
57
58/// Handle to the committer task. The task ends once every handle is gone.
59#[derive(Debug, Clone)]
60pub(super) struct Committer {
61    jobs: mpsc::Sender<Job>,
62}
63
64impl Committer {
65    /// Start the committer on the current Tokio runtime.
66    pub(super) fn spawn(pool: SqlitePool, gate: WriteGate, config: GroupCommit) -> Self {
67        let max_batch = config.max_batch.max(1);
68        let (jobs, queue) = mpsc::channel(config.queue.max(1));
69        tokio::spawn(run(pool, gate, queue, max_batch));
70        Self { jobs }
71    }
72
73    /// Queue one write and wait for its committed result.
74    pub(super) async fn submit(&self, compiled: CompiledQuery, returning: bool) -> Reply {
75        let (reply, answer) = oneshot::channel();
76        let job = Job {
77            compiled,
78            returning,
79            reply,
80        };
81        self.jobs.send(job).await.map_err(|_| stopped())?;
82        answer.await.map_err(|_| stopped())?
83    }
84}
85
86fn stopped() -> OrmError {
87    BackendError::Database("SQLite group committer stopped".into()).into()
88}
89
90async fn run(pool: SqlitePool, gate: WriteGate, mut queue: mpsc::Receiver<Job>, max: usize) {
91    let mut batch = Vec::with_capacity(max);
92    while queue.recv_many(&mut batch, max).await > 0 {
93        // A caller that gave up before its write started never gets one,
94        // as with an ungrouped write cancelled before it was sent.
95        batch.retain(|job| !job.reply.is_closed());
96        if batch.is_empty() {
97            continue;
98        }
99        let replies = match gate.acquire().await {
100            Ok(_permit) => match pool.acquire().await {
101                Ok(mut conn) => write(&mut conn, &mut batch).await,
102                Err(err) => fail_all(batch.len(), &OrmError::from(map_error(err))),
103            },
104            Err(err) => fail_all(batch.len(), &err),
105        };
106        for (job, reply) in batch.drain(..).zip(replies) {
107            drop(job.reply.send(reply));
108        }
109    }
110}
111
112async fn write(conn: &mut SqliteConnection, batch: &mut [Job]) -> Vec<Reply> {
113    if let [job] = batch {
114        let compiled = take(&mut job.compiled);
115        return vec![run_write(&mut *conn, compiled, job.returning).await];
116    }
117    if let Err(err) = SqliteTransactionManager::begin(conn, Some("BEGIN IMMEDIATE".into())).await {
118        return fail_all(batch.len(), &OrmError::from(map_error(err)));
119    }
120    let mut replies = Vec::with_capacity(batch.len());
121    for job in batch.iter_mut() {
122        if let Err(err) = savepoint(conn, "SAVEPOINT siderite_group").await {
123            return abort(conn, batch.len(), &err).await;
124        }
125        let compiled = take(&mut job.compiled);
126        let reply = run_write(&mut *conn, compiled, job.returning).await;
127        let close = if reply.is_ok() {
128            savepoint(conn, "RELEASE siderite_group").await
129        } else {
130            // Fails when the statement ended the whole transaction.
131            match savepoint(conn, "ROLLBACK TO siderite_group").await {
132                Ok(()) => savepoint(conn, "RELEASE siderite_group").await,
133                Err(err) => Err(err),
134            }
135        };
136        if let Err(err) = close {
137            return abort(conn, batch.len(), &err).await;
138        }
139        replies.push(reply);
140    }
141    match SqliteTransactionManager::commit(conn).await {
142        Ok(()) => replies,
143        Err(err) => abort(conn, batch.len(), &OrmError::from(map_error(err))).await,
144    }
145}
146
147/// Roll the batch back and fail every caller with `err`.
148async fn abort(conn: &mut SqliteConnection, len: usize, err: &OrmError) -> Vec<Reply> {
149    drop(SqliteTransactionManager::rollback(conn).await);
150    fail_all(len, err)
151}
152
153fn fail_all(len: usize, err: &OrmError) -> Vec<Reply> {
154    let message = format!("SQLite group commit failed: {err}");
155    (0..len)
156        .map(|_| Err(BackendError::Database(message.clone()).into()))
157        .collect()
158}
159
160async fn savepoint(conn: &mut SqliteConnection, sql: &str) -> Result<(), OrmError> {
161    sqlx::Executor::execute(&mut *conn, sql)
162        .await
163        .map(|_| ())
164        .map_err(|e| map_error(e).into())
165}
166
167fn take(compiled: &mut CompiledQuery) -> CompiledQuery {
168    CompiledQuery {
169        sql: std::mem::take(&mut compiled.sql),
170        params: std::mem::take(&mut compiled.params),
171    }
172}
siderite_backends/sqlite/group.rs

siderite_backends/sqlite/
group.rs
