laminar_core/cluster/control/
chitchat_kv.rs1#![allow(clippy::disallowed_types)] use std::sync::Arc;
6
7use async_trait::async_trait;
8
9use super::barrier::ClusterKv;
10use crate::cluster::discovery::NodeId;
11
12pub 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 #[must_use]
26 pub fn from_handle(handle: &chitchat::ChitchatHandle) -> Self {
27 Self {
28 chitchat: handle.chitchat(),
29 }
30 }
31
32 #[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 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 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}