feat: add wal replay benchmark

src/benchmarks/Cargo.toml

@@ -54,6 +54,7 @@ size_ext = { workspace = true }
 snafu = { workspace = true }
 table_engine = { workspace = true }
 table_kv = { workspace = true }
+tempfile = { workspace = true }
 time_ext = { workspace = true }
 tokio = { workspace = true }
 toml_ext = { workspace = true }
@@ -63,6 +64,7 @@ zstd = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true }
+tempfile = { workspace = true }
 
 [[bench]]
 name = "bench"

src/benchmarks/bench.toml

@@ -64,3 +64,7 @@ bench_measurement_time = "60s"
 bench_sample_size = 60
 batch_size = 512
 value_size = 1024
+
+[replay_bench]
+bench_measurement_time = "3s"
+bench_sample_size = 10

src/benchmarks/benches/bench.rs

@@ -17,13 +17,14 @@
 
 //! Benchmarks
 
-use std::sync::Once;
+use std::{cell::RefCell, sync::Once};
 
 use benchmarks::{
     config::{self, BenchConfig},
     merge_memtable_bench::MergeMemTableBench,
     merge_sst_bench::MergeSstBench,
     parquet_bench::ParquetBench,
+    replay_bench::ReplayBench,
     scan_memtable_bench::ScanMemTableBench,
     sst_bench::SstBench,
     wal_write_bench::WalWriteBench,
@@ -208,6 +209,24 @@ fn bench_wal_write(c: &mut Criterion) {
     group.finish();
 }
 
+fn bench_replay_iter(b: &mut Bencher<'_>, bench: &RefCell<ReplayBench>) {
+    let mut bench = bench.borrow_mut();
+    b.iter(|| bench.run_bench())
+}
+
+fn bench_replay(c: &mut Criterion) {
+    let config = init_bench();
+
+    let mut group = c.benchmark_group("replay");
+
+    group.measurement_time(config.replay_bench.bench_measurement_time.0);
+    group.sample_size(config.replay_bench.bench_sample_size);
+
+    let bench = RefCell::new(ReplayBench::new(config.replay_bench));
+    group.bench_with_input(BenchmarkId::new("replay", 0), &bench, bench_replay_iter);
+    group.finish();
+}
+
 criterion_group!(
     name = benches;
     config = Criterion::default().with_profiler(PProfProfiler::new(100, Output::Flamegraph(None)));
@@ -217,6 +236,7 @@ criterion_group!(
     bench_scan_memtable,
     bench_merge_memtable,
     bench_wal_write,
+    bench_replay,
 );
 
 criterion_main!(benches);

src/benchmarks/config/bench.toml

@@ -71,3 +71,8 @@ bench_measurement_time = "60s"
 bench_sample_size = 60
 batch_size = 512
 value_size = 1024
+
+[replay_bench]
+bench_measurement_time = "3s"
+bench_sample_size = 10
+batch_size = 10000

src/benchmarks/src/config.rs

@@ -38,6 +38,7 @@ pub struct BenchConfig {
     pub scan_memtable_bench: ScanMemTableBenchConfig,
     pub merge_memtable_bench: MergeMemTableBenchConfig,
     pub wal_write_bench: WalWriteBenchConfig,
+    pub replay_bench: ReplayConfig,
 }
 
 // TODO(yingwen): Maybe we can use layze static to load config first.
@@ -147,3 +148,10 @@ pub struct WalWriteBenchConfig {
     pub batch_size: usize,
     pub value_size: usize,
 }
+
+#[derive(Deserialize)]
+pub struct ReplayConfig {
+    pub bench_measurement_time: ReadableDuration,
+    pub bench_sample_size: usize,
+    pub batch_size: usize,
+}

src/benchmarks/src/lib.rs

@@ -23,9 +23,11 @@ pub mod config;
 pub mod merge_memtable_bench;
 pub mod merge_sst_bench;
 pub mod parquet_bench;
+pub mod replay_bench;
 pub mod scan_memtable_bench;
 pub mod sst_bench;
 pub mod sst_tools;
+pub mod table;
 pub mod util;
 pub mod wal_write_bench;
 

src/benchmarks/src/replay_bench.rs

@@ -0,0 +1,97 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Replay bench.
+
+use std::sync::Arc;
+
+use analytic_engine::RecoverMode;
+use runtime::Runtime;
+use util::{OpenTablesMethod, RocksDBEngineBuildContext, TestContext, TestEnv};
+use wal::rocksdb_impl::manager::RocksDBWalsOpener;
+
+use crate::{config::ReplayConfig, table::FixedSchemaTable, util};
+
+pub struct ReplayBench {
+    runtime: Arc<Runtime>,
+    test_ctx: TestContext<RocksDBWalsOpener>,
+    table: FixedSchemaTable,
+    batch_size: usize,
+}
+
+impl ReplayBench {
+    pub fn new(config: ReplayConfig) -> Self {
+        let runtime = util::new_runtime(1);
+        let engine_context = RocksDBEngineBuildContext::new(
+            RecoverMode::TableBased,
+            OpenTablesMethod::WithOpenShard,
+        );
+        let env: TestEnv = TestEnv::builder().build();
+
+        let (test_ctx, fixed_schema_table) = env.block_on(async {
+            let mut test_ctx = env.new_context(&engine_context);
+            test_ctx.open().await;
+
+            let fixed_schema_table = test_ctx
+                .create_fixed_schema_table("test_replay_table1")
+                .await;
+            let _ = test_ctx
+                .create_fixed_schema_table("test_replay_table2")
+                .await;
+            let _ = test_ctx
+                .create_fixed_schema_table("test_replay_table3")
+                .await;
+
+            (test_ctx, fixed_schema_table)
+        });
+
+        ReplayBench {
+            runtime: Arc::new(runtime),
+            test_ctx,
+            table: fixed_schema_table,
+            batch_size: config.batch_size,
+        }
+    }
+
+    pub fn run_bench(&mut self) {
+        self.runtime.block_on(async {
+            self.table.prepare_write_requests(self.batch_size);
+            let rows = self.table.row_tuples();
+
+            // Write data to table.
+            let mut table_names = Vec::new();
+            for (table_name, _) in self.test_ctx.name_to_tables().iter() {
+                let row_group = self.table.rows_to_row_group(&rows);
+                self.test_ctx
+                    .write_to_table(table_name.as_str(), row_group)
+                    .await;
+                table_names.push(table_name.clone());
+            }
+
+            // Reopen db.
+            self.test_ctx
+                .reopen_with_tables(
+                    table_names
+                        .iter()
+                        .map(|s| s.as_str())
+                        .collect::<Vec<_>>()
+                        .as_slice(),
+                )
+                .await;
+        });
+    }
+}