fix(forester): retain proofs after warmup timeout

Author: sergeytimoshin

forester/src/processor/v2/proof_cache.rs

@@ -26,6 +26,10 @@ pub struct ProofCache {
     tree: Pubkey,
     base_root: [u8; 32],
     proofs: VecDeque<CachedProof>,
+    /// 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<CachedProof>,
     warming_proofs: BTreeMap<u64, CachedProof>,
     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,62 +165,103 @@ impl ProofCache {
         );
     }
 
-    pub fn take_if_valid(&mut self, current_root: &[u8; 32]) -> Option<Vec<CachedProof>> {
-        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<Vec<CachedProof>> {
+        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)",
                 skipped, self.tree
             );
         }
 
-        if self.proofs.is_empty() {
-            debug!(
-                "Cache empty after skipping stale proofs for tree {} (current_root {:?})",
-                self.tree,
-                &current_root[..4]
-            );
-            return None;
-        }
-
         let mut expected = *current_root;
         let mut taken: Vec<CachedProof> = 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 = ?&current_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);
+    }
 }

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<R: Rpc> TxSender<R> {
             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<ProofJobResult>,
     tree: solana_sdk::pubkey::Pubkey,
     collection_timeout: Duration,
+    retention_timeout: Duration,
 ) -> JoinHandle<usize> {
     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);
     }
 }