From 05a3bb7d3f93ebcfbf0be5f25a7787438b4af8ca Mon Sep 17 00:00:00 2001
From: Robert Habermeier <rphmeier@gmail.com>
Date: Tue, 11 Feb 2025 12:56:28 +0100
Subject: [PATCH] feat(benchtop): expose page cache upper levels /
 prepopulation in CLI

---
 benchtop/src/backend.rs |  4 ++++
 benchtop/src/cli.rs     | 12 ++++++++++++
 benchtop/src/main.rs    |  4 ++++
 benchtop/src/nomt.rs    |  4 ++++
 4 files changed, 24 insertions(+)

diff --git a/benchtop/src/backend.rs b/benchtop/src/backend.rs
index dd2d8f28..afb68fd2 100644
--- a/benchtop/src/backend.rs
+++ b/benchtop/src/backend.rs
@@ -23,6 +23,8 @@ impl Backend {
         hashtable_buckets: Option<u32>,
         page_cache_size: Option<usize>,
         leaf_cache_size: Option<usize>,
+        page_cache_upper_levels: usize,
+        prepopulate_page_cache: bool,
         overlay_window_length: usize,
     ) -> DB {
         match self {
@@ -34,6 +36,8 @@ impl Backend {
                 hashtable_buckets,
                 page_cache_size,
                 leaf_cache_size,
+                page_cache_upper_levels,
+                prepopulate_page_cache,
                 overlay_window_length,
             )),
             Backend::SpTrie => DB::SpTrie(SpTrieDB::open(reset)),
diff --git a/benchtop/src/cli.rs b/benchtop/src/cli.rs
index 788dbfc7..013349f1 100644
--- a/benchtop/src/cli.rs
+++ b/benchtop/src/cli.rs
@@ -150,6 +150,18 @@ pub struct WorkloadParams {
     #[arg(long = "page-cache-size")]
     pub page_cache_size: Option<usize>,
 
+    /// The number of upper levels of the NOMT page tree to keep permanently cached.
+    /// Only used with the Nomt backend.
+    #[arg(long = "page-cache-upper-levels")]
+    #[clap(default_value = "3")]
+    pub page_cache_upper_levels: usize,
+
+    /// Whether to prepopulate the page cache with the upper levels of the page tree in NOMT.
+    /// Only used with the Nomt backend.
+    #[arg(long = "prepopulate-page-cache")]
+    #[clap(default_value = "true")]
+    pub prepopulate_page_cache: bool,
+
     /// The size of the leaf cache used in NOMT to store Beatree leaves, measured in MiB.
     /// Only used with the Nomt backend.
     #[arg(long = "leaf-cache-size")]
diff --git a/benchtop/src/main.rs b/benchtop/src/main.rs
index cae76e2f..8c7f6313 100644
--- a/benchtop/src/main.rs
+++ b/benchtop/src/main.rs
@@ -33,6 +33,8 @@ pub fn init(params: InitParams) -> Result<()> {
         workload_params.hashtable_buckets,
         workload_params.page_cache_size,
         workload_params.leaf_cache_size,
+        workload_params.page_cache_upper_levels,
+        workload_params.prepopulate_page_cache,
         0,
     );
     db.execute(None, &mut *init, None);
@@ -54,6 +56,8 @@ pub fn run(params: RunParams) -> Result<()> {
         workload_params.hashtable_buckets,
         workload_params.page_cache_size,
         workload_params.leaf_cache_size,
+        workload_params.page_cache_upper_levels,
+        workload_params.prepopulate_page_cache,
         workload_params.overlay_window_length,
     );
 
diff --git a/benchtop/src/nomt.rs b/benchtop/src/nomt.rs
index 61a09cb2..be5a7714 100644
--- a/benchtop/src/nomt.rs
+++ b/benchtop/src/nomt.rs
@@ -26,6 +26,8 @@ impl NomtDB {
         hashtable_buckets: Option<u32>,
         page_cache_size: Option<usize>,
         leaf_cache_size: Option<usize>,
+        page_cache_upper_levels: usize,
+        prepopulate_page_cache: bool,
         overlay_window_capacity: usize,
     ) -> Self {
         let nomt_db_folder =
@@ -50,6 +52,8 @@ impl NomtDB {
         if let Some(buckets) = hashtable_buckets {
             opts.hashtable_buckets(buckets);
         }
+        opts.page_cache_upper_levels(page_cache_upper_levels);
+        opts.prepopulate_page_cache(prepopulate_page_cache);
 
         let nomt = Nomt::open(opts).unwrap();
         Self {