chore(test-suite): gw-stress support batch of ct handles (#1721)

Author: eudelins-zama

kms-connector/connector-db/docker-compose.yml

@@ -21,6 +21,8 @@ services:
     build:
       context: ../..
       dockerfile: kms-connector/connector-db/Dockerfile
+      args:
+        RUST_IMAGE_VERSION: stable
     environment:
       DATABASE_URL: postgresql://postgres:postgres@kms-connector-db:5432/kms-connector
     depends_on:

test-suite/gateway-stress/src/blockchain/manager.rs

@@ -160,12 +160,16 @@ impl GatewayTestManager {
         }
 
         burst_tasks.join_all().await;
+
         let elapsed = session_start.elapsed().as_secs_f64();
-        info!(
-            "Handled all burst in {:.2}s. Throughput: {:.2} tps",
-            elapsed,
-            (self.config.parallel_requests * (burst_index - 1) as u32) as f64 / elapsed
-        );
+        let handle_batch_size = match args.decryption_type {
+            DecryptionType::Public => self.config.public_ct.len() as u32,
+            DecryptionType::User => self.config.user_ct.len() as u32,
+        };
+        let total_decryption =
+            self.config.parallel_requests * handle_batch_size * (burst_index - 1) as u32;
+        let throughput = total_decryption as f64 / elapsed;
+        info!("Handled all burst in {elapsed:.2}s. Throughput: {throughput:.2} tps");
         Ok(())
     }
 

test-suite/gateway-stress/src/db/manager.rs

@@ -3,7 +3,10 @@ use crate::{
     cli::{DbBenchmarkArgs, DbTestArgs},
     config::Config,
     db::{DbConnector, RequestBuilder, ResponseTracker},
-    decryption::{BurstResult, types::DecryptionRequest},
+    decryption::{
+        BurstResult,
+        types::{DecryptionRequest, DecryptionType},
+    },
 };
 use anyhow::anyhow;
 use std::sync::Arc;
@@ -40,6 +43,8 @@ impl DatabaseTestManager {
             let connector = DbConnector::connect(db_config, i, request_sender).await?;
             let tracker = ResponseTracker::new(
                 connector.name.clone(),
+                config.public_ct.len(),
+                config.user_ct.len(),
                 request_receiver,
                 connector.db_pool.clone(),
             );
@@ -130,12 +135,16 @@ impl DatabaseTestManager {
         }
 
         burst_tasks.join_all().await;
+
         let elapsed = session_start.elapsed().as_secs_f64();
-        info!(
-            "Handled all burst in {:.2}s. Throughput: {:.2} tps",
-            elapsed,
-            (self.config.parallel_requests * (burst_index - 1) as u32) as f64 / elapsed
-        );
+        let handle_batch_size = match args.decryption_type {
+            DecryptionType::Public => self.config.public_ct.len() as u32,
+            DecryptionType::User => self.config.user_ct.len() as u32,
+        };
+        let total_decryption =
+            self.config.parallel_requests * handle_batch_size * (burst_index - 1) as u32;
+        let throughput = total_decryption as f64 / elapsed;
+        info!("Handled all burst in {elapsed:.2}s. Throughput: {throughput:.2} tps");
 
         if !args.skip_clear_db {
             self.clear_databases().await?;

test-suite/gateway-stress/src/db/response_tracker.rs

@@ -13,18 +13,24 @@ const RESPONSE_POLLING: Duration = Duration::from_millis(500);
 
 pub struct ResponseTracker {
     name: String,
+    public_decrypt_ct_handles_batch_size: usize,
+    user_decrypt_ct_handles_batch_size: usize,
     request_receiver: UnboundedReceiver<Vec<DecryptionRequestDbMetadata>>,
     db_pool: Pool<Postgres>,
 }
 
 impl ResponseTracker {
     pub fn new(
         name: String,
+        public_decrypt_ct_handles_batch_size: usize,
+        user_decrypt_ct_handles_batch_size: usize,
         request_receiver: UnboundedReceiver<Vec<DecryptionRequestDbMetadata>>,
         db_pool: Pool<Postgres>,
     ) -> Self {
         Self {
             name,
+            public_decrypt_ct_handles_batch_size,
+            user_decrypt_ct_handles_batch_size,
             request_receiver,
             db_pool,
         }
@@ -44,6 +50,7 @@ impl ResponseTracker {
 
         let burst_start = requests.iter().map(|r| r.created_at).min().unwrap();
         let mut interval = interval(RESPONSE_POLLING);
+
         let mut ids_to_process: Vec<Vec<u8>> = requests
             .into_iter()
             .map(|r| r.id.to_le_bytes_vec())
@@ -60,10 +67,11 @@ impl ResponseTracker {
 
         let burst_end = responses.iter().map(|r| r.created_at).max().unwrap();
         let latency = (burst_end - burst_start).as_seconds_f64();
+        let decryption_number: usize = responses.iter().map(|r| r.handle_batch_size).sum();
 
         let result = BurstResult {
             latency,
-            throughput: responses.len() as f64 / latency,
+            throughput: decryption_number as f64 / latency,
         };
         debug!(
             latency = result.latency,
@@ -91,6 +99,7 @@ impl ResponseTracker {
             responses.push(DecryptionResponseDbMetadata {
                 id: U256::from_le_slice(&row.decryption_id),
                 created_at: row.created_at,
+                handle_batch_size: self.public_decrypt_ct_handles_batch_size,
             });
         }
 
@@ -104,6 +113,7 @@ impl ResponseTracker {
             responses.push(DecryptionResponseDbMetadata {
                 id: U256::from_le_slice(&row.decryption_id),
                 created_at: row.created_at,
+                handle_batch_size: self.user_decrypt_ct_handles_batch_size,
             });
         }
 

test-suite/gateway-stress/src/db/types.rs

@@ -54,6 +54,7 @@ pub struct DecryptionRequestDbMetadata {
 pub struct DecryptionResponseDbMetadata {
     pub id: U256,
     pub created_at: PrimitiveDateTime,
+    pub handle_batch_size: usize,
 }
 
 impl From<PgRow> for DecryptionRequestDbMetadata {

test-suite/gateway-stress/src/decryption/public.rs

@@ -215,7 +215,7 @@ where
     let latency = burst_start.elapsed().as_secs_f64();
     let result = BurstResult {
         latency,
-        throughput: config.parallel_requests as f64 / latency,
+        throughput: (config.parallel_requests * config.public_ct.len() as u32) as f64 / latency,
     };
     progress_bar.finish_with_message(format!(
         "Handled burst #{} of {} in {:.2}s. Throughput: {:.2} tps",

test-suite/gateway-stress/src/decryption/user.rs

@@ -300,7 +300,7 @@ where
     let latency = burst_start.elapsed().as_secs_f64();
     let result = BurstResult {
         latency,
-        throughput: config.parallel_requests as f64 / latency,
+        throughput: (config.parallel_requests * config.user_ct.len() as u32) as f64 / latency,
     };
     progress_bar.finish_with_message(format!(
         "Handled burst #{} of {} in {:.2}s. Throughput: {:.2} tps",