Skip to content

Commit

Permalink
Merge branch 'stefan/1525-get_successors_optimisation' into 'master'
Browse files Browse the repository at this point in the history 
perf[NET-1525]: bitcoin adapter get_successors optimisation

According to the Criterion benchmark, the change speeds up get_successors calls from ~18ms to ~8ms. 

See merge request dfinity-lab/public/ic!14426
  • Loading branch information
@stefanneamtu
stefanneamtu committed Sep 13, 2023
2 parents 5dbf937 + 7529b65 commit 7b1cede
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 104 deletions.
123 changes: 33 additions & 90 deletions rs/bitcoin/adapter/src/blockchainstate.rs
View file
Expand Up @@ -4,15 +4,12 @@ use crate::{common::BlockHeight, config::Config, metrics::BlockchainStateMetrics
use bitcoin::{blockdata::constants::genesis_block, Block, BlockHash, BlockHeader, Network};
use ic_btc_validation::{validate_header, HeaderStore, ValidateHeaderError};
use ic_metrics::MetricsRegistry;
use parking_lot::Mutex;
use std::{collections::HashMap, sync::Arc};
use std::collections::HashMap;
use thiserror::Error;

/// This field contains the datatype used to store "work" of a Bitcoin blockchain
pub type Work = bitcoin::util::uint::Uint256;

pub type CachedHeader = Arc<HeaderNode>;

/// Contains the necessary information about a tip.
#[derive(Debug, Clone)]
pub struct Tip {
Expand All @@ -25,70 +22,18 @@ pub struct Tip {
pub work: Work,
}

/// A possible error that the header cache may raise.
#[derive(Debug, Error)]
enum HeaderCacheError {
/// This variant is used when the predecessor of the input header is not part of header cache.
#[error("Received a block header where we do not have the previous header in the cache: {0}")]
NotFound(BlockHash),
/// This variant is used when the header is already part of header cache.
#[error("Received a block header that already exists in the cache")]
AlreadyExists,
}

/// Used to maintain the discovered headers from peers.
#[derive(Debug)]
struct HeaderCache(HashMap<BlockHash, CachedHeader>);

impl HeaderCache {
/// Creates a new `HeaderCache` with a set genesis header determined by the
/// provided network.
fn new(genesis_block_header: BlockHeader) -> Self {
let cached_header = Arc::new(HeaderNode {
header: genesis_block_header,
height: 0,
work: genesis_block_header.work(),
children: Mutex::new(vec![]),
});
let mut headers = HashMap::new();
headers.insert(genesis_block_header.block_hash(), cached_header);
Self(headers)
}

/// Retrieves a cached header entry from internal HashMap. If not found,
/// returns None.
fn get(&self, hash: &BlockHash) -> Option<&CachedHeader> {
self.0.get(hash)
}

/// Adds a header to the internal HashMap if it does not already exist in the
/// map and the previous header's hash is a key in the HashMap.
fn insert(&mut self, header: BlockHeader) -> Result<(), HeaderCacheError> {
let block_hash = header.block_hash();
if self.0.contains_key(&block_hash) {
return Err(HeaderCacheError::AlreadyExists);
}

let cached_header = {
let parent = self
.get(&header.prev_blockhash)
.ok_or(HeaderCacheError::NotFound(header.prev_blockhash))?;
let next_height = parent.height + 1;
let next_work = parent.work + header.work();
let cached_header = Arc::new(HeaderNode {
header,
height: next_height,
work: next_work,
children: Mutex::new(vec![]),
});
parent.children.lock().push(cached_header.clone());
cached_header
};

self.0.insert(block_hash, cached_header);

Ok(())
}
/// Creates a new cache with a set genesis header determined by the
/// provided network.
fn init_cache_with_genesis(genesis_block_header: BlockHeader) -> HashMap<BlockHash, HeaderNode> {
let cached_header = HeaderNode {
header: genesis_block_header,
height: 0,
work: genesis_block_header.work(),
children: vec![],
};
let mut headers = HashMap::new();
headers.insert(genesis_block_header.block_hash(), cached_header);
headers
}

/// This struct stores a BlockHeader along with its height in the Bitcoin Blockchain.
Expand All @@ -102,7 +47,7 @@ pub struct HeaderNode {
/// That is, this field is the sum of work of the above header and all its ancestors.
pub work: Work,
/// This field contains this node's successor headers.
pub children: Mutex<Vec<CachedHeader>>,
pub children: Vec<BlockHash>,
}

/// The result when `BlockchainState::add_header(...)` is called.
Expand Down Expand Up @@ -143,8 +88,8 @@ pub struct BlockchainState {
/// The starting point of the blockchain
genesis_block_header: BlockHeader,

/// This field stores all the Bitcoin headers using a HashMap containing BlockHash and the corresponding header.
header_cache: HeaderCache,
/// This field stores all the Bitcoin headers using a HashMap containining BlockHash and the corresponding header.
header_cache: HashMap<BlockHash, HeaderNode>,

/// This field stores a hashmap containing BlockHash and the corresponding Block.
block_cache: HashMap<BlockHash, Block>,
Expand All @@ -162,7 +107,7 @@ impl BlockchainState {
pub fn new(config: &Config, metrics_registry: &MetricsRegistry) -> Self {
// Create a header cache and inserting dummy header corresponding the `adapter_genesis_hash`.
let genesis_block_header = genesis_block(config.network).header;
let header_cache = HeaderCache::new(genesis_block_header);
let header_cache = init_cache_with_genesis(genesis_block_header);
let block_cache = HashMap::new();
let tips = vec![Tip {
header: genesis_block_header,
Expand All @@ -186,7 +131,7 @@ impl BlockchainState {
}

/// Returns the header for the given block hash.
pub fn get_cached_header(&self, hash: &BlockHash) -> Option<&CachedHeader> {
pub fn get_cached_header(&self, hash: &BlockHash) -> Option<&HeaderNode> {
self.header_cache.get(hash)
}

Expand Down Expand Up @@ -238,24 +183,22 @@ impl BlockchainState {
return Err(AddHeaderError::InvalidHeader(block_hash, err));
}

let prev_hash = header.prev_blockhash;
let cached_header = match self.header_cache.insert(header) {
Ok(_) => self
.header_cache
.get(&block_hash)
.expect("Header should exist in the cache at this point."),
Err(err) => match err {
HeaderCacheError::NotFound(prev_hash) => {
return Err(AddHeaderError::PrevHeaderNotCached(prev_hash));
}
HeaderCacheError::AlreadyExists => {
return Ok(AddHeaderResult::HeaderAlreadyExists(block_hash));
}
},
let parent = self
.header_cache
.get_mut(&header.prev_blockhash)
.ok_or(AddHeaderError::PrevHeaderNotCached(header.prev_blockhash))?;

let cached_header = HeaderNode {
header,
height: parent.height + 1,
work: parent.work + header.work(),
children: vec![],
};
parent.children.push(header.block_hash());

// Update the tip headers.
// If the previous header already exists in `tips`, then update it with the new tip.
let prev_hash = header.prev_blockhash;
let maybe_cached_header_idx = self
.tips
.iter()
Expand All @@ -276,10 +219,10 @@ impl BlockchainState {
}
};

self.header_cache.insert(block_hash, cached_header);

self.metrics.header_cache_size.inc();
Ok(AddHeaderResult::HeaderAdded(
cached_header.header.block_hash(),
))
Ok(AddHeaderResult::HeaderAdded(header.block_hash()))
}

/// This method adds a new block to the `block_cache`
Expand Down
33 changes: 19 additions & 14 deletions rs/bitcoin/adapter/src/get_successors_handler.rs
View file
Expand Up @@ -10,8 +10,8 @@ use tokio::sync::{mpsc::Sender, Mutex};
use tonic::{Code, Status};

use crate::{
blockchainstate::CachedHeader, common::BlockHeight, config::Config,
metrics::GetSuccessorMetrics, BlockchainManagerRequest, BlockchainState,
common::BlockHeight, config::Config, metrics::GetSuccessorMetrics, BlockchainManagerRequest,
BlockchainState,
};

// Max size of the `GetSuccessorsResponse` message.
Expand Down Expand Up @@ -188,16 +188,15 @@ fn get_successor_blocks(
let mut successor_blocks = vec![];
// Block hashes that should be looked at in subsequent breadth-first searches.
let mut response_block_size: usize = 0;
let mut queue: VecDeque<CachedHeader> = state
let mut queue: VecDeque<BlockHash> = state
.get_cached_header(anchor)
.map(|c| c.children.lock().clone())
.map(|c| c.children.clone())
.unwrap_or_default()
.into_iter()
.collect();

// Compute the blocks by starting a breadth-first search.
while let Some(cached_header) = queue.pop_front() {
let block_hash = cached_header.header.block_hash();
while let Some(block_hash) = queue.pop_front() {
if !seen.contains(&block_hash) {
// Retrieve the block from the cache.
match state.get_block(&block_hash) {
Expand All @@ -222,7 +221,12 @@ fn get_successor_blocks(
}
}

queue.extend(cached_header.children.lock().clone());
queue.extend(
state
.get_cached_header(&block_hash)
.map(|header| header.children.clone())
.unwrap_or_default(),
);
}

successor_blocks
Expand All @@ -240,23 +244,24 @@ fn get_next_headers(
.copied()
.chain(blocks.iter().map(|b| b.block_hash()))
.collect();
let mut queue: VecDeque<CachedHeader> = state
let mut queue: VecDeque<BlockHash> = state
.get_cached_header(anchor)
.map(|c| c.children.lock().clone())
.map(|c| c.children.clone())
.unwrap_or_default()
.into_iter()
.collect();
let mut next_headers = vec![];
while let Some(cached_header) = queue.pop_front() {
while let Some(block_hash) = queue.pop_front() {
if next_headers.len() >= MAX_NEXT_BLOCK_HEADERS_LENGTH {
break;
}

let block_hash = cached_header.header.block_hash();
if !seen.contains(&block_hash) {
next_headers.push(cached_header.header);
if let Some(header_node) = state.get_cached_header(&block_hash) {
if !seen.contains(&block_hash) {
next_headers.push(header_node.header);
}
queue.extend(header_node.children.clone());
}
queue.extend(cached_header.children.lock().clone());
}
next_headers
}
Expand Down

0 comments on commit 7b1cede

Please sign in to comment.