ack time, logging (#465)

worldcoin · Sep 30, 2024 · 59a1330 · 59a1330
1 parent 402d565
commit 59a1330
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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;