diff --git a/src/raft/node/candidate.rs b/src/raft/node/candidate.rs index 82137af32..034e0028e 100644 --- a/src/raft/node/candidate.rs +++ b/src/raft/node/candidate.rs @@ -1,12 +1,12 @@ use super::super::{Address, Event, Message}; -use super::{rand_election_timeout, Follower, Leader, Node, NodeID, NodeState, Term, Ticks}; +use super::{rand_election_timeout, Follower, Leader, Node, NodeID, NodeState, Role, Term, Ticks}; use crate::error::{Error, Result}; use ::log::{debug, info, warn}; use std::collections::HashSet; /// A candidate is campaigning to become a leader. -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub struct Candidate { /// Votes received (including ourself). votes: HashSet, @@ -27,6 +27,8 @@ impl Candidate { } } +impl Role for Candidate {} + impl NodeState { /// Asserts internal invariants. fn assert(&mut self) -> Result<()> { diff --git a/src/raft/node/follower.rs b/src/raft/node/follower.rs index 3bf51764d..4c89f453f 100644 --- a/src/raft/node/follower.rs +++ b/src/raft/node/follower.rs @@ -1,12 +1,12 @@ use super::super::{Address, Event, Instruction, Message, RequestID, Response}; -use super::{rand_election_timeout, Candidate, Node, NodeID, NodeState, Term, Ticks}; +use super::{rand_election_timeout, Candidate, Node, NodeID, NodeState, Role, Term, Ticks}; use crate::error::{Error, Result}; use ::log::{debug, error, info, warn}; use std::collections::HashSet; // A follower replicates state from a leader. -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub struct Follower { /// The leader, or None if just initialized. leader: Option, @@ -34,6 +34,8 @@ impl Follower { } } +impl Role for Follower {} + impl NodeState { /// Asserts internal invariants. fn assert(&mut self) -> Result<()> { diff --git a/src/raft/node/leader.rs b/src/raft/node/leader.rs index 57ac6dbf5..8c6c01213 100644 --- a/src/raft/node/leader.rs +++ b/src/raft/node/leader.rs @@ -1,12 +1,12 @@ use super::super::{Address, Event, Index, Instruction, Message, Request, Response, Status}; -use super::{Follower, Node, NodeID, NodeState, Term, Ticks, HEARTBEAT_INTERVAL}; +use super::{Follower, Node, NodeID, NodeState, Role, Term, Ticks, HEARTBEAT_INTERVAL}; use crate::error::Result; use ::log::{debug, info}; use std::collections::{HashMap, HashSet}; /// Peer replication progress. -#[derive(Debug)] +#[derive(Debug, PartialEq)] struct Progress { /// The next index to replicate to the peer. next: Index, @@ -15,7 +15,7 @@ struct Progress { } // A leader serves requests and replicates the log to followers. -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub struct Leader { /// Peer replication progress. progress: HashMap, @@ -32,6 +32,8 @@ impl Leader { } } +impl Role for Leader {} + impl NodeState { /// Asserts internal invariants. fn assert(&mut self) -> Result<()> { diff --git a/src/raft/node/mod.rs b/src/raft/node/mod.rs index 1cf821460..bcc5cd455 100644 --- a/src/raft/node/mod.rs +++ b/src/raft/node/mod.rs @@ -133,8 +133,11 @@ impl From> for Node { } } +/// A Raft role: leader, follower, or candidate. +pub trait Role: std::fmt::Debug + PartialEq {} + // A Raft node with role R -pub struct NodeState { +pub struct NodeState { id: NodeID, peers: HashSet, term: Term, @@ -144,9 +147,9 @@ pub struct NodeState { role: R, } -impl NodeState { +impl NodeState { /// Transforms the node into another role. - fn become_role(self, role: T) -> NodeState { + fn become_role(self, role: T) -> NodeState { NodeState { id: self.id, peers: self.peers, @@ -369,17 +372,17 @@ mod tests { NodeAsserter::new(node) } - fn setup_rolenode() -> Result<(NodeState<()>, mpsc::UnboundedReceiver)> { + fn setup_rolenode() -> Result<(NodeState, mpsc::UnboundedReceiver)> { setup_rolenode_peers(vec![2, 3]) } fn setup_rolenode_peers( peers: Vec, - ) -> Result<(NodeState<()>, mpsc::UnboundedReceiver)> { + ) -> Result<(NodeState, mpsc::UnboundedReceiver)> { let (node_tx, node_rx) = mpsc::unbounded_channel(); let (state_tx, _) = mpsc::unbounded_channel(); let node = NodeState { - role: (), + role: Follower::new(None, None), id: 1, peers: HashSet::from_iter(peers), term: 1, @@ -488,11 +491,11 @@ mod tests { #[test] fn become_role() -> Result<()> { let (node, _) = setup_rolenode()?; - let new = node.become_role("role"); + let new = node.become_role(Candidate::new()); assert_eq!(new.id, 1); assert_eq!(new.term, 1); assert_eq!(new.peers, HashSet::from([2, 3])); - assert_eq!(new.role, "role"); + assert_eq!(new.role, Candidate::new()); Ok(()) }