use crate::node::NodeId;
use std::{
collections::{
BTreeSet,
btree_set::{Iter, Union},
},
iter::Peekable,
};
#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
pub struct ClusterConfig {
pub voters: BTreeSet<NodeId>,
pub new_voters: BTreeSet<NodeId>,
pub non_voters: BTreeSet<NodeId>,
}
impl ClusterConfig {
pub fn new() -> Self {
Self::default()
}
pub fn contains(&self, id: NodeId) -> bool {
self.voters.contains(&id) || self.new_voters.contains(&id) || self.non_voters.contains(&id)
}
pub fn is_joint_consensus(&self) -> bool {
!self.new_voters.is_empty()
}
pub fn unique_nodes(&self) -> impl '_ + Iterator<Item = NodeId> {
UniqueNodes {
voters: self.voters.union(&self.new_voters).peekable(),
non_voters: self.non_voters.iter().peekable(),
}
}
pub(crate) fn unique_voters(&self) -> impl '_ + Iterator<Item = NodeId> {
self.voters.union(&self.new_voters).copied()
}
pub(crate) fn is_voter(&self, id: NodeId) -> bool {
self.voters.contains(&id) || self.new_voters.contains(&id)
}
pub fn to_joint_consensus(&self, adding_voters: &[NodeId], removing_voters: &[NodeId]) -> Self {
let mut config = self.clone();
config.new_voters = config.voters.clone();
config.new_voters.extend(adding_voters.iter().copied());
config.new_voters.retain(|id| !removing_voters.contains(id));
config
}
pub(crate) fn voter_majority_count(&self) -> usize {
self.voters.len() / 2 + 1
}
pub(crate) fn new_voter_majority_count(&self) -> usize {
if self.new_voters.is_empty() {
0
} else {
self.new_voters.len() / 2 + 1
}
}
}
#[derive(Debug)]
pub struct UniqueNodes<'a> {
voters: Peekable<Union<'a, NodeId>>,
non_voters: Peekable<Iter<'a, NodeId>>,
}
impl Iterator for UniqueNodes<'_> {
type Item = NodeId;
fn next(&mut self) -> Option<Self::Item> {
let a = self.voters.peek().copied();
let b = self.non_voters.peek().copied();
match (a, b) {
(None, None) => None,
(Some(a), None) => {
self.voters.next();
Some(*a)
}
(None, Some(b)) => {
self.non_voters.next();
Some(*b)
}
(Some(a), Some(b)) if a == b => {
self.voters.next();
self.non_voters.next();
Some(*a)
}
(Some(a), Some(b)) if a < b => {
self.voters.next();
Some(*a)
}
(Some(_), Some(b)) => {
self.non_voters.next();
Some(*b)
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn unique_nodes() {
let mut config = ClusterConfig::new();
config.voters.insert(id(1));
config.voters.insert(id(2));
config.new_voters.insert(id(2));
config.new_voters.insert(id(3));
config.non_voters.insert(id(4));
config.non_voters.insert(id(5));
config.non_voters.insert(id(6));
let nodes: Vec<_> = config.unique_nodes().map(|n| n.get()).collect();
assert_eq!(nodes, vec![1, 2, 3, 4, 5, 6]);
}
fn id(n: u64) -> NodeId {
NodeId::new(n)
}
}