Guard against allocation errors in `libstress`

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

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

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, true).unwrap(),
-            generate_random_string_vector(COUNT, LEN, &mut rng, false),
+            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, true).unwrap(),
-            generate_random_string_vector(COUNT, LEN, &mut rng, false),
+            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, true).unwrap(),
-            generate_random_string_vector(COUNT, LEN, &mut rng, false),
+            generate_random_string_vector(COUNT, LEN, &mut rng, false).unwrap(),
         );
         let cmap: Coremap<Data, Data> = keys
             .iter()

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`

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 keys = generate_random_string_vector(num, size, &mut rand, true);
-    let values = libstress::utils::generate_random_string_vector(num, size, &mut rand, false);
+    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)

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<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();