test

.github/workflows/temp-build-and-push.yaml

@@ -0,0 +1,45 @@
+name: Branch - Build and push docker image
+
+on:
+  push:
+
+concurrency:
+  group: '${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}'
+  cancel-in-progress: true
+
+env:
+  REGISTRY: ghcr.io
+  IMAGE_NAME: ${{ github.repository }}
+
+jobs:
+  docker:
+    runs-on:
+      labels: ubuntu-22.04-64core
+    permissions:
+      packages: write
+      contents: read
+      attestations: write
+      id-token: write
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v3
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      - name: Log in to the Container registry
+        uses: docker/login-action@v3
+        with:
+          registry: ${{ env.REGISTRY }}
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+      - name: Build and Push
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          push: true
+          tags: |
+            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ github.sha }}
+          platforms: linux/amd64
+          cache-from: type=gha
+          cache-to: type=gha,mode=max

Cargo.toml

@@ -44,6 +44,8 @@ telemetry-batteries = { git = "https://github.com/worldcoin/telemetry-batteries.
 thiserror = "1"
 tokio = { version = "1.40", features = ["full", "rt-multi-thread"] }
 uuid = { version = "1", features = ["v4"] }
+opentelemetry = { version = "0.24.0" }
+tracing-opentelemetry = { version = "0.25.0"}
 
 # Abort on panics rather than unwinding.
 # This improves performance and makes panic propagation more reliable.

deploy/stage/common-values-iris-mpc.yaml

@@ -1,4 +1,4 @@
-image: "ghcr.io/worldcoin/iris-mpc:v0.8.26"
+image: "ghcr.io/worldcoin/iris-mpc:ecfcc4be02cbed9260c70e6baeab64948df7d59a"
 
 environment: stage
 replicaCount: 1

iris-mpc-common/Cargo.toml

@@ -28,6 +28,8 @@ config = "0.14.0"
 tokio.workspace = true
 tracing.workspace = true
 tracing-subscriber.workspace = true
+opentelemetry.workspace = true
+tracing-opentelemetry.workspace = true
 
 reqwest = { workspace = true, features = ["blocking", "json"] }
 sodiumoxide = "0.2.7"

iris-mpc-common/src/helpers/aws.rs

@@ -1,8 +1,5 @@
 use aws_sdk_sns::types::MessageAttributeValue;
 use std::collections::HashMap;
-use telemetry_batteries::reexports::opentelemetry::trace::{
-    SpanContext, SpanId, TraceFlags, TraceId, TraceState,
-};
 
 pub const TRACE_ID_MESSAGE_ATTRIBUTE_NAME: &str = "TraceID";
 pub const SPAN_ID_MESSAGE_ATTRIBUTE_NAME: &str = "SpanID";
@@ -36,41 +33,3 @@ pub fn construct_message_attributes(
 
     Ok(message_attributes)
 }
-
-// This would only ever be leveraged if the code had isolated flows for every
-// message, leaving for now, maybe it will happen in the future
-pub fn trace_from_message_attributes(
-    message_attributes: &HashMap<String, MessageAttributeValue>,
-    receipt_handle: &str,
-) -> eyre::Result<()> {
-    if let Some(trace_id) = message_attributes.get(TRACE_ID_MESSAGE_ATTRIBUTE_NAME) {
-        if let Some(span_id) = message_attributes.get(SPAN_ID_MESSAGE_ATTRIBUTE_NAME) {
-            let trace_id = trace_id
-                .string_value()
-                .expect("Could not parse TraceID")
-                .parse::<u128>()?;
-
-            let span_id = span_id
-                .string_value()
-                .expect("Could not parse SpanID")
-                .parse::<u64>()?;
-
-            // Create and set the span parent context
-            let parent_ctx = SpanContext::new(
-                TraceId::from(trace_id),
-                SpanId::from(span_id),
-                TraceFlags::default(),
-                true,
-                TraceState::default(),
-            );
-
-            telemetry_batteries::tracing::trace_from_ctx(parent_ctx);
-        } else {
-            tracing::warn!(?receipt_handle, "SQS message missing SpanID");
-        }
-    } else {
-        tracing::warn!(?receipt_handle, "SQS message missing TraceID");
-    }
-
-    Ok(())
-}

iris-mpc-common/src/helpers/mod.rs

@@ -7,3 +7,4 @@ pub mod smpc_request;
 pub mod sqs_s3_helper;
 pub mod sync;
 pub mod task_monitor;
+pub mod tracing;

iris-mpc-common/src/helpers/tracing.rs

@@ -0,0 +1,38 @@
+use opentelemetry::{
+    trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState},
+    Context,
+};
+use tracing::info_span;
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+
+pub fn trace_from_message_attributes(trace_id: &str, span_id: &str) -> eyre::Result<SpanContext> {
+    tracing::info!(
+        "Creating span context from message attributes. trace id: {}, span id: {}",
+        trace_id,
+        span_id
+    );
+
+    // Create and set the span parent context
+    let parent_span_ctx = SpanContext::new(
+        TraceId::from(trace_id.parse::<u128>()?),
+        SpanId::from(span_id.parse::<u64>()?),
+        TraceFlags::default(),
+        true,
+        TraceState::default(),
+    );
+    // let parent_ctx = Context::new().with_remote_span_context(parent_span_ctx);
+    // let item_span = info_span!("item", trace_id = trace_id, span_id = span_id);
+    // item_span.set_parent(parent_ctx.clone());
+    Ok(parent_span_ctx)
+}
+
+pub fn link_batch_spans(span_contexts: Vec<SpanContext>, batch_span: &tracing::Span) {
+    tracing::info!("Linking batch spans to item spans");
+    for span_ctx in span_contexts {
+        let parent_ctx = Context::new().with_remote_span_context(span_ctx);
+        let item_span = info_span!("item");
+        item_span.set_parent(parent_ctx.clone());
+        batch_span.follows_from(&item_span);
+        item_span.follows_from(batch_span);
+    }
+}

iris-mpc-gpu/Cargo.toml

@@ -31,6 +31,7 @@ iris-mpc-common = { path = "../iris-mpc-common" }
 base64 = "0.22.1"
 metrics = "0.22.1"
 metrics-exporter-statsd = "0.7"
+opentelemetry.workspace = true
 
 [dev-dependencies]
 criterion = "0.5"

iris-mpc-gpu/src/server/actor.rs

@@ -55,6 +55,7 @@ impl ServerActorHandle {
         &mut self,
         batch: BatchQuery,
     ) -> impl Future<Output = ServerJobResult> {
+        tracing::info!("Submitting batch query to the Actor job queue");
         let (tx, rx) = oneshot::channel();
         let job = ServerJob {
             batch,
@@ -464,6 +465,7 @@ impl ServerActor {
             .preprocess_db(&mut self.right_mask_db_slices, &self.current_db_sizes);
     }
 
+    #[tracing::instrument(skip(self, batch, return_channel))]
     fn process_batch_query(
         &mut self,
         batch: BatchQuery,
@@ -712,17 +714,7 @@ impl ServerActor {
 
         self.device_manager.await_streams(&self.streams[0]);
 
-        // Iterate over a list of tracing payloads, and create logs with mappings to
-        // payloads Log at least a "start" event using a log with trace.id
-        // and parent.trace.id
-        for tracing_payload in batch.metadata.iter() {
-            tracing::info!(
-                node_id = tracing_payload.node_id,
-                dd.trace_id = tracing_payload.trace_id,
-                dd.span_id = tracing_payload.span_id,
-                "Protocol finished",
-            );
-        }
+        tracing::info!("Protocol finished");
 
         // Fetch the final results (blocking)
         let mut host_results = self
@@ -907,6 +899,7 @@ impl ServerActor {
         Ok(())
     }
 
+    #[tracing::instrument(skip(self, compact_device_queries, compact_device_sums, events, eye_db))]
     fn compare_query_against_db_and_self(
         &mut self,
         compact_device_queries: &DeviceCompactQuery,
@@ -1190,6 +1183,7 @@ impl ServerActor {
         }
     }
 
+    #[tracing::instrument(skip(self, valid_entries))]
     fn sync_batch_entries(&mut self, valid_entries: &[bool]) -> eyre::Result<Vec<bool>> {
         let mut buffer = self
             .device_manager
@@ -1223,6 +1217,7 @@ impl ServerActor {
         Ok(valid_merged)
     }
 
+    #[tracing::instrument(skip(self))]
     fn prepare_deletion_shares(&self) -> eyre::Result<(DeviceCompactQuery, DeviceCompactSums)> {
         let (dummy_code_share, dummy_mask_share) = get_dummy_shares_for_deletion(self.party_id);
         let compact_query = {

iris-mpc-gpu/src/server/mod.rs

@@ -7,6 +7,7 @@ pub use actor::{get_dummy_shares_for_deletion, ServerActor, ServerActorHandle};
 use iris_mpc_common::galois_engine::degree4::{
     GaloisRingIrisCodeShare, GaloisRingTrimmedMaskCodeShare,
 };
+use opentelemetry::trace::SpanContext;
 use std::collections::HashSet;
 use tokio::sync::oneshot;
 
@@ -27,6 +28,7 @@ pub struct BatchMetadata {
 pub struct BatchQuery {
     pub request_ids:                Vec<String>,
     pub metadata:                   Vec<BatchMetadata>,
+    pub span_contexts:              Vec<SpanContext>,
     pub query_left:                 BatchQueryEntries,
     pub db_left:                    BatchQueryEntries,
     pub store_left:                 BatchQueryEntries,

iris-mpc/Cargo.toml

@@ -26,7 +26,8 @@ dotenvy.workspace = true
 rand.workspace = true
 base64.workspace = true
 uuid.workspace = true
-
+opentelemetry = { version = "0.24.0", features = ["trace"] }
+tracing-opentelemetry.workspace = true
 sodiumoxide = "0.2.7"
 iris-mpc-gpu = { path = "../iris-mpc-gpu" }
 iris-mpc-common = { path = "../iris-mpc-common" }