1use std::collections::HashMap;
25use std::sync::{Arc, Mutex, OnceLock};
26
27use beamr::atom::{Atom, AtomTable};
28use beamr::distribution::connection::ConnectionManager;
29use beamr::distribution::control::{encode_pg_update_frame, encode_send_frame};
30use beamr::distribution::pg::{PgRegistry, PgUpdate, RemoteMember};
31use beamr::native::ProcessContext;
32use beamr::term::Term;
33
34use crate::cluster::discovery::ClusterResolver;
35use liminal::channel::{ClusterObserver, encode_envelope};
36use liminal::envelope::Envelope;
37
38#[derive(Clone)]
41pub struct ClusterSync {
42 inner: Arc<SyncInner>,
43}
44
45struct SyncInner {
46 pg: Arc<PgRegistry>,
47 atoms: Arc<AtomTable>,
48 connections: ConnectionManager,
49 local_node: Atom,
52 _resolver: Arc<ClusterResolver>,
56 local: Mutex<HashMap<Atom, Vec<u64>>>,
59 peers: OnceLock<PeerListener>,
62}
63
64pub type PeerListener = Arc<dyn Fn(&str, bool) + Send + Sync>;
66
67impl std::fmt::Debug for ClusterSync {
68 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69 formatter
70 .debug_struct("ClusterSync")
71 .field("local_node", &self.inner.local_node)
72 .finish_non_exhaustive()
73 }
74}
75
76impl ClusterSync {
77 #[must_use]
81 pub fn new(
82 pg: Arc<PgRegistry>,
83 atoms: Arc<AtomTable>,
84 connections: ConnectionManager,
85 local_node: Atom,
86 resolver: Arc<ClusterResolver>,
87 ) -> Self {
88 Self {
89 inner: Arc::new(SyncInner {
90 pg,
91 atoms,
92 connections,
93 local_node,
94 _resolver: resolver,
95 local: Mutex::new(HashMap::new()),
96 peers: OnceLock::new(),
97 }),
98 }
99 }
100
101 pub fn set_peer_listener(&self, listener: PeerListener) -> bool {
104 self.inner.peers.set(listener).is_ok()
105 }
106
107 pub(crate) fn peer_changed(&self, name: &str, present: bool) {
109 if let Some(listener) = self.inner.peers.get() {
110 listener(name, present);
111 }
112 }
113
114 fn scope(&self) -> Atom {
116 self.inner.pg.default_scope()
117 }
118
119 fn group(&self, channel: &str) -> Atom {
121 self.inner.atoms.intern(channel)
122 }
123
124 #[must_use]
126 pub fn remote_targets(&self, channel: &str) -> Vec<RemoteMember> {
127 let group = self.group(channel);
128 self.inner.pg.remote_members(self.scope(), group)
129 }
130
131 fn record_local(&self, group: Atom, pid: u64) {
133 let mut local = self.lock_local();
134 let pids = local.entry(group).or_default();
135 if !pids.contains(&pid) {
136 pids.push(pid);
137 }
138 drop(local);
139 }
140
141 fn forget_local(&self, group: Atom, pid: u64) {
143 let mut local = self.lock_local();
144 if let Some(pids) = local.get_mut(&group) {
145 pids.retain(|candidate| *candidate != pid);
146 if pids.is_empty() {
147 local.remove(&group);
148 }
149 }
150 }
151
152 fn local_memberships(&self) -> Vec<(Atom, u64)> {
154 let local = self.lock_local();
155 local
156 .iter()
157 .flat_map(|(group, pids)| pids.iter().map(move |pid| (*group, *pid)))
158 .collect()
159 }
160
161 fn send_to_member(&self, member: RemoteMember, frame_bytes: &[u8]) {
163 let Some(to_pid) = Term::try_pid(member.pid_number) else {
167 tracing::warn!(
168 pid_number = member.pid_number,
169 "remote member pid out of immediate range; skipping cross-node delivery"
170 );
171 return;
172 };
173 let mut context = ProcessContext::new();
174 let Ok(payload) = context.alloc_binary(frame_bytes) else {
175 tracing::warn!("failed to allocate cross-node envelope payload");
176 return;
177 };
178 let Ok(frame) = encode_send_frame(
179 Term::atom(beamr::atom::Atom::OK),
180 to_pid,
181 payload,
182 &self.inner.atoms,
183 ) else {
184 tracing::warn!("failed to encode cross-node send frame");
185 return;
186 };
187 self.write_frame(member.node, &frame);
188 }
189
190 fn write_frame(&self, node: Atom, frame: &[u8]) {
192 let Some(connection) = self.inner.connections.get_connection(node) else {
193 return;
203 };
204 write_raw_blocking(&connection, frame);
205 }
206
207 fn backfill_member(&self, node: Atom, group: Atom, pid: u64) {
211 let update = PgUpdate::Join {
212 scope: self.scope(),
213 group,
214 pid,
215 };
216 if let Ok(frame) = encode_pg_update_frame(update, self.inner.local_node, &self.inner.atoms)
217 {
218 self.write_frame(node, &frame);
219 } else {
220 tracing::warn!("failed to encode cluster backfill frame");
221 }
222 }
223
224 fn lock_local(&self) -> std::sync::MutexGuard<'_, HashMap<Atom, Vec<u64>>> {
225 self.inner
226 .local
227 .lock()
228 .unwrap_or_else(std::sync::PoisonError::into_inner)
229 }
230}
231
232impl ClusterObserver for ClusterSync {
233 fn on_subscribe(&self, channel: &str, subscriber_pid: u64) {
234 let group = self.group(channel);
235 self.inner.pg.join(self.scope(), group, subscriber_pid);
237 self.record_local(group, subscriber_pid);
238 tracing::debug!(
239 channel = %channel,
240 pid = subscriber_pid,
241 "advertised local subscription to cluster"
242 );
243 }
244
245 fn on_unsubscribe(&self, channel: &str, subscriber_pid: u64) {
246 let group = self.group(channel);
247 self.inner.pg.leave(self.scope(), group, subscriber_pid);
248 self.forget_local(group, subscriber_pid);
249 tracing::debug!(
250 channel = %channel,
251 pid = subscriber_pid,
252 "withdrew local subscription from cluster"
253 );
254 }
255
256 fn on_publish(&self, channel: &str, envelope: &Envelope) {
257 if !crate::console::tappable(channel) {
260 return;
261 }
262 let targets = self.remote_targets(channel);
263 if targets.is_empty() {
264 return;
265 }
266 let frame_bytes = encode_envelope(envelope);
267 for member in targets {
268 self.send_to_member(member, &frame_bytes);
269 }
270 }
271}
272
273impl ClusterSync {
274 pub fn on_peer_join(&self, node: Atom) {
276 for (group, pid) in self.local_memberships() {
277 self.backfill_member(node, group, pid);
278 }
279 }
280
281 pub fn on_peer_leave(&self, node: Atom) {
284 let name = self
285 .inner
286 .atoms
287 .resolve(node)
288 .map_or_else(|| format!("<atom {node:?}>"), str::to_owned);
289 tracing::info!(
290 peer = %name,
291 "peer departed; its remote subscriptions were purged by beamr"
292 );
293 }
294}
295
296fn write_raw_blocking(
301 connection: &Arc<beamr::distribution::connection::DistConnection>,
302 frame: &[u8],
303) {
304 let connection = Arc::clone(connection);
305 let frame = frame.to_vec();
306 let write = async move {
307 if let Err(error) = connection.write_raw(&frame).await {
311 tracing::warn!(%error, "cross-node distribution write failed");
312 }
313 };
314 if let Ok(handle) = tokio::runtime::Handle::try_current() {
315 if matches!(
316 handle.runtime_flavor(),
317 tokio::runtime::RuntimeFlavor::MultiThread
318 ) {
319 tokio::task::block_in_place(|| handle.block_on(write));
320 return;
321 }
322 }
323 match tokio::runtime::Builder::new_current_thread()
324 .enable_all()
325 .build()
326 {
327 Ok(runtime) => runtime.block_on(write),
328 Err(error) => tracing::warn!(error = %error, "failed to build cluster send runtime"),
329 }
330}
331
332#[cfg(test)]
333#[allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)]
334mod tests {
335 use super::ClusterSync;
336 use crate::cluster::discovery::ClusterResolver;
337 use beamr::atom::AtomTable;
338 use beamr::distribution::connection::ConnectionManager;
339 use beamr::distribution::pg::PgRegistry;
340 use beamr::distribution::resolver::StaticResolver;
341 use liminal::channel::ClusterObserver;
342 use std::collections::HashMap;
343 use std::sync::Arc;
344
345 fn sync_fixture() -> (ClusterSync, Arc<PgRegistry>, Arc<AtomTable>) {
346 let atoms = Arc::new(AtomTable::with_common_atoms());
347 let pg = Arc::new(PgRegistry::new(&atoms));
348 let connections = ConnectionManager::new(
349 Arc::clone(&atoms),
350 Arc::new(StaticResolver::new(HashMap::new())),
351 "test-cookie",
352 "local@127.0.0.1",
353 1,
354 );
355 let local_node = atoms.intern("local@127.0.0.1");
356 let resolver = Arc::new(ClusterResolver::new());
357 let sync = ClusterSync::new(
358 Arc::clone(&pg),
359 Arc::clone(&atoms),
360 connections,
361 local_node,
362 resolver,
363 );
364 (sync, pg, atoms)
365 }
366
367 #[test]
368 fn subscribe_joins_the_channel_pg_group() {
369 let (sync, pg, atoms) = sync_fixture();
370 sync.on_subscribe("orders", 42);
371 let group = atoms.intern("orders");
372 assert_eq!(pg.local_members(pg.default_scope(), group), vec![42]);
373 }
374
375 #[test]
376 fn unsubscribe_leaves_the_channel_pg_group() {
377 let (sync, pg, atoms) = sync_fixture();
378 sync.on_subscribe("orders", 42);
379 sync.on_unsubscribe("orders", 42);
380 let group = atoms.intern("orders");
381 assert!(pg.local_members(pg.default_scope(), group).is_empty());
382 }
383
384 #[test]
385 fn local_memberships_track_subscriptions_for_backfill() {
386 let (sync, _pg, _atoms) = sync_fixture();
387 sync.on_subscribe("orders", 1);
388 sync.on_subscribe("orders", 2);
389 sync.on_subscribe("events", 3);
390 let mut memberships = sync.local_memberships();
391 memberships.sort_by_key(|(group, pid)| (*group, *pid));
392 assert_eq!(memberships.len(), 3);
393 sync.on_unsubscribe("events", 3);
395 let remaining = sync.local_memberships();
396 assert_eq!(remaining.len(), 2);
397 assert!(remaining.iter().all(|(_, pid)| *pid == 1 || *pid == 2));
398 }
399
400 #[test]
401 fn remote_targets_empty_without_remote_members() {
402 let (sync, _pg, _atoms) = sync_fixture();
403 sync.on_subscribe("orders", 1);
404 assert!(sync.remote_targets("orders").is_empty());
405 }
406
407 #[test]
414 fn distribution_write_failure_warns() {
415 const COOKIE: &str = "sync-write-cookie";
416 const PEER_NAME: &str = "sync-peer@127.0.0.1";
417 let runtime = tokio::runtime::Builder::new_multi_thread()
418 .worker_threads(2)
419 .enable_all()
420 .build()
421 .expect("build loopback test runtime");
422
423 let peer_atoms = Arc::new(AtomTable::with_common_atoms());
424 let peer = ConnectionManager::new(
425 peer_atoms,
426 Arc::new(StaticResolver::new(HashMap::new())),
427 COOKIE,
428 PEER_NAME,
429 1,
430 );
431 peer.set_runtime_handle(runtime.handle().clone());
432 let accept = runtime
435 .block_on(peer.listen("127.0.0.1:0".parse().expect("loopback addr")))
436 .expect("peer binds a loopback listener");
437
438 let mut routes = HashMap::new();
439 routes.insert(PEER_NAME.to_owned(), accept.local_addr());
440 let dialer_atoms = Arc::new(AtomTable::with_common_atoms());
441 let dialer = ConnectionManager::new(
442 Arc::clone(&dialer_atoms),
443 Arc::new(StaticResolver::new(routes)),
444 COOKIE,
445 "sync-dialer@127.0.0.1",
446 1,
447 );
448 dialer.set_runtime_handle(runtime.handle().clone());
449 runtime
450 .block_on(dialer.connect(PEER_NAME))
451 .expect("loopback handshake succeeds");
452
453 let peer_node = dialer_atoms.intern(PEER_NAME);
454 let connection = dialer
455 .get_connection(peer_node)
456 .expect("dialer holds the live connection");
457 connection.mark_down_write_timeout();
459
460 let log = crate::test_log::CapturedLog::default();
461 log.capture(|| super::write_raw_blocking(&connection, b"frame-bytes"));
462
463 let output = log.contents();
464 assert!(
465 output.contains("WARN") && output.contains("cross-node distribution write failed"),
466 "expected the discarded write error to be warned about, got: {output}"
467 );
468 }
469
470 #[test]
475 fn console_publishes_never_reach_the_fan_out() {
476 let (sync, pg, atoms) = sync_fixture();
477 let remote_node = atoms.intern("node-b@127.0.0.1");
478 for channel in ["orders", "console.topology"] {
479 let group = atoms.intern(channel);
480 pg.apply_remote_join(pg.default_scope(), group, remote_node, u64::MAX, 0);
481 assert_eq!(sync.remote_targets(channel).len(), 1);
482 }
483 let envelope = || {
484 liminal::envelope::Envelope::new(
485 b"{}".to_vec(),
486 None,
487 liminal::channel::SchemaId::new(),
488 liminal::envelope::PublisherId::default(),
489 )
490 };
491 let skipped = "remote member pid out of immediate range";
492
493 let console = crate::test_log::CapturedLog::default();
494 console.capture(|| sync.on_publish("console.topology", &envelope()));
495 assert!(
496 !console.contents().contains(skipped),
497 "a console record fanned out"
498 );
499
500 let ordinary = crate::test_log::CapturedLog::default();
501 ordinary.capture(|| sync.on_publish("orders", &envelope()));
502 assert!(ordinary.contents().contains(skipped), "positive control");
503 }
504
505 #[test]
506 fn remote_targets_reflect_applied_remote_joins() {
507 let (sync, pg, atoms) = sync_fixture();
508 let group = atoms.intern("orders");
509 let remote_node = atoms.intern("node-b@127.0.0.1");
510 pg.apply_remote_join(pg.default_scope(), group, remote_node, 99, 0);
511 let targets = sync.remote_targets("orders");
512 assert_eq!(targets.len(), 1);
513 assert_eq!(targets[0].node, remote_node);
514 assert_eq!(targets[0].pid_number, 99);
515 }
516}