feat: orphan chunk state witness pool

chain/chain/src/test_utils/kv_runtime.rs

@@ -13,6 +13,7 @@ use near_epoch_manager::types::BlockHeaderInfo;
 use near_epoch_manager::{EpochManagerAdapter, RngSeed};
 use near_pool::types::TransactionGroupIterator;
 use near_primitives::account::{AccessKey, Account};
+use near_primitives::block::Tip;
 use near_primitives::block_header::{Approval, ApprovalInner};
 use near_primitives::epoch_manager::block_info::BlockInfo;
 use near_primitives::epoch_manager::epoch_info::EpochInfo;
@@ -942,6 +943,14 @@ impl EpochManagerAdapter for MockEpochManager {
         Ok(true)
     }
 
+    fn verify_chunk_state_witness_signature_in_epoch(
+        &self,
+        _state_witness: &ChunkStateWitness,
+        _epoch_id: &EpochId,
+    ) -> Result<bool, Error> {
+        Ok(true)
+    }
+
     fn cares_about_shard_from_prev_block(
         &self,
         parent_hash: &CryptoHash,
@@ -990,6 +999,14 @@ impl EpochManagerAdapter for MockEpochManager {
         Ok(shard_layout != next_shard_layout)
     }
 
+    fn possible_epochs_of_height_around_tip(
+        &self,
+        _tip: &Tip,
+        _height: BlockHeight,
+    ) -> Result<Vec<EpochId>, EpochError> {
+        unimplemented!();
+    }
+
     #[cfg(feature = "new_epoch_sync")]
     fn get_all_epoch_hashes(
         &self,

chain/client/Cargo.toml

@@ -41,6 +41,7 @@ thiserror.workspace = true
 tokio.workspace = true
 tracing.workspace = true
 yansi.workspace = true
+bytesize.workspace = true
 
 near-async.workspace = true
 near-cache.workspace = true

chain/client/src/client.rs

@@ -353,6 +353,7 @@ impl Client {
             network_adapter.clone().into_sender(),
             runtime_adapter.clone(),
             chunk_endorsement_tracker.clone(),
+            config.orphan_state_witness_pool_size,
         );
         let chunk_distribution_network = ChunkDistributionNetwork::from_config(&config);
         Ok(Self {
@@ -1643,6 +1644,8 @@ impl Client {
 
         self.shards_manager_adapter
             .send(ShardsManagerRequestFromClient::CheckIncompleteChunks(*block.hash()));
+
+        self.process_ready_orphan_chunk_state_witnesses(&block);
     }
 
     /// Reconcile the transaction pool after processing a block.

chain/client/src/client_actor.rs

@@ -2006,32 +2006,11 @@ impl Handler<WithSpanContext<ChunkStateWitnessMessage>> for ClientActor {
     fn handle(
         &mut self,
         msg: WithSpanContext<ChunkStateWitnessMessage>,
-        ctx: &mut Context<Self>,
+        _ctx: &mut Context<Self>,
     ) -> Self::Result {
         let (_span, msg) = handler_debug_span!(target: "client", msg);
-        let peer_id = msg.peer_id.clone();
-        let attempts_remaining = msg.attempts_remaining;
-        match self.client.process_chunk_state_witness(msg.witness, msg.peer_id, None) {
-            Err(err) => {
-                tracing::error!(target: "client", ?err, "Error processing chunk state witness");
-            }
-            Ok(Some(witness)) => {
-                if attempts_remaining > 0 {
-                    ctx.run_later(Duration::from_millis(100), move |_, ctx| {
-                        ctx.address().do_send(
-                            ChunkStateWitnessMessage {
-                                witness,
-                                peer_id,
-                                attempts_remaining: attempts_remaining - 1,
-                            }
-                            .with_span_context(),
-                        );
-                    });
-                } else {
-                    tracing::error!(target: "client", "Failed to process chunk state witness even after 5 tries due to missing parent block");
-                }
-            }
-            Ok(None) => {}
+        if let Err(err) = self.client.process_chunk_state_witness(msg.witness, msg.peer_id, None) {
+            tracing::error!(target: "client", ?err, "Error processing chunk state witness");
         }
     }
 }

chain/client/src/lib.rs

@@ -13,12 +13,16 @@ pub use crate::client::{Client, ProduceChunkResult};
 pub use crate::client_actor::NetworkAdversarialMessage;
 pub use crate::client_actor::{start_client, ClientActor};
 pub use crate::config_updater::ConfigUpdater;
+pub use crate::stateless_validation::chunk_validator::orphan_witness_handling::{
+    HandleOrphanWitnessOutcome, MAX_ORPHAN_WITNESS_SIZE,
+};
 pub use crate::sync::adapter::{SyncAdapter, SyncMessage};
 pub use crate::view_client::{start_view_client, ViewClientActor};
 pub use near_client_primitives::debug::DebugStatus;
 pub use near_network::client::{
     BlockApproval, BlockResponse, ProcessTxRequest, ProcessTxResponse, SetNetworkInfo,
 };
+pub use stateless_validation::processing_tracker::{ProcessingDoneTracker, ProcessingDoneWaiter};
 
 pub mod adapter;
 pub mod adversarial;

chain/client/src/metrics.rs

@@ -583,6 +583,34 @@ pub(crate) static CHUNK_STATE_WITNESS_TOTAL_SIZE: Lazy<HistogramVec> = Lazy::new
     .unwrap()
 });
 
+pub(crate) static ORPHAN_CHUNK_STATE_WITNESSES_TOTAL_COUNT: Lazy<IntCounterVec> = Lazy::new(|| {
+    try_create_int_counter_vec(
+        "near_orphan_chunk_state_witness_total_count",
+        "Total number of orphaned chunk state witnesses that were saved for later processing",
+        &["shard_id"],
+    )
+    .unwrap()
+});
+
+pub(crate) static ORPHAN_CHUNK_STATE_WITNESS_POOL_SIZE: Lazy<IntGaugeVec> = Lazy::new(|| {
+    try_create_int_gauge_vec(
+        "near_orphan_chunk_state_witness_pool_size",
+        "Number of orphaned witnesses kept in OrphanStateWitnessPool (by shard_id)",
+        &["shard_id"],
+    )
+    .unwrap()
+});
+
+pub(crate) static ORPHAN_CHUNK_STATE_WITNESS_POOL_MEMORY_USED: Lazy<IntGaugeVec> =
+    Lazy::new(|| {
+        try_create_int_gauge_vec(
+            "near_orphan_chunk_state_witness_pool_memory_used",
+            "Memory in bytes consumed by the OrphanStateWitnessPool (by shard_id)",
+            &["shard_id"],
+        )
+        .unwrap()
+    });
+
 pub(crate) static BLOCK_PRODUCER_ENDORSED_STAKE_RATIO: Lazy<HistogramVec> = Lazy::new(|| {
     try_create_histogram_vec(
         "near_block_producer_endorsed_stake_ratio",

chain/client/src/stateless_validation/chunk_validator.rs → chain/client/src/stateless_validation/chunk_validator/mod.rs

@@ -1,3 +1,6 @@
+pub mod orphan_witness_handling;
+pub mod orphan_witness_pool;
+
 use super::processing_tracker::ProcessingDoneTracker;
 use crate::stateless_validation::chunk_endorsement_tracker::ChunkEndorsementTracker;
 use crate::{metrics, Client};
@@ -31,6 +34,7 @@ use near_primitives::types::chunk_extra::ChunkExtra;
 use near_primitives::types::ShardId;
 use near_primitives::validator_signer::ValidatorSigner;
 use near_store::PartialStorage;
+use orphan_witness_pool::OrphanStateWitnessPool;
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -53,6 +57,7 @@ pub struct ChunkValidator {
     network_sender: Sender<PeerManagerMessageRequest>,
     runtime_adapter: Arc<dyn RuntimeAdapter>,
     chunk_endorsement_tracker: Arc<ChunkEndorsementTracker>,
+    orphan_witness_pool: OrphanStateWitnessPool,
 }
 
 impl ChunkValidator {
@@ -62,13 +67,15 @@ impl ChunkValidator {
         network_sender: Sender<PeerManagerMessageRequest>,
         runtime_adapter: Arc<dyn RuntimeAdapter>,
         chunk_endorsement_tracker: Arc<ChunkEndorsementTracker>,
+        orphan_witness_pool_size: usize,
     ) -> Self {
         Self {
             my_signer,
             epoch_manager,
             network_sender,
             runtime_adapter,
             chunk_endorsement_tracker,
+            orphan_witness_pool: OrphanStateWitnessPool::new(orphan_witness_pool_size),
         }
     }
 
@@ -588,18 +595,44 @@ impl Client {
         witness: ChunkStateWitness,
         peer_id: PeerId,
         processing_done_tracker: Option<ProcessingDoneTracker>,
-    ) -> Result<Option<ChunkStateWitness>, Error> {
+    ) -> Result<(), Error> {
         // TODO(#10502): Handle production of state witness for first chunk after genesis.
         // Properly handle case for chunk right after genesis.
         // Context: We are currently unable to handle production of the state witness for the
         // first chunk after genesis as it's not possible to run the genesis chunk in runtime.
         let prev_block_hash = witness.inner.chunk_header.prev_block_hash();
         let prev_block = match self.chain.get_block(prev_block_hash) {
             Ok(block) => block,
-            Err(_) => {
-                return Ok(Some(witness));
+            Err(Error::DBNotFoundErr(_)) => {
+                // Previous block isn't available at the moment, add this witness to the orphan pool.
+                self.handle_orphan_state_witness(witness)?;
+                return Ok(());
             }
+            Err(err) => return Err(err),
         };
+        self.process_chunk_state_witness_with_prev_block(
+            witness,
+            peer_id,
+            &prev_block,
+            processing_done_tracker,
+        )
+    }
+
+    pub fn process_chunk_state_witness_with_prev_block(
+        &mut self,
+        witness: ChunkStateWitness,
+        peer_id: PeerId,
+        prev_block: &Block,
+        processing_done_tracker: Option<ProcessingDoneTracker>,
+    ) -> Result<(), Error> {
+        if witness.inner.chunk_header.prev_block_hash() != prev_block.hash() {
+            return Err(Error::Other(format!(
+                "process_chunk_state_witness_with_prev_block - prev_block doesn't match ({} != {})",
+                witness.inner.chunk_header.prev_block_hash(),
+                prev_block.hash()
+            )));
+        }
+
         let prev_chunk_header = Chain::get_prev_chunk_header(
             self.epoch_manager.as_ref(),
             &prev_block,
@@ -616,11 +649,9 @@ impl Client {
                 &self.chunk_validator.network_sender,
                 self.chunk_endorsement_tracker.as_ref(),
             );
-            return Ok(None);
+            return Ok(());
         }
 
-        // TODO(#10265): If the previous block does not exist, we should
-        // queue this (similar to orphans) to retry later.
         let result = self.chunk_validator.start_validating_chunk(
             witness,
             &self.chain,
@@ -635,6 +666,6 @@ impl Client {
                 },
             ));
         }
-        result.map(|_| None)
+        result
     }
 }