diff --git a/forester/src/epoch_manager.rs b/forester/src/epoch_manager.rs index 9b342d2f30..2d933ee9ba 100644 --- a/forester/src/epoch_manager.rs +++ b/forester/src/epoch_manager.rs @@ -2261,43 +2261,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 @@ -2307,6 +2270,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 { @@ -2346,6 +2310,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 @@ -3435,6 +3430,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!( @@ -3502,6 +3498,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!( @@ -3618,12 +3615,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; @@ -3644,7 +3645,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, @@ -3701,7 +3702,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( @@ -3775,9 +3778,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 @@ -3813,15 +3817,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 { @@ -3838,7 +3839,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!( @@ -3868,8 +3870,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?; @@ -3905,19 +3910,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; @@ -3967,7 +4003,7 @@ impl EpochManager { } } } - chunk_items += proof.items; + chunk_items += proof.items * zkp_batch_size; } if !instructions.is_empty() { @@ -3997,22 +4033,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()); } } } @@ -4737,6 +4784,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 4441cab3c5..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) { @@ -340,6 +339,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!( @@ -572,7 +580,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 +626,7 @@ where } } - cache + warmup .add_proof(result.seq, result.old_root, result.new_root, instruction) .await; proofs_cached += 1; @@ -635,7 +643,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..c288cc51a0 100644 --- a/forester/src/processor/v2/proof_cache.rs +++ b/forester/src/processor/v2/proof_cache.rs @@ -1,7 +1,13 @@ -use std::collections::{BTreeMap, VecDeque}; +use std::{ + 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; @@ -23,8 +29,8 @@ pub struct ProofCache { tree: Pubkey, base_root: [u8; 32], proofs: VecDeque, - warming_proofs: BTreeMap, is_warming: bool, + warming_generation: u64, max_proofs: usize, } @@ -34,127 +40,128 @@ impl ProofCache { tree, base_root: [0u8; 32], 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(); + // A new collection must not erase completed work from an earlier one. 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"); - return; + if self.warming_generation != generation { + debug!( + tree = %self.tree, + generation, + active_generation = self.warming_generation, + "Retaining proof from an older cache collection" + ); } - if self.warming_proofs.contains_key(&seq) { - warn!( - "Duplicate cached proof seq={} for tree {}, ignoring", - seq, self.tree + self.add_late_proof(seq, old_root, new_root, instruction); + } + + 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; } - let items = instruction.items_count(); - self.warming_proofs.insert( - seq, - CachedProof { - seq, - old_root, - new_root, - instruction, - items, - }, - ); + self.is_warming = false; - 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 - ); - } - debug!( - "Cached proof seq={} for tree {} (total cached: {})", - seq, + info!( + "Cache warm-up complete for tree {}: {} proofs cached with root {:?}", self.tree, - self.warming_proofs.len() + self.proofs.len(), + &self.base_root[..4] ); } - pub fn finish_warming(&mut self) { - 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] + pub fn add_late_proof( + &mut self, + seq: u64, + old_root: [u8; 32], + new_root: [u8; 32], + instruction: BatchInstruction, + ) { + let duplicate = self + .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; } - if let Some(first) = self.warming_proofs.values().next() { - if self.base_root != [0u8; 32] && self.base_root != first.old_root { + let items = instruction.items_count(); + self.proofs.push_back(CachedProof { + seq, + old_root, + new_root, + instruction, + items, + }); + + while self.proofs.len() > self.max_proofs { + if let Some(dropped) = self.proofs.pop_front() { 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 + tree = %self.tree, + seq = dropped.seq, + max = self.max_proofs, + "Late proof cache limit reached; dropping oldest candidate" ); } } - 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, - self.proofs.len(), - &self.base_root[..4] + debug!( + tree = %self.tree, + seq, + 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.proofs.is_empty() || self.is_warming { + /// 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; } - 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; - } - self.proofs.pop_front(); - skipped += 1; - } + // 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!( @@ -163,36 +170,27 @@ 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); } if taken.is_empty() { + debug!( + tree = %self.tree, + current_root = ?¤t_root[..4], + retained_candidates = self.proofs.len(), + "No cached proof currently links to the on-chain root" + ); return None; } @@ -219,6 +217,14 @@ impl ProofCache { 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 { self.is_warming } @@ -229,13 +235,28 @@ impl ProofCache { pub fn clear(&mut self) { self.proofs.clear(); - self.warming_proofs.clear(); self.is_warming = false; } + + fn abort_warming(&mut self, generation: u64) { + if !self.is_warming || self.warming_generation != generation { + return; + } + + // Cancellation releases scheduling state, never already completed proofs. + self.is_warming = false; + warn!( + tree = %self.tree, + generation, + "Cache warm-up was cancelled; releasing warming state" + ); + } } pub struct SharedProofCache { inner: Mutex, + sending: Mutex<()>, + pending_collections: AtomicUsize, } impl std::fmt::Debug for SharedProofCache { @@ -248,14 +269,43 @@ impl SharedProofCache { pub fn new(tree: Pubkey) -> Self { Self { inner: Mutex::new(ProofCache::new(tree)), + sending: Mutex::new(()), + pending_collections: AtomicUsize::new(0), } } - pub async fn start_warming(&self, base_root: [u8; 32]) { - self.inner.lock().await.start_warming(base_root); + 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 add_proof( + 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( &self, seq: u64, old_root: [u8; 32], @@ -265,15 +315,7 @@ impl SharedProofCache { 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 take_if_valid(&self, current_root: &[u8; 32]) -> Option> { - self.inner.lock().await.take_if_valid(current_root) + .add_late_proof(seq, old_root, new_root, instruction); } pub async fn is_warming(&self) -> bool { @@ -292,3 +334,281 @@ impl SharedProofCache { self.inner.lock().await.clear(); } } + +/// 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 +/// 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::*; + + 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())); + 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); + } + + #[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.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 98a81f24fb..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() } } @@ -194,8 +195,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 { @@ -209,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(); @@ -350,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 @@ -516,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 7c15c29397..120204f9c4 100644 --- a/forester/src/processor/v2/tx_sender.rs +++ b/forester/src/processor/v2/tx_sender.rs @@ -5,7 +5,21 @@ 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); +/// 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, @@ -22,8 +36,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, }, }; @@ -166,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 { @@ -386,7 +414,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(); @@ -514,18 +562,17 @@ 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 + for (seq, entry) in self.buffer.drain_all() { + warmup .add_proof(seq, entry.old_root, entry.new_root, entry.instruction) .await; saved += 1; @@ -534,7 +581,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; @@ -544,7 +591,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; @@ -566,9 +613,12 @@ 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, + LATE_PROOF_COLLECTION_TIMEOUT, + LATE_PROOF_RETENTION_TIMEOUT, )); saved @@ -576,32 +626,101 @@ impl TxSender { } fn spawn_late_proof_collector( + warmup: ProofCacheWarmup, 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; - 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; + let collection = async { + while let Some(result) = proof_rx.recv().await { + match result.result { + Ok(instruction) => { + warmup + .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" + ); + } } - Err(error) => { - warn!( - tree = %tree, - seq = result.seq, - error = %error, - "Late proof failed while warming cache" - ); + } + }; + + let collection_timed_out = tokio::time::timeout(collection_timeout, collection) + .await + .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; 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" + ); } } - cache.finish_warming().await; let total_cached_proofs = cache.len().await; info!( tree = %tree, @@ -617,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, @@ -635,10 +790,17 @@ 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(cache.clone(), proof_rx, tree); + let collector = spawn_late_proof_collector( + warmup, + cache.clone(), + proof_rx, + tree, + LATE_PROOF_COLLECTION_TIMEOUT, + LATE_PROOF_RETENTION_TIMEOUT, + ); assert!(cache.is_warming().await); proof_tx @@ -649,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); @@ -660,10 +822,17 @@ 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(cache.clone(), proof_rx, tree); + let collector = spawn_late_proof_collector( + warmup, + cache.clone(), + proof_rx, + tree, + LATE_PROOF_COLLECTION_TIMEOUT, + LATE_PROOF_RETENTION_TIMEOUT, + ); proof_tx .send(ProofJobResult { seq: 0, @@ -682,4 +851,41 @@ 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)); + 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, + Duration::from_millis(20), + Duration::from_secs(1), + ); + + assert!(cache.is_warming().await); + 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 + .send(proof_result(0, base_root, next_root)) + .await + .unwrap(); + drop(proof_tx); + + assert_eq!(collector.await.unwrap(), 1); + 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); + } }