apollo_consensus: move vote storage to SM; dedupe in SHC

Author: asmaastarkware

crates/apollo_consensus/src/single_height_consensus.rs

@@ -139,8 +139,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>,
     height_voted_storage: Arc<Mutex<dyn HeightVotedStorageTrait>>,
@@ -165,8 +163,6 @@ impl SingleHeightConsensus {
             timeouts,
             state_machine,
             proposals: HashMap::new(),
-            prevotes: HashMap::new(),
-            precommits: HashMap::new(),
             last_prevote: None,
             last_precommit: None,
             height_voted_storage,
@@ -259,7 +255,7 @@ impl SingleHeightConsensus {
                 context.broadcast(last_vote.clone()).await?;
                 Ok(ShcReturn::Tasks(vec![ShcTask::Prevote(
                     self.timeouts.get_prevote_timeout(0),
-                    StateMachineEvent::Prevote(last_vote.clone()),
+                    StateMachineEvent::Prevote(vote),
                 )]))
             }
             StateMachineEvent::Precommit(vote) => {
@@ -276,7 +272,7 @@ impl SingleHeightConsensus {
                 context.broadcast(last_vote.clone()).await?;
                 Ok(ShcReturn::Tasks(vec![ShcTask::Precommit(
                     self.timeouts.get_precommit_timeout(0),
-                    StateMachineEvent::Precommit(last_vote.clone()),
+                    StateMachineEvent::Precommit(vote),
                 )]))
             }
             StateMachineEvent::Proposal(proposal_id, round, valid_round) => {
@@ -354,29 +350,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) };
@@ -507,29 +501,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(0),
+                    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(0),
+                    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()),
@@ -583,18 +572,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())
@@ -610,4 +588,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/single_height_consensus_test.rs

@@ -54,7 +54,7 @@ fn get_proposal_init_for_height(height: BlockNumber) -> ProposalInit {
 }
 
 fn prevote_task(block_felt: Option<Felt>, height: u64, round: u32, voter: ValidatorId) -> ShcTask {
-    let duration = TIMEOUTS.get_prevote_timeout(round);
+    let duration = TIMEOUTS.get_prevote_timeout(0);
     ShcTask::Prevote(
         duration,
         StateMachineEvent::Prevote(prevote(block_felt, height, round, voter)),
@@ -67,7 +67,7 @@ fn precommit_task(
     round: u32,
     voter: ValidatorId,
 ) -> ShcTask {
-    let duration = TIMEOUTS.get_precommit_timeout(round);
+    let duration = TIMEOUTS.get_precommit_timeout(0);
     ShcTask::Precommit(
         duration,
         StateMachineEvent::Precommit(precommit(block_felt, height, round, voter)),
@@ -357,7 +357,7 @@ async fn rebroadcast_votes() {
         Arc::new(Mutex::new(NoOpHeightVotedStorage)),
     );
 
-    context.expect_proposer().times(1).returning(move |_, _| *PROPOSER_ID);
+    context.expect_proposer().returning(move |_, _| *PROPOSER_ID);
     context.expect_build_proposal().times(1).returning(move |_, _| {
         let (block_sender, block_receiver) = oneshot::channel();
         block_sender.send(BLOCK.id).unwrap();
@@ -383,27 +383,72 @@ async fn rebroadcast_votes() {
     // 3 of 4 Prevotes is enough to send a Precommit.
     context
         .expect_broadcast()
-        .times(2) // vote rebroadcast
-        .withf(move |msg: &Vote| {
-            msg == &precommit(Some(BLOCK.id.0), 0, 0, *PROPOSER_ID)
-        })
+        .times(1) // just the first precommit at round 0.
+        .withf(move |msg: &Vote| msg == &precommit(Some(BLOCK.id.0), 0, 0, *PROPOSER_ID))
         .returning(move |_| Ok(()));
     // The Node got a Prevote quorum.
     assert_eq!(
         shc.handle_vote(&mut context, prevote(Some(BLOCK.id.0), 0, 0, *VALIDATOR_ID_2)).await,
         Ok(ShcReturn::Tasks(vec![
             timeout_prevote_task(0),
+            // initiate timeout for rebroadcast at round 0.
             precommit_task(Some(BLOCK.id.0), 0, 0, *PROPOSER_ID),
         ]))
     );
-    // Re-broadcast vote.
+    // Advance to the next round with NIL precommits.
+    shc.handle_vote(&mut context, precommit(None, 0, 0, *VALIDATOR_ID_1)).await.unwrap();
+    shc.handle_vote(&mut context, precommit(None, 0, 0, *VALIDATOR_ID_2)).await.unwrap();
+    context.expect_repropose().returning(move |id, init| {
+        assert_eq!(init.height, BlockNumber(0));
+        assert_eq!(id, BLOCK.id);
+        assert_eq!(init.round, 1);
+    });
+    context
+        .expect_broadcast()
+        .times(1)
+        .withf(move |msg: &Vote| msg == &prevote(Some(BLOCK.id.0), 0, 1, *PROPOSER_ID))
+        .returning(move |_| Ok(()));
+    shc.handle_event(&mut context, StateMachineEvent::TimeoutPrecommit(0)).await.unwrap();
+    context
+        .expect_broadcast()
+        .times(2) // first: initial precommit vote at r1, second: rebroadcast at r1
+        .withf(move |msg: &Vote| msg == &precommit(Some(BLOCK.id.0), 0, 1, *PROPOSER_ID))
+        .returning(move |_| Ok(()));
+    assert_eq!(
+        shc.handle_vote(&mut context, prevote(Some(BLOCK.id.0), 0, 1, *VALIDATOR_ID_2)).await,
+        Ok(ShcReturn::Tasks(Vec::new()))
+    );
+    assert_eq!(
+        shc.handle_vote(&mut context, prevote(Some(BLOCK.id.0), 0, 1, *VALIDATOR_ID_3)).await,
+        Ok(ShcReturn::Tasks(vec![
+            timeout_prevote_task(1),
+            // initiate timeout for rebroadcast at round 1.
+            precommit_task(Some(BLOCK.id.0), 0, 1, *PROPOSER_ID),
+        ]))
+    );
+
+    // Re-broadcast with older vote (round 0) - should be ignored (no broadcast, no task).
     assert_eq!(
         shc.handle_event(
             &mut context,
             StateMachineEvent::Precommit(precommit(Some(BLOCK.id.0), 0, 0, *PROPOSER_ID))
         )
         .await,
-        Ok(ShcReturn::Tasks(vec![precommit_task(Some(BLOCK.id.0), 0, 0, *PROPOSER_ID),]))
+        Ok(ShcReturn::Tasks(Vec::new()))
+    );
+
+    // Re-broadcast with current round (round 1) - should broadcast and schedule another timeout for
+    // rebroadcast at round 1.
+    assert_eq!(
+        shc.handle_event(
+            &mut context,
+            StateMachineEvent::Precommit(precommit(Some(BLOCK.id.0), 0, 1, *PROPOSER_ID))
+        )
+        .await,
+        Ok(ShcReturn::Tasks(vec![ShcTask::Precommit(
+            TIMEOUTS.get_precommit_timeout(0),
+            StateMachineEvent::Precommit(precommit(Some(BLOCK.id.0), 0, 1, *PROPOSER_ID)),
+        )]))
     );
 }
 

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,15 @@ 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;
+        let voter = vote.voter;
+        let inserted = self.prevotes.insert((round, voter), (vote, 1)).is_none();
+        assert!(
+            inserted,
+            "SHC should handle conflicts & replays: duplicate prevote for round={round}, \
+             voter={voter}",
+        );
+        self.map_round_to_upons(round, leader_fn)
     }
 
     pub(crate) fn handle_timeout_prevote(&mut self, round: u32) -> VecDeque<StateMachineEvent> {
@@ -360,15 +368,15 @@ 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;
+        let voter = vote.voter;
+        let inserted = self.precommits.insert((round, voter), (vote, 1)).is_none();
+        assert!(
+            inserted,
+            "SHC should handle conflicts & replays: duplicate precommit for round={round}, \
+             voter={voter}"
+        );
+        self.map_round_to_upons(round, leader_fn)
     }
 
     fn handle_timeout_precommit<LeaderFn>(
@@ -640,24 +648,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)
     }
 }