From 88a637575a6ffe0511d35ce1238908a3a6d6c490 Mon Sep 17 00:00:00 2001 From: Sergey Timoshin Date: Mon, 3 Aug 2026 12:09:56 +0000 Subject: [PATCH 01/10] fix(forester): retain proofs after eligibility ends --- forester/src/epoch_manager.rs | 26 +++++ forester/src/processor/v2/tx_sender.rs | 134 +++++++++++++++++++++++-- 2 files changed, 153 insertions(+), 7 deletions(-) diff --git a/forester/src/epoch_manager.rs b/forester/src/epoch_manager.rs index 241797732e..7e48ffebb7 100644 --- a/forester/src/epoch_manager.rs +++ b/forester/src/epoch_manager.rs @@ -2334,6 +2334,22 @@ impl EpochManager { return Ok(()); } + if let Some(cache) = self + .proof_caches + .get(&tree_pubkey) + .map(|cache| cache.clone()) + { + if cache.is_warming().await { + debug!( + event = "v2_proof_work_deferred_cache_warming", + run_id = %self.run_id, + tree = %tree_pubkey, + "Deferring V2 proof work while late proofs are collected" + ); + return Ok(()); + } + } + // Try to send any cached proofs first let cached_send_start = Instant::now(); if let Some(items_sent) = self @@ -3701,6 +3717,16 @@ impl EpochManager { .or_insert_with(|| Arc::new(SharedProofCache::new(tree_pubkey))) .clone(); + if cache.is_warming().await { + info!( + event = "prewarm_skipped_cache_warming", + run_id = %self_clone.run_id, + tree = %tree_pubkey, + "Tree cache is already collecting proofs; skipping pre-warm" + ); + return; + } + let cache_len = cache.len().await; if cache_len > 0 && !cache.is_warming().await { let mut rpc = match self_clone.rpc_pool.get_connection().await { diff --git a/forester/src/processor/v2/tx_sender.rs b/forester/src/processor/v2/tx_sender.rs index 5b61559b17..9550e019ef 100644 --- a/forester/src/processor/v2/tx_sender.rs +++ b/forester/src/processor/v2/tx_sender.rs @@ -367,7 +367,7 @@ impl TxSender { let current_slot = self.context.slot_tracker.estimated_current_slot(); if !self.is_still_eligible_at(current_slot) { - let proofs_saved = self.save_proofs_to_cache(&mut proof_rx, None).await; + let proofs_saved = self.handoff_proofs_to_cache(proof_rx, None).await; info!( "Active phase ended for epoch {}, stopping tx sender before recv (saved {} proofs to cache)", self.context.epoch, proofs_saved @@ -392,7 +392,7 @@ impl TxSender { let current_slot = self.context.slot_tracker.estimated_current_slot(); if !self.is_still_eligible_at(current_slot) { - let proofs_saved = self.save_proofs_to_cache(&mut proof_rx, Some(result)).await; + let proofs_saved = self.handoff_proofs_to_cache(proof_rx, Some(result)).await; info!( "Active phase ended for epoch {}, stopping tx sender (saved {} proofs to cache)", self.context.epoch, proofs_saved @@ -493,9 +493,15 @@ impl TxSender { }) } - async fn save_proofs_to_cache( + /// Moves proof result ownership to the per-tree cache when eligibility ends. + /// + /// Results that are already available are cached synchronously. The receiver is + /// then kept alive by a background task so proofs that finish after this sender + /// returns are not lost. The cache remains in the warming state until every + /// submitted proof job has dropped its sender. + async fn handoff_proofs_to_cache( &mut self, - proof_rx: &mut mpsc::Receiver, + mut proof_rx: mpsc::Receiver, current_result: Option, ) -> usize { let cache = match &self.proof_cache { @@ -545,16 +551,130 @@ impl TxSender { } } - cache.finish_warming().await; - if saved > 0 { info!( - "Saved {} proofs to cache for potential reuse (root: {:?})", + "Saved {} available proofs to cache and waiting for late results (root: {:?})", saved, &self.last_seen_root[..4] ); + } else { + info!( + "Waiting for late proof results to warm cache (root: {:?})", + &self.last_seen_root[..4] + ); + } + + let _ = spawn_late_proof_collector(cache.clone(), proof_rx, self.context.merkle_tree); + + saved + } +} + +fn spawn_late_proof_collector( + cache: Arc, + mut proof_rx: mpsc::Receiver, + tree: solana_sdk::pubkey::Pubkey, +) -> JoinHandle { + tokio::spawn(async move { + let mut saved = 0usize; + while let Some(result) = proof_rx.recv().await { + match result.result { + Ok(instruction) => { + cache + .add_proof(result.seq, result.old_root, result.new_root, instruction) + .await; + saved += 1; + } + Err(error) => { + warn!( + tree = %tree, + seq = result.seq, + error = %error, + "Late proof failed while warming cache" + ); + } + } } + cache.finish_warming().await; + let total_cached_proofs = cache.len().await; + info!( + tree = %tree, + late_proofs_cached = saved, + total_cached_proofs, + "Late proof collection completed" + ); saved + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn proof_result(seq: u64, old_root: [u8; 32], new_root: [u8; 32]) -> ProofJobResult { + ProofJobResult { + seq, + result: Ok(BatchInstruction::Append(Vec::new())), + old_root, + new_root, + proof_duration_ms: 1, + round_trip_ms: 1, + submitted_at: std::time::Instant::now(), + } + } + + #[tokio::test] + async fn late_proof_result_finishes_warming_cache() { + let tree = solana_sdk::pubkey::Pubkey::new_unique(); + let base_root = [1u8; 32]; + let next_root = [2u8; 32]; + let cache = Arc::new(SharedProofCache::new(tree)); + cache.start_warming(base_root).await; + + let (proof_tx, proof_rx) = mpsc::channel(1); + let collector = spawn_late_proof_collector(cache.clone(), proof_rx, tree); + + assert!(cache.is_warming().await); + proof_tx + .send(proof_result(0, base_root, next_root)) + .await + .unwrap(); + drop(proof_tx); + + assert_eq!(collector.await.unwrap(), 1); + assert!(!cache.is_warming().await); + let cached = cache.take_if_valid(&base_root).await.unwrap(); + assert_eq!(cached.len(), 1); + assert_eq!(cached[0].old_root, base_root); + assert_eq!(cached[0].new_root, next_root); + } + + #[tokio::test] + async fn failed_late_proof_does_not_block_cache_completion() { + let tree = solana_sdk::pubkey::Pubkey::new_unique(); + let base_root = [3u8; 32]; + let cache = Arc::new(SharedProofCache::new(tree)); + cache.start_warming(base_root).await; + + let (proof_tx, proof_rx) = mpsc::channel(1); + let collector = spawn_late_proof_collector(cache.clone(), proof_rx, tree); + proof_tx + .send(ProofJobResult { + seq: 0, + result: Err("proof failed".to_string()), + old_root: base_root, + new_root: [4u8; 32], + proof_duration_ms: 1, + round_trip_ms: 1, + submitted_at: std::time::Instant::now(), + }) + .await + .unwrap(); + drop(proof_tx); + + assert_eq!(collector.await.unwrap(), 0); + assert!(!cache.is_warming().await); + assert!(cache.is_empty().await); } } From 55b752c86b3b51d21c551f0209f0bf71177acb79 Mon Sep 17 00:00:00 2001 From: Sergey Timoshin Date: Tue, 8 Sep 2026 12:42:58 +0000 Subject: [PATCH 02/10] fix: explicitly detach the late-proof collector --- forester/src/processor/v2/tx_sender.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/forester/src/processor/v2/tx_sender.rs b/forester/src/processor/v2/tx_sender.rs index 9550e019ef..7c15c29397 100644 --- a/forester/src/processor/v2/tx_sender.rs +++ b/forester/src/processor/v2/tx_sender.rs @@ -564,7 +564,12 @@ impl TxSender { ); } - let _ = spawn_late_proof_collector(cache.clone(), proof_rx, self.context.merkle_tree); + // Dropping the JoinHandle detaches the collector so it can finish warming the cache. + drop(spawn_late_proof_collector( + cache.clone(), + proof_rx, + self.context.merkle_tree, + )); saved } From 900044f5a266215766b78cdbe2ce3325ed345acf Mon Sep 17 00:00:00 2001 From: Sergey Timoshin Date: Wed, 23 Sep 2026 11:52:25 +0000 Subject: [PATCH 03/10] fix(forester): bound late proof cache warming --- forester/src/processor/v2/tx_sender.rs | 82 ++++++++++++++++++++------ 1 file changed, 65 insertions(+), 17 deletions(-) diff --git a/forester/src/processor/v2/tx_sender.rs b/forester/src/processor/v2/tx_sender.rs index 7c15c29397..c9ebdd6723 100644 --- a/forester/src/processor/v2/tx_sender.rs +++ b/forester/src/processor/v2/tx_sender.rs @@ -6,6 +6,9 @@ const MAX_BUFFER_SIZE: usize = 1000; const V2_IXS_PER_TX_WITH_LUT: usize = 5; const V2_IXS_PER_TX_WITHOUT_LUT: usize = 4; const FLUSH_MARGIN_SLOTS: u64 = 2; +/// Late proofs are an optimization. Do not let a proof job that never releases +/// its sender keep the tree cache in the warming state indefinitely. +const LATE_PROOF_COLLECTION_TIMEOUT: Duration = Duration::from_secs(30); use light_batched_merkle_tree::merkle_tree::{ InstructionDataBatchAppendInputs, InstructionDataBatchNullifyInputs, @@ -569,6 +572,7 @@ impl TxSender { cache.clone(), proof_rx, self.context.merkle_tree, + LATE_PROOF_COLLECTION_TIMEOUT, )); saved @@ -579,28 +583,45 @@ fn spawn_late_proof_collector( cache: Arc, mut proof_rx: mpsc::Receiver, tree: solana_sdk::pubkey::Pubkey, + collection_timeout: Duration, ) -> JoinHandle { tokio::spawn(async move { let mut saved = 0usize; - while let Some(result) = proof_rx.recv().await { - match result.result { - Ok(instruction) => { - cache - .add_proof(result.seq, result.old_root, result.new_root, instruction) - .await; - saved += 1; - } - Err(error) => { - warn!( - tree = %tree, - seq = result.seq, - error = %error, - "Late proof failed while warming cache" - ); + let collection = async { + while let Some(result) = proof_rx.recv().await { + match result.result { + Ok(instruction) => { + cache + .add_proof(result.seq, result.old_root, result.new_root, instruction) + .await; + saved += 1; + } + Err(error) => { + warn!( + tree = %tree, + seq = result.seq, + error = %error, + "Late proof failed while warming cache" + ); + } } } + }; + + if tokio::time::timeout(collection_timeout, collection) + .await + .is_err() + { + warn!( + tree = %tree, + timeout_ms = collection_timeout.as_millis(), + late_proofs_cached = saved, + "Late proof collection timed out; releasing cache warming state" + ); } + // This must run on channel closure and timeout so a retained sender can + // never make the tree permanently ineligible for future proof work. cache.finish_warming().await; let total_cached_proofs = cache.len().await; info!( @@ -638,7 +659,12 @@ mod tests { cache.start_warming(base_root).await; let (proof_tx, proof_rx) = mpsc::channel(1); - let collector = spawn_late_proof_collector(cache.clone(), proof_rx, tree); + let collector = spawn_late_proof_collector( + cache.clone(), + proof_rx, + tree, + LATE_PROOF_COLLECTION_TIMEOUT, + ); assert!(cache.is_warming().await); proof_tx @@ -663,7 +689,12 @@ mod tests { cache.start_warming(base_root).await; let (proof_tx, proof_rx) = mpsc::channel(1); - let collector = spawn_late_proof_collector(cache.clone(), proof_rx, tree); + let collector = spawn_late_proof_collector( + cache.clone(), + proof_rx, + tree, + LATE_PROOF_COLLECTION_TIMEOUT, + ); proof_tx .send(ProofJobResult { seq: 0, @@ -682,4 +713,21 @@ mod tests { assert!(!cache.is_warming().await); assert!(cache.is_empty().await); } + + #[tokio::test] + async fn retained_proof_sender_cannot_block_cache_completion() { + let tree = solana_sdk::pubkey::Pubkey::new_unique(); + let base_root = [5u8; 32]; + let cache = Arc::new(SharedProofCache::new(tree)); + cache.start_warming(base_root).await; + + let (proof_tx, proof_rx) = mpsc::channel(1); + let collector = + spawn_late_proof_collector(cache.clone(), proof_rx, tree, Duration::from_millis(20)); + + assert!(cache.is_warming().await); + assert_eq!(collector.await.unwrap(), 0); + assert!(!cache.is_warming().await); + assert!(proof_tx.is_closed()); + } } From 05a85575d57c94eafce8ae7100322126c9ed2649 Mon Sep 17 00:00:00 2001 From: Sergey Timoshin Date: Thu, 24 Sep 2026 12:22:59 +0000 Subject: [PATCH 04/10] fix(forester): release cancelled cache warmups --- forester/src/processor/v2/processor.rs | 6 +- forester/src/processor/v2/proof_cache.rs | 170 +++++++++++++++++++---- forester/src/processor/v2/tx_sender.rs | 39 ++++-- 3 files changed, 174 insertions(+), 41 deletions(-) diff --git a/forester/src/processor/v2/processor.rs b/forester/src/processor/v2/processor.rs index 4441cab3c5..e05244618f 100644 --- a/forester/src/processor/v2/processor.rs +++ b/forester/src/processor/v2/processor.rs @@ -572,7 +572,7 @@ where let num_batches = queue_data.num_batches; let num_workers = self.context.num_proof_workers.max(1); - cache.start_warming(initial_root).await; + let warmup = cache.start_warming(initial_root).await; let (proof_tx, mut proof_rx) = mpsc::channel(num_workers * 2); @@ -618,7 +618,7 @@ where } } - cache + warmup .add_proof(result.seq, result.old_root, result.new_root, instruction) .await; proofs_cached += 1; @@ -635,7 +635,7 @@ where } } - cache.finish_warming().await; + warmup.finish().await; if proofs_cached < jobs_sent { warn!( diff --git a/forester/src/processor/v2/proof_cache.rs b/forester/src/processor/v2/proof_cache.rs index 123b4acf98..b7c97567ca 100644 --- a/forester/src/processor/v2/proof_cache.rs +++ b/forester/src/processor/v2/proof_cache.rs @@ -1,4 +1,7 @@ -use std::collections::{BTreeMap, VecDeque}; +use std::{ + collections::{BTreeMap, VecDeque}, + sync::Arc, +}; use solana_sdk::pubkey::Pubkey; use tokio::sync::Mutex; @@ -25,6 +28,7 @@ pub struct ProofCache { proofs: VecDeque, warming_proofs: BTreeMap, is_warming: bool, + warming_generation: u64, max_proofs: usize, } @@ -36,31 +40,40 @@ impl ProofCache { proofs: VecDeque::new(), warming_proofs: BTreeMap::new(), is_warming: false, + warming_generation: 0, max_proofs: DEFAULT_MAX_CACHED_PROOFS, } } - pub fn start_warming(&mut self, base_root: [u8; 32]) { + pub fn start_warming(&mut self, base_root: [u8; 32]) -> u64 { debug!( "Starting cache warm-up for tree {} with root {:?}", self.tree, &base_root[..4] ); + self.warming_generation = self.warming_generation.wrapping_add(1); self.base_root = base_root; self.proofs.clear(); self.warming_proofs.clear(); self.is_warming = true; + self.warming_generation } pub fn add_proof( &mut self, + generation: u64, seq: u64, old_root: [u8; 32], new_root: [u8; 32], instruction: BatchInstruction, ) { - if !self.is_warming { - warn!("Attempted to add proof to cache that is not warming"); + if !self.is_warming || self.warming_generation != generation { + warn!( + tree = %self.tree, + generation, + active_generation = self.warming_generation, + "Attempted to add proof to an inactive cache warm-up" + ); return; } if self.warming_proofs.contains_key(&seq) { @@ -101,7 +114,17 @@ impl ProofCache { ); } - pub fn finish_warming(&mut self) { + pub fn finish_warming(&mut self, generation: u64) { + if !self.is_warming || self.warming_generation != generation { + debug!( + tree = %self.tree, + generation, + active_generation = self.warming_generation, + "Ignoring completion from an inactive cache warm-up" + ); + return; + } + self.is_warming = false; if self.warming_proofs.is_empty() { @@ -232,6 +255,20 @@ impl ProofCache { self.warming_proofs.clear(); self.is_warming = false; } + + fn abort_warming(&mut self, generation: u64) { + if !self.is_warming || self.warming_generation != generation { + return; + } + + self.warming_proofs.clear(); + self.is_warming = false; + warn!( + tree = %self.tree, + generation, + "Cache warm-up was cancelled; releasing warming state" + ); + } } pub struct SharedProofCache { @@ -251,25 +288,13 @@ impl SharedProofCache { } } - pub async fn start_warming(&self, base_root: [u8; 32]) { - self.inner.lock().await.start_warming(base_root); - } - - pub async fn add_proof( - &self, - seq: u64, - old_root: [u8; 32], - new_root: [u8; 32], - instruction: BatchInstruction, - ) { - self.inner - .lock() - .await - .add_proof(seq, old_root, new_root, instruction); - } - - pub async fn finish_warming(&self) { - self.inner.lock().await.finish_warming(); + pub async fn start_warming(self: &Arc, base_root: [u8; 32]) -> ProofCacheWarmup { + let generation = self.inner.lock().await.start_warming(base_root); + ProofCacheWarmup { + cache: self.clone(), + generation, + finished: false, + } } pub async fn take_if_valid(&self, current_root: &[u8; 32]) -> Option> { @@ -292,3 +317,100 @@ impl SharedProofCache { self.inner.lock().await.clear(); } } + +/// Owns one cache warm-up session and releases it if its future is cancelled. +/// +/// The generation prevents a cancelled, older session from clearing or +/// completing a newer session for the same tree. +pub struct ProofCacheWarmup { + cache: Arc, + generation: u64, + finished: bool, +} + +impl ProofCacheWarmup { + pub async fn add_proof( + &self, + seq: u64, + old_root: [u8; 32], + new_root: [u8; 32], + instruction: BatchInstruction, + ) { + self.cache.inner.lock().await.add_proof( + self.generation, + seq, + old_root, + new_root, + instruction, + ); + } + + pub async fn finish(mut self) { + self.cache + .inner + .lock() + .await + .finish_warming(self.generation); + self.finished = true; + } +} + +impl Drop for ProofCacheWarmup { + fn drop(&mut self) { + if self.finished { + return; + } + + if let Ok(mut cache) = self.cache.inner.try_lock() { + cache.abort_warming(self.generation); + return; + } + + let cache = self.cache.clone(); + let generation = self.generation; + tokio::spawn(async move { + cache.inner.lock().await.abort_warming(generation); + }); + } +} + +#[cfg(test)] +mod tests { + use std::future::pending; + + use super::*; + + #[tokio::test] + async fn cancelled_warmup_releases_warming_state() { + let cache = Arc::new(SharedProofCache::new(Pubkey::new_unique())); + let task_cache = cache.clone(); + + let task = tokio::spawn(async move { + let _warmup = task_cache.start_warming([1u8; 32]).await; + pending::<()>().await; + }); + + while !cache.is_warming().await { + tokio::task::yield_now().await; + } + + task.abort(); + let _ = task.await; + tokio::task::yield_now().await; + + assert!(!cache.is_warming().await); + } + + #[tokio::test] + async fn cancelled_old_warmup_does_not_clear_new_session() { + let cache = Arc::new(SharedProofCache::new(Pubkey::new_unique())); + let old_warmup = cache.start_warming([1u8; 32]).await; + let new_warmup = cache.start_warming([2u8; 32]).await; + + drop(old_warmup); + + assert!(cache.is_warming().await); + new_warmup.finish().await; + assert!(!cache.is_warming().await); + } +} diff --git a/forester/src/processor/v2/tx_sender.rs b/forester/src/processor/v2/tx_sender.rs index c9ebdd6723..af92e43926 100644 --- a/forester/src/processor/v2/tx_sender.rs +++ b/forester/src/processor/v2/tx_sender.rs @@ -25,8 +25,10 @@ use tracing::{debug, info, warn}; use crate::{ errors::ForesterError, processor::v2::{ - common::send_transaction_batch, proof_cache::SharedProofCache, - proof_worker::ProofJobResult, BatchContext, + common::send_transaction_batch, + proof_cache::{ProofCacheWarmup, SharedProofCache}, + proof_worker::ProofJobResult, + BatchContext, }, }; @@ -517,18 +519,18 @@ impl TxSender { let mut saved = 0; - cache.start_warming(self.last_seen_root).await; + let warmup = cache.start_warming(self.last_seen_root).await; // Save proofs from pending_batch (already processed but not yet sent) for (instruction, seq, old_root, new_root) in self.pending_batch.drain(..) { - cache.add_proof(seq, old_root, new_root, instruction).await; + warmup.add_proof(seq, old_root, new_root, instruction).await; saved += 1; } // Save proofs from the reorder buffer (received but waiting for in-order processing) while let Some(entry) = self.buffer.pop_next() { let seq = self.buffer.expected_seq() - 1; - cache + warmup .add_proof(seq, entry.old_root, entry.new_root, entry.instruction) .await; saved += 1; @@ -537,7 +539,7 @@ impl TxSender { // Save the current result if provided if let Some(result) = current_result { if let Ok(instruction) = result.result { - cache + warmup .add_proof(result.seq, result.old_root, result.new_root, instruction) .await; saved += 1; @@ -547,7 +549,7 @@ impl TxSender { // Drain remaining proofs from the channel while let Ok(result) = proof_rx.try_recv() { if let Ok(instruction) = result.result { - cache + warmup .add_proof(result.seq, result.old_root, result.new_root, instruction) .await; saved += 1; @@ -569,6 +571,7 @@ impl TxSender { // Dropping the JoinHandle detaches the collector so it can finish warming the cache. drop(spawn_late_proof_collector( + warmup, cache.clone(), proof_rx, self.context.merkle_tree, @@ -580,6 +583,7 @@ impl TxSender { } fn spawn_late_proof_collector( + warmup: ProofCacheWarmup, cache: Arc, mut proof_rx: mpsc::Receiver, tree: solana_sdk::pubkey::Pubkey, @@ -591,7 +595,7 @@ fn spawn_late_proof_collector( while let Some(result) = proof_rx.recv().await { match result.result { Ok(instruction) => { - cache + warmup .add_proof(result.seq, result.old_root, result.new_root, instruction) .await; saved += 1; @@ -622,7 +626,7 @@ fn spawn_late_proof_collector( // This must run on channel closure and timeout so a retained sender can // never make the tree permanently ineligible for future proof work. - cache.finish_warming().await; + warmup.finish().await; let total_cached_proofs = cache.len().await; info!( tree = %tree, @@ -656,10 +660,11 @@ mod tests { let base_root = [1u8; 32]; let next_root = [2u8; 32]; let cache = Arc::new(SharedProofCache::new(tree)); - cache.start_warming(base_root).await; + let warmup = cache.start_warming(base_root).await; let (proof_tx, proof_rx) = mpsc::channel(1); let collector = spawn_late_proof_collector( + warmup, cache.clone(), proof_rx, tree, @@ -686,10 +691,11 @@ mod tests { let tree = solana_sdk::pubkey::Pubkey::new_unique(); let base_root = [3u8; 32]; let cache = Arc::new(SharedProofCache::new(tree)); - cache.start_warming(base_root).await; + let warmup = cache.start_warming(base_root).await; let (proof_tx, proof_rx) = mpsc::channel(1); let collector = spawn_late_proof_collector( + warmup, cache.clone(), proof_rx, tree, @@ -719,11 +725,16 @@ mod tests { let tree = solana_sdk::pubkey::Pubkey::new_unique(); let base_root = [5u8; 32]; let cache = Arc::new(SharedProofCache::new(tree)); - cache.start_warming(base_root).await; + let warmup = cache.start_warming(base_root).await; let (proof_tx, proof_rx) = mpsc::channel(1); - let collector = - spawn_late_proof_collector(cache.clone(), proof_rx, tree, Duration::from_millis(20)); + let collector = spawn_late_proof_collector( + warmup, + cache.clone(), + proof_rx, + tree, + Duration::from_millis(20), + ); assert!(cache.is_warming().await); assert_eq!(collector.await.unwrap(), 0); From 998876cc0118610a3e553fb429f8f29ff8cb9290 Mon Sep 17 00:00:00 2001 From: Sergey Timoshin Date: Thu, 24 Sep 2026 13:15:14 +0000 Subject: [PATCH 05/10] fix(forester): flush partial proof batches before eligibility ends --- forester/src/processor/v2/proof_worker.rs | 28 +++++++++++++++++++- forester/src/processor/v2/tx_sender.rs | 31 +++++++++++++++++++++-- 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/forester/src/processor/v2/proof_worker.rs b/forester/src/processor/v2/proof_worker.rs index 98a81f24fb..01a74d1c91 100644 --- a/forester/src/processor/v2/proof_worker.rs +++ b/forester/src/processor/v2/proof_worker.rs @@ -194,8 +194,34 @@ async fn run_proof_pipeline( semaphore: Arc, ) -> crate::Result<()> { while let Ok(job) = job_rx.recv().await { + if job.result_tx.is_closed() { + debug!( + "Skipping cancelled proof job seq={}: result channel closed", + job.seq + ); + continue; + } + let clients = clients.clone(); - let permit = semaphore.clone().acquire_owned().await; + let permit = tokio::select! { + permit = semaphore.clone().acquire_owned() => permit, + _ = job.result_tx.closed() => { + debug!( + "Cancelling queued proof job seq={} while waiting for prover capacity", + job.seq + ); + continue; + } + }; + + if job.result_tx.is_closed() { + debug!( + "Skipping cancelled proof job seq={}: result channel closed", + job.seq + ); + continue; + } + // Spawn immediately so we don't block receiving the next job // while waiting for HTTP submission. Semaphore bounds concurrency. tokio::spawn(async move { diff --git a/forester/src/processor/v2/tx_sender.rs b/forester/src/processor/v2/tx_sender.rs index af92e43926..460dd83dfd 100644 --- a/forester/src/processor/v2/tx_sender.rs +++ b/forester/src/processor/v2/tx_sender.rs @@ -5,7 +5,14 @@ use borsh::BorshSerialize; const MAX_BUFFER_SIZE: usize = 1000; const V2_IXS_PER_TX_WITH_LUT: usize = 5; const V2_IXS_PER_TX_WITHOUT_LUT: usize = 4; -const FLUSH_MARGIN_SLOTS: u64 = 2; +/// Flush an incomplete transaction with enough time left for confirmation. +/// +/// `send_transaction_batch` requires at least four slots. The old two-slot +/// margin could therefore never send a partial batch: proof results would sit +/// in `pending_batch` until eligibility ended and then be handed back to the +/// cache. Ten slots leaves roughly four seconds on the default slot schedule +/// and gives the confirmation loop useful headroom. +const FLUSH_MARGIN_SLOTS: u64 = 10; /// Late proofs are an optimization. Do not let a proof job that never releases /// its sender keep the tree cache in the warming state indefinitely. const LATE_PROOF_COLLECTION_TIMEOUT: Duration = Duration::from_secs(30); @@ -391,7 +398,27 @@ impl TxSender { let result = match tokio::time::timeout(Duration::from_secs(1), proof_rx.recv()).await { Ok(Some(r)) => r, Ok(None) => break, - Err(_) => continue, + Err(_) => { + // A partial batch may have been waiting since the previous + // proof result. Re-check the slot on every receive timeout; + // otherwise it is only flushed when another proof arrives, + // which may be after the eligibility window has ended. + let current_slot = self.context.slot_tracker.estimated_current_slot(); + if !self.pending_batch.is_empty() + && self.should_flush_due_to_time_at(current_slot) + { + let batch = std::mem::replace( + &mut self.pending_batch, + Vec::with_capacity(self.ixs_per_tx), + ); + let earliest = self.pending_batch_earliest_submit.take(); + + if batch_tx.send((batch, earliest)).is_err() { + break; + } + } + continue; + } }; let current_slot = self.context.slot_tracker.estimated_current_slot(); From 30601b03f2366dd143c9fc44e4a511b592412312 Mon Sep 17 00:00:00 2001 From: Sergey Timoshin Date: Thu, 24 Sep 2026 15:36:55 +0000 Subject: [PATCH 06/10] fix(forester): refinalize stale v2 eligibility --- forester/src/epoch_manager.rs | 2 ++ forester/src/processor/v2/processor.rs | 9 +++++++++ 2 files changed, 11 insertions(+) diff --git a/forester/src/epoch_manager.rs b/forester/src/epoch_manager.rs index 7e48ffebb7..9673b49e97 100644 --- a/forester/src/epoch_manager.rs +++ b/forester/src/epoch_manager.rs @@ -3508,6 +3508,7 @@ impl EpochManager { let mut proc = processor.lock().await; match proc.process().await { Ok(res) => Ok(res), + Err(error) if error.is_forester_not_eligible() => Err(error), Err(error) if matches!(&error, ForesterError::V2(v2_error) if v2_error.is_constraint()) => { warn!( @@ -3575,6 +3576,7 @@ impl EpochManager { let mut proc = processor.lock().await; match proc.process().await { Ok(res) => Ok(res), + Err(error) if error.is_forester_not_eligible() => Err(error), Err(error) if matches!(&error, ForesterError::V2(v2_error) if v2_error.is_constraint()) => { warn!( diff --git a/forester/src/processor/v2/processor.rs b/forester/src/processor/v2/processor.rs index e05244618f..460a088b94 100644 --- a/forester/src/processor/v2/processor.rs +++ b/forester/src/processor/v2/processor.rs @@ -340,6 +340,15 @@ where drop(proof_tx); let tx_result = match tx_sender_handle.await.map_err(ForesterError::from)? { + Err(error) if error.is_forester_not_eligible() => { + warn!( + event = "v2_tx_sender_stale_eligibility", + tree = %self.context.merkle_tree, + error = %error, + "Tx sender detected stale forester eligibility" + ); + return Err(error); + } Err(error) if matches!(&error, ForesterError::V2(v2_error) if v2_error.is_constraint()) => { warn!( From 898a781e0e01b6823f2d0fa3f8411768a6842452 Mon Sep 17 00:00:00 2001 From: Sergey Timoshin Date: Thu, 24 Sep 2026 15:48:20 +0000 Subject: [PATCH 07/10] fix(forester): bound speculative proof warming --- forester/src/epoch_manager.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/forester/src/epoch_manager.rs b/forester/src/epoch_manager.rs index 9673b49e97..62e2153546 100644 --- a/forester/src/epoch_manager.rs +++ b/forester/src/epoch_manager.rs @@ -3693,12 +3693,16 @@ impl EpochManager { return; } - if slots_until_active < 15 { + // The remote prover may need 20-35 seconds for an uncached proof. Starting + // speculative work with less time than that only leaves orphaned requests + // in the prover queue after this timeout expires, delaying active work. + const MIN_PREWARM_SLOTS: u64 = 90; + if slots_until_active < MIN_PREWARM_SLOTS { info!( event = "prewarm_skipped_not_enough_time", run_id = %self.run_id, slots_until_active, - min_required_slots = 15, + min_required_slots = MIN_PREWARM_SLOTS, "Skipping pre-warming; not enough slots until active phase" ); return; @@ -3776,7 +3780,9 @@ impl EpochManager { } }; - const PREWARM_MAX_BATCHES: usize = 4; + // One speculative batch is sufficient to warm this tree. More + // batches can monopolize a small shared prover deployment. + const PREWARM_MAX_BATCHES: usize = 1; let mut p = processor.lock().await; match p .prewarm_from_indexer( From 988849f2c123ebdf877a2f4ebf9283fffac6dfef Mon Sep 17 00:00:00 2001 From: Sergey Timoshin Date: Thu, 24 Sep 2026 16:13:26 +0000 Subject: [PATCH 08/10] fix(forester): retain proofs after warmup timeout --- forester/src/processor/v2/proof_cache.rs | 159 +++++++++++++++++------ forester/src/processor/v2/tx_sender.rs | 87 +++++++++++-- 2 files changed, 199 insertions(+), 47 deletions(-) diff --git a/forester/src/processor/v2/proof_cache.rs b/forester/src/processor/v2/proof_cache.rs index b7c97567ca..c7d5ac0d7b 100644 --- a/forester/src/processor/v2/proof_cache.rs +++ b/forester/src/processor/v2/proof_cache.rs @@ -26,6 +26,10 @@ pub struct ProofCache { tree: Pubkey, base_root: [u8; 32], proofs: VecDeque, + /// Proofs that completed after the cache warm-up deadline. These are kept + /// separately because a newer warm-up must not discard potentially useful + /// work from an older eligibility window. + late_proofs: VecDeque, warming_proofs: BTreeMap, is_warming: bool, warming_generation: u64, @@ -38,6 +42,7 @@ impl ProofCache { tree, base_root: [0u8; 32], proofs: VecDeque::new(), + late_proofs: VecDeque::new(), warming_proofs: BTreeMap::new(), is_warming: false, warming_generation: 0, @@ -160,25 +165,73 @@ impl ProofCache { ); } - pub fn take_if_valid(&mut self, current_root: &[u8; 32]) -> Option> { - if self.proofs.is_empty() || self.is_warming { - return None; + pub fn add_late_proof( + &mut self, + seq: u64, + old_root: [u8; 32], + new_root: [u8; 32], + instruction: BatchInstruction, + ) { + let duplicate = self + .proofs + .iter() + .chain(self.warming_proofs.values()) + .chain(self.late_proofs.iter()) + .any(|proof| proof.old_root == old_root && proof.new_root == new_root); + if duplicate { + debug!( + tree = %self.tree, + seq, + "Ignoring duplicate late proof" + ); + return; } - let mut skipped = 0; - while let Some(proof) = self.proofs.front() { - if proof.old_root == *current_root { - break; - } - if proof.new_root == *current_root { - self.proofs.pop_front(); - skipped += 1; - continue; + let items = instruction.items_count(); + self.late_proofs.push_back(CachedProof { + seq, + old_root, + new_root, + instruction, + items, + }); + + while self.late_proofs.len() > self.max_proofs { + if let Some(dropped) = self.late_proofs.pop_front() { + warn!( + tree = %self.tree, + seq = dropped.seq, + max = self.max_proofs, + "Late proof cache limit reached; dropping oldest candidate" + ); } - self.proofs.pop_front(); - skipped += 1; } + debug!( + tree = %self.tree, + seq, + late_proofs_cached = self.late_proofs.len(), + "Retained proof that completed after cache warm-up deadline" + ); + } + + pub fn take_if_valid(&mut self, current_root: &[u8; 32]) -> Option> { + if self.is_warming || (self.proofs.is_empty() && self.late_proofs.is_empty()) { + return None; + } + + // Treat both synchronously warmed and later results as candidates. Proof + // completion order is not guaranteed, so build the usable chain by roots + // instead of sequence number. Unmatched candidates stay cached: a missing + // predecessor may still arrive, or the on-chain root may advance to them. + let mut candidates = VecDeque::with_capacity(self.proofs.len() + self.late_proofs.len()); + candidates.append(&mut self.proofs); + candidates.append(&mut self.late_proofs); + + let before = candidates.len(); + candidates.retain(|proof| proof.new_root != *current_root); + let skipped = before - candidates.len(); + if skipped > 0 { debug!( "Skipped {} stale cached proofs for tree {} (on-chain already advanced)", @@ -186,36 +239,29 @@ impl ProofCache { ); } - if self.proofs.is_empty() { - debug!( - "Cache empty after skipping stale proofs for tree {} (current_root {:?})", - self.tree, - ¤t_root[..4] - ); - return None; - } - let mut expected = *current_root; let mut taken: Vec = Vec::new(); - while let Some(proof) = self.proofs.pop_front() { - if proof.old_root != expected { - warn!( - "Cache chain broken for tree {} at seq {}: expected root {:?}, got {:?}. Dropping remaining {} proofs.", - self.tree, - proof.seq, - &expected[..4], - &proof.old_root[..4], - self.proofs.len() - ); - self.proofs.clear(); - break; - } + while let Some(position) = candidates + .iter() + .position(|proof| proof.old_root == expected) + { + let proof = candidates + .remove(position) + .expect("candidate position was found"); expected = proof.new_root; taken.push(proof); } + self.late_proofs = candidates; + if taken.is_empty() { + debug!( + tree = %self.tree, + current_root = ?¤t_root[..4], + retained_candidates = self.late_proofs.len(), + "No cached proof currently links to the on-chain root" + ); return None; } @@ -235,11 +281,11 @@ impl ProofCache { } pub fn len(&self) -> usize { - self.proofs.len() + self.proofs.len() + self.late_proofs.len() } pub fn is_empty(&self) -> bool { - self.proofs.is_empty() + self.proofs.is_empty() && self.late_proofs.is_empty() } pub fn is_warming(&self) -> bool { @@ -252,6 +298,7 @@ impl ProofCache { pub fn clear(&mut self) { self.proofs.clear(); + self.late_proofs.clear(); self.warming_proofs.clear(); self.is_warming = false; } @@ -301,6 +348,19 @@ impl SharedProofCache { self.inner.lock().await.take_if_valid(current_root) } + pub async fn add_late_proof( + &self, + seq: u64, + old_root: [u8; 32], + new_root: [u8; 32], + instruction: BatchInstruction, + ) { + self.inner + .lock() + .await + .add_late_proof(seq, old_root, new_root, instruction); + } + pub async fn is_warming(&self) -> bool { self.inner.lock().await.is_warming() } @@ -413,4 +473,27 @@ mod tests { new_warmup.finish().await; assert!(!cache.is_warming().await); } + + #[tokio::test] + async fn late_proofs_are_linked_by_root_when_they_arrive_out_of_order() { + let cache = Arc::new(SharedProofCache::new(Pubkey::new_unique())); + let root_1 = [1u8; 32]; + let root_2 = [2u8; 32]; + let root_3 = [3u8; 32]; + + cache + .add_late_proof(1, root_2, root_3, BatchInstruction::Append(Vec::new())) + .await; + cache + .add_late_proof(0, root_1, root_2, BatchInstruction::Append(Vec::new())) + .await; + + let proofs = cache.take_if_valid(&root_1).await.unwrap(); + assert_eq!(proofs.len(), 2); + assert_eq!(proofs[0].old_root, root_1); + assert_eq!(proofs[0].new_root, root_2); + assert_eq!(proofs[1].old_root, root_2); + assert_eq!(proofs[1].new_root, root_3); + assert!(cache.is_empty().await); + } } diff --git a/forester/src/processor/v2/tx_sender.rs b/forester/src/processor/v2/tx_sender.rs index 460dd83dfd..9b345a09bb 100644 --- a/forester/src/processor/v2/tx_sender.rs +++ b/forester/src/processor/v2/tx_sender.rs @@ -16,6 +16,10 @@ const FLUSH_MARGIN_SLOTS: u64 = 10; /// Late proofs are an optimization. Do not let a proof job that never releases /// its sender keep the tree cache in the warming state indefinitely. const LATE_PROOF_COLLECTION_TIMEOUT: Duration = Duration::from_secs(30); +/// Keep receiving proof results after releasing the warming lock. This matches +/// the default maximum prover wait and bounds detached collector lifetime when +/// a producer retains its sender indefinitely. +const LATE_PROOF_RETENTION_TIMEOUT: Duration = Duration::from_secs(600); use light_batched_merkle_tree::merkle_tree::{ InstructionDataBatchAppendInputs, InstructionDataBatchNullifyInputs, @@ -603,6 +607,7 @@ impl TxSender { proof_rx, self.context.merkle_tree, LATE_PROOF_COLLECTION_TIMEOUT, + LATE_PROOF_RETENTION_TIMEOUT, )); saved @@ -615,6 +620,7 @@ fn spawn_late_proof_collector( mut proof_rx: mpsc::Receiver, tree: solana_sdk::pubkey::Pubkey, collection_timeout: Duration, + retention_timeout: Duration, ) -> JoinHandle { tokio::spawn(async move { let mut saved = 0usize; @@ -639,21 +645,69 @@ fn spawn_late_proof_collector( } }; - if tokio::time::timeout(collection_timeout, collection) + let collection_timed_out = tokio::time::timeout(collection_timeout, collection) .await - .is_err() - { + .is_err(); + + // Release the tree immediately at the scheduling deadline. The + // receiver remains alive below so completed work is still retained. + warmup.finish().await; + + if collection_timed_out { warn!( tree = %tree, timeout_ms = collection_timeout.as_millis(), late_proofs_cached = saved, - "Late proof collection timed out; releasing cache warming state" + "Late proof collection timed out; released cache warming state and retaining later results" ); + + let mut retained = 0usize; + let retention = async { + while let Some(result) = proof_rx.recv().await { + match result.result { + Ok(instruction) => { + cache + .add_late_proof( + result.seq, + result.old_root, + result.new_root, + instruction, + ) + .await; + saved += 1; + retained += 1; + } + Err(error) => { + warn!( + tree = %tree, + seq = result.seq, + error = %error, + "Proof failed after cache warm-up deadline" + ); + } + } + } + }; + + if tokio::time::timeout(retention_timeout, retention) + .await + .is_err() + { + warn!( + tree = %tree, + timeout_ms = retention_timeout.as_millis(), + retained_late_proofs = retained, + "Late proof retention window expired" + ); + } else if retained > 0 { + info!( + tree = %tree, + retained_late_proofs = retained, + "Retained proofs that completed after cache warm-up deadline" + ); + } } - // This must run on channel closure and timeout so a retained sender can - // never make the tree permanently ineligible for future proof work. - warmup.finish().await; let total_cached_proofs = cache.len().await; info!( tree = %tree, @@ -696,6 +750,7 @@ mod tests { proof_rx, tree, LATE_PROOF_COLLECTION_TIMEOUT, + LATE_PROOF_RETENTION_TIMEOUT, ); assert!(cache.is_warming().await); @@ -727,6 +782,7 @@ mod tests { proof_rx, tree, LATE_PROOF_COLLECTION_TIMEOUT, + LATE_PROOF_RETENTION_TIMEOUT, ); proof_tx .send(ProofJobResult { @@ -761,11 +817,24 @@ mod tests { proof_rx, tree, Duration::from_millis(20), + Duration::from_secs(1), ); assert!(cache.is_warming().await); - assert_eq!(collector.await.unwrap(), 0); + tokio::time::sleep(Duration::from_millis(40)).await; assert!(!cache.is_warming().await); - assert!(proof_tx.is_closed()); + assert!(!proof_tx.is_closed()); + + let next_root = [6u8; 32]; + proof_tx + .send(proof_result(0, base_root, next_root)) + .await + .unwrap(); + drop(proof_tx); + + assert_eq!(collector.await.unwrap(), 1); + let cached = cache.take_if_valid(&base_root).await.unwrap(); + assert_eq!(cached.len(), 1); + assert_eq!(cached[0].new_root, next_root); } } From 15fc4ab28eec0a04e3c9a9afec1c294db9aabb1f Mon Sep 17 00:00:00 2001 From: Sergey Timoshin Date: Thu, 24 Sep 2026 17:20:00 +0000 Subject: [PATCH 09/10] fix(forester): stream ready proofs and retain unconfirmed chains --- forester/src/epoch_manager.rs | 149 ++++++---- forester/src/processor/v2/processor.rs | 5 +- forester/src/processor/v2/proof_cache.rs | 323 +++++++++++++++------- forester/src/processor/v2/proof_worker.rs | 80 +++++- forester/src/processor/v2/tx_sender.rs | 59 +++- 5 files changed, 448 insertions(+), 168 deletions(-) diff --git a/forester/src/epoch_manager.rs b/forester/src/epoch_manager.rs index 62e2153546..79db0cf60d 100644 --- a/forester/src/epoch_manager.rs +++ b/forester/src/epoch_manager.rs @@ -2334,43 +2334,6 @@ impl EpochManager { return Ok(()); } - if let Some(cache) = self - .proof_caches - .get(&tree_pubkey) - .map(|cache| cache.clone()) - { - if cache.is_warming().await { - debug!( - event = "v2_proof_work_deferred_cache_warming", - run_id = %self.run_id, - tree = %tree_pubkey, - "Deferring V2 proof work while late proofs are collected" - ); - return Ok(()); - } - } - - // Try to send any cached proofs first - let cached_send_start = Instant::now(); - if let Some(items_sent) = self - .try_send_cached_proofs(epoch_info, tree_accounts, consecutive_eligibility_end) - .await? - { - if items_sent > 0 { - let cached_send_duration = cached_send_start.elapsed(); - info!( - event = "cached_proofs_sent", - run_id = %self.run_id, - tree = %tree_pubkey, - items = items_sent, - duration_ms = cached_send_duration.as_millis() as u64, - "Sent items from proof cache" - ); - self.update_metrics_and_counts(epoch_info.epoch, items_sent, cached_send_duration) - .await; - } - } - let mut estimated_slot = self.slot_tracker.estimated_current_slot(); // Adaptive queue polling: start responsive, then back off (capped) while the @@ -2380,6 +2343,7 @@ impl EpochManager { // load (and can exhaust a shared RPC credit budget). const POLL_INTERVAL_MIN: Duration = Duration::from_millis(200); const POLL_INTERVAL_MAX: Duration = Duration::from_secs(10); + const PENDING_PROOF_POLL_INTERVAL: Duration = Duration::from_secs(1); let mut poll_interval = POLL_INTERVAL_MIN; 'inner_processing_loop: loop { @@ -2419,6 +2383,37 @@ impl EpochManager { break 'inner_processing_loop; } + // Consume ready prefixes even while later proofs are still arriving. + // Recheck throughout the slot instead of forfeiting it on a warming cache. + if self + .try_send_cached_proofs( + epoch_info, + epoch_pda, + tree_accounts, + consecutive_eligibility_end, + ) + .await? + .is_some() + { + tokio::time::sleep(POLL_INTERVAL_MIN).await; + estimated_slot = self.slot_tracker.estimated_current_slot(); + continue; + } + + let cache = self + .proof_caches + .get(&tree_pubkey) + .map(|cache| cache.clone()); + if let Some(cache) = cache { + if cache.has_pending_proofs().await { + // Do not generate another speculative chain for the same tree. + // Completed prefixes remain eligible for sending on every tick. + tokio::time::sleep(PENDING_PROOF_POLL_INTERVAL).await; + estimated_slot = self.slot_tracker.estimated_current_slot(); + continue; + } + } + // Process directly - the processor fetches queue data from the indexer let processing_start_time = Instant::now(); match self @@ -3723,7 +3718,7 @@ impl EpochManager { .or_insert_with(|| Arc::new(SharedProofCache::new(tree_pubkey))) .clone(); - if cache.is_warming().await { + if cache.has_pending_proofs().await { info!( event = "prewarm_skipped_cache_warming", run_id = %self_clone.run_id, @@ -3856,9 +3851,10 @@ impl EpochManager { async fn try_send_cached_proofs( &self, epoch_info: &Epoch, + epoch_pda: &ForesterEpochPda, tree_accounts: &TreeAccounts, consecutive_eligibility_end: u64, - ) -> Result> { + ) -> std::result::Result, ForesterError> { let tree_pubkey = tree_accounts.merkle_tree; // Check eligibility window before attempting to send cached proofs @@ -3894,15 +3890,12 @@ impl EpochManager { None => return Ok(None), }; - if cache.is_warming().await { - debug!( - event = "cached_proofs_skipped_cache_warming", - run_id = %self.run_id, - tree = %tree_pubkey, - "Skipping cached proofs because cache is still warming" - ); + if cache.is_empty().await { return Ok(None); } + let Some(_send_guard) = cache.try_lock_for_sending() else { + return Ok(Some(0)); + }; let mut rpc = self.rpc_pool.get_connection().await?; let current_root = match self.fetch_current_root(&mut *rpc, tree_accounts).await { @@ -3919,7 +3912,8 @@ impl EpochManager { } }; - let cached_proofs = match cache.take_if_valid(¤t_root).await { + drop(rpc); + let cached_proofs = match cache.ready_chain(¤t_root).await { Some(proofs) => proofs, None => { debug!( @@ -3949,8 +3943,11 @@ impl EpochManager { let items_sent = self .send_cached_proofs_as_transactions( epoch_info, + epoch_pda, tree_accounts, + &cache, cached_proofs, + consecutive_eligibility_end, confirmation_deadline, ) .await?; @@ -3986,19 +3983,50 @@ impl EpochManager { Ok(root) } + #[allow(clippy::too_many_arguments)] async fn send_cached_proofs_as_transactions( &self, epoch_info: &Epoch, + epoch_pda: &ForesterEpochPda, tree_accounts: &TreeAccounts, + cache: &SharedProofCache, cached_proofs: Vec, + consecutive_eligibility_end: u64, confirmation_deadline: Instant, - ) -> Result { + ) -> std::result::Result { let mut total_items = 0; let authority = self.config.payer_keypair.pubkey(); let derivation = self.config.derivation_pubkey; + let zkp_batch_size = self + .zkp_batch_sizes + .get(&tree_accounts.merkle_tree) + .map(|size| *size as usize) + .ok_or_else(|| anyhow!("Missing ZKP batch size for cached proof tree"))?; const PROOFS_PER_TX: usize = 4; for chunk in cached_proofs.chunks(PROOFS_PER_TX) { + let current_slot = self.slot_tracker.estimated_current_slot(); + if current_slot < epoch_info.phases.active.start + || current_slot >= consecutive_eligibility_end + || Instant::now() >= confirmation_deadline + { + break; + } + let light_slot = (current_slot - epoch_info.phases.active.start) + / epoch_pda.protocol_config.slot_length; + if !self + .check_forester_eligibility( + epoch_pda, + light_slot, + &tree_accounts.merkle_tree, + epoch_info.epoch, + epoch_info, + ) + .await? + { + return Err(ForesterError::NotEligible); + } + let send_started = Instant::now(); let mut instructions = Vec::new(); let mut chunk_items = 0; @@ -4048,7 +4076,7 @@ impl EpochManager { } } } - chunk_items += proof.items; + chunk_items += proof.items * zkp_batch_size; } if !instructions.is_empty() { @@ -4078,22 +4106,33 @@ impl EpochManager { .map_err(RpcError::from) { Ok(sig) => { + cache.confirm(chunk).await; info!( event = "cached_proofs_tx_sent", run_id = %self.run_id, signature = %sig, instruction_count, + queue_items = chunk_items, "Sent cached proofs transaction" ); total_items += chunk_items; + self.update_metrics_and_counts( + epoch_info.epoch, + chunk_items, + send_started.elapsed(), + ) + .await; } Err(e) => { warn!( event = "cached_proofs_tx_send_failed", run_id = %self.run_id, error = ?e, - "Failed to send cached proofs transaction" + "Cached proof send failed; preserving unconfirmed chain for retry" ); + // Stop the dependent chain. In particular, preserve the + // typed 6004 error so process_queue re-finalizes eligibility. + return Err(e.into()); } } } @@ -4818,6 +4857,16 @@ mod tests { ForesterConfig, }; + #[test] + fn cached_send_rpc_error_preserves_eligibility_recovery_signal() { + let error = RpcError::TransactionError(TransactionError::InstructionError( + 2, + InstructionError::Custom(6004), + )); + let error = ForesterError::from(error); + assert!(error.is_forester_not_eligible()); + } + fn create_test_config_with_skip_flags( skip_v1_state: bool, skip_v1_address: bool, diff --git a/forester/src/processor/v2/processor.rs b/forester/src/processor/v2/processor.rs index 460a088b94..72fb286b22 100644 --- a/forester/src/processor/v2/processor.rs +++ b/forester/src/processor/v2/processor.rs @@ -257,9 +257,8 @@ where pub async fn clear_cache(&mut self) { self.cached_state = None; - if let Some(proof_cache) = &self.proof_cache { - proof_cache.clear().await; - } + // Staging state is optimistic, but completed proofs are independently + // root-validated. A staging reset must not erase unconfirmed work. } pub fn update_eligibility(&mut self, end_slot: u64) { diff --git a/forester/src/processor/v2/proof_cache.rs b/forester/src/processor/v2/proof_cache.rs index c7d5ac0d7b..c288cc51a0 100644 --- a/forester/src/processor/v2/proof_cache.rs +++ b/forester/src/processor/v2/proof_cache.rs @@ -1,10 +1,13 @@ use std::{ - collections::{BTreeMap, VecDeque}, - sync::Arc, + collections::VecDeque, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, }; use solana_sdk::pubkey::Pubkey; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, MutexGuard}; use tracing::{debug, info, warn}; use super::tx_sender::BatchInstruction; @@ -26,11 +29,6 @@ pub struct ProofCache { tree: Pubkey, base_root: [u8; 32], proofs: VecDeque, - /// Proofs that completed after the cache warm-up deadline. These are kept - /// separately because a newer warm-up must not discard potentially useful - /// work from an older eligibility window. - late_proofs: VecDeque, - warming_proofs: BTreeMap, is_warming: bool, warming_generation: u64, max_proofs: usize, @@ -42,8 +40,6 @@ impl ProofCache { tree, base_root: [0u8; 32], proofs: VecDeque::new(), - late_proofs: VecDeque::new(), - warming_proofs: BTreeMap::new(), is_warming: false, warming_generation: 0, max_proofs: DEFAULT_MAX_CACHED_PROOFS, @@ -58,8 +54,7 @@ impl ProofCache { ); self.warming_generation = self.warming_generation.wrapping_add(1); self.base_root = base_root; - self.proofs.clear(); - self.warming_proofs.clear(); + // A new collection must not erase completed work from an earlier one. self.is_warming = true; self.warming_generation } @@ -72,51 +67,15 @@ impl ProofCache { new_root: [u8; 32], instruction: BatchInstruction, ) { - if !self.is_warming || self.warming_generation != generation { - warn!( + if self.warming_generation != generation { + debug!( tree = %self.tree, generation, active_generation = self.warming_generation, - "Attempted to add proof to an inactive cache warm-up" - ); - return; - } - if self.warming_proofs.contains_key(&seq) { - warn!( - "Duplicate cached proof seq={} for tree {}, ignoring", - seq, self.tree - ); - return; - } - - let items = instruction.items_count(); - self.warming_proofs.insert( - seq, - CachedProof { - seq, - old_root, - new_root, - instruction, - items, - }, - ); - - while self.warming_proofs.len() > self.max_proofs { - let Some((&last_seq, _)) = self.warming_proofs.last_key_value() else { - break; - }; - self.warming_proofs.remove(&last_seq); - warn!( - "Proof cache warm-up limit reached for tree {} (max={}), dropping newest seq={}", - self.tree, self.max_proofs, last_seq + "Retaining proof from an older cache collection" ); } - debug!( - "Cached proof seq={} for tree {} (total cached: {})", - seq, - self.tree, - self.warming_proofs.len() - ); + self.add_late_proof(seq, old_root, new_root, instruction); } pub fn finish_warming(&mut self, generation: u64) { @@ -132,31 +91,6 @@ impl ProofCache { self.is_warming = false; - if self.warming_proofs.is_empty() { - self.proofs.clear(); - info!( - "Cache warm-up complete for tree {}: 0 proofs cached with root {:?}", - self.tree, - &self.base_root[..4] - ); - return; - } - - if let Some(first) = self.warming_proofs.values().next() { - if self.base_root != [0u8; 32] && self.base_root != first.old_root { - warn!( - "First cached proof root mismatch for tree {}: base_root={:?}, proof.old_root={:?} (seq={})", - self.tree, - &self.base_root[..4], - &first.old_root[..4], - first.seq - ); - } - } - - self.proofs = self.warming_proofs.values().cloned().collect(); - self.warming_proofs.clear(); - info!( "Cache warm-up complete for tree {}: {} proofs cached with root {:?}", self.tree, @@ -175,8 +109,6 @@ impl ProofCache { let duplicate = self .proofs .iter() - .chain(self.warming_proofs.values()) - .chain(self.late_proofs.iter()) .any(|proof| proof.old_root == old_root && proof.new_root == new_root); if duplicate { debug!( @@ -188,7 +120,7 @@ impl ProofCache { } let items = instruction.items_count(); - self.late_proofs.push_back(CachedProof { + self.proofs.push_back(CachedProof { seq, old_root, new_root, @@ -196,8 +128,8 @@ impl ProofCache { items, }); - while self.late_proofs.len() > self.max_proofs { - if let Some(dropped) = self.late_proofs.pop_front() { + while self.proofs.len() > self.max_proofs { + if let Some(dropped) = self.proofs.pop_front() { warn!( tree = %self.tree, seq = dropped.seq, @@ -210,13 +142,15 @@ impl ProofCache { debug!( tree = %self.tree, seq, - late_proofs_cached = self.late_proofs.len(), - "Retained proof that completed after cache warm-up deadline" + cached_proofs = self.proofs.len(), + "Cached proof is ready for root-linked reuse" ); } - pub fn take_if_valid(&mut self, current_root: &[u8; 32]) -> Option> { - if self.is_warming || (self.proofs.is_empty() && self.late_proofs.is_empty()) { + /// Snapshot the usable prefix without consuming it. Failed, timed-out or + /// cancelled sends leave these proofs available until confirmation. + pub fn ready_chain(&mut self, current_root: &[u8; 32]) -> Option> { + if self.proofs.is_empty() { return None; } @@ -224,13 +158,10 @@ impl ProofCache { // completion order is not guaranteed, so build the usable chain by roots // instead of sequence number. Unmatched candidates stay cached: a missing // predecessor may still arrive, or the on-chain root may advance to them. - let mut candidates = VecDeque::with_capacity(self.proofs.len() + self.late_proofs.len()); - candidates.append(&mut self.proofs); - candidates.append(&mut self.late_proofs); - - let before = candidates.len(); - candidates.retain(|proof| proof.new_root != *current_root); - let skipped = before - candidates.len(); + let before = self.proofs.len(); + self.proofs.retain(|proof| proof.new_root != *current_root); + let skipped = before - self.proofs.len(); + let mut candidates = self.proofs.clone(); if skipped > 0 { debug!( @@ -253,13 +184,11 @@ impl ProofCache { taken.push(proof); } - self.late_proofs = candidates; - if taken.is_empty() { debug!( tree = %self.tree, current_root = ?¤t_root[..4], - retained_candidates = self.late_proofs.len(), + retained_candidates = self.proofs.len(), "No cached proof currently links to the on-chain root" ); return None; @@ -281,11 +210,19 @@ impl ProofCache { } pub fn len(&self) -> usize { - self.proofs.len() + self.late_proofs.len() + self.proofs.len() } pub fn is_empty(&self) -> bool { - self.proofs.is_empty() && self.late_proofs.is_empty() + self.proofs.is_empty() + } + + pub fn confirm(&mut self, confirmed: &[CachedProof]) { + self.proofs.retain(|candidate| { + !confirmed.iter().any(|proof| { + proof.old_root == candidate.old_root && proof.new_root == candidate.new_root + }) + }); } pub fn is_warming(&self) -> bool { @@ -298,8 +235,6 @@ impl ProofCache { pub fn clear(&mut self) { self.proofs.clear(); - self.late_proofs.clear(); - self.warming_proofs.clear(); self.is_warming = false; } @@ -308,7 +243,7 @@ impl ProofCache { return; } - self.warming_proofs.clear(); + // Cancellation releases scheduling state, never already completed proofs. self.is_warming = false; warn!( tree = %self.tree, @@ -320,6 +255,8 @@ impl ProofCache { pub struct SharedProofCache { inner: Mutex, + sending: Mutex<()>, + pending_collections: AtomicUsize, } impl std::fmt::Debug for SharedProofCache { @@ -332,6 +269,8 @@ impl SharedProofCache { pub fn new(tree: Pubkey) -> Self { Self { inner: Mutex::new(ProofCache::new(tree)), + sending: Mutex::new(()), + pending_collections: AtomicUsize::new(0), } } @@ -344,8 +283,26 @@ impl SharedProofCache { } } - pub async fn take_if_valid(&self, current_root: &[u8; 32]) -> Option> { - self.inner.lock().await.take_if_valid(current_root) + pub async fn ready_chain(&self, current_root: &[u8; 32]) -> Option> { + self.inner.lock().await.ready_chain(current_root) + } + + pub async fn confirm(&self, confirmed: &[CachedProof]) { + self.inner.lock().await.confirm(confirmed); + } + + /// Serialize cached sends without blocking proof collection. + pub fn try_lock_for_sending(&self) -> Option> { + self.sending.try_lock().ok() + } + + pub fn start_collecting(self: &Arc) -> ProofCollection { + self.pending_collections.fetch_add(1, Ordering::AcqRel); + ProofCollection(self.clone()) + } + + pub async fn has_pending_proofs(&self) -> bool { + self.pending_collections.load(Ordering::Acquire) > 0 || self.is_warming().await } pub async fn add_late_proof( @@ -378,6 +335,16 @@ impl SharedProofCache { } } +/// Bounds speculative work for a tree across the late-result retention period. +/// Dropping or aborting a collector always releases the scheduling gate. +pub struct ProofCollection(Arc); + +impl Drop for ProofCollection { + fn drop(&mut self) { + self.0.pending_collections.fetch_sub(1, Ordering::AcqRel); + } +} + /// Owns one cache warm-up session and releases it if its future is cancelled. /// /// The generation prevents a cancelled, older session from clearing or @@ -440,6 +407,152 @@ mod tests { use super::*; + fn add(cache: &mut ProofCache, seq: u64, old: u8, new: u8) { + cache.add_late_proof( + seq, + [old; 32], + [new; 32], + BatchInstruction::Append(Vec::new()), + ); + } + + #[test] + fn ready_prefix_is_available_during_collection_and_until_confirmed() { + let mut cache = ProofCache::new(Pubkey::new_unique()); + let generation = cache.start_warming([1; 32]); + cache.add_proof( + generation, + 0, + [1; 32], + [2; 32], + BatchInstruction::Append(Vec::new()), + ); + let first = cache.ready_chain(&[1; 32]).unwrap(); + assert!(cache.is_warming()); + assert_eq!(first.len(), 1); + // A failed or timed-out send doesn't consume the prefix. + assert_eq!(cache.ready_chain(&[1; 32]).unwrap().len(), 1); + cache.confirm(&first); + cache.add_proof( + generation, + 1, + [2; 32], + [3; 32], + BatchInstruction::Append(Vec::new()), + ); + cache.finish_warming(generation); + assert_eq!(cache.ready_chain(&[2; 32]).unwrap()[0].new_root, [3; 32]); + assert!(cache.ready_chain(&[1; 32]).is_none()); + } + + #[test] + fn new_or_cancelled_warmup_preserves_completed_work() { + let mut cache = ProofCache::new(Pubkey::new_unique()); + let old = cache.start_warming([1; 32]); + add(&mut cache, 0, 1, 2); + let new = cache.start_warming([2; 32]); + // An older result remains useful even after a new collection starts. + cache.add_proof( + old, + 1, + [2; 32], + [3; 32], + BatchInstruction::Append(Vec::new()), + ); + cache.abort_warming(old); + assert!(cache.is_warming()); + cache.abort_warming(new); + assert!(!cache.is_warming()); + assert_eq!(cache.ready_chain(&[1; 32]).unwrap().len(), 2); + let empty = cache.start_warming([1; 32]); + cache.finish_warming(empty); + assert_eq!(cache.ready_chain(&[1; 32]).unwrap().len(), 2); + } + + #[test] + fn only_confirmed_prefix_is_removed_after_partial_send() { + let mut cache = ProofCache::new(Pubkey::new_unique()); + for i in 1..=6 { + add(&mut cache, i as u64, i, i + 1); + } + let chain = cache.ready_chain(&[1; 32]).unwrap(); + cache.confirm(&chain[..4]); + // The second transaction failed: its two proofs are still retryable. + let retry = cache.ready_chain(&[5; 32]).unwrap(); + assert_eq!(retry.len(), 2); + assert_eq!(retry[0].old_root, [5; 32]); + assert_eq!(retry[1].new_root, [7; 32]); + } + + #[test] + fn confirmed_on_chain_proof_is_not_replayed_after_ambiguous_timeout() { + let mut cache = ProofCache::new(Pubkey::new_unique()); + add(&mut cache, 0, 1, 2); + add(&mut cache, 1, 2, 3); + let _unconfirmed_send = cache.ready_chain(&[1; 32]).unwrap(); + let retry = cache.ready_chain(&[2; 32]).unwrap(); + assert_eq!(retry.len(), 1); + assert_eq!(retry[0].old_root, [2; 32]); + } + + #[tokio::test] + async fn cancelled_cached_send_preserves_proofs_and_releases_send_lock() { + let cache = Arc::new(SharedProofCache::new(Pubkey::new_unique())); + cache + .add_late_proof(0, [1; 32], [2; 32], BatchInstruction::Append(Vec::new())) + .await; + let send_cache = cache.clone(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(async move { + let _guard = send_cache.try_lock_for_sending().unwrap(); + let _proofs = send_cache.ready_chain(&[1; 32]).await.unwrap(); + started_tx.send(()).unwrap(); + pending::<()>().await; + }); + started_rx.await.unwrap(); + assert!(cache.try_lock_for_sending().is_none()); + // Receiving more proofs never waits for the sending lock. + cache + .add_late_proof(1, [2; 32], [3; 32], BatchInstruction::Append(Vec::new())) + .await; + task.abort(); + let _ = task.await; + assert!(cache.try_lock_for_sending().is_some()); + assert_eq!(cache.ready_chain(&[1; 32]).await.unwrap().len(), 2); + } + + #[tokio::test] + async fn pending_collectors_bound_new_work_without_hiding_ready_proofs() { + let cache = Arc::new(SharedProofCache::new(Pubkey::new_unique())); + let first = cache.start_collecting(); + let second = cache.start_collecting(); + let warmup = cache.start_warming([1; 32]).await; + warmup + .add_proof(0, [1; 32], [2; 32], BatchInstruction::Append(Vec::new())) + .await; + warmup.finish().await; + assert!(cache.has_pending_proofs().await); + assert!(cache.ready_chain(&[1; 32]).await.is_some()); + drop(first); + assert!(cache.has_pending_proofs().await); + drop(second); + assert!(!cache.has_pending_proofs().await); + } + + #[test] + fn cache_deduplicates_by_roots_and_enforces_one_capacity_limit() { + let mut cache = ProofCache::new(Pubkey::new_unique()); + cache.max_proofs = 2; + add(&mut cache, 0, 1, 2); + add(&mut cache, 99, 1, 2); + assert_eq!(cache.len(), 1); + add(&mut cache, 1, 2, 3); + add(&mut cache, 2, 3, 4); + assert_eq!(cache.len(), 2); + assert!(cache.ready_chain(&[1; 32]).is_none()); + assert_eq!(cache.ready_chain(&[2; 32]).unwrap().len(), 2); + } + #[tokio::test] async fn cancelled_warmup_releases_warming_state() { let cache = Arc::new(SharedProofCache::new(Pubkey::new_unique())); @@ -488,12 +601,14 @@ mod tests { .add_late_proof(0, root_1, root_2, BatchInstruction::Append(Vec::new())) .await; - let proofs = cache.take_if_valid(&root_1).await.unwrap(); + let proofs = cache.ready_chain(&root_1).await.unwrap(); assert_eq!(proofs.len(), 2); assert_eq!(proofs[0].old_root, root_1); assert_eq!(proofs[0].new_root, root_2); assert_eq!(proofs[1].old_root, root_2); assert_eq!(proofs[1].new_root, root_3); + assert_eq!(cache.len().await, 2); + cache.confirm(&proofs).await; assert!(cache.is_empty().await); } } diff --git a/forester/src/processor/v2/proof_worker.rs b/forester/src/processor/v2/proof_worker.rs index 01a74d1c91..8c59780e5a 100644 --- a/forester/src/processor/v2/proof_worker.rs +++ b/forester/src/processor/v2/proof_worker.rs @@ -43,24 +43,25 @@ impl ProofInput { } } - fn to_json(&self, tree_id: &str, batch_index: u64) -> String { + fn to_json(&self, tree_id: &str) -> String { + // The prover hashes the entire payload for deduplication. Local sequence + // numbers change across retries/foresters and must stay out of that key. + // Keep treeId for fair queuing; jobs within a tree use the prover's FIFO + // fallback, while TxSender orders results using the local seq below. match self { ProofInput::Append(inputs) => BatchAppendInputsJson::from_inputs(inputs) .with_tree_id(tree_id.to_string()) - .with_batch_index(batch_index) .to_string(), ProofInput::Nullify(inputs) => { use light_prover_client::proof_types::batch_update::BatchUpdateProofInputsJson; BatchUpdateProofInputsJson::from_update_inputs(inputs) .with_tree_id(tree_id.to_string()) - .with_batch_index(batch_index) .to_string() } ProofInput::AddressAppend(inputs) => { use light_prover_client::proof_types::batch_address_append::BatchAddressAppendInputsJson; BatchAddressAppendInputsJson::from_inputs(inputs) .with_tree_id(tree_id.to_string()) - .with_batch_index(batch_index) .to_string() } } @@ -235,8 +236,7 @@ async fn run_proof_pipeline( async fn submit_and_poll_proof(clients: Arc, job: ProofJob) { let client = clients.get_client(&job.inputs); - // Use seq as batch_index for ordering in the prover queue - let inputs_json = job.inputs.to_json(&job.tree_id, job.seq); + let inputs_json = job.inputs.to_json(&job.tree_id); let circuit_type = job.inputs.circuit_type(); let round_trip_start = std::time::Instant::now(); @@ -376,7 +376,7 @@ async fn poll_and_send_result( ); tokio::time::sleep(Duration::from_millis(200)).await; - let inputs_json = inputs.to_json(&tree_id, seq); + let inputs_json = inputs.to_json(&tree_id); let circuit_type = inputs.circuit_type(); let Some(submit_result) = submit_with_backpressure(client, &inputs_json, circuit_type, seq, &result_tx).await @@ -542,6 +542,72 @@ mod tests { use super::*; + #[test] + fn proof_requests_keep_stable_inputs_and_exclude_local_sequence_metadata() { + let inputs = [ + ProofInput::Append(BatchAppendsCircuitInputs { + public_input_hash: 1.into(), + old_root: 2.into(), + new_root: 3.into(), + leaves_hashchain_hash: 4.into(), + start_index: 5, + old_leaves: vec![], + leaves: vec![], + merkle_proofs: vec![], + height: 32, + batch_size: 500, + }), + ProofInput::Nullify(BatchUpdateCircuitInputs { + public_input_hash: 1.into(), + old_root: 2.into(), + new_root: 3.into(), + leaves_hashchain_hash: 4.into(), + tx_hashes: vec![], + leaves: vec![], + old_leaves: vec![], + merkle_proofs: vec![], + path_indices: vec![], + height: 32, + batch_size: 500, + }), + ProofInput::AddressAppend(BatchAddressAppendInputs { + batch_size: 250, + hashchain_hash: 4u32.into(), + low_element_values: vec![], + low_element_indices: vec![], + low_element_next_indices: vec![], + low_element_next_values: vec![], + low_element_proofs: vec![], + new_element_values: vec![], + new_element_proofs: vec![], + new_root: 3u32.into(), + old_root: 2u32.into(), + public_input_hash: 1u32.into(), + start_index: 5, + tree_height: 40, + }), + ]; + for input in inputs { + let (result_tx, _rx) = mpsc::channel(1); + let mut job = ProofJob { + seq: 1, + inputs: input, + result_tx, + tree_id: "tree-a".into(), + }; + let first = job.inputs.to_json(&job.tree_id); + job.seq = 999; + let retry = job.inputs.to_json(&job.tree_id); + assert_eq!(first, retry); + let json: serde_json::Value = serde_json::from_str(&first).unwrap(); + assert!(json.get("batchIndex").is_none()); + assert_eq!(json["treeId"], "tree-a"); + assert!(json.get("oldRoot").is_some()); + assert!(json.get("newRoot").is_some()); + assert_ne!(first, job.inputs.to_json("tree-b")); + } + } + #[test] fn queue_full_detection_is_specific() { let queue_full = ProverClientError::ProverServerError( diff --git a/forester/src/processor/v2/tx_sender.rs b/forester/src/processor/v2/tx_sender.rs index 9b345a09bb..120204f9c4 100644 --- a/forester/src/processor/v2/tx_sender.rs +++ b/forester/src/processor/v2/tx_sender.rs @@ -182,6 +182,18 @@ impl OrderedProofBuffer { fn expected_seq(&self) -> u64 { self.base_seq } + + fn drain_all(&mut self) -> Vec<(u64, BufferEntry)> { + let mut entries = Vec::with_capacity(self.len); + for offset in 0..self.buffer.len() { + let index = (self.head + offset) % self.buffer.len(); + if let Some(entry) = self.buffer[index].take() { + entries.push((self.base_seq + offset as u64, entry)); + } + } + self.len = 0; + entries + } } pub struct TxSender { @@ -559,8 +571,7 @@ impl TxSender { } // Save proofs from the reorder buffer (received but waiting for in-order processing) - while let Some(entry) = self.buffer.pop_next() { - let seq = self.buffer.expected_seq() - 1; + for (seq, entry) in self.buffer.drain_all() { warmup .add_proof(seq, entry.old_root, entry.new_root, entry.instruction) .await; @@ -622,7 +633,9 @@ fn spawn_late_proof_collector( collection_timeout: Duration, retention_timeout: Duration, ) -> JoinHandle { + let collection = cache.start_collecting(); tokio::spawn(async move { + let _collection = collection; let mut saved = 0usize; let collection = async { while let Some(result) = proof_rx.recv().await { @@ -723,6 +736,42 @@ fn spawn_late_proof_collector( mod tests { use super::*; + #[test] + fn handoff_drains_out_of_order_proofs_across_sequence_gaps() { + let mut buffer = OrderedProofBuffer::new(4, 10); + let now = std::time::Instant::now(); + assert!(buffer.insert( + 10, + BatchInstruction::Append(Vec::new()), + [1; 32], + [2; 32], + now + )); + assert!(buffer.pop_next().is_some()); + // Includes a wrapped ring-buffer entry, but seq=11 hasn't arrived. + assert!(buffer.insert( + 12, + BatchInstruction::Append(Vec::new()), + [3; 32], + [4; 32], + now + )); + assert!(buffer.insert( + 14, + BatchInstruction::Append(Vec::new()), + [5; 32], + [6; 32], + now + )); + assert!(buffer.pop_next().is_none()); + let entries = buffer.drain_all(); + assert_eq!( + entries.iter().map(|(seq, _)| *seq).collect::>(), + vec![12, 14] + ); + assert_eq!(buffer.len(), 0); + } + fn proof_result(seq: u64, old_root: [u8; 32], new_root: [u8; 32]) -> ProofJobResult { ProofJobResult { seq, @@ -762,7 +811,7 @@ mod tests { assert_eq!(collector.await.unwrap(), 1); assert!(!cache.is_warming().await); - let cached = cache.take_if_valid(&base_root).await.unwrap(); + let cached = cache.ready_chain(&base_root).await.unwrap(); assert_eq!(cached.len(), 1); assert_eq!(cached[0].old_root, base_root); assert_eq!(cached[0].new_root, next_root); @@ -824,6 +873,7 @@ mod tests { tokio::time::sleep(Duration::from_millis(40)).await; assert!(!cache.is_warming().await); assert!(!proof_tx.is_closed()); + assert!(cache.has_pending_proofs().await); let next_root = [6u8; 32]; proof_tx @@ -833,8 +883,9 @@ mod tests { drop(proof_tx); assert_eq!(collector.await.unwrap(), 1); - let cached = cache.take_if_valid(&base_root).await.unwrap(); + let cached = cache.ready_chain(&base_root).await.unwrap(); assert_eq!(cached.len(), 1); assert_eq!(cached[0].new_root, next_root); + assert!(!cache.has_pending_proofs().await); } } From 4b61b93f8cb541d9510d387e3771331bbd1a1776 Mon Sep 17 00:00:00 2001 From: Sergey Timoshin Date: Wed, 14 Oct 2026 18:47:22 +0000 Subject: [PATCH 10/10] fix(forester): address cached-proof review feedback - retry recoverable cached-send failures within the light slot; only ForesterNotEligible still propagates so process_queue re-finalizes - evict cached proofs rejected by the verifier (13006) instead of retrying them every eligible slot - prune cached proofs whose old root is already in the on-chain root history, and expire unlinked proofs after the collector window - make the pending-proof scheduling gate root-aware: collections on a chain the tree has left keep retaining results but no longer block fresh proof generation - skip pre-warming only when a cached proof links to the current root - hand proofs off to the cache once the send budget is gone, including after a slot jump during the receive wait and batches the sender rejects for lack of budget Co-Authored-By: Claude Opus 5.5 --- forester/src/epoch_manager.rs | 206 +++++++++++----- forester/src/processor/v2/common.rs | 2 +- forester/src/processor/v2/processor.rs | 4 + forester/src/processor/v2/proof_cache.rs | 301 +++++++++++++++++++---- forester/src/processor/v2/tx_sender.rs | 291 ++++++++++++++++++---- 5 files changed, 650 insertions(+), 154 deletions(-) diff --git a/forester/src/epoch_manager.rs b/forester/src/epoch_manager.rs index 79db0cf60d..e302493e0a 100644 --- a/forester/src/epoch_manager.rs +++ b/forester/src/epoch_manager.rs @@ -70,6 +70,7 @@ use crate::{ tx_builder::EpochManagerTransactions, }, v2::{ + common::map_send_error, strategy::{AddressTreeStrategy, StateTreeStrategy}, BatchContext, BatchInstruction, ProcessingResult, ProverConfig, QueueProcessor, SharedProofCache, @@ -263,6 +264,19 @@ pub enum MerkleProofType { StateProof(MerkleProof), } +/// On-chain root plus the recent root history used to prune superseded proofs. +struct TreeRoots { + current: [u8; 32], + recent: Vec<[u8; 32]>, +} + +enum CachedProofSend { + /// Cached proofs were sent, or another task holds the tree's send lock. + Progress, + /// Nothing was sent. Carries the on-chain roots when they were fetched. + Idle(Option), +} + #[derive(Debug)] pub struct EpochManager { config: Arc, @@ -2385,20 +2399,38 @@ impl EpochManager { // Consume ready prefixes even while later proofs are still arriving. // Recheck throughout the slot instead of forfeiting it on a warming cache. - if self + let tree_roots = match self .try_send_cached_proofs( epoch_info, epoch_pda, tree_accounts, consecutive_eligibility_end, ) - .await? - .is_some() + .await { - tokio::time::sleep(POLL_INTERVAL_MIN).await; - estimated_slot = self.slot_tracker.estimated_current_slot(); - continue; - } + Ok(CachedProofSend::Progress) => { + tokio::time::sleep(POLL_INTERVAL_MIN).await; + estimated_slot = self.slot_tracker.estimated_current_slot(); + continue; + } + Ok(CachedProofSend::Idle(tree_roots)) => tree_roots, + Err(e) if e.is_forester_not_eligible() => return Err(e), + Err(e) => { + // Returning would forfeit the rest of the slot. Retry with a + // fresh root check; rejected proofs have already been evicted. + warn!( + event = "cached_proofs_send_retrying", + run_id = %self.run_id, + tree = %tree_pubkey, + error = ?e, + "Cached proof send failed; retrying within the light slot" + ); + tokio::time::sleep(poll_interval).await; + poll_interval = (poll_interval * 2).min(POLL_INTERVAL_MAX); + estimated_slot = self.slot_tracker.estimated_current_slot(); + continue; + } + }; let cache = self .proof_caches @@ -2406,11 +2438,22 @@ impl EpochManager { .map(|cache| cache.clone()); if let Some(cache) = cache { if cache.has_pending_proofs().await { - // Do not generate another speculative chain for the same tree. + let current_root = match tree_roots { + Some(roots) => Some(roots.current), + None => self + .fetch_tree_roots_pooled(tree_accounts) + .await + .map(|roots| roots.current), + }; + // Do not generate another speculative chain while pending work + // can still extend the on-chain root. Pending work on a chain + // the tree has left keeps collecting, but no longer blocks. // Completed prefixes remain eligible for sending on every tick. - tokio::time::sleep(PENDING_PROOF_POLL_INTERVAL).await; - estimated_slot = self.slot_tracker.estimated_current_slot(); - continue; + if current_root.is_none_or(|root| cache.awaits_proof_for(&root)) { + tokio::time::sleep(PENDING_PROOF_POLL_INTERVAL).await; + estimated_slot = self.slot_tracker.estimated_current_slot(); + continue; + } } } @@ -3718,18 +3761,8 @@ impl EpochManager { .or_insert_with(|| Arc::new(SharedProofCache::new(tree_pubkey))) .clone(); - if cache.has_pending_proofs().await { - info!( - event = "prewarm_skipped_cache_warming", - run_id = %self_clone.run_id, - tree = %tree_pubkey, - "Tree cache is already collecting proofs; skipping pre-warm" - ); - return; - } - let cache_len = cache.len().await; - if cache_len > 0 && !cache.is_warming().await { + if cache_len > 0 || cache.has_pending_proofs().await { let mut rpc = match self_clone.rpc_pool.get_connection().await { Ok(r) => r, Err(e) => { @@ -3743,15 +3776,38 @@ impl EpochManager { return; } }; - if let Ok(current_root) = - self_clone.fetch_current_root(&mut *rpc, &tree_accounts).await + let roots = match self_clone.fetch_tree_roots(&mut *rpc, &tree_accounts).await { + Ok(roots) => roots, + Err(e) => { + warn!( + event = "prewarm_cache_validation_root_fetch_failed", + run_id = %self_clone.run_id, + tree = %tree_pubkey, + error = ?e, + "Failed to fetch tree root for cache validation" + ); + return; + } + }; + drop(rpc); + if cache.awaits_proof_for(&roots.current) { + info!( + event = "prewarm_skipped_cache_warming", + run_id = %self_clone.run_id, + tree = %tree_pubkey, + "Tree cache is already collecting proofs; skipping pre-warm" + ); + return; + } + // Only proofs that link to the current root make pre-warming redundant. + if cache.links_to(&roots.current, &roots.recent).await { info!( event = "prewarm_skipped_cache_already_warm", run_id = %self_clone.run_id, tree = %tree_pubkey, cached_proofs = cache_len, - root_prefix = ?¤t_root[..4], + root_prefix = ?&roots.current[..4], "Tree already has cached proofs from previous epoch; skipping pre-warm" ); return; @@ -3854,7 +3910,7 @@ impl EpochManager { epoch_pda: &ForesterEpochPda, tree_accounts: &TreeAccounts, consecutive_eligibility_end: u64, - ) -> std::result::Result, ForesterError> { + ) -> std::result::Result { let tree_pubkey = tree_accounts.merkle_tree; // Check eligibility window before attempting to send cached proofs @@ -3868,7 +3924,7 @@ impl EpochManager { eligibility_end_slot = consecutive_eligibility_end, "Skipping cached proof send because eligibility window has ended" ); - return Ok(None); + return Ok(CachedProofSend::Idle(None)); } let Some(confirmation_deadline) = scheduled_confirmation_deadline( @@ -3882,24 +3938,24 @@ impl EpochManager { eligibility_end_slot = consecutive_eligibility_end, "Skipping cached proofs because not enough eligible slots remain for confirmation" ); - return Ok(None); + return Ok(CachedProofSend::Idle(None)); }; let cache = match self.proof_caches.get(&tree_pubkey) { Some(c) => c.clone(), - None => return Ok(None), + None => return Ok(CachedProofSend::Idle(None)), }; if cache.is_empty().await { - return Ok(None); + return Ok(CachedProofSend::Idle(None)); } let Some(_send_guard) = cache.try_lock_for_sending() else { - return Ok(Some(0)); + return Ok(CachedProofSend::Progress); }; let mut rpc = self.rpc_pool.get_connection().await?; - let current_root = match self.fetch_current_root(&mut *rpc, tree_accounts).await { - Ok(root) => root, + let roots = match self.fetch_tree_roots(&mut *rpc, tree_accounts).await { + Ok(roots) => roots, Err(e) => { warn!( event = "cached_proofs_root_fetch_failed", @@ -3908,12 +3964,13 @@ impl EpochManager { error = ?e, "Failed to fetch current root for tree" ); - return Ok(None); + return Ok(CachedProofSend::Idle(None)); } }; drop(rpc); - let cached_proofs = match cache.ready_chain(¤t_root).await { + let current_root = roots.current; + let cached_proofs = match cache.ready_chain(¤t_root, &roots.recent).await { Some(proofs) => proofs, None => { debug!( @@ -3923,12 +3980,12 @@ impl EpochManager { root_prefix = ?¤t_root[..4], "No valid cached proofs for tree" ); - return Ok(None); + return Ok(CachedProofSend::Idle(Some(roots))); } }; if cached_proofs.is_empty() { - return Ok(Some(0)); + return Ok(CachedProofSend::Progress); } info!( @@ -3940,26 +3997,45 @@ impl EpochManager { "Sending cached proofs for tree" ); - let items_sent = self - .send_cached_proofs_as_transactions( - epoch_info, - epoch_pda, - tree_accounts, - &cache, - cached_proofs, - consecutive_eligibility_end, - confirmation_deadline, - ) - .await?; + self.send_cached_proofs_as_transactions( + epoch_info, + epoch_pda, + tree_accounts, + &cache, + cached_proofs, + consecutive_eligibility_end, + confirmation_deadline, + ) + .await?; - Ok(Some(items_sent)) + Ok(CachedProofSend::Progress) + } + + async fn fetch_tree_roots_pooled(&self, tree_accounts: &TreeAccounts) -> Option { + let roots = async { + let mut rpc = self.rpc_pool.get_connection().await?; + self.fetch_tree_roots(&mut *rpc, tree_accounts).await + }; + match roots.await { + Ok(roots) => Some(roots), + Err(e) => { + warn!( + event = "tree_root_fetch_failed", + run_id = %self.run_id, + tree = %tree_accounts.merkle_tree, + error = ?e, + "Failed to fetch current root for tree" + ); + None + } + } } - async fn fetch_current_root( + async fn fetch_tree_roots( &self, rpc: &mut impl Rpc, tree_accounts: &TreeAccounts, - ) -> Result<[u8; 32]> { + ) -> Result { use light_batched_merkle_tree::merkle_tree::BatchedMerkleTreeAccount; let mut account = rpc @@ -3979,8 +4055,10 @@ impl EpochManager { _ => return Err(anyhow!("Unsupported tree type for root fetch")), }; - let root = tree.root_history.last().copied().unwrap_or([0u8; 32]); - Ok(root) + Ok(TreeRoots { + current: tree.root_history.last().copied().unwrap_or([0u8; 32]), + recent: tree.root_history.iter().copied().collect(), + }) } #[allow(clippy::too_many_arguments)] @@ -4103,7 +4181,7 @@ impl EpochManager { }, ) .await - .map_err(RpcError::from) + .map_err(|error| map_send_error(tree_accounts.merkle_tree, error)) { Ok(sig) => { cache.confirm(chunk).await; @@ -4124,15 +4202,21 @@ impl EpochManager { .await; } Err(e) => { - warn!( - event = "cached_proofs_tx_send_failed", - run_id = %self.run_id, - error = ?e, - "Cached proof send failed; preserving unconfirmed chain for retry" - ); + if matches!(&e, ForesterError::V2(v2_error) if v2_error.is_constraint()) { + // Retrying a proof the verifier rejected can never succeed + // and would block fresh proofs for this root. + cache.evict(chunk).await; + } else { + warn!( + event = "cached_proofs_tx_send_failed", + run_id = %self.run_id, + error = ?e, + "Cached proof send failed; preserving unconfirmed chain for retry" + ); + } // Stop the dependent chain. In particular, preserve the // typed 6004 error so process_queue re-finalizes eligibility. - return Err(e.into()); + return Err(e); } } } diff --git a/forester/src/processor/v2/common.rs b/forester/src/processor/v2/common.rs index 80e2abc29f..24e4dac463 100644 --- a/forester/src/processor/v2/common.rs +++ b/forester/src/processor/v2/common.rs @@ -214,7 +214,7 @@ pub(crate) async fn send_transaction_batch( Ok(signature.to_string()) } -fn map_send_error(tree: Pubkey, error: SmartTransactionError) -> ForesterError { +pub(crate) fn map_send_error(tree: Pubkey, error: SmartTransactionError) -> ForesterError { if let Some(transaction_error) = error.transaction_error() { let v2_error = V2Error::from_transaction_error(tree, &transaction_error); if v2_error.is_batch_not_ready() { diff --git a/forester/src/processor/v2/processor.rs b/forester/src/processor/v2/processor.rs index 72fb286b22..bf4731da09 100644 --- a/forester/src/processor/v2/processor.rs +++ b/forester/src/processor/v2/processor.rs @@ -581,6 +581,7 @@ where let num_workers = self.context.num_proof_workers.max(1); let warmup = cache.start_warming(initial_root).await; + let collection = cache.start_collecting(initial_root); let (proof_tx, mut proof_rx) = mpsc::channel(num_workers * 2); @@ -626,12 +627,14 @@ where } } + collection.record_proof(result.old_root, result.new_root); warmup .add_proof(result.seq, result.old_root, result.new_root, instruction) .await; proofs_cached += 1; } Err(e) => { + collection.record_failure(result.old_root); warn!( event = "v2_prewarm_proof_generation_failed", tree = %self.context.merkle_tree, @@ -644,6 +647,7 @@ where } warmup.finish().await; + drop(collection); if proofs_cached < jobs_sent { warn!( diff --git a/forester/src/processor/v2/proof_cache.rs b/forester/src/processor/v2/proof_cache.rs index c288cc51a0..56ff464fdc 100644 --- a/forester/src/processor/v2/proof_cache.rs +++ b/forester/src/processor/v2/proof_cache.rs @@ -1,18 +1,25 @@ use std::{ - collections::VecDeque, + collections::{HashMap, HashSet, VecDeque}, sync::{ - atomic::{AtomicUsize, Ordering}, - Arc, + atomic::{AtomicU64, Ordering}, + Arc, Mutex as StdMutex, MutexGuard as StdMutexGuard, }, + time::{Duration, Instant}, }; use solana_sdk::pubkey::Pubkey; use tokio::sync::{Mutex, MutexGuard}; use tracing::{debug, info, warn}; -use super::tx_sender::BatchInstruction; +use super::tx_sender::{ + BatchInstruction, LATE_PROOF_COLLECTION_TIMEOUT, LATE_PROOF_RETENTION_TIMEOUT, +}; const DEFAULT_MAX_CACHED_PROOFS: usize = 256; +/// A collector can deliver a missing predecessor only until its retention +/// window ends. Proofs that still don't link to the on-chain root are orphaned. +const ORPHANED_PROOF_TTL: Duration = + LATE_PROOF_COLLECTION_TIMEOUT.saturating_add(LATE_PROOF_RETENTION_TIMEOUT); #[derive(Debug, Clone)] pub struct CachedProof { @@ -22,6 +29,7 @@ pub struct CachedProof { pub instruction: BatchInstruction, /// Number of ZKP batch instructions represented by this proof. pub items: usize, + pub cached_at: Instant, } #[derive(Debug)] @@ -126,6 +134,7 @@ impl ProofCache { new_root, instruction, items, + cached_at: Instant::now(), }); while self.proofs.len() > self.max_proofs { @@ -149,27 +158,31 @@ impl ProofCache { /// Snapshot the usable prefix without consuming it. Failed, timed-out or /// cancelled sends leave these proofs available until confirmation. - pub fn ready_chain(&mut self, current_root: &[u8; 32]) -> Option> { + pub fn ready_chain( + &mut self, + current_root: &[u8; 32], + recent_roots: &[[u8; 32]], + ) -> Option> { + self.ready_chain_at(current_root, recent_roots, Instant::now()) + } + + fn ready_chain_at( + &mut self, + current_root: &[u8; 32], + recent_roots: &[[u8; 32]], + now: Instant, + ) -> Option> { if self.proofs.is_empty() { return None; } + let skipped = self.prune_superseded(current_root, recent_roots); + // Treat both synchronously warmed and later results as candidates. Proof // completion order is not guaranteed, so build the usable chain by roots // instead of sequence number. Unmatched candidates stay cached: a missing // predecessor may still arrive, or the on-chain root may advance to them. - let before = self.proofs.len(); - self.proofs.retain(|proof| proof.new_root != *current_root); - let skipped = before - self.proofs.len(); let mut candidates = self.proofs.clone(); - - if skipped > 0 { - debug!( - "Skipped {} stale cached proofs for tree {} (on-chain already advanced)", - skipped, self.tree - ); - } - let mut expected = *current_root; let mut taken: Vec = Vec::new(); @@ -184,6 +197,19 @@ impl ProofCache { taken.push(proof); } + let orphaned: Vec = candidates + .into_iter() + .filter(|proof| now.saturating_duration_since(proof.cached_at) >= ORPHANED_PROOF_TTL) + .collect(); + if !orphaned.is_empty() { + self.remove(&orphaned); + debug!( + tree = %self.tree, + expired = orphaned.len(), + "Dropped cached proofs that never linked to the on-chain root" + ); + } + if taken.is_empty() { debug!( tree = %self.tree, @@ -209,6 +235,33 @@ impl ProofCache { Some(taken) } + /// Whether a cached proof starts at the on-chain root, after dropping + /// proofs the tree has already moved past. + pub fn links_to(&mut self, current_root: &[u8; 32], recent_roots: &[[u8; 32]]) -> bool { + self.prune_superseded(current_root, recent_roots); + self.proofs + .iter() + .any(|proof| proof.old_root == *current_root) + } + + /// A tree never returns to a root it has left. Proofs ending at the current + /// root, or starting from an earlier root in its history, can never land. + fn prune_superseded(&mut self, current_root: &[u8; 32], recent_roots: &[[u8; 32]]) -> usize { + let before = self.proofs.len(); + self.proofs.retain(|proof| { + proof.new_root != *current_root + && (proof.old_root == *current_root || !recent_roots.contains(&proof.old_root)) + }); + let skipped = before - self.proofs.len(); + if skipped > 0 { + debug!( + "Skipped {} stale cached proofs for tree {} (on-chain already advanced)", + skipped, self.tree + ); + } + skipped + } + pub fn len(&self) -> usize { self.proofs.len() } @@ -218,8 +271,22 @@ impl ProofCache { } pub fn confirm(&mut self, confirmed: &[CachedProof]) { + self.remove(confirmed); + } + + /// Drop proofs the program rejected so the tree can generate fresh ones. + pub fn evict(&mut self, rejected: &[CachedProof]) { + self.remove(rejected); + warn!( + tree = %self.tree, + evicted = rejected.len(), + "Evicted cached proofs rejected on-chain" + ); + } + + fn remove(&mut self, proofs: &[CachedProof]) { self.proofs.retain(|candidate| { - !confirmed.iter().any(|proof| { + !proofs.iter().any(|proof| { proof.old_root == candidate.old_root && proof.new_root == candidate.new_root }) }); @@ -253,10 +320,26 @@ impl ProofCache { } } +/// Roots from which a pending collection may still deliver the next proof: the +/// root it expected next at handoff plus every root a received proof reached, +/// minus roots already settled by a received or failed proof. +#[derive(Debug, Default)] +struct PendingRoots { + reached: HashSet<[u8; 32]>, + settled: HashSet<[u8; 32]>, +} + +impl PendingRoots { + fn awaits(&self, root: &[u8; 32]) -> bool { + self.reached.contains(root) && !self.settled.contains(root) + } +} + pub struct SharedProofCache { inner: Mutex, sending: Mutex<()>, - pending_collections: AtomicUsize, + collections: StdMutex>, + next_collection_id: AtomicU64, } impl std::fmt::Debug for SharedProofCache { @@ -270,7 +353,8 @@ impl SharedProofCache { Self { inner: Mutex::new(ProofCache::new(tree)), sending: Mutex::new(()), - pending_collections: AtomicUsize::new(0), + collections: StdMutex::new(HashMap::new()), + next_collection_id: AtomicU64::new(0), } } @@ -283,26 +367,70 @@ impl SharedProofCache { } } - pub async fn ready_chain(&self, current_root: &[u8; 32]) -> Option> { - self.inner.lock().await.ready_chain(current_root) + pub async fn ready_chain( + &self, + current_root: &[u8; 32], + recent_roots: &[[u8; 32]], + ) -> Option> { + self.inner + .lock() + .await + .ready_chain(current_root, recent_roots) + } + + pub async fn links_to(&self, current_root: &[u8; 32], recent_roots: &[[u8; 32]]) -> bool { + self.inner.lock().await.links_to(current_root, recent_roots) } pub async fn confirm(&self, confirmed: &[CachedProof]) { self.inner.lock().await.confirm(confirmed); } + pub async fn evict(&self, rejected: &[CachedProof]) { + self.inner.lock().await.evict(rejected); + } + /// Serialize cached sends without blocking proof collection. pub fn try_lock_for_sending(&self) -> Option> { self.sending.try_lock().ok() } - pub fn start_collecting(self: &Arc) -> ProofCollection { - self.pending_collections.fetch_add(1, Ordering::AcqRel); - ProofCollection(self.clone()) + /// Register pending proof work whose next proof starts at `next_root`. + pub fn start_collecting(self: &Arc, next_root: [u8; 32]) -> ProofCollection { + let id = self.next_collection_id.fetch_add(1, Ordering::Relaxed); + let mut roots = PendingRoots::default(); + roots.reached.insert(next_root); + self.lock_collections().insert(id, roots); + ProofCollection { + cache: self.clone(), + id, + } } pub async fn has_pending_proofs(&self) -> bool { - self.pending_collections.load(Ordering::Acquire) > 0 || self.is_warming().await + let collecting = !self.lock_collections().is_empty(); + collecting || self.is_warming().await + } + + /// Whether pending work may still deliver the proof that extends `root`. + /// Collections on a chain the tree has left keep retaining late results, + /// but no longer hold back fresh proof generation. + pub fn awaits_proof_for(&self, root: &[u8; 32]) -> bool { + self.lock_collections() + .values() + .any(|roots| roots.awaits(root)) + } + + fn lock_collections(&self) -> StdMutexGuard<'_, HashMap> { + self.collections + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + fn update_collection(&self, id: u64, update: impl FnOnce(&mut PendingRoots)) { + if let Some(roots) = self.lock_collections().get_mut(&id) { + update(roots); + } } pub async fn add_late_proof( @@ -337,11 +465,30 @@ impl SharedProofCache { /// Bounds speculative work for a tree across the late-result retention period. /// Dropping or aborting a collector always releases the scheduling gate. -pub struct ProofCollection(Arc); +pub struct ProofCollection { + cache: Arc, + id: u64, +} + +impl ProofCollection { + pub fn record_proof(&self, old_root: [u8; 32], new_root: [u8; 32]) { + self.cache.update_collection(self.id, |roots| { + roots.settled.insert(old_root); + roots.reached.insert(new_root); + }); + } + + /// A failed job will never deliver the proof starting at `old_root`. + pub fn record_failure(&self, old_root: [u8; 32]) { + self.cache.update_collection(self.id, |roots| { + roots.settled.insert(old_root); + }); + } +} impl Drop for ProofCollection { fn drop(&mut self) { - self.0.pending_collections.fetch_sub(1, Ordering::AcqRel); + self.cache.lock_collections().remove(&self.id); } } @@ -427,11 +574,11 @@ mod tests { [2; 32], BatchInstruction::Append(Vec::new()), ); - let first = cache.ready_chain(&[1; 32]).unwrap(); + let first = cache.ready_chain(&[1; 32], &[]).unwrap(); assert!(cache.is_warming()); assert_eq!(first.len(), 1); // A failed or timed-out send doesn't consume the prefix. - assert_eq!(cache.ready_chain(&[1; 32]).unwrap().len(), 1); + assert_eq!(cache.ready_chain(&[1; 32], &[]).unwrap().len(), 1); cache.confirm(&first); cache.add_proof( generation, @@ -441,8 +588,11 @@ mod tests { BatchInstruction::Append(Vec::new()), ); cache.finish_warming(generation); - assert_eq!(cache.ready_chain(&[2; 32]).unwrap()[0].new_root, [3; 32]); - assert!(cache.ready_chain(&[1; 32]).is_none()); + assert_eq!( + cache.ready_chain(&[2; 32], &[]).unwrap()[0].new_root, + [3; 32] + ); + assert!(cache.ready_chain(&[1; 32], &[]).is_none()); } #[test] @@ -463,10 +613,10 @@ mod tests { assert!(cache.is_warming()); cache.abort_warming(new); assert!(!cache.is_warming()); - assert_eq!(cache.ready_chain(&[1; 32]).unwrap().len(), 2); + assert_eq!(cache.ready_chain(&[1; 32], &[]).unwrap().len(), 2); let empty = cache.start_warming([1; 32]); cache.finish_warming(empty); - assert_eq!(cache.ready_chain(&[1; 32]).unwrap().len(), 2); + assert_eq!(cache.ready_chain(&[1; 32], &[]).unwrap().len(), 2); } #[test] @@ -475,10 +625,10 @@ mod tests { for i in 1..=6 { add(&mut cache, i as u64, i, i + 1); } - let chain = cache.ready_chain(&[1; 32]).unwrap(); + let chain = cache.ready_chain(&[1; 32], &[]).unwrap(); cache.confirm(&chain[..4]); // The second transaction failed: its two proofs are still retryable. - let retry = cache.ready_chain(&[5; 32]).unwrap(); + let retry = cache.ready_chain(&[5; 32], &[]).unwrap(); assert_eq!(retry.len(), 2); assert_eq!(retry[0].old_root, [5; 32]); assert_eq!(retry[1].new_root, [7; 32]); @@ -489,8 +639,8 @@ mod tests { let mut cache = ProofCache::new(Pubkey::new_unique()); add(&mut cache, 0, 1, 2); add(&mut cache, 1, 2, 3); - let _unconfirmed_send = cache.ready_chain(&[1; 32]).unwrap(); - let retry = cache.ready_chain(&[2; 32]).unwrap(); + let _unconfirmed_send = cache.ready_chain(&[1; 32], &[]).unwrap(); + let retry = cache.ready_chain(&[2; 32], &[]).unwrap(); assert_eq!(retry.len(), 1); assert_eq!(retry[0].old_root, [2; 32]); } @@ -505,7 +655,7 @@ mod tests { let (started_tx, started_rx) = tokio::sync::oneshot::channel(); let task = tokio::spawn(async move { let _guard = send_cache.try_lock_for_sending().unwrap(); - let _proofs = send_cache.ready_chain(&[1; 32]).await.unwrap(); + let _proofs = send_cache.ready_chain(&[1; 32], &[]).await.unwrap(); started_tx.send(()).unwrap(); pending::<()>().await; }); @@ -518,21 +668,21 @@ mod tests { task.abort(); let _ = task.await; assert!(cache.try_lock_for_sending().is_some()); - assert_eq!(cache.ready_chain(&[1; 32]).await.unwrap().len(), 2); + assert_eq!(cache.ready_chain(&[1; 32], &[]).await.unwrap().len(), 2); } #[tokio::test] async fn pending_collectors_bound_new_work_without_hiding_ready_proofs() { let cache = Arc::new(SharedProofCache::new(Pubkey::new_unique())); - let first = cache.start_collecting(); - let second = cache.start_collecting(); + let first = cache.start_collecting([1; 32]); + let second = cache.start_collecting([1; 32]); let warmup = cache.start_warming([1; 32]).await; warmup .add_proof(0, [1; 32], [2; 32], BatchInstruction::Append(Vec::new())) .await; warmup.finish().await; assert!(cache.has_pending_proofs().await); - assert!(cache.ready_chain(&[1; 32]).await.is_some()); + assert!(cache.ready_chain(&[1; 32], &[]).await.is_some()); drop(first); assert!(cache.has_pending_proofs().await); drop(second); @@ -549,8 +699,8 @@ mod tests { add(&mut cache, 1, 2, 3); add(&mut cache, 2, 3, 4); assert_eq!(cache.len(), 2); - assert!(cache.ready_chain(&[1; 32]).is_none()); - assert_eq!(cache.ready_chain(&[2; 32]).unwrap().len(), 2); + assert!(cache.ready_chain(&[1; 32], &[]).is_none()); + assert_eq!(cache.ready_chain(&[2; 32], &[]).unwrap().len(), 2); } #[tokio::test] @@ -601,7 +751,7 @@ mod tests { .add_late_proof(0, root_1, root_2, BatchInstruction::Append(Vec::new())) .await; - let proofs = cache.ready_chain(&root_1).await.unwrap(); + let proofs = cache.ready_chain(&root_1, &[]).await.unwrap(); assert_eq!(proofs.len(), 2); assert_eq!(proofs[0].old_root, root_1); assert_eq!(proofs[0].new_root, root_2); @@ -611,4 +761,65 @@ mod tests { cache.confirm(&proofs).await; assert!(cache.is_empty().await); } + + #[test] + fn proofs_from_superseded_roots_are_pruned() { + let mut cache = ProofCache::new(Pubkey::new_unique()); + add(&mut cache, 0, 1, 2); + add(&mut cache, 1, 2, 3); + add(&mut cache, 2, 3, 4); + // A proof on a chain that never reached the tree stays a candidate. + add(&mut cache, 3, 9, 10); + let history = [[1; 32], [2; 32], [3; 32], [4; 32]]; + + assert!(cache.ready_chain(&[4; 32], &history).is_none()); + assert_eq!(cache.len(), 1); + assert!(!cache.links_to(&[4; 32], &history)); + assert!(cache.links_to(&[9; 32], &[])); + } + + #[test] + fn unlinked_proofs_expire_after_retention_window() { + let mut cache = ProofCache::new(Pubkey::new_unique()); + add(&mut cache, 0, 1, 2); + add(&mut cache, 5, 6, 7); + let later = Instant::now() + ORPHANED_PROOF_TTL + Duration::from_secs(1); + + // Linked proofs stay usable regardless of age. + let chain = cache.ready_chain_at(&[1; 32], &[], later).unwrap(); + assert_eq!(chain.len(), 1); + assert_eq!(cache.len(), 1); + assert!(cache.links_to(&[1; 32], &[])); + } + + #[test] + fn rejected_proofs_are_evicted() { + let mut cache = ProofCache::new(Pubkey::new_unique()); + add(&mut cache, 0, 1, 2); + add(&mut cache, 1, 2, 3); + let chain = cache.ready_chain(&[1; 32], &[]).unwrap(); + cache.evict(&chain[..1]); + assert!(cache.ready_chain(&[1; 32], &[]).is_none()); + assert_eq!(cache.ready_chain(&[2; 32], &[]).unwrap().len(), 1); + } + + #[tokio::test] + async fn collection_only_holds_back_roots_it_can_still_extend() { + let cache = Arc::new(SharedProofCache::new(Pubkey::new_unique())); + let collection = cache.start_collecting([1; 32]); + assert!(cache.awaits_proof_for(&[1; 32])); + + collection.record_proof([1; 32], [2; 32]); + assert!(!cache.awaits_proof_for(&[1; 32])); + assert!(cache.awaits_proof_for(&[2; 32])); + // Another forester moved the tree onto a chain this collection can't extend. + assert!(!cache.awaits_proof_for(&[7; 32])); + assert!(cache.has_pending_proofs().await); + + collection.record_failure([2; 32]); + assert!(!cache.awaits_proof_for(&[2; 32])); + + drop(collection); + assert!(!cache.has_pending_proofs().await); + } } diff --git a/forester/src/processor/v2/tx_sender.rs b/forester/src/processor/v2/tx_sender.rs index 120204f9c4..03e2daf6d9 100644 --- a/forester/src/processor/v2/tx_sender.rs +++ b/forester/src/processor/v2/tx_sender.rs @@ -15,11 +15,11 @@ const V2_IXS_PER_TX_WITHOUT_LUT: usize = 4; const FLUSH_MARGIN_SLOTS: u64 = 10; /// Late proofs are an optimization. Do not let a proof job that never releases /// its sender keep the tree cache in the warming state indefinitely. -const LATE_PROOF_COLLECTION_TIMEOUT: Duration = Duration::from_secs(30); +pub(super) const LATE_PROOF_COLLECTION_TIMEOUT: Duration = Duration::from_secs(30); /// Keep receiving proof results after releasing the warming lock. This matches /// the default maximum prover wait and bounds detached collector lifetime when /// a producer retains its sender indefinitely. -const LATE_PROOF_RETENTION_TIMEOUT: Duration = Duration::from_secs(600); +pub(super) const LATE_PROOF_RETENTION_TIMEOUT: Duration = Duration::from_secs(600); use light_batched_merkle_tree::merkle_tree::{ InstructionDataBatchAppendInputs, InstructionDataBatchNullifyInputs, @@ -37,10 +37,11 @@ use crate::{ errors::ForesterError, processor::v2::{ common::send_transaction_batch, - proof_cache::{ProofCacheWarmup, SharedProofCache}, + proof_cache::{ProofCacheWarmup, ProofCollection, SharedProofCache}, proof_worker::ProofJobResult, BatchContext, }, + transaction_timing::scheduled_confirmation_deadline, }; #[derive(Debug, Clone, Default)] @@ -86,6 +87,8 @@ pub struct TxSenderResult { } type TxSenderTaskResult = std::result::Result; +/// (instruction, seq, old_root, new_root) +type PendingProof = (BatchInstruction, u64, [u8; 32], [u8; 32]); pub(crate) type TxSenderTask = JoinHandle; #[derive(Debug, Clone)] @@ -201,7 +204,9 @@ pub struct TxSender { buffer: OrderedProofBuffer, zkp_batch_size: u64, last_seen_root: [u8; 32], - pending_batch: Vec<(BatchInstruction, u64, [u8; 32], [u8; 32])>, // (instruction, seq, old_root, new_root) + /// Old root of the next proof expected in sequence. + next_expected_root: [u8; 32], + pending_batch: Vec, /// Earliest submission time in the pending batch (for end-to-end latency) pending_batch_earliest_submit: Option, proof_timings: ProofTimings, @@ -231,6 +236,7 @@ impl TxSender { buffer: OrderedProofBuffer::new(MAX_BUFFER_SIZE, initial_seq), zkp_batch_size, last_seen_root, + next_expected_root: last_seen_root, pending_batch: Vec::with_capacity(ixs_per_tx), pending_batch_earliest_submit: None, proof_timings: ProofTimings::default(), @@ -250,20 +256,25 @@ impl TxSender { slots_remaining <= FLUSH_MARGIN_SLOTS } + /// `send_transaction_batch` rejects batches without eligibility and the + /// minimum confirmation budget, so proofs must be handed off before then. #[inline] - fn is_still_eligible_at(&self, current_slot: u64) -> bool { - current_slot < self.context.eligibility_end_slot() + fn has_send_budget_at(&self, current_slot: u64) -> bool { + let slots_remaining = self + .context + .eligibility_end_slot() + .saturating_sub(current_slot); + scheduled_confirmation_deadline(slots_remaining).is_some() } async fn run(mut self, mut proof_rx: mpsc::Receiver) -> TxSenderTaskResult { - let (batch_tx, mut batch_rx) = mpsc::unbounded_channel::<( - Vec<(BatchInstruction, u64, [u8; 32], [u8; 32])>, - Option, - )>(); + let (batch_tx, mut batch_rx) = + mpsc::unbounded_channel::<(Vec, Option)>(); let sender_context = self.context.clone(); let mut sender_last_root = self.last_seen_root; let zkp_batch_size_val = self.zkp_batch_size; + let sender_cache = self.proof_cache.clone(); let sender_handle = tokio::spawn(async move { let mut sender_processed = 0usize; @@ -378,14 +389,20 @@ impl TxSender { warn!("tx error {} epoch {}", e, sender_context.epoch); if e.is_not_in_active_phase() { warn!("Active phase ended while sending tx, stopping sender loop"); - return Ok((sender_processed, total_tx_sending_duration)); + let mut unsent = batch; + while let Ok((queued, _)) = batch_rx.try_recv() { + unsent.extend(queued); + } + let returned = + return_unsent_to_cache(sender_cache.as_ref(), unsent).await; + return Ok((sender_processed, total_tx_sending_duration, returned)); } else { return Err(e); } } } } - Ok((sender_processed, total_tx_sending_duration)) + Ok((sender_processed, total_tx_sending_duration, 0)) }); loop { @@ -394,19 +411,19 @@ impl TxSender { } let current_slot = self.context.slot_tracker.estimated_current_slot(); - if !self.is_still_eligible_at(current_slot) { + if !self.has_send_budget_at(current_slot) { let proofs_saved = self.handoff_proofs_to_cache(proof_rx, None).await; info!( "Active phase ended for epoch {}, stopping tx sender before recv (saved {} proofs to cache)", self.context.epoch, proofs_saved ); drop(batch_tx); - let (items_processed, tx_sending_duration) = + let (items_processed, tx_sending_duration, returned) = sender_handle.await.map_err(ForesterError::from)??; return Ok(TxSenderResult { items_processed, proof_timings: self.proof_timings, - proofs_saved_to_cache: proofs_saved, + proofs_saved_to_cache: proofs_saved + returned, tx_sending_duration, }); } @@ -420,7 +437,10 @@ impl TxSender { // otherwise it is only flushed when another proof arrives, // which may be after the eligibility window has ended. let current_slot = self.context.slot_tracker.estimated_current_slot(); + // The slot may have jumped past the send budget during the + // wait. The check at the top of the loop hands the batch off. if !self.pending_batch.is_empty() + && self.has_send_budget_at(current_slot) && self.should_flush_due_to_time_at(current_slot) { let batch = std::mem::replace( @@ -429,7 +449,11 @@ impl TxSender { ); let earliest = self.pending_batch_earliest_submit.take(); - if batch_tx.send((batch, earliest)).is_err() { + if let Err(mpsc::error::SendError((batch, earliest))) = + batch_tx.send((batch, earliest)) + { + self.pending_batch = batch; + self.pending_batch_earliest_submit = earliest; break; } } @@ -439,19 +463,19 @@ impl TxSender { let current_slot = self.context.slot_tracker.estimated_current_slot(); - if !self.is_still_eligible_at(current_slot) { + if !self.has_send_budget_at(current_slot) { let proofs_saved = self.handoff_proofs_to_cache(proof_rx, Some(result)).await; info!( "Active phase ended for epoch {}, stopping tx sender (saved {} proofs to cache)", self.context.epoch, proofs_saved ); drop(batch_tx); - let (items_processed, tx_sending_duration) = + let (items_processed, tx_sending_duration, returned) = sender_handle.await.map_err(ForesterError::from)??; return Ok(TxSenderResult { items_processed, proof_timings: self.proof_timings, - proofs_saved_to_cache: proofs_saved, + proofs_saved_to_cache: proofs_saved + returned, tx_sending_duration, }); } @@ -496,6 +520,7 @@ impl TxSender { while let Some(entry) = self.buffer.pop_next() { let seq = self.buffer.expected_seq() - 1; + self.next_expected_root = entry.new_root; self.pending_batch .push((entry.instruction, seq, entry.old_root, entry.new_root)); self.pending_batch_earliest_submit = @@ -515,13 +540,32 @@ impl TxSender { ); let earliest = self.pending_batch_earliest_submit.take(); - if batch_tx.send((batch, earliest)).is_err() { + if let Err(mpsc::error::SendError((batch, earliest))) = + batch_tx.send((batch, earliest)) + { + self.pending_batch = batch; + self.pending_batch_earliest_submit = earliest; break; } } } } + if sender_handle.is_finished() { + // The sender only stops while batches can still be queued when the + // send budget ended (or on error). Keep the proofs it never received. + drop(batch_tx); + let (items_processed, tx_sending_duration, returned) = + sender_handle.await.map_err(ForesterError::from)??; + let proofs_saved = self.handoff_proofs_to_cache(proof_rx, None).await; + return Ok(TxSenderResult { + items_processed, + proof_timings: self.proof_timings, + proofs_saved_to_cache: proofs_saved + returned, + tx_sending_duration, + }); + } + if !self.pending_batch.is_empty() { let batch = std::mem::replace(&mut self.pending_batch, Vec::with_capacity(self.ixs_per_tx)); @@ -530,13 +574,13 @@ impl TxSender { } drop(batch_tx); - let (items_processed, tx_sending_duration) = + let (items_processed, tx_sending_duration, returned) = sender_handle.await.map_err(ForesterError::from)??; Ok(TxSenderResult { items_processed, proof_timings: self.proof_timings, - proofs_saved_to_cache: 0, + proofs_saved_to_cache: returned, tx_sending_duration, }) } @@ -562,39 +606,40 @@ impl TxSender { let mut saved = 0; - let warmup = cache.start_warming(self.last_seen_root).await; + let warmup = cache.start_warming(self.next_expected_root).await; + let collection = cache.start_collecting(self.next_expected_root); // Save proofs from pending_batch (already processed but not yet sent) for (instruction, seq, old_root, new_root) in self.pending_batch.drain(..) { + collection.record_proof(old_root, new_root); warmup.add_proof(seq, old_root, new_root, instruction).await; saved += 1; } // Save proofs from the reorder buffer (received but waiting for in-order processing) for (seq, entry) in self.buffer.drain_all() { + collection.record_proof(entry.old_root, entry.new_root); warmup .add_proof(seq, entry.old_root, entry.new_root, entry.instruction) .await; saved += 1; } - // Save the current result if provided - if let Some(result) = current_result { - if let Ok(instruction) = result.result { - warmup - .add_proof(result.seq, result.old_root, result.new_root, instruction) - .await; - saved += 1; - } - } - - // Drain remaining proofs from the channel + // Save the current result if provided, then drain remaining proofs from the channel + let mut available: Vec<_> = current_result.into_iter().collect(); while let Ok(result) = proof_rx.try_recv() { - if let Ok(instruction) = result.result { - warmup - .add_proof(result.seq, result.old_root, result.new_root, instruction) - .await; - saved += 1; + available.push(result); + } + for result in available { + match result.result { + Ok(instruction) => { + collection.record_proof(result.old_root, result.new_root); + warmup + .add_proof(result.seq, result.old_root, result.new_root, instruction) + .await; + saved += 1; + } + Err(_) => collection.record_failure(result.old_root), } } @@ -602,18 +647,19 @@ impl TxSender { info!( "Saved {} available proofs to cache and waiting for late results (root: {:?})", saved, - &self.last_seen_root[..4] + &self.next_expected_root[..4] ); } else { info!( "Waiting for late proof results to warm cache (root: {:?})", - &self.last_seen_root[..4] + &self.next_expected_root[..4] ); } // Dropping the JoinHandle detaches the collector so it can finish warming the cache. drop(spawn_late_proof_collector( warmup, + collection, cache.clone(), proof_rx, self.context.merkle_tree, @@ -625,28 +671,50 @@ impl TxSender { } } +/// Keep proofs from batches rejected because the send budget ended. +async fn return_unsent_to_cache( + cache: Option<&Arc>, + unsent: Vec, +) -> usize { + let Some(cache) = cache else { + debug!( + "No proof cache available, discarding {} unsent proofs", + unsent.len() + ); + return 0; + }; + let returned = unsent.len(); + for (instruction, seq, old_root, new_root) in unsent { + cache + .add_late_proof(seq, old_root, new_root, instruction) + .await; + } + returned +} + fn spawn_late_proof_collector( warmup: ProofCacheWarmup, + collection: ProofCollection, cache: Arc, mut proof_rx: mpsc::Receiver, tree: solana_sdk::pubkey::Pubkey, collection_timeout: Duration, retention_timeout: Duration, ) -> JoinHandle { - let collection = cache.start_collecting(); tokio::spawn(async move { - let _collection = collection; let mut saved = 0usize; - let collection = async { + let warming = async { while let Some(result) = proof_rx.recv().await { match result.result { Ok(instruction) => { + collection.record_proof(result.old_root, result.new_root); warmup .add_proof(result.seq, result.old_root, result.new_root, instruction) .await; saved += 1; } Err(error) => { + collection.record_failure(result.old_root); warn!( tree = %tree, seq = result.seq, @@ -658,7 +726,7 @@ fn spawn_late_proof_collector( } }; - let collection_timed_out = tokio::time::timeout(collection_timeout, collection) + let collection_timed_out = tokio::time::timeout(collection_timeout, warming) .await .is_err(); @@ -679,6 +747,7 @@ fn spawn_late_proof_collector( while let Some(result) = proof_rx.recv().await { match result.result { Ok(instruction) => { + collection.record_proof(result.old_root, result.new_root); cache .add_late_proof( result.seq, @@ -691,6 +760,7 @@ fn spawn_late_proof_collector( retained += 1; } Err(error) => { + collection.record_failure(result.old_root); warn!( tree = %tree, seq = result.seq, @@ -734,7 +804,21 @@ fn spawn_late_proof_collector( #[cfg(test)] mod tests { + use std::sync::atomic::AtomicU64; + + use forester_utils::{ + forester_epoch::{EpochPhases, Phase}, + rpc_pool::SolanaRpcPoolBuilder, + }; + use light_client::rpc::LightClient; + use solana_sdk::{commitment_config::CommitmentConfig, pubkey::Pubkey, signature::Keypair}; + use super::*; + use crate::{ + processor::{tx_cache::ProcessedHashCache, v2::ProverConfig}, + slot_tracker::SlotTracker, + smart_transaction::TransactionPolicy, + }; #[test] fn handoff_drains_out_of_order_proofs_across_sequence_gaps() { @@ -795,6 +879,7 @@ mod tests { let (proof_tx, proof_rx) = mpsc::channel(1); let collector = spawn_late_proof_collector( warmup, + cache.start_collecting(base_root), cache.clone(), proof_rx, tree, @@ -811,7 +896,7 @@ mod tests { assert_eq!(collector.await.unwrap(), 1); assert!(!cache.is_warming().await); - let cached = cache.ready_chain(&base_root).await.unwrap(); + let cached = cache.ready_chain(&base_root, &[]).await.unwrap(); assert_eq!(cached.len(), 1); assert_eq!(cached[0].old_root, base_root); assert_eq!(cached[0].new_root, next_root); @@ -827,6 +912,7 @@ mod tests { let (proof_tx, proof_rx) = mpsc::channel(1); let collector = spawn_late_proof_collector( warmup, + cache.start_collecting(base_root), cache.clone(), proof_rx, tree, @@ -862,6 +948,7 @@ mod tests { let (proof_tx, proof_rx) = mpsc::channel(1); let collector = spawn_late_proof_collector( warmup, + cache.start_collecting(base_root), cache.clone(), proof_rx, tree, @@ -883,9 +970,119 @@ mod tests { drop(proof_tx); assert_eq!(collector.await.unwrap(), 1); - let cached = cache.ready_chain(&base_root).await.unwrap(); + let cached = cache.ready_chain(&base_root, &[]).await.unwrap(); assert_eq!(cached.len(), 1); assert_eq!(cached[0].new_root, next_root); assert!(!cache.has_pending_proofs().await); } + + const ELIGIBILITY_END_SLOT: u64 = 1_000; + + async fn test_sender_context( + current_slot: u64, + ) -> (BatchContext, Arc) { + let rpc_pool = SolanaRpcPoolBuilder::::new() + .url("http://127.0.0.1:1".to_string()) + .commitment(CommitmentConfig::confirmed()) + .build() + .await + .unwrap(); + let slot_tracker = Arc::new(SlotTracker::new(current_slot, Duration::from_secs(60))); + let prover_url = "http://127.0.0.1:1".to_string(); + let active = Phase { + start: 0, + end: 1_000_000, + }; + let context = BatchContext { + rpc_pool: Arc::new(rpc_pool), + authority: Arc::new(Keypair::new()), + run_id: Arc::from("test"), + derivation: Pubkey::new_unique(), + epoch: 0, + merkle_tree: Pubkey::new_unique(), + output_queue: Pubkey::new_unique(), + prover_config: Arc::new(ProverConfig { + append_url: prover_url.clone(), + update_url: prover_url.clone(), + address_append_url: prover_url, + api_key: None, + polling_interval: Duration::from_secs(1), + max_wait_time: Duration::from_secs(1), + max_concurrent_jobs: 1, + network: "test".to_string(), + }), + ops_cache: Arc::new(tokio::sync::Mutex::new(ProcessedHashCache::new(60))), + epoch_phases: EpochPhases { + registration: active.clone(), + active: active.clone(), + report_work: active.clone(), + post: active, + }, + slot_tracker: slot_tracker.clone(), + input_queue_hint: None, + output_queue_hint: None, + num_proof_workers: 1, + forester_eligibility_end_slot: Arc::new(AtomicU64::new(ELIGIBILITY_END_SLOT)), + address_lookup_tables: Arc::new(Vec::new()), + transaction_policy: TransactionPolicy::default(), + max_batches_per_tree: 1, + }; + (context, slot_tracker) + } + + #[tokio::test] + async fn partial_batch_is_cached_when_slot_jumps_during_receive_wait() { + let (context, slot_tracker) = test_sender_context(ELIGIBILITY_END_SLOT - 500).await; + let cache = Arc::new(SharedProofCache::new(context.merkle_tree)); + let root = [1u8; 32]; + let (proof_tx, proof_rx) = mpsc::channel(4); + let sender = TxSender::spawn(context, proof_rx, 10, root, Some(cache.clone()), 0); + + // A single proof stays in the partial batch: far from the flush margin. + proof_tx + .send(proof_result(0, root, [2u8; 32])) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + // The slot jumps past eligibility while the sender waits for the next proof. + slot_tracker.update(ELIGIBILITY_END_SLOT); + + let result = tokio::time::timeout(Duration::from_secs(5), sender) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(result.items_processed, 0); + assert_eq!(result.proofs_saved_to_cache, 1); + let cached = cache.ready_chain(&root, &[]).await.unwrap(); + assert_eq!(cached.len(), 1); + assert_eq!(cached[0].new_root, [2u8; 32]); + drop(proof_tx); + } + + #[tokio::test] + async fn proofs_are_cached_once_confirmation_budget_is_exhausted() { + // Still eligible, but fewer slots remain than a send needs to confirm. + let (context, _slot_tracker) = test_sender_context(ELIGIBILITY_END_SLOT - 2).await; + let cache = Arc::new(SharedProofCache::new(context.merkle_tree)); + let root = [3u8; 32]; + let (proof_tx, proof_rx) = mpsc::channel(4); + proof_tx + .send(proof_result(0, root, [4u8; 32])) + .await + .unwrap(); + let sender = TxSender::spawn(context, proof_rx, 10, root, Some(cache.clone()), 0); + + let result = tokio::time::timeout(Duration::from_secs(5), sender) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(result.items_processed, 0); + assert_eq!(result.proofs_saved_to_cache, 1); + assert_eq!(cache.ready_chain(&root, &[]).await.unwrap().len(), 1); + // The next proof is still pending, so the collection holds the tree. + assert!(cache.awaits_proof_for(&[4u8; 32])); + drop(proof_tx); + } }