mod common;
use std::collections::BTreeMap;
use grebedb::{Database, Error};
use indexmap::IndexSet;
use rand::{distributions::Uniform, prelude::*, SeedableRng};
use rand_xorshift::XorShiftRng;
enum OperationChoice {
Get,
Put,
Remove,
Flush,
}
enum Operation {
Get(Vec<u8>),
Put(Vec<u8>, Vec<u8>),
Remove(Vec<u8>),
Flush,
}
struct OperationGenerator {
rng: XorShiftRng,
key_range: Uniform<u64>,
existing_keys: IndexSet<u64>,
choices: Vec<OperationChoice>,
}
impl OperationGenerator {
fn new(max_keys: usize) -> Self {
let mut choices = Vec::new();
for _ in 0..200 {
choices.push(OperationChoice::Get);
}
for _ in 0..700 {
choices.push(OperationChoice::Put);
}
for _ in 0..100 {
choices.push(OperationChoice::Remove);
}
choices.push(OperationChoice::Flush);
Self {
rng: XorShiftRng::seed_from_u64(1),
key_range: Uniform::from(0..max_keys as u64),
existing_keys: IndexSet::new(),
choices,
}
}
pub fn get(&mut self) -> Operation {
let operation_seed: u64 = self.rng.gen();
let mut operation_rng = XorShiftRng::seed_from_u64(operation_seed);
let choice = self.choices.choose(&mut operation_rng).unwrap();
match choice {
OperationChoice::Get => {
if !self.existing_keys.is_empty() && self.rng.gen_bool(0.6) {
let key = self.get_existing_key(&mut operation_rng);
let key_vec = Self::derive_key_bytes(key);
Operation::Get(key_vec)
} else {
let key = self.get_nonexisting_key(&mut operation_rng);
let key_vec = Self::derive_key_bytes(key);
Operation::Get(key_vec)
}
}
OperationChoice::Put => {
if !self.existing_keys.is_empty() && self.rng.gen_bool(0.6) {
let key = self.get_existing_key(&mut operation_rng);
let key_vec = Self::derive_key_bytes(key);
let value_vec = Self::derive_value_bytes(key, operation_seed);
Operation::Put(key_vec, value_vec)
} else {
let key = self.generate_new_key(&mut operation_rng);
let key_vec = Self::derive_key_bytes(key);
let value_vec = Self::derive_value_bytes(key, operation_seed);
Operation::Put(key_vec, value_vec)
}
}
OperationChoice::Remove => {
if !self.existing_keys.is_empty() && self.rng.gen_bool(0.8) {
let index = self.get_existing_index(&mut operation_rng);
let key = self.existing_keys.swap_remove_index(index).unwrap();
let key_vec = Self::derive_key_bytes(key);
Operation::Remove(key_vec)
} else {
let key = self.get_nonexisting_key(&mut operation_rng);
let key_vec = Self::derive_key_bytes(key);
Operation::Remove(key_vec)
}
}
OperationChoice::Flush => Operation::Flush,
}
}
fn get_existing_index(&mut self, rng: &mut impl Rng) -> usize {
rng.gen_range(0..self.existing_keys.len())
}
fn get_existing_key(&mut self, rng: &mut impl Rng) -> u64 {
let index = self.get_existing_index(rng);
*self.existing_keys.get_index(index).unwrap()
}
fn get_nonexisting_key(&mut self, rng: &mut impl Rng) -> u64 {
loop {
let key = self.key_range.sample(rng);
if !self.existing_keys.contains(&key) {
return key;
}
}
}
fn generate_new_key(&mut self, rng: &mut impl Rng) -> u64 {
let key = self.key_range.sample(rng);
self.existing_keys.insert(key);
key
}
fn derive_key_bytes(key: u64) -> Vec<u8> {
format!("my key {:016x}", key).into_bytes()
}
fn derive_value_bytes(key: u64, seed: u64) -> Vec<u8> {
let mut value_rng = XorShiftRng::seed_from_u64(key ^ seed);
let size = value_rng.gen_range(1..=4096);
let mut value = vec![0; size];
value_rng.fill_bytes(&mut value);
value
}
}
fn rand_operation(mut database: Database, rounds: usize) -> Result<(), Error> {
let mut generator = OperationGenerator::new((rounds / 4).max(10));
let mut std_map = BTreeMap::<Vec<u8>, Vec<u8>>::new();
let mut value_buffer = Vec::new();
for num in 0..rounds {
match generator.get() {
Operation::Get(key) => {
let has_key = database.get_buf(&key, &mut value_buffer)?;
match std_map.get(&key) {
Some(value) => {
assert!(has_key);
assert_eq!(&value_buffer, value);
}
None => {
assert!(!has_key);
}
}
}
Operation::Put(key, value) => {
database.put(key.clone(), value.clone())?;
assert!(database.contains_key(&key)?);
std_map.insert(key, value);
}
Operation::Remove(key) => {
database.remove(key.clone())?;
assert!(!database.contains_key(&key)?);
std_map.remove(&key);
}
Operation::Flush => database.flush()?,
}
if num % 1000 == 0 {
database.verify(|_, _| {})?;
}
}
database.verify(|_, _| {})?;
println!(
"current len={}, expected len={}",
database.cursor()?.count(),
std_map.len()
);
let mut cursor = database.cursor()?;
let mut std_iter = std_map.iter();
loop {
let current = cursor.next();
let expected = std_iter.next();
if current.is_none() && expected.is_none() {
break;
} else {
let (key, value) = current.unwrap();
let (expected_key, expected_value) = expected.unwrap();
assert_eq!(&key, expected_key);
assert_eq!(&value, expected_value);
}
}
Ok(())
}
fn rand_operation_10000(database: Database) -> Result<(), Error> {
rand_operation(database, 10000)
}
fn rand_operation_100000(database: Database) -> Result<(), Error> {
rand_operation(database, 100000)
}
#[cfg(debug_assertions)]
mod a {
use grebedb::Options;
use super::*;
#[test]
fn rand_operation_10000_fast() {
let database = Database::open_memory(Options::default()).unwrap();
rand_operation(database, 10000).unwrap();
}
matrix_test_ignore!(rand_operation_10000);
matrix_test_ignore!(rand_operation_100000);
}
#[cfg(not(debug_assertions))]
mod a {
use super::*;
matrix_test!(rand_operation_10000);
matrix_test_ignore!(rand_operation_100000);
}