fix(state-sync): set the sync hash only when it's final (#12720)

There's currently a bug that can happen if the sync hash that we set
doesn't end up on the canonical chain. Here we add a couple more test
cases for that, and fix it by considering only finalized blocks when we
set the state sync hash, and changing the snapshot logic to check
whether the current head is the prev block of what will be the sync hash
block. This means the sync hash is found a little later than it was
before, but on a chain with long epoch lengths, it shouldn't be a
problem.

Author: marcelo-gonzalez

chain/chain/src/chain.rs

@@ -3952,13 +3952,15 @@ impl Chain {
                         Ok(SnapshotAction::None)
                     }
                 } else {
-                    let Some(sync_hash) = self.get_sync_hash(&head.last_block_hash)? else {
-                        return Ok(SnapshotAction::None);
-                    };
-                    if sync_hash == head.last_block_hash {
-                        // note that here we're returning prev_block_hash instead of last_block_hash because in this case
-                        // we can't detect the right sync hash until it is actually applied as the head block
-                        Ok(SnapshotAction::MakeSnapshot(head.prev_block_hash))
+                    let is_sync_prev = crate::state_sync::is_sync_prev_hash(
+                        &self.chain_store.store(),
+                        &head.last_block_hash,
+                        &head.prev_block_hash,
+                    )?;
+                    if is_sync_prev {
+                        // Here the head block is the prev block of what the sync hash will be, and the previous
+                        // block is the point in the chain we want to snapshot state for
+                        Ok(SnapshotAction::MakeSnapshot(head.last_block_hash))
                     } else {
                         Ok(SnapshotAction::None)
                     }
@@ -4074,6 +4076,12 @@ fn shard_id_out_of_bounds(shard_id: ShardId) -> Error {
 /// ApplyChunksMode::NotCaughtUp once with ApplyChunksMode::CatchingUp. Note
 /// that it does not guard whether the children shards are ready or not, see the
 /// comments before `need_to_reshard`
+// TODO(state-sync): After the changes in https://github.com/near/nearcore/pull/12617,
+// this needs to be changed to be aware of what shards can be applied now. Otherwise we have
+// a bug in the rare case where we have something like this sequence of tracked shards in consecutive epochs:
+// (s0) -> (s1) -> (s0, s2)
+// In this case we don't state sync s0 since we already have the state, but we apply chunks with mode `NotCaughtUp`
+// in the middle epoch there because we're downloading state for s2.
 fn get_should_apply_chunk(
     mode: ApplyChunksMode,
     cares_about_shard_this_epoch: bool,

chain/chain/src/state_sync.rs

@@ -15,14 +15,6 @@ fn get_state_sync_new_chunks(
     Ok(store.get_ser(DBCol::StateSyncNewChunks, block_hash.as_ref())?)
 }
 
-fn iter_state_sync_new_chunks_keys<'a>(
-    store: &'a Store,
-) -> impl Iterator<Item = Result<CryptoHash, std::io::Error>> + 'a {
-    store
-        .iter(DBCol::StateSyncNewChunks)
-        .map(|item| item.and_then(|(k, _v)| CryptoHash::try_from_slice(&k)))
-}
-
 fn iter_state_sync_hashes_keys<'a>(
     store: &'a Store,
 ) -> impl Iterator<Item = Result<EpochId, std::io::Error>> + 'a {
@@ -31,17 +23,18 @@ fn iter_state_sync_hashes_keys<'a>(
         .map(|item| item.and_then(|(k, _v)| EpochId::try_from_slice(&k)))
 }
 
+/// Saves new chunk info and returns whether there are at least 2 chunks per shard in the epoch for header.prev_hash()
 fn save_epoch_new_chunks<T: ChainStoreAccess>(
     chain_store: &T,
     store_update: &mut StoreUpdate,
     header: &BlockHeader,
-) -> Result<(), Error> {
+) -> Result<bool, Error> {
     let Some(mut num_new_chunks) =
         get_state_sync_new_chunks(&chain_store.store(), header.prev_hash())?
     else {
         // This might happen in the case of epoch sync where we save individual headers without having all
         // headers that belong to the epoch.
-        return Ok(());
+        return Ok(false);
     };
 
     // This shouldn't happen because block headers in the same epoch should have chunks masks
@@ -53,17 +46,11 @@ fn save_epoch_new_chunks<T: ChainStoreAccess>(
             block_hash=%header.hash(), chunk_mask_len=%header.chunk_mask().len(), stored_len=%num_new_chunks.len(),
             "block header's chunk mask not of the same length as stored value in DBCol::StateSyncNewChunks",
         );
-        return Ok(());
+        return Ok(false);
     }
 
     let done = num_new_chunks.iter().all(|num_chunks| *num_chunks >= 2);
-    if done {
-        // TODO(current_epoch_state_sync): this will not be correct if this block doesn't end up finalized on the main chain.
-        // We should fix it by setting the sync hash when it's finalized, which requires making changes to how we take state snapshots.
-        store_update.set_ser(DBCol::StateSyncHashes, header.epoch_id().as_ref(), header.hash())?;
-        store_update.delete_all(DBCol::StateSyncNewChunks);
-        return Ok(());
-    }
+
     for (num_new_chunks, new_chunk) in num_new_chunks.iter_mut().zip(header.chunk_mask().iter()) {
         // Only need to reach 2, so don't bother adding more than that
         if *new_chunk && *num_new_chunks < 2 {
@@ -72,7 +59,7 @@ fn save_epoch_new_chunks<T: ChainStoreAccess>(
     }
 
     store_update.set_ser(DBCol::StateSyncNewChunks, header.hash().as_ref(), &num_new_chunks)?;
-    Ok(())
+    Ok(done)
 }
 
 fn on_new_epoch(store_update: &mut StoreUpdate, header: &BlockHeader) -> Result<(), Error> {
@@ -96,31 +83,89 @@ fn remove_old_epochs(
     Ok(())
 }
 
-fn remove_old_blocks<T: ChainStoreAccess>(
+/// Helper to turn DBNotFoundErr() into None. We might get DBNotFoundErr() in the case of epoch sync
+/// where we save individual headers without having all headers that belong to the epoch.
+fn maybe_get_block_header<T: ChainStoreAccess>(
+    chain_store: &T,
+    block_hash: &CryptoHash,
+) -> Result<Option<BlockHeader>, Error> {
+    match chain_store.get_block_header(block_hash) {
+        Ok(block_header) => Ok(Some(block_header)),
+        // This might happen in the case of epoch sync where we save individual headers without having all
+        // headers that belong to the epoch.
+        Err(Error::DBNotFoundErr(_)) => Ok(None),
+        Err(e) => Err(e),
+    }
+}
+
+fn has_enough_new_chunks(store: &Store, block_hash: &CryptoHash) -> Result<Option<bool>, Error> {
+    let Some(num_new_chunks) = get_state_sync_new_chunks(store, block_hash)? else {
+        // This might happen in the case of epoch sync where we save individual headers without having all
+        // headers that belong to the epoch.
+        return Ok(None);
+    };
+    Ok(Some(num_new_chunks.iter().all(|num_chunks| *num_chunks >= 2)))
+}
+
+/// Save num new chunks info and store the state sync hash if it has been found. We store it only
+/// once it becomes final.
+/// This should only be called if DBCol::StateSyncHashes does not yet have an entry for header.epoch_id().
+/// The logic should still be correct if it is, but it's unnecessary and will waste a lot of time if called
+/// on a header far away from the epoch start.
+fn on_new_header<T: ChainStoreAccess>(
     chain_store: &T,
     store_update: &mut StoreUpdate,
     header: &BlockHeader,
 ) -> Result<(), Error> {
-    if header.last_final_block() == &CryptoHash::default() {
+    let done = save_epoch_new_chunks(chain_store, store_update, header)?;
+    if !done {
         return Ok(());
     }
-    // We don't need to keep info for old blocks around. After a block is finalized, we don't need anything before it
-    let last_final_header = match chain_store.get_block_header(header.last_final_block()) {
-        Ok(h) => h,
-        // This might happen in the case of epoch sync where we save individual headers without having all
-        // headers that belong to the epoch.
-        Err(Error::DBNotFoundErr(_)) => return Ok(()),
-        Err(e) => return Err(e),
+
+    // Now check if the sync hash is known and finalized. The sync hash is the block after the first block with at least 2
+    // chunks per shard in the epoch. Note that we cannot just check if the current header.last_final_block() is the sync
+    // hash, because even though this function is called for each header, it is not guaranteed that we'll see every block
+    // by checking header.last_final_block(), because it is possible for the final block to jump by more than one upon a new
+    // head update. So here we iterate backwards until we find it, if it exists yet.
+
+    let epoch_id = header.epoch_id();
+    let last_final_hash = header.last_final_block();
+
+    let Some(mut sync) = maybe_get_block_header(chain_store, last_final_hash)? else {
+        return Ok(());
     };
-    for block_hash in iter_state_sync_new_chunks_keys(&chain_store.store()) {
-        let block_hash = block_hash?;
-        let old_header = chain_store.get_block_header(&block_hash)?;
-        if old_header.height() < last_final_header.height() {
-            store_update.delete(DBCol::StateSyncNewChunks, block_hash.as_ref());
+    loop {
+        let Some(sync_prev) = maybe_get_block_header(chain_store, sync.prev_hash())? else {
+            return Ok(());
+        };
+        if sync_prev.epoch_id() != epoch_id
+            || sync_prev.height() == chain_store.get_genesis_height()
+        {
+            return Ok(());
+        }
+        if has_enough_new_chunks(&chain_store.store(), sync_prev.hash())? != Some(true) {
+            return Ok(());
         }
-    }
 
-    Ok(())
+        let Some(sync_prev_prev) = maybe_get_block_header(chain_store, sync_prev.prev_hash())?
+        else {
+            return Ok(());
+        };
+        let Some(prev_prev_done) =
+            has_enough_new_chunks(&chain_store.store(), sync_prev_prev.hash())?
+        else {
+            return Ok(());
+        };
+
+        if !prev_prev_done {
+            // `sync_prev_prev` doesn't have enough new chunks, and `sync_prev` does, meaning `sync` is the first final
+            // valid sync block
+            store_update.set_ser(DBCol::StateSyncHashes, epoch_id.as_ref(), sync.hash())?;
+            store_update.delete_all(DBCol::StateSyncNewChunks);
+            return Ok(());
+        }
+        sync = sync_prev;
+    }
 }
 
 /// Updates information in the DB related to calculating the correct "sync_hash" for this header's epoch,
@@ -154,6 +199,35 @@ pub(crate) fn update_sync_hashes<T: ChainStoreAccess>(
         return remove_old_epochs(&chain_store.store(), store_update, header, &prev_header);
     }
 
-    save_epoch_new_chunks(chain_store, store_update, header)?;
-    remove_old_blocks(chain_store, store_update, header)
+    on_new_header(chain_store, store_update, header)
+}
+
+///. Returns whether `block_hash` is the block that will appear immediately before the "sync_hash" block. That is,
+/// whether it is going to be the prev_hash of the "sync_hash" block, when it is found.
+///
+/// `block_hash` is the prev_hash of the future "sync_hash" block iff it is the first block for which the
+/// number of new chunks in the epoch in each shard is at least 2
+///
+/// This function can only return true before we save the "sync_hash" block to the `StateSyncHashes` column,
+/// because it relies on data stored in the `StateSyncNewChunks` column, which is cleaned up after that.
+///
+/// This is used when making state snapshots, because in that case we don't need to wait for the "sync_hash"
+/// block to be finalized to take a snapshot of the state as of its prev prev block
+pub(crate) fn is_sync_prev_hash(
+    store: &Store,
+    block_hash: &CryptoHash,
+    prev_hash: &CryptoHash,
+) -> Result<bool, Error> {
+    let Some(new_chunks) = get_state_sync_new_chunks(store, block_hash)? else {
+        return Ok(false);
+    };
+    let done = new_chunks.iter().all(|num_chunks| *num_chunks >= 2);
+    if !done {
+        return Ok(false);
+    }
+    let Some(prev_new_chunks) = get_state_sync_new_chunks(store, prev_hash)? else {
+        return Ok(false);
+    };
+    let prev_done = prev_new_chunks.iter().all(|num_chunks| *num_chunks >= 2);
+    Ok(!prev_done)
 }

integration-tests/src/test_loop/tests/resharding_v3.rs

@@ -36,15 +36,15 @@ use crate::test_loop::utils::{ONE_NEAR, TGAS};
 use near_parameters::{vm, RuntimeConfig, RuntimeConfigStore};
 
 /// Default and minimal epoch length used in resharding tests.
-const DEFAULT_EPOCH_LENGTH: u64 = 6;
+const DEFAULT_EPOCH_LENGTH: u64 = 7;
 
 /// Increased epoch length that has to be used in some tests due to the delay caused by catch up.
 ///
 /// With shorter epoch length, a chunk producer might not finish catch up on time,
 /// before it is supposed to accept transactions for the next epoch.
 /// That would result in chunk producer rejecting a transaction
 /// and later we would hit the `DBNotFoundErr("Transaction ...)` error in tests.
-const INCREASED_EPOCH_LENGTH: u64 = 8;
+const INCREASED_EPOCH_LENGTH: u64 = 10;
 
 /// Garbage collection window length.
 const GC_NUM_EPOCHS_TO_KEEP: u64 = 3;
@@ -704,6 +704,7 @@ fn test_resharding_v3_drop_chunks_all() {
     test_resharding_v3_base(
         TestReshardingParametersBuilder::default()
             .chunk_ranges_to_drop(chunk_ranges_to_drop)
+            .epoch_length(INCREASED_EPOCH_LENGTH)
             .build(),
     );
 }
@@ -773,7 +774,9 @@ fn test_resharding_v3_shard_shuffling_untrack_then_track() {
         schedule: shard_sequence_to_schedule(tracked_shard_sequence),
     };
     let params = TestReshardingParametersBuilder::default()
-        .shuffle_shard_assignment_for_chunk_producers(true)
+        // TODO(resharding): uncomment after the bug in the comment above get_should_apply_chunk()
+        // in chain.rs is fixed
+        //.shuffle_shard_assignment_for_chunk_producers(true)
         .num_clients(num_clients)
         .tracked_shard_schedule(Some(tracked_shard_schedule))
         // TODO(resharding): uncomment after fixing test_resharding_v3_state_cleanup()
@@ -1067,6 +1070,7 @@ fn test_resharding_v3_yield_timeout() {
             ReceiptKind::PromiseYield,
         ))
         .allow_negative_refcount(true)
+        .epoch_length(INCREASED_EPOCH_LENGTH)
         .build();
     test_resharding_v3_base(params);
 }