fix(coprocessor): restore tracing telemetry parity

Author: Eikix

coprocessor/fhevm-engine/fhevm-engine-common/src/telemetry.rs

@@ -206,6 +206,14 @@ pub fn register_histogram(config: Option<&MetricsConfig>, name: &str, desc: &str
         .unwrap_or_else(|_| panic!("Failed to register latency histogram: {}", name))
 }
 
+/// Returns the legacy short-form transaction id used by older telemetry helpers.
+pub fn short_txn_id(transaction_id: &[u8]) -> String {
+    to_hex(transaction_id)
+        .get(0..10)
+        .unwrap_or_default()
+        .to_owned()
+}
+
 pub(crate) static TXN_METRICS_MANAGER: LazyLock<TransactionMetrics> =
     LazyLock::new(|| TransactionMetrics::new(NonZeroUsize::new(100).unwrap()));
 

coprocessor/fhevm-engine/host-listener/src/database/tfhe_event_propagate.rs

@@ -274,7 +274,7 @@ impl Database {
     }
 
     #[rustfmt::skip]
-    #[tracing::instrument(skip_all, fields(operation = "handle_tfhe_event"))]
+    #[tracing::instrument(skip_all, fields(operation = "handle_tfhe_event", txn_id = tracing::field::Empty))]
     pub async fn insert_tfhe_event(
         &self,
         tx: &mut Transaction<'_>,
@@ -289,6 +289,12 @@ impl Database {
         let ty = |to_type: &ToType| vec![*to_type];
         let as_bytes = |x: &ClearConst| x.to_be_bytes_vec();
         let fhe_operation = event_to_op_int(event);
+        if let Some(transaction_hash) = log.transaction_hash.as_ref() {
+            tracing::Span::current().record(
+                "txn_id",
+                tracing::field::display(telemetry::short_txn_id(transaction_hash.as_ref())),
+            );
+        }
         let insert_computation = |tx, result, dependencies, scalar_byte| {
             self.insert_computation(tx, result, dependencies, fhe_operation, scalar_byte, log)
         };
@@ -436,7 +442,7 @@ impl Database {
     }
 
     /// Handles all types of ACL events
-    #[tracing::instrument(skip_all, fields(operation = "handle_acl_event"))]
+    #[tracing::instrument(skip_all, fields(operation = "handle_acl_event", txn_id = tracing::field::Empty))]
     pub async fn handle_acl_event(
         &self,
         tx: &mut Transaction<'_>,
@@ -447,6 +453,14 @@ impl Database {
         block_number: u64,
     ) -> Result<bool, SqlxError> {
         let data = &event.data;
+        if let Some(transaction_hash) = transaction_hash.as_ref() {
+            tracing::Span::current().record(
+                "txn_id",
+                tracing::field::display(telemetry::short_txn_id(
+                    transaction_hash.as_ref(),
+                )),
+            );
+        }
 
         let transaction_hash = transaction_hash.map(|h| h.to_vec());
 

coprocessor/fhevm-engine/sns-worker/src/aws_upload.rs

@@ -10,7 +10,7 @@ use aws_sdk_s3::Client;
 use bytesize::ByteSize;
 use fhevm_engine_common::chain_id::ChainId;
 use fhevm_engine_common::pg_pool::{PostgresPoolManager, ServiceError};
-use fhevm_engine_common::utils::to_hex;
+use fhevm_engine_common::{telemetry, utils::to_hex};
 use futures::future::join_all;
 use opentelemetry::trace::{Status, TraceContextExt};
 use sha3::{Digest, Keccak256};
@@ -537,18 +537,25 @@ async fn fetch_pending_uploads(
         };
 
         if !ct64_compressed.is_empty() || !is_ct128_empty {
+            let recovery_span = tracing::info_span!(
+                "recovery_task",
+                operation = "recovery_task",
+                txn_id = tracing::field::Empty
+            );
+            if let Some(transaction_id) = transaction_id.as_deref() {
+                recovery_span.record(
+                    "txn_id",
+                    tracing::field::display(telemetry::short_txn_id(transaction_id)),
+                );
+            }
             let item = HandleItem {
                 host_chain_id: ChainId::try_from(row.host_chain_id)
                     .map_err(|e| ExecutionError::ConversionError(e.into()))?,
                 key_id_gw: row.key_id_gw,
                 handle: handle.clone(),
                 ct64_compressed,
                 ct128: Arc::new(ct128),
-                otel: tracing::info_span!(
-                    "recovery_task",
-                    operation = "recovery_task",
-                    handle = %to_hex(&handle)
-                ),
+                otel: recovery_span,
                 transaction_id,
             };
 

coprocessor/fhevm-engine/sns-worker/src/executor.rs

@@ -308,11 +308,17 @@ pub async fn garbage_collect(pool: &PgPool, limit: u32) -> Result<(), ExecutionE
 
     info!(count, "Starting garbage collection of ciphertexts128");
 
-    // Limit the number of rows to update in case of a large backlog due to catchup or burst
-    // Skip Locked to prevent concurrent updates
-    let _start = SystemTime::now();
-    let rows_affected: u64 = sqlx::query!(
-        "
+    // Limit the number of rows to update in case of a large backlog due to catchup or burst.
+    // Skip locked to prevent concurrent updates.
+    let cleanup_span = tracing::info_span!(
+        "cleanup_ct128",
+        operation = "cleanup_ct128",
+        rows_affected = tracing::field::Empty
+    );
+    let rows_affected: u64 = async {
+        Ok::<u64, sqlx::Error>(
+            sqlx::query!(
+                "
         WITH uploaded_ct128 AS (
             SELECT c.handle
             FROM ciphertexts128 c
@@ -327,20 +333,22 @@ pub async fn garbage_collect(pool: &PgPool, limit: u32) -> Result<(), ExecutionE
         USING uploaded_ct128 r
         WHERE c.handle = r.handle;
         ",
-        limit as i32
-    )
-    .execute(pool)
-    .await?
-    .rows_affected();
+                limit as i32
+            )
+            .execute(pool)
+            .await?
+            .rows_affected(),
+        )
+    }
+    .instrument(cleanup_span.clone())
+    .await?;
+    cleanup_span.record("rows_affected", rows_affected as i64);
 
     if rows_affected > 0 {
-        let _cleanup_span = tracing::info_span!(
-            "cleanup_ct128",
-            operation = "cleanup_ct128",
-            rows_affected = rows_affected
-        )
-        .entered();
-        info!("Cleaning up old ciphertexts128");
+        info!(parent: &cleanup_span,
+            rows_affected = rows_affected,
+            "Cleaning up old ciphertexts128"
+        );
     }
 
     Ok(())
@@ -446,7 +454,11 @@ pub async fn query_sns_tasks(
     order: Order,
     key_id_gw: &DbKeyId,
 ) -> Result<Option<Vec<HandleItem>>, ExecutionError> {
-    let _start_time = SystemTime::now();
+    let fetch_span = tracing::info_span!(
+        "db_fetch_tasks",
+        operation = "db_fetch_tasks",
+        count = tracing::field::Empty
+    );
 
     let query = format!(
         "
@@ -466,7 +478,9 @@ pub async fn query_sns_tasks(
     let records = sqlx::query(&query)
         .bind(limit as i64)
         .fetch_all(db_txn.as_mut())
+        .instrument(fetch_span.clone())
         .await?;
+    fetch_span.record("count", records.len() as i64);
 
     info!(target: "worker", { count = records.len(), order = order.to_string() }, "Fetched SnS tasks");
     tracing::Span::current().record("count", records.len());
@@ -485,6 +499,14 @@ pub async fn query_sns_tasks(
             let handle: Vec<u8> = record.try_get("handle")?;
             let ciphertext: Vec<u8> = record.try_get("ciphertext")?;
             let transaction_id: Option<Vec<u8>> = record.try_get("transaction_id")?;
+            let task_span =
+                tracing::info_span!("task", operation = "task", txn_id = tracing::field::Empty);
+            if let Some(transaction_id) = transaction_id.as_deref() {
+                task_span.record(
+                    "txn_id",
+                    tracing::field::display(telemetry::short_txn_id(transaction_id)),
+                );
+            }
 
             Ok(HandleItem {
                 // TODO: During key rotation, ensure all coprocessors pin the same key_id_gw for a batch
@@ -494,11 +516,7 @@ pub async fn query_sns_tasks(
                 handle: handle.clone(),
                 ct64_compressed: Arc::new(ciphertext),
                 ct128: Arc::new(BigCiphertext::default()), // to be computed
-                otel: tracing::info_span!(
-                    "task",
-                    operation = "task",
-                    handle = %to_hex(&handle)
-                ),
+                otel: task_span,
                 transaction_id,
             })
         })