Add Raft Role trait

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<NodeID>,
@@ -27,6 +27,8 @@ impl Candidate {
     }
 }
 
+impl Role for Candidate {}
+
 impl NodeState<Candidate> {
     /// Asserts internal invariants.
     fn assert(&mut self) -> Result<()> {

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<NodeID>,
@@ -34,6 +34,8 @@ impl Follower {
     }
 }
 
+impl Role for Follower {}
+
 impl NodeState<Follower> {
     /// Asserts internal invariants.
     fn assert(&mut self) -> Result<()> {

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<NodeID, Progress>,
@@ -32,6 +32,8 @@ impl Leader {
     }
 }
 
+impl Role for Leader {}
+
 impl NodeState<Leader> {
     /// Asserts internal invariants.
     fn assert(&mut self) -> Result<()> {

src/raft/node/mod.rs

@@ -133,8 +133,11 @@ impl From<NodeState<Leader>> 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<R> {
+pub struct NodeState<R: Role> {
     id: NodeID,
     peers: HashSet<NodeID>,
     term: Term,
@@ -144,9 +147,9 @@ pub struct NodeState<R> {
     role: R,
 }
 
-impl<R> NodeState<R> {
+impl<R: Role> NodeState<R> {
     /// Transforms the node into another role.
-    fn become_role<T>(self, role: T) -> NodeState<T> {
+    fn become_role<T: Role>(self, role: T) -> NodeState<T> {
         NodeState {
             id: self.id,
             peers: self.peers,
@@ -369,17 +372,17 @@ mod tests {
         NodeAsserter::new(node)
     }
 
-    fn setup_rolenode() -> Result<(NodeState<()>, mpsc::UnboundedReceiver<Message>)> {
+    fn setup_rolenode() -> Result<(NodeState<Follower>, mpsc::UnboundedReceiver<Message>)> {
         setup_rolenode_peers(vec![2, 3])
     }
 
     fn setup_rolenode_peers(
         peers: Vec<NodeID>,
-    ) -> Result<(NodeState<()>, mpsc::UnboundedReceiver<Message>)> {
+    ) -> Result<(NodeState<Follower>, mpsc::UnboundedReceiver<Message>)> {
         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(())
     }