#![allow(dead_code)]
use std::collections::HashMap;
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ElectionState {
Idle,
Candidate,
Follower,
Leader,
}
impl ElectionState {
#[must_use]
pub fn is_candidate(self) -> bool {
self == ElectionState::Candidate
}
#[must_use]
pub fn is_leader(self) -> bool {
self == ElectionState::Leader
}
#[must_use]
pub fn label(self) -> &'static str {
match self {
ElectionState::Idle => "idle",
ElectionState::Candidate => "candidate",
ElectionState::Follower => "follower",
ElectionState::Leader => "leader",
}
}
}
#[derive(Debug, Clone)]
pub struct NodeVote {
pub voter_id: String,
pub candidate_id: String,
pub term: u64,
pub cast_at: Instant,
pub reason: Option<String>,
}
impl NodeVote {
#[must_use]
pub fn new(voter_id: impl Into<String>, candidate_id: impl Into<String>, term: u64) -> Self {
Self {
voter_id: voter_id.into(),
candidate_id: candidate_id.into(),
term,
cast_at: Instant::now(),
reason: None,
}
}
#[must_use]
pub fn is_valid(&self, expected_term: u64) -> bool {
self.term == expected_term
&& !self.voter_id.is_empty()
&& !self.candidate_id.is_empty()
&& self.voter_id != self.candidate_id
}
#[must_use]
pub fn age_ms(&self, now: Instant) -> u64 {
now.saturating_duration_since(self.cast_at).as_millis() as u64
}
}
#[derive(Debug)]
pub struct ElectionManager {
pub node_id: String,
pub term: u64,
pub state: ElectionState,
votes: HashMap<String, NodeVote>,
cluster_size: usize,
election_started_at: Option<Instant>,
election_timeout: Duration,
}
impl ElectionManager {
#[must_use]
pub fn new(
node_id: impl Into<String>,
cluster_size: usize,
election_timeout: Duration,
) -> Self {
Self {
node_id: node_id.into(),
term: 0,
state: ElectionState::Idle,
votes: HashMap::new(),
cluster_size,
election_started_at: None,
election_timeout,
}
}
pub fn start_election(&mut self) {
self.term += 1;
self.state = ElectionState::Candidate;
self.votes.clear();
self.election_started_at = Some(Instant::now());
}
pub fn record_vote(&mut self, vote: NodeVote) -> bool {
if !vote.is_valid(self.term) {
return false;
}
if self.votes.contains_key(&vote.voter_id) {
return false;
}
self.votes.insert(vote.voter_id.clone(), vote);
self.update_state();
true
}
#[must_use]
pub fn winner(&self) -> Option<&str> {
let mut tally: HashMap<&str, usize> = HashMap::new();
for vote in self.votes.values() {
*tally.entry(vote.candidate_id.as_str()).or_insert(0) += 1;
}
tally
.into_iter()
.max_by_key(|(_, count)| *count)
.map(|(id, _)| id)
}
#[must_use]
pub fn quorum(&self) -> usize {
self.cluster_size / 2 + 1
}
#[must_use]
pub fn vote_count(&self) -> usize {
self.votes.len()
}
#[must_use]
pub fn is_timed_out(&self, now: Instant) -> bool {
match self.election_started_at {
None => false,
Some(started) => now.saturating_duration_since(started) >= self.election_timeout,
}
}
pub fn become_follower(&mut self) {
self.state = ElectionState::Follower;
self.votes.clear();
self.election_started_at = None;
}
pub fn become_leader(&mut self) {
self.state = ElectionState::Leader;
}
fn update_state(&mut self) {
let my_id = self.node_id.clone();
let my_votes = self
.votes
.values()
.filter(|v| v.candidate_id == my_id)
.count();
if my_votes >= self.quorum() {
self.state = ElectionState::Leader;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
fn make_manager(node_id: &str, cluster_size: usize) -> ElectionManager {
ElectionManager::new(node_id, cluster_size, Duration::from_secs(5))
}
fn cast_vote(manager: &mut ElectionManager, voter: &str, candidate: &str) -> bool {
let vote = NodeVote::new(voter, candidate, manager.term);
manager.record_vote(vote)
}
#[test]
fn test_election_state_labels() {
assert_eq!(ElectionState::Idle.label(), "idle");
assert_eq!(ElectionState::Candidate.label(), "candidate");
assert_eq!(ElectionState::Follower.label(), "follower");
assert_eq!(ElectionState::Leader.label(), "leader");
}
#[test]
fn test_is_candidate() {
assert!(ElectionState::Candidate.is_candidate());
assert!(!ElectionState::Leader.is_candidate());
assert!(!ElectionState::Idle.is_candidate());
}
#[test]
fn test_is_leader() {
assert!(ElectionState::Leader.is_leader());
assert!(!ElectionState::Candidate.is_leader());
}
#[test]
fn test_node_vote_is_valid() {
let vote = NodeVote::new("voter1", "node2", 3);
assert!(vote.is_valid(3));
assert!(!vote.is_valid(2)); }
#[test]
fn test_node_vote_invalid_self_vote() {
let vote = NodeVote::new("node1", "node1", 1);
assert!(!vote.is_valid(1)); }
#[test]
fn test_node_vote_invalid_empty_ids() {
let vote = NodeVote::new("", "node2", 1);
assert!(!vote.is_valid(1));
}
#[test]
fn test_start_election_increments_term() {
let mut mgr = make_manager("n1", 5);
assert_eq!(mgr.term, 0);
mgr.start_election();
assert_eq!(mgr.term, 1);
assert!(mgr.state.is_candidate());
}
#[test]
fn test_start_election_clears_votes() {
let mut mgr = make_manager("n1", 3);
mgr.start_election();
cast_vote(&mut mgr, "n2", "n1");
mgr.start_election(); assert_eq!(mgr.vote_count(), 0);
}
#[test]
fn test_record_vote_accepted() {
let mut mgr = make_manager("n1", 3);
mgr.start_election();
assert!(cast_vote(&mut mgr, "n2", "n1"));
assert_eq!(mgr.vote_count(), 1);
}
#[test]
fn test_record_vote_duplicate_rejected() {
let mut mgr = make_manager("n1", 5);
mgr.start_election();
assert!(cast_vote(&mut mgr, "n2", "n1"));
assert!(!cast_vote(&mut mgr, "n2", "n1")); assert_eq!(mgr.vote_count(), 1);
}
#[test]
fn test_winner_after_majority() {
let mut mgr = make_manager("n1", 3);
mgr.start_election();
cast_vote(&mut mgr, "n2", "n1");
cast_vote(&mut mgr, "n3", "n1");
assert_eq!(mgr.winner(), Some("n1"));
assert!(mgr.state.is_leader());
}
#[test]
fn test_quorum_calculation() {
let mgr3 = make_manager("n1", 3);
assert_eq!(mgr3.quorum(), 2);
let mgr5 = make_manager("n1", 5);
assert_eq!(mgr5.quorum(), 3);
}
#[test]
fn test_become_follower() {
let mut mgr = make_manager("n1", 3);
mgr.start_election();
mgr.become_follower();
assert_eq!(mgr.state, ElectionState::Follower);
assert_eq!(mgr.vote_count(), 0);
}
#[test]
fn test_is_timed_out() {
let mut mgr = ElectionManager::new("n1", 3, Duration::from_millis(1));
mgr.start_election();
std::thread::sleep(Duration::from_millis(5));
assert!(mgr.is_timed_out(Instant::now()));
}
#[test]
fn test_not_timed_out_before_deadline() {
let mut mgr = ElectionManager::new("n1", 3, Duration::from_secs(60));
mgr.start_election();
assert!(!mgr.is_timed_out(Instant::now()));
}
}