feat: MUB benchmarks (#635) (#1271)

Cargo.lock

@@ -1872,6 +1872,7 @@ dependencies = [
  "criterion",
  "data_types",
  "flatbuffers",
+ "flate2",
  "generated_types",
  "influxdb_line_protocol",
  "internal_types",

mutable_buffer/Cargo.toml

@@ -33,8 +33,17 @@ tracker = { path = "../tracker" }
 [dev-dependencies] # In alphabetical order
 test_helpers = { path = "../test_helpers" }
 criterion = "0.3"
+flate2 = "1.0.20"
 
 [features]
 default = []
 # Disables snapshot caching
 nocache = []
+
+[[bench]]
+name = "snapshot"
+harness = false
+
+[[bench]]
+name = "write"
+harness = false

mutable_buffer/benches/snapshot.rs

@@ -0,0 +1,48 @@
+use std::io::Read;
+
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use flate2::read::GzDecoder;
+
+use internal_types::entry::{test_helpers::lp_to_entries, ClockValue};
+use mutable_buffer::chunk::Chunk;
+use tracker::MemRegistry;
+
+#[inline]
+fn snapshot_chunk(chunk: &Chunk) {
+    let _ = chunk.snapshot();
+}
+
+fn chunk(count: usize) -> Chunk {
+    let mut chunk = Chunk::new(0, &MemRegistry::new());
+
+    let raw = include_bytes!("../../tests/fixtures/lineproto/tag_values.lp.gz");
+    let mut gz = GzDecoder::new(&raw[..]);
+    let mut lp = String::new();
+    gz.read_to_string(&mut lp).unwrap();
+
+    for _ in 0..count {
+        for entry in lp_to_entries(&lp) {
+            for write in entry.partition_writes().iter().flatten() {
+                chunk
+                    .write_table_batches(ClockValue::new(0), 0, write.table_batches().as_slice())
+                    .unwrap();
+            }
+        }
+    }
+
+    chunk
+}
+
+pub fn snapshot_mb(c: &mut Criterion) {
+    let mut group = c.benchmark_group("snapshot_mb");
+    for count in &[1, 2, 3, 4, 5] {
+        let chunk = chunk(*count as _);
+        group.bench_function(BenchmarkId::from_parameter(count), |b| {
+            b.iter(|| snapshot_chunk(&chunk));
+        });
+    }
+    group.finish();
+}
+
+criterion_group!(benches, snapshot_mb);
+criterion_main!(benches);

mutable_buffer/benches/write.rs

@@ -0,0 +1,45 @@
+use std::io::Read;
+
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use flate2::read::GzDecoder;
+
+use internal_types::entry::{test_helpers::lp_to_entries, ClockValue, Entry};
+use mutable_buffer::chunk::Chunk;
+use tracker::MemRegistry;
+
+#[inline]
+fn write_chunk(count: usize, entries: &[Entry]) {
+    let mut chunk = Chunk::new(0, &MemRegistry::new());
+
+    for _ in 0..count {
+        for entry in entries {
+            for write in entry.partition_writes().iter().flatten() {
+                chunk
+                    .write_table_batches(ClockValue::new(0), 0, write.table_batches().as_slice())
+                    .unwrap();
+            }
+        }
+    }
+}
+
+fn load_entries() -> Vec<Entry> {
+    let raw = include_bytes!("../../tests/fixtures/lineproto/tag_values.lp.gz");
+    let mut gz = GzDecoder::new(&raw[..]);
+    let mut lp = String::new();
+    gz.read_to_string(&mut lp).unwrap();
+    lp_to_entries(&lp)
+}
+
+pub fn write_mb(c: &mut Criterion) {
+    let mut group = c.benchmark_group("write_mb");
+    let entries = load_entries();
+    for count in &[1, 2, 3, 4, 5] {
+        group.bench_function(BenchmarkId::from_parameter(count), |b| {
+            b.iter(|| write_chunk(*count, &entries));
+        });
+    }
+    group.finish();
+}
+
+criterion_group!(benches, write_mb);
+criterion_main!(benches);