From 9fec85b9d9cd9e5f35e2c4078c2df76d090a1d30 Mon Sep 17 00:00:00 2001
From: Ertugrul Aypek <ertugrul.aypek@toolsforhumanity.com>
Date: Tue, 24 Sep 2024 16:48:05 +0200
Subject: [PATCH] Send trace info back to results sns (#445)

* send trace info back to results sns

* bump version
---
 deploy/stage/common-values-iris-mpc.yaml |  2 +-
 iris-mpc-common/src/helpers/aws.rs       | 11 ++++++-----
 iris-mpc-gpu/src/server/actor.rs         |  1 +
 iris-mpc-gpu/src/server/mod.rs           |  1 +
 iris-mpc/src/bin/server.rs               | 20 ++++++++++++++++----
 5 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/deploy/stage/common-values-iris-mpc.yaml b/deploy/stage/common-values-iris-mpc.yaml
index 6cc856552..9bbc7028e 100644
--- a/deploy/stage/common-values-iris-mpc.yaml
+++ b/deploy/stage/common-values-iris-mpc.yaml
@@ -1,4 +1,4 @@
-image: "ghcr.io/worldcoin/iris-mpc:v0.8.9"
+image: "ghcr.io/worldcoin/iris-mpc:v0.8.10"
 
 environment: stage
 replicaCount: 1
diff --git a/iris-mpc-common/src/helpers/aws.rs b/iris-mpc-common/src/helpers/aws.rs
index a5d687e5b..1447ece39 100644
--- a/iris-mpc-common/src/helpers/aws.rs
+++ b/iris-mpc-common/src/helpers/aws.rs
@@ -6,14 +6,15 @@ pub const TRACE_ID_MESSAGE_ATTRIBUTE_NAME: &str = "TraceID";
 pub const SPAN_ID_MESSAGE_ATTRIBUTE_NAME: &str = "SpanID";
 pub const NODE_ID_MESSAGE_ATTRIBUTE_NAME: &str = "NodeID";
 
-pub fn construct_message_attributes() -> eyre::Result<HashMap<String, MessageAttributeValue>> {
-    let (trace_id, span_id) = telemetry_batteries::tracing::extract_span_ids();
-
+pub fn construct_message_attributes(
+    trace_id: &String,
+    span_id: &String,
+) -> eyre::Result<HashMap<String, MessageAttributeValue>> {
     let mut message_attributes = HashMap::new();
 
     let trace_id_message_attribute = MessageAttributeValue::builder()
         .data_type("String")
-        .string_value(trace_id.to_string())
+        .string_value(trace_id)
         .build()?;
 
     message_attributes.insert(
@@ -23,7 +24,7 @@ pub fn construct_message_attributes() -> eyre::Result<HashMap<String, MessageAtt
 
     let span_id_message_attribute = MessageAttributeValue::builder()
         .data_type("String")
-        .string_value(span_id.to_string())
+        .string_value(span_id)
         .build()?;
 
     message_attributes.insert(
diff --git a/iris-mpc-gpu/src/server/actor.rs b/iris-mpc-gpu/src/server/actor.rs
index e227f50ff..3578e00aa 100644
--- a/iris-mpc-gpu/src/server/actor.rs
+++ b/iris-mpc-gpu/src/server/actor.rs
@@ -834,6 +834,7 @@ impl ServerActor {
             .send(ServerJobResult {
                 merged_results,
                 request_ids: batch.request_ids,
+                metadata: batch.metadata,
                 matches,
                 match_ids,
                 store_left: query_store_left,
diff --git a/iris-mpc-gpu/src/server/mod.rs b/iris-mpc-gpu/src/server/mod.rs
index 11248b0bf..51ed3beb7 100644
--- a/iris-mpc-gpu/src/server/mod.rs
+++ b/iris-mpc-gpu/src/server/mod.rs
@@ -91,6 +91,7 @@ pub struct ServerJob {
 pub struct ServerJobResult {
     pub merged_results: Vec<u32>,
     pub request_ids:    Vec<String>,
+    pub metadata:       Vec<BatchMetadata>,
     pub matches:        Vec<bool>,
     pub match_ids:      Vec<Vec<u32>>,
     pub store_left:     BatchQueryEntries,
diff --git a/iris-mpc/src/bin/server.rs b/iris-mpc/src/bin/server.rs
index f7b0cf341..ef44943b3 100644
--- a/iris-mpc/src/bin/server.rs
+++ b/iris-mpc/src/bin/server.rs
@@ -10,7 +10,10 @@ use iris_mpc_common::{
     config::{json_wrapper::JsonStrWrapper, Config, Opt},
     galois_engine::degree4::{GaloisRingIrisCodeShare, GaloisRingTrimmedMaskCodeShare},
     helpers::{
-        aws::{SPAN_ID_MESSAGE_ATTRIBUTE_NAME, TRACE_ID_MESSAGE_ATTRIBUTE_NAME},
+        aws::{
+            construct_message_attributes, SPAN_ID_MESSAGE_ATTRIBUTE_NAME,
+            TRACE_ID_MESSAGE_ATTRIBUTE_NAME,
+        },
         key_pair::SharesEncryptionKeyPairs,
         kms_dh::derive_shared_secret,
         smpc_request::{
@@ -448,17 +451,22 @@ async fn initialize_chacha_seeds(
 
 async fn send_results_to_sns(
     result_events: Vec<String>,
+    metadata: &[BatchMetadata],
     sns_client: &SNSClient,
     config: &Config,
-    message_attributes: &HashMap<String, MessageAttributeValue>,
+    base_message_attributes: &HashMap<String, MessageAttributeValue>,
 ) -> eyre::Result<()> {
-    for result_event in result_events {
+    for (i, result_event) in result_events.iter().enumerate() {
+        let trace_attributes =
+            construct_message_attributes(&metadata[i].trace_id, &metadata[i].span_id)?;
+        let mut message_attributes = base_message_attributes.clone();
+        message_attributes.extend(trace_attributes);
         sns_client
             .publish()
             .topic_arn(&config.results_topic_arn)
             .message(result_event)
             .message_group_id(format!("party-id-{}", config.party_id))
-            .set_message_attributes(Some(message_attributes.clone()))
+            .set_message_attributes(Some(message_attributes))
             .send()
             .await?;
     }
@@ -531,6 +539,7 @@ async fn server_main(config: Config) -> eyre::Result<()> {
     tracing::info!("Replaying results");
     send_results_to_sns(
         store.last_results(max_sync_lookback).await?,
+        &Vec::new(),
         &sns_client,
         &config,
         &uniqueness_result_attributes,
@@ -707,6 +716,7 @@ async fn server_main(config: Config) -> eyre::Result<()> {
         while let Some(ServerJobResult {
             merged_results,
             request_ids,
+            metadata,
             matches,
             match_ids,
             store_left,
@@ -774,6 +784,7 @@ async fn server_main(config: Config) -> eyre::Result<()> {
             tracing::info!("Sending {} uniqueness results", uniqueness_results.len());
             send_results_to_sns(
                 uniqueness_results,
+                &metadata,
                 &sns_client_bg,
                 &config_bg,
                 &uniqueness_result_attributes,
@@ -796,6 +807,7 @@ async fn server_main(config: Config) -> eyre::Result<()> {
             );
             send_results_to_sns(
                 identity_deletion_results,
+                &metadata,
                 &sns_client_bg,
                 &config_bg,
                 &identity_deletion_result_attributes,