Skip to main content

laminar_core/cluster/control/
chitchat_kv.rs

1//! Chitchat-backed [`ClusterKv`]. Reads filter by `live_nodes()` so
2//! suspected peers don't count toward quorum.
3#![allow(clippy::disallowed_types)] // cold control-plane path
4
5use std::sync::Arc;
6
7use async_trait::async_trait;
8
9use super::barrier::ClusterKv;
10use crate::cluster::discovery::NodeId;
11
12/// Chitchat-backed cluster KV.
13pub struct ChitchatKv {
14    chitchat: Arc<tokio::sync::Mutex<chitchat::Chitchat>>,
15}
16
17impl std::fmt::Debug for ChitchatKv {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        f.debug_struct("ChitchatKv").finish_non_exhaustive()
20    }
21}
22
23impl ChitchatKv {
24    /// Construct from a chitchat handle.
25    #[must_use]
26    pub fn from_handle(handle: &chitchat::ChitchatHandle) -> Self {
27        Self {
28            chitchat: handle.chitchat(),
29        }
30    }
31
32    /// Construct from a shared chitchat state.
33    #[must_use]
34    pub fn from_chitchat(chitchat: Arc<tokio::sync::Mutex<chitchat::Chitchat>>) -> Self {
35        Self { chitchat }
36    }
37}
38
39fn encode_node_id(node_id: NodeId) -> String {
40    format!("node-{}", node_id.0)
41}
42
43fn decode_chitchat_id(id: &chitchat::ChitchatId) -> Option<NodeId> {
44    id.node_id
45        .strip_prefix("node-")?
46        .parse::<u64>()
47        .ok()
48        .map(NodeId)
49}
50
51#[async_trait]
52impl ClusterKv for ChitchatKv {
53    async fn write(&self, key: &str, value: String) {
54        let mut guard = self.chitchat.lock().await;
55        guard.self_node_state().set(key, value);
56    }
57
58    async fn read_from(&self, who: NodeId, key: &str) -> Option<String> {
59        let target = encode_node_id(who);
60        let guard = self.chitchat.lock().await;
61
62        // Read the key from the node's newest generation — an older generation
63        // must not mask a key the newest generation dropped.
64        let mut best: Option<(u64, Option<String>)> = None;
65        for (cc_id, state) in guard.node_states() {
66            if cc_id.node_id == target
67                && best.as_ref().is_none_or(|&(g, _)| cc_id.generation_id > g)
68            {
69                best = Some((cc_id.generation_id, state.get(key).map(str::to_string)));
70            }
71        }
72        best.and_then(|(_, val)| val)
73    }
74
75    async fn scan(&self, key: &str) -> Vec<(NodeId, String)> {
76        let guard = self.chitchat.lock().await;
77        let live: Vec<&chitchat::ChitchatId> = guard.live_nodes().collect();
78        // Newest generation per node wins, even if it lacks the key — an older
79        // generation must not resurrect a value the newest one dropped.
80        let mut best: std::collections::HashMap<NodeId, (u64, Option<String>)> =
81            std::collections::HashMap::new();
82        for (cc_id, state) in guard.node_states() {
83            if !live.contains(&cc_id) {
84                continue;
85            }
86            let Some(node_id) = decode_chitchat_id(cc_id) else {
87                continue;
88            };
89            if best
90                .get(&node_id)
91                .is_none_or(|&(g, _)| cc_id.generation_id > g)
92            {
93                best.insert(
94                    node_id,
95                    (cc_id.generation_id, state.get(key).map(str::to_string)),
96                );
97            }
98        }
99        best.into_iter()
100            .filter_map(|(id, (_, val))| val.map(|v| (id, v)))
101            .collect()
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn node_id_encoding_roundtrips() {
111        for &v in &[1u64, 42, u64::MAX] {
112            let s = encode_node_id(NodeId(v));
113            assert!(s.starts_with("node-"), "got: {s}");
114            let parsed: u64 = s.strip_prefix("node-").unwrap().parse().unwrap();
115            assert_eq!(parsed, v);
116        }
117    }
118
119    #[test]
120    fn decode_rejects_unexpected_formats() {
121        let bad = chitchat::ChitchatId::new("foo".to_string(), 0, "127.0.0.1:1".parse().unwrap());
122        assert_eq!(decode_chitchat_id(&bad), None);
123    }
124
125    #[test]
126    fn decode_accepts_valid_format() {
127        let good =
128            chitchat::ChitchatId::new("node-42".to_string(), 0, "127.0.0.1:1".parse().unwrap());
129        assert_eq!(decode_chitchat_id(&good), Some(NodeId(42)));
130    }
131}