Add bench feature

2026-02-11 14:40:36 +00:00 · 2023-05-31 21:15:45 -04:00
5 changed files with 319 additions and 290 deletions
--- a/pageserver/Cargo.toml
+++ b/pageserver/Cargo.toml
@@ -9,6 +9,9 @@ default = []
 # Enables test-only APIs, incuding failpoints. In particular, enables the `fail_point!` macro,
 # which adds some runtime cost to run tests on outage conditions
 testing = ["fail/failpoints"]
 # Just a marker that compiles mock structs that are used in both tests and benchmarks. We
 # hide them behind a feature flag so that we can apply stronger lints to prod-only code.
 bench = []
 [dependencies]
 anyhow.workspace = true
--- a/pageserver/benches/README.md
+++ b/pageserver/benches/README.md
@@ -3,10 +3,10 @@
 # How to run
 To run all benchmarks:
-`cargo bench`
+`cargo bench --features bench`
 To run a specific file:
-`cargo bench --bench bench_layer_map`
+`cargo bench --features bench --bench bench_layer_map`
 To run a specific function:
-`cargo bench --bench bench_layer_map -- real_map_uniform_queries`
+`cargo bench --features bench --bench bench_layer_map -- real_map_uniform_queries`
--- a/pageserver/benches/bench_layer_map.rs
+++ b/pageserver/benches/bench_layer_map.rs
@@ -1,245 +1,264 @@
-use pageserver::keyspace::{KeyPartitioning, KeySpace};
+// Hiding this code under a compilation flag allows us to lint it differently than prod code
-use pageserver::repository::Key;
+#[cfg(feature = "bench")]
-use pageserver::tenant::layer_map::LayerMap;
+pub mod bench {
-use pageserver::tenant::storage_layer::{Layer, LayerDescriptor, LayerFileName};
+    use pageserver::keyspace::{KeyPartitioning, KeySpace};
-use rand::prelude::{SeedableRng, SliceRandom, StdRng};
+    use pageserver::repository::Key;
-use std::cmp::{max, min};
+    use pageserver::tenant::layer_map::LayerMap;
-use std::fs::File;
+    use pageserver::tenant::storage_layer::mock::LayerDescriptor;
-use std::io::{BufRead, BufReader};
+    use pageserver::tenant::storage_layer::{Layer, LayerFileName};
-use std::path::PathBuf;
+    use rand::prelude::{SeedableRng, SliceRandom, StdRng};
-use std::str::FromStr;
+    use std::cmp::{max, min};
-use std::sync::Arc;
+    use std::fs::File;
-use std::time::Instant;
+    use std::io::{BufRead, BufReader};
    use std::path::PathBuf;
    use std::str::FromStr;
    use std::sync::Arc;
    use std::time::Instant;
-use utils::lsn::Lsn;
+    use utils::lsn::Lsn;
-use criterion::{black_box, criterion_group, criterion_main, Criterion};
+    use criterion::{black_box, criterion_group, criterion_main, Criterion};
-fn build_layer_map(filename_dump: PathBuf) -> LayerMap<LayerDescriptor> {
+    fn build_layer_map(filename_dump: PathBuf) -> LayerMap<LayerDescriptor> {
-    let mut layer_map = LayerMap::<LayerDescriptor>::default();
+        let mut layer_map = LayerMap::<LayerDescriptor>::default();
-    let mut min_lsn = Lsn(u64::MAX);
+        let mut min_lsn = Lsn(u64::MAX);
-    let mut max_lsn = Lsn(0);
+        let mut max_lsn = Lsn(0);
-    let filenames = BufReader::new(File::open(filename_dump).unwrap()).lines();
+        let filenames = BufReader::new(File::open(filename_dump).unwrap()).lines();
-    let mut updates = layer_map.batch_update();
+        let mut updates = layer_map.batch_update();
-    for fname in filenames {
+        for fname in filenames {
-        let fname = fname.unwrap();
+            let fname = fname.unwrap();
-        let fname = LayerFileName::from_str(&fname).unwrap();
+            let fname = LayerFileName::from_str(&fname).unwrap();
-        let layer = LayerDescriptor::from(fname);
+            let layer = LayerDescriptor::from(fname);
-        let lsn_range = layer.get_lsn_range();
+            let lsn_range = layer.get_lsn_range();
-        min_lsn = min(min_lsn, lsn_range.start);
+            min_lsn = min(min_lsn, lsn_range.start);
-        max_lsn = max(max_lsn, Lsn(lsn_range.end.0 - 1));
+            max_lsn = max(max_lsn, Lsn(lsn_range.end.0 - 1));
-        updates.insert_historic(Arc::new(layer));
+            updates.insert_historic(Arc::new(layer));
    }
    println!("min: {min_lsn}, max: {max_lsn}");
    updates.flush();
    layer_map
 }
 /// Construct a layer map query pattern for benchmarks
 fn uniform_query_pattern(layer_map: &LayerMap<LayerDescriptor>) -> Vec<(Key, Lsn)> {
    // For each image layer we query one of the pages contained, at LSN right
    // before the image layer was created. This gives us a somewhat uniform
    // coverage of both the lsn and key space because image layers have
    // approximately equal sizes and cover approximately equal WAL since
    // last image.
    layer_map
        .iter_historic_layers()
        .filter_map(|l| {
            if l.is_incremental() {
                None
            } else {
                let kr = l.get_key_range();
                let lr = l.get_lsn_range();
                let key_inside = kr.start.next();
                let lsn_before = Lsn(lr.start.0 - 1);
                Some((key_inside, lsn_before))
            }
        })
        .collect()
 }
 // Construct a partitioning for testing get_difficulty map when we
 // don't have an exact result of `collect_keyspace` to work with.
 fn uniform_key_partitioning(layer_map: &LayerMap<LayerDescriptor>, _lsn: Lsn) -> KeyPartitioning {
    let mut parts = Vec::new();
    // We add a partition boundary at the start of each image layer,
    // no matter what lsn range it covers. This is just the easiest
    // thing to do. A better thing to do would be to get a real
    // partitioning from some database. Even better, remove the need
    // for key partitions by deciding where to create image layers
    // directly based on a coverage-based difficulty map.
    let mut keys: Vec<_> = layer_map
        .iter_historic_layers()
        .filter_map(|l| {
            if l.is_incremental() {
                None
            } else {
                let kr = l.get_key_range();
                Some(kr.start.next())
            }
        })
        .collect();
    keys.sort();
    let mut current_key = Key::from_hex("000000000000000000000000000000000000").unwrap();
    for key in keys {
        parts.push(KeySpace {
            ranges: vec![current_key..key],
        });
        current_key = key;
    }
    KeyPartitioning { parts }
 }
 // Benchmark using metadata extracted from our performance test environment, from
 // a project where we have run pgbench many timmes. The pgbench database was initialized
 // between each test run.
 fn bench_from_captest_env(c: &mut Criterion) {
    // TODO consider compressing this file
    let layer_map = build_layer_map(PathBuf::from("benches/odd-brook-layernames.txt"));
    let queries: Vec<(Key, Lsn)> = uniform_query_pattern(&layer_map);
    // Test with uniform query pattern
    c.bench_function("captest_uniform_queries", |b| {
        b.iter(|| {
            for q in queries.clone().into_iter() {
                black_box(layer_map.search(q.0, q.1));
            }
        });
    });
    // test with a key that corresponds to the RelDir entry. See pgdatadir_mapping.rs.
    c.bench_function("captest_rel_dir_query", |b| {
        b.iter(|| {
            let result = black_box(layer_map.search(
                Key::from_hex("000000067F00008000000000000000000001").unwrap(),
                // This LSN is higher than any of the LSNs in the tree
                Lsn::from_str("D0/80208AE1").unwrap(),
            ));
            result.unwrap();
        });
    });
 }
 // Benchmark using metadata extracted from a real project that was taknig
 // too long processing layer map queries.
 fn bench_from_real_project(c: &mut Criterion) {
    // Init layer map
    let now = Instant::now();
    let layer_map = build_layer_map(PathBuf::from("benches/odd-brook-layernames.txt"));
    println!("Finished layer map init in {:?}", now.elapsed());
    // Choose uniformly distributed queries
    let queries: Vec<(Key, Lsn)> = uniform_query_pattern(&layer_map);
    // Choose inputs for get_difficulty_map
    let latest_lsn = layer_map
        .iter_historic_layers()
        .map(|l| l.get_lsn_range().end)
        .max()
        .unwrap();
    let partitioning = uniform_key_partitioning(&layer_map, latest_lsn);
    // Check correctness of get_difficulty_map
    // TODO put this in a dedicated test outside of this mod
    {
        println!("running correctness check");
        let now = Instant::now();
        let result_bruteforce = layer_map.get_difficulty_map_bruteforce(latest_lsn, &partitioning);
        assert!(result_bruteforce.len() == partitioning.parts.len());
        println!("Finished bruteforce in {:?}", now.elapsed());
        let now = Instant::now();
        let result_fast = layer_map.get_difficulty_map(latest_lsn, &partitioning, None);
        assert!(result_fast.len() == partitioning.parts.len());
        println!("Finished fast in {:?}", now.elapsed());
        // Assert results are equal. Manually iterate for easier debugging.
        let zip = std::iter::zip(
            &partitioning.parts,
            std::iter::zip(result_bruteforce, result_fast),
        );
        for (_part, (bruteforce, fast)) in zip {
            assert_eq!(bruteforce, fast);
        }
-        println!("No issues found");
+        println!("min: {min_lsn}, max: {max_lsn}");
        updates.flush();
        layer_map
    }
-    // Define and name the benchmark function
+    /// Construct a layer map query pattern for benchmarks
-    let mut group = c.benchmark_group("real_map");
+    fn uniform_query_pattern(layer_map: &LayerMap<LayerDescriptor>) -> Vec<(Key, Lsn)> {
-    group.bench_function("uniform_queries", |b| {
+        // For each image layer we query one of the pages contained, at LSN right
-        b.iter(|| {
+        // before the image layer was created. This gives us a somewhat uniform
-            for q in queries.clone().into_iter() {
+        // coverage of both the lsn and key space because image layers have
-                black_box(layer_map.search(q.0, q.1));
+        // approximately equal sizes and cover approximately equal WAL since
-            }
+        // last image.
-        });
+        layer_map
-    });
+            .iter_historic_layers()
-    group.bench_function("get_difficulty_map", |b| {
+            .filter_map(|l| {
-        b.iter(|| {
+                if l.is_incremental() {
-            layer_map.get_difficulty_map(latest_lsn, &partitioning, Some(3));
+                    None
-        });
+                } else {
-    });
+                    let kr = l.get_key_range();
-    group.finish();
+                    let lr = l.get_lsn_range();
 }
-// Benchmark using synthetic data. Arrange image layers on stacked diagonal lines.
+                    let key_inside = kr.start.next();
-fn bench_sequential(c: &mut Criterion) {
+                    let lsn_before = Lsn(lr.start.0 - 1);
-    // Init layer map. Create 100_000 layers arranged in 1000 diagonal lines.
+
-    //
+                    Some((key_inside, lsn_before))
-    // TODO This code is pretty slow and runs even if we're only running other
+                }
-    //      benchmarks. It needs to be somewhere else, but it's not clear where.
+            })
-    //      Putting it inside the `bench_function` closure is not a solution
+            .collect()
    //      because then it runs multiple times during warmup.
    let now = Instant::now();
    let mut layer_map = LayerMap::default();
    let mut updates = layer_map.batch_update();
    for i in 0..100_000 {
        let i32 = (i as u32) % 100;
        let zero = Key::from_hex("000000000000000000000000000000000000").unwrap();
        let layer = LayerDescriptor {
            key: zero.add(10 * i32)..zero.add(10 * i32 + 1),
            lsn: Lsn(i)..Lsn(i + 1),
            is_incremental: false,
            short_id: format!("Layer {}", i),
        };
        updates.insert_historic(Arc::new(layer));
    }
    updates.flush();
    println!("Finished layer map init in {:?}", now.elapsed());
-    // Choose 100 uniformly random queries
+    // Construct a partitioning for testing get_difficulty map when we
-    let rng = &mut StdRng::seed_from_u64(1);
+    // don't have an exact result of `collect_keyspace` to work with.
-    let queries: Vec<(Key, Lsn)> = uniform_query_pattern(&layer_map)
+    fn uniform_key_partitioning(
-        .choose_multiple(rng, 100)
+        layer_map: &LayerMap<LayerDescriptor>,
-        .copied()
+        _lsn: Lsn,
-        .collect();
+    ) -> KeyPartitioning {
        let mut parts = Vec::new();
-    // Define and name the benchmark function
+        // We add a partition boundary at the start of each image layer,
-    let mut group = c.benchmark_group("sequential");
+        // no matter what lsn range it covers. This is just the easiest
-    group.bench_function("uniform_queries", |b| {
+        // thing to do. A better thing to do would be to get a real
-        b.iter(|| {
+        // partitioning from some database. Even better, remove the need
-            for q in queries.clone().into_iter() {
+        // for key partitions by deciding where to create image layers
-                black_box(layer_map.search(q.0, q.1));
+        // directly based on a coverage-based difficulty map.
-            }
+        let mut keys: Vec<_> = layer_map
            .iter_historic_layers()
            .filter_map(|l| {
                if l.is_incremental() {
                    None
                } else {
                    let kr = l.get_key_range();
                    Some(kr.start.next())
                }
            })
            .collect();
        keys.sort();
        let mut current_key = Key::from_hex("000000000000000000000000000000000000").unwrap();
        for key in keys {
            parts.push(KeySpace {
                ranges: vec![current_key..key],
            });
            current_key = key;
        }
        KeyPartitioning { parts }
    }
    // Benchmark using metadata extracted from our performance test environment, from
    // a project where we have run pgbench many timmes. The pgbench database was initialized
    // between each test run.
    fn bench_from_captest_env(c: &mut Criterion) {
        // TODO consider compressing this file
        let layer_map = build_layer_map(PathBuf::from("benches/odd-brook-layernames.txt"));
        let queries: Vec<(Key, Lsn)> = uniform_query_pattern(&layer_map);
        // Test with uniform query pattern
        c.bench_function("captest_uniform_queries", |b| {
            b.iter(|| {
                for q in queries.clone().into_iter() {
                    black_box(layer_map.search(q.0, q.1));
                }
            });
        });
-    });
+
-    group.finish();
+        // test with a key that corresponds to the RelDir entry. See pgdatadir_mapping.rs.
        c.bench_function("captest_rel_dir_query", |b| {
            b.iter(|| {
                let result = black_box(layer_map.search(
                    Key::from_hex("000000067F00008000000000000000000001").unwrap(),
                    // This LSN is higher than any of the LSNs in the tree
                    Lsn::from_str("D0/80208AE1").unwrap(),
                ));
                result.unwrap();
            });
        });
    }
    // Benchmark using metadata extracted from a real project that was taknig
    // too long processing layer map queries.
    fn bench_from_real_project(c: &mut Criterion) {
        // Init layer map
        let now = Instant::now();
        let layer_map = build_layer_map(PathBuf::from("benches/odd-brook-layernames.txt"));
        println!("Finished layer map init in {:?}", now.elapsed());
        // Choose uniformly distributed queries
        let queries: Vec<(Key, Lsn)> = uniform_query_pattern(&layer_map);
        // Choose inputs for get_difficulty_map
        let latest_lsn = layer_map
            .iter_historic_layers()
            .map(|l| l.get_lsn_range().end)
            .max()
            .unwrap();
        let partitioning = uniform_key_partitioning(&layer_map, latest_lsn);
        // Check correctness of get_difficulty_map
        // TODO put this in a dedicated test outside of this mod
        {
            println!("running correctness check");
            let now = Instant::now();
            let result_bruteforce =
                layer_map.get_difficulty_map_bruteforce(latest_lsn, &partitioning);
            assert!(result_bruteforce.len() == partitioning.parts.len());
            println!("Finished bruteforce in {:?}", now.elapsed());
            let now = Instant::now();
            let result_fast = layer_map.get_difficulty_map(latest_lsn, &partitioning, None);
            assert!(result_fast.len() == partitioning.parts.len());
            println!("Finished fast in {:?}", now.elapsed());
            // Assert results are equal. Manually iterate for easier debugging.
            let zip = std::iter::zip(
                &partitioning.parts,
                std::iter::zip(result_bruteforce, result_fast),
            );
            for (_part, (bruteforce, fast)) in zip {
                assert_eq!(bruteforce, fast);
            }
            println!("No issues found");
        }
        // Define and name the benchmark function
        let mut group = c.benchmark_group("real_map");
        group.bench_function("uniform_queries", |b| {
            b.iter(|| {
                for q in queries.clone().into_iter() {
                    black_box(layer_map.search(q.0, q.1));
                }
            });
        });
        group.bench_function("get_difficulty_map", |b| {
            b.iter(|| {
                layer_map.get_difficulty_map(latest_lsn, &partitioning, Some(3));
            });
        });
        group.finish();
    }
    // Benchmark using synthetic data. Arrange image layers on stacked diagonal lines.
    fn bench_sequential(c: &mut Criterion) {
        // Init layer map. Create 100_000 layers arranged in 1000 diagonal lines.
        //
        // TODO This code is pretty slow and runs even if we're only running other
        //      benchmarks. It needs to be somewhere else, but it's not clear where.
        //      Putting it inside the `bench_function` closure is not a solution
        //      because then it runs multiple times during warmup.
        let now = Instant::now();
        let mut layer_map = LayerMap::default();
        let mut updates = layer_map.batch_update();
        for i in 0..100_000 {
            let i32 = (i as u32) % 100;
            let zero = Key::from_hex("000000000000000000000000000000000000").unwrap();
            let layer = LayerDescriptor {
                key: zero.add(10 * i32)..zero.add(10 * i32 + 1),
                lsn: Lsn(i)..Lsn(i + 1),
                is_incremental: false,
                short_id: format!("Layer {}", i),
            };
            updates.insert_historic(Arc::new(layer));
        }
        updates.flush();
        println!("Finished layer map init in {:?}", now.elapsed());
        // Choose 100 uniformly random queries
        let rng = &mut StdRng::seed_from_u64(1);
        let queries: Vec<(Key, Lsn)> = uniform_query_pattern(&layer_map)
            .choose_multiple(rng, 100)
            .copied()
            .collect();
        // Define and name the benchmark function
        let mut group = c.benchmark_group("sequential");
        group.bench_function("uniform_queries", |b| {
            b.iter(|| {
                for q in queries.clone().into_iter() {
                    black_box(layer_map.search(q.0, q.1));
                }
            });
        });
        group.finish();
    }
    criterion_group!(group_1, bench_from_captest_env);
    criterion_group!(group_2, bench_from_real_project);
    criterion_group!(group_3, bench_sequential);
 }
-criterion_group!(group_1, bench_from_captest_env);
+#[cfg(feature = "bench")]
-criterion_group!(group_2, bench_from_real_project);
+use criterion::criterion_main;
-criterion_group!(group_3, bench_sequential);
+
-criterion_main!(group_1, group_2, group_3);
+#[cfg(feature = "bench")]
 criterion_main!(bench::group_1, bench::group_2, bench::group_3);
 #[cfg(not(feature = "bench"))]
 fn main() {
    panic!("Use `--features bench` to run benchmarks")
 }
--- a/pageserver/src/tenant/layer_map.rs
+++ b/pageserver/src/tenant/layer_map.rs
@@ -762,7 +762,8 @@ where
 #[cfg(test)]
 mod tests {
    use super::{LayerMap, Replacement};
-    use crate::tenant::storage_layer::{Layer, LayerDescriptor, LayerFileName};
+    use crate::tenant::storage_layer::mock::LayerDescriptor;
    use crate::tenant::storage_layer::{Layer, LayerFileName};
    use std::str::FromStr;
    use std::sync::Arc;
--- a/pageserver/src/tenant/storage_layer.rs
+++ b/pageserver/src/tenant/storage_layer.rs
@@ -460,80 +460,86 @@ pub fn downcast_remote_layer(
    }
 }
-/// Holds metadata about a layer without any content. Used mostly for testing.
+// Hiding this code under a compilation flag allows us to lint it differently than prod code
-///
+#[cfg(any(test, feature = "bench"))]
-/// To use filenames as fixtures, parse them as [`LayerFileName`] then convert from that to a
+pub mod mock {
-/// LayerDescriptor.
+    use super::*;
 #[derive(Clone, Debug)]
 pub struct LayerDescriptor {
    pub key: Range<Key>,
    pub lsn: Range<Lsn>,
    pub is_incremental: bool,
    pub short_id: String,
 }
-impl Layer for LayerDescriptor {
+    /// Holds metadata about a layer without any content. Used mostly for testing.
-    fn get_key_range(&self) -> Range<Key> {
+    ///
-        self.key.clone()
+    /// To use filenames as fixtures, parse them as [`LayerFileName`] then convert from that to a
    /// LayerDescriptor.
    #[derive(Clone, Debug)]
    pub struct LayerDescriptor {
        pub key: Range<Key>,
        pub lsn: Range<Lsn>,
        pub is_incremental: bool,
        pub short_id: String,
    }
-    fn get_lsn_range(&self) -> Range<Lsn> {
+    impl Layer for LayerDescriptor {
-        self.lsn.clone()
+        fn get_key_range(&self) -> Range<Key> {
-    }
+            self.key.clone()
        }
-    fn is_incremental(&self) -> bool {
+        fn get_lsn_range(&self) -> Range<Lsn> {
-        self.is_incremental
+            self.lsn.clone()
-    }
+        }
-    fn get_value_reconstruct_data(
+        fn is_incremental(&self) -> bool {
-        &self,
+            self.is_incremental
-        _key: Key,
+        }
        _lsn_range: Range<Lsn>,
        _reconstruct_data: &mut ValueReconstructState,
        _ctx: &RequestContext,
    ) -> Result<ValueReconstructResult> {
        todo!("This method shouldn't be part of the Layer trait")
    }
-    fn short_id(&self) -> String {
+        fn get_value_reconstruct_data(
-        self.short_id.clone()
+            &self,
-    }
+            _key: Key,
            _lsn_range: Range<Lsn>,
            _reconstruct_data: &mut ValueReconstructState,
            _ctx: &RequestContext,
        ) -> Result<ValueReconstructResult> {
            todo!("This method shouldn't be part of the Layer trait")
        }
-    fn dump(&self, _verbose: bool, _ctx: &RequestContext) -> Result<()> {
+        fn short_id(&self) -> String {
-        todo!()
+            self.short_id.clone()
-    }
+        }
 }
-impl From<DeltaFileName> for LayerDescriptor {
+        fn dump(&self, _verbose: bool, _ctx: &RequestContext) -> Result<()> {
-    fn from(value: DeltaFileName) -> Self {
+            todo!()
        let short_id = value.to_string();
        LayerDescriptor {
            key: value.key_range,
            lsn: value.lsn_range,
            is_incremental: true,
            short_id,
        }
    }
 }
-impl From<ImageFileName> for LayerDescriptor {
+    impl From<DeltaFileName> for LayerDescriptor {
-    fn from(value: ImageFileName) -> Self {
+        fn from(value: DeltaFileName) -> Self {
-        let short_id = value.to_string();
+            let short_id = value.to_string();
-        let lsn = value.lsn_as_range();
+            LayerDescriptor {
-        LayerDescriptor {
+                key: value.key_range,
-            key: value.key_range,
+                lsn: value.lsn_range,
-            lsn,
+                is_incremental: true,
-            is_incremental: false,
+                short_id,
-            short_id,
+            }
        }
    }
 }
-impl From<LayerFileName> for LayerDescriptor {
+    impl From<ImageFileName> for LayerDescriptor {
-    fn from(value: LayerFileName) -> Self {
+        fn from(value: ImageFileName) -> Self {
-        match value {
+            let short_id = value.to_string();
-            LayerFileName::Delta(d) => Self::from(d),
+            let lsn = value.lsn_as_range();
-            LayerFileName::Image(i) => Self::from(i),
+            LayerDescriptor {
                key: value.key_range,
                lsn,
                is_incremental: false,
                short_id,
            }
        }
    }
    impl From<LayerFileName> for LayerDescriptor {
        fn from(value: LayerFileName) -> Self {
            match value {
                LayerFileName::Delta(d) => Self::from(d),
                LayerFileName::Image(i) => Self::from(i),
            }
        }
    }
 }