From 11e0cf842628685036265a6164295dd18d543978 Mon Sep 17 00:00:00 2001
From: Sayan Nandan <nandansayan@outlook.com>
Date: Sat, 7 Aug 2021 19:53:48 -0700
Subject: [PATCH] Use read_exact for all tests to ensure verity

---
 libstress/src/lib.rs       | 13 +++++++++++++
 sky-bench/src/benchtool.rs | 24 ++++++++++++++++++++----
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/libstress/src/lib.rs b/libstress/src/lib.rs
index a9a039b4..7852e3cc 100644
--- a/libstress/src/lib.rs
+++ b/libstress/src/lib.rs
@@ -180,6 +180,19 @@ where
             self.needs_iterator_pool,
         )
     }
+    /// Get a [`Workpool`] with the base config but with a custom loop-stage closure
+    pub fn with_loop_closure<Dlp>(&self, lp: Dlp) -> Workpool<Inp, UIn, Lv, Dlp, Ex>
+    where
+        Dlp: Fn(&mut Inp, UIn) + Clone + Send + Sync + 'static,
+    {
+        Workpool::new(
+            self.count,
+            self.init_pre_loop_var.clone(),
+            lp,
+            self.on_exit.clone(),
+            self.needs_iterator_pool,
+        )
+    }
 }
 
 impl<Inp: 'static, UIn, Lp, Lv, Ex> Clone for Workpool<Inp, UIn, Lv, Lp, Ex>
diff --git a/sky-bench/src/benchtool.rs b/sky-bench/src/benchtool.rs
index 26e27811..f79c29dd 100644
--- a/sky-bench/src/benchtool.rs
+++ b/sky-bench/src/benchtool.rs
@@ -45,6 +45,7 @@ const SIMPLE_QUERY_SIZE: usize = 3;
 /// &<n>\n
 /// (<tsymbol><size>\n<element>)*
 /// ```
+#[allow(dead_code)] // TODO(@ohsayan): Remove this lint
 pub fn calculate_array_dataframe_size(element_count: usize, per_element_size: usize) -> usize {
     let mut s = 0;
     s += 1; // `&`
@@ -90,6 +91,7 @@ fn test_monoelement_calculation() {
 /// ```text
 /// *<n>\n
 /// ```
+#[allow(dead_code)] // TODO(@ohsayan): Remove this lint
 pub fn calculate_metaframe_size(queries: usize) -> usize {
     let mut s = 0;
     s += 1; // `*`
@@ -98,6 +100,11 @@ pub fn calculate_metaframe_size(queries: usize) -> usize {
     s
 }
 
+#[test]
+fn test_simple_query_metaframe_size() {
+    assert_eq!(calculate_metaframe_size(1), SIMPLE_QUERY_SIZE);
+}
+
 /// Run the benchmark tool
 pub fn runner(
     host: String,
@@ -144,10 +151,11 @@ pub fn runner(
             let _ = stream.read_exact(&mut v).unwrap();
             stream
         },
-        |sock, packet: Vec<u8>| {
+        move |sock, packet: Vec<u8>| {
             sock.write_all(&packet).unwrap();
-            // we don't care much about what's returned
-            let _ = sock.read(&mut [0; 1024]).unwrap();
+            // all `okay`s are returned (for both update and set)
+            let mut v = vec![0; response_okay_size];
+            let _ = sock.read_exact(&mut v).unwrap();
         },
         |socket| {
             socket.shutdown(std::net::Shutdown::Both).unwrap();
@@ -196,8 +204,16 @@ pub fn runner(
     setpool.execute_and_finish_iter(set_packs);
     dt.stop_timer("SET").unwrap();
 
+    // TODO: Update the getpool to use correct sizes
     // bench GET
-    let getpool = pool_config.get_pool();
+    let get_response_packet_size =
+        calculate_monoelement_dataframe_size(per_kv_size) + SIMPLE_QUERY_SIZE;
+    let getpool = pool_config.with_loop_closure(move |sock: &mut TcpStream, packet: Vec<u8>| {
+        sock.write_all(&packet).unwrap();
+        // all `okay`s are returned (for both update and set)
+        let mut v = vec![0; get_response_packet_size];
+        let _ = sock.read(&mut v).unwrap();
+    });
     dt.create_timer("GET").unwrap();
     dt.start_timer("GET").unwrap();
     getpool.execute_and_finish_iter(get_packs);