From 7f8d649f71a19fb279cbc7a62f14c09096afc73c Mon Sep 17 00:00:00 2001
From: robin-near <robin@near.org>
Date: Tue, 30 Jan 2024 12:16:27 -0800
Subject: [PATCH] Try retry orphan witness

---
 chain/client/src/adapter.rs                   |  6 ++++-
 chain/client/src/client_actor.rs              | 27 ++++++++++++++++---
 .../stateless_validation/chunk_validator.rs   | 13 ++++++---
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/chain/client/src/adapter.rs b/chain/client/src/adapter.rs
index 40f1d6eb102..ad1254c1290 100644
--- a/chain/client/src/adapter.rs
+++ b/chain/client/src/adapter.rs
@@ -141,6 +141,7 @@ pub enum ProcessTxResponse {
 pub struct ChunkStateWitnessMessage {
     pub witness: ChunkStateWitness,
     pub peer_id: PeerId,
+    pub attempts_remaining: usize,
 }
 
 #[derive(actix::Message, Debug)]
@@ -349,7 +350,10 @@ impl near_network::client::Client for Adapter {
     async fn chunk_state_witness(&self, witness: ChunkStateWitness, peer_id: PeerId) {
         match self
             .client_addr
-            .send(ChunkStateWitnessMessage { witness, peer_id }.with_span_context())
+            .send(
+                ChunkStateWitnessMessage { witness, peer_id, attempts_remaining: 5 }
+                    .with_span_context(),
+            )
             .await
         {
             Ok(()) => {}
diff --git a/chain/client/src/client_actor.rs b/chain/client/src/client_actor.rs
index 5617dabeea8..8946dc1957f 100644
--- a/chain/client/src/client_actor.rs
+++ b/chain/client/src/client_actor.rs
@@ -2006,11 +2006,32 @@ impl Handler<WithSpanContext<ChunkStateWitnessMessage>> for ClientActor {
     fn handle(
         &mut self,
         msg: WithSpanContext<ChunkStateWitnessMessage>,
-        _: &mut Context<Self>,
+        ctx: &mut Context<Self>,
     ) -> Self::Result {
         let (_span, msg) = handler_debug_span!(target: "client", msg);
-        if let Err(err) = self.client.process_chunk_state_witness(msg.witness, msg.peer_id) {
-            tracing::error!(target: "client", ?err, "Error processing chunk state witness");
+        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) {
+            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) => {}
         }
     }
 }
diff --git a/chain/client/src/stateless_validation/chunk_validator.rs b/chain/client/src/stateless_validation/chunk_validator.rs
index 055ed1f5575..fb8b161f8fa 100644
--- a/chain/client/src/stateless_validation/chunk_validator.rs
+++ b/chain/client/src/stateless_validation/chunk_validator.rs
@@ -574,13 +574,18 @@ impl Client {
         &mut self,
         witness: ChunkStateWitness,
         peer_id: PeerId,
-    ) -> Result<(), Error> {
+    ) -> Result<Option<ChunkStateWitness>, 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 = self.chain.get_block(prev_block_hash)?;
+        let prev_block = match self.chain.get_block(prev_block_hash) {
+            Ok(block) => block,
+            Err(_) => {
+                return Ok(Some(witness));
+            }
+        };
         let prev_chunk_header = Chain::get_prev_chunk_header(
             self.epoch_manager.as_ref(),
             &prev_block,
@@ -597,7 +602,7 @@ impl Client {
                 &self.chunk_validator.network_sender,
                 self.chunk_endorsement_tracker.chunk_endorsements.as_ref(),
             );
-            return Ok(());
+            return Ok(None);
         }
 
         // TODO(#10265): If the previous block does not exist, we should
@@ -616,6 +621,6 @@ impl Client {
                 },
             ));
         }
-        result
+        result.map(|_| None)
     }
 }