feat(state): sequence branch switches as evaluate, body rollback, then header rewrite

zebra-state/src/lib.rs

@@ -55,8 +55,8 @@ pub use request::{
 pub use request::Spend;
 
 pub use response::{
-    AnyTx, GetBlockTemplateChainInfo, KnownBlock, MinedTx, NonFinalizedBlocksListener,
-    ReadResponse, Response,
+    AnyTx, GetBlockTemplateChainInfo, HeaderRangeCommitOutcome, KnownBlock, MinedTx,
+    NonFinalizedBlocksListener, ReadResponse, Response,
 };
 pub use service::{
     chain_tip::{ChainTipBlock, ChainTipChange, ChainTipSender, LatestChainTip, TipAction},

zebra-state/src/response.rs

@@ -27,13 +27,32 @@ use crate::{ReadRequest, Request};
 
 use crate::{service::read::AddressUtxos, NonFinalizedState, TransactionLocation, WatchReceiver};
 
+/// The outcome of a successful [`Request::CommitHeaderRange`].
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub struct HeaderRangeCommitOutcome {
+    /// Hash of the last header in the committed range.
+    pub tip_hash: block::Hash,
+    /// First height where the committed range replaced a previously stored
+    /// conflicting header suffix, when the commit reorged the stored header
+    /// chain. `None` when the range only extended or duplicated stored headers.
+    pub reorged_at: Option<block::Height>,
+    /// The new branch's hash at `reorged_at`, when the commit reorged the
+    /// stored header chain. Used by the switch orchestration to recognize the
+    /// stranded old-branch body suffix without re-reading state.
+    pub reorged_to_hash: Option<block::Hash>,
+}
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 /// A response to a [`StateService`](crate::service::StateService) [`Request`].
 pub enum Response {
     /// Response to [`Request::CommitSemanticallyVerifiedBlock`] and [`Request::CommitCheckpointVerifiedBlock`]
     /// indicating that a block was successfully committed to the state.
     Committed(block::Hash),
 
+    /// Response to [`Request::CommitHeaderRange`] indicating that a header
+    /// range was successfully committed to the state.
+    CommittedHeaderRange(HeaderRangeCommitOutcome),
+
     /// Response to [`Request::InvalidateBlock`] indicating that a block was found and
     /// invalidated in the state.
     Invalidated(block::Hash),

zebra-state/src/service.rs

@@ -59,7 +59,8 @@ use crate::{
         watch_receiver::WatchReceiver,
     },
     BoxError, CheckpointVerifiedBlock, CommitHeaderRangeError, CommitSemanticallyVerifiedError,
-    Config, KnownBlock, ReadRequest, ReadResponse, Request, Response, SemanticallyVerifiedBlock,
+    Config, HeaderRangeCommitOutcome, KnownBlock, ReadRequest, ReadResponse, Request, Response,
+    SemanticallyVerifiedBlock,
 };
 
 pub mod block_iter;
@@ -988,7 +989,7 @@ impl StateService {
         headers: Vec<Arc<block::Header>>,
         body_sizes: Vec<u32>,
         tree_aux_roots: Vec<BlockCommitmentRoots>,
-    ) -> oneshot::Receiver<Result<block::Hash, CommitHeaderRangeError>> {
+    ) -> oneshot::Receiver<Result<HeaderRangeCommitOutcome, CommitHeaderRangeError>> {
         let (rsp_tx, rsp_rx) = oneshot::channel();
 
         let Some(sender) = &self.block_write_sender.non_finalized else {
@@ -1255,7 +1256,7 @@ impl Service<Request> for StateService {
                         .map_err(|_recv_error| CommitHeaderRangeError::CommitResponseDropped)
                         .and_then(|result| result)
                         .map_err(BoxError::from)
-                        .map(Response::Committed)
+                        .map(Response::CommittedHeaderRange)
                 }
                 .instrument(span)
                 .boxed()

zebra-state/src/service/finalized_state/zebra_db/block.rs

@@ -51,7 +51,7 @@ use crate::{
         FromDisk, IntoDisk, RawBytes, COMMITMENT_ROOTS_BY_HEIGHT, PRUNING_METADATA,
         VCT_SYNC_METADATA, VCT_UPGRADE_METADATA,
     },
-    HashOrHeight,
+    HashOrHeight, HeaderRangeCommitOutcome,
 };
 
 #[cfg(feature = "indexer")]
@@ -75,6 +75,7 @@ pub(crate) struct AdvertisedBodySize(u32);
 
 /// Builds an [`AdvertisedBodySize`] for tests that plant raw rows.
 #[cfg(test)]
+#[allow(dead_code)] // used by the coherence suite, dropped from canary picks
 pub(crate) fn test_body_size(size: u32) -> AdvertisedBodySize {
     AdvertisedBodySize(size)
 }
@@ -2019,7 +2020,7 @@ impl DiskWriteBatch {
         anchor: block::Hash,
         headers: &[Arc<block::Header>],
         body_sizes: &[u32],
-    ) -> Result<block::Hash, CommitHeaderRangeError> {
+    ) -> Result<HeaderRangeCommitOutcome, CommitHeaderRangeError> {
         let roots = inferred_header_range_roots(zebra_db, anchor, headers.len())?;
         self.prepare_header_range_batch_with_roots(zebra_db, anchor, headers, body_sizes, &roots)
     }
@@ -2034,7 +2035,7 @@ impl DiskWriteBatch {
         headers: &[Arc<block::Header>],
         body_sizes: &[u32],
         tree_aux_roots: &[BlockCommitmentRoots],
-    ) -> Result<block::Hash, CommitHeaderRangeError> {
+    ) -> Result<HeaderRangeCommitOutcome, CommitHeaderRangeError> {
         if headers.is_empty() {
             return Err(CommitHeaderRangeError::EmptyRange);
         }
@@ -2328,9 +2329,17 @@ impl DiskWriteBatch {
             self.set_canonical_suffix(zebra_db, (fork_height, fork_hash), &new_rows)?;
         }
 
-        Ok(block::Hash::from(
-            &**headers.last().expect("headers is non-empty"),
-        ))
+        Ok(HeaderRangeCommitOutcome {
+            tip_hash: block::Hash::from(&**headers.last().expect("headers is non-empty")),
+            reorged_at: first_conflicting_height,
+            reorged_to_hash: first_conflicting_height.map(|reorged_at| {
+                validated_headers
+                    .iter()
+                    .find(|(height, ..)| *height == reorged_at)
+                    .expect("the first conflicting height is inside the validated range")
+                    .1
+            }),
+        })
     }
 
     /// Deletes the block header at `height`.

zebra-state/src/service/finalized_state/zebra_db/block/tests/common.rs

@@ -125,13 +125,13 @@ pub(super) fn commit_header_range(
 ) -> block::Hash {
     let mut batch = DiskWriteBatch::new();
     let body_sizes = vec![0; headers.len()];
-    let committed_hash = batch
+    let outcome = batch
         .prepare_header_range_batch(state, anchor, headers, &body_sizes)
         .expect("header range is valid");
     state
         .write_batch(batch)
         .expect("header range batch writes successfully");
-    committed_hash
+    outcome.tip_hash
 }
 
 /// Commits `block`'s header and transaction data (the body-commit batch shape),

zebra-state/src/service/finalized_state/zebra_db/block/tests/header_store_coherence/reads.rs

@@ -61,7 +61,9 @@ fn redeliver_trunk(
         .collect();
     let body_sizes = vec![0; headers.len()];
     let mut batch = DiskWriteBatch::new();
-    let result = batch.prepare_header_range_batch(state, anchor, &headers, &body_sizes);
+    let result = batch
+        .prepare_header_range_batch(state, anchor, &headers, &body_sizes)
+        .map(|outcome| outcome.tip_hash);
     if result.is_ok() {
         state
             .write_batch(batch)

zebra-state/src/service/finalized_state/zebra_db/block/tests/vectors.rs

@@ -106,7 +106,7 @@ fn header_range_commit_keeps_body_availability_separate() {
     let (state, genesis, block1) = mainnet_state_with_genesis();
 
     let mut batch = DiskWriteBatch::new();
-    let committed_hash = batch
+    let outcome = batch
         .prepare_header_range_batch(
             &state,
             genesis.hash(),
@@ -118,7 +118,7 @@ fn header_range_commit_keeps_body_availability_separate() {
         .write_batch(batch)
         .expect("header range batch writes successfully");
 
-    assert_eq!(committed_hash, block1.hash());
+    assert_eq!(outcome.tip_hash, block1.hash());
     assert_eq!(state.best_header_tip(), Some((Height(1), block1.hash())));
     assert_eq!(state.finalized_tip_height(), Some(Height(0)));
     assert_eq!(state.tip(), Some((Height(0), genesis.hash())));

zebra-state/src/service/non_finalized_state.rs

@@ -370,21 +370,25 @@ impl NonFinalizedState {
         Ok(())
     }
 
-    /// Invalidate block with hash `block_hash` and all descendants from the non-finalized state. Insert
-    /// the new chain into the chain_set and discard the previous.
+    /// Removes the block with hash `block_hash` and all its descendants from
+    /// the non-finalized state, returning the removed blocks. Inserts the
+    /// truncated chain into the chain_set and discards the previous.
     #[allow(clippy::unwrap_in_result)]
-    pub fn invalidate_block(&mut self, block_hash: Hash) -> Result<block::Hash, InvalidateError> {
+    fn remove_block_and_descendants(
+        &mut self,
+        block_hash: Hash,
+    ) -> Result<Vec<ContextuallyVerifiedBlock>, InvalidateError> {
         let Some(chain) = self.find_chain(|chain| chain.contains_block_hash(block_hash)) else {
             return Err(InvalidateError::BlockNotFound(block_hash));
         };
 
-        let invalidated_blocks = if chain.non_finalized_root_hash() == block_hash {
+        let removed_blocks = if chain.non_finalized_root_hash() == block_hash {
             let chain_tip_hash = chain.non_finalized_tip_hash();
             self.chain_set
                 .retain(|chain| chain.non_finalized_tip_hash() != chain_tip_hash);
             chain.blocks.values().cloned().collect()
         } else {
-            let (new_chain, invalidated_blocks) = chain
+            let (new_chain, removed_blocks) = chain
                 .invalidate_block(block_hash)
                 .expect("already checked that chain contains hash");
 
@@ -394,9 +398,25 @@ impl NonFinalizedState {
                 chain_set.retain(|c| !c.contains_block_hash(block_hash))
             });
 
-            invalidated_blocks
+            removed_blocks
         };
 
+        self.update_metrics_for_chains();
+        self.update_metrics_bars();
+
+        Ok(removed_blocks)
+    }
+
+    /// Invalidate block with hash `block_hash` and all descendants from the non-finalized state. Insert
+    /// the new chain into the chain_set and discard the previous.
+    ///
+    /// The removed blocks are recorded in the invalidated list: they are
+    /// refused on re-delivery until explicitly reconsidered. For a reorg
+    /// rollback, use [`Self::rollback_block`] instead.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn invalidate_block(&mut self, block_hash: Hash) -> Result<block::Hash, InvalidateError> {
+        let invalidated_blocks = self.remove_block_and_descendants(block_hash)?;
+
         // TODO: Allow for invalidating multiple block hashes at a given height (#9552).
         self.invalidated_blocks.insert(
             invalidated_blocks
@@ -411,9 +431,20 @@ impl NonFinalizedState {
             self.invalidated_blocks.shift_remove_index(0);
         }
 
-        self.update_metrics_for_chains();
-        self.update_metrics_bars();
+        Ok(block_hash)
+    }
 
+    /// Rolls back the block with hash `block_hash` and all its descendants
+    /// from the non-finalized state, without recording them as invalidated.
+    ///
+    /// A branch-switch rollback is a chain-selection outcome, not a verdict
+    /// on the blocks: the rolled-back branch must revalidate and recommit
+    /// normally the moment it wins the work race again. Recording it in the
+    /// invalidated list would refuse the re-delivery
+    /// (`BlockPreviouslyInvalidated`) and strand the node if the winning
+    /// branch later evaporates.
+    pub fn rollback_block(&mut self, block_hash: Hash) -> Result<block::Hash, InvalidateError> {
+        self.remove_block_and_descendants(block_hash)?;
         Ok(block_hash)
     }
 

zebra-state/src/service/tests.rs

@@ -35,8 +35,8 @@ use crate::{
         setup::{partial_nu5_chain_strategy, transaction_v4_from_coinbase},
         FakeChainHelper,
     },
-    BoxError, CheckpointVerifiedBlock, Config, ReadRequest, ReadResponse, Request, Response,
-    SemanticallyVerifiedBlock,
+    BoxError, CheckpointVerifiedBlock, Config, HeaderRangeCommitOutcome, ReadRequest, ReadResponse,
+    Request, Response, SemanticallyVerifiedBlock,
 };
 
 const LAST_BLOCK_HEIGHT: u32 = 10;
@@ -538,7 +538,11 @@ async fn header_only_service_requests_preserve_body_boundary() -> std::result::R
                 tree_aux_roots: roots_from_height(Height(1), 2),
             })
             .await?,
-        Response::Committed(block2_hash),
+        Response::CommittedHeaderRange(HeaderRangeCommitOutcome {
+            tip_hash: block2_hash,
+            reorged_at: None,
+            reorged_to_hash: None,
+        }),
     );
 
     assert_eq!(
@@ -793,7 +797,14 @@ async fn commit_header_range_completes_while_in_finalized_write_phase(
     .await
     .expect("CommitHeaderRange must not deadlock while in the finalized write phase")?;
 
-    assert_eq!(committed, Response::Committed(block2_hash));
+    assert_eq!(
+        committed,
+        Response::CommittedHeaderRange(HeaderRangeCommitOutcome {
+            tip_hash: block2_hash,
+            reorged_at: None,
+            reorged_to_hash: None,
+        })
+    );
 
     assert_eq!(
         state