Change: Add: RaftStateMachine::save_snapshot()

Before this commit, `RaftStateMachine::install_snapshot()` is
responsible to 1) replace the state machine with the provided snapshot
and 2) persist the snapshot on disk.
Since this commit, a new method `RaftStateMachine::save_snapshot()` is
responsible to persist a snapshot. `install_snapshot()` does not have to
do this. But doing this has no harm.

- Part of #1333

Upgrade tip:

Implement persisting snapshot in `save_snapshot()` and it is recommended
to remove persisting snpashot from `install_snapshot()`.

Author: drmingdrmer

Diff for: cluster_benchmark/tests/benchmark/store.rs

@@ -162,22 +162,14 @@ impl RaftSnapshotBuilder<TypeConfig> for Arc<StateMachineStore> {
             snapshot_id,
         };
 
-        let snapshot = StoredSnapshot {
-            meta: meta.clone(),
-            data: data.clone(),
+        let snapshot = Snapshot {
+            meta,
+            snapshot: Cursor::new(data),
         };
 
-        {
-            let mut current_snapshot = self.current_snapshot.write().await;
-            *current_snapshot = Some(snapshot);
-        }
-
-        tracing::info!(snapshot_size, "log compaction complete");
+        self.save_snapshot(&snapshot).await?;
 
-        Ok(Snapshot {
-            meta,
-            snapshot: Cursor::new(data),
-        })
+        Ok(snapshot)
     }
 }
 
@@ -304,9 +296,25 @@ impl RaftStateMachine<TypeConfig> for Arc<StateMachineStore> {
             *sm = new_sm;
         }
 
-        // Update current snapshot.
-        let mut current_snapshot = self.current_snapshot.write().await;
-        *current_snapshot = Some(new_snapshot);
+        Ok(())
+    }
+
+    async fn save_snapshot(&mut self, snapshot: &Snapshot<TypeConfig>) -> Result<(), StorageError<TypeConfig>> {
+        let new_snapshot = StoredSnapshot {
+            meta: snapshot.meta.clone(),
+            data: snapshot.snapshot.clone().into_inner(),
+        };
+
+        let mut current = self.current_snapshot.write().await;
+
+        // Only save it if the new snapshot contains more recent data than the current snapshot.
+
+        let current_last = current.as_ref().and_then(|s| s.meta.last_log_id.clone());
+        if new_snapshot.meta.last_log_id <= current_last {
+            return Ok(());
+        }
+
+        *current = Some(new_snapshot);
         Ok(())
     }
 

Diff for: examples/raft-kv-memstore-opendal-snapshot-data/src/store.rs

@@ -123,20 +123,14 @@ impl RaftSnapshotBuilder<TypeConfig> for Arc<StateMachineStore> {
             snapshot_id: snapshot_id.clone(),
         };
 
-        let snapshot = StoredSnapshot {
-            meta: meta.clone(),
-            data: snapshot_id.clone(),
+        let snapshot = Snapshot {
+            meta,
+            snapshot: snapshot_id,
         };
 
-        {
-            let mut current_snapshot = self.current_snapshot.lock().unwrap();
-            *current_snapshot = Some(snapshot);
-        }
+        self.save_snapshot(&snapshot).await?;
 
-        Ok(Snapshot {
-            meta,
-            snapshot: snapshot_id,
-        })
+        Ok(snapshot)
     }
 }
 
@@ -201,9 +195,26 @@ impl RaftStateMachine<TypeConfig> for Arc<StateMachineStore> {
             *state_machine = updated_state_machine;
         }
 
-        // Update current snapshot.
-        let mut current_snapshot = self.current_snapshot.lock().unwrap();
-        *current_snapshot = Some(new_snapshot);
+        Ok(())
+    }
+
+    async fn save_snapshot(&mut self, snapshot: &Snapshot) -> Result<(), StorageError> {
+        let new_snapshot = StoredSnapshot {
+            meta: snapshot.meta.clone(),
+            data: snapshot.snapshot.clone(),
+        };
+
+        let mut current = self.current_snapshot.lock().unwrap();
+
+        // Only save it if the new snapshot contains more recent data than the current snapshot.
+
+        let current_last = current.as_ref().and_then(|s| s.meta.last_log_id.clone());
+
+        if new_snapshot.meta.last_log_id <= current_last {
+            return Ok(());
+        }
+
+        *current = Some(new_snapshot);
         Ok(())
     }
 

Diff for: examples/raft-kv-memstore-singlethreaded/src/store.rs

@@ -157,20 +157,14 @@ impl RaftSnapshotBuilder<TypeConfig> for Rc<StateMachineStore> {
             snapshot_id,
         };
 
-        let snapshot = StoredSnapshot {
-            meta: meta.clone(),
-            data: data.clone(),
+        let snapshot = Snapshot {
+            meta,
+            snapshot: Cursor::new(data),
         };
 
-        {
-            let mut current_snapshot = self.current_snapshot.borrow_mut();
-            *current_snapshot = Some(snapshot);
-        }
+        self.save_snapshot(&snapshot).await?;
 
-        Ok(Snapshot {
-            meta,
-            snapshot: Cursor::new(data),
-        })
+        Ok(snapshot)
     }
 }
 
@@ -238,9 +232,26 @@ impl RaftStateMachine<TypeConfig> for Rc<StateMachineStore> {
             *state_machine = updated_state_machine;
         }
 
-        // Update current snapshot.
-        let mut current_snapshot = self.current_snapshot.borrow_mut();
-        *current_snapshot = Some(new_snapshot);
+        Ok(())
+    }
+
+    async fn save_snapshot(&mut self, snapshot: &Snapshot) -> Result<(), StorageError> {
+        let new_snapshot = StoredSnapshot {
+            meta: snapshot.meta.clone(),
+            data: snapshot.snapshot.clone().into_inner(),
+        };
+
+        let mut current = self.current_snapshot.borrow_mut();
+
+        // Only save it if the new snapshot contains more recent data than the current snapshot.
+
+        let current_last = current.as_ref().and_then(|s| s.meta.last_log_id.clone());
+        if new_snapshot.meta.last_log_id <= current_last {
+            return Ok(());
+        }
+
+        *current = Some(new_snapshot);
+
         Ok(())
     }
 

Diff for: examples/raft-kv-memstore/src/store/mod.rs

@@ -98,9 +98,6 @@ impl RaftSnapshotBuilder<TypeConfig> for Arc<StateMachineStore> {
         let last_applied_log = state_machine.last_applied_log;
         let last_membership = state_machine.last_membership.clone();
 
-        // Lock the current snapshot before releasing the lock on the state machine, to avoid a race
-        // condition on the written snapshot
-        let mut current_snapshot = self.current_snapshot.write().await;
         drop(state_machine);
 
         let snapshot_idx = self.snapshot_idx.fetch_add(1, Ordering::Relaxed) + 1;
@@ -116,17 +113,14 @@ impl RaftSnapshotBuilder<TypeConfig> for Arc<StateMachineStore> {
             snapshot_id,
         };
 
-        let snapshot = StoredSnapshot {
+        let snapshot = Snapshot {
             meta: meta.clone(),
-            data: data.clone(),
+            snapshot: Cursor::new(data.clone()),
         };
 
-        *current_snapshot = Some(snapshot);
+        self.save_snapshot(&snapshot).await?;
 
-        Ok(Snapshot {
-            meta,
-            snapshot: Cursor::new(data),
-        })
+        Ok(snapshot)
     }
 }
 
@@ -203,13 +197,26 @@ impl RaftStateMachine<TypeConfig> for Arc<StateMachineStore> {
         let mut state_machine = self.state_machine.write().await;
         *state_machine = updated_state_machine;
 
-        // Lock the current snapshot before releasing the lock on the state machine, to avoid a race
-        // condition on the written snapshot
-        let mut current_snapshot = self.current_snapshot.write().await;
-        drop(state_machine);
+        Ok(())
+    }
+
+    async fn save_snapshot(&mut self, snapshot: &Snapshot<TypeConfig>) -> Result<(), StorageError<TypeConfig>> {
+        let new_snapshot = StoredSnapshot {
+            meta: snapshot.meta.clone(),
+            data: snapshot.snapshot.clone().into_inner(),
+        };
+
+        let mut current = self.current_snapshot.write().await;
+
+        // Only save it if the new snapshot contains more recent data than the current snapshot.
+
+        let current_last = current.as_ref().and_then(|s| s.meta.last_log_id.clone());
+        if new_snapshot.meta.last_log_id <= current_last {
+            return Ok(());
+        }
+
+        *current = Some(new_snapshot);
 
-        // Update current snapshot.
-        *current_snapshot = Some(new_snapshot);
         Ok(())
     }
 

Diff for: examples/raft-kv-rocksdb/src/store.rs

@@ -13,6 +13,8 @@ use openraft::RaftSnapshotBuilder;
 use openraft_rocksstore::log_store::RocksLogStore;
 use rocksdb::ColumnFamily;
 use rocksdb::ColumnFamilyDescriptor;
+use rocksdb::Direction;
+use rocksdb::IteratorMode;
 use rocksdb::Options;
 use rocksdb::DB;
 use serde::Deserialize;
@@ -100,17 +102,14 @@ impl RaftSnapshotBuilder<TypeConfig> for StateMachineStore {
             snapshot_id,
         };
 
-        let snapshot = StoredSnapshot {
-            meta: meta.clone(),
-            data: kv_json.clone(),
+        let snapshot = Snapshot {
+            meta,
+            snapshot: Cursor::new(kv_json),
         };
 
-        self.set_current_snapshot_(snapshot)?;
+        self.save_snapshot(&snapshot).await?;
 
-        Ok(Snapshot {
-            meta,
-            snapshot: Cursor::new(kv_json),
-        })
+        Ok(snapshot)
     }
 }
 
@@ -146,18 +145,49 @@ impl StateMachineStore {
         Ok(())
     }
 
+    /// List all snapshots in the db, returns the last
     fn get_current_snapshot_(&self) -> StorageResult<Option<StoredSnapshot>> {
-        Ok(self
-            .db
-            .get_cf(self.store(), b"snapshot")
-            .map_err(|e| StorageError::read(&e))?
-            .and_then(|v| serde_json::from_slice(&v).ok()))
+        let mut last_snapshot_key = None;
+
+        let it = self.db.iterator_cf(self.store(), IteratorMode::From(b"snapshot-", Direction::Forward));
+
+        for kv in it {
+            let (key, _value) = kv.map_err(|e| StorageError::read(&e))?;
+            if key.starts_with(b"snapshot-") {
+                last_snapshot_key = Some(key.to_vec());
+            } else {
+                break;
+            }
+        }
+        let Some(key) = last_snapshot_key else {
+            return Ok(None);
+        };
+
+        let data = self.db.get_cf(self.store(), &key).map_err(|e| StorageError::read(&e))?.unwrap();
+
+        let snap: StoredSnapshot = serde_json::from_slice(&data).map_err(|e| StorageError::read_snapshot(None, &e))?;
+
+        Ok(Some(snap))
     }
 
+    /// Save snapshot by last-log-id and when reading, get the last one as the current.
+    ///
+    /// So that writing an old one won't overwrite the newer one.
+    /// In a real world application, the old ones should be cleaned.
     fn set_current_snapshot_(&self, snap: StoredSnapshot) -> StorageResult<()> {
+        let last_log_id = snap.meta.last_log_id.as_ref();
+
+        let id_str = Self::order_preserved_log_id_string(last_log_id);
+
+        let key = format!("snapshot-{id_str}");
         self.db
-            .put_cf(self.store(), b"snapshot", serde_json::to_vec(&snap).unwrap().as_slice())
+            .put_cf(
+                self.store(),
+                key.as_bytes(),
+                serde_json::to_vec(&snap).unwrap().as_slice(),
+            )
             .map_err(|e| StorageError::write_snapshot(Some(snap.meta.signature()), &e))?;
+
         self.flush(ErrorSubject::Snapshot(Some(snap.meta.signature())), ErrorVerb::Write)?;
         Ok(())
     }
@@ -170,6 +200,14 @@ impl StateMachineStore {
     fn store(&self) -> &ColumnFamily {
         self.db.cf_handle("store").unwrap()
     }
+
+    fn order_preserved_log_id_string(log_id: Option<&LogId>) -> String {
+        let term = log_id.map(|l| l.committed_leader_id().term).unwrap_or_default();
+        let node_id = log_id.map(|l| l.committed_leader_id().node_id).unwrap_or_default();
+        let index = log_id.map(|l| l.index).unwrap_or_default();
+
+        format!("{:020}-{:020}-{:020}", term, node_id, index)
+    }
 }
 
 impl RaftStateMachine<TypeConfig> for StateMachineStore {
@@ -227,7 +265,16 @@ impl RaftStateMachine<TypeConfig> for StateMachineStore {
             data: snapshot.into_inner(),
         };
 
-        self.update_state_machine_(new_snapshot.clone()).await?;
+        self.update_state_machine_(new_snapshot).await?;
+
+        Ok(())
+    }
+
+    async fn save_snapshot(&mut self, snapshot: &Snapshot) -> Result<(), StorageError> {
+        let new_snapshot = StoredSnapshot {
+            meta: snapshot.meta.clone(),
+            data: snapshot.snapshot.clone().into_inner(),
+        };
 
         self.set_current_snapshot_(new_snapshot)?;