Merge pull request #47 from butaneprotocol/ensure-healthy-leaders

Make raft elections more robust

Author: SupernaviX

src/main.rs

@@ -44,43 +44,49 @@ struct Node {
 
 impl Node {
     pub fn new(config: Arc<OracleConfig>) -> Result<Self> {
-        let (leader_tx, leader_rx) = watch::channel(RaftLeader::Unknown);
-        let (raft_client, raft_sink) = RaftClient::new();
+        let (leader_sink, leader_source) = watch::channel(RaftLeader::Unknown);
+        let (price_feed_sink, price_feed_source) = watch::channel(vec![]);
+        let (price_audit_sink, price_audit_source) = watch::channel(vec![]);
+        let (raft_client, raft_source) = RaftClient::new();
+
         let (health_server, health_sink) = HealthServer::new(
             config.frost_address.as_ref(),
             &config.network,
-            leader_rx.clone(),
+            leader_source.clone(),
             raft_client.clone(),
         );
 
         // Construct a peer-to-peer network that can connect to peers, and dispatch messages to the correct state machine
         let mut network = Network::new(&config.network, health_sink.clone());
 
-        let raft = Raft::new(&config, &mut network, leader_tx, raft_sink);
-
-        let (price_feed_tx, price_feed_rx) = watch::channel(vec![]);
-        let (price_audit_tx, price_audit_rx) = watch::channel(vec![]);
+        let raft = Raft::new(
+            &config,
+            &mut network,
+            leader_sink,
+            price_feed_source.clone(),
+            raft_source,
+        );
 
         let (signature_aggregator, payload_source) = if config.consensus {
             SignatureAggregator::consensus(
                 &config,
                 &mut network,
                 raft_client,
-                price_feed_rx,
-                leader_rx,
+                price_feed_source,
+                leader_source,
             )?
         } else {
-            SignatureAggregator::single(&config, raft_client, price_feed_rx, leader_rx)?
+            SignatureAggregator::single(&config, raft_client, price_feed_source, leader_source)?
         };
 
         let price_aggregator = PriceAggregator::new(
-            price_feed_tx,
-            price_audit_tx,
+            price_feed_sink,
+            price_audit_sink,
             payload_source.clone(),
             config.clone(),
         )?;
 
-        let api_server = APIServer::new(&config, payload_source.clone(), price_audit_rx);
+        let api_server = APIServer::new(&config, payload_source.clone(), price_audit_source);
 
         let publisher = Publisher::new(&config, payload_source)?;
 

src/price_aggregator/source_adapter.rs

@@ -71,12 +71,15 @@ impl SourceAdapter {
                 let info = &snapshot.info;
                 let as_of = snapshot.as_of;
                 let name = snapshot.name.as_ref();
-                let span = info_span!("update_price", token = info.token, unit = info.unit, name);
                 if info.value.is_zero() {
+                    let span =
+                        info_span!("update_price", token = info.token, unit = info.unit, name);
                     span.in_scope(|| warn!("ignoring reported value of 0"));
                     continue;
                 }
                 if info.reliability.is_zero() {
+                    let span =
+                        info_span!("update_price", token = info.token, unit = info.unit, name);
                     span.in_scope(|| warn!("ignoring reported value with reliability 0"));
                     continue;
                 }

src/raft.rs

@@ -13,12 +13,14 @@ use tracing::{debug, info, info_span};
 use crate::{
     config::OracleConfig,
     network::{Network, NetworkChannel, NodeId},
+    price_feed::PriceFeedEntry,
 };
 
 pub use client::RaftClient;
 use logic::RaftState;
 pub use logic::{RaftLeader, RaftMessage};
 
+#[derive(Debug)]
 pub enum RaftCommand {
     /// If we are currently leader, step down.
     /// Stop sending heartbeats, so that another node triggers election and wins.
@@ -29,6 +31,8 @@ pub enum RaftCommand {
 
 pub struct Raft {
     pub id: NodeId,
+    expected_payloads: usize,
+    price_feed_source: watch::Receiver<Vec<PriceFeedEntry>>,
     command_source: mpsc::Receiver<RaftCommand>,
     channel: NetworkChannel<RaftMessage>,
     state: RaftState,
@@ -39,12 +43,14 @@ impl Raft {
         config: &OracleConfig,
         network: &mut Network,
         leader_sink: watch::Sender<RaftLeader>,
+        price_feed_source: watch::Receiver<Vec<PriceFeedEntry>>,
         command_source: mpsc::Receiver<RaftCommand>,
     ) -> Self {
         // quorum is set to a majority of expected nodes (which includes ourself!)
         let quorum = ((config.network.peers.len() + 1) / 2) + 1;
         let heartbeat_freq = config.heartbeat;
         let timeout_freq = config.timeout;
+        let expected_payloads = config.synthetics.len();
 
         info!(
             quorum = quorum,
@@ -64,6 +70,8 @@ impl Raft {
         );
         Self {
             id,
+            expected_payloads,
+            price_feed_source,
             command_source,
             channel,
             state,
@@ -72,6 +80,7 @@ impl Raft {
 
     pub async fn handle_messages(self) {
         let (sender, mut receiver) = self.channel.split();
+        let mut price_feed_source = self.price_feed_source;
         let mut command_source = self.command_source;
         let mut state = self.state;
         loop {
@@ -85,14 +94,19 @@ impl Raft {
                     })
                 },
                 Some(command) = command_source.recv() => {
-                    let span = info_span!("raft_command");
+                    let span = info_span!("raft_command", ?command);
                     span.in_scope(|| {
                         match command {
                             RaftCommand::Abdicate => state.abdicate(),
                             RaftCommand::ForceElection(term) => state.run_election(term, Instant::now()),
                         }
                     })
                 }
+                Ok(()) = price_feed_source.changed() => {
+                    let present_payloads = price_feed_source.borrow().len();
+                    let missing_payloads = self.expected_payloads - present_payloads;
+                    state.set_missing_payloads(missing_payloads)
+                }
                 _ = sleep_until(next_event) => {
                     let span = info_span!("raft_tick");
                     span.in_scope(|| {

src/raft/logic.rs

@@ -1,4 +1,4 @@
-use std::{collections::BTreeSet, ops::Sub};
+use std::collections::BTreeSet;
 
 use minicbor::{Decode, Encode};
 use rand::Rng;
@@ -69,6 +69,7 @@ pub struct RaftState {
 
     pub status: RaftStatus,
     pub term: usize,
+    missing_payloads: usize,
 
     pub leader_sink: watch::Sender<RaftLeader>,
 }
@@ -96,6 +97,7 @@ impl RaftState {
                 voted_for: None,
             },
             term: 0,
+            missing_payloads: 0,
             leader_sink,
         };
         state.clear_status();
@@ -112,11 +114,11 @@ impl RaftState {
         let can_reach_quorum = (self.peers.len() + 1) >= self.quorum;
         if can_reach_quorum {
             // If we can reach quorum, this is when we'll hold our next election.
-            let election_time = self.last_event + self.timeout_freq.sub(self.jitter);
+            let election_time = self.last_event + self.time_until_next_election();
             next_event = next_event.min(election_time);
         }
 
-        if self.is_leader() {
+        if matches!(self.status, RaftStatus::Leader { abdicating: false }) {
             // If we're the leader, this is when we'll send our next heartbeat.
             let heartbeat_time = self.last_event + self.heartbeat_freq;
             next_event = next_event.min(heartbeat_time);
@@ -136,17 +138,18 @@ impl RaftState {
         }
     }
 
-    fn is_leader(&self) -> bool {
-        matches!(self.status, RaftStatus::Leader { .. })
-    }
-
     pub fn abdicate(&mut self) -> Vec<(NodeId, RaftMessage)> {
         if let RaftStatus::Leader { abdicating } = &mut self.status {
             *abdicating = true;
         }
         vec![]
     }
 
+    pub fn set_missing_payloads(&mut self, missing_payloads: usize) -> Vec<(NodeId, RaftMessage)> {
+        self.missing_payloads = missing_payloads;
+        vec![]
+    }
+
     pub fn receive(
         &mut self,
         timestamp: Instant,
@@ -322,7 +325,7 @@ impl RaftState {
     }
 
     pub fn tick(&mut self, timestamp: Instant) -> Vec<(NodeId, RaftMessage)> {
-        let actual_timeout = self.timeout_freq.sub(self.jitter);
+        let actual_timeout = self.time_until_next_election();
         let elapsed_time = timestamp.duration_since(self.last_event);
         let heartbeat_timeout = elapsed_time > self.heartbeat_freq;
         let election_timeout = elapsed_time > actual_timeout;
@@ -399,6 +402,12 @@ impl RaftState {
             .collect()
     }
 
+    fn time_until_next_election(&self) -> Duration {
+        // The more payloads this node fails to produce, the longer it waits to start an election.
+        // This prevents unhealthy nodes from becoming leader.
+        (self.timeout_freq * (self.missing_payloads + 1) as u32) - self.jitter
+    }
+
     fn clear_status(&mut self) {
         self.set_status(RaftStatus::Follower {
             leader: None,

src/raft/tests.rs

@@ -512,6 +512,33 @@ async fn should_start_new_election_if_current_election_times_out() {
     );
 }
 
+#[tokio::test]
+async fn should_delay_starting_new_election_if_missing_payloads() {
+    let start_time = Instant::now();
+    let heartbeat_freq = Duration::from_millis(1000);
+    let timeout_freq = Duration::from_millis(2000);
+
+    let (mut state, [other_id1, other_id2], _) =
+        generate_state(start_time, heartbeat_freq, timeout_freq);
+    assert_eq!(state.receive(&other_id1, RaftMessage::Connect), vec![]);
+    assert_eq!(state.receive(&other_id2, RaftMessage::Connect), vec![]);
+
+    state.become_follower(&other_id1, 1);
+
+    state.set_missing_payloads(2);
+
+    // we're missing two payloads, so we should wait two timeout cycles before running for leader.
+    assert_eq!(state.tick(timeout_freq), vec![]);
+    assert_eq!(state.tick(timeout_freq), vec![]);
+    assert_eq!(
+        state.tick(timeout_freq),
+        vec![
+            (other_id1.clone(), RaftMessage::RequestVote { term: 2 }),
+            (other_id2.clone(), RaftMessage::RequestVote { term: 2 }),
+        ]
+    );
+}
+
 #[tokio::test]
 async fn should_stop_sending_messages_if_we_abdicate() {
     let start_time = Instant::now();
@@ -672,6 +699,10 @@ impl Participant {
         self.state.abdicate();
     }
 
+    pub fn set_missing_payloads(&mut self, missing_payloads: usize) {
+        self.state.set_missing_payloads(missing_payloads);
+    }
+
     pub fn begin_election(&mut self) -> (Vec<NodeId>, usize) {
         // "run" for long enough to start another election
         let begin_election_messages = self.tick(self.timeout_freq);