feat(chain)!: Use a more granular state machine for IBD

crates/floresta-chain/src/extensions.rs

@@ -293,6 +293,7 @@ mod tests {
     use crate::BlockConsumer;
     use crate::BlockchainError;
     use crate::UtxoData;
+    use crate::pruned_utreexo::IBDState;
 
     #[derive(Debug)]
     pub enum MockBlockchainError {
@@ -440,6 +441,10 @@ mod tests {
         fn acc(&self) -> Stump {
             unimplemented!()
         }
+
+        fn ibd_state(&self) -> IBDState {
+            unimplemented!()
+        }
     }
 
     fn get_genesis_header() -> Header {

crates/floresta-chain/src/pruned_utreexo/chain_state.rs

@@ -66,6 +66,7 @@ use crate::BestChain;
 use crate::ChainStore;
 use crate::extensions::WorkExt;
 use crate::prelude::*;
+use crate::pruned_utreexo::IBDState;
 use crate::pruned_utreexo::utxo_data::UtxoData;
 use crate::read_lock;
 use crate::write_lock;
@@ -132,8 +133,8 @@ pub struct ChainStateInner<PersistedState: ChainStore> {
     subscribers: Vec<Arc<dyn BlockConsumer>>,
     /// Fee estimation for 1, 10 and 20 blocks
     fee_estimation: (f64, f64, f64),
-    /// Are we in Initial Block Download?
-    ibd: bool,
+    /// What is our current IBD state?
+    ibd: IBDState,
     /// Parameters for the chain and functions that verify the chain.
     consensus: Consensus,
     /// Assume valid is a Core-specific config that tells the node to not validate signatures
@@ -605,7 +606,7 @@ impl<PersistedState: ChainStore> ChainState<PersistedState> {
                 },
                 subscribers: Vec::new(),
                 fee_estimation: (1_f64, 1_f64, 1_f64),
-                ibd: true,
+                ibd: IBDState::HeadersSync,
                 consensus: Consensus { parameters },
                 assume_valid,
             }),
@@ -767,7 +768,7 @@ impl<PersistedState: ChainStore> ChainState<PersistedState> {
             chainstore,
             fee_estimation: (1_f64, 1_f64, 1_f64),
             subscribers: Vec::new(),
-            ibd: true,
+            ibd: IBDState::HeadersSync,
             consensus: Consensus {
                 parameters: network.into(),
             },
@@ -1131,7 +1132,7 @@ impl<PersistedState: ChainStore> BlockchainInterface for ChainState<PersistedSta
     }
 
     fn is_in_ibd(&self) -> bool {
-        self.inner.read().ibd
+        self.inner.read().ibd != IBDState::Done
     }
 
     fn get_block_height(&self, hash: &BlockHash) -> Result<Option<u32>, Self::Error> {
@@ -1231,6 +1232,11 @@ impl<PersistedState: ChainStore> BlockchainInterface for ChainState<PersistedSta
 
         Ok(height + chain_params.coinbase_maturity <= current_height)
     }
+
+    fn ibd_state(&self) -> IBDState {
+        let inner = read_lock!(self);
+        inner.ibd
+    }
 }
 
 impl<PersistedState: ChainStore> UpdatableChainstate for ChainState<PersistedState> {
@@ -1298,9 +1304,9 @@ impl<PersistedState: ChainStore> UpdatableChainstate for ChainState<PersistedSta
         Ok(())
     }
 
-    fn toggle_ibd(&self, is_ibd: bool) {
+    fn update_ibd(&self, ibd_state: IBDState) {
         let mut inner = write_lock!(self);
-        inner.ibd = is_ibd;
+        inner.ibd = ibd_state;
     }
 
     fn connect_block(
@@ -1489,7 +1495,7 @@ impl<T: ChainStore> TryFrom<ChainStateBuilder<T>> for ChainState<T> {
             chainstore: builder.chainstore()?,
             best_block: builder.best_block()?,
             assume_valid: builder.assume_valid(),
-            ibd: builder.ibd(),
+            ibd: builder.ibd_state(),
             subscribers: Vec::new(),
             fee_estimation: (1_f64, 1_f64, 1_f64),
             consensus: Consensus {

crates/floresta-chain/src/pruned_utreexo/chain_state_builder.rs

@@ -25,6 +25,7 @@ use crate::DatabaseError;
 use crate::DiskBlockHeader;
 use crate::prelude::Vec;
 use crate::pruned_utreexo::Box;
+use crate::pruned_utreexo::IBDState;
 
 #[derive(Debug)]
 /// Represents errors that can occur during the construction of a ChainState instance.
@@ -55,7 +56,7 @@ pub struct ChainStateBuilder<PersistedState: ChainStore> {
     chainstore: Option<PersistedState>,
 
     /// Indicates whether the builder is in initial block download mode.
-    ibd: bool,
+    ibd: IBDState,
 
     /// The chain parameters.
     chain_params: Option<ChainParams>,
@@ -76,7 +77,7 @@ impl<T: ChainStore> ChainStateBuilder<T> {
         ChainStateBuilder {
             acc: None,
             chainstore: None,
-            ibd: true,
+            ibd: IBDState::HeadersSync,
             chain_params: None,
             assume_valid: None,
             tip: None,
@@ -115,7 +116,7 @@ impl<T: ChainStore> ChainStateBuilder<T> {
     }
 
     /// Enable or disable Initial Block Download (IBD) mode.
-    pub fn toggle_ibd(mut self, ibd: bool) -> Self {
+    pub fn select_ibd(mut self, ibd: IBDState) -> Self {
         self.ibd = ibd;
         self
     }
@@ -159,11 +160,6 @@ impl<T: ChainStore> ChainStateBuilder<T> {
             .ok_or(BlockchainBuilderError::MissingChainstore)
     }
 
-    /// Returns whether Initial Block Download (IBD) mode is enabled.
-    pub(super) fn ibd(&self) -> bool {
-        self.ibd
-    }
-
     /// Get the chain parameters, returning an error if they haven't been set.
     pub(super) fn chain_params(&self) -> Result<ChainParams, BlockchainBuilderError> {
         self.chain_params
@@ -196,4 +192,9 @@ impl<T: ChainStore> ChainStateBuilder<T> {
     pub(super) fn assume_valid(&self) -> Option<BlockHash> {
         self.assume_valid
     }
+
+    /// Returns the inner [`IBDState`]
+    pub(super) fn ibd_state(&self) -> IBDState {
+        self.ibd
+    }
 }

crates/floresta-chain/src/pruned_utreexo/mod.rs

@@ -41,6 +41,30 @@ use crate::BlockchainError;
 use crate::prelude::*;
 use crate::pruned_utreexo::utxo_data::UtxoData;
 
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Default)]
+/// Our current IBD state, meaning which startup phase are we, if any.
+///
+/// During startup, our node will go through a bootstrap process called Initial Block Download,
+/// where it will catch up with the network. This enum is a simple state machine that represents
+/// which state we are currently in.
+pub enum IBDState {
+    #[default]
+    /// Downloading headers to stablish which is the most-work chain.
+    ///
+    /// During this phase, we only download and check headers. We will finish when we are convinced
+    /// this is the most work chain available.
+    HeadersSync,
+
+    /// Downloading and checking blocks.
+    ///
+    /// After we find the most work chain, we start downloading blocks and connecting them to our
+    /// chain. This step usually takes the longest time.
+    DownloadingBlocks,
+
+    /// We've finished IBD and are now listening for new blocks as they are found.
+    Done,
+}
+
 /// This trait is the main interface between our blockchain backend and other services.
 /// It'll be useful for transitioning from rpc to a p2p based node
 pub trait BlockchainInterface {
@@ -128,6 +152,9 @@ pub trait BlockchainInterface {
 
     /// Returns the amount of [`Work`] associated with a given chain tip
     fn get_work(&self, tip: BlockHash) -> Result<Work, Self::Error>;
+
+    /// Returns the current state of our chain.
+    fn ibd_state(&self) -> IBDState;
 }
 
 /// [UpdatableChainstate] is a contract that a is expected from a chainstate
@@ -158,8 +185,8 @@ pub trait UpdatableChainstate {
     fn handle_transaction(&self) -> Result<(), BlockchainError>;
     /// Persists our data. Should be invoked periodically.
     fn flush(&self) -> Result<(), BlockchainError>;
-    /// Toggle IBD on/off
-    fn toggle_ibd(&self, is_ibd: bool);
+    /// Update IBD state
+    fn update_ibd(&self, ibd_state: IBDState);
     /// Tells this blockchain to consider this block invalid, and not build on top of it
     fn invalidate_block(&self, block: BlockHash) -> Result<(), BlockchainError>;
     /// Marks one block as being fully validated, this overrides a block that was explicitly
@@ -209,8 +236,8 @@ impl<T: UpdatableChainstate> UpdatableChainstate for Arc<T> {
         T::get_acc(self)
     }
 
-    fn toggle_ibd(&self, is_ibd: bool) {
-        T::toggle_ibd(self, is_ibd)
+    fn update_ibd(&self, ibd_state: IBDState) {
+        T::update_ibd(self, ibd_state)
     }
 
     fn connect_block(
@@ -361,6 +388,10 @@ impl<T: BlockchainInterface> BlockchainInterface for Arc<T> {
     fn get_fork_point(&self, block: BlockHash) -> Result<BlockHash, Self::Error> {
         T::get_fork_point(self, block)
     }
+
+    fn ibd_state(&self) -> IBDState {
+        T::ibd_state(self)
+    }
 }
 
 /// This module defines an [UtxoData] struct, helpful for transaction validation

crates/floresta-chain/src/pruned_utreexo/partial_chain.rs

@@ -39,6 +39,7 @@ use super::chainparams::ChainParams;
 use super::consensus::Consensus;
 use super::error::BlockValidationErrors;
 use super::error::BlockchainError;
+use crate::pruned_utreexo::IBDState;
 use crate::pruned_utreexo::utxo_data::UtxoData;
 
 #[doc(hidden)]
@@ -284,8 +285,9 @@ impl UpdatableChainstate for PartialChainState {
         Ok(())
     }
 
-    fn toggle_ibd(&self, _is_ibd: bool) {
-        // no-op: we know if we finished by looking at our current and end height
+    fn update_ibd(&self, _ibd_state: IBDState) {
+        // no-op: we are only used for IBD, so we are always in IBD, and we don't need to update
+        // anything
     }
 
     // these are unimplemented, and will panic if called
@@ -329,6 +331,10 @@ impl UpdatableChainstate for PartialChainState {
 impl BlockchainInterface for PartialChainState {
     type Error = BlockchainError;
 
+    fn ibd_state(&self) -> IBDState {
+        IBDState::DownloadingBlocks
+    }
+
     fn get_params(&self) -> bitcoin::params::Params {
         self.inner().chain_params().params
     }

crates/floresta-wire/src/p2p_wire/node/chain_selector_ctx.rs

@@ -58,6 +58,7 @@ use bitcoin::p2p::ServiceFlags;
 use floresta_chain::ChainBackend;
 use floresta_chain::CompactLeafData;
 use floresta_chain::proof_util;
+use floresta_chain::pruned_utreexo::IBDState;
 use floresta_common::service_flags;
 use rand::rng;
 use rand::seq::IndexedRandom;
@@ -729,7 +730,7 @@ where
 
                 self.context.state = ChainSelectorState::Done;
                 self.chain.mark_chain_as_assumed(acc, tips[0]).unwrap();
-                self.chain.toggle_ibd(false);
+                self.chain.update_ibd(IBDState::Done);
             }
             // if we have more than one tip, we need to check if our best chain has an invalid block
             tips.remove(0); // no need to check our best one
@@ -882,6 +883,7 @@ where
         // We downloaded all headers in the most-pow chain, and all our peers agree
         // this is the most-pow chain, we're done!
         if self.context.state == ChainSelectorState::Done {
+            self.chain.update_ibd(IBDState::DownloadingBlocks);
             try_and_log!(self.chain.flush());
             return Ok(LoopControl::Break);
         }

crates/floresta-wire/src/p2p_wire/node/sync_ctx.rs

@@ -8,6 +8,7 @@ use std::time::Instant;
 use bitcoin::p2p::ServiceFlags;
 use floresta_chain::ThreadSafeChain;
 use floresta_chain::proof_util;
+use floresta_chain::pruned_utreexo::IBDState;
 use floresta_common::service_flags;
 use rand::rng;
 use rand::seq::IteratorRandom;
@@ -242,7 +243,7 @@ where
 
         if validation_index == best_block {
             info!("IBD is finished, switching to normal operation mode");
-            self.chain.toggle_ibd(false);
+            self.chain.update_ibd(IBDState::Done);
             return LoopControl::Break;
         }