orch/cluster/
storage.rs

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

orch/cluster/
storage.rs
