feat(benches): add benchmarks for varied forms of db lookups

Author: rymnc

CHANGELOG.md

@@ -6,6 +6,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ## [Unreleased]
 
+### Added
+- [2142](https://github.com/FuelLabs/fuel-core/pull/2142): Added benchmarks for varied forms of db lookups to assist in optimizations.
+
 ## [Version 0.35.0]
 
 ### Added

Cargo.lock

@@ -34,6 +34,7 @@ p256 = { version = "0.13", default-features = false, features = [
   "digest",
   "ecdsa",
 ] }
+postcard = { workspace = true }
 primitive-types = { workspace = true, default-features = false }
 quanta = "0.12"
 rand = { workspace = true }
@@ -66,3 +67,7 @@ name = "block_target_gas"
 [[bench]]
 harness = false
 name = "transaction_throughput"
+
+[[bench]]
+harness = false
+name = "db_lookup_times"

benches/Cargo.toml

@@ -0,0 +1,93 @@
+use crate::db_lookup_times_utils::{
+    matrix::matrix,
+    utils::{
+        get_full_block,
+        get_random_block_height,
+        multi_get_block,
+        open_db,
+        open_raw_rocksdb,
+    },
+};
+use criterion::{
+    criterion_group,
+    criterion_main,
+    Criterion,
+};
+use db_lookup_times_utils::seed::{
+    seed_compressed_blocks_and_transactions_matrix,
+    seed_full_block_matrix,
+};
+use fuel_core_storage::transactional::AtomicView;
+
+mod db_lookup_times_utils;
+
+pub fn header_and_tx_lookup(c: &mut Criterion) {
+    let method = "header_and_tx";
+
+    seed_compressed_blocks_and_transactions_matrix(method);
+    let mut group = c.benchmark_group(method);
+
+    for (block_count, tx_count) in matrix() {
+        let database = open_db(block_count, tx_count, method);
+        let height = get_random_block_height(block_count);
+        let view = database.latest_view().unwrap();
+        group.bench_function(format!("{block_count}/{tx_count}"), |b| {
+            b.iter(|| {
+                let block = (&view).get_full_block(&height);
+                assert!(block.is_ok());
+                assert!(block.unwrap().is_some());
+            });
+        });
+    }
+
+    group.finish();
+}
+
+pub fn multi_get_lookup(c: &mut Criterion) {
+    let method = "multi_get";
+
+    seed_compressed_blocks_and_transactions_matrix(method);
+    let mut group = c.benchmark_group(method);
+
+    for (block_count, tx_count) in matrix() {
+        let database = open_raw_rocksdb(block_count, tx_count, method);
+        let height = get_random_block_height(block_count);
+        group.bench_function(format!("{block_count}/{tx_count}"), |b| {
+            b.iter(|| {
+                // todo: clean up
+                multi_get_block(&database, height);
+            });
+        });
+    }
+
+    group.finish();
+}
+
+pub fn full_block_lookup(c: &mut Criterion) {
+    let method = "full_block";
+
+    seed_full_block_matrix();
+    let mut group = c.benchmark_group(method);
+
+    for (block_count, tx_count) in matrix() {
+        let database = open_db(block_count, tx_count, method);
+        let height = get_random_block_height(block_count);
+        let view = database.latest_view().unwrap();
+        group.bench_function(format!("{block_count}/{tx_count}"), |b| {
+            b.iter(|| {
+                let full_block = get_full_block(&view, &height);
+                assert!(full_block.is_ok());
+                assert!(full_block.unwrap().is_some());
+            });
+        });
+    }
+
+    group.finish();
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default().sample_size(10).measurement_time(std::time::Duration::from_secs(10));
+    targets = header_and_tx_lookup, multi_get_lookup, full_block_lookup
+}
+criterion_main!(benches);

benches/benches/db_lookup_times.rs

@@ -0,0 +1,13 @@
+pub const BLOCK_COUNT_MATRIX: [u32; 1] = [1];
+pub const TX_COUNT_MATRIX: [u32; 1] = [1];
+
+// todo: we can make this lazy loaded
+pub fn matrix() -> Box<dyn Iterator<Item = (u32, u32)>> {
+    let iter = BLOCK_COUNT_MATRIX.iter().flat_map(|&block_count| {
+        TX_COUNT_MATRIX
+            .iter()
+            .map(move |&tx_count| (block_count, tx_count))
+    });
+
+    Box::new(iter)
+}

benches/benches/db_lookup_times_utils/matrix.rs

@@ -0,0 +1,3 @@
+pub mod matrix;
+pub mod seed;
+pub mod utils;

benches/benches/db_lookup_times_utils/mod.rs

@@ -0,0 +1,155 @@
+use crate::db_lookup_times_utils::{
+    matrix::matrix,
+    utils::open_raw_rocksdb,
+};
+use fuel_core::{
+    database::database_description::on_chain::OnChain,
+    state::rocks_db::RocksDb,
+};
+use fuel_core_storage::{
+    column::Column,
+    kv_store::{
+        KeyValueMutate,
+        Value,
+    },
+};
+use fuel_core_types::{
+    blockchain::{
+        block::{
+            Block,
+            PartialFuelBlock,
+        },
+        header::{
+            ConsensusHeader,
+            PartialBlockHeader,
+        },
+        primitives::Empty,
+    },
+    fuel_tx::{
+        Transaction,
+        UniqueIdentifier,
+    },
+    fuel_types::{
+        BlockHeight,
+        ChainId,
+    },
+};
+
+fn generate_bench_block(height: u32, tx_count: u32) -> Block {
+    let header = PartialBlockHeader {
+        application: Default::default(),
+        consensus: ConsensusHeader::<Empty> {
+            height: BlockHeight::from(height),
+            ..Default::default()
+        },
+    };
+    let txes = generate_bench_transactions(tx_count);
+    let block = PartialFuelBlock::new(header, txes);
+    block.generate(&[], Default::default()).unwrap()
+}
+
+fn generate_bench_transactions(tx_count: u32) -> Vec<Transaction> {
+    let mut txes = vec![];
+    for _ in 0..tx_count {
+        txes.push(Transaction::default_test_tx());
+    }
+    txes
+}
+
+fn height_key(block_height: u32) -> Vec<u8> {
+    BlockHeight::from(block_height).to_bytes().to_vec()
+}
+
+fn insert_compressed_block(
+    database: &mut RocksDb<OnChain>,
+    height: u32,
+    tx_count: u32,
+) -> Block {
+    let block = generate_bench_block(height, tx_count);
+
+    let compressed_block = block.compress(&ChainId::default());
+    let height_key = height_key(height);
+
+    let raw_compressed_block = postcard::to_allocvec(&compressed_block).unwrap().to_vec();
+    let raw_transactions = block
+        .transactions()
+        .into_iter()
+        .map(|tx| {
+            (
+                tx.id(&ChainId::default()),
+                postcard::to_allocvec(tx).unwrap().to_vec(),
+            )
+        })
+        .collect::<Vec<_>>();
+
+    // 1. insert into CompressedBlocks table
+    database
+        .put(
+            height_key.as_slice(),
+            Column::FuelBlocks,
+            Value::new(raw_compressed_block),
+        )
+        .unwrap();
+    // 2. insert into Transactions table
+    for (tx_id, tx) in raw_transactions {
+        database
+            .put(tx_id.as_slice(), Column::Transactions, Value::new(tx))
+            .unwrap();
+    }
+
+    block
+}
+
+fn insert_full_block(database: &mut RocksDb<OnChain>, height: u32, tx_count: u32) {
+    let block = match insert_compressed_block(database, height, tx_count) {
+        Block::V1(b) => b,
+    };
+
+    let height_key = height_key(height);
+    let raw_full_block = postcard::to_allocvec(&block).unwrap().to_vec();
+
+    // 3. insert into FullFuelBlocks table
+    database
+        .put(
+            height_key.as_slice(),
+            Column::FullFuelBlocks,
+            Value::new(raw_full_block),
+        )
+        .unwrap();
+}
+
+fn seed_compressed_blocks_and_transactions(
+    database: &mut RocksDb<OnChain>,
+    block_count: u32,
+    tx_count: u32,
+) -> Vec<Block> {
+    let mut blocks = vec![];
+    for block_number in 0..block_count {
+        blocks.push(insert_compressed_block(database, block_number, tx_count));
+    }
+    blocks
+}
+
+pub fn seed_compressed_blocks_and_transactions_matrix(method: &str) {
+    for (block_count, tx_count) in matrix() {
+        let mut database = open_raw_rocksdb(block_count, tx_count, method);
+        let _ =
+            seed_compressed_blocks_and_transactions(&mut database, block_count, tx_count);
+    }
+}
+
+fn seed_full_blocks(database: &mut RocksDb<OnChain>, block_count: u32, tx_count: u32) {
+    // we seed compressed blocks and transactions to not affect individual
+    // lookup times
+
+    for block_number in 0..block_count {
+        insert_full_block(database, block_number, tx_count);
+    }
+}
+
+pub fn seed_full_block_matrix() {
+    for (block_count, tx_count) in matrix() {
+        let mut database = open_raw_rocksdb(block_count, tx_count, "full_block");
+        seed_full_blocks(&mut database, block_count, tx_count);
+    }
+}

benches/benches/db_lookup_times_utils/seed.rs

@@ -0,0 +1,89 @@
+use fuel_core::{
+    database::{
+        database_description::on_chain::OnChain,
+        Database,
+    },
+    state::{
+        historical_rocksdb::StateRewindPolicy,
+        rocks_db::RocksDb,
+        IterableKeyValueView,
+    },
+};
+use fuel_core_storage::{
+    column::Column,
+    kv_store::{
+        KeyValueInspect,
+        StorageColumn,
+    },
+    tables::FullFuelBlocks,
+    StorageAsRef,
+};
+use fuel_core_types::{
+    blockchain::block::{
+        Block,
+        CompressedBlock,
+    },
+    fuel_tx::Transaction,
+    fuel_types::BlockHeight,
+};
+use rand::Rng;
+use std::{
+    borrow::Cow,
+    path::Path,
+};
+
+pub fn get_random_block_height(block_count: u32) -> BlockHeight {
+    BlockHeight::from(rand::thread_rng().gen_range(0..block_count))
+}
+
+pub fn open_db(block_count: u32, tx_count: u32, method: &str) -> Database {
+    Database::open_rocksdb(
+        Path::new(format!("./{block_count}/{method}/{tx_count}").as_str()),
+        None, // no caching
+        StateRewindPolicy::NoRewind,
+    )
+    .unwrap()
+}
+
+pub fn open_raw_rocksdb(
+    block_count: u32,
+    tx_count: u32,
+    method: &str,
+) -> RocksDb<OnChain> {
+    RocksDb::default_open(
+        Path::new(format!("./{block_count}/{method}/{tx_count}").as_str()),
+        None,
+    )
+    .unwrap()
+}
+
+pub fn get_full_block(
+    view: &IterableKeyValueView<Column>,
+    height: &BlockHeight,
+) -> Result<Option<Block>, anyhow::Error> {
+    let db_block = view.storage::<FullFuelBlocks>().get(height)?;
+    if let Some(block) = db_block {
+        Ok(Some(Block::V1(Cow::into_owned(block))))
+    } else {
+        Ok(None)
+    }
+}
+
+pub fn multi_get_block(database: &RocksDb<OnChain>, height: BlockHeight) {
+    let height_key = height.to_bytes();
+
+    let raw_block = database
+        .get(&height_key, Column::FuelBlocks)
+        .unwrap()
+        .unwrap();
+    let block: CompressedBlock = postcard::from_bytes(raw_block.as_slice()).unwrap();
+    let tx_ids = block.transactions().into_iter();
+    let raw_txs = database
+        .multi_get(Column::Transactions.id(), tx_ids)
+        .unwrap();
+    for raw_tx in raw_txs {
+        if let Some(tx) = raw_tx {
+            let _: Transaction = postcard::from_bytes(tx.as_slice()).unwrap();
+        }
+    }
+}