257 lines
7.8 KiB
Rust
257 lines
7.8 KiB
Rust
use clap::Parser;
|
|
#[cfg(target_os = "linux")]
|
|
use logger::Logger;
|
|
use mace::{Mace, Options};
|
|
#[cfg(feature = "custom_alloc")]
|
|
use myalloc::{MyAlloc, print_filtered_trace};
|
|
use rand::prelude::*;
|
|
use std::path::Path;
|
|
use std::process::exit;
|
|
use std::sync::Arc;
|
|
use std::thread::JoinHandle;
|
|
use std::time::Instant;
|
|
|
|
#[cfg(feature = "custom_alloc")]
|
|
#[global_allocator]
|
|
static GLOBAL: MyAlloc = MyAlloc;
|
|
|
|
#[derive(Parser, Debug)]
|
|
#[command(author, version, about, long_about = None)]
|
|
struct Args {
|
|
#[arg(short = 'p', long, default_value = "/tmp/mace")]
|
|
path: String,
|
|
|
|
#[arg(short = 'm', long, default_value = "insert")]
|
|
mode: String,
|
|
|
|
#[arg(short = 'k', long, default_value = "16")]
|
|
key_size: usize,
|
|
|
|
#[arg(short = 'v', long, default_value = "1024")]
|
|
value_size: usize,
|
|
|
|
#[arg(short = 't', long, default_value = "4")]
|
|
threads: usize,
|
|
|
|
#[arg(short = 'i', long, default_value = "10000")]
|
|
iterations: usize,
|
|
|
|
#[arg(short = 'r', long, default_value = "30")]
|
|
insert_ratio: u8,
|
|
|
|
#[arg(long, default_value = "false")]
|
|
random: bool,
|
|
|
|
#[arg(long, default_value = "8192")]
|
|
blob_size: usize,
|
|
}
|
|
|
|
fn main() {
|
|
#[cfg(target_os = "linux")]
|
|
{
|
|
Logger::init().add_file("/tmp/x.log", true);
|
|
log::set_max_level(log::LevelFilter::Info);
|
|
}
|
|
let mut args = Args::parse();
|
|
|
|
let path = Path::new(&args.path);
|
|
|
|
if args.path.is_empty() {
|
|
eprintln!("path is empty");
|
|
exit(1);
|
|
}
|
|
|
|
if path.exists() {
|
|
eprintln!("path {:?} already exists", args.path);
|
|
exit(1);
|
|
}
|
|
|
|
if args.threads == 0 {
|
|
eprintln!("Error: threads must be greater than 0");
|
|
exit(1);
|
|
}
|
|
|
|
if !matches!(args.mode.as_str(), "insert" | "get" | "mixed" | "scan") {
|
|
eprintln!("Error: Invalid mode");
|
|
exit(1);
|
|
}
|
|
|
|
if args.key_size < 16 || args.value_size < 16 {
|
|
eprintln!("Error: key_size or value_size too small, must >= 16");
|
|
exit(1);
|
|
}
|
|
|
|
if args.insert_ratio > 100 {
|
|
eprintln!("Error: Insert ratio must be between 0 and 100");
|
|
exit(1);
|
|
}
|
|
|
|
let mut keys: Vec<Vec<Vec<u8>>> = Vec::with_capacity(args.threads);
|
|
let mut opt = Options::new(path);
|
|
opt.sync_on_write = false;
|
|
opt.over_provision = true; // large value will use lots of memeory
|
|
opt.inline_size = args.blob_size;
|
|
opt.tmp_store = args.mode != "get" && args.mode != "scan";
|
|
opt.cache_capacity = 3 << 30;
|
|
let mut saved = opt.clone();
|
|
saved.tmp_store = false;
|
|
let mut db = Mace::new(opt.validate().unwrap()).unwrap();
|
|
db.disable_gc();
|
|
let mut bkt = db.new_bucket("default").unwrap();
|
|
|
|
let mut rng = rand::rng();
|
|
let value = Arc::new(vec![b'0'; args.value_size]);
|
|
let mut key_counts = vec![args.iterations / args.threads; args.threads];
|
|
for cnt in key_counts.iter_mut().take(args.iterations % args.threads) {
|
|
*cnt += 1;
|
|
}
|
|
for tid in 0..args.threads {
|
|
let mut tk = Vec::with_capacity(key_counts[tid]);
|
|
for i in 0..key_counts[tid] {
|
|
let mut key = format!("key_{tid}_{i}").into_bytes();
|
|
key.resize(args.key_size, b'x');
|
|
tk.push(key);
|
|
}
|
|
if args.random || args.mode == "get" {
|
|
tk.shuffle(&mut rng);
|
|
}
|
|
keys.push(tk);
|
|
}
|
|
|
|
if args.mode == "get" || args.mode == "scan" {
|
|
let pre_tx = bkt.begin().unwrap();
|
|
(0..args.threads).for_each(|tid| {
|
|
for k in &keys[tid] {
|
|
pre_tx.put(k, &*value).unwrap();
|
|
}
|
|
});
|
|
pre_tx.commit().unwrap();
|
|
drop(bkt);
|
|
drop(db);
|
|
// re-open db
|
|
saved.tmp_store = true;
|
|
db = Mace::new(saved.validate().unwrap()).unwrap();
|
|
bkt = db.get_bucket("default").unwrap();
|
|
|
|
// simulate common use cases
|
|
for _ in 0..args.iterations {
|
|
let tid = rng.random_range(0..args.threads);
|
|
let Some(k) = keys[tid].choose(&mut rng) else {
|
|
continue;
|
|
};
|
|
let view = bkt.view().unwrap();
|
|
view.get(k).unwrap();
|
|
}
|
|
}
|
|
|
|
let ready_barrier = Arc::new(std::sync::Barrier::new(args.threads + 1));
|
|
let start_barrier = Arc::new(std::sync::Barrier::new(args.threads + 1));
|
|
let total_ops = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
|
|
|
let h: Vec<JoinHandle<()>> = (0..args.threads)
|
|
.map(|tid| {
|
|
let db = bkt.clone();
|
|
let tk: &Vec<Vec<u8>> = unsafe { std::mem::transmute(&keys[tid]) };
|
|
let total_ops = total_ops.clone();
|
|
let ready_barrier = Arc::clone(&ready_barrier);
|
|
let start_barrier = Arc::clone(&start_barrier);
|
|
let mode = args.mode.clone();
|
|
let insert_ratio = args.insert_ratio;
|
|
let val = value.clone();
|
|
let prefix = format!("key_{tid}_");
|
|
|
|
std::thread::spawn(move || {
|
|
coreid::bind_core(tid);
|
|
let mut round = 0;
|
|
ready_barrier.wait();
|
|
start_barrier.wait();
|
|
match mode.as_str() {
|
|
"insert" => {
|
|
for key in tk {
|
|
round += 1;
|
|
let tx = db.begin().unwrap();
|
|
tx.put(key.as_slice(), val.as_slice()).unwrap();
|
|
tx.commit().unwrap();
|
|
}
|
|
}
|
|
"get" => {
|
|
for key in tk {
|
|
round += 1;
|
|
let tx = db.view().unwrap();
|
|
let x = tx.get(key).unwrap();
|
|
std::hint::black_box(x);
|
|
}
|
|
}
|
|
"mixed" => {
|
|
for key in tk {
|
|
let is_insert = rand::random_range(0..100) < insert_ratio;
|
|
round += 1;
|
|
|
|
if is_insert {
|
|
let tx = db.begin().unwrap();
|
|
tx.put(key, &*val).unwrap();
|
|
tx.commit().unwrap();
|
|
} else {
|
|
let tx = db.view().unwrap();
|
|
let x = tx.get(key); // not found
|
|
let _ = std::hint::black_box(x);
|
|
}
|
|
}
|
|
}
|
|
"scan" => {
|
|
let view = db.view().unwrap();
|
|
let iter = view.seek(prefix);
|
|
for x in iter {
|
|
round += 1;
|
|
std::hint::black_box(x);
|
|
}
|
|
}
|
|
_ => panic!("Invalid mode"),
|
|
}
|
|
|
|
total_ops.fetch_add(round, std::sync::atomic::Ordering::Relaxed);
|
|
})
|
|
})
|
|
.collect();
|
|
|
|
ready_barrier.wait();
|
|
let start_time = Instant::now();
|
|
start_barrier.wait();
|
|
|
|
for x in h {
|
|
x.join().unwrap();
|
|
}
|
|
|
|
let duration = start_time.elapsed();
|
|
let total = total_ops.load(std::sync::atomic::Ordering::Relaxed);
|
|
let ops = (total as f64 / duration.as_secs_f64()) as usize;
|
|
|
|
let ratio = if args.mode == "mixed" {
|
|
args.insert_ratio
|
|
} else if args.mode == "insert" {
|
|
100
|
|
} else {
|
|
0
|
|
};
|
|
if args.mode == "insert" {
|
|
if args.random {
|
|
args.mode = "random_insert".into();
|
|
} else {
|
|
args.mode = "sequential_insert".into();
|
|
}
|
|
}
|
|
eprintln!(
|
|
"{},{},{},{},{},{},{}",
|
|
args.mode,
|
|
args.threads,
|
|
args.key_size,
|
|
args.value_size,
|
|
ratio,
|
|
ops,
|
|
duration.as_millis()
|
|
);
|
|
drop(db);
|
|
#[cfg(feature = "custom_alloc")]
|
|
print_filtered_trace(|x, y| log::info!("{}{}", x, y));
|
|
}
|