From b3b236d9aa99c767079e265ca32d8624767ccf93 Mon Sep 17 00:00:00 2001 From: Sayan Nandan Date: Tue, 31 May 2022 06:58:27 -0700 Subject: [PATCH] Guard against allocation errors in `libstress` --- Makefile | 2 ++ libstress/src/lib.rs | 25 +++++++++++++++++-------- server/src/storage/v1/tests.rs | 12 ++++++------ sky-bench/src/benchtool/mod.rs | 4 ++++ sky-bench/src/testkey.rs | 9 +++++++-- stress-test/src/linearity_client.rs | 21 +++++++++++++++++---- 6 files changed, 53 insertions(+), 20 deletions(-) diff --git a/Makefile b/Makefile index ddc335e9..b195a8ce 100644 --- a/Makefile +++ b/Makefile @@ -2,6 +2,8 @@ BUILD_HARNESS:=cargo build -p harness RUN_HARNESS:= target/debug/harness .harness: @${BUILD_HARNESS} +build: + @cargo build test: .harness @${RUN_HARNESS} test bundle: .harness diff --git a/libstress/src/lib.rs b/libstress/src/lib.rs index 22bee825..71348860 100644 --- a/libstress/src/lib.rs +++ b/libstress/src/lib.rs @@ -365,6 +365,7 @@ pub mod utils { const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; use rand::distributions::{Alphanumeric, Standard}; use std::collections::HashSet; + use std::collections::TryReserveError; /// Generate a random UTF-8 string pub fn ran_string(len: usize, rand: impl rand::Rng) -> String { @@ -385,9 +386,10 @@ pub mod utils { size: usize, mut rng: impl rand::Rng, unique: bool, - ) -> Vec> { + ) -> Result>, TryReserveError> { if unique { - let mut keys: HashSet> = HashSet::with_capacity(count); + let mut keys = HashSet::new(); + keys.try_reserve(size)?; (0..count).into_iter().for_each(|_| { let mut ran = ran_bytes(size, &mut rng); while keys.contains(&ran) { @@ -395,12 +397,15 @@ pub mod utils { } keys.insert(ran); }); - keys.into_iter().collect() + Ok(keys.into_iter().collect()) } else { + let mut keys = Vec::new(); + keys.try_reserve_exact(size)?; (0..count) .into_iter() .map(|_| ran_bytes(size, &mut rng)) - .collect() + .for_each(|bytes| keys.push(bytes)); + Ok(keys) } } /// Generate a vector of random UTF-8 valid strings @@ -409,9 +414,10 @@ pub mod utils { size: usize, mut rng: impl rand::Rng, unique: bool, - ) -> Vec { + ) -> Result, TryReserveError> { if unique { - let mut keys: HashSet = HashSet::with_capacity(count); + let mut keys = HashSet::new(); + keys.try_reserve(size)?; (0..count).into_iter().for_each(|_| { let mut ran = ran_string(size, &mut rng); while keys.contains(&ran) { @@ -419,12 +425,15 @@ pub mod utils { } keys.insert(ran); }); - keys.into_iter().collect() + Ok(keys.into_iter().collect()) } else { + let mut keys = Vec::new(); + keys.try_reserve_exact(size)?; (0..count) .into_iter() .map(|_| ran_string(size, &mut rng)) - .collect() + .for_each(|bytes| keys.push(bytes)); + Ok(keys) } } pub fn rand_alphastring(len: usize, rng: &mut impl rand::Rng) -> String { diff --git a/server/src/storage/v1/tests.rs b/server/src/storage/v1/tests.rs index 7b69292e..7337f291 100644 --- a/server/src/storage/v1/tests.rs +++ b/server/src/storage/v1/tests.rs @@ -69,8 +69,8 @@ cfg_test!( const LEN: usize = 8_usize; let mut rng = thread_rng(); let (keys, values) = ( - generate_random_string_vector(COUNT, LEN, &mut rng, true), - generate_random_string_vector(COUNT, LEN, &mut rng, false), + generate_random_string_vector(COUNT, LEN, &mut rng, true).unwrap(), + generate_random_string_vector(COUNT, LEN, &mut rng, false).unwrap(), ); let cmap: Coremap = keys .iter() @@ -91,8 +91,8 @@ cfg_test!( const LEN: usize = 8_usize; let mut rng = thread_rng(); let (keys, values) = ( - generate_random_string_vector(COUNT, LEN, &mut rng, true), - generate_random_string_vector(COUNT, LEN, &mut rng, false), + generate_random_string_vector(COUNT, LEN, &mut rng, true).unwrap(), + generate_random_string_vector(COUNT, LEN, &mut rng, false).unwrap(), ); let cmap: Coremap = keys .iter() @@ -115,8 +115,8 @@ cfg_test!( const LEN: usize = 8_usize; let mut rng = thread_rng(); let (keys, values) = ( - generate_random_string_vector(COUNT, LEN, &mut rng, true), - generate_random_string_vector(COUNT, LEN, &mut rng, false), + generate_random_string_vector(COUNT, LEN, &mut rng, true).unwrap(), + generate_random_string_vector(COUNT, LEN, &mut rng, false).unwrap(), ); let cmap: Coremap = keys .iter() diff --git a/sky-bench/src/benchtool/mod.rs b/sky-bench/src/benchtool/mod.rs index f2b23ac5..73ff2ba1 100644 --- a/sky-bench/src/benchtool/mod.rs +++ b/sky-bench/src/benchtool/mod.rs @@ -103,6 +103,10 @@ pub fn runner( let keys = generate_random_byte_vector(max_queries, per_kv_size, &mut rand, true); let values = generate_random_byte_vector(max_queries, per_kv_size, &mut rand, false); + let (keys, values) = match (keys, values) { + (Ok(k), Ok(v)) => (k, v), + _ => err!("Allocation error"), + }; /* We create three vectors of vectors: `set_packs`, `get_packs` and `del_packs` diff --git a/sky-bench/src/testkey.rs b/sky-bench/src/testkey.rs index 44bb47dc..0f4fb5e4 100644 --- a/sky-bench/src/testkey.rs +++ b/sky-bench/src/testkey.rs @@ -25,6 +25,7 @@ */ use crate::{benchtool::validation::SQ_RESPCODE_SIZE, hoststr, sanity_test}; +use libstress::utils::generate_random_string_vector; use libstress::Workpool; use rand::thread_rng; use skytable::Query; @@ -53,8 +54,12 @@ pub fn create_testkeys(host: &str, port: u16, num: usize, connections: usize, si Some(connections), ); println!("Generating keys ..."); - let keys = libstress::utils::generate_random_string_vector(num, size, &mut rand, true); - let values = libstress::utils::generate_random_string_vector(num, size, &mut rand, false); + let keys = generate_random_string_vector(num, size, &mut rand, true); + let values = generate_random_string_vector(num, size, &mut rand, false); + let (keys, values) = match (keys, values) { + (Ok(k), Ok(v)) => (k, v), + _ => err!("Allocation error"), + }; { let np = np; (0..num) diff --git a/stress-test/src/linearity_client.rs b/stress-test/src/linearity_client.rs index c98f0e9f..752ff119 100644 --- a/stress-test/src/linearity_client.rs +++ b/stress-test/src/linearity_client.rs @@ -112,9 +112,15 @@ pub fn stress_linearity_concurrent_clients_set( // generate the random k/v pairs let keys = generate_random_string_vector(DEFAULT_QUERY_COUNT, DEFAULT_SIZE_KV, &mut rng, true); - let values: Vec = + let values = generate_random_string_vector(DEFAULT_QUERY_COUNT, DEFAULT_SIZE_KV, &mut rng, false); - + let (keys, values) = match (keys, values) { + (Ok(k), Ok(v)) => (k, v), + _ => { + eprintln!("Allocation error"); + std::process::exit(0x01); + } + }; // make sure the database is empty temp_con.flushdb().unwrap(); @@ -166,10 +172,17 @@ pub fn stress_linearity_concurrent_clients_get( ); let mut current_thread_count = 1usize; - // Generate the random k/v pairs + // generate the random k/v pairs let keys = generate_random_string_vector(DEFAULT_QUERY_COUNT, DEFAULT_SIZE_KV, &mut rng, true); - let values: Vec = + let values = generate_random_string_vector(DEFAULT_QUERY_COUNT, DEFAULT_SIZE_KV, &mut rng, false); + let (keys, values) = match (keys, values) { + (Ok(k), Ok(v)) => (k, v), + _ => { + eprintln!("Allocation error"); + std::process::exit(0x01); + } + }; // Make sure that the database is empty temp_con.flushdb().unwrap();