fix: wait for correct RPC epoch state when processing epoch events (cherry-pick #3005) (#3013)

fix: wait for correct RPC epoch state when processing epoch events (#3005)

Co-authored-by: Markus Legner <mlegner@users.noreply.github.com>

Author: ebmifa

crates/walrus-service/src/node.rs

@@ -46,6 +46,8 @@ use sui_macros::{fail_point_arg, fail_point_async};
 use sui_types::{base_types::ObjectID, event::EventID};
 use system_events::{CompletableHandle, EVENT_ID_FOR_CHECKPOINT_EVENTS, EventHandle};
 use thread_pool::{BoundedThreadPool, ThreadPoolBuilder};
+#[cfg(not(any(test, msim)))]
+use tokio::time::sleep;
 use tokio::{
     select,
     sync::{Notify, RwLock, watch},
@@ -132,7 +134,7 @@ use walrus_storage_node_client::{
         StoredOnNodeStatus,
     },
 };
-use walrus_sui::client::FixedSystemParameters;
+use walrus_sui::{client::FixedSystemParameters, types::move_structs::EpochState};
 use walrus_utils::metrics::{Registry, TaskMonitorFamily, monitored_scope};
 
 use self::{
@@ -1646,13 +1648,18 @@ impl StorageNode {
                 let _scope = monitored_scope::monitored_scope(
                     "ProcessEvent::EpochChangeEvent::EpochParametersSelected",
                 );
+                self.wait_for_epoch_state(event.next_epoch.saturating_sub(1), |state| {
+                    matches!(state, EpochState::NextParamsSelected(_))
+                })
+                .await?;
                 self.handle_epoch_parameters_selected(event);
                 event_handle.mark_as_complete();
             }
             EpochChangeEvent::EpochChangeStart(event) => {
                 let _scope = monitored_scope::monitored_scope(
                     "ProcessEvent::EpochChangeEvent::EpochChangeStart",
                 );
+                self.wait_for_epoch_state(event.epoch, |_| true).await?;
                 fail_point_async!("epoch_change_start_entry");
                 self.process_epoch_change_start_event(blob_event_processor, event_handle, &event)
                     .await?;
@@ -1661,6 +1668,13 @@ impl StorageNode {
                 let _scope = monitored_scope::monitored_scope(
                     "ProcessEvent::EpochChangeEvent::EpochChangeDone",
                 );
+                self.wait_for_epoch_state(event.epoch, |state| {
+                    matches!(
+                        state,
+                        EpochState::EpochChangeDone(_) | EpochState::NextParamsSelected(_)
+                    )
+                })
+                .await?;
                 self.process_epoch_change_done_event(&event).await?;
                 event_handle.mark_as_complete();
             }
@@ -1680,6 +1694,51 @@ impl StorageNode {
         Ok(())
     }
 
+    /// Repeatedly checks until the current Sui epoch state matches the expectation.
+    ///
+    /// Returns `Ok(())` if the current epoch is equal to the `expected_epoch` and the
+    /// `state_matches` function returns `true` or the current epoch is greater than the
+    /// `expected_epoch` (irrespective of the state).
+    #[cfg(not(any(test, msim)))]
+    async fn wait_for_epoch_state(
+        &self,
+        expected_epoch: Epoch,
+        state_matches: impl Fn(&EpochState) -> bool,
+    ) -> anyhow::Result<()> {
+        const EPOCH_STATE_WAIT_TIMEOUT: Duration = Duration::from_secs(5);
+        const EPOCH_STATE_WAIT_SLEEP: Duration = Duration::from_millis(100);
+        let deadline = Instant::now() + EPOCH_STATE_WAIT_TIMEOUT;
+        while Instant::now() < deadline {
+            self.inner.contract_service.flush_cache().await;
+            let Ok((epoch, state)) = self.inner.contract_service.get_epoch_and_state().await else {
+                tracing::warn!("failed to get current epoch and state");
+                continue;
+            };
+            if epoch == expected_epoch && state_matches(&state) || epoch > expected_epoch {
+                return Ok(());
+            }
+            tracing::debug!(
+                expected_epoch,
+                current_epoch = epoch,
+                current_state = ?state,
+                "waiting for expected epoch state",
+            );
+            sleep(EPOCH_STATE_WAIT_SLEEP).await;
+        }
+        bail!("timed out after waiting for expected epoch state")
+    }
+
+    #[cfg(any(test, msim))]
+    #[allow(clippy::unused_async)]
+    async fn wait_for_epoch_state(
+        &self,
+        _expected_epoch: Epoch,
+        _state_matches: impl Fn(&EpochState) -> bool,
+    ) -> anyhow::Result<()> {
+        tracing::info!("waiting for epoch state is not supported in tests, skipping");
+        Ok(())
+    }
+
     /// Handles the epoch parameters selected event.
     ///
     /// This function cancels the scheduled voting end and initiates the epoch change.