[zero] perf(zebra-consensus): overlap block-path UTXO lookups

zebra/zebra-consensus/Cargo.toml

@@ -113,5 +113,9 @@ harness = false
 name = "sapling"
 harness = false
 
+[[bench]]
+name = "utxo_lookup"
+harness = false
+
 [lints]
 workspace = true

zebra/zebra-consensus/benches/utxo_lookup.rs

@@ -0,0 +1,144 @@
+//! Benchmarks the block-path spent-UTXO fetch under state-service latency.
+//!
+//! Every `AwaitUtxo` response is delayed ~1ms (the tokio timer resolution),
+//! standing in for a state service under load. Run this bench on the base
+//! commit for the serial baseline: one awaited round trip per input, against
+//! this branch's overlapped lookups.
+
+// Disabled due to warnings in criterion macros
+#![allow(missing_docs)]
+
+use std::{collections::HashMap, sync::Arc, time::Duration};
+
+use chrono::{DateTime, Utc};
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use tower::{service_fn, ServiceExt};
+
+use zebra_chain::{
+    amount::Amount,
+    parameters::{Network, NetworkUpgrade},
+    transaction::{LockTime, Transaction},
+    transparent,
+};
+use zebra_consensus::transaction::{BlockRequest, BlockTxVerifier};
+
+const INPUTS: usize = 1001;
+const LOOKUP_LATENCY: Duration = Duration::from_millis(1);
+
+/// A signature-free 1001-input standard P2SH consolidation whose spent UTXOs
+/// all come from the state service.
+fn consolidation() -> (Arc<Transaction>, transparent::Output) {
+    let network = Network::new_default_testnet();
+    let block_height = (NetworkUpgrade::Nu5
+        .activation_height(&network)
+        .expect("NU5 activation height is specified")
+        + 10)
+        .expect("height in range");
+
+    const OP_TRUE: u8 = 0x51;
+    let unlock_script = transparent::Script::new(&[0x01, OP_TRUE]);
+    // OP_HASH160 <RIPEMD160(SHA256([OP_TRUE]))> OP_EQUAL
+    let mut p2sh_lock_bytes = vec![0xa9, 0x14];
+    p2sh_lock_bytes.extend_from_slice(&[
+        0xda, 0x17, 0x45, 0xe9, 0xb5, 0x49, 0xbd, 0x0b, 0xfa, 0x1a, 0x56, 0x99, 0x71, 0xc7, 0x7e,
+        0xba, 0x30, 0xcd, 0x5a, 0x4b,
+    ]);
+    p2sh_lock_bytes.push(0x87);
+    let lock_script = transparent::Script::new(&p2sh_lock_bytes);
+
+    let spent_output = transparent::Output {
+        value: Amount::try_from(10_000).expect("valid amount"),
+        lock_script,
+    };
+
+    let source_hash = zebra_chain::transaction::Hash([7u8; 32]);
+    let inputs: Vec<transparent::Input> = (0..INPUTS)
+        .map(|index| transparent::Input::PrevOut {
+            outpoint: transparent::OutPoint {
+                hash: source_hash,
+                // Bounded by INPUTS, so the cast cannot truncate.
+                index: index as u32,
+            },
+            unlock_script: unlock_script.clone(),
+            sequence: 0,
+        })
+        .collect();
+
+    let output = transparent::Output {
+        value: Amount::try_from(5_000).expect("valid amount"),
+        lock_script: transparent::Script::new(&[0]),
+    };
+
+    let transaction = Arc::new(Transaction::V5 {
+        inputs,
+        outputs: vec![output],
+        lock_time: LockTime::unlocked(),
+        expiry_height: (block_height + 1).expect("height in range"),
+        sapling_shielded_data: None,
+        orchard_shielded_data: None,
+        network_upgrade: NetworkUpgrade::Nu5,
+    });
+
+    (transaction, spent_output)
+}
+
+fn utxo_fetch(c: &mut Criterion) {
+    let rt = tokio::runtime::Runtime::new().expect("runtime");
+    let network = Network::new_default_testnet();
+    let (transaction, spent_output) = consolidation();
+
+    let block_height = (NetworkUpgrade::Nu5
+        .activation_height(&network)
+        .expect("NU5 activation height is specified")
+        + 10)
+        .expect("height in range");
+    let fund_height = (block_height - 1).expect("height in range");
+
+    let state = move || {
+        let spent_output = spent_output.clone();
+        service_fn(move |request: zebra_state::Request| {
+            let spent_output = spent_output.clone();
+            async move {
+                match request {
+                    zebra_state::Request::AwaitUtxo(_) => {
+                        tokio::time::sleep(LOOKUP_LATENCY).await;
+                        Ok::<_, zebra_consensus::BoxError>(zebra_state::Response::Utxo(
+                            transparent::Utxo::new(spent_output, fund_height, false),
+                        ))
+                    }
+                    other => unreachable!("unexpected state request: {other:?}"),
+                }
+            }
+        })
+    };
+
+    let mut group = c.benchmark_group("block_verification_1001_inputs");
+    group.sample_size(10);
+
+    group.bench_function("state_latency_1ms_per_utxo", |b| {
+        b.iter_batched(
+            || {
+                (
+                    BlockTxVerifier::new(&network, state()),
+                    BlockRequest {
+                        transaction_hash: transaction.hash(),
+                        transaction: transaction.clone(),
+                        known_utxos: Arc::new(HashMap::new()),
+                        height: block_height,
+                        time: DateTime::<Utc>::MAX_UTC,
+                    },
+                )
+            },
+            |(verifier, request)| {
+                rt.block_on(verifier.oneshot(request))
+                    .expect("transaction verifies")
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    group.finish();
+}
+
+criterion_group!(benches, utxo_fetch);
+criterion_main!(benches);

zebra/zebra-consensus/src/transaction.rs

@@ -61,6 +61,17 @@ mod tests;
 ///     chain in the correct order.)
 const UTXO_LOOKUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(6 * 60);
 
+/// The maximum number of in-flight state UTXO lookups per transaction: enough
+/// to overlap the per-lookup latency without flooding the state service's buffer.
+///
+/// Batched lookups start their [`UTXO_LOOKUP_TIMEOUT`] clocks together instead
+/// of serially, so an `AwaitUtxo` answered more than 6 minutes into the phase
+/// now fails it. Accepted narrowing: that case only occurs in out-of-order
+/// sync, it already failed serially whenever such an input was among the first
+/// in flight, and the designed recovery is a sync restart (see the timeout's
+/// docs).
+const MAX_CONCURRENT_UTXO_LOOKUPS: usize = 64;
+
 /// A timeout applied to output lookup requests sent to the mempool. This is shorter than the
 /// timeout for the state UTXO lookups because a block is likely to be mined every 75 seconds
 /// after Blossom is active, changing the best chain tip and requiring re-verification of transactions
@@ -347,9 +358,9 @@ where
             // Load spent UTXOs from the block context and state.
             // The UTXOs are required for almost all the async checks.
             //
-            // This phase is one awaited state round trip per transparent input, so its share of
-            // the total tells operators whether transaction verification is bound by state
-            // lookups or by the cryptographic checks timed below.
+            // This phase overlaps up to `MAX_CONCURRENT_UTXO_LOOKUPS` state lookups, so its
+            // share of the total tells operators whether transaction verification is bound by
+            // state lookups or by the cryptographic checks timed below.
             let utxo_fetch_start = Instant::now();
             let spent_utxos_result =
                 Self::block_spent_utxos(tx.clone(), known_utxos, state.clone()).await;
@@ -448,33 +459,57 @@ where
         // Pre-allocate with None so we can fill each slot by input index, preserving input order.
         let mut spent_outputs: Vec<Option<transparent::Output>> = vec![None; inputs.len()];
 
+        // Serve block-internal spends from `known_utxos` inline; only the
+        // remaining inputs go to the state service. Those lookups run
+        // concurrently: with thousands of transparent inputs, one awaited
+        // round trip per input turns this phase into a serial latency chain.
+        // Results carry their input index, so completion order cannot affect
+        // the `spent_outputs` order, which v5 sighashes commit to.
+        //
+        // The futures are collected eagerly so each owns its captures; a lazy
+        // iterator would hold `&state` inside the stream and require `ZS: Sync`.
+        let mut lookups = Vec::new();
         for (input_idx, input) in inputs.iter().enumerate() {
-            if let transparent::Input::PrevOut { outpoint, .. } = input {
+            let transparent::Input::PrevOut { outpoint, .. } = input else {
+                continue;
+            };
+            let outpoint = *outpoint;
+
+            if let Some(ordered) = known_utxos.get(&outpoint) {
+                tracing::trace!("UTXO in known_utxos, discarding query");
+                spent_outputs[input_idx] = Some(ordered.utxo.output.clone());
+                spent_utxos.insert(outpoint, ordered.utxo.clone());
+                continue;
+            }
+
+            let state = state.clone();
+            lookups.push(async move {
                 tracing::trace!("awaiting outpoint lookup");
 
-                let utxo = if let Some(output) = known_utxos.get(outpoint) {
-                    tracing::trace!("UTXO in known_utxos, discarding query");
-                    output.utxo.clone()
-                } else {
-                    let response = state
-                        .clone()
-                        .oneshot(zebra_state::Request::AwaitUtxo(*outpoint))
-                        .await
-                        .map_err(|boxed_error| match boxed_error.downcast::<Elapsed>() {
-                            Ok(_) => TransactionError::TransparentInputNotFound,
-                            Err(boxed_error) => TransactionError::from(boxed_error),
-                        })?;
-
-                    if let zebra_state::Response::Utxo(utxo) = response {
-                        utxo
-                    } else {
-                        unreachable!("AwaitUtxo always responds with Utxo")
-                    }
+                let response = state
+                    .oneshot(zebra_state::Request::AwaitUtxo(outpoint))
+                    .await
+                    .map_err(|boxed_error| match boxed_error.downcast::<Elapsed>() {
+                        Ok(_) => TransactionError::TransparentInputNotFound,
+                        Err(boxed_error) => TransactionError::from(boxed_error),
+                    })?;
+
+                let zebra_state::Response::Utxo(utxo) = response else {
+                    unreachable!("AwaitUtxo always responds with Utxo")
                 };
-                tracing::trace!(?utxo, "got UTXO");
-                spent_outputs[input_idx] = Some(utxo.output.clone());
-                spent_utxos.insert(*outpoint, utxo);
-            }
+
+                Ok::<_, TransactionError>((input_idx, outpoint, utxo))
+            });
+        }
+
+        let mut lookups =
+            futures::stream::iter(lookups).buffer_unordered(MAX_CONCURRENT_UTXO_LOOKUPS);
+
+        while let Some(lookup) = lookups.next().await {
+            let (input_idx, outpoint, utxo) = lookup?;
+            tracing::trace!(?utxo, "got UTXO");
+            spent_outputs[input_idx] = Some(utxo.output.clone());
+            spent_utxos.insert(outpoint, utxo);
         }
 
         let spent_outputs: Vec<transparent::Output> = spent_outputs.into_iter().flatten().collect();