refactor: replay state as enum, #5971

Author: fdefelici

stacks-signer/src/v0/signer_state.rs

@@ -103,22 +103,25 @@ pub type ReplayScopeOpt = Option<ReplayScope>;
 /// Represents the Tx Replay state
 pub enum ReplayState {
     /// No replay has started yet, or the previous replay was cleared.
-    NotStartedOrCleared,
+    Unset,
     /// A replay is currently in progress, with an associated transaction set and scope.
     InProgress(ReplayTransactionSet, ReplayScope),
-    /// An invalid state — a transaction set was provided without a valid scope.
+    /// An invalid state where a replay set was provided without a valid scope.
+    /// Possibly caused by the scope being a local state in the `Signer` struct.
     Invalid,
 }
 
 impl ReplayState {
     /// Constructs a `ReplayState` from a given replay transaction set and an optional scope.
     ///
-    /// - Returns `NotStartedOrCleared` if the replay set is empty.
-    /// - Returns `Invalid` if the replay set exists but the scope is missing.
-    /// - Returns `InProgress` if both a non-empty replay set and a valid scope are provided.
-    fn from(replay_set: &ReplayTransactionSet, scope_opt: &ReplayScopeOpt) -> Self {
+    /// # Returns
+    ///
+    /// - `Unset` if the replay set is empty.
+    /// - `Invalid` if the replay set exists but the scope is missing.
+    /// - `InProgress` if both a non-empty replay set and a valid scope are provided.
+    fn infer_state(replay_set: &ReplayTransactionSet, scope_opt: &ReplayScopeOpt) -> Self {
         if replay_set.is_empty() {
-            return Self::NotStartedOrCleared;
+            return Self::Unset;
         }
 
         match scope_opt {
@@ -599,7 +602,7 @@ impl LocalStateMachine {
                 return Err(ClientError::InvalidResponse(err_msg).into());
             }
 
-            let replay_state = ReplayState::from(&tx_replay_set, &tx_replay_scope);
+            let replay_state = ReplayState::infer_state(&tx_replay_set, &tx_replay_scope);
             if let Some(new_replay_state) = self.handle_possible_bitcoin_fork(
                 db,
                 client,
@@ -608,15 +611,18 @@ impl LocalStateMachine {
                 replay_state,
             )? {
                 match new_replay_state {
-                    ReplayState::NotStartedOrCleared => {
+                    ReplayState::Unset => {
                         tx_replay_set = ReplayTransactionSet::none();
                         *tx_replay_scope = None;
                     }
                     ReplayState::InProgress(new_txs_set, new_scope) => {
                         tx_replay_set = new_txs_set;
                         *tx_replay_scope = Some(new_scope);
                     }
-                    ReplayState::Invalid => warn!("Tx Replay: Invalid is not possible here!"),
+                    ReplayState::Invalid => {
+                        warn!("Tx Replay: Invalid state due to scope being not set while in replay mode!");
+                        return Err(SignerChainstateError::LocalStateMachineNotReady);
+                    }
                 }
             }
         }
@@ -954,8 +960,13 @@ impl LocalStateMachine {
         state.tx_replay_set.clone_as_optional()
     }
 
-    /// Handle a possible bitcoin fork. If a fork is detetected,
-    /// return the transactions that should be replayed.
+    /// Handle a possible bitcoin fork. If a fork is detected,
+    /// try to handle the possible replay state.
+    ///
+    /// # Returns
+    /// - `Ok(None)` if nothing need to be done about replay
+    /// - `Ok(Some(ReplayState))` if a change (new or update) to the replay state is required
+    /// - `Err(SignerChainstateError)` in case of chain state errors
     pub fn handle_possible_bitcoin_fork(
         &self,
         db: &SignerDb,
@@ -973,51 +984,37 @@ impl LocalStateMachine {
             return Ok(None);
         }
 
-        if matches!(replay_state, ReplayState::Invalid) {
-            warn!("Tx Replay: Invalid state due to scope being not set while in replay mode!");
-            return Err(SignerChainstateError::LocalStateMachineNotReady);
-        }
-
-        if let ReplayState::InProgress(_txs_set, scope) = replay_state {
-            info!("Tx Replay: detected bitcoin fork while in replay mode. Tryng to handle the fork";
-                "expected_burn_block.height" => expected_burn_block.burn_block_height,
-                "expected_burn_block.hash" => %expected_burn_block.consensus_hash,
-                "prior_state_machine.burn_block_height" => prior_state_machine.burn_block_height,
-                "prior_state_machine.burn_block" => %prior_state_machine.burn_block,
-            );
-
-            let is_deepest_fork =
-                expected_burn_block.burn_block_height < scope.fork_origin.burn_block_height;
-            if !is_deepest_fork {
-                //if it is within the scope or after - this is not a new fork, but the continue of a reorg
-                info!("Tx Replay: nothing todo. Reorg in progress!");
-                return Ok(None);
-            }
-
-            let replay_state;
-            if let Some(replay_set) = self.compute_forked_txs_set_in_same_cycle(
+        match replay_state {
+            ReplayState::Unset => self.handle_fork_for_new_replay(
                 db,
                 client,
                 expected_burn_block,
-                &scope.past_tip,
-            )? {
-                let scope = ReplayScope {
-                    fork_origin: expected_burn_block.clone(),
-                    past_tip: scope.past_tip.clone(),
-                };
-
-                info!("Tx Replay: replay set updated with {} tx(s)", replay_set.len();
-                    "tx_replay_set" => ?replay_set,
-                    "tx_replay_scope" => ?scope);
-                replay_state =
-                    ReplayState::InProgress(ReplayTransactionSet::new(replay_set), scope);
-            } else {
-                info!("Tx Replay: replay set will be cleared, because the fork involves the previous reward cycle.");
-                replay_state = ReplayState::NotStartedOrCleared;
-            }
-            return Ok(Some(replay_state));
+                prior_state_machine,
+            ),
+            ReplayState::InProgress(_, scope) => self.handle_fork_on_in_progress_replay(
+                db,
+                client,
+                expected_burn_block,
+                prior_state_machine,
+                scope,
+            ),
+            ReplayState::Invalid => Ok(Some(ReplayState::Invalid)),
         }
+    }
 
+    /// Understand if the fork produces a replay set to be managed
+    ///
+    /// # Returns
+    ///
+    /// - `Ok(None)` if nothing need to be done
+    /// - `Ok(Some(ReplayState::InProgress(..)))` in case a replay need to be started
+    fn handle_fork_for_new_replay(
+        &self,
+        db: &SignerDb,
+        client: &StacksClient,
+        expected_burn_block: &NewBurnBlock,
+        prior_state_machine: &SignerStateMachine,
+    ) -> Result<Option<ReplayState>, SignerChainstateError> {
         info!("Signer State: fork detected";
             "expected_burn_block.height" => expected_burn_block.burn_block_height,
             "expected_burn_block.hash" => %expected_burn_block.consensus_hash,
@@ -1061,8 +1058,8 @@ impl LocalStateMachine {
                         past_tip: potential_replay_tip,
                     };
                     info!("Tx Replay: replay set updated with {} tx(s)", replay_set.len();
-                    "tx_replay_set" => ?replay_set,
-                    "tx_replay_scope" => ?scope);
+                        "tx_replay_set" => ?replay_set,
+                        "tx_replay_scope" => ?scope);
                     let replay_state =
                         ReplayState::InProgress(ReplayTransactionSet::new(replay_set), scope);
                     Ok(Some(replay_state))
@@ -1071,6 +1068,59 @@ impl LocalStateMachine {
         }
     }
 
+    /// Understand if the fork produces changes over an in-progress replay
+    ///
+    /// # Returns
+    ///
+    /// - `Ok(None)` if nothing need to be done
+    /// - `Ok(Some(ReplayState::Unset))` in case a replay set need to be cleared
+    /// - `Ok(Some(ReplayState::InProgress(..)))` in case a replay set need to be updated
+    fn handle_fork_on_in_progress_replay(
+        &self,
+        db: &SignerDb,
+        client: &StacksClient,
+        expected_burn_block: &NewBurnBlock,
+        prior_state_machine: &SignerStateMachine,
+        scope: ReplayScope,
+    ) -> Result<Option<ReplayState>, SignerChainstateError> {
+        info!("Tx Replay: detected bitcoin fork while in replay mode. Tryng to handle the fork";
+            "expected_burn_block.height" => expected_burn_block.burn_block_height,
+            "expected_burn_block.hash" => %expected_burn_block.consensus_hash,
+            "prior_state_machine.burn_block_height" => prior_state_machine.burn_block_height,
+            "prior_state_machine.burn_block" => %prior_state_machine.burn_block,
+        );
+
+        let is_deepest_fork =
+            expected_burn_block.burn_block_height < scope.fork_origin.burn_block_height;
+        if !is_deepest_fork {
+            //if it is within the scope or after - this is not a new fork, but the continue of a reorg
+            info!("Tx Replay: nothing todo. Reorg in progress!");
+            return Ok(None);
+        }
+
+        let replay_state;
+        if let Some(replay_set) = self.compute_forked_txs_set_in_same_cycle(
+            db,
+            client,
+            expected_burn_block,
+            &scope.past_tip,
+        )? {
+            let scope = ReplayScope {
+                fork_origin: expected_burn_block.clone(),
+                past_tip: scope.past_tip.clone(),
+            };
+
+            info!("Tx Replay: replay set updated with {} tx(s)", replay_set.len();
+                    "tx_replay_set" => ?replay_set,
+                    "tx_replay_scope" => ?scope);
+            replay_state = ReplayState::InProgress(ReplayTransactionSet::new(replay_set), scope);
+        } else {
+            info!("Tx Replay: replay set will be cleared, because the fork involves the previous reward cycle.");
+            replay_state = ReplayState::Unset;
+        }
+        return Ok(Some(replay_state));
+    }
+
     /// Retrieves the set of transactions that were part of a Bitcoin fork within the same reward cycle.
     ///
     /// This method identifies the range of Tenures affected by a fork, from the `fork_tip` down to the `fork_origin`