apollo_consensus: unify timeout scheduling requests into single enum variant

Author: asmaastarkware

crates/apollo_consensus/src/manager.rs

@@ -39,7 +39,7 @@ use crate::metrics::{
     CONSENSUS_REPROPOSALS,
 };
 use crate::single_height_consensus::{ShcReturn, SingleHeightConsensus};
-use crate::state_machine::{SMRequest, StateMachineEvent};
+use crate::state_machine::{SMRequest, StateMachineEvent, Step};
 use crate::storage::HeightVotedStorageTrait;
 use crate::types::{
     BroadcastVoteChannel,
@@ -785,29 +785,24 @@ impl<ContextT: ConsensusContext> MultiHeightManager<ContextT> {
                 .boxed();
                 Ok(Some(fut))
             }
-            SMRequest::ScheduleTimeoutPropose(round) => {
-                let duration = timeouts.get_proposal_timeout(round);
-                let fut = async move {
-                    tokio::time::sleep(duration).await;
-                    StateMachineEvent::TimeoutPropose(round)
-                }
-                .boxed();
-                Ok(Some(fut))
-            }
-            SMRequest::ScheduleTimeoutPrevote(round) => {
-                let duration = timeouts.get_prevote_timeout(round);
-                let fut = async move {
-                    tokio::time::sleep(duration).await;
-                    StateMachineEvent::TimeoutPrevote(round)
-                }
-                .boxed();
-                Ok(Some(fut))
-            }
-            SMRequest::ScheduleTimeoutPrecommit(round) => {
-                let duration = timeouts.get_precommit_timeout(round);
+            SMRequest::ScheduleTimeout(step, round) => {
+                let (duration, event) = match step {
+                    Step::Propose => (
+                        timeouts.get_proposal_timeout(round),
+                        StateMachineEvent::TimeoutPropose(round),
+                    ),
+                    Step::Prevote => (
+                        timeouts.get_prevote_timeout(round),
+                        StateMachineEvent::TimeoutPrevote(round),
+                    ),
+                    Step::Precommit => (
+                        timeouts.get_precommit_timeout(round),
+                        StateMachineEvent::TimeoutPrecommit(round),
+                    ),
+                };
                 let fut = async move {
                     tokio::time::sleep(duration).await;
-                    StateMachineEvent::TimeoutPrecommit(round)
+                    event
                 }
                 .boxed();
                 Ok(Some(fut))

crates/apollo_consensus/src/single_height_consensus_test.rs

@@ -7,7 +7,7 @@ use test_case::test_case;
 
 use super::SingleHeightConsensus;
 use crate::single_height_consensus::ShcReturn;
-use crate::state_machine::{SMRequest, StateMachineEvent};
+use crate::state_machine::{SMRequest, StateMachineEvent, Step};
 use crate::test_utils::{precommit, prevote, TestBlock};
 use crate::types::{ProposalCommitment, ValidatorId};
 use crate::votes_threshold::QuorumType;
@@ -70,7 +70,7 @@ fn proposer() {
         ShcReturn::Requests(r) => r,
         _ => panic!("expected requests"),
     };
-    assert!(matches!(reqs.pop_front(), Some(SMRequest::ScheduleTimeoutPrevote { .. })));
+    assert!(matches!(reqs.pop_front(), Some(SMRequest::ScheduleTimeout(Step::Prevote, _))));
     assert!(matches!(
         reqs.pop_front(),
         Some(SMRequest::BroadcastVote(v)) if v.vote_type == VoteType::Precommit
@@ -143,7 +143,7 @@ fn validator(repeat_proposal: bool) {
         ShcReturn::Requests(r) => r,
         _ => panic!("expected requests"),
     };
-    assert!(matches!(reqs.pop_front(), Some(SMRequest::ScheduleTimeoutPrevote { .. })));
+    assert!(matches!(reqs.pop_front(), Some(SMRequest::ScheduleTimeout(Step::Prevote, _))));
     assert!(
         matches!(reqs.pop_front(), Some(SMRequest::BroadcastVote(v)) if v.vote_type == VoteType::Precommit)
     );
@@ -191,7 +191,7 @@ fn vote_twice(same_vote: bool) {
         ShcReturn::Requests(r) => r,
         _ => panic!("expected requests"),
     };
-    assert!(matches!(reqs.pop_front(), Some(SMRequest::ScheduleTimeoutPrevote { .. })));
+    assert!(matches!(reqs.pop_front(), Some(SMRequest::ScheduleTimeout(Step::Prevote, _))));
     assert!(matches!(
         reqs.pop_front(),
         Some(SMRequest::BroadcastVote(v)) if v.vote_type == VoteType::Precommit
@@ -350,7 +350,9 @@ fn repropose() {
     // Expect ScheduleTimeoutPrevote{round:0} and BroadcastVote(Precommit).
     match ret {
         ShcReturn::Requests(reqs) => {
-            assert!(reqs.iter().any(|req| matches!(req, SMRequest::ScheduleTimeoutPrevote(0))));
+            assert!(
+                reqs.iter().any(|req| matches!(req, SMRequest::ScheduleTimeout(Step::Prevote, 0)))
+            );
             assert!(reqs.iter().any(|req| matches!(
                 req,
                 SMRequest::BroadcastVote(v) if v.vote_type == VoteType::Precommit && v.round == 0
@@ -365,7 +367,10 @@ fn repropose() {
     // assert that ret is ScheduleTimeoutPrecommit
     match ret {
         ShcReturn::Requests(reqs) => {
-            assert!(reqs.iter().any(|req| matches!(req, SMRequest::ScheduleTimeoutPrecommit(0))));
+            assert!(
+                reqs.iter()
+                    .any(|req| matches!(req, SMRequest::ScheduleTimeout(Step::Precommit, 0)))
+            );
         }
         _ => panic!("expected requests"),
     }

crates/apollo_consensus/src/state_machine.rs

@@ -67,12 +67,8 @@ pub(crate) enum SMRequest {
     StartValidateProposal(ProposalInit),
     /// Request to broadcast a Prevote or Precommit vote.
     BroadcastVote(Vote),
-    /// Request to schedule a TimeoutPropose.
-    ScheduleTimeoutPropose(Round),
-    /// Request to schedule a TimeoutPrevote.
-    ScheduleTimeoutPrevote(Round),
-    /// Request to schedule a TimeoutPrecommit.
-    ScheduleTimeoutPrecommit(Round),
+    /// Request to schedule a timeout for a specific step and round.
+    ScheduleTimeout(Step, Round),
     /// Decision reached for the given proposal and round.
     DecisionReached(ProposalCommitment, Round),
     /// Request to re-propose (sent by the leader after advancing to a new round
@@ -511,7 +507,7 @@ impl StateMachine {
             }
         } else {
             info!("START_ROUND_VALIDATOR: Starting round {round} as Validator");
-            VecDeque::from([SMRequest::ScheduleTimeoutPropose(self.round)])
+            VecDeque::from([SMRequest::ScheduleTimeout(Step::Propose, self.round)])
         };
         output.append(&mut self.current_round_upons());
         output
@@ -625,7 +621,7 @@ impl StateMachine {
         if !self.mixed_prevote_quorum.insert(self.round) {
             return VecDeque::new();
         }
-        VecDeque::from([SMRequest::ScheduleTimeoutPrevote(self.round)])
+        VecDeque::from([SMRequest::ScheduleTimeout(Step::Prevote, self.round)])
     }
 
     // LOC 36 in the paper.
@@ -684,7 +680,7 @@ impl StateMachine {
         if !self.mixed_precommit_quorum.insert(self.round) {
             return VecDeque::new();
         }
-        VecDeque::from([SMRequest::ScheduleTimeoutPrecommit(self.round)])
+        VecDeque::from([SMRequest::ScheduleTimeout(Step::Precommit, self.round)])
     }
 
     // LOC 49 in the paper.