use std::path::PathBuf;
use bable::{
cache::{BlockCache, CacheOptions, IndexCache},
BableIterator, Builder, Flag, MergeTableIterator, Options, RefCounter, SimpleBuilder, Table,
TableBuilder, TableIterator,
};
use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
use rand::{thread_rng, Rng, RngCore};
use scopeguard::defer;
use rand::distributions::Alphanumeric;
use vpb::kvstructs::{Key, Value, ValueExt};
#[cfg(not(target_env = "msvc"))]
use tikv_jemallocator::Jemalloc;
use zallocator::pool::AllocatorPool;
#[cfg(not(target_env = "msvc"))]
#[global_allocator]
static GLOBAL: Jemalloc = Jemalloc;
#[allow(dead_code)]
pub fn rand_value() -> Vec<u8> {
rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(32)
.collect::<Vec<_>>()
}
fn get_table_for_bench<B: TableBuilder>(count: usize) -> Table {
let opts = Options::default_with_pool(AllocatorPool::new(1))
.set_compression(vpb::Compression::new())
.set_block_cache(BlockCache::new(CacheOptions::new(1000000 * 10, 1000000)).unwrap());
let mut builder = B::new(RefCounter::new(opts)).unwrap();
let mut rng = thread_rng();
let mut filename = std::env::temp_dir();
filename.push(rng.gen::<u32>().to_string());
filename.set_extension("sst");
defer!(std::fs::remove_file(&filename).unwrap());
(0..count).for_each(|i| {
builder.insert(
&Key::from(format!("{:016x}", i)),
&Value::from(i.to_string()),
0,
);
});
Table::create_table(&filename, builder).unwrap()
}
fn bench_table_builder(c: &mut Criterion) {
c.bench_function("bench table simple builder", |b| {
const KEY_COUNT: usize = 1300000;
let mut key_list = vec![];
for i in 0..KEY_COUNT {
let k = Key::from(format!("{:032}", i));
key_list.push(k);
}
let vs = Value::from(rand_value());
let opt = RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_table_size(5 << 20)
.set_block_size(4 * 1024)
.set_compression(vpb::Compression::new()),
);
b.iter(|| {
let mut builder = SimpleBuilder::new(opt.clone()).unwrap();
key_list.iter().take(KEY_COUNT).for_each(|k| {
builder.insert(k, &vs, 0);
});
builder.build().unwrap();
});
});
c.bench_function("bench table builder no compression", |b| {
const KEY_COUNT: usize = 1300000;
let mut key_list = vec![];
for i in 0..KEY_COUNT {
let k = Key::from(format!("{:032}", i));
key_list.push(k);
}
let vs = Value::from(rand_value());
let opt = RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_table_size(5 << 20)
.set_block_size(4 * 1024)
.set_compression(vpb::Compression::new()),
);
b.iter(|| {
let mut builder = Builder::new(opt.clone()).unwrap();
key_list.iter().take(KEY_COUNT).for_each(|k| {
builder.insert(k, &vs, 0);
});
builder.build().unwrap();
});
});
c.bench_function("bench table builder snappy compression", |b| {
const KEY_COUNT: usize = 1300000;
let mut key_list = vec![];
for i in 0..KEY_COUNT {
let k = Key::from(format!("{:032}", i));
key_list.push(k);
}
let vs = Value::from(rand_value());
let opt = RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_compression(vpb::Compression::snappy()),
);
b.iter(|| {
let mut builder = Builder::new(opt.clone()).unwrap();
key_list.iter().take(KEY_COUNT).for_each(|k| {
builder.insert(k, &vs, 0);
});
builder.build().unwrap();
});
});
c.bench_function("bench table builder lz4 compression", |b| {
const KEY_COUNT: usize = 1300000;
let mut key_list = vec![];
for i in 0..KEY_COUNT {
let k = Key::from(format!("{:032}", i));
key_list.push(k);
}
let vs = Value::from(rand_value());
let opt = RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_compression(vpb::Compression::lz4()),
);
b.iter(|| {
let mut builder = Builder::new(opt.clone()).unwrap();
key_list.iter().take(KEY_COUNT).for_each(|k| {
builder.insert(k, &vs, 0);
});
builder.build().unwrap();
});
});
c.bench_function("bench table builder zstd compression level 1", |b| {
const KEY_COUNT: usize = 1300000;
let mut key_list = vec![];
for i in 0..KEY_COUNT {
let k = Key::from(format!("{:032}", i));
key_list.push(k);
}
let vs = Value::from(rand_value());
let opt = RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_compression(vpb::Compression::zstd(1)),
);
b.iter(|| {
let mut builder = Builder::new(opt.clone()).unwrap();
key_list.iter().take(KEY_COUNT).for_each(|k| {
builder.insert(k, &vs, 0);
});
builder.build().unwrap();
});
});
c.bench_function("bench table builder zstd compression level 3", |b| {
const KEY_COUNT: usize = 1300000;
let mut key_list = vec![];
for i in 0..KEY_COUNT {
let k = Key::from(format!("{:032}", i));
key_list.push(k);
}
let vs = Value::from(rand_value());
let opt = RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_compression(vpb::Compression::zstd(3)),
);
b.iter(|| {
let mut builder = Builder::new(opt.clone()).unwrap();
key_list.iter().take(KEY_COUNT).for_each(|k| {
builder.insert(k, &vs, 0);
});
builder.build().unwrap();
});
});
c.bench_function("bench table builder zstd compression level 15", |b| {
const KEY_COUNT: usize = 1300000;
let mut key_list = vec![];
for i in 0..KEY_COUNT {
let k = Key::from(format!("{:032}", i));
key_list.push(k);
}
let vs = Value::from(rand_value());
let opt = RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_compression(vpb::Compression::zstd(15)),
);
b.iter(|| {
let mut builder = Builder::new(opt.clone()).unwrap();
key_list.iter().take(KEY_COUNT).for_each(|k| {
builder.insert(k, &vs, 0);
});
builder.build().unwrap();
});
});
c.bench_function("bench table builder encryption", |b| {
const KEY_COUNT: usize = 1300000;
let mut key_list = vec![];
for i in 0..KEY_COUNT {
let k = Key::from(format!("{:032}", i));
key_list.push(k);
}
let vs = Value::from(rand_value());
let mut rng = thread_rng();
let mut key = vec![0u8; 32];
rng.fill_bytes(key.as_mut_slice());
let opt = RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_index_cache(IndexCache::new(CacheOptions::new(1000, 1 << 20)).unwrap())
.set_encryption(vpb::Encryption::aes(key)),
);
b.iter(|| {
let mut builder = Builder::new(opt.clone()).unwrap();
key_list.iter().take(KEY_COUNT).for_each(|k| {
builder.insert(k, &vs, 0);
});
builder.build().unwrap();
});
});
}
fn bench_table(c: &mut Criterion) {
c.bench_function("bench table read", |b| {
let n = 5 * (1e6 as usize);
let tbl = get_table_for_bench::<Builder>(n);
b.iter(|| {
let mut it = tbl.iter(Flag::NONE);
it.seek_to_first();
while it.valid() {
it.next();
}
});
});
c.bench_function("bench table read and build (simple builder)", |b| {
let n = 5 * (1e6 as usize);
let tbl = get_table_for_bench::<SimpleBuilder>(n);
b.iter(|| {
let mut it = tbl.iter(Flag::NONE);
let mut builder = SimpleBuilder::new(RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_compression(vpb::Compression::new())
.set_block_cache(
BlockCache::new(CacheOptions::new(1000000 * 10, 1000000)).unwrap(),
),
))
.unwrap();
it.seek_to_first();
while it.valid() {
builder.insert(
&it.key().unwrap().to_key(),
&it.val().unwrap().to_value(),
0,
);
it.next();
}
builder.build().unwrap();
});
});
c.bench_function("bench table read and build", |b| {
let n = 5 * (1e6 as usize);
let tbl = get_table_for_bench::<Builder>(n);
b.iter(|| {
let mut it = tbl.iter(Flag::NONE);
let mut builder = Builder::new(RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_compression(vpb::Compression::new())
.set_block_cache(
BlockCache::new(CacheOptions::new(1000000 * 10, 1000000)).unwrap(),
),
))
.unwrap();
it.seek_to_first();
while it.valid() {
builder.insert(
&it.key().unwrap().to_key(),
&it.val().unwrap().to_value(),
0,
);
it.next();
}
builder.build().unwrap();
});
});
let mut rng = thread_rng();
const M: usize = 5; const N: usize = 5 * (1e6 as usize);
const TABLE_SIZE: usize = N / M;
c.bench_function("bench table merged (simple builder)", |b| {
b.iter_batched(
|| {
let cache = BlockCache::new(CacheOptions::new(1000000 * 10, 1000000)).unwrap();
let mut tables = vec![];
for i in 0..M {
let mut filename = PathBuf::new();
filename.push(std::env::temp_dir());
filename.push(rng.next_u32().to_string());
filename.set_extension("sst");
let opts = RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_compression(vpb::Compression::new())
.set_block_size(4 * 1024)
.set_block_cache(cache.clone()),
);
let mut builder = SimpleBuilder::new(opts).unwrap();
for j in 0..TABLE_SIZE {
let id = j * M + i;
let k = format!("{:016x}", id);
let v = format!("{}", id);
builder.insert(&Key::from(k), &Value::from(v), 0);
}
let tbl = Table::create_table(filename, builder).unwrap();
tables.push(tbl);
}
tables
},
|tables| {
let mut iters = vec![];
for tbl in tables {
iters.push(Box::new(TableIterator::Uni(tbl.iter(Flag::NONE))));
}
let mut iter = MergeTableIterator::new(iters, false).unwrap();
iter.rewind();
while iter.valid() {
iter.next();
}
},
BatchSize::SmallInput,
);
});
c.bench_function("bench table merged", |b| {
b.iter_batched(
|| {
let cache = BlockCache::new(CacheOptions::new(1000000 * 10, 1000000)).unwrap();
let mut tables = vec![];
for i in 0..M {
let mut filename = PathBuf::new();
filename.push(std::env::temp_dir());
filename.push(rng.next_u32().to_string());
filename.set_extension("sst");
let opts = RefCounter::new(
Options::default_with_pool(AllocatorPool::new(1))
.set_compression(vpb::Compression::new())
.set_block_size(4 * 1024)
.set_block_cache(cache.clone()),
);
let mut builder = Builder::new(opts).unwrap();
for j in 0..TABLE_SIZE {
let id = j * M + i;
let k = format!("{:016x}", id);
let v = format!("{}", id);
builder.insert(&Key::from(k), &Value::from(v), 0);
}
let tbl = Table::create_table(filename, builder).unwrap();
tables.push(tbl);
}
tables
},
|tables| {
let mut iters = vec![];
for tbl in tables {
iters.push(Box::new(TableIterator::Uni(tbl.iter(Flag::NONE))));
}
let mut iter = MergeTableIterator::new(iters, false).unwrap();
iter.rewind();
while iter.valid() {
iter.next();
}
},
BatchSize::SmallInput,
);
});
c.bench_function("bench random read", |b| {
let n = 5 * (1e6 as usize);
let tbl = get_table_for_bench::<Builder>(n);
let mut it = tbl.iter(Flag::NONE);
b.iter_batched(
|| {
let i: usize = rng.gen_range(0..n);
(Key::from(format!("{:016x}", i)), Value::from(i.to_string()))
},
|(k, v)| {
it.seek(&k);
assert!(it.valid());
assert_eq!(it.val().unwrap().parse_value(), v.parse_value())
},
BatchSize::SmallInput,
);
});
}
criterion_group! {
name = benches_table;
config = Criterion::default().sample_size(10);
targets = bench_table_builder, bench_table
}
criterion_main!(benches_table);