Guard against allocation errors in `libstress`

next
Sayan Nandan 2 years ago
parent b4e1a94c9c
commit b3b236d9aa
No known key found for this signature in database
GPG Key ID: 8BC07A0A4D41DD52

@ -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

@ -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<Vec<u8>> {
) -> Result<Vec<Vec<u8>>, TryReserveError> {
if unique {
let mut keys: HashSet<Vec<u8>> = 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<String> {
) -> Result<Vec<String>, TryReserveError> {
if unique {
let mut keys: HashSet<String> = 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 {

@ -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<Data, Data> = 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<Data, Data> = 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<Data, Data> = keys
.iter()

@ -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`

@ -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)

@ -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<String> =
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<String> =
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();

Loading…
Cancel
Save