apollo_consensus: move vote storage to SM; dedupe in SHC

Author: asmaastarkware

crates/apollo_consensus/src/single_height_consensus.rs

@@ -137,8 +137,6 @@ pub(crate) struct SingleHeightConsensus {
     timeouts: TimeoutsConfig,
     state_machine: StateMachine,
     proposals: HashMap<Round, Option<ProposalCommitment>>,
-    prevotes: HashMap<(Round, ValidatorId), Vote>,
-    precommits: HashMap<(Round, ValidatorId), Vote>,
     last_prevote: Option<Vote>,
     last_precommit: Option<Vote>,
 }
@@ -161,8 +159,6 @@ impl SingleHeightConsensus {
             timeouts,
             state_machine,
             proposals: HashMap::new(),
-            prevotes: HashMap::new(),
-            precommits: HashMap::new(),
             last_prevote: None,
             last_precommit: None,
         }
@@ -349,29 +345,27 @@ impl SingleHeightConsensus {
             return Ok(ShcReturn::Tasks(Vec::new()));
         }
 
-        let (votes, sm_vote) = match vote.vote_type {
-            VoteType::Prevote => (&mut self.prevotes, StateMachineEvent::Prevote(vote.clone())),
+        // Check duplicates/conflicts from SM stored votes.
+        let (votes_map, sm_vote) = match vote.vote_type {
+            VoteType::Prevote => {
+                (self.state_machine.prevotes_ref(), StateMachineEvent::Prevote(vote.clone()))
+            }
             VoteType::Precommit => {
-                (&mut self.precommits, StateMachineEvent::Precommit(vote.clone()))
+                (self.state_machine.precommits_ref(), StateMachineEvent::Precommit(vote.clone()))
             }
         };
-
-        match votes.entry((vote.round, vote.voter)) {
-            Entry::Vacant(entry) => {
-                entry.insert(vote.clone());
-            }
-            Entry::Occupied(entry) => {
-                let old = entry.get();
-                if old.proposal_commitment != vote.proposal_commitment {
-                    warn!("Conflicting votes: old={:?}, new={:?}", old, vote);
-                    CONSENSUS_CONFLICTING_VOTES.increment(1);
-                    return Ok(ShcReturn::Tasks(Vec::new()));
-                } else {
-                    // Replay, ignore.
-                    return Ok(ShcReturn::Tasks(Vec::new()));
-                }
+        if let Some((old_vote, _)) = votes_map.get(&(vote.round, vote.voter)) {
+            if old_vote.proposal_commitment == vote.proposal_commitment {
+                // Duplicate - ignore.
+                return Ok(ShcReturn::Tasks(Vec::new()));
+            } else {
+                // Conflict - ignore and record.
+                warn!("Conflicting votes: old={old_vote:?}, new={vote:?}");
+                CONSENSUS_CONFLICTING_VOTES.increment(1);
+                return Ok(ShcReturn::Tasks(Vec::new()));
             }
         }
+
         info!("Accepting {:?}", vote);
         let height = self.state_machine.height();
         let leader_fn = |round: Round| -> ValidatorId { context.proposer(height, round) };
@@ -502,29 +496,24 @@ impl SingleHeightConsensus {
         context: &mut ContextT,
         vote: Vote,
     ) -> Result<Vec<ShcTask>, ConsensusError> {
-        let prevote_timeout = self.timeouts.get_prevote_timeout(vote.round);
-        let precommit_timeout = self.timeouts.get_precommit_timeout(vote.round);
-        let (votes, last_vote, task) = match vote.vote_type {
+        let (last_vote, task) = match vote.vote_type {
             VoteType::Prevote => (
-                &mut self.prevotes,
                 &mut self.last_prevote,
-                ShcTask::Prevote(prevote_timeout, StateMachineEvent::Prevote(vote.clone())),
+                ShcTask::Prevote(
+                    self.timeouts.get_prevote_timeout(vote.round),
+                    StateMachineEvent::Prevote(vote.clone()),
+                ),
             ),
             VoteType::Precommit => (
-                &mut self.precommits,
                 &mut self.last_precommit,
-                ShcTask::Precommit(precommit_timeout, StateMachineEvent::Precommit(vote.clone())),
+                ShcTask::Precommit(
+                    self.timeouts.get_precommit_timeout(vote.round),
+                    StateMachineEvent::Precommit(vote.clone()),
+                ),
             ),
         };
         // Ensure the voter matches this node.
         assert_eq!(vote.voter, self.state_machine.validator_id());
-        if let Some(old) =
-            votes.insert((vote.round, self.state_machine.validator_id()), vote.clone())
-        {
-            return Err(ConsensusError::InternalInconsistency(format!(
-                "State machine should not send repeat votes: old={old:?}, new={vote:?}"
-            )));
-        }
         *last_vote = match last_vote {
             None => Some(vote.clone()),
             Some(last_vote) if vote.round > last_vote.round => Some(vote.clone()),
@@ -569,18 +558,7 @@ impl SingleHeightConsensus {
                  {block}"
             )));
         }
-        let supporting_precommits: Vec<Vote> = self
-            .validators
-            .iter()
-            .filter_map(|v| {
-                let vote = self.precommits.get(&(round, *v))?;
-                if vote.proposal_commitment == Some(proposal_id) {
-                    Some(vote.clone())
-                } else {
-                    None
-                }
-            })
-            .collect();
+        let supporting_precommits = self.precommit_votes_for_value(round, Some(proposal_id));
 
         // TODO(matan): Check actual weights.
         let vote_weight = u64::try_from(supporting_precommits.len())
@@ -596,4 +574,18 @@ impl SingleHeightConsensus {
         }
         Ok(ShcReturn::Decision(Decision { precommits: supporting_precommits, block }))
     }
+
+    fn precommit_votes_for_value(
+        &self,
+        round: Round,
+        value: Option<ProposalCommitment>,
+    ) -> Vec<Vote> {
+        self.state_machine
+            .precommits_ref()
+            .iter()
+            .filter_map(|(&(r, _voter), (v, _w))| {
+                if r == round && v.proposal_commitment == value { Some(v.clone()) } else { None }
+            })
+            .collect()
+    }
 }

crates/apollo_consensus/src/state_machine.rs

@@ -78,9 +78,9 @@ pub(crate) struct StateMachine {
     is_observer: bool,
     // {round: (proposal_id, valid_round)}
     proposals: HashMap<Round, (Option<ProposalCommitment>, Option<Round>)>,
-    // {round: {proposal_id: vote_count}
-    prevotes: HashMap<Round, HashMap<Option<ProposalCommitment>, u32>>,
-    precommits: HashMap<Round, HashMap<Option<ProposalCommitment>, u32>>,
+    // {(round, voter): (vote, weight)}
+    prevotes: HashMap<(Round, ValidatorId), (Vote, u32)>,
+    precommits: HashMap<(Round, ValidatorId), (Vote, u32)>,
     // When true, the state machine will wait for a GetProposal event, buffering all other input
     // events in `events_queue`.
     awaiting_get_proposal: bool,
@@ -145,6 +145,14 @@ impl StateMachine {
         self.height
     }
 
+    pub(crate) fn prevotes_ref(&self) -> &HashMap<(Round, ValidatorId), (Vote, u32)> {
+        &self.prevotes
+    }
+
+    pub(crate) fn precommits_ref(&self) -> &HashMap<(Round, ValidatorId), (Vote, u32)> {
+        &self.precommits
+    }
+
     fn make_self_vote(
         &self,
         vote_type: VoteType,
@@ -328,15 +336,9 @@ impl StateMachine {
     where
         LeaderFn: Fn(Round) -> ValidatorId,
     {
-        let prevote_count = self
-            .prevotes
-            .entry(vote.round)
-            .or_default()
-            .entry(vote.proposal_commitment)
-            .or_insert(0);
-        // TODO(matan): Use variable weight.
-        *prevote_count += 1;
-        self.map_round_to_upons(vote.round, leader_fn)
+        let round = vote.round;
+        self.prevotes.entry((round, vote.voter)).or_insert((vote, 1));
+        self.map_round_to_upons(round, leader_fn)
     }
 
     pub(crate) fn handle_timeout_prevote(&mut self, round: u32) -> VecDeque<StateMachineEvent> {
@@ -360,15 +362,9 @@ impl StateMachine {
     where
         LeaderFn: Fn(Round) -> ValidatorId,
     {
-        let precommit_count = self
-            .precommits
-            .entry(vote.round)
-            .or_default()
-            .entry(vote.proposal_commitment)
-            .or_insert(0);
-        // TODO(matan): Use variable weight.
-        *precommit_count += 1;
-        self.map_round_to_upons(vote.round, leader_fn)
+        let round = vote.round;
+        self.precommits.entry((round, vote.voter)).or_insert((vote, 1));
+        self.map_round_to_upons(round, leader_fn)
     }
 
     fn handle_timeout_precommit<LeaderFn>(
@@ -640,24 +636,36 @@ impl StateMachine {
 
     fn round_has_enough_votes(
         &self,
-        votes: &HashMap<u32, HashMap<Option<ProposalCommitment>, u32>>,
+        votes: &HashMap<(Round, ValidatorId), (Vote, u32)>,
         round: u32,
         threshold: &VotesThreshold,
     ) -> bool {
-        threshold
-            .is_met(votes.get(&round).map_or(0, |v| v.values().sum()).into(), self.total_weight)
+        let weight_sum = votes
+            .iter()
+            .filter_map(
+                |(&(r, _voter), (_v, w))| if r == round { Some(u64::from(*w)) } else { None },
+            )
+            .sum();
+        threshold.is_met(weight_sum, self.total_weight)
     }
 
     fn value_has_enough_votes(
         &self,
-        votes: &HashMap<u32, HashMap<Option<ProposalCommitment>, u32>>,
+        votes: &HashMap<(Round, ValidatorId), (Vote, u32)>,
         round: u32,
         value: &Option<ProposalCommitment>,
         threshold: &VotesThreshold,
     ) -> bool {
-        threshold.is_met(
-            votes.get(&round).map_or(0, |v| *v.get(value).unwrap_or(&0)).into(),
-            self.total_weight,
-        )
+        let weight_sum = votes
+            .iter()
+            .filter_map(|(&(r, _voter), (v, w))| {
+                if r == round && &v.proposal_commitment == value {
+                    Some(u64::from(*w))
+                } else {
+                    None
+                }
+            })
+            .sum();
+        threshold.is_met(weight_sum, self.total_weight)
     }
 }

crates/apollo_consensus/src/state_machine_test.rs

@@ -14,6 +14,8 @@ use crate::votes_threshold::QuorumType;
 lazy_static! {
     static ref PROPOSER_ID: ValidatorId = DEFAULT_VALIDATOR_ID.into();
     static ref VALIDATOR_ID: ValidatorId = (DEFAULT_VALIDATOR_ID + 1).into();
+    static ref VALIDATOR_ID_2: ValidatorId = (DEFAULT_VALIDATOR_ID + 2).into();
+    static ref VALIDATOR_ID_3: ValidatorId = (DEFAULT_VALIDATOR_ID + 3).into();
 }
 
 const PROPOSAL_ID: Option<ProposalCommitment> = Some(ProposalCommitment(Felt::ONE));
@@ -33,6 +35,8 @@ struct TestWrapper<LeaderFn: Fn(Round) -> ValidatorId> {
     state_machine: StateMachine,
     leader_fn: LeaderFn,
     events: VecDeque<StateMachineEvent>,
+    peer_voters: Vec<ValidatorId>,
+    next_peer_idx: usize,
 }
 
 impl<LeaderFn: Fn(Round) -> ValidatorId> TestWrapper<LeaderFn> {
@@ -43,13 +47,27 @@ impl<LeaderFn: Fn(Round) -> ValidatorId> TestWrapper<LeaderFn> {
         is_observer: bool,
         quorum_type: QuorumType,
     ) -> Self {
+        let mut peer_voters = vec![*PROPOSER_ID, *VALIDATOR_ID, *VALIDATOR_ID_2, *VALIDATOR_ID_3]
+            .into_iter()
+            .filter(|v| *v != id)
+            .collect::<Vec<_>>();
+        // Ensure deterministic order.
+        peer_voters.sort();
         Self {
             state_machine: StateMachine::new(HEIGHT, id, total_weight, is_observer, quorum_type),
             leader_fn,
             events: VecDeque::new(),
+            peer_voters,
+            next_peer_idx: 0,
         }
     }
 
+    fn next_peer(&mut self) -> ValidatorId {
+        let voter = self.peer_voters[self.next_peer_idx % self.peer_voters.len()];
+        self.next_peer_idx += 1;
+        voter
+    }
+
     pub fn next_event(&mut self) -> Option<StateMachineEvent> {
         self.events.pop_front()
     }
@@ -67,20 +85,22 @@ impl<LeaderFn: Fn(Round) -> ValidatorId> TestWrapper<LeaderFn> {
     }
 
     pub fn send_prevote(&mut self, proposal_id: Option<ProposalCommitment>, round: Round) {
+        let voter = self.next_peer();
         self.send_event(StateMachineEvent::Prevote(mk_vote(
             VoteType::Prevote,
             round,
             proposal_id,
-            *PROPOSER_ID,
+            voter,
         )))
     }
 
     pub fn send_precommit(&mut self, proposal_id: Option<ProposalCommitment>, round: Round) {
+        let voter = self.next_peer();
         self.send_event(StateMachineEvent::Precommit(mk_vote(
             VoteType::Precommit,
             round,
             proposal_id,
-            *PROPOSER_ID,
+            voter,
         )))
     }