From 59a13309efbf33c4e6891db3a193cada2644bdb6 Mon Sep 17 00:00:00 2001
From: Carlo Mazzaferro <carlo.mazzaferro@gmail.com>
Date: Mon, 30 Sep 2024 15:27:31 +0200
Subject: [PATCH] ack time, logging (#465)

---
 iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs | 13 ++++++++-----
 iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs |  6 ++++++
 iris-mpc-upgrade/src/config.rs                 |  1 +
 3 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs b/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs
index 95394c7ab..5bce2521b 100644
--- a/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs
+++ b/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs
@@ -7,7 +7,10 @@ use iris_mpc_common::{
     IRIS_CODE_LENGTH,
 };
 use iris_mpc_upgrade::{
-    config::{UpgradeClientConfig, BATCH_SUCCESSFUL_ACK, FINAL_BATCH_SUCCESSFUL_ACK},
+    config::{
+        UpgradeClientConfig, BATCH_SUCCESSFUL_ACK, BATCH_TIMEOUT_SECONDS,
+        FINAL_BATCH_SUCCESSFUL_ACK,
+    },
     db::V1Db,
     packets::{MaskShareMessage, TwoToThreeIrisCodeMessage},
     OldIrisShareSource,
@@ -329,7 +332,7 @@ async fn send_batch_and_wait_for_ack(
 }
 
 async fn wait_for_ack(server: &mut TlsStream<TcpStream>) -> eyre::Result<()> {
-    match timeout(Duration::from_secs(10), server.read_u8()).await {
+    match timeout(Duration::from_secs(BATCH_TIMEOUT_SECONDS), server.read_u8()).await {
         Ok(Ok(BATCH_SUCCESSFUL_ACK)) => {
             // Ack received successfully
             tracing::info!("ACK received for batch");
@@ -505,7 +508,7 @@ impl OldIrisShareSource for MockOldDbParty1 {
     fn stream_shares(
         &self,
         share_id_range: std::ops::Range<u64>,
-    ) -> eyre::Result<impl futures::Stream<Item = eyre::Result<(u64, EncodedBits)>>> {
+    ) -> eyre::Result<impl Stream<Item = eyre::Result<(u64, EncodedBits)>>> {
         let mut id = share_id_range.start;
         Ok(futures::stream::poll_fn(move |_| {
             if id < share_id_range.end {
@@ -564,7 +567,7 @@ impl OldIrisShareSource for MockOldDbParty2 {
     fn stream_shares(
         &self,
         share_id_range: std::ops::Range<u64>,
-    ) -> eyre::Result<impl futures::Stream<Item = eyre::Result<(u64, EncodedBits)>>> {
+    ) -> eyre::Result<impl Stream<Item = eyre::Result<(u64, EncodedBits)>>> {
         let mut id = share_id_range.start;
         Ok(futures::stream::poll_fn(move |_| {
             if id < share_id_range.end {
@@ -584,7 +587,7 @@ impl OldIrisShareSource for MockOldDbParty2 {
     fn stream_masks(
         &self,
         share_id_range: std::ops::Range<u64>,
-    ) -> eyre::Result<impl futures::Stream<Item = eyre::Result<(u64, Bits)>>> {
+    ) -> eyre::Result<impl Stream<Item = eyre::Result<(u64, Bits)>>> {
         let mut id = share_id_range.start;
         Ok(futures::stream::poll_fn(move |_| {
             if id >= share_id_range.end {
diff --git a/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs b/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
index 7e7e80e85..88bbe63ee 100644
--- a/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
+++ b/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
@@ -189,10 +189,16 @@ async fn main() -> eyre::Result<()> {
         let duration = start_time.elapsed();
         tracing::info!("Processed batch in {:.2?}", duration);
     }
+
+    tracing::info!("Finalizing upgrade");
     client_stream2.write_u8(FINAL_BATCH_SUCCESSFUL_ACK).await?;
+    tracing::info!("Sent final ACK to client2");
     client_stream1.write_u8(FINAL_BATCH_SUCCESSFUL_ACK).await?;
+    tracing::info!("Sent final ACK to client1");
 
+    tracing::info!("Updating iris id sequence");
     sink.update_iris_id_sequence().await?;
+    tracing::info!("Iris id sequence updated");
 
     Ok(())
 }
diff --git a/iris-mpc-upgrade/src/config.rs b/iris-mpc-upgrade/src/config.rs
index d9effa5bf..d6a1e0a9e 100644
--- a/iris-mpc-upgrade/src/config.rs
+++ b/iris-mpc-upgrade/src/config.rs
@@ -6,6 +6,7 @@ use std::{
     str::FromStr,
 };
 
+pub const BATCH_TIMEOUT_SECONDS: u64 = 30;
 pub const BATCH_SUCCESSFUL_ACK: u8 = 1;
 pub const FINAL_BATCH_SUCCESSFUL_ACK: u8 = 42;