liminal_server/cluster/
discovery.rs1use std::collections::HashMap;
23use std::net::SocketAddr;
24use std::sync::{Arc, RwLock};
25
26use beamr::atom::AtomTable;
27use beamr::distribution::connection::ConnectionManager;
28use beamr::distribution::resolver::{NodeResolver, ResolveError, ResolveFuture, Resolver};
29
30#[derive(Debug, Default)]
35pub struct ClusterResolver {
36 nodes: RwLock<HashMap<String, SocketAddr>>,
37}
38
39impl ClusterResolver {
40 #[must_use]
42 pub fn new() -> Self {
43 Self {
44 nodes: RwLock::new(HashMap::new()),
45 }
46 }
47
48 pub fn register(&self, name: impl Into<String>, address: SocketAddr) {
50 self.nodes
51 .write()
52 .unwrap_or_else(std::sync::PoisonError::into_inner)
53 .insert(name.into(), address);
54 }
55
56 fn lookup(&self, name: &str) -> Option<SocketAddr> {
57 self.nodes
58 .read()
59 .unwrap_or_else(std::sync::PoisonError::into_inner)
60 .get(name)
61 .copied()
62 }
63}
64
65impl NodeResolver for ClusterResolver {
66 fn resolve<'a>(&'a self, name: &'a str) -> ResolveFuture<'a> {
67 let result = self.lookup(name).ok_or(ResolveError::NotFound);
68 Box::pin(async move { result })
69 }
70}
71
72#[must_use]
74fn seed_label(index: usize, address: SocketAddr) -> String {
75 format!("seed-{index}@{address}")
76}
77
78#[must_use]
85pub fn seed_resolver(seeds: &[SocketAddr]) -> (Arc<ClusterResolver>, Vec<String>) {
86 let resolver = Arc::new(ClusterResolver::new());
87 let labels = register_seed_labels(&resolver, seeds);
88 (resolver, labels)
89}
90
91pub fn register_seed_labels(resolver: &ClusterResolver, seeds: &[SocketAddr]) -> Vec<String> {
97 let mut labels = Vec::with_capacity(seeds.len());
98 for (index, address) in seeds.iter().enumerate() {
99 let label = seed_label(index, *address);
100 resolver.register(label.clone(), *address);
101 labels.push(label);
102 }
103 labels
104}
105
106#[must_use]
109pub fn as_resolver(resolver: Arc<ClusterResolver>) -> Resolver {
110 resolver
111}
112
113#[derive(Clone, Debug, Default, PartialEq, Eq)]
115pub struct SeedConnectOutcome {
116 pub attempted: usize,
118 pub connected: usize,
120}
121
122impl SeedConnectOutcome {
123 #[must_use]
126 pub const fn is_satisfied(&self) -> bool {
127 self.attempted == 0 || self.connected > 0
128 }
129}
130
131pub async fn connect_seeds(
138 connections: &ConnectionManager,
139 resolver: &Arc<ClusterResolver>,
140 atoms: &AtomTable,
141 labels: &[String],
142) -> SeedConnectOutcome {
143 let mut outcome = SeedConnectOutcome {
144 attempted: labels.len(),
145 connected: 0,
146 };
147 for label in labels {
148 match connections.connect(label).await {
149 Ok(connection) => {
150 let address = connection.peer_addr();
151 if let Some(name) = atoms.resolve(connection.node()).map(str::to_owned) {
152 resolver.register(name.clone(), address);
153 tracing::info!(
154 seed_label = %label,
155 peer = %name,
156 peer_addr = %address,
157 "connected to cluster seed node"
158 );
159 } else {
160 tracing::info!(
161 seed_label = %label,
162 peer_addr = %address,
163 "connected to cluster seed node"
164 );
165 }
166 outcome.connected += 1;
167 }
168 Err(error) => {
169 tracing::warn!(
170 seed_label = %label,
171 error = %error,
172 "cluster seed node unreachable at startup; continuing with reachable seeds"
173 );
174 }
175 }
176 }
177 outcome
178}
179
180#[cfg(test)]
181#[allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)]
182mod tests {
183 use super::{ClusterResolver, SeedConnectOutcome, as_resolver, seed_label, seed_resolver};
184 use beamr::distribution::resolver::{NodeResolver, ResolveError};
185 use std::net::SocketAddr;
186 use std::sync::Arc;
187 use std::task::{Context, Poll, Waker};
188
189 fn resolve_now(resolver: &ClusterResolver, name: &str) -> Result<SocketAddr, ResolveError> {
190 let mut context = Context::from_waker(Waker::noop());
191 let mut future = resolver.resolve(name);
192 match future.as_mut().poll(&mut context) {
193 Poll::Ready(result) => result,
194 Poll::Pending => panic!("cluster resolver future should be ready immediately"),
195 }
196 }
197
198 fn socket(address: &str) -> SocketAddr {
199 address.parse().expect("valid socket address")
200 }
201
202 #[test]
203 fn seed_resolver_maps_each_seed_to_a_synthetic_label() {
204 let seeds = vec![socket("127.0.0.1:9000"), socket("127.0.0.1:9001")];
205 let (resolver, labels) = seed_resolver(&seeds);
206
207 assert_eq!(labels.len(), 2);
208 assert_eq!(labels[0], seed_label(0, seeds[0]));
209 assert_eq!(labels[1], seed_label(1, seeds[1]));
210 assert_eq!(resolve_now(&resolver, &labels[0]), Ok(seeds[0]));
211 assert_eq!(resolve_now(&resolver, &labels[1]), Ok(seeds[1]));
212 }
213
214 #[test]
215 fn resolver_learns_real_peer_names() {
216 let resolver = ClusterResolver::new();
217 assert_eq!(
218 resolve_now(&resolver, "node-b@host"),
219 Err(ResolveError::NotFound)
220 );
221 resolver.register("node-b@host", socket("127.0.0.1:9100"));
222 assert_eq!(
223 resolve_now(&resolver, "node-b@host"),
224 Ok(socket("127.0.0.1:9100"))
225 );
226 }
227
228 #[test]
229 fn as_resolver_coerces_to_shared_handle() {
230 let (resolver, _labels) = seed_resolver(&[socket("127.0.0.1:9000")]);
231 let shared = as_resolver(Arc::clone(&resolver));
232 let mut context = Context::from_waker(Waker::noop());
234 let mut future = shared.resolve("seed-0@127.0.0.1:9000");
235 let outcome = match future.as_mut().poll(&mut context) {
236 Poll::Ready(result) => result,
237 Poll::Pending => panic!("future should be ready"),
238 };
239 assert_eq!(outcome, Ok(socket("127.0.0.1:9000")));
240 }
241
242 #[test]
243 fn outcome_is_satisfied_when_no_seeds_or_some_connected() {
244 assert!(
245 SeedConnectOutcome {
246 attempted: 0,
247 connected: 0
248 }
249 .is_satisfied()
250 );
251 assert!(
252 SeedConnectOutcome {
253 attempted: 3,
254 connected: 1
255 }
256 .is_satisfied()
257 );
258 assert!(
259 !SeedConnectOutcome {
260 attempted: 2,
261 connected: 0
262 }
263 .is_satisfied()
264 );
265 }
266}