From a64a5cb56c5a70ac9a56fcdd9583ebd794fee36e Mon Sep 17 00:00:00 2001
From: Bojan Serafimov <bojan.serafimov7@gmail.com>
Date: Thu, 18 May 2023 14:43:29 -0400
Subject: [PATCH] increase size

---
 pageserver/benches/bench_disk_lookup.rs | 62 +++++++++++++++----------
 1 file changed, 38 insertions(+), 24 deletions(-)

diff --git a/pageserver/benches/bench_disk_lookup.rs b/pageserver/benches/bench_disk_lookup.rs
index 3deaf1d87e..664662394e 100644
--- a/pageserver/benches/bench_disk_lookup.rs
+++ b/pageserver/benches/bench_disk_lookup.rs
@@ -1,5 +1,6 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use pageserver::{tenant::{disk_btree::{DiskBtreeBuilder, DiskBtreeReader, VisitDirection}, block_io::{BlockBuf, FileBlockReader}, storage_layer::DeltaLayerWriter}, repository::Key, virtual_file::{VirtualFile, self}, page_cache};
+use std::time::Instant;
 use rand::prelude::{SeedableRng, SliceRandom, StdRng};
 use utils::{id::{TimelineId, TenantId}, lsn::Lsn};
 use std::{io::{Read, Write}, path::PathBuf};
@@ -12,6 +13,34 @@ struct MockLayer {
 }
 
 impl MockLayer {
+    fn make_simple(n_keys: i128) -> Self {
+        let block_buf = BlockBuf::new();
+        let mut writer = DiskBtreeBuilder::<_, 24>::new(block_buf);
+        for i in 0..n_keys {
+            let key: Key = Key::from_i128(i);
+            let value: u64 = i as u64;
+
+            let mut key_bytes: [u8; 24] = [8u8; 24];
+            key.write_to_byte_slice(&mut key_bytes);
+            writer.append(&key_bytes, value).unwrap();
+        }
+        let (index_root_blk, block_buf) = writer.finish().unwrap();
+        let index_start_blk = 0; // ???
+        let path = std::env::current_dir().unwrap().join("btree.tmp");
+        let layer = MockLayer {
+            path: path.clone(),
+            index_start_blk,
+            index_root_blk,
+        };
+
+        let mut file = VirtualFile::create(&path).unwrap();
+        for buf in block_buf.blocks {
+            file.write_all(buf.as_ref()).unwrap();
+        }
+
+        layer
+    }
+
     fn read(&self, key: i128) -> Option<u64> {
         // Read from disk btree
         let file = FileBlockReader::new(VirtualFile::open(&self.path).unwrap());
@@ -42,38 +71,23 @@ fn bench_disk_lookup(c: &mut Criterion) {
     virtual_file::init(10);
     page_cache::init(10000);
 
-    let n_keys = 250000;
+    // Results in a 40MB index
+    let n_keys = 4_000_000;
 
     // Write to disk btree
-    let block_buf = BlockBuf::new();
-    let mut writer = DiskBtreeBuilder::<_, 24>::new(block_buf);
-    for i in 0..n_keys {
-        let key: Key = Key::from_i128(i);
-        let value: u64 = i as u64;
-
-        let mut key_bytes: [u8; 24] = [8u8; 24];
-        key.write_to_byte_slice(&mut key_bytes);
-        writer.append(&key_bytes, value).unwrap();
-    }
-    let (index_root_blk, block_buf) = writer.finish().unwrap();
-    let index_start_blk = 0; // ???
-    let path = std::env::current_dir().unwrap().join("btree.tmp");
-    let layer = MockLayer {
-        path: path.clone(),
-        index_start_blk,
-        index_root_blk,
-    };
-
-    let mut file = VirtualFile::create(&path).unwrap();
-    for buf in block_buf.blocks {
-        file.write_all(buf.as_ref()).unwrap();
-    }
+    let layer = MockLayer::make_simple(n_keys);
 
     // Pick queries
     let rng = &mut StdRng::seed_from_u64(1);
     let queries: Vec<_> = (0..n_keys).collect();
     let queries: Vec<_> = queries.choose_multiple(rng, 100).copied().collect();
 
+    let now = Instant::now();
+    black_box({
+        assert_eq!(layer.read(queries[0]), Some(queries[0] as u64));
+    });
+    println!("Finished one query in {:?}", now.elapsed());
+
     // Define and name the benchmark function
     let mut group = c.benchmark_group("g1");
     group.bench_function("f1", |b| {