Add a snapshot size test.

kvraft/src/testing_utils/config.rs

@@ -11,7 +11,7 @@ use ruaft::{Persister, RpcClient};
 
 use crate::client::Clerk;
 use crate::server::KVServer;
-use crate::testing_utils::memory_persister::MemoryStorage;
+use crate::testing_utils::memory_persister::{MemoryPersister, MemoryStorage};
 use crate::testing_utils::rpcs::register_kv_server;
 
 struct ConfigState {
@@ -251,24 +251,32 @@ impl Config {
 }
 
 impl Config {
-    pub fn check_log_size(&self, upper: usize) -> Result<()> {
+    fn check_size(
+        &self,
+        upper: usize,
+        size_fn: impl Fn(&MemoryPersister) -> usize,
+    ) -> Result<()> {
         let mut over_limits = String::new();
         for (index, p) in self.storage.lock().all().iter().enumerate() {
-            if p.state_size() > upper {
-                let str =
-                    format!(" (index {}, size {})", index, p.state_size());
+            let size = size_fn(p);
+            if size > upper {
+                let str = format!(" (index {}, size {})", index, size);
                 over_limits.push_str(&str);
             }
         }
         if over_limits.len() != 0 {
-            anyhow::bail!(
-                "logs were not trimmed to {}: {}",
-                upper,
-                over_limits
-            );
+            anyhow::bail!("logs were not trimmed to {}:{}", upper, over_limits);
         }
         Ok(())
     }
+
+    pub fn check_log_size(&self, upper: usize) -> Result<()> {
+        self.check_size(upper, MemoryPersister::state_size)
+    }
+
+    pub fn check_snapshot_size(&self, upper: usize) -> Result<()> {
+        self.check_size(upper, MemoryPersister::snapshot_size)
+    }
 }
 
 pub fn make_config(

kvraft/src/testing_utils/memory_persister.rs

@@ -51,6 +51,10 @@ impl MemoryPersister {
     pub fn restore(&self, state: State) {
         *self.state.lock() = state;
     }
+
+    pub fn snapshot_size(&self) -> usize {
+        self.state.lock().snapshot.len()
+    }
 }
 
 #[derive(Default)]

tests/snapshot_tests.rs

@@ -30,7 +30,8 @@ fn install_snapshot_rpc() {
         clerk.put("b", "B");
     }
 
-    cfg.check_log_size(MAX_RAFT_STATE * 2).expect("Hi");
+    cfg.check_log_size(MAX_RAFT_STATE * 2)
+        .expect("Log does not seem to be trimmed:");
 
     // Swap majority and minority.
     let (mut majority, mut minority) = (minority, majority);
@@ -62,3 +63,30 @@ fn install_snapshot_rpc() {
 
     cfg.end();
 }
+
+#[test]
+fn snapshot_size() {
+    const SERVERS: usize = 3;
+    const MAX_RAFT_STATE: usize = 1000;
+    const MAX_SNAPSHOT_STATE: usize = 500;
+    let cfg = Arc::new(make_config(SERVERS, false, MAX_RAFT_STATE));
+    defer!(cfg.clean_up());
+
+    let mut clerk = cfg.make_clerk();
+
+    cfg.begin("Test: snapshot size is reasonable (3B)");
+
+    for _ in 0..200 {
+        clerk.put("x", "0");
+        assert_eq!(clerk.get("x"), Some("0".to_owned()));
+        clerk.put("x", "1");
+        assert_eq!(clerk.get("x"), Some("1".to_owned()));
+    }
+
+    cfg.check_log_size(MAX_RAFT_STATE * 2)
+        .expect("Log does not seem to be trimmed:");
+    cfg.check_snapshot_size(MAX_SNAPSHOT_STATE)
+        .expect("Snapshot size is too big:");
+
+    cfg.end();
+}