From fe4fa29930c30b0bdf50dbd4bbc82ad6fa7541fd Mon Sep 17 00:00:00 2001
From: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>
Date: Thu, 22 Apr 2021 12:15:32 +0100
Subject: [PATCH] feat: MUB benchmarks (#635) (#1271)

---
 Cargo.lock                         |  1 +
 mutable_buffer/Cargo.toml          |  9 ++++++
 mutable_buffer/benches/snapshot.rs | 48 ++++++++++++++++++++++++++++++
 mutable_buffer/benches/write.rs    | 45 ++++++++++++++++++++++++++++
 4 files changed, 103 insertions(+)
 create mode 100644 mutable_buffer/benches/snapshot.rs
 create mode 100644 mutable_buffer/benches/write.rs

diff --git a/Cargo.lock b/Cargo.lock
index 2f1df033f6..a8f8b0d760 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1872,6 +1872,7 @@ dependencies = [
  "criterion",
  "data_types",
  "flatbuffers",
+ "flate2",
  "generated_types",
  "influxdb_line_protocol",
  "internal_types",
diff --git a/mutable_buffer/Cargo.toml b/mutable_buffer/Cargo.toml
index b37b990402..f1fbb70d1e 100644
--- a/mutable_buffer/Cargo.toml
+++ b/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
diff --git a/mutable_buffer/benches/snapshot.rs b/mutable_buffer/benches/snapshot.rs
new file mode 100644
index 0000000000..f89c79de41
--- /dev/null
+++ b/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);
diff --git a/mutable_buffer/benches/write.rs b/mutable_buffer/benches/write.rs
new file mode 100644
index 0000000000..60284d590d
--- /dev/null
+++ b/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);