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

orch/cluster/
storage.rs
