use divan::counter::BytesCount;
use divan::Bencher;
use rand::prelude::SliceRandom;
use rand::{
distributions::{Distribution, Uniform},
Rng,
};
use rustcask::Rustcask;
use tempfile::TempDir;
fn main() {
divan::main();
}
const COUNT_KV_PAIRS: usize = 1000;
const OVERWRITE_COUNT: usize = 30;
const KEY_SIZE: usize = 1024; const VAL_SIZE: usize = 8096;
#[derive(Clone)]
struct KeyValuePair(Vec<u8>, Vec<u8>);
impl KeyValuePair {
fn random<R: Rng>(rng: &mut R, key_size: usize, value_size: usize) -> KeyValuePair {
let key: Vec<u8> = (0..key_size).map(|_| rng.gen::<u8>()).collect();
let value: Vec<u8> = (0..value_size).map(|_| rng.gen::<u8>()).collect();
KeyValuePair(key, value)
}
fn random_many<R: Rng>(
rng: &mut R,
count_kv_pairs: usize,
max_key_size: usize,
max_val_size: usize,
) -> Vec<KeyValuePair> {
let key_dist = Uniform::from(1..max_key_size);
let val_dist = Uniform::from(1..max_val_size);
(0..count_kv_pairs)
.map(|_| {
let key_size = key_dist.sample(rng);
let value_size = val_dist.sample(rng);
KeyValuePair::random(rng, key_size, value_size)
})
.collect()
}
}
#[divan::bench]
fn bench_writes(bencher: Bencher) {
let temp_dir = TempDir::new().expect("unable to create temporary working directory");
let store = Rustcask::builder().open(temp_dir.path()).unwrap();
bencher
.with_inputs(move || {
let mut rng = rand::thread_rng();
let store = store.clone();
(KeyValuePair::random(&mut rng, KEY_SIZE, VAL_SIZE), store)
})
.input_counter(|(kv_pair, _)| BytesCount::new(kv_pair.0.len() + kv_pair.1.len()))
.bench_values(|(kv_pair, mut store)| store.set(kv_pair.0, kv_pair.1));
}
#[divan::bench]
fn bench_writes_sync_mode(bencher: Bencher) {
let temp_dir = TempDir::new().expect("unable to create temporary working directory");
let store = Rustcask::builder()
.set_sync_mode(true)
.open(temp_dir.path())
.unwrap();
bencher
.with_inputs(move || {
let mut rng = rand::thread_rng();
let store = store.clone();
(KeyValuePair::random(&mut rng, KEY_SIZE, VAL_SIZE), store)
})
.input_counter(|(kv_pair, _)| BytesCount::new(kv_pair.0.len() + kv_pair.1.len()))
.bench_values(|(kv_pair, mut store)| store.set(kv_pair.0, kv_pair.1));
}
#[divan::bench]
fn bench_random_reads(bencher: Bencher) {
let temp_dir = TempDir::new().expect("unable to create temporary working directory");
let mut store = Rustcask::builder().open(temp_dir.path()).unwrap();
let mut rng: rand::prelude::ThreadRng = rand::thread_rng();
let kv_pairs = KeyValuePair::random_many(&mut rng, COUNT_KV_PAIRS, KEY_SIZE, VAL_SIZE);
for kv_pair in kv_pairs.clone() {
store.set(kv_pair.0, kv_pair.1).unwrap();
}
bencher
.with_inputs(move || {
let store = store.clone();
let mut rng = rand::thread_rng();
let rand_index = rng.gen_range(0..kv_pairs.len());
let rand_kv_pair = kv_pairs[rand_index].clone();
(rand_kv_pair, store)
})
.input_counter(|(rand_kv_pair, _)| {
BytesCount::new(rand_kv_pair.0.len() + rand_kv_pair.1.len())
})
.bench_values(|(kv_pair, mut store)| store.get(&kv_pair.0).unwrap());
}
#[divan::bench()]
fn bench_open_hint_files_disabled(bencher: Bencher) {
let temp_dir = TempDir::new().expect("unable to create temporary working directory");
let mut store = Rustcask::builder().open(temp_dir.path()).unwrap();
let mut rng: rand::prelude::ThreadRng = rand::thread_rng();
let kv_pairs = KeyValuePair::random_many(&mut rng, COUNT_KV_PAIRS, KEY_SIZE, VAL_SIZE);
for kv_pair in kv_pairs.clone() {
store.set(kv_pair.0, kv_pair.1).unwrap();
}
let sample_kv_pairs: Vec<&KeyValuePair> = kv_pairs
.choose_multiple(&mut rng, OVERWRITE_COUNT)
.collect();
for kv_pair in sample_kv_pairs {
let rand_value = (0..VAL_SIZE).map(|_| rng.gen::<u8>()).collect();
store.set(kv_pair.0.clone(), rand_value).unwrap();
}
drop(store);
bencher.bench_local(|| {
Rustcask::builder().open(temp_dir.path()).unwrap();
});
}