From 35b360976346cb1db86af339603c7445a90b7416 Mon Sep 17 00:00:00 2001
From: Robin Nabel <rnabel@ucdavis.edu>
Date: Mon, 21 Oct 2024 12:52:41 +0200
Subject: [PATCH] Remove NodeStateSnapshot (#151)

* Remove NodeStateSnapshot

* fixing unit test and rustfmt

---------

Co-authored-by: Paul Masurel <paul@quickwit.io>
---
 chitchat-test/tests/cli.rs |  8 ++++----
 chitchat/src/lib.rs        |  4 ++--
 chitchat/src/state.rs      | 27 +++++++--------------------
 3 files changed, 13 insertions(+), 26 deletions(-)

diff --git a/chitchat-test/tests/cli.rs b/chitchat-test/tests/cli.rs
index a142a61..93ff6a0 100644
--- a/chitchat-test/tests/cli.rs
+++ b/chitchat-test/tests/cli.rs
@@ -75,11 +75,11 @@ fn test_multiple_nodes() {
     assert_eq!(info.live_nodes.len(), 5);
     assert_eq!(info.dead_nodes.len(), 0);
 
-    assert!(info.cluster_state.node_state_snapshots.get(3).is_some());
+    assert!(info.cluster_state.node_states.get(3).is_some());
     // Check that "some_key" we set on this local node (localhost:13001) is
     // indeed set to be "some_value"
-    let node = info.cluster_state.node_state_snapshots.get(1).unwrap();
-    let versioned_value = node.node_state.get_versioned("some_key").unwrap();
+    let node_state = info.cluster_state.node_states.get(1).unwrap();
+    let versioned_value = node_state.get_versioned("some_key").unwrap();
     assert_eq!(versioned_value.value, "some_value");
 }
 
@@ -88,7 +88,7 @@ fn test_multiple_nodes_with_dns_resolution_for_seed() {
     let _child_handles = setup_nodes(12_000, 5, 5, true);
     // Check node states through api.
     let info = get_node_info("http://127.0.0.1:12001").unwrap();
-    assert!(info.cluster_state.node_state_snapshots.get(3).is_some());
+    assert!(info.cluster_state.node_states.get(3).is_some());
     assert_eq!(info.cluster_id, "testing");
     assert_eq!(info.live_nodes.len(), 5);
     assert_eq!(info.dead_nodes.len(), 0);
diff --git a/chitchat/src/lib.rs b/chitchat/src/lib.rs
index 95dfb84..de26ad8 100644
--- a/chitchat/src/lib.rs
+++ b/chitchat/src/lib.rs
@@ -361,8 +361,8 @@ impl Chitchat {
     /// Disclaimer:
     /// The callback is required to be as light as possible.
     /// In particular,
-    /// - it should not access the cluster state (as it is locked at the moment of the
-    /// execution of the callback.
+    /// - it should not access the cluster state (as it is locked at the moment of the execution of
+    ///   the callback.
     /// - it should be fast: the callback is executed in an async context.
     ///
     /// The callback is called with a [`KeyChangeEvent`] that contains the key stripped of the
diff --git a/chitchat/src/state.rs b/chitchat/src/state.rs
index b2a09c2..ecc0a89 100644
--- a/chitchat/src/state.rs
+++ b/chitchat/src/state.rs
@@ -793,30 +793,17 @@ impl<'a> StaleNode<'a> {
     }
 }
 
-#[derive(Debug, Serialize, Deserialize)]
-pub struct NodeStateSnapshot {
-    pub chitchat_id: ChitchatId,
-    pub node_state: NodeState,
-}
-
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ClusterStateSnapshot {
-    pub node_state_snapshots: Vec<NodeStateSnapshot>,
+    pub node_states: Vec<NodeState>,
     pub seed_addrs: HashSet<SocketAddr>,
 }
 
 impl From<&ClusterState> for ClusterStateSnapshot {
     fn from(cluster_state: &ClusterState) -> Self {
-        let node_state_snapshots = cluster_state
-            .node_states
-            .iter()
-            .map(|(chitchat_id, node_state)| NodeStateSnapshot {
-                chitchat_id: chitchat_id.clone(),
-                node_state: node_state.clone(),
-            })
-            .collect();
+        let node_states = cluster_state.node_states.values().cloned().collect();
         Self {
-            node_state_snapshots,
+            node_states,
             seed_addrs: cluster_state.seed_addrs(),
         }
     }
@@ -1226,12 +1213,11 @@ mod tests {
         // GC if tombstone (=100) + grace_period > heartbeat (=110).
         tokio::time::advance(Duration::from_secs(5)).await;
         cluster_state.gc_keys_marked_for_deletion(Duration::from_secs(10));
-        assert!(cluster_state
+        assert!(!cluster_state
             .node_state(&node1)
             .unwrap()
             .key_values
-            .get("key_a")
-            .is_none());
+            .contains_key("key_a"));
         cluster_state
             .node_state(&node1)
             .unwrap()
@@ -1455,7 +1441,8 @@ mod tests {
             node1_state.set_with_version("key_b".to_string(), "2".to_string(), 2); // 2
 
             let node2_state = cluster_state.node_state_mut(&node2);
-            node2_state.set_with_version("key_c".to_string(), "3".to_string(), 2); // 2
+            node2_state.set_with_version("key_c".to_string(), "3".to_string(), 2);
+            // 2
         }
 
         {