Change: Add RaftStateMachine::save_snapshot() method

Previously, `RaftStateMachine::install_snapshot()` was responsible for
two tasks:

- Replacing the state machine with the given snapshot.

- Persisting the snapshot to disk.

This commit introduces a new method,
`RaftStateMachine::save_snapshot()`, to handle the task of
persisting snapshots. With this change:

- `install_snapshot()` no longer needs to persist snapshots, though it
  can still do so without causing harm.
- The responsibility for persisting snapshots is now clearly separated.

- Part of #1333

Upgrade tip:

- Implement snapshot persistence logic in the new `save_snapshot()`
  method.
- It is recommended to remove snapshot persistence logic from
  `install_snapshot()` for better separation of concerns.

Author: drmingdrmer

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);
+        if new_snapshot.meta.last_log_id <= current_last {
+            return Ok(());
+        }
+
+        *current = Some(new_snapshot);
         Ok(())
     }
 

examples/raft-kv-memstore-grpc/src/store/mod.rs

@@ -141,23 +141,42 @@ impl RaftStateMachine<TypeConfig> for Arc<StateMachineStore> {
     async fn install_snapshot(&mut self, meta: &SnapshotMeta, snapshot: SnapshotData) -> Result<(), StorageError> {
         tracing::info!("install snapshot");
 
-        let new_snapshot = StoredSnapshot {
-            meta: meta.clone(),
-            data: snapshot,
-        };
-
         // Update the state machine.
         {
-            let d: pb::StateMachineData = prost::Message::decode(new_snapshot.data.as_ref())
-                .map_err(|e| StorageError::read_snapshot(None, &e))?;
+            let d: pb::StateMachineData =
+                prost::Message::decode(snapshot.as_ref()).map_err(|e| StorageError::read_snapshot(None, &e))?;
 
             let mut state_machine = self.state_machine.lock().unwrap();
             *state_machine = d;
         }
 
-        // Update current snapshot.
-        let mut current_snapshot = self.current_snapshot.lock().unwrap();
-        *current_snapshot = Some(new_snapshot);
+        let snapshot = Snapshot {
+            meta: meta.clone(),
+            snapshot,
+        };
+
+        self.save_snapshot(&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(),
+        };
+
+        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);
+        if new_snapshot.meta.last_log_id <= current_last {
+            return Ok(());
+        }
+
+        *current = Some(new_snapshot);
+
         Ok(())
     }
 

examples/raft-kv-memstore-network-v2/src/store.rs

@@ -165,21 +165,40 @@ impl RaftStateMachine<TypeConfig> for Arc<StateMachineStore> {
     async fn install_snapshot(&mut self, meta: &SnapshotMeta, snapshot: SnapshotData) -> Result<(), StorageError> {
         tracing::info!("install snapshot");
 
-        let new_snapshot = StoredSnapshot {
-            meta: meta.clone(),
-            data: snapshot,
-        };
-
         // Update the state machine.
         {
-            let updated_state_machine: StateMachineData = new_snapshot.data.clone();
+            let updated_state_machine: StateMachineData = snapshot.clone();
             let mut state_machine = self.state_machine.lock().unwrap();
             *state_machine = updated_state_machine;
         }
 
-        // Update current snapshot.
-        let mut current_snapshot = self.current_snapshot.lock().unwrap();
-        *current_snapshot = Some(new_snapshot);
+        let snapshot = Snapshot {
+            meta: meta.clone(),
+            snapshot,
+        };
+
+        self.save_snapshot(&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(),
+        };
+
+        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);
+        if new_snapshot.meta.last_log_id <= current_last {
+            return Ok(());
+        }
+
+        *current = Some(new_snapshot);
+
         Ok(())
     }
 

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);
+
+        if new_snapshot.meta.last_log_id <= current_last {
+            return Ok(());
+        }
+
+        *current = Some(new_snapshot);
         Ok(())
     }
 

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);
+        if new_snapshot.meta.last_log_id <= current_last {
+            return Ok(());
+        }
+
+        *current = Some(new_snapshot);
+
         Ok(())
     }
 

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);
+        if new_snapshot.meta.last_log_id <= current_last {
+            return Ok(());
+        }
+
+        *current = Some(new_snapshot);
 
-        // Update current snapshot.
-        *current_snapshot = Some(new_snapshot);
         Ok(())
     }