feat(btc-adapter): implement a hybrid header cache (#6797)

Implement a hybrid header cache that would persist anchor header and its
ancestors to disk.
This leads to a simpler on-disk cache implementation, since it only has
to persist a single chain of headers.
The downside is that headers since the anchor point would be lost if the
adapter is restarted, in which case it would sync more headers in order
to catch up.

---------

Co-authored-by: mducroux <[email protected]>

Author: ninegua

rs/bitcoin/adapter/benches/e2e.rs

@@ -189,8 +189,12 @@ fn add_800k_block_headers(criterion: &mut Criterion) {
     group.bench_function("add_headers", |bench| {
         let rt = tokio::runtime::Runtime::new().unwrap();
         bench.iter(|| {
-            let blockchain_state =
-                BlockchainState::new(Network::Bitcoin, &MetricsRegistry::default());
+            let blockchain_state = BlockchainState::new(
+                Network::Bitcoin,
+                None,
+                &MetricsRegistry::default(),
+                no_op_logger(),
+            );
             // Headers are processed in chunks of at most MAX_HEADERS_SIZE entries
             for chunk in bitcoin_headers_to_add.chunks(MAX_HEADERS_SIZE) {
                 let (added_headers, error) =

rs/bitcoin/adapter/src/blockchainmanager.rs

@@ -801,7 +801,8 @@ pub mod test {
     where
         BlockchainState<Network>: HeaderValidator<Network>,
     {
-        let blockchain_state = BlockchainState::new(network, &MetricsRegistry::default());
+        let blockchain_state =
+            BlockchainState::new(network, None, &MetricsRegistry::default(), no_op_logger());
         (
             blockchain_state.genesis(),
             BlockchainManager::new(

rs/bitcoin/adapter/src/blockchainstate.rs

@@ -3,10 +3,7 @@
 use crate::common::HeaderValidator;
 use crate::{
     common::{BlockHeight, BlockchainBlock, BlockchainHeader, BlockchainNetwork},
-    header_cache::{
-        AddHeaderCacheError, AddHeaderResult, HeaderCache, HeaderNode, InMemoryHeaderCache,
-        LMDBHeaderCache, Tip,
-    },
+    header_cache::{AddHeaderCacheError, AddHeaderResult, HeaderNode, HybridHeaderCache, Tip},
     metrics::BlockchainStateMetrics,
 };
 use bitcoin::{BlockHash, block::Header, consensus::Encodable, dogecoin::Header as DogecoinHeader};
@@ -60,7 +57,7 @@ pub type SerializedBlock = Vec<u8>;
 /// The BlockChainState also maintains the child relationhips between the headers.
 pub struct BlockchainState<Network: BlockchainNetwork> {
     /// This field stores all the Bitcoin headers using a HashMap containing BlockHash and the corresponding header.
-    header_cache: Arc<dyn HeaderCache<Header = Network::Header> + Send>,
+    pub(crate) header_cache: Arc<HybridHeaderCache<Network::Header>>,
 
     /// This field stores a hashmap containing BlockHash and the corresponding SerializedBlock.
     block_cache: RwLock<HashMap<BlockHash, Arc<SerializedBlock>>>,
@@ -99,28 +96,15 @@ impl<Network: BlockchainNetwork> BlockchainState<Network>
 where
     Network::Header: Send + Sync,
 {
-    /// Create a new BlockChainState object with in-memory cache.
-    pub fn new(network: Network, metrics_registry: &MetricsRegistry) -> Self {
-        let genesis_block_header = network.genesis_block_header();
-        let header_cache = Arc::new(InMemoryHeaderCache::new(genesis_block_header));
-        let block_cache = RwLock::new(HashMap::new());
-        BlockchainState {
-            header_cache,
-            block_cache,
-            network,
-            metrics: BlockchainStateMetrics::new(metrics_registry),
-        }
-    }
-
-    /// Create a new BlockChainState with on-disk cache.
-    pub fn new_with_cache_dir(
+    /// Create a new BlockChainState with an optional on-disk cache if cache_dir is specified.
+    pub fn new(
         network: Network,
-        cache_dir: PathBuf,
+        cache_dir: Option<PathBuf>,
         metrics_registry: &MetricsRegistry,
         logger: ReplicaLogger,
     ) -> Self {
         let genesis_block_header = network.genesis_block_header();
-        let header_cache = Arc::new(LMDBHeaderCache::new(
+        let header_cache = Arc::new(HybridHeaderCache::new(
             genesis_block_header,
             cache_dir,
             logger,
@@ -387,17 +371,22 @@ mod test {
     use std::collections::HashSet;
 
     fn run_in_memory<R>(network: Network, test_fn: impl Fn(BlockchainState<Network>) -> R) -> R {
-        test_fn(BlockchainState::new(network, &MetricsRegistry::default()))
+        test_fn(BlockchainState::new(
+            network,
+            None,
+            &MetricsRegistry::default(),
+            no_op_logger(),
+        ))
     }
 
     fn run_with_cache_dir<R>(
         network: Network,
         test_fn: impl Fn(BlockchainState<Network>) -> R,
     ) -> R {
         let dir = tempdir().unwrap();
-        test_fn(BlockchainState::new_with_cache_dir(
+        test_fn(BlockchainState::new(
             network,
-            dir.path().to_path_buf(),
+            Some(dir.path().to_path_buf()),
             &MetricsRegistry::default(),
             no_op_logger(),
         ))
@@ -555,7 +544,7 @@ mod test {
             "unsuccessfully added fork chain: {maybe_err:?}"
         );
 
-        let mut tips = state.header_cache.get_tips();
+        let mut tips = crate::header_cache::test::get_tips(&state.header_cache);
         tips.sort_by(|x, y| y.work.cmp(&x.work));
         assert_eq!(tips.len(), 2);
         assert_eq!(tips[0].header.block_hash(), *last_fork_hash);

rs/bitcoin/adapter/src/get_successors_handler.rs

@@ -1,6 +1,6 @@
 use std::{
     collections::{HashSet, VecDeque},
-    sync::Arc,
+    sync::{Arc, Mutex},
 };
 
 use bitcoin::{BlockHash, consensus::Encodable};
@@ -90,9 +90,10 @@ pub struct GetSuccessorsHandler<Network: BlockchainNetwork> {
     blockchain_manager_tx: Sender<BlockchainManagerRequest>,
     network: Network,
     metrics: GetSuccessorMetrics,
+    pruning_task_handle: Mutex<Option<tokio::task::JoinHandle<()>>>,
 }
 
-impl<Network: BlockchainNetwork> GetSuccessorsHandler<Network> {
+impl<Network: BlockchainNetwork + Send + Sync> GetSuccessorsHandler<Network> {
     /// Creates a GetSuccessorsHandler to be used to access the blockchain state
     /// inside of the adapter when a `GetSuccessorsRequest` is received.
     pub fn new(
@@ -106,6 +107,7 @@ impl<Network: BlockchainNetwork> GetSuccessorsHandler<Network> {
             blockchain_manager_tx,
             network,
             metrics: GetSuccessorMetrics::new(metrics_registry),
+            pruning_task_handle: Mutex::new(None),
         }
     }
 
@@ -123,6 +125,21 @@ impl<Network: BlockchainNetwork> GetSuccessorsHandler<Network> {
             .processed_block_hashes
             .observe(request.processed_block_hashes.len() as f64);
 
+        // Spawn persist-to-disk task without waiting for it to finish, and make sure there
+        // is only one task running at a time.
+        let cache = self.state.header_cache.clone();
+        let mut handle = self.pruning_task_handle.lock().unwrap();
+        let is_finished = handle
+            .as_ref()
+            .map(|handle| handle.is_finished())
+            .unwrap_or(true);
+        if is_finished {
+            *handle = Some(tokio::task::spawn_blocking(move || {
+                // Error is ignored, since it is a background task
+                let _ = cache.persist_and_prune_headers_below_anchor(request.anchor);
+            }));
+        }
+
         let (blocks, next, obsolete_blocks) = {
             let anchor_height = self
                 .state
@@ -312,6 +329,7 @@ mod test {
     use super::*;
 
     use bitcoin::{Block, consensus::Decodable};
+    use ic_logger::no_op_logger;
     use ic_metrics::MetricsRegistry;
     use tokio::sync::mpsc::channel;
 
@@ -323,12 +341,18 @@ mod test {
         Block::consensus_decode(&mut (*serialized_block).as_slice()).unwrap()
     }
 
+    fn new_blockchain_state<Network: BlockchainNetwork>(
+        network: Network,
+    ) -> BlockchainState<Network> {
+        BlockchainState::new(network, None, &MetricsRegistry::default(), no_op_logger())
+    }
+
     /// This tests ensures that `BlockchainManager::get_successors(...)` will return relevant blocks
     /// with the next headers of many forks and enqueue missing block hashes.
     #[tokio::test]
     async fn test_get_successors() {
         let network = bitcoin::Network::Regtest;
-        let blockchain_state = BlockchainState::new(network, &MetricsRegistry::default());
+        let blockchain_state = new_blockchain_state(network);
         let genesis = blockchain_state.genesis();
         let genesis_hash = genesis.block_hash();
         let (blockchain_manager_tx, _blockchain_manager_rx) = channel(10);
@@ -440,7 +464,7 @@ mod test {
     #[tokio::test]
     async fn test_get_successors_wait_header_sync_regtest() {
         let network = bitcoin::Network::Regtest;
-        let blockchain_state = BlockchainState::new(network, &MetricsRegistry::default());
+        let blockchain_state = new_blockchain_state(network);
         let genesis = blockchain_state.genesis();
         let genesis_hash = genesis.block_hash();
         let (blockchain_manager_tx, _blockchain_manager_rx) = channel(10);
@@ -490,7 +514,7 @@ mod test {
     #[tokio::test]
     async fn test_get_successors_multiple_blocks() {
         let network = bitcoin::Network::Regtest;
-        let blockchain_state = BlockchainState::new(network, &MetricsRegistry::default());
+        let blockchain_state = new_blockchain_state(network);
         let genesis = blockchain_state.genesis();
         let genesis_hash = genesis.block_hash();
         let (blockchain_manager_tx, _blockchain_manager_rx) = channel(10);
@@ -564,7 +588,7 @@ mod test {
     #[tokio::test]
     async fn test_get_successors_max_num_blocks() {
         let network = bitcoin::Network::Regtest;
-        let blockchain_state = BlockchainState::new(network, &MetricsRegistry::default());
+        let blockchain_state = new_blockchain_state(network);
         let genesis = blockchain_state.genesis();
         let genesis_hash = genesis.block_hash();
         let (blockchain_manager_tx, _blockchain_manager_rx) = channel(10);
@@ -600,7 +624,7 @@ mod test {
     #[tokio::test]
     async fn test_get_successors_multiple_blocks_out_of_order() {
         let network = bitcoin::Network::Regtest;
-        let blockchain_state = BlockchainState::new(network, &MetricsRegistry::default());
+        let blockchain_state = new_blockchain_state(network);
         let genesis = blockchain_state.genesis();
         let genesis_hash = genesis.block_hash();
         let (blockchain_manager_tx, _blockchain_manager_rx) = channel(10);
@@ -694,7 +718,7 @@ mod test {
     #[tokio::test]
     async fn test_get_successors_large_block() {
         let network = bitcoin::Network::Regtest;
-        let blockchain_state = BlockchainState::new(network, &MetricsRegistry::default());
+        let blockchain_state = new_blockchain_state(network);
         let genesis = blockchain_state.genesis();
         let genesis_hash = genesis.block_hash();
         let (blockchain_manager_tx, _blockchain_manager_rx) = channel(10);
@@ -760,7 +784,7 @@ mod test {
     #[tokio::test]
     async fn test_get_successors_many_blocks_until_size_cap_is_met() {
         let network = bitcoin::Network::Regtest;
-        let blockchain_state = BlockchainState::new(network, &MetricsRegistry::default());
+        let blockchain_state = new_blockchain_state(network);
         let genesis = blockchain_state.genesis();
         let genesis_hash = genesis.block_hash();
         let (blockchain_manager_tx, _blockchain_manager_rx) = channel(10);