From ad7ceea0e82d62a61dbe6c23d7ae45acf2f63a57 Mon Sep 17 00:00:00 2001
From: Ertugrul Aypek <ertugrul.aypek@toolsforhumanity.com>
Date: Tue, 21 Jan 2025 12:53:25 +0100
Subject: [PATCH] move memory benchmark behind a flag

---
 iris-mpc-common/src/config/mod.rs |  3 +++
 iris-mpc-store/src/s3_importer.rs |  2 +-
 iris-mpc/src/bin/server.rs        | 31 +++++++++++++++++--------------
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/iris-mpc-common/src/config/mod.rs b/iris-mpc-common/src/config/mod.rs
index 83888abb0..980c1ff9a 100644
--- a/iris-mpc-common/src/config/mod.rs
+++ b/iris-mpc-common/src/config/mod.rs
@@ -120,6 +120,9 @@ pub struct Config {
 
     #[serde(default)]
     pub load_chunks_static_ips: usize,
+
+    #[serde(default)]
+    pub test_load_into_memory: bool,
 }
 
 fn default_load_chunks_parallelism() -> usize {
diff --git a/iris-mpc-store/src/s3_importer.rs b/iris-mpc-store/src/s3_importer.rs
index abc5ae0e5..97b7c61db 100644
--- a/iris-mpc-store/src/s3_importer.rs
+++ b/iris-mpc-store/src/s3_importer.rs
@@ -334,7 +334,7 @@ pub async fn fetch_to_memory(
                     match result.read_exact(&mut slice).await {
                         Ok(_) => {
                             let iris = S3StoredIris::from_bytes(&slice)?;
-                            if iris.id % 10000 == 0 {
+                            if iris.id % 100000 == 0 {
                                 tracing::info!("Fetched iris with id: {}", iris.id);
                             }
                         }
diff --git a/iris-mpc/src/bin/server.rs b/iris-mpc/src/bin/server.rs
index 983183927..940fa97b4 100644
--- a/iris-mpc/src/bin/server.rs
+++ b/iris-mpc/src/bin/server.rs
@@ -1181,20 +1181,23 @@ async fn server_main(config: Config) -> eyre::Result<()> {
                             db_chunks_folder_name.clone(),
                         )
                         .await?;
-                        let fetch_test_ts = Instant::now();
-                        fetch_to_memory(
-                            Arc::new(db_chunks_s3_store),
-                            load_chunks_parallelism,
-                            db_chunks_folder_name.clone(),
-                            last_snapshot_details.clone(),
-                        )
-                        .await?;
-                        let elapsed = fetch_test_ts.elapsed();
-                        tracing::info!(
-                            "Fetch to memory took {:?} with {} par",
-                            elapsed,
-                            load_chunks_parallelism
-                        );
+
+                        if config.test_load_into_memory {
+                            let fetch_test_ts = Instant::now();
+                            fetch_to_memory(
+                                Arc::new(db_chunks_s3_store),
+                                load_chunks_parallelism,
+                                db_chunks_folder_name.clone(),
+                                last_snapshot_details.clone(),
+                            )
+                            .await?;
+                            let elapsed = fetch_test_ts.elapsed();
+                            tracing::info!(
+                                "Fetch to memory took {:?} with {} par",
+                                elapsed,
+                                load_chunks_parallelism
+                            );
+                        }
 
                         let s3_store = S3Store::new(db_chunks_s3_client, db_chunks_bucket_name);
                         let s3_arc = Arc::new(s3_store);