From a4e7bcd8646ad751a2d2b36006ccf13c11860544 Mon Sep 17 00:00:00 2001
From: Herr Seppia <seppia@dusk.network>
Date: Tue, 22 Oct 2024 17:08:32 +0200
Subject: [PATCH] consensus: fix proposal early timeout

Fix the edge case where next_slot is greater than the event_loop duration
---
 consensus/src/execution_ctx.rs     |  5 ++--
 consensus/src/proposal/step.rs     | 38 +++++++++++++++++-------------
 consensus/src/ratification/step.rs |  2 +-
 consensus/src/validation/step.rs   |  2 +-
 4 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/consensus/src/execution_ctx.rs b/consensus/src/execution_ctx.rs
index 9a1a380d4..cc0309fb0 100644
--- a/consensus/src/execution_ctx.rs
+++ b/consensus/src/execution_ctx.rs
@@ -129,6 +129,7 @@ impl<'a, T: Operations + 'static, DB: Database> ExecutionCtx<'a, T, DB> {
     pub async fn event_loop<C: MsgHandler>(
         &mut self,
         phase: Arc<Mutex<C>>,
+        additional_timeout: Option<Duration>,
     ) -> Result<Message, ConsensusError> {
         let open_consensus_mode = self.last_step_running();
 
@@ -141,8 +142,8 @@ impl<'a, T: Operations + 'static, DB: Database> ExecutionCtx<'a, T, DB> {
             dur
         } else {
             let dur = self.iter_ctx.get_timeout(self.step_name());
-            debug!(event = "run event_loop", ?dur);
-            dur
+            debug!(event = "run event_loop", ?dur, ?additional_timeout);
+            dur + additional_timeout.unwrap_or_default()
         };
 
         let deadline = Instant::now().checked_add(timeout).unwrap();
diff --git a/consensus/src/proposal/step.rs b/consensus/src/proposal/step.rs
index 0ad85ab50..f971e9cc6 100644
--- a/consensus/src/proposal/step.rs
+++ b/consensus/src/proposal/step.rs
@@ -60,6 +60,8 @@ impl<T: Operations + 'static, D: Database> ProposalStep<T, D> {
             .get_current_committee()
             .expect("committee to be created before run");
 
+        let tip_timestamp = ctx.round_update.timestamp();
+
         if ctx.am_member(committee) {
             let iteration =
                 cmp::min(config::RELAX_ITERATION_THRESHOLD, ctx.iteration);
@@ -88,10 +90,7 @@ impl<T: Operations + 'static, D: Database> ProposalStep<T, D> {
                     .await
                 {
                     Ok(HandleMsgOutput::Ready(msg)) => {
-                        Self::wait_until_next_slot(
-                            ctx.round_update.timestamp(),
-                        )
-                        .await;
+                        Self::wait_until_next_slot(tip_timestamp).await;
                         return Ok(msg);
                     }
                     Err(e) => {
@@ -106,14 +105,17 @@ impl<T: Operations + 'static, D: Database> ProposalStep<T, D> {
 
         // handle queued messages for current round and step.
         if let Some(m) = ctx.handle_future_msgs(self.handler.clone()).await {
-            Self::wait_until_next_slot(ctx.round_update.timestamp()).await;
+            Self::wait_until_next_slot(tip_timestamp).await;
             return Ok(m);
         }
 
-        match ctx.event_loop(self.handler.clone()).await {
+        let additional_timeout = Self::next_slot_in(tip_timestamp);
+        match ctx
+            .event_loop(self.handler.clone(), additional_timeout)
+            .await
+        {
             Ok(msg) => {
-                Self::wait_until_next_slot(ctx.round_update.timestamp()).await;
-
+                Self::wait_until_next_slot(tip_timestamp).await;
                 Ok(msg)
             }
             Err(err) => Err(err),
@@ -122,19 +124,23 @@ impl<T: Operations + 'static, D: Database> ProposalStep<T, D> {
 
     /// Waits until the next slot is reached
     async fn wait_until_next_slot(tip_timestamp: u64) {
+        if let Some(delay) = Self::next_slot_in(tip_timestamp) {
+            info!(event = "next_slot", ?delay);
+            tokio::time::sleep(delay).await;
+        }
+    }
+
+    /// Calculate the duration needed to the next slot
+    fn next_slot_in(tip_timestamp: u64) -> Option<Duration> {
         let current_time_secs = get_current_timestamp();
 
         let next_slot_timestamp = tip_timestamp + MINIMUM_BLOCK_TIME;
         if current_time_secs >= next_slot_timestamp {
-            return;
+            None
+        } else {
+            // block_timestamp - localtime
+            Some(Duration::from_secs(next_slot_timestamp - current_time_secs))
         }
-
-        // block_timestamp - localtime
-        let delay =
-            Duration::from_secs(next_slot_timestamp - current_time_secs);
-
-        info!(event = "next_slot", ?delay);
-        tokio::time::sleep(delay).await;
     }
 
     pub fn name(&self) -> &'static str {
diff --git a/consensus/src/ratification/step.rs b/consensus/src/ratification/step.rs
index 948354acb..3c218f244 100644
--- a/consensus/src/ratification/step.rs
+++ b/consensus/src/ratification/step.rs
@@ -148,7 +148,7 @@ impl RatificationStep {
             return Ok(m);
         }
 
-        ctx.event_loop(self.handler.clone()).await
+        ctx.event_loop(self.handler.clone(), None).await
     }
 
     pub fn name(&self) -> &'static str {
diff --git a/consensus/src/validation/step.rs b/consensus/src/validation/step.rs
index 939006a77..a54ee3df0 100644
--- a/consensus/src/validation/step.rs
+++ b/consensus/src/validation/step.rs
@@ -274,7 +274,7 @@ impl<T: Operations + 'static> ValidationStep<T> {
             return Ok(m);
         }
 
-        ctx.event_loop(self.handler.clone()).await
+        ctx.event_loop(self.handler.clone(), None).await
     }
 
     pub fn name(&self) -> &'static str {