refactor(coprocessor): address review feedback on tracing style

Author: Eikix

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

@@ -50,9 +50,8 @@ impl TryFrom<u64> for ChainId {
     type Error = InvalidChainId;
 
     fn try_from(value: u64) -> Result<Self, Self::Error> {
-        #[allow(clippy::checked_conversions)]
-        if value <= i64::MAX as u64 {
-            Ok(ChainId(value as i64))
+        if let Ok(value) = i64::try_from(value) {
+            Ok(ChainId(value))
         } else {
             Err(InvalidChainId {
                 value: value.to_string(),

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

@@ -544,7 +544,11 @@ async fn fetch_pending_uploads(
                 handle: handle.clone(),
                 ct64_compressed,
                 ct128: Arc::new(ct128),
-                otel: tracing::info_span!("recovery_task", operation = "recovery_task"),
+                otel: tracing::info_span!(
+                    "recovery_task",
+                    operation = "recovery_task",
+                    handle = %to_hex(&handle)
+                ),
                 transaction_id,
             };
 

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

@@ -176,11 +176,11 @@ impl SwitchNSquashService {
     }
 }
 
+#[tracing::instrument(skip_all, fields(operation = "fetch_keyset"))]
 async fn get_keyset(
     pool: PgPool,
     keys_cache: Arc<RwLock<lru::LruCache<DbKeyId, KeySet>>>,
 ) -> Result<Option<(DbKeyId, KeySet)>, ExecutionError> {
-    let _t = tracing::info_span!("fetch_keyset", operation = "fetch_keyset");
     fetch_latest_keyset(&keys_cache, &pool).await
 }
 
@@ -334,11 +334,13 @@ pub async fn garbage_collect(pool: &PgPool, limit: u32) -> Result<(), ExecutionE
     .rows_affected();
 
     if rows_affected > 0 {
-        let _s = tracing::info_span!("cleanup_ct128", operation = "cleanup_ct128");
-        info!(
-            rows_affected = rows_affected,
-            "Cleaning up old ciphertexts128"
-        );
+        let _cleanup_span = tracing::info_span!(
+            "cleanup_ct128",
+            operation = "cleanup_ct128",
+            rows_affected = rows_affected
+        )
+        .entered();
+        info!("Cleaning up old ciphertexts128");
     }
 
     Ok(())
@@ -377,25 +379,29 @@ async fn fetch_and_execute_sns_tasks(
         maybe_remaining = conf.db.batch_limit as usize == tasks.len();
         tasks_processed = tasks.len();
 
-        let t = tracing::info_span!(
+        let batch_exec_span = tracing::info_span!(
             "batch_execution",
             operation = "batch_execution",
             count = tasks.len()
         );
 
-        process_tasks(
-            &mut tasks,
-            keys,
-            tx,
-            conf.enable_compression,
-            conf.schedule_policy,
-            token.clone(),
-        )?;
+        batch_exec_span.in_scope(|| {
+            process_tasks(
+                &mut tasks,
+                keys,
+                tx,
+                conf.enable_compression,
+                conf.schedule_policy,
+                token.clone(),
+            )
+        })?;
 
-        update_computations_status(trx, &tasks).await?;
+        update_computations_status(trx, &tasks)
+            .instrument(batch_exec_span.clone())
+            .await?;
 
         let batch_store_span = tracing::info_span!(
-            parent: &t,
+            parent: &batch_exec_span,
             "batch_store_ciphertext128",
             operation = "batch_store_ciphertext128"
         );
@@ -433,6 +439,7 @@ async fn fetch_and_execute_sns_tasks(
 }
 
 /// Queries the database for a fixed number of tasks.
+#[tracing::instrument(skip_all, fields(operation = "db_fetch_tasks", count = tracing::field::Empty))]
 pub async fn query_sns_tasks(
     db_txn: &mut Transaction<'_, Postgres>,
     limit: u32,
@@ -462,19 +469,12 @@ pub async fn query_sns_tasks(
         .await?;
 
     info!(target: "worker", { count = records.len(), order = order.to_string() }, "Fetched SnS tasks");
+    tracing::Span::current().record("count", records.len());
 
     if records.is_empty() {
         return Ok(None);
     }
 
-    {
-        let _t = tracing::info_span!(
-            "db_fetch_tasks",
-            operation = "db_fetch_tasks",
-            count = records.len()
-        );
-    }
-
     // Convert the records into HandleItem structs
     let tasks = records
         .into_iter()
@@ -494,7 +494,11 @@ 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"),
+                otel: tracing::info_span!(
+                    "task",
+                    operation = "task",
+                    handle = %to_hex(&handle)
+                ),
                 transaction_id,
             })
         })

coprocessor/fhevm-engine/tfhe-worker/src/tfhe_worker.rs

@@ -164,16 +164,26 @@ async fn tfhe_worker_cycle(
 
             // Lock another dependence chain if available and
             // continue processing without waiting for notification
-            let _dcid_span = tracing::info_span!(parent: &loop_span, "query_dependence_chain", operation = "query_dependence_chain");
+            let _dcid_span = tracing::info_span!(
+                parent: &loop_span,
+                "query_dependence_chain",
+                operation = "query_dependence_chain",
+                dependence_chain_id = tracing::field::Empty
+            );
 
             let (dependence_chain_id, _) = dcid_mngr.acquire_next_lock().await?;
             immediately_poll_more_work = dependence_chain_id.is_some();
 
-            tracing::info!(
-                parent: &_dcid_span,
-                dependence_chain_id = ?dependence_chain_id.as_ref().map(hex::encode),
-                "acquired dependence chain lock"
+            _dcid_span.record(
+                "dependence_chain_id",
+                tracing::field::display(
+                    dependence_chain_id
+                        .as_ref()
+                        .map(hex::encode)
+                        .unwrap_or_else(|| "none".to_string()),
+                ),
             );
+            tracing::info!(parent: &_dcid_span, "acquired dependence chain lock");
             drop(_dcid_span);
 
             continue;
@@ -277,7 +287,8 @@ async fn query_for_work<'a>(
 {
     let _s_dcid = tracing::info_span!(
         "query_dependence_chain",
-        operation = "query_dependence_chain"
+        operation = "query_dependence_chain",
+        dependence_chain_id = tracing::field::Empty
     );
     // Lock dependence chain
     let (dependence_chain_id, locking_reason) =
@@ -302,14 +313,23 @@ async fn query_for_work<'a>(
         info!(target: "tfhe_worker", "No dcid found to process");
         return Ok((vec![], PrimitiveDateTime::MAX, false));
     }
-    tracing::info!(
-        parent: &_s_dcid,
-        dependence_chain_id = ?dependence_chain_id.as_ref().map(hex::encode),
-        "query dependence chain result"
+    _s_dcid.record(
+        "dependence_chain_id",
+        tracing::field::display(
+            dependence_chain_id
+                .as_ref()
+                .map(hex::encode)
+                .unwrap_or_else(|| "none".to_string()),
+        ),
     );
+    tracing::info!(parent: &_s_dcid, "query dependence chain result");
     drop(_s_dcid);
 
-    let _s_work = tracing::info_span!("query_work_items", operation = "query_work_items");
+    let _s_work = tracing::info_span!(
+        "query_work_items",
+        operation = "query_work_items",
+        count = tracing::field::Empty
+    );
     let transaction_batch_size = args.work_items_batch_size;
     let started_at = SystemTime::now();
     let the_work = query!(
@@ -351,7 +371,8 @@ WHERE c.transaction_id IN (
     })?;
 
     WORK_ITEMS_QUERY_HISTOGRAM.observe(started_at.elapsed().unwrap_or_default().as_secs_f64());
-    tracing::info!(parent: &_s_work, count = the_work.len(), "work items queried");
+    _s_work.record("count", the_work.len());
+    tracing::info!(parent: &_s_work, "work items queried");
     drop(_s_work);
     health_check.update_db_access();
     if the_work.is_empty() {

coprocessor/fhevm-engine/transaction-sender/src/ops/allow_handle.rs

@@ -361,6 +361,8 @@ where
             let src_transaction_id = row.transaction_id.clone();
             let _span =
                 tracing::info_span!("prepare_allow_account", operation = "prepare_allow_account");
+            // Use `enter()` in async loops to avoid keeping a non-Send
+            // entered guard across await points.
             let _enter = _span.enter();
 
             let handle = row.handle.clone();
@@ -430,9 +432,6 @@ where
                 event_type,
             };
 
-            drop(_enter);
-            drop(_span);
-
             let operation = self.clone();
             join_set.spawn(async move {
                 operation

coprocessor/fhevm-engine/transaction-sender/src/ops/verify_proof.rs

@@ -251,6 +251,8 @@ where
                 "prepare_verify_proof_resp",
                 operation = "prepare_verify_proof_resp"
             );
+            // Use `enter()` in async loops to avoid keeping a non-Send
+            // entered guard across await points.
             let _enter = _span.enter();
 
             let txn_request = match row.verified {
@@ -354,9 +356,6 @@ where
                 }
             };
 
-            drop(_enter);
-            drop(_span);
-
             let self_clone = self.clone();
             let src_transaction_id = transaction_id;
             join_set.spawn(async move {

coprocessor/fhevm-engine/zkproof-worker/src/verifier.rs

@@ -521,7 +521,7 @@ fn create_ciphertext(
     handle[30] = serialized_type as u8;
     handle[31] = current_ciphertext_version() as u8;
 
-    tracing::Span::current().record("ct_type", serialized_type as i64);
+    tracing::Span::current().record("ct_type", tracing::field::display(serialized_type));
     tracing::info!(
         request_id,
         handle = %hex::encode(&handle),