use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use parking_lot::Mutex;
use crate::cluster::peer::PeerState;
use crate::cluster::pool::ServerPool;
pub type Incarnation = u64;
pub type Tick = u64;
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum Status {
Alive,
Suspect {
since: Tick,
suspectors: u32,
},
Dead,
}
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct Member {
pub incarnation: Incarnation,
pub status: Status,
}
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum ProbeResult {
Acked,
IndirectAcked,
Failed,
}
#[derive(Copy, Clone, Debug)]
pub struct SwimConfig {
pub indirect_probes: u32,
pub suspicion_periods_base: Tick,
pub ns_dilation: Tick,
pub ns_max: u32,
pub refutation_enabled: bool,
}
impl Default for SwimConfig {
fn default() -> Self {
Self {
indirect_probes: 3,
suspicion_periods_base: 4,
ns_dilation: 1,
ns_max: 8,
refutation_enabled: true,
}
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct Update {
pub member: usize,
pub incarnation: Incarnation,
pub status: Status,
}
#[derive(Clone, Debug)]
pub struct SwimState {
me: usize,
my_incarnation: Incarnation,
members: BTreeMap<usize, Member>,
nack_score: u32,
cfg: SwimConfig,
}
impl SwimState {
#[must_use]
pub fn new(me: usize, n: usize, cfg: SwimConfig) -> Self {
assert!(me < n, "self index {me} must be < group size {n}");
let mut members = BTreeMap::new();
for i in 0..n {
if i != me {
members.insert(
i,
Member {
incarnation: 0,
status: Status::Alive,
},
);
}
}
Self {
me,
my_incarnation: 0,
members,
nack_score: 0,
cfg,
}
}
#[must_use]
pub fn me(&self) -> usize {
self.me
}
#[must_use]
pub fn my_incarnation(&self) -> Incarnation {
self.my_incarnation
}
#[must_use]
pub fn nack_score(&self) -> u32 {
self.nack_score
}
#[must_use]
pub fn status(&self, member: usize) -> Status {
if member == self.me {
return Status::Alive;
}
self.members
.get(&member)
.map_or(Status::Alive, |m| m.status)
}
#[must_use]
pub fn incarnation(&self, member: usize) -> Incarnation {
if member == self.me {
return self.my_incarnation;
}
self.members.get(&member).map_or(0, |m| m.incarnation)
}
#[must_use]
pub fn member_state(&self, member: usize) -> PeerState {
match self.status(member) {
Status::Alive => PeerState::Normal,
Status::Suspect { .. } | Status::Dead => PeerState::Down,
}
}
pub fn on_probe(
&mut self,
tick: Tick,
target: usize,
result: ProbeResult,
) -> Vec<(u32, PeerState)> {
if target == self.me {
return Vec::new();
}
match result {
ProbeResult::Acked => {
self.nack_score = self.nack_score.saturating_sub(1);
self.mark_alive_local(tick, target)
}
ProbeResult::IndirectAcked => {
self.nack_score = (self.nack_score + 1).min(self.cfg.ns_max);
self.mark_alive_local(tick, target)
}
ProbeResult::Failed => self.begin_suspicion(tick, target),
}
}
pub fn on_update(&mut self, tick: Tick, update: Update) -> Vec<(u32, PeerState)> {
if update.member == self.me {
self.handle_update_about_self(update);
return Vec::new();
}
let entry = self.members.entry(update.member).or_insert(Member {
incarnation: 0,
status: Status::Alive,
});
let prev_state = status_to_peer_state(entry.status);
let take = match update.incarnation.cmp(&entry.incarnation) {
std::cmp::Ordering::Greater => true,
std::cmp::Ordering::Less => false,
std::cmp::Ordering::Equal => {
status_rank(update.status) > status_rank(entry.status)
|| matches!(
(update.status, entry.status),
(Status::Suspect { .. }, Status::Suspect { .. })
)
}
};
if take {
let merged = merge_suspect(entry.status, update.status, tick);
entry.incarnation = update.incarnation;
entry.status = merged;
}
let new_state = status_to_peer_state(self.status(update.member));
transition(update.member, prev_state, new_state)
}
pub fn tick(&mut self, tick: Tick) -> Vec<(u32, PeerState)> {
let mut transitions = Vec::new();
let members: Vec<usize> = self.members.keys().copied().collect();
for m in members {
let Status::Suspect { since, suspectors } = self.members[&m].status else {
continue;
};
if tick >= self.confirm_deadline_for(since, suspectors) {
let prev = status_to_peer_state(self.members[&m].status);
self.members.get_mut(&m).unwrap().status = Status::Dead;
let now = status_to_peer_state(Status::Dead);
transitions.extend(transition(m, prev, now));
}
}
transitions
}
#[must_use]
pub fn confirm_deadline(&self, member: usize) -> Option<Tick> {
match self.members.get(&member)?.status {
Status::Suspect { since, suspectors } => {
Some(self.confirm_deadline_for(since, suspectors))
}
_ => None,
}
}
#[must_use]
pub fn snapshot(&self) -> Vec<(usize, PeerState)> {
self.members
.keys()
.map(|&m| (m, self.member_state(m)))
.collect()
}
fn confirm_deadline_for(&self, since: Tick, suspectors: u32) -> Tick {
let dilated = self
.cfg
.suspicion_periods_base
.saturating_mul(1 + u64::from(self.nack_score) * self.cfg.ns_dilation);
let shift = (suspectors.saturating_sub(1)).min(63);
let shrunk = (dilated >> shift).max(1);
since.saturating_add(shrunk)
}
fn mark_alive_local(&mut self, _tick: Tick, target: usize) -> Vec<(u32, PeerState)> {
let entry = self.members.entry(target).or_insert(Member {
incarnation: 0,
status: Status::Alive,
});
let prev = status_to_peer_state(entry.status);
if !matches!(entry.status, Status::Dead) {
entry.status = Status::Alive;
}
let now = status_to_peer_state(self.status(target));
transition(target, prev, now)
}
fn begin_suspicion(&mut self, tick: Tick, target: usize) -> Vec<(u32, PeerState)> {
let entry = self.members.entry(target).or_insert(Member {
incarnation: 0,
status: Status::Alive,
});
let prev = status_to_peer_state(entry.status);
match entry.status {
Status::Alive => {
entry.status = Status::Suspect {
since: tick,
suspectors: 1,
};
}
Status::Suspect { since, suspectors } => {
entry.status = Status::Suspect {
since,
suspectors: suspectors.saturating_add(1),
};
}
Status::Dead => {}
}
let now = status_to_peer_state(self.status(target));
transition(target, prev, now)
}
fn handle_update_about_self(&mut self, update: Update) {
let suspects_us = matches!(update.status, Status::Suspect { .. } | Status::Dead);
if suspects_us && self.cfg.refutation_enabled && update.incarnation >= self.my_incarnation {
self.my_incarnation = update.incarnation + 1;
}
}
}
fn status_rank(s: Status) -> u8 {
match s {
Status::Alive => 0,
Status::Suspect { .. } => 1,
Status::Dead => 2,
}
}
fn merge_suspect(existing: Status, incoming: Status, tick: Tick) -> Status {
match (existing, incoming) {
(
Status::Suspect {
since: s1,
suspectors: c1,
},
Status::Suspect { suspectors: c2, .. },
) => Status::Suspect {
since: s1,
suspectors: c1.saturating_add(c2),
},
(_, Status::Suspect { suspectors, .. }) => Status::Suspect {
since: tick,
suspectors: suspectors.max(1),
},
(_, other) => other,
}
}
fn status_to_peer_state(s: Status) -> PeerState {
match s {
Status::Alive => PeerState::Normal,
Status::Suspect { .. } | Status::Dead => PeerState::Down,
}
}
#[allow(
clippy::cast_possible_truncation,
reason = "member count is bounded by the configured peer array, well under u32::MAX"
)]
fn transition(member: usize, prev: PeerState, now: PeerState) -> Vec<(u32, PeerState)> {
if prev == now {
Vec::new()
} else {
vec![(member as u32, now)]
}
}
#[derive(Clone)]
pub struct SwimHandler {
pool: Arc<ServerPool>,
state: Arc<Mutex<SwimState>>,
period: Duration,
epoch: Instant,
}
impl SwimHandler {
#[must_use]
pub fn new(
pool: Arc<ServerPool>,
me: usize,
n: usize,
cfg: SwimConfig,
period: Duration,
) -> Self {
Self {
pool,
state: Arc::new(Mutex::new(SwimState::new(me, n, cfg))),
period: if period.is_zero() {
Duration::from_millis(1)
} else {
period
},
epoch: Instant::now(),
}
}
#[must_use]
pub fn pool(&self) -> &Arc<ServerPool> {
&self.pool
}
#[allow(
clippy::cast_possible_truncation,
reason = "period is >= 1ms so the tick count stays far below u64::MAX for any realistic uptime"
)]
fn tick_of(&self, now: Instant) -> Tick {
let elapsed = now.saturating_duration_since(self.epoch);
(elapsed.as_nanos() / self.period.as_nanos().max(1)) as Tick
}
pub fn on_probe(
&self,
now: Instant,
target: usize,
result: ProbeResult,
) -> Vec<(u32, PeerState)> {
let tick = self.tick_of(now);
let t = self.state.lock().on_probe(tick, target, result);
self.apply(&t);
t
}
pub fn on_update(&self, now: Instant, update: Update) -> Vec<(u32, PeerState)> {
let tick = self.tick_of(now);
let t = self.state.lock().on_update(tick, update);
self.apply(&t);
t
}
pub fn evaluate(&self, now: Instant) -> Vec<(u32, PeerState)> {
let tick = self.tick_of(now);
let t = self.state.lock().tick(tick);
self.apply(&t);
t
}
fn apply(&self, transitions: &[(u32, PeerState)]) {
if transitions.is_empty() {
return;
}
let mut peers = self.pool.peers().write();
for &(idx, state) in transitions {
if let Some(p) = peers.iter_mut().find(|p| p.idx() == idx && !p.is_local()) {
if p.state() != state {
p.set_state(state, now_secs_wall());
}
}
}
}
}
fn now_secs_wall() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |d| d.as_secs())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn dead_node_is_confirmed_after_timeout() {
let mut s = SwimState::new(0, 3, SwimConfig::default());
let t = s.on_probe(1, 2, ProbeResult::Failed);
assert_eq!(t, vec![(2, PeerState::Down)]);
assert!(matches!(s.status(2), Status::Suspect { .. }));
let deadline = s.confirm_deadline(2).unwrap();
let none = s.tick(deadline - 1);
assert!(none.is_empty());
assert!(matches!(s.status(2), Status::Suspect { .. }));
s.tick(deadline);
assert_eq!(s.status(2), Status::Dead);
}
#[test]
fn refutation_clears_a_false_suspicion() {
let mut s0 = SwimState::new(0, 2, SwimConfig::default());
s0.on_probe(1, 1, ProbeResult::Failed);
assert_eq!(s0.member_state(1), PeerState::Down);
let mut s1 = SwimState::new(1, 2, SwimConfig::default());
s1.on_update(
1,
Update {
member: 1,
incarnation: 0,
status: Status::Suspect {
since: 1,
suspectors: 1,
},
},
);
assert_eq!(s1.my_incarnation(), 1, "should bump to refute");
let back = s0.on_update(
2,
Update {
member: 1,
incarnation: s1.my_incarnation(),
status: Status::Alive,
},
);
assert_eq!(back, vec![(1, PeerState::Normal)]);
assert_eq!(s0.member_state(1), PeerState::Normal);
}
#[test]
fn indirect_ack_raises_nack_score_and_keeps_target_alive() {
let mut s = SwimState::new(0, 3, SwimConfig::default());
s.on_probe(1, 1, ProbeResult::IndirectAcked);
assert_eq!(s.nack_score(), 1);
assert_eq!(s.member_state(1), PeerState::Normal);
}
#[test]
fn nack_score_dilates_suspicion_timeout() {
let cfg = SwimConfig::default();
let mut healthy = SwimState::new(0, 2, cfg);
let mut slow = SwimState::new(0, 2, cfg);
for _ in 0..3 {
slow.on_probe(0, 1, ProbeResult::IndirectAcked);
}
assert!(slow.nack_score() > 0);
healthy.on_probe(1, 1, ProbeResult::Failed);
slow.on_probe(1, 1, ProbeResult::Failed);
let dh = healthy.confirm_deadline(1).unwrap();
let ds = slow.confirm_deadline(1).unwrap();
assert!(ds > dh, "slow observer must wait longer (ds={ds}, dh={dh})");
}
#[test]
fn dogpile_shortens_confirm_deadline() {
let cfg = SwimConfig::default();
let mut lone = SwimState::new(0, 3, cfg);
let mut piled = SwimState::new(0, 3, cfg);
lone.on_probe(1, 2, ProbeResult::Failed);
piled.on_probe(1, 2, ProbeResult::Failed);
for _ in 0..2 {
piled.on_update(
1,
Update {
member: 2,
incarnation: 0,
status: Status::Suspect {
since: 1,
suspectors: 1,
},
},
);
}
let dl = lone.confirm_deadline(2).unwrap();
let dp = piled.confirm_deadline(2).unwrap();
assert!(dp < dl, "dogpile must shorten confirm (dp={dp}, dl={dl})");
}
#[test]
fn higher_incarnation_wins_merge() {
let mut s = SwimState::new(0, 2, SwimConfig::default());
s.on_update(
1,
Update {
member: 1,
incarnation: 0,
status: Status::Suspect {
since: 1,
suspectors: 1,
},
},
);
assert_eq!(s.member_state(1), PeerState::Down);
s.on_update(
2,
Update {
member: 1,
incarnation: 1,
status: Status::Alive,
},
);
assert_eq!(s.member_state(1), PeerState::Normal);
s.on_update(
3,
Update {
member: 1,
incarnation: 0,
status: Status::Dead,
},
);
assert_eq!(s.member_state(1), PeerState::Normal);
}
#[test]
fn refutation_disabled_lets_false_death_stick() {
let cfg = SwimConfig {
refutation_enabled: false,
..SwimConfig::default()
};
let mut s1 = SwimState::new(1, 2, cfg);
s1.on_update(
1,
Update {
member: 1,
incarnation: 0,
status: Status::Suspect {
since: 1,
suspectors: 1,
},
},
);
assert_eq!(s1.my_incarnation(), 0, "refutation disabled: no bump");
}
#[test]
fn self_is_always_alive() {
let s = SwimState::new(1, 3, SwimConfig::default());
assert_eq!(s.member_state(1), PeerState::Normal);
assert_eq!(s.status(1), Status::Alive);
}
}