From 4e7d3b357cf3b7edec95ea2859df63cf14459a35 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5rten=20Blankfors?= <marten@blankfors.se>
Date: Sat, 9 Nov 2024 15:19:29 +0700
Subject: [PATCH] feat: simple bench target for measuring coin batch lookup
 times

---
 Cargo.lock                                |   1 +
 benches/Cargo.toml                        |   5 +
 benches/benches/end_to_end_query_times.rs | 145 ++++++++++++++++++++++
 3 files changed, 151 insertions(+)
 create mode 100644 benches/benches/end_to_end_query_times.rs

diff --git a/Cargo.lock b/Cargo.lock
index 67a97f72c3e..33c27003ca7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3263,6 +3263,7 @@ dependencies = [
  "ethnum",
  "fuel-core",
  "fuel-core-chain-config",
+ "fuel-core-client",
  "fuel-core-database",
  "fuel-core-services",
  "fuel-core-storage",
diff --git a/benches/Cargo.toml b/benches/Cargo.toml
index e4ee7db927b..d2546fd5876 100644
--- a/benches/Cargo.toml
+++ b/benches/Cargo.toml
@@ -23,6 +23,7 @@ fuel-core = { path = "../crates/fuel-core", default-features = false, features =
   "rocksdb-production",
 ] }
 fuel-core-chain-config = { workspace = true }
+fuel-core-client = { path = "./../crates/client" }
 fuel-core-database = { path = "./../crates/database" }
 fuel-core-services = { path = "./../crates/services" }
 fuel-core-storage = { path = "./../crates/storage", features = ["smt"] }
@@ -76,3 +77,7 @@ name = "transaction_throughput"
 [[bench]]
 harness = false
 name = "db_lookup_times"
+
+[[bench]]
+harness = false
+name = "end_to_end_query_times"
diff --git a/benches/benches/end_to_end_query_times.rs b/benches/benches/end_to_end_query_times.rs
new file mode 100644
index 00000000000..703518cce29
--- /dev/null
+++ b/benches/benches/end_to_end_query_times.rs
@@ -0,0 +1,145 @@
+use fuel_core::{
+    combined_database::CombinedDatabase,
+    fuel_core_graphql_api::database::ReadDatabase,
+};
+
+use fuel_core_chain_config::Randomize;
+use fuel_core_storage::{
+    tables::{
+        Coins,
+        FuelBlocks,
+    },
+    transactional::WriteTransaction,
+    StorageAsMut,
+};
+use fuel_core_types::{
+    blockchain::block::CompressedBlock,
+    entities::coins::coin::CompressedCoin,
+    fuel_tx::{
+        Address,
+        AssetId,
+        UtxoId,
+    },
+    fuel_types::BlockHeight,
+};
+use rand::{
+    rngs::StdRng,
+    SeedableRng,
+};
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    println!("Setting up bench harness.");
+    let mut harness = Harness::new(StdRng::seed_from_u64(2322)).await?;
+
+    println!("Populating storage with transactions.");
+    let utxo_ids = harness.populate_database().await?;
+
+    println!("Querying transactions from storage.");
+    let total = harness.query_database_many_times(&utxo_ids).await?;
+
+    println!("Queried a total of {total} coins");
+
+    Ok(())
+}
+
+struct Harness<Rng> {
+    rng: Rng,
+    params: Parameters,
+    db: CombinedDatabase,
+}
+
+impl<Rng: rand::RngCore + rand::CryptoRng + Send> Harness<Rng> {
+    async fn new(rng: Rng) -> anyhow::Result<Self> {
+        let params = Parameters::hard_coded();
+        let db = CombinedDatabase::default();
+
+        Ok(Self { rng, params, db })
+    }
+
+    async fn populate_database(&mut self) -> anyhow::Result<Vec<UtxoId>> {
+        let mut utxo_ids = Vec::new();
+
+        let mut coins: Vec<_> = (0..(self.params.utxo_count_per_block
+            * self.params.num_blocks))
+            .map(|_| self.generate_random_coin())
+            .collect();
+
+        for block_height in 0..self.params.num_blocks {
+            let block_height = BlockHeight::from(block_height as u32);
+            let mut compressed_block = CompressedBlock::default();
+            compressed_block.header_mut().set_block_height(block_height);
+
+            let mut transaction = self.db.on_chain_mut().write_transaction();
+
+            transaction
+                .storage::<FuelBlocks>()
+                .insert(&block_height, &compressed_block)
+                .unwrap();
+
+            for _ in 0..self.params.utxo_count_per_block {
+                let (utxo_id, coin) = coins.pop().unwrap(); // TODO: Cleanup
+
+                transaction
+                    .storage::<Coins>()
+                    .insert(&utxo_id, &coin)
+                    .unwrap();
+
+                utxo_ids.push(utxo_id);
+            }
+
+            transaction.commit().unwrap();
+        }
+
+        Ok(utxo_ids)
+    }
+
+    async fn query_database_many_times(
+        &mut self,
+        utxo_ids: &[UtxoId],
+    ) -> anyhow::Result<usize> {
+        let read_database = ReadDatabase::new(
+            0,
+            BlockHeight::new(0),
+            self.db.on_chain().clone(),
+            self.db.off_chain().clone(),
+        );
+
+        let read_view = read_database.view()?;
+
+        let mut total = 0;
+        for _ in 0..self.params.num_queries {
+            let res: Vec<_> = read_view.coins(utxo_ids.to_vec()).await.collect();
+            total += res.len();
+        }
+
+        Ok(total)
+    }
+
+    fn generate_random_coin(&mut self) -> (UtxoId, CompressedCoin) {
+        let utxo_id = UtxoId::randomize(&mut self.rng);
+
+        let mut coin = CompressedCoin::default();
+        coin.set_amount(self.rng.next_u64());
+        coin.set_asset_id(AssetId::randomize(&mut self.rng));
+        coin.set_owner(Address::randomize(&mut self.rng));
+
+        (utxo_id, coin)
+    }
+}
+
+struct Parameters {
+    num_queries: usize,
+    num_blocks: usize,
+    utxo_count_per_block: usize,
+}
+
+impl Parameters {
+    fn hard_coded() -> Self {
+        Self {
+            num_queries: 100,
+            num_blocks: 100,
+            utxo_count_per_block: 1000,
+        }
+    }
+}