use std::cell::RefCell;
use std::collections::{BTreeMap, BTreeSet, VecDeque};
use crate::Context;
use crate::cell::CellHandle;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PeerState {
Alive,
Suspect,
Dead,
Left,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PeerChangeEvent<P> {
Joined(P),
Left(P),
StateChanged {
peer: P,
from: PeerState,
to: PeerState,
},
}
#[derive(Debug, Clone, Copy)]
pub struct MembershipConfig {
pub phi_threshold: f64,
pub suspect_timeout: u64,
pub max_samples: usize,
pub min_std: f64,
}
impl Default for MembershipConfig {
fn default() -> Self {
Self {
phi_threshold: 8.0,
suspect_timeout: 5,
max_samples: 100,
min_std: 0.1,
}
}
}
#[derive(Debug, Clone)]
pub struct PhiAccrual {
window: VecDeque<f64>,
max_samples: usize,
min_std: f64,
last_heartbeat: Option<u64>,
}
impl PhiAccrual {
pub fn new(max_samples: usize, min_std: f64) -> Self {
Self {
window: VecDeque::new(),
max_samples: max_samples.max(1),
min_std,
last_heartbeat: None,
}
}
pub fn heartbeat(&mut self, now: u64) {
if let Some(last) = self.last_heartbeat {
let interval = now.saturating_sub(last) as f64;
self.window.push_back(interval);
while self.window.len() > self.max_samples {
self.window.pop_front();
}
}
self.last_heartbeat = Some(now);
}
fn mean(&self) -> f64 {
let n = self.window.len() as f64;
self.window.iter().sum::<f64>() / n
}
fn std(&self, mean: f64) -> f64 {
let n = self.window.len() as f64;
let var = self
.window
.iter()
.map(|x| (x - mean) * (x - mean))
.sum::<f64>()
/ n;
var.sqrt().max(self.min_std)
}
pub fn phi(&self, now: u64) -> f64 {
let Some(last) = self.last_heartbeat else {
return 0.0;
};
if self.window.is_empty() {
return 0.0;
}
let elapsed = now.saturating_sub(last) as f64;
let mean = self.mean();
let std = self.std(mean);
let y = (elapsed - mean) / std;
let e = (-y * (1.5976 + 0.070566 * y * y)).exp();
if elapsed > mean {
-(e / (1.0 + e)).log10()
} else {
-(1.0 - 1.0 / (1.0 + e)).log10()
}
}
}
struct PeerRecord {
state: PeerState,
detector: PhiAccrual,
suspect_since: Option<u64>,
}
pub struct MembershipCore<P> {
config: MembershipConfig,
peers: BTreeMap<P, PeerRecord>,
}
impl<P: Ord + Clone> MembershipCore<P> {
pub fn new(config: MembershipConfig) -> Self {
Self {
config,
peers: BTreeMap::new(),
}
}
fn new_detector(&self) -> PhiAccrual {
PhiAccrual::new(self.config.max_samples, self.config.min_std)
}
pub fn alive_set(&self) -> BTreeSet<P> {
self.peers
.iter()
.filter(|(_, r)| r.state == PeerState::Alive)
.map(|(p, _)| p.clone())
.collect()
}
pub fn state(&self, peer: &P) -> Option<PeerState> {
self.peers.get(peer).map(|r| r.state)
}
pub fn join(&mut self, peer: P, now: u64) -> Vec<PeerChangeEvent<P>> {
let mut detector = self.new_detector();
detector.heartbeat(now);
let known = self.peers.contains_key(&peer);
let prev = self.peers.get(&peer).map(|r| r.state);
self.peers.insert(
peer.clone(),
PeerRecord {
state: PeerState::Alive,
detector,
suspect_since: None,
},
);
match (known, prev) {
(false, _) => vec![PeerChangeEvent::Joined(peer)],
(true, Some(PeerState::Alive)) => vec![],
(true, Some(from)) => vec![PeerChangeEvent::StateChanged {
peer,
from,
to: PeerState::Alive,
}],
(true, None) => vec![],
}
}
pub fn heartbeat(&mut self, peer: P, now: u64) -> Vec<PeerChangeEvent<P>> {
let Some(record) = self.peers.get_mut(&peer) else {
return self.join(peer, now);
};
record.detector.heartbeat(now);
let from = record.state;
if from != PeerState::Alive && from != PeerState::Left {
record.state = PeerState::Alive;
record.suspect_since = None;
return vec![PeerChangeEvent::StateChanged {
peer,
from,
to: PeerState::Alive,
}];
}
vec![]
}
pub fn leave(&mut self, peer: P, _now: u64) -> Vec<PeerChangeEvent<P>> {
let Some(record) = self.peers.get_mut(&peer) else {
return vec![];
};
if record.state == PeerState::Left {
return vec![];
}
record.state = PeerState::Left;
record.suspect_since = None;
vec![PeerChangeEvent::Left(peer)]
}
pub fn tick(&mut self, now: u64) -> Vec<PeerChangeEvent<P>> {
let threshold = self.config.phi_threshold;
let timeout = self.config.suspect_timeout;
let mut events = Vec::new();
for (peer, record) in self.peers.iter_mut() {
match record.state {
PeerState::Alive => {
if record.detector.phi(now) > threshold {
record.state = PeerState::Suspect;
record.suspect_since = Some(now);
events.push(PeerChangeEvent::StateChanged {
peer: peer.clone(),
from: PeerState::Alive,
to: PeerState::Suspect,
});
}
}
PeerState::Suspect => {
let expired = record
.suspect_since
.is_some_and(|since| now.saturating_sub(since) >= timeout);
if expired {
record.state = PeerState::Dead;
events.push(PeerChangeEvent::StateChanged {
peer: peer.clone(),
from: PeerState::Suspect,
to: PeerState::Dead,
});
}
}
PeerState::Dead | PeerState::Left => {}
}
}
events
}
}
pub struct MembershipCell<P> {
core: RefCell<MembershipCore<P>>,
peer_set: CellHandle<BTreeSet<P>>,
}
impl<P: Ord + Clone + 'static> MembershipCell<P> {
pub fn new(ctx: &Context, config: MembershipConfig) -> Self {
Self {
core: RefCell::new(MembershipCore::new(config)),
peer_set: ctx.cell(BTreeSet::new()),
}
}
fn refresh(&self, ctx: &Context) {
let set = self.core.borrow().alive_set();
self.peer_set.set(ctx, set);
}
pub fn join(&self, ctx: &Context, peer: P, now: u64) -> Vec<PeerChangeEvent<P>> {
let events = self.core.borrow_mut().join(peer, now);
self.refresh(ctx);
events
}
pub fn heartbeat(&self, ctx: &Context, peer: P, now: u64) -> Vec<PeerChangeEvent<P>> {
let events = self.core.borrow_mut().heartbeat(peer, now);
self.refresh(ctx);
events
}
pub fn leave(&self, ctx: &Context, peer: P, now: u64) -> Vec<PeerChangeEvent<P>> {
let events = self.core.borrow_mut().leave(peer, now);
self.refresh(ctx);
events
}
pub fn tick(&self, ctx: &Context, now: u64) -> Vec<PeerChangeEvent<P>> {
let events = self.core.borrow_mut().tick(now);
self.refresh(ctx);
events
}
pub fn peer_set(&self, ctx: &Context) -> BTreeSet<P> {
self.peer_set.get(ctx)
}
pub fn peer_set_cell(&self) -> CellHandle<BTreeSet<P>> {
self.peer_set
}
pub fn state(&self, peer: &P) -> Option<PeerState> {
self.core.borrow().state(peer)
}
}
pub type PeerSet<P> = CellHandle<BTreeSet<P>>;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn phi_low_when_current_high_after_gap() {
let mut d = PhiAccrual::new(100, 0.1);
d.heartbeat(0);
d.heartbeat(1);
d.heartbeat(2);
d.heartbeat(3);
assert!(d.phi(3) < 8.0, "phi at last heartbeat should be low");
assert!(d.phi(100) > 8.0, "phi after a long gap should be high");
}
#[test]
fn lifecycle_transitions() {
let mut m = MembershipCore::<u64>::new(MembershipConfig::default());
assert_eq!(m.join(1, 0), vec![PeerChangeEvent::Joined(1)]);
m.heartbeat(1, 1);
m.heartbeat(1, 2);
m.heartbeat(1, 3);
assert_eq!(m.tick(3), vec![]);
assert_eq!(m.state(&1), Some(PeerState::Alive));
assert_eq!(
m.tick(100),
vec![PeerChangeEvent::StateChanged {
peer: 1,
from: PeerState::Alive,
to: PeerState::Suspect
}]
);
assert_eq!(
m.tick(106),
vec![PeerChangeEvent::StateChanged {
peer: 1,
from: PeerState::Suspect,
to: PeerState::Dead
}]
);
assert!(m.alive_set().is_empty());
}
#[test]
fn heartbeat_refutes_suspicion() {
let mut m = MembershipCore::<u64>::new(MembershipConfig::default());
m.join(1, 0);
m.heartbeat(1, 1);
m.heartbeat(1, 2);
m.tick(100); assert_eq!(m.state(&1), Some(PeerState::Suspect));
let ev = m.heartbeat(1, 101); assert_eq!(m.state(&1), Some(PeerState::Alive));
assert_eq!(
ev,
vec![PeerChangeEvent::StateChanged {
peer: 1,
from: PeerState::Suspect,
to: PeerState::Alive
}]
);
}
}