test(tfhe-worker): filter pattern spans by synthetic tx id

Author: Eikix

coprocessor/fhevm-engine/scheduler/Cargo.toml

@@ -27,3 +27,4 @@ fhevm-engine-common = { path = "../fhevm-engine-common" }
 [features]
 nightly-avx512 = ["tfhe/nightly-avx512"]
 gpu = ["tfhe/gpu"]
+test-span-attrs = []

coprocessor/fhevm-engine/scheduler/src/dfg/scheduler.rs

@@ -27,6 +27,27 @@ use super::{pattern, DFComponentGraph, DFGraph, OpNode};
 const OPERATION_RERANDOMISATION_DOMAIN_SEPARATOR: [u8; 8] = *b"TFHE_Rrd";
 const COMPACT_PUBLIC_ENCRYPTION_DOMAIN_SEPARATOR: [u8; 8] = *b"TFHE_Enc";
 
+#[cfg(feature = "test-span-attrs")]
+macro_rules! execute_transaction_span {
+    ($tx_pattern_id:expr, $transaction_id:expr) => {
+        tracing::info_span!(
+            "execute_transaction",
+            transaction_pattern_id = %$tx_pattern_id,
+            test_transaction_id = %hex::encode($transaction_id),
+        )
+    };
+}
+
+#[cfg(not(feature = "test-span-attrs"))]
+macro_rules! execute_transaction_span {
+    ($tx_pattern_id:expr, $transaction_id:expr) => {
+        tracing::info_span!(
+            "execute_transaction",
+            transaction_pattern_id = %$tx_pattern_id,
+        )
+    };
+}
+
 pub enum PartitionStrategy {
     MaxParallelism,
     MaxLocality,
@@ -334,11 +355,8 @@ fn execute_partition(
         }
 
         // Prime the scheduler with ready ops from the transaction's subgraph
-        let _exec_guard = tracing::info_span!(
-            "execute_transaction",
-            transaction_pattern_id = %pattern::pattern_to_base64url(transaction_pattern_id),
-        )
-        .entered();
+        let transaction_pattern_id = pattern::pattern_to_base64url(transaction_pattern_id);
+        let _exec_guard = execute_transaction_span!(transaction_pattern_id, &tid).entered();
         let started_at = std::time::Instant::now();
 
         let Ok(ts) = daggy::petgraph::algo::toposort(&dfg.graph, None) else {
@@ -483,6 +501,7 @@ fn try_execute_node(
 }
 
 type OpResult = Result<CompressedCiphertext>;
+
 fn run_computation(
     operation: i32,
     inputs: Vec<SupportedFheCiphertexts>,

coprocessor/fhevm-engine/tfhe-worker/Cargo.toml

@@ -56,6 +56,7 @@ test-harness = { path = "../test-harness" }
 serde = { workspace = true }
 serial_test = { workspace = true }
 opentelemetry_sdk = { workspace = true, features = ["testing"] }
+scheduler = { path = "../scheduler", features = ["test-span-attrs"] }
 
 [[bin]]
 name = "tfhe_worker"

coprocessor/fhevm-engine/tfhe-worker/src/tests/pattern_integration.rs

@@ -36,12 +36,19 @@ fn span_id_key(id: SpanId) -> u64 {
 /// Multiple `execute_transaction` spans may share the same
 /// `transaction_pattern_id` when a transaction is split into several
 /// scheduler components.
-fn group_ops_by_transaction(spans: &[SpanData]) -> HashMap<String, Vec<(u64, Vec<String>)>> {
+fn group_ops_by_transaction(
+    spans: &[SpanData],
+    test_transaction_ids: &HashSet<String>,
+) -> HashMap<String, Vec<(u64, Vec<String>)>> {
     // execute_transaction span_id → transaction_pattern_id
     let exec_spans: HashMap<u64, String> = spans
         .iter()
         .filter(|s| s.name == "execute_transaction")
         .filter_map(|s| {
+            let tx_id = get_string_attr(s, "test_transaction_id")?;
+            if !test_transaction_ids.contains(&tx_id) {
+                return None;
+            }
             let tp = get_string_attr(s, "transaction_pattern_id")?;
             Some((span_id_key(s.span_context.span_id()), tp))
         })
@@ -439,8 +446,20 @@ async fn test_erc20_transaction_pattern_ids() -> Result<(), Box<dyn std::error::
 
     // ── Span assertions ─────────────────────────────────────────────────
 
+    let test_transaction_ids = HashSet::from([
+        hex::encode(tx_single_id),
+        hex::encode(tx_triple_id),
+        hex::encode(tx_nocmux_id),
+    ]);
     let spans = finished_test_spans().expect("failed to read spans");
-    let groups = group_ops_by_transaction(&spans);
+    assert!(
+        spans.iter().any(|s| {
+            s.name == "execute_transaction" && get_string_attr(s, "test_transaction_id").is_some()
+        }),
+        "execute_transaction spans are missing test_transaction_id; \
+         scheduler test-span-attrs feature is likely disabled"
+    );
+    let groups = group_ops_by_transaction(&spans, &test_transaction_ids);
 
     // We expect exactly 3 distinct transaction_pattern_ids.
     //