fix(sandbox): prevent state patch loss during block processing (#15534)

## Background

`sandbox_patch_state` sometimes permanently loses patches — the patched
account never appears, even on retry. This was reported in the [Zulip
thread](https://near.zulipchat.com/#narrow/channel/308695-nearone.2Fprivate/topic/Sandbox.20patch_state.20issues)
and reproduced with
[near-sandbox-rs#51](near/near-sandbox-rs#51).

Three interrelated bugs:

- **Patch lost on block failure**: `pending_state_patch.take()` drains
the queue at the start of block processing. If preprocessing fails
(orphan, missing chunks, optimistic deferral), the patch is gone
forever.
- **RPC returns too early**: `patch_state_in_progress()` checks
`is_empty()`, which is true right after `take()` — before the patch is
actually committed to state.
- **Startup race**: Block 1 after genesis has no new chunks. The
runtime's `apply()` early-returns for old chunks and silently drops the
`state_patch`.

## What changed

Introduces `SandboxPatchTracker` — a type that encapsulates the pending
patch, a generation counter, and a committed-generation counter. Like
`SandboxStatePatch`, it uses the ZST pattern: real struct on sandbox
builds, zero-sized no-op on non-sandbox builds. This replaces the three
separate fields on `Chain` (`pending_state_patch`,
`sandbox_patch_generation`, `sandbox_patch_committed_gen`) with a single
`sandbox_patches: SandboxPatchTracker`.

The key design change is **taking the patch late** — at point of use
inside the per-shard loop of `apply_chunks_preprocessing`, instead of at
the top of `start_process_block_impl`. This eliminates the need for
backup/clone/restore on error paths, since any error before the take
leaves the patch in the tracker. The generation counter (`in_progress()`
checks `generation != committed_gen`, not `is_empty()`) ensures the RPC
doesn't return prematurely despite the late take.

This removes the `state_patch` parameter from `preprocess_block` and
`apply_chunks_preprocessing`, and results in zero `#[cfg(feature =
"sandbox")]` blocks in `chain.rs`.

## Related

- #15536 — alternative fix by @r-near that addresses the same three bugs
with a backup/restore approach and 22 cfg-gates. This PR takes a
different approach (take-late + tracker) to achieve the same result with
less code deviation.
- #14893 — fixed a different patch loss path (`clear()` in
`postprocess_ready_block`)

Author: shreyan-gupta

chain/chain/src/block_processing_utils.rs

@@ -36,6 +36,10 @@ pub(crate) struct BlockPreprocessInfo {
     pub(crate) apply_chunks_done_waiter: ApplyChunksDoneWaiter,
     /// Used to calculate block processing time metric.
     pub(crate) block_start_processing_time: Instant,
+    /// The sandbox patch generation observed when this block was preprocessed.
+    /// Used after `chain_update.commit()` to advance the tracker's committed
+    /// generation.  Always 0 on non-sandbox builds.
+    pub(crate) sandbox_patch_generation: u64,
 }
 
 pub(crate) struct OptimisticBlockInfo {

chain/chain/src/chain.rs

@@ -69,7 +69,7 @@ use near_primitives::optimistic_block::{
     BlockToApply, CachedShardUpdateKey, OptimisticBlock, OptimisticBlockKeySource,
 };
 use near_primitives::receipt::Receipt;
-use near_primitives::sandbox::state_patch::SandboxStatePatch;
+use near_primitives::sandbox::state_patch::{SandboxPatchTracker, SandboxStatePatch};
 use near_primitives::shard_layout::{ShardLayout, ShardUId};
 use near_primitives::sharding::{
     ChunkHash, ReceiptProof, ShardChunk, ShardChunkHeader, ShardProof, StateSyncInfo,
@@ -302,19 +302,9 @@ pub struct Chain {
     /// Prevents re-application of known-to-be-invalid blocks, so that in case of a
     /// protocol issue we can recover faster by focusing on correct blocks.
     invalid_blocks: LruCache<CryptoHash, ()>,
-    /// Support for sandbox's patch_state requests.
-    ///
-    /// Sandbox needs ability to arbitrary modify the state. Blockchains
-    /// naturally prevent state tampering, so we can't *just* modify data in
-    /// place in the database. Instead, we will include this "bonus changes" in
-    /// the next block we'll be processing, keeping them in this field in the
-    /// meantime.
-    ///
-    /// Note that without `sandbox` feature enabled, `SandboxStatePatch` is
-    /// a ZST.  All methods of the type are no-ops which behave as if the object
-    /// was empty and could not hold any records (which it cannot).  It's
-    /// impossible to have non-empty state patch on non-sandbox builds.
-    pending_state_patch: SandboxStatePatch,
+    /// Tracks sandbox state patches through the block processing pipeline.
+    /// On non-sandbox builds this is a ZST and every method is a no-op.
+    sandbox_patches: SandboxPatchTracker,
     /// A callback to initiate state snapshot.
     snapshot_callbacks: Option<SnapshotCallbacks>,
     /// Manages all tasks related to resharding.
@@ -458,7 +448,7 @@ impl Chain {
             last_time_head_updated: clock.now(),
             processed_hashes: LruCache::new(NonZeroUsize::new(PROCESSED_HASHES_POOL_SIZE).unwrap()),
             invalid_blocks: LruCache::new(NonZeroUsize::new(INVALID_CHUNKS_POOL_SIZE).unwrap()),
-            pending_state_patch: Default::default(),
+            sandbox_patches: Default::default(),
             snapshot_callbacks: None,
             resharding_manager,
             validator_signer,
@@ -640,7 +630,7 @@ impl Chain {
             memtrie_loading_spawner: memtrie_loading_spawner.into_spawner(),
             apply_chunk_results_cache: ApplyChunksResultCache::new(APPLY_CHUNK_RESULTS_CACHE_SIZE),
             last_time_head_updated: clock.now(),
-            pending_state_patch: Default::default(),
+            sandbox_patches: Default::default(),
             snapshot_callbacks,
             resharding_manager,
             validator_signer,
@@ -1650,14 +1640,12 @@ impl Chain {
 
         // 1) preprocess the block where we verify that the block is valid and ready to be processed
         //    No chain updates are applied at this step.
-        let state_patch = self.pending_state_patch.take();
         let preprocess_timer = metrics::BLOCK_PREPROCESSING_TIME.start_timer();
         let preprocess_res = self.preprocess_block(
             &block,
             &provenance,
             &mut block_processing_artifact.invalid_chunks,
             block_received_time,
-            state_patch,
         );
         let preprocess_res = match preprocess_res {
             Ok(preprocess_res) => {
@@ -1835,6 +1823,7 @@ impl Chain {
         let should_save_state_transition_data =
             self.should_produce_state_witness_for_this_or_next_epoch(block.header())?;
         let epoch_to_check = self.protocol_version_check;
+        let sandbox_patch_gen = block_preprocess_info.sandbox_patch_generation;
         let mut chain_update = self.chain_update();
         let block_hash = *block.hash();
         let new_head = chain_update.postprocess_block(
@@ -1847,6 +1836,7 @@ impl Chain {
             chain_update.check_protocol_version(&block_hash, epoch_to_check)?;
         }
         chain_update.commit()?;
+        self.sandbox_patches.mark_committed(sandbox_patch_gen);
         Ok(new_head)
     }
 
@@ -2348,7 +2338,6 @@ impl Chain {
         provenance: &Provenance,
         invalid_chunks: &mut Vec<ShardChunkHeader>,
         block_received_time: Instant,
-        state_patch: SandboxStatePatch,
     ) -> Result<PreprocessBlockResult, Error> {
         let header = block.header();
 
@@ -2535,7 +2524,6 @@ impl Chain {
             // for these states as well
             // otherwise put the block into the permanent storage, waiting for be caught up
             if is_caught_up { ApplyChunksMode::IsCaughtUp } else { ApplyChunksMode::NotCaughtUp },
-            state_patch,
             invalid_chunks,
         )?;
 
@@ -2550,6 +2538,7 @@ impl Chain {
                 provenance: provenance.clone(),
                 apply_chunks_done_waiter,
                 block_start_processing_time: block_received_time,
+                sandbox_patch_generation: self.sandbox_patches.generation_for_block(),
             },
             apply_chunks_still_applying,
         ))
@@ -2775,7 +2764,6 @@ impl Chain {
                 &prev_block,
                 &receipts_by_shard,
                 ApplyChunksMode::CatchingUp,
-                Default::default(),
                 &mut Vec::new(),
             )?;
             metrics::SCHEDULED_CATCHUP_BLOCK.set(block.header().height() as i64);
@@ -3200,7 +3188,6 @@ impl Chain {
         prev_block: &Block,
         incoming_receipts: &HashMap<ShardId, Vec<ReceiptProof>>,
         mode: ApplyChunksMode,
-        mut state_patch: SandboxStatePatch,
         invalid_chunks: &mut Vec<ShardChunkHeader>,
     ) -> Result<Vec<UpdateShardJob>, Error> {
         let protocol_version =
@@ -3248,12 +3235,21 @@ impl Chain {
             return Err(Error::BlockPendingOptimisticExecution);
         }
 
+        let block_height = block.header().height();
         for (shard_index, (block_context, cached_shard_update_key)) in
             update_shard_args.into_iter().enumerate()
         {
-            // XXX: This is a bit questionable -- sandbox state patching works
-            // only for a single shard. This so far has been enough.
-            let state_patch = state_patch.take();
+            // Take the sandbox patch from the tracker for this shard.  Only
+            // shards with a new chunk receive the patch — old-chunk shards
+            // would silently drop it in the runtime.  The CatchingUp mode
+            // never consumes sandbox patches.
+            let shard_has_new_chunk =
+                chunk_headers.get(shard_index).map_or(false, |h| h.is_new_chunk(block_height));
+            let state_patch = if mode != ApplyChunksMode::CatchingUp {
+                self.sandbox_patches.take_for_shard(shard_has_new_chunk)
+            } else {
+                SandboxStatePatch::default()
+            };
 
             let shard_id = shard_layout.get_shard_id(shard_index)?;
             let prev_chunk_header =
@@ -4030,14 +4026,12 @@ impl Chain {
 
 /// Sandbox node specific operations
 impl Chain {
-    // NB: `SandboxStatePatch` can only be created in `#[cfg(feature =
-    // "sandbox")]`, so we don't need extra cfg-gating here.
     pub fn patch_state(&mut self, patch: SandboxStatePatch) {
-        self.pending_state_patch.merge(patch);
+        self.sandbox_patches.submit(patch);
     }
 
     pub fn patch_state_in_progress(&self) -> bool {
-        !self.pending_state_patch.is_empty()
+        self.sandbox_patches.in_progress()
     }
 }
 

core/primitives/src/sandbox.rs

@@ -5,10 +5,7 @@ pub mod state_patch {
     /// Changes to the state to be applied via sandbox-only state patching
     /// feature.
     ///
-    /// As we only expose this functionality for sandbox, we make sure that the
-    /// object can be non-empty only if `sandbox` feature is enabled.  On
-    /// non-sandbox build, this struct is ZST and its methods are essentially
-    /// short-circuited by treating the type as always empty.
+    /// On non-sandbox builds this struct is a ZST whose methods are no-ops.
     #[derive(Default, Clone)]
     pub struct SandboxStatePatch {
         records: Vec<StateRecord>,
@@ -40,6 +37,69 @@ pub mod state_patch {
             self.records.into_iter()
         }
     }
+
+    /// Tracks sandbox state patches through the block processing pipeline.
+    ///
+    /// Wraps `SandboxStatePatch` with a generation counter so the RPC can
+    /// accurately report when a patch has been committed to the DB (not just
+    /// consumed from the pending queue).
+    #[derive(Default)]
+    pub struct SandboxPatchTracker {
+        pending: SandboxStatePatch,
+        /// Incremented each time `submit()` is called with a non-empty patch.
+        generation: u64,
+        /// Advanced to match `generation` after the block carrying the patch
+        /// is committed to the DB.
+        committed_gen: u64,
+    }
+
+    impl SandboxPatchTracker {
+        /// Queue a state patch for inclusion in a future block.
+        pub fn submit(&mut self, patch: SandboxStatePatch) {
+            if patch.is_empty() {
+                return;
+            }
+            self.pending.merge(patch);
+            self.generation += 1;
+        }
+
+        /// Whether there is a submitted patch that hasn't been committed yet.
+        pub fn in_progress(&self) -> bool {
+            self.generation != self.committed_gen
+        }
+
+        /// Take the pending patch for a shard, but only if the shard has a new
+        /// chunk.  The first shard with a new chunk gets the entire patch;
+        /// subsequent shards get an empty one.
+        pub fn take_for_shard(&mut self, shard_has_new_chunk: bool) -> SandboxStatePatch {
+            if shard_has_new_chunk && !self.pending.is_empty() {
+                self.pending.take()
+            } else {
+                SandboxStatePatch::default()
+            }
+        }
+
+        /// Returns the generation to stamp on `BlockPreprocessInfo`.
+        ///
+        /// If the pending queue was drained during this block's preprocessing,
+        /// returns the current generation so that `mark_committed` will advance
+        /// `committed_gen`.  Otherwise returns a sentinel that won't advance.
+        pub fn generation_for_block(&self) -> u64 {
+            if self.pending.is_empty() && self.generation > self.committed_gen {
+                self.generation
+            } else {
+                self.committed_gen
+            }
+        }
+
+        /// Called after `chain_update.commit()`.  Advances `committed_gen` if
+        /// the block carried a patch whose generation is newer.
+        pub fn mark_committed(&mut self, generation: u64) {
+            if generation > self.committed_gen {
+                self.committed_gen = generation;
+            }
+        }
+    }
 }
 
 #[cfg(not(feature = "sandbox"))]
@@ -59,7 +119,7 @@ pub mod state_patch {
             Self
         }
         #[inline(always)]
-        pub fn merge(&self, _other: SandboxStatePatch) {}
+        pub fn merge(&mut self, _other: SandboxStatePatch) {}
     }
 
     impl IntoIterator for SandboxStatePatch {
@@ -71,4 +131,26 @@ pub mod state_patch {
             std::iter::empty()
         }
     }
+
+    #[derive(Default)]
+    pub struct SandboxPatchTracker;
+
+    impl SandboxPatchTracker {
+        #[inline(always)]
+        pub fn submit(&mut self, _patch: SandboxStatePatch) {}
+        #[inline(always)]
+        pub fn in_progress(&self) -> bool {
+            false
+        }
+        #[inline(always)]
+        pub fn take_for_shard(&mut self, _shard_has_new_chunk: bool) -> SandboxStatePatch {
+            SandboxStatePatch
+        }
+        #[inline(always)]
+        pub fn generation_for_block(&self) -> u64 {
+            0
+        }
+        #[inline(always)]
+        pub fn mark_committed(&mut self, _generation: u64) {}
+    }
 }