orch/cluster/
storage.rs

1#![expect(
2    clippy::result_large_err,
3    reason = "openraft's trait methods return its own StorageError, so its size is not ours to change"
4)]
5
6use std::fmt::Debug;
7use std::io::Cursor;
8use std::ops::{Bound, RangeBounds};
9use std::path::Path;
10
11use openraft::storage::{LogState, RaftLogReader, RaftSnapshotBuilder, Snapshot};
12use openraft::{
13    BasicNode, Entry, EntryPayload, LogId, OptionalSend, RaftLogId, RaftStorage, SnapshotMeta,
14    StorageError, StorageIOError, StoredMembership, Vote,
15};
16use serde::de::DeserializeOwned;
17use serde::{Deserialize, Serialize};
18use sqlx::sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions};
19use sqlx::{Acquire, SqliteConnection};
20
21use super::{NodeId, Reply, TypeConfig};
22use crate::store::{TaskStore, kv, replica, tasks};
23
24type Err = StorageError<NodeId>;
25
26const APPLIED_KEY: &str = "raft:applied";
27const MEMBERSHIP_KEY: &str = "raft:membership";
28
29fn io<E: std::error::Error + 'static>(e: E) -> Err {
30    StorageIOError::read(&e).into()
31}
32
33#[expect(
34    clippy::needless_pass_by_value,
35    reason = "map_err passes the error by value"
36)]
37fn io_any(e: anyhow::Error) -> Err {
38    io(std::io::Error::other(format!("{e:#}")))
39}
40
41/// The replicated rows plus the Raft position they reflect
42#[derive(Serialize, Deserialize)]
43struct SnapshotData {
44    last_applied: Option<LogId<NodeId>>,
45    membership: StoredMembership<NodeId, BasicNode>,
46    tasks: Vec<serde_json::Value>,
47    tables: std::collections::BTreeMap<String, Vec<serde_json::Value>>,
48    kv: Vec<kv::Row>,
49}
50
51#[derive(Serialize, Deserialize)]
52struct StoredSnapshot {
53    meta: SnapshotMeta<NodeId, BasicNode>,
54    data: Vec<u8>,
55}
56
57/**
58 * The log and vote live in `raft.db`, the state machine is orch's own database. Apply writes
59 * each entry and `raft:applied` in one transaction on `orch.db`, so a crash never applies an
60 * entry twice, and a log that runs ahead of the state is normal for Raft
61 **/
62#[derive(Clone)]
63pub struct RaftStore {
64    log: SqlitePool,
65    store: TaskStore,
66}
67
68impl Debug for RaftStore {
69    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70        f.debug_struct("RaftStore")
71            .field("log", &self.log)
72            .finish_non_exhaustive()
73    }
74}
75
76impl RaftStore {
77    pub async fn open(raft_db: &Path, store: TaskStore) -> anyhow::Result<Self> {
78        let options = SqliteConnectOptions::new()
79            .filename(raft_db)
80            .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
81            .create_if_missing(true);
82        let log = SqlitePoolOptions::new()
83            .max_connections(1)
84            .connect_with(options)
85            .await?;
86        sqlx::query(
87            "CREATE TABLE IF NOT EXISTS log (idx INTEGER PRIMARY KEY, entry TEXT NOT NULL);
88             CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);",
89        )
90        .execute(&log)
91        .await?;
92        Ok(Self { log, store })
93    }
94
95    async fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, Err> {
96        let v: Option<(String,)> = sqlx::query_as("SELECT value FROM meta WHERE key = ?")
97            .bind(key)
98            .fetch_optional(&self.log)
99            .await
100            .map_err(io)?;
101        v.map(|(s,)| serde_json::from_str(&s).map_err(io))
102            .transpose()
103    }
104
105    async fn put<T: Serialize>(&self, key: &str, value: &T) -> Result<(), Err> {
106        let s = serde_json::to_string(value).map_err(io)?;
107        sqlx::query("INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)")
108            .bind(key)
109            .bind(s)
110            .execute(&self.log)
111            .await
112            .map_err(io)?;
113        Ok(())
114    }
115
116    async fn state_get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, Err> {
117        let v = self.store.kv_get(key).await.map_err(io_any)?;
118        v.map(|s| serde_json::from_str(&s).map_err(io)).transpose()
119    }
120}
121
122async fn state_put<T: Serialize>(
123    conn: &mut SqliteConnection,
124    now: &str,
125    key: &str,
126    value: &T,
127) -> Result<(), Err> {
128    let s = serde_json::to_string(value).map_err(io)?;
129    kv::set(conn, now, key, &s).await.map_err(io_any)
130}
131
132impl RaftLogReader<TypeConfig> for RaftStore {
133    async fn try_get_log_entries<RB: RangeBounds<u64> + Clone + Debug + OptionalSend>(
134        &mut self,
135        range: RB,
136    ) -> Result<Vec<Entry<TypeConfig>>, Err> {
137        let start = match range.start_bound() {
138            Bound::Included(&s) => s,
139            Bound::Excluded(&s) => s + 1,
140            Bound::Unbounded => 0,
141        };
142        let end = match range.end_bound() {
143            Bound::Included(&e) => e.saturating_add(1),
144            Bound::Excluded(&e) => e,
145            Bound::Unbounded => u64::MAX,
146        };
147        // SQLite integers are i64, so the open end is clamped
148        let rows: Vec<(String,)> =
149            sqlx::query_as("SELECT entry FROM log WHERE idx >= ? AND idx < ? ORDER BY idx")
150                .bind(i64::try_from(start).unwrap_or(i64::MAX))
151                .bind(i64::try_from(end).unwrap_or(i64::MAX))
152                .fetch_all(&self.log)
153                .await
154                .map_err(io)?;
155        rows.into_iter()
156            .map(|(s,)| serde_json::from_str(&s).map_err(io))
157            .collect()
158    }
159}
160
161impl RaftSnapshotBuilder<TypeConfig> for RaftStore {
162    async fn build_snapshot(&mut self) -> Result<Snapshot<TypeConfig>, Err> {
163        let mut tables = std::collections::BTreeMap::new();
164        for table in replica::TABLES {
165            let rows = replica::table_rows(self.store.pool(), table)
166                .await
167                .map_err(io_any)?;
168            tables.insert(table.to_string(), rows);
169        }
170
171        let snap = SnapshotData {
172            tables,
173            last_applied: self.state_get(APPLIED_KEY).await?,
174            membership: self.state_get(MEMBERSHIP_KEY).await?.unwrap_or_default(),
175            tasks: tasks::replicated_rows(self.store.pool())
176                .await
177                .map_err(io_any)?,
178            kv: kv::replicated_rows(self.store.pool())
179                .await
180                .map_err(io_any)?,
181        };
182        let data = serde_json::to_vec(&snap).map_err(io)?;
183        let last = snap.last_applied;
184        let meta = SnapshotMeta {
185            last_log_id: last,
186            last_membership: snap.membership,
187            snapshot_id: last
188                .map(|l| format!("{}-{}", l.leader_id, l.index))
189                .unwrap_or_default(),
190        };
191        self.put(
192            "snapshot",
193            &StoredSnapshot {
194                meta: meta.clone(),
195                data: data.clone(),
196            },
197        )
198        .await?;
199        Ok(Snapshot {
200            meta,
201            snapshot: Box::new(Cursor::new(data)),
202        })
203    }
204}
205
206impl RaftStorage<TypeConfig> for RaftStore {
207    type LogReader = Self;
208    type SnapshotBuilder = Self;
209
210    async fn save_vote(&mut self, vote: &Vote<NodeId>) -> Result<(), Err> {
211        self.put("vote", vote).await
212    }
213
214    async fn read_vote(&mut self) -> Result<Option<Vote<NodeId>>, Err> {
215        self.get("vote").await
216    }
217
218    async fn save_committed(&mut self, committed: Option<LogId<NodeId>>) -> Result<(), Err> {
219        self.put("committed", &committed).await
220    }
221
222    async fn read_committed(&mut self) -> Result<Option<LogId<NodeId>>, Err> {
223        Ok(self
224            .get::<Option<LogId<NodeId>>>("committed")
225            .await?
226            .flatten())
227    }
228
229    async fn get_log_state(&mut self) -> Result<LogState<TypeConfig>, Err> {
230        let purged = self.get::<Option<LogId<NodeId>>>("purged").await?.flatten();
231        let last: Option<(String,)> =
232            sqlx::query_as("SELECT entry FROM log ORDER BY idx DESC LIMIT 1")
233                .fetch_optional(&self.log)
234                .await
235                .map_err(io)?;
236        let last = match last {
237            Some((s,)) => Some(
238                *serde_json::from_str::<Entry<TypeConfig>>(&s)
239                    .map_err(io)?
240                    .get_log_id(),
241            ),
242            None => purged,
243        };
244        Ok(LogState {
245            last_purged_log_id: purged,
246            last_log_id: last,
247        })
248    }
249
250    async fn get_log_reader(&mut self) -> Self::LogReader {
251        self.clone()
252    }
253
254    async fn append_to_log<I>(&mut self, entries: I) -> Result<(), Err>
255    where
256        I: IntoIterator<Item = Entry<TypeConfig>> + OptionalSend,
257    {
258        let rows = entries
259            .into_iter()
260            .map(|e| {
261                Ok((
262                    i64::try_from(e.log_id.index).map_err(io)?,
263                    serde_json::to_string(&e).map_err(io)?,
264                ))
265            })
266            .collect::<Result<Vec<_>, Err>>()?;
267        let mut tx = self.log.begin().await.map_err(io)?;
268        for (idx, s) in rows {
269            sqlx::query("INSERT OR REPLACE INTO log (idx, entry) VALUES (?, ?)")
270                .bind(idx)
271                .bind(s)
272                .execute(&mut *tx)
273                .await
274                .map_err(io)?;
275        }
276
277        tx.commit().await.map_err(io)
278    }
279
280    async fn delete_conflict_logs_since(&mut self, log_id: LogId<NodeId>) -> Result<(), Err> {
281        sqlx::query("DELETE FROM log WHERE idx >= ?")
282            .bind(i64::try_from(log_id.index).map_err(io)?)
283            .execute(&self.log)
284            .await
285            .map_err(io)?;
286        Ok(())
287    }
288
289    async fn purge_logs_upto(&mut self, log_id: LogId<NodeId>) -> Result<(), Err> {
290        self.put("purged", &Some(log_id)).await?;
291        sqlx::query("DELETE FROM log WHERE idx <= ?")
292            .bind(i64::try_from(log_id.index).map_err(io)?)
293            .execute(&self.log)
294            .await
295            .map_err(io)?;
296        Ok(())
297    }
298
299    async fn last_applied_state(
300        &mut self,
301    ) -> Result<(Option<LogId<NodeId>>, StoredMembership<NodeId, BasicNode>), Err> {
302        Ok((
303            self.state_get(APPLIED_KEY).await?,
304            self.state_get(MEMBERSHIP_KEY).await?.unwrap_or_default(),
305        ))
306    }
307
308    async fn apply_to_state_machine(
309        &mut self,
310        entries: &[Entry<TypeConfig>],
311    ) -> Result<Vec<Reply>, Err> {
312        let mut tx = self.store.begin_write().await.map_err(io)?;
313        let mut out = Vec::with_capacity(entries.len());
314        for e in entries {
315            let reply = match &e.payload {
316                EntryPayload::Blank => Reply::Done,
317                EntryPayload::Normal(req) => {
318                    // A savepoint per entry, so a failed command rolls back alone
319                    let mut sp = tx.begin().await.map_err(io)?;
320                    match req.cmd.clone().apply(&mut sp, &req.now, true).await {
321                        Ok(reply) => {
322                            sp.commit().await.map_err(io)?;
323                            reply
324                        }
325                        Err(err) => {
326                            sp.rollback().await.map_err(io)?;
327                            Reply::Err(format!("{err:#}"))
328                        }
329                    }
330                }
331                EntryPayload::Membership(m) => {
332                    let stored = StoredMembership::new(Some(e.log_id), m.clone());
333                    state_put(&mut tx, "", MEMBERSHIP_KEY, &stored).await?;
334                    Reply::Done
335                }
336            };
337            state_put(&mut tx, "", APPLIED_KEY, &e.log_id).await?;
338            out.push(reply);
339        }
340
341        tx.commit().await.map_err(io)?;
342        Ok(out)
343    }
344
345    async fn get_snapshot_builder(&mut self) -> Self::SnapshotBuilder {
346        self.clone()
347    }
348
349    async fn begin_receiving_snapshot(&mut self) -> Result<Box<Cursor<Vec<u8>>>, Err> {
350        Ok(Box::new(Cursor::new(Vec::new())))
351    }
352
353    async fn install_snapshot(
354        &mut self,
355        meta: &SnapshotMeta<NodeId, BasicNode>,
356        snapshot: Box<Cursor<Vec<u8>>>,
357    ) -> Result<(), Err> {
358        let data = snapshot.into_inner();
359        let snap: SnapshotData = serde_json::from_slice(&data).map_err(io)?;
360        let mut tx = self.store.begin_write().await.map_err(io)?;
361        tasks::replace_replicated(&mut tx, &snap.tasks)
362            .await
363            .map_err(io_any)?;
364        for table in replica::TABLES {
365            let rows = snap
366                .tables
367                .get(table)
368                .map(Vec::as_slice)
369                .unwrap_or_default();
370            replica::replace_table(&mut tx, table, rows)
371                .await
372                .map_err(io_any)?;
373        }
374        kv::replace_replicated(&mut tx, &snap.kv)
375            .await
376            .map_err(io_any)?;
377        state_put(&mut tx, "", APPLIED_KEY, &meta.last_log_id).await?;
378        state_put(&mut tx, "", MEMBERSHIP_KEY, &meta.last_membership).await?;
379        tx.commit().await.map_err(io)?;
380        self.put(
381            "snapshot",
382            &StoredSnapshot {
383                meta: meta.clone(),
384                data,
385            },
386        )
387        .await
388    }
389
390    async fn get_current_snapshot(&mut self) -> Result<Option<Snapshot<TypeConfig>>, Err> {
391        let s: Option<StoredSnapshot> = self.get("snapshot").await?;
392        Ok(s.map(|s| Snapshot {
393            meta: s.meta,
394            snapshot: Box::new(Cursor::new(s.data)),
395        }))
396    }
397}
orch/cluster/storage.rs

orch/cluster/
storage.rs
