use std::collections::{BTreeSet, VecDeque};
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, Mutex, PoisonError, Weak};
use std::thread::JoinHandle;
use beamr::atom::{Atom, AtomTable};
use beamr::distribution::connection::{AcceptHandle, ConnectionManager};
use beamr::distribution::connection_events::{ConnectionEvent, SubscriberId};
use beamr::scheduler::Scheduler;
use crate::ServerError;
use crate::cluster::discovery::{self, ClusterResolver};
use crate::cluster::sync::ClusterSync;
use crate::config::types::ClusterConfig;
const EFFECT_QUEUE_WARN_DEPTH: usize = 1024;
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct MembershipDelta {
pub joined: Vec<Atom>,
pub left: Vec<Atom>,
}
impl MembershipDelta {
#[must_use]
pub fn is_empty(&self) -> bool {
self.joined.is_empty() && self.left.is_empty()
}
}
#[derive(Clone)]
pub struct Membership {
inner: Arc<MembershipInner>,
}
#[derive(Default)]
struct EffectQueue {
pending: VecDeque<MembershipDelta>,
shutdown: bool,
}
struct MembershipInner {
connections: ConnectionManager,
atoms: Arc<AtomTable>,
peers: Mutex<BTreeSet<Atom>>,
queue: Mutex<EffectQueue>,
wake: Condvar,
subscription: Mutex<Option<SubscriberId>>,
events_observed: AtomicU64,
consumer_wakes: AtomicU64,
source_snapshots: AtomicU64,
queue_high_water: AtomicUsize,
depth_warned: AtomicBool,
}
impl MembershipInner {
fn observe(&self, event: ConnectionEvent) {
self.events_observed.fetch_add(1, Ordering::Relaxed);
let node = event.node();
let delta = {
let mut tracked = self.peers.lock().unwrap_or_else(PoisonError::into_inner);
match event {
ConnectionEvent::Up(_) if tracked.insert(node) => MembershipDelta {
joined: vec![node],
left: Vec::new(),
},
ConnectionEvent::Down(_) if tracked.remove(&node) => MembershipDelta {
joined: Vec::new(),
left: vec![node],
},
_ => return,
}
};
let depth = {
let mut queue = self.queue.lock().unwrap_or_else(PoisonError::into_inner);
queue.pending.push_back(delta);
queue.pending.len()
};
self.queue_high_water.fetch_max(depth, Ordering::Relaxed);
self.wake.notify_one();
}
}
impl Drop for MembershipInner {
fn drop(&mut self) {
let id = self
.subscription
.lock()
.unwrap_or_else(PoisonError::into_inner)
.take();
if let Some(id) = id {
self.connections.unsubscribe_connection_events(id);
}
}
}
impl std::fmt::Debug for Membership {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("Membership")
.field("peer_count", &self.peers().len())
.finish()
}
}
impl Membership {
#[must_use]
pub fn new(connections: ConnectionManager, atoms: Arc<AtomTable>) -> Self {
let inner = Arc::new(MembershipInner {
connections,
atoms,
peers: Mutex::new(BTreeSet::new()),
queue: Mutex::new(EffectQueue::default()),
wake: Condvar::new(),
subscription: Mutex::new(None),
events_observed: AtomicU64::new(0),
consumer_wakes: AtomicU64::new(0),
source_snapshots: AtomicU64::new(0),
queue_high_water: AtomicUsize::new(0),
depth_warned: AtomicBool::new(false),
});
let weak: Weak<MembershipInner> = Arc::downgrade(&inner);
inner.source_snapshots.fetch_add(1, Ordering::Relaxed);
let id = inner
.connections
.subscribe_connection_events_with_snapshot(move |event| {
if let Some(inner) = weak.upgrade() {
inner.observe(event);
}
});
*inner
.subscription
.lock()
.unwrap_or_else(PoisonError::into_inner) = Some(id);
Self { inner }
}
#[must_use]
pub fn peers(&self) -> Vec<Atom> {
self.lock_peers().iter().copied().collect()
}
#[must_use]
pub fn peer_names(&self) -> Vec<String> {
self.peers()
.into_iter()
.filter_map(|peer| self.inner.atoms.resolve(peer).map(str::to_owned))
.collect()
}
#[must_use]
pub fn events_observed(&self) -> u64 {
self.inner.events_observed.load(Ordering::Relaxed)
}
#[must_use]
pub fn consumer_wakes(&self) -> u64 {
self.inner.consumer_wakes.load(Ordering::Relaxed)
}
#[must_use]
pub fn source_snapshots(&self) -> u64 {
self.inner.source_snapshots.load(Ordering::Relaxed)
}
#[must_use]
pub fn queue_high_water(&self) -> usize {
self.inner.queue_high_water.load(Ordering::Relaxed)
}
#[must_use]
pub fn pending_effects(&self) -> usize {
self.lock_queue().pending.len()
}
fn wait_for_effects(&self) -> Option<Vec<MembershipDelta>> {
let mut queue = self
.inner
.wake
.wait_while(self.lock_queue(), |queue| {
queue.pending.is_empty() && !queue.shutdown
})
.unwrap_or_else(PoisonError::into_inner);
self.inner.consumer_wakes.fetch_add(1, Ordering::Relaxed);
if queue.pending.is_empty() {
return None;
}
Some(queue.pending.drain(..).collect())
}
fn take_pending(&self) -> Vec<MembershipDelta> {
self.inner.consumer_wakes.fetch_add(1, Ordering::Relaxed);
self.lock_queue().pending.drain(..).collect()
}
fn signal_shutdown(&self) {
self.lock_queue().shutdown = true;
self.inner.wake.notify_all();
}
fn unsubscribe(&self) {
let id = self
.inner
.subscription
.lock()
.unwrap_or_else(PoisonError::into_inner)
.take();
if let Some(id) = id {
self.inner.connections.unsubscribe_connection_events(id);
}
}
fn warn_once_on_depth(&self, depth: usize) {
if depth >= EFFECT_QUEUE_WARN_DEPTH
&& !self.inner.depth_warned.swap(true, Ordering::Relaxed)
{
tracing::warn!(
depth,
high_water = self.queue_high_water(),
"cluster membership effect queue is unusually deep; the queue is \
lossless so nothing was discarded, but the consumer is being starved"
);
}
}
fn name(&self, peer: Atom) -> String {
self.inner
.atoms
.resolve(peer)
.map_or_else(|| format!("<atom {peer:?}>"), str::to_owned)
}
fn lock_peers(&self) -> std::sync::MutexGuard<'_, BTreeSet<Atom>> {
self.inner
.peers
.lock()
.unwrap_or_else(PoisonError::into_inner)
}
fn lock_queue(&self) -> std::sync::MutexGuard<'_, EffectQueue> {
self.inner
.queue
.lock()
.unwrap_or_else(PoisonError::into_inner)
}
}
pub struct ClusterHandle {
accept: AcceptHandle,
consumer: Option<MembershipConsumer>,
membership: Membership,
_runtime: Arc<tokio::runtime::Runtime>,
}
impl std::fmt::Debug for ClusterHandle {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("ClusterHandle")
.field("listen_addr", &self.accept.local_addr())
.field("membership", &self.membership)
.finish_non_exhaustive()
}
}
impl ClusterHandle {
#[must_use]
pub fn listen_addr(&self) -> SocketAddr {
self.accept.local_addr()
}
#[must_use]
pub const fn membership(&self) -> &Membership {
&self.membership
}
pub fn shutdown(&mut self) {
if let Some(consumer) = self.consumer.take() {
consumer.stop();
}
self.membership.unsubscribe();
self.accept.shutdown();
}
}
impl Drop for ClusterHandle {
fn drop(&mut self) {
self.shutdown();
}
}
struct MembershipConsumer {
membership: Membership,
handle: Option<JoinHandle<()>>,
}
impl MembershipConsumer {
fn prime(membership: &Membership, sync: &ClusterSync) {
for delta in membership.take_pending() {
apply_delta(membership, sync, delta);
}
}
fn start(membership: Membership, sync: ClusterSync) -> Self {
let membership_for_thread = membership.clone();
let handle = std::thread::Builder::new()
.name("liminal-cluster-membership".to_owned())
.spawn(move || {
run_consumer(&membership_for_thread, &sync);
})
.ok();
Self { membership, handle }
}
fn stop(mut self) {
self.membership.signal_shutdown();
if let Some(handle) = self.handle.take() {
let _ = handle.join();
}
}
}
fn run_consumer(membership: &Membership, sync: &ClusterSync) {
while let Some(batch) = membership.wait_for_effects() {
membership.warn_once_on_depth(batch.len());
for delta in batch {
apply_delta(membership, sync, delta);
}
}
}
fn apply_delta(membership: &Membership, sync: &ClusterSync, delta: MembershipDelta) {
for peer in delta.joined {
let name = membership.name(peer);
tracing::info!(peer = %name, peers = ?membership.peer_names(), "cluster peer joined");
sync.on_peer_join(peer);
}
for peer in delta.left {
let name = membership.name(peer);
tracing::warn!(peer = %name, peers = ?membership.peer_names(), "cluster peer left");
sync.on_peer_leave(peer);
}
}
pub fn start(
scheduler: &Arc<Scheduler>,
resolver: Arc<ClusterResolver>,
config: &ClusterConfig,
install_observer: impl FnOnce(ClusterSync),
on_established: impl FnOnce(),
) -> Result<ClusterHandle, ServerError> {
let connections =
scheduler
.try_distribution_connections()
.ok_or_else(|| ServerError::ClusterJoin {
message: "scheduler was composed without a distribution service; \
cluster membership requires one"
.to_owned(),
})?;
let atoms = Arc::clone(scheduler.atom_table());
let pg = scheduler.pg_registry();
let local_node = atoms.intern(&config.node_name);
let labels = discovery::register_seed_labels(&resolver, &config.seed_nodes);
let runtime = Arc::new(
tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.enable_all()
.build()
.map_err(|error| ServerError::ClusterJoin {
message: format!("failed to build cluster runtime: {error}"),
})?,
);
connections.set_runtime_handle(runtime.handle().clone());
let accept = runtime
.block_on(scheduler.start_distribution_listener(config.listen_address))
.map_err(|error| ServerError::ClusterJoin {
message: format!(
"failed to bind cluster distribution listener on {}: {error}",
config.listen_address
),
})?;
let outcome = runtime.block_on(discovery::connect_seeds(
&connections,
&resolver,
&atoms,
&labels,
));
if !outcome.is_satisfied() {
return Err(ServerError::ClusterJoin {
message: format!(
"no configured seed node was reachable ({} attempted)",
outcome.attempted
),
});
}
let membership = Membership::new(connections.clone(), Arc::clone(&atoms));
let sync = ClusterSync::new(pg, Arc::clone(&atoms), connections, local_node, resolver);
install_observer(sync.clone());
MembershipConsumer::prime(&membership, &sync);
tracing::info!(
node_name = %config.node_name,
peers = ?membership.peer_names(),
"cluster membership established"
);
on_established();
let consumer = MembershipConsumer::start(membership.clone(), sync);
Ok(ClusterHandle {
accept,
consumer: Some(consumer),
membership,
_runtime: runtime,
})
}
#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)]
mod tests {
use super::{Membership, MembershipDelta};
use beamr::atom::AtomTable;
use beamr::distribution::connection::{AcceptHandle, ConnectionManager};
use beamr::distribution::connection_events::{ConnectionEvent, ConnectionGeneration};
use beamr::distribution::resolver::StaticResolver;
use std::collections::HashMap;
use std::sync::Arc;
const COOKIE: &str = "srv008-membership-cookie";
const DIALER_NAME: &str = "dialer@127.0.0.1";
const PEER_NAME: &str = "peer@127.0.0.1";
fn empty_manager(atoms: &Arc<AtomTable>) -> ConnectionManager {
ConnectionManager::new(
Arc::clone(atoms),
Arc::new(StaticResolver::new(HashMap::new())),
"test-cookie",
"local@127.0.0.1",
1,
)
}
struct LivePair {
runtime: tokio::runtime::Runtime,
dialer: ConnectionManager,
dialer_atoms: Arc<AtomTable>,
_peer: ConnectionManager,
_accept: AcceptHandle,
}
impl LivePair {
fn new() -> Self {
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.enable_all()
.build()
.expect("build loopback test runtime");
let peer_atoms = Arc::new(AtomTable::with_common_atoms());
let peer = ConnectionManager::new(
peer_atoms,
Arc::new(StaticResolver::new(HashMap::new())),
COOKIE,
PEER_NAME,
1,
);
peer.set_runtime_handle(runtime.handle().clone());
let accept = runtime
.block_on(peer.listen("127.0.0.1:0".parse().expect("loopback addr")))
.expect("peer binds a loopback listener");
let mut routes = HashMap::new();
routes.insert(PEER_NAME.to_owned(), accept.local_addr());
let dialer_atoms = Arc::new(AtomTable::with_common_atoms());
let dialer = ConnectionManager::new(
Arc::clone(&dialer_atoms),
Arc::new(StaticResolver::new(routes)),
COOKIE,
DIALER_NAME,
1,
);
dialer.set_runtime_handle(runtime.handle().clone());
Self {
runtime,
dialer,
dialer_atoms,
_peer: peer,
_accept: accept,
}
}
fn connect(&self) {
self.runtime
.block_on(self.dialer.connect(PEER_NAME))
.expect("loopback handshake succeeds");
}
fn disconnect(&self) {
let node = self.dialer_atoms.intern(PEER_NAME);
assert!(
self.dialer.disconnect_node(node),
"the live loopback link must be closable"
);
}
}
#[test]
fn delta_is_empty_by_default() {
assert!(MembershipDelta::default().is_empty());
}
#[test]
fn empty_initial_view_yields_no_peers() {
let atoms = Arc::new(AtomTable::with_common_atoms());
let membership = Membership::new(empty_manager(&atoms), Arc::clone(&atoms));
assert!(membership.peers().is_empty());
assert_eq!(
membership.events_observed(),
0,
"an empty table synthesizes no catch-up event"
);
assert_eq!(membership.pending_effects(), 0);
assert_eq!(
membership.source_snapshots(),
1,
"arming asks for the initial view exactly once"
);
}
#[test]
fn a_late_arm_sees_a_live_peer_in_its_initial_view() {
let pair = LivePair::new();
pair.connect();
let membership = Membership::new(pair.dialer.clone(), Arc::clone(&pair.dialer_atoms));
assert_eq!(
membership.peers().len(),
1,
"the initial view must contain the peer that was already live"
);
assert_eq!(membership.events_observed(), 1);
assert_eq!(
membership.pending_effects(),
1,
"the initial view's join effect is queued for the consumer, not run \
on beamr's delivery thread"
);
}
#[test]
fn join_leave_and_rejoin_are_observed_in_order() {
let pair = LivePair::new();
let membership = Membership::new(pair.dialer.clone(), Arc::clone(&pair.dialer_atoms));
let peer = pair.dialer_atoms.intern(PEER_NAME);
pair.connect();
assert_eq!(membership.peers(), vec![peer], "join is visible at once");
pair.disconnect();
assert!(
membership.peers().is_empty(),
"leave is visible at once, with no sampling in between"
);
pair.connect();
assert_eq!(membership.peers(), vec![peer], "rejoin is visible at once");
let effects = membership.take_pending();
assert_eq!(
effects,
vec![
MembershipDelta {
joined: vec![peer],
left: Vec::new(),
},
MembershipDelta {
joined: Vec::new(),
left: vec![peer],
},
MembershipDelta {
joined: vec![peer],
left: Vec::new(),
},
],
"the consumer receives the transitions in the order they happened"
);
assert_eq!(
membership.events_observed(),
3,
"exactly one event per transition — no duplicates to dedupe"
);
}
#[test]
fn queued_effects_are_lossless_and_report_their_depth() {
let atoms = Arc::new(AtomTable::with_common_atoms());
let membership = Membership::new(empty_manager(&atoms), Arc::clone(&atoms));
for index in 0_u64..64 {
let peer = atoms.intern(&format!("peer-{index}@127.0.0.1"));
membership.inner.observe(ConnectionEvent::up(
peer,
ConnectionGeneration::from_raw(index + 1),
1,
));
}
assert_eq!(
membership.pending_effects(),
64,
"no membership effect may be discarded"
);
assert_eq!(membership.queue_high_water(), 64);
assert_eq!(membership.peers().len(), 64);
let drained = membership.take_pending();
assert_eq!(
drained.len(),
64,
"every queued effect reaches the consumer"
);
assert_eq!(membership.pending_effects(), 0);
assert_eq!(
membership.queue_high_water(),
64,
"the high-water mark is a disclosure, not a counter that resets"
);
}
#[test]
fn a_redundant_event_queues_no_effect() {
let atoms = Arc::new(AtomTable::with_common_atoms());
let membership = Membership::new(empty_manager(&atoms), Arc::clone(&atoms));
let peer = atoms.intern(PEER_NAME);
membership.inner.observe(ConnectionEvent::up(
peer,
ConnectionGeneration::from_raw(1),
1,
));
membership.inner.observe(ConnectionEvent::up(
peer,
ConnectionGeneration::from_raw(2),
1,
));
assert_eq!(membership.peers(), vec![peer]);
assert_eq!(
membership.pending_effects(),
1,
"a repeated Up for a tracked peer must not queue a second join effect"
);
}
#[test]
fn shutdown_wakes_a_consumer_with_nothing_pending() {
let atoms = Arc::new(AtomTable::with_common_atoms());
let membership = Membership::new(empty_manager(&atoms), Arc::clone(&atoms));
assert_eq!(membership.pending_effects(), 0, "nothing is pending");
let waiter = membership.clone();
let joined = std::thread::spawn(move || waiter.wait_for_effects());
membership.signal_shutdown();
assert!(
joined.join().expect("the waiter thread joins").is_none(),
"shutdown must release a consumer that has no work"
);
}
#[test]
fn membership_source_has_no_retired_poll_family() {
const SOURCE: &str = include_str!("membership.rs");
let production = SOURCE.split("mod tests").next().unwrap_or(SOURCE);
for forbidden in [
"POLL_INTERVAL",
"poll_once",
"thread::sleep",
"PollLoop",
"run_poll_loop",
"connected_nodes",
] {
assert!(
!production.contains(forbidden),
"retired membership poll-family source `{forbidden}` reappeared"
);
}
}
#[test]
fn armed_membership_observes_a_join_without_sampling() {
let pair = LivePair::new();
let membership = Membership::new(pair.dialer.clone(), Arc::clone(&pair.dialer_atoms));
pair.connect();
assert_eq!(
membership.peers().len(),
1,
"an armed membership source must observe the join the instant beamr \
installs it, without any sampling of the connection table"
);
}
#[test]
fn peer_names_resolve_through_the_atom_table() {
let atoms = Arc::new(AtomTable::with_common_atoms());
let membership = Membership::new(empty_manager(&atoms), Arc::clone(&atoms));
assert!(membership.peer_names().is_empty());
}
}