From 3cd0079fcc93d185d45638fba5b414fa2d818d8b Mon Sep 17 00:00:00 2001 From: Dev Ojha Date: Tue, 30 Jun 2026 18:25:51 -0600 Subject: [PATCH 1/2] perf(consensus): flush crypto batches at block boundary --- CHANGELOG.md | 3 + tower-batch-control/CHANGELOG.md | 7 ++ tower-batch-control/src/message.rs | 18 ++- tower-batch-control/src/service.rs | 15 ++- tower-batch-control/src/worker.rs | 32 ++++- tower-batch-control/tests/ed25519.rs | 43 +++++++ tower-fallback/CHANGELOG.md | 6 + tower-fallback/src/service.rs | 8 ++ zebra-consensus/src/block.rs | 6 +- zebra-consensus/src/primitives.rs | 176 +++++++++++++++++++++++++++ zebra-consensus/src/transaction.rs | 58 +++++++-- 11 files changed, 351 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ed4c0a3f453..4557e949671 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org). ### Performance +- Explicitly flush consensus crypto batches at the semantic block verification + boundary, reducing near-tip block verification latency without changing the + global batch latency setting. - Compute the v5 ZIP-244 txid and authorizing-data digest natively. Both previously routed through `Transaction::to_librustzcash`, which re-serializes and reparses the whole transaction — decompressing every Jubjub and Pallas diff --git a/tower-batch-control/CHANGELOG.md b/tower-batch-control/CHANGELOG.md index b0a66c5f214..1bc8e78753a 100644 --- a/tower-batch-control/CHANGELOG.md +++ b/tower-batch-control/CHANGELOG.md @@ -5,6 +5,13 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added + +- Added `Batch::flush()` for explicitly queueing a pending batch flush before + the configured maximum latency expires. + ## [0.2.42] - 2025-10-15 ### Added diff --git a/tower-batch-control/src/message.rs b/tower-batch-control/src/message.rs index 05415e2b4ca..21f746bb9a4 100644 --- a/tower-batch-control/src/message.rs +++ b/tower-batch-control/src/message.rs @@ -4,13 +4,19 @@ use tokio::sync::{oneshot, OwnedSemaphorePermit}; use super::error::ServiceError; -/// Message sent to the batch worker +/// Message sent to the batch worker. #[derive(Debug)] -pub(crate) struct Message { - pub(crate) request: Request, - pub(crate) tx: Tx, - pub(crate) span: tracing::Span, - pub(super) _permit: OwnedSemaphorePermit, +pub(crate) enum Message { + /// A batch item request. + Item { + request: Request, + tx: Tx, + span: tracing::Span, + _permit: OwnedSemaphorePermit, + }, + + /// An explicit request to flush the pending batch. + Flush { span: tracing::Span }, } /// Response sender diff --git a/tower-batch-control/src/service.rs b/tower-batch-control/src/service.rs index 4fdf7f53361..c2ccec09893 100644 --- a/tower-batch-control/src/service.rs +++ b/tower-batch-control/src/service.rs @@ -223,6 +223,19 @@ where fn get_worker_error(&self) -> crate::BoxError { self.error_handle.get_error_on_closed() } + + /// Explicitly flushes the current pending batch. + /// + /// The flush is queued after item requests that have already been sent to + /// this batch worker. This method returns when the command has been queued, + /// not when the underlying batch has completed. + pub fn flush(&self) -> Result<(), crate::BoxError> { + self.tx + .send(Message::Flush { + span: tracing::Span::current(), + }) + .map_err(|_| self.get_worker_error()) + } } impl Service for Batch @@ -330,7 +343,7 @@ where // acquired, so we can freely allocate a oneshot. let (tx, rx) = oneshot::channel(); - match self.tx.send(Message { + match self.tx.send(Message::Item { request, tx, span, diff --git a/tower-batch-control/src/worker.rs b/tower-batch-control/src/worker.rs index 50bfacf92b7..e9d27599eca 100644 --- a/tower-batch-control/src/worker.rs +++ b/tower-batch-control/src/worker.rs @@ -166,6 +166,12 @@ where /// Waits until the inner service is ready, /// then stores a future which resolves when the batch finishes. async fn flush_service(&mut self) { + if self.pending_items_weight == 0 { + tracing::trace!("skipping flush for empty batch"); + self.pending_batch_timer = None; + return; + } + if self.failed.is_some() { tracing::trace!("worker failure: skipping flush"); return; @@ -233,7 +239,12 @@ where }, maybe_msg = self.rx.recv(), if self.can_spawn_new_batches() => match maybe_msg { - Some(msg) => { + Some(Message::Item { + request, + tx, + span, + .. + }) => { tracing::trace!( pending_items_weight = self.pending_items_weight, batch_deadline = ?self.pending_batch_timer.as_ref().map(|sleep| sleep.deadline()), @@ -241,10 +252,9 @@ where "batch message received", ); - let span = msg.span; let is_new_batch = self.pending_items_weight == 0; - self.process_req(msg.request, msg.tx) + self.process_req(request, tx) // Apply the provided span to request processing. .instrument(span) .await; @@ -279,6 +289,16 @@ where ); } } + Some(Message::Flush { span }) => { + tracing::trace!( + pending_items_weight = self.pending_items_weight, + batch_deadline = ?self.pending_batch_timer.as_ref().map(|sleep| sleep.deadline()), + running_batches = self.concurrent_batches.len(), + "explicit batch flush received", + ); + + self.flush_service().instrument(span).await; + } None => { tracing::trace!("batch channel closed and emptied, exiting worker task"); @@ -375,9 +395,9 @@ where // Fail queued requests while let Ok(msg) = self.rx.try_recv() { - let _ = msg - .tx - .send(Err(self.failed.as_ref().expect("just set failed").clone())); + if let Message::Item { tx, .. } = msg { + let _ = tx.send(Err(self.failed.as_ref().expect("just set failed").clone())); + } } // Clear any finished batches, ignoring any errors. diff --git a/tower-batch-control/tests/ed25519.rs b/tower-batch-control/tests/ed25519.rs index 764da9cd930..5bea37cb41f 100644 --- a/tower-batch-control/tests/ed25519.rs +++ b/tower-batch-control/tests/ed25519.rs @@ -223,6 +223,30 @@ where Ok(()) } +async fn sign_and_verify_after_explicit_flush( + mut verifier: Batch, + n: usize, +) -> Result<(), BoxError> { + let mut results = FuturesOrdered::new(); + for _ in 0..n { + let sk = SigningKey::new(thread_rng()); + let vk_bytes = VerificationKeyBytes::from(&sk); + let msg = b"BatchVerifyTest"; + let sig = sk.sign(&msg[..]); + + verifier.ready().await?; + results.push_back(verifier.call((vk_bytes, sig, msg).into())); + } + + verifier.flush()?; + + while let Some(result) = results.next().await { + result?; + } + + Ok(()) +} + #[tokio::test(flavor = "multi_thread")] async fn batch_flushes_on_max_items_weight() -> Result<(), Report> { use tokio::time::timeout; @@ -259,6 +283,25 @@ async fn batch_flushes_on_max_latency() -> Result<(), Report> { Ok(()) } +#[tokio::test(flavor = "multi_thread")] +async fn batch_flushes_on_explicit_flush() -> Result<(), Report> { + use tokio::time::timeout; + let _init_guard = zebra_test::init(); + + // Use a very high max_items and long max_latency. Without the explicit + // flush, this verification would wait for the latency timer. + let verifier = Batch::new(Verifier::default(), 100, 10, Duration::from_secs(1000)); + timeout( + Duration::from_secs(1), + sign_and_verify_after_explicit_flush(verifier, 10), + ) + .await + .map_err(|e| eyre!(e))? + .map_err(|e| eyre!(e))?; + + Ok(()) +} + #[tokio::test(flavor = "multi_thread")] async fn fallback_verification() -> Result<(), Report> { let _init_guard = zebra_test::init(); diff --git a/tower-fallback/CHANGELOG.md b/tower-fallback/CHANGELOG.md index b1abe936f2b..a0195c25939 100644 --- a/tower-fallback/CHANGELOG.md +++ b/tower-fallback/CHANGELOG.md @@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added + +- Added `Fallback::primary()` for read-only access to the primary service. + ## [0.2.41] - 2025-07-11 First "stable" release. However, be advised that the API may still greatly diff --git a/tower-fallback/src/service.rs b/tower-fallback/src/service.rs index 99c5c53ab07..df8842a8e56 100644 --- a/tower-fallback/src/service.rs +++ b/tower-fallback/src/service.rs @@ -31,6 +31,14 @@ impl Fallback { pub fn new(svc1: S1, svc2: S2) -> Self { Self { svc1, svc2 } } + + /// Returns the primary service. + /// + /// Requests are first attempted using this service, before retrying on the + /// fallback service. + pub fn primary(&self) -> &S1 { + &self.svc1 + } } impl Service for Fallback diff --git a/zebra-consensus/src/block.rs b/zebra-consensus/src/block.rs index 1a1217d44c9..6575c128792 100644 --- a/zebra-consensus/src/block.rs +++ b/zebra-consensus/src/block.rs @@ -31,7 +31,7 @@ use zebra_chain::{ }; use zebra_state as zs; -use crate::{error::*, transaction as tx, BoxError}; +use crate::{error::*, primitives, transaction as tx, BoxError}; pub mod check; pub mod request; @@ -276,6 +276,10 @@ where let known_outpoint_hashes: Arc> = Arc::new(known_utxos.keys().map(|outpoint| outpoint.hash).collect()); + let _block_batch_flush = primitives::register_block_verifier_batch_flush( + &known_outpoint_hashes, + block.transactions.len(), + ); for (&transaction_hash, transaction) in transaction_hashes.iter().zip(block.transactions.iter()) diff --git a/zebra-consensus/src/primitives.rs b/zebra-consensus/src/primitives.rs index 822349e18d1..3b19e40395b 100644 --- a/zebra-consensus/src/primitives.rs +++ b/zebra-consensus/src/primitives.rs @@ -1,5 +1,11 @@ //! Asynchronous verification of cryptographic primitives. +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; + +use once_cell::sync::Lazy; use tokio::sync::oneshot::error::RecvError; use crate::BoxError; @@ -17,6 +23,147 @@ const MAX_BATCH_SIZE: usize = 64; /// The maximum latency bound for any of the batch verifiers. const MAX_BATCH_LATENCY: std::time::Duration = std::time::Duration::from_millis(100); +/// Registered block transaction-verification batches waiting for an explicit +/// cryptographic flush. +static BLOCK_VERIFIER_BATCH_FLUSHES: Lazy>> = + Lazy::new(|| Mutex::new(HashMap::new())); + +/// Opaque identity for a semantic block verifier's transaction batch. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) struct BlockVerifierBatchFlushKey(usize); + +impl BlockVerifierBatchFlushKey { + /// Returns a key for a value shared by all transaction requests in one + /// block verification. + fn new(shared_block_context: &Arc) -> Self { + // This cast stores the Arc allocation address as an opaque identity + // key. The address is never dereferenced or used for pointer arithmetic. + Self(Arc::as_ptr(shared_block_context) as usize) + } +} + +/// Registration guard for one semantic block verifier's crypto batch flush. +#[derive(Debug)] +pub(crate) struct BlockVerifierBatchFlushGuard { + key: BlockVerifierBatchFlushKey, +} + +impl Drop for BlockVerifierBatchFlushGuard { + fn drop(&mut self) { + BLOCK_VERIFIER_BATCH_FLUSHES + .lock() + .expect("block verifier batch flush registry mutex should not be poisoned") + .remove(&self.key); + } +} + +/// Tracks transaction verifier progress toward one block-level flush. +#[derive(Debug)] +struct BlockFlush { + expected_transactions: usize, + started_transactions: usize, + flush_queued: bool, +} + +/// Registers a block's transaction verifier batch for one explicit crypto +/// flush once every transaction has started its asynchronous checks. +pub(crate) fn register_block_verifier_batch_flush( + shared_block_context: &Arc, + expected_transactions: usize, +) -> BlockVerifierBatchFlushGuard { + let key = BlockVerifierBatchFlushKey::new(shared_block_context); + + BLOCK_VERIFIER_BATCH_FLUSHES + .lock() + .expect("block verifier batch flush registry mutex should not be poisoned") + .insert( + key, + BlockFlush { + expected_transactions, + started_transactions: 0, + flush_queued: expected_transactions == 0, + }, + ); + + BlockVerifierBatchFlushGuard { key } +} + +/// Returns the block-level batch flush key for `shared_block_context`. +pub(crate) fn block_verifier_batch_flush_key( + shared_block_context: &Arc, +) -> BlockVerifierBatchFlushKey { + BlockVerifierBatchFlushKey::new(shared_block_context) +} + +/// Records that one block transaction has started its async checks, and +/// flushes the shared crypto batches when all transactions in the block have +/// reached that boundary. +pub(crate) fn start_block_transaction_async_checks(key: BlockVerifierBatchFlushKey) { + if block_verifier_batch_flush_ready(key) { + flush_block_verifier_batches(); + } +} + +/// Returns `true` if recording `key` made its block ready to flush. +fn block_verifier_batch_flush_ready(key: BlockVerifierBatchFlushKey) -> bool { + let mut flushes = BLOCK_VERIFIER_BATCH_FLUSHES + .lock() + .expect("block verifier batch flush registry mutex should not be poisoned"); + + let Some(flush) = flushes.get_mut(&key) else { + return false; + }; + + flush.started_transactions = flush.started_transactions.saturating_add(1); + + if flush.flush_queued || flush.started_transactions < flush.expected_transactions { + return false; + } + + flush.flush_queued = true; + true +} + +/// Explicitly flushes the shared crypto batch services used by semantic block +/// transaction verification. +/// +/// The block verifier still awaits every transaction verifier future before +/// committing a block; this just starts pending batch work before +/// [`MAX_BATCH_LATENCY`] expires. +fn flush_block_verifier_batches() { + if let Some(verifier) = Lazy::get(&ed25519::VERIFIER) { + queue_batch_flush("ed25519", verifier.primary().flush()); + } + + if let Some(verifier) = Lazy::get(&sapling::VERIFIER) { + queue_batch_flush("sapling", verifier.primary().flush()); + } + + if let Some(verifier) = Lazy::get(&halo2::VERIFIER_PRE_NU6_2) { + queue_batch_flush("halo2_pre_nu6_2", verifier.primary().flush()); + } + + if let Some(verifier) = Lazy::get(&halo2::VERIFIER_NU6_2) { + queue_batch_flush("halo2_nu6_2", verifier.primary().flush()); + } + + if let Some(verifier) = Lazy::get(&halo2::VERIFIER_NU6_3_ONWARD) { + queue_batch_flush("halo2_nu6_3_onward", verifier.primary().flush()); + } +} + +/// Logs best-effort flush queueing failures without changing verification +/// semantics. +fn queue_batch_flush(verifier: &'static str, result: Result<(), BoxError>) { + if let Err(error) = result { + tracing::trace!( + ?error, + verifier, + "could not queue explicit block verifier batch flush" + ); + } +} + /// Fires off a task into the Rayon threadpool, awaits the result through a oneshot channel, /// then converts the error to a [`BoxError`]. pub async fn spawn_fifo_and_convert< @@ -47,3 +194,32 @@ pub async fn spawn_fifo T + Send>( rsp_rx.await } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn block_flush_is_ready_when_expected_transactions_start_checks() { + let shared_block_context = Arc::new(()); + let key = block_verifier_batch_flush_key(&shared_block_context); + let _guard = register_block_verifier_batch_flush(&shared_block_context, 2); + + assert!(!block_verifier_batch_flush_ready(key)); + assert!(block_verifier_batch_flush_ready(key)); + assert!(!block_verifier_batch_flush_ready(key)); + } + + #[test] + fn block_flush_guard_unregisters_batch() { + let shared_block_context = Arc::new(()); + let key = block_verifier_batch_flush_key(&shared_block_context); + + { + let _guard = register_block_verifier_batch_flush(&shared_block_context, 1); + assert!(block_verifier_batch_flush_ready(key)); + } + + assert!(!block_verifier_batch_flush_ready(key)); + } +} diff --git a/zebra-consensus/src/transaction.rs b/zebra-consensus/src/transaction.rs index 10a1c670e99..9fe126a5944 100644 --- a/zebra-consensus/src/transaction.rs +++ b/zebra-consensus/src/transaction.rs @@ -319,6 +319,19 @@ impl Request { pub fn is_mempool(&self) -> bool { matches!(self, Request::Mempool { .. }) } + + /// Returns the block verifier batch flush key, if this is a block request. + fn block_verifier_batch_flush_key(&self) -> Option { + match self { + Request::Block { + known_outpoint_hashes, + .. + } => Some(primitives::block_verifier_batch_flush_key( + known_outpoint_hashes, + )), + Request::Mempool { .. } => None, + } + } } impl Response { @@ -572,9 +585,11 @@ where async_checks.push(check_anchors_and_revealed_nullifiers_query); } + let block_batch_flush_key = req.block_verifier_batch_flush_key(); + tracing::trace!(?tx_id, "awaiting async checks..."); - async_checks.check().await?; + async_checks.check(block_batch_flush_key).await?; tracing::trace!(?tx_id, "finished async checks"); @@ -1262,15 +1277,44 @@ impl AsyncChecks { /// /// If any of the checks fail, this method immediately returns the error and cancels all other /// checks by dropping them. - async fn check(mut self) -> Result<(), BoxError> { + async fn check( + mut self, + block_batch_flush_key: Option, + ) -> Result<(), BoxError> { + let mut needs_block_batch_flush = block_batch_flush_key.is_some(); + let mut block_batch_flush = match block_batch_flush_key { + Some(key) => async move { + tokio::task::yield_now().await; + primitives::start_block_transaction_async_checks(key); + } + .boxed(), + None => futures::future::pending().boxed(), + }; + // Wait for all asynchronous checks to complete // successfully, or fail verification if they error. - while let Some(check) = self.0.next().await { - tracing::trace!(?check, remaining = self.0.len()); - check?; - } + loop { + tokio::select! { + biased; + + check = self.0.next() => { + let Some(check) = check else { + if needs_block_batch_flush { + block_batch_flush.await; + } + + return Ok(()); + }; - Ok(()) + tracing::trace!(?check, remaining = self.0.len()); + check?; + } + + () = &mut block_batch_flush, if needs_block_batch_flush => { + needs_block_batch_flush = false; + } + } + } } } From 4c6d075bb9b5e3cf4d360d70b1b5c8078f67e015 Mon Sep 17 00:00:00 2001 From: Dev Ojha Date: Sun, 5 Jul 2026 02:31:43 -0500 Subject: [PATCH 2/2] test(consensus): document flush guard drop order and add flush edge-case tests - Document that the block batch flush guard must be declared after known_outpoint_hashes, so the registry entry is removed while the Arc address is still pinned and cannot be reused by another block's key. - Test that an explicit flush on an empty batch is a no-op and later batches still verify. - Test that re-registering a reused flush key starts a fresh transaction count instead of inheriting a stale one. --- tower-batch-control/tests/ed25519.rs | 22 ++++++++++++++++++++++ zebra-consensus/src/block.rs | 4 ++++ zebra-consensus/src/primitives.rs | 18 ++++++++++++++++++ 3 files changed, 44 insertions(+) diff --git a/tower-batch-control/tests/ed25519.rs b/tower-batch-control/tests/ed25519.rs index 5bea37cb41f..ceaacba80c1 100644 --- a/tower-batch-control/tests/ed25519.rs +++ b/tower-batch-control/tests/ed25519.rs @@ -302,6 +302,28 @@ async fn batch_flushes_on_explicit_flush() -> Result<(), Report> { Ok(()) } +#[tokio::test(flavor = "multi_thread")] +async fn batch_flush_on_empty_batch_is_noop() -> Result<(), Report> { + use tokio::time::timeout; + let _init_guard = zebra_test::init(); + + // Flushing with no pending items must not fail the worker or disturb later + // batches: items queued after the empty flush are still verified by the + // next explicit flush. + let verifier = Batch::new(Verifier::default(), 100, 10, Duration::from_secs(1000)); + verifier.flush().map_err(|e| eyre!(e))?; + + timeout( + Duration::from_secs(1), + sign_and_verify_after_explicit_flush(verifier, 10), + ) + .await + .map_err(|e| eyre!(e))? + .map_err(|e| eyre!(e))?; + + Ok(()) +} + #[tokio::test(flavor = "multi_thread")] async fn fallback_verification() -> Result<(), Report> { let _init_guard = zebra_test::init(); diff --git a/zebra-consensus/src/block.rs b/zebra-consensus/src/block.rs index 6575c128792..e413e01a309 100644 --- a/zebra-consensus/src/block.rs +++ b/zebra-consensus/src/block.rs @@ -276,6 +276,10 @@ where let known_outpoint_hashes: Arc> = Arc::new(known_utxos.keys().map(|outpoint| outpoint.hash).collect()); + // This guard must be declared after `known_outpoint_hashes`: the flush + // registry is keyed by the `Arc` allocation address, so the guard must be + // dropped first, removing the registry entry while the address is still + // pinned by this `Arc` and cannot be reused by another block's key. let _block_batch_flush = primitives::register_block_verifier_batch_flush( &known_outpoint_hashes, block.transactions.len(), diff --git a/zebra-consensus/src/primitives.rs b/zebra-consensus/src/primitives.rs index 3b19e40395b..e478d63ccee 100644 --- a/zebra-consensus/src/primitives.rs +++ b/zebra-consensus/src/primitives.rs @@ -222,4 +222,22 @@ mod tests { assert!(!block_verifier_batch_flush_ready(key)); } + + #[test] + fn block_flush_registration_starts_fresh_after_guard_drop() { + let shared_block_context = Arc::new(()); + let key = block_verifier_batch_flush_key(&shared_block_context); + + { + let _guard = register_block_verifier_batch_flush(&shared_block_context, 2); + assert!(!block_verifier_batch_flush_ready(key)); + } + + // Re-registering the same key starts a fresh transaction count, so a + // partially counted earlier registration for a reused `Arc` address + // can't make a later block's flush fire early. + let _guard = register_block_verifier_batch_flush(&shared_block_context, 2); + assert!(!block_verifier_batch_flush_ready(key)); + assert!(block_verifier_batch_flush_ready(key)); + } }