From dfa35d42df7be1f05a4a6e62d6f1218ba3cdf679 Mon Sep 17 00:00:00 2001
From: Robin Nabel <rnabel@ucdavis.edu>
Date: Sat, 24 Aug 2024 16:13:20 +0100
Subject: [PATCH] Remove NodeStateSnapshot

---
 chitchat-test/tests/cli.rs |  8 ++++----
 chitchat/src/state.rs      | 19 +++----------------
 2 files changed, 7 insertions(+), 20 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/state.rs b/chitchat/src/state.rs
index b2a09c2..3756792 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(),
         }
     }