diff --git a/kvstore/benches/basic.rs b/kvstore/benches/basic.rs index 19aee83d9a..db96c52392 100644 --- a/kvstore/benches/basic.rs +++ b/kvstore/benches/basic.rs @@ -4,11 +4,11 @@ extern crate test; use std::fs; use std::path::{Path, PathBuf}; -use rand::{self, thread_rng, Rng}; +use rand::{self, Rng}; use test::Bencher; -use solana_kvstore::{Config, Key, KvStore}; +use solana_kvstore::{test::gen, Config, Key, KvStore}; const SMALL_SIZE: usize = 512; const LARGE_SIZE: usize = 32 * 1024; @@ -44,7 +44,7 @@ fn bench_write_partitioned(bench: &mut Bencher, rows: &[(Key, Vec)], ledger_ fn bench_write_small(bench: &mut Bencher) { let ledger_path = setup("bench_write_small"); let num_entries = 32 * 1024; - let rows = gen_pairs(SMALL_SIZE).take(num_entries).collect::>(); + let rows = gen::pairs(SMALL_SIZE).take(num_entries).collect::>(); bench_write(bench, &rows, &ledger_path.to_string_lossy()); } @@ -53,7 +53,7 @@ fn bench_write_small(bench: &mut Bencher) { fn bench_write_small_partitioned(bench: &mut Bencher) { let ledger_path = setup("bench_write_small_partitioned"); let num_entries = 32 * 1024; - let rows = gen_pairs(SMALL_SIZE).take(num_entries).collect::>(); + let rows = gen::pairs(SMALL_SIZE).take(num_entries).collect::>(); bench_write_partitioned(bench, &rows, &ledger_path.to_string_lossy()); } @@ -62,7 +62,7 @@ fn bench_write_small_partitioned(bench: &mut Bencher) { fn bench_write_large(bench: &mut Bencher) { let ledger_path = setup("bench_write_large"); let num_entries = 32 * 1024; - let rows = gen_pairs(LARGE_SIZE).take(num_entries).collect::>(); + let rows = gen::pairs(LARGE_SIZE).take(num_entries).collect::>(); bench_write(bench, &rows, &ledger_path.to_string_lossy()); } @@ -71,7 +71,7 @@ fn bench_write_large(bench: &mut Bencher) { fn bench_write_huge(bench: &mut Bencher) { let ledger_path = setup("bench_write_huge"); let num_entries = 32 * 1024; - let rows = gen_pairs(HUGE_SIZE).take(num_entries).collect::>(); + let rows = gen::pairs(HUGE_SIZE).take(num_entries).collect::>(); bench_write(bench, &rows, &ledger_path.to_string_lossy()); } @@ -86,8 +86,8 @@ fn bench_read_sequential(bench: &mut Bencher) { let num_large_blobs = 32 * 1024; let total_blobs = num_small_blobs + num_large_blobs; - let small = gen_data(SMALL_SIZE).take(num_small_blobs); - let large = gen_data(LARGE_SIZE).take(num_large_blobs); + let small = gen::data(SMALL_SIZE).take(num_small_blobs); + let large = gen::data(LARGE_SIZE).take(num_large_blobs); let rows = gen_seq_keys().zip(small.chain(large)); let _ = store.put_many(rows); @@ -119,8 +119,8 @@ fn bench_read_random(bench: &mut Bencher) { let num_large_blobs = 32 * 1024; let total_blobs = num_small_blobs + num_large_blobs; - let small = gen_data(SMALL_SIZE).take(num_small_blobs); - let large = gen_data(LARGE_SIZE).take(num_large_blobs); + let small = gen::data(SMALL_SIZE).take(num_small_blobs); + let large = gen::data(LARGE_SIZE).take(num_large_blobs); let rows = gen_seq_keys().zip(small.chain(large)); let _ = store.put_many(rows); @@ -165,24 +165,6 @@ fn gen_seq_keys() -> impl Iterator { }) } -fn gen_keys() -> impl Iterator { - let mut rng = thread_rng(); - - std::iter::repeat_with(move || { - let buf = rng.gen(); - - Key(buf) - }) -} - -fn gen_data(size: usize) -> impl Iterator> { - std::iter::repeat(vec![1u8; size]) -} - -fn gen_pairs(data_size: usize) -> impl Iterator)> { - gen_keys().zip(gen_data(data_size)) -} - fn teardown>(p: P) { KvStore::destroy(p).expect("Expect successful store destruction"); } diff --git a/kvstore/src/lib.rs b/kvstore/src/lib.rs index f3ee853f69..8c629d6525 100644 --- a/kvstore/src/lib.rs +++ b/kvstore/src/lib.rs @@ -368,3 +368,41 @@ fn is_lvl0_full(tables: &[BTreeMap], config: &Config) -> bool { tables[0].len() > config.max_tables } } + +#[cfg(any(test, debug_assertions))] +pub mod test { + pub mod gen { + use crate::Key; + use rand::distributions::Uniform; + use rand::{rngs::SmallRng, FromEntropy, Rng}; + use std::iter; + use std::ops::Range; + + pub fn keys() -> impl Iterator { + let mut rng = SmallRng::from_entropy(); + iter::repeat_with(move || Key(rng.gen())) + } + + pub fn data(size: usize) -> impl Iterator> { + iter::repeat(vec![0; size]) + } + + pub fn data_vary(range: Range) -> impl Iterator> { + let dist = Uniform::from(range); + let mut rng = SmallRng::from_entropy(); + + iter::repeat_with(move || { + let size: u64 = rng.sample(dist); + vec![0; size as usize] + }) + } + + pub fn pairs(size: usize) -> impl Iterator)> { + keys().zip(data(size)) + } + + pub fn pairs_vary(range: Range) -> impl Iterator)> { + keys().zip(data_vary(range)) + } + } +} diff --git a/kvstore/src/sstable.rs b/kvstore/src/sstable.rs index e319bd5eb9..07bc39ae2c 100644 --- a/kvstore/src/sstable.rs +++ b/kvstore/src/sstable.rs @@ -410,9 +410,9 @@ fn overlapping(r1: &RangeInclusive, r2: &RangeInclusive) -> b } #[cfg(test)] -mod test { +pub mod test { use super::*; - use rand::{thread_rng, Rng}; + use crate::test::gen; use std::sync::{Arc, RwLock}; #[test] @@ -557,17 +557,8 @@ mod test { } fn gen_records() -> impl Iterator { - let mut rng = thread_rng(); - let commit = rng.gen(); - - std::iter::repeat_with(move || { - let buf: [u8; KEY_LEN] = rng.gen(); - let data_size: u8 = buf[0]; - - let val = Some(vec![0; data_size as usize]); - - (Key(buf), Value::new(commit, val)) - }) + gen::pairs_vary(0..255) + .map(|(key, bytes)| (key, Value::new(bytes.len() as i64, Some(bytes)))) } } diff --git a/kvstore/src/storage.rs b/kvstore/src/storage.rs index 186fa1d79d..ad4ac6640e 100644 --- a/kvstore/src/storage.rs +++ b/kvstore/src/storage.rs @@ -158,7 +158,7 @@ fn opt_bytes_memory(bytes: &Option>) -> usize { #[cfg(test)] mod test { use super::*; - use rand::{self, thread_rng, Rng}; + use crate::test::gen; const COMMIT: i64 = -1; @@ -168,7 +168,7 @@ mod test { let mut table = MemTable::default(); - for (key, data) in gen_pairs(DATA_SIZE).take(1024) { + for (key, data) in gen::pairs(DATA_SIZE).take(1024) { table.put(&key, COMMIT, &data); } @@ -181,7 +181,7 @@ mod test { const DATA_SIZE: usize = 32; let mut table = MemTable::default(); - let input = gen_pairs(DATA_SIZE).take(1024).collect::>(); + let input = gen::pairs(DATA_SIZE).take(1024).collect::>(); for (key, data) in &input { table.put(key, COMMIT, data); @@ -196,7 +196,7 @@ mod test { assert_eq!(table.mem_size, expected_size); // Deletes of things not in the memory table must be recorded - for key in gen_keys().take(512) { + for key in gen::keys().take(512) { table.delete(&key, COMMIT); } @@ -207,8 +207,8 @@ mod test { #[test] fn test_put_order_irrelevant() { let (mut table_1, mut table_2) = (MemTable::default(), MemTable::default()); - let big_input: Vec<_> = gen_pairs(1024).take(128).collect(); - let small_input: Vec<_> = gen_pairs(16).take(128).collect(); + let big_input: Vec<_> = gen::pairs(1024).take(128).collect(); + let small_input: Vec<_> = gen::pairs(16).take(128).collect(); for (key, data) in big_input.iter().chain(small_input.iter()) { table_1.put(key, COMMIT, data); @@ -237,8 +237,8 @@ mod test { #[test] fn test_delete_order_irrelevant() { let (mut table_1, mut table_2) = (MemTable::default(), MemTable::default()); - let big_input: Vec<_> = gen_pairs(1024).take(128).collect(); - let small_input: Vec<_> = gen_pairs(16).take(128).collect(); + let big_input: Vec<_> = gen::pairs(1024).take(128).collect(); + let small_input: Vec<_> = gen::pairs(16).take(128).collect(); for (key, data) in big_input.iter().chain(small_input.iter()) { table_1.put(key, COMMIT, data); @@ -277,22 +277,4 @@ mod test { assert_eq!(table_1.mem_size, table_2.mem_size); assert_eq!(table_1.values, table_2.values); } - - fn gen_keys() -> impl Iterator { - let mut rng = thread_rng(); - - std::iter::repeat_with(move || { - let buf = rng.gen(); - - Key(buf) - }) - } - - fn gen_data(size: usize) -> impl Iterator> { - std::iter::repeat(vec![1u8; size]) - } - - fn gen_pairs(data_size: usize) -> impl Iterator)> { - gen_keys().zip(gen_data(data_size)) - } } diff --git a/kvstore/src/writebatch.rs b/kvstore/src/writebatch.rs index 7e710f09ca..774cf7ad75 100644 --- a/kvstore/src/writebatch.rs +++ b/kvstore/src/writebatch.rs @@ -118,15 +118,15 @@ impl Default for Config { #[cfg(test)] mod test { use super::*; + use crate::test::gen; use crate::writelog::Config as WalConfig; - use rand::{self, thread_rng, Rng}; const CAPACITY: usize = 10 * 1024; #[test] fn test_put_associative() { let mut writebatch = setup(); - let input: Vec<_> = gen_pairs(32).take(100).collect(); + let input: Vec<_> = gen::pairs(32).take(100).collect(); writebatch.put_many(input.iter()).unwrap(); @@ -146,7 +146,7 @@ mod test { #[test] fn test_delete_associative() { let (mut writebatch, mut writebatch2) = (setup(), setup()); - let input: Vec<_> = gen_pairs(32).take(100).collect(); + let input: Vec<_> = gen::pairs(32).take(100).collect(); writebatch.put_many(input.iter()).unwrap(); writebatch2.put_many(input.iter()).unwrap(); @@ -172,7 +172,7 @@ mod test { let mut writebatch = setup(); let space_per_record = CAPACITY / AMT_RECORDS - MemTable::OVERHEAD_PER_RECORD; - let input: Vec<_> = gen_pairs(space_per_record).take(AMT_RECORDS).collect(); + let input: Vec<_> = gen::pairs(space_per_record).take(AMT_RECORDS).collect(); writebatch.put_many(input.iter()).unwrap(); @@ -181,7 +181,7 @@ mod test { _ => panic!("Writebatch should be exactly at capacity"), } - let (key, data) = gen_pairs(space_per_record).next().unwrap(); + let (key, data) = gen::pairs(space_per_record).next().unwrap(); let result = writebatch.put(&key, &data); assert!(result.is_err()); @@ -206,22 +206,4 @@ mod test { log: Arc::new(RwLock::new(log)), } } - - fn gen_keys() -> impl Iterator { - let mut rng = thread_rng(); - - std::iter::repeat_with(move || { - let buf = rng.gen(); - - Key(buf) - }) - } - - fn gen_data(size: usize) -> impl Iterator> { - std::iter::repeat(vec![1u8; size]) - } - - fn gen_pairs(data_size: usize) -> impl Iterator)> { - gen_keys().zip(gen_data(data_size)) - } } diff --git a/kvstore/tests/basic.rs b/kvstore/tests/basic.rs index ad4a408682..5ee0b42009 100644 --- a/kvstore/tests/basic.rs +++ b/kvstore/tests/basic.rs @@ -1,8 +1,7 @@ -use rand::{thread_rng, Rng}; - use std::fs; use std::path::{Path, PathBuf}; +use solana_kvstore::test::gen; use solana_kvstore::{Config, Key, KvStore}; const KB: usize = 1024; @@ -20,7 +19,7 @@ fn test_put_get() { }; let lsm = KvStore::open(&path, cfg).unwrap(); - let (key, bytes) = gen_pairs(HALF_KB).take(1).next().unwrap(); + let (key, bytes) = gen::pairs(HALF_KB).take(1).next().unwrap(); lsm.put(&key, &bytes).expect("put fail"); let out_bytes = lsm.get(&key).expect("get fail").expect("missing"); @@ -42,7 +41,7 @@ fn test_put_get_many() { }; let lsm = KvStore::open(&path, cfg).unwrap(); - let mut pairs: Vec<_> = gen_pairs(HALF_KB).take(1024).collect(); + let mut pairs: Vec<_> = gen::pairs(HALF_KB).take(1024).collect(); pairs.sort_unstable_by_key(|(k, _)| *k); lsm.put_many(pairs.clone().drain(..)) @@ -70,7 +69,7 @@ fn test_delete() { }; let lsm = KvStore::open(&path, cfg).unwrap(); - let mut pairs: Vec<_> = gen_pairs(HALF_KB).take(64 * 6).collect(); + let mut pairs: Vec<_> = gen::pairs(HALF_KB).take(64 * 6).collect(); pairs.sort_unstable_by_key(|(k, _)| *k); for (k, i) in pairs.iter() { @@ -104,7 +103,7 @@ fn test_delete_many() { }; let lsm = KvStore::open(&path, cfg).unwrap(); - let mut pairs: Vec<_> = gen_pairs(HALF_KB).take(64 * 6).collect(); + let mut pairs: Vec<_> = gen::pairs(HALF_KB).take(64 * 6).collect(); pairs.sort_unstable_by_key(|(k, _)| *k); for (k, i) in pairs.iter() { @@ -132,7 +131,7 @@ fn test_close_reopen() { let cfg = Config::default(); let lsm = KvStore::open(&path, cfg).unwrap(); - let mut pairs: Vec<_> = gen_pairs(KB).take(1024).collect(); + let mut pairs: Vec<_> = gen::pairs(KB).take(1024).collect(); pairs.sort_unstable_by_key(|(k, _)| *k); for (k, i) in pairs.iter() { @@ -174,7 +173,7 @@ fn test_partitioned() { let lsm = KvStore::partitioned(&path, &storage_dirs, cfg).unwrap(); - let mut pairs: Vec<_> = gen_pairs(HALF_KB).take(64 * 12).collect(); + let mut pairs: Vec<_> = gen::pairs(HALF_KB).take(64 * 12).collect(); pairs.sort_unstable_by_key(|(k, _)| *k); lsm.put_many(pairs.iter()).expect("put_many fail"); @@ -207,7 +206,7 @@ fn test_in_memory() { }; let lsm = KvStore::open(&path, cfg).unwrap(); - let mut pairs: Vec<_> = gen_pairs(HALF_KB).take(64 * 12).collect(); + let mut pairs: Vec<_> = gen::pairs(HALF_KB).take(64 * 12).collect(); pairs.sort_unstable_by_key(|(k, _)| *k); lsm.put_many(pairs.iter()).expect("put_many fail"); @@ -239,14 +238,3 @@ fn setup(test_name: &str) -> PathBuf { fn teardown(p: &Path) { KvStore::destroy(p).expect("Expect successful store destruction"); } - -fn gen_pairs(data_size: usize) -> impl Iterator)> { - let mut rng = thread_rng(); - - std::iter::repeat_with(move || { - let data = vec![0u8; data_size]; - let buf = rng.gen(); - - (Key(buf), data) - }) -}