add integration tests and get shard_uid with runtime adapter

chain/chain/src/store.rs

@@ -12,7 +12,9 @@ use near_primitives::errors::InvalidTxError;
 use near_primitives::hash::CryptoHash;
 use near_primitives::merkle::{MerklePath, PartialMerkleTree};
 use near_primitives::receipt::Receipt;
-use near_primitives::shard_layout::{account_id_to_shard_id, get_block_shard_uid, ShardLayout, ShardUId};
+use near_primitives::shard_layout::{
+    account_id_to_shard_id, get_block_shard_uid, ShardUId,
+};
 use near_primitives::sharding::{
     ChunkHash, EncodedShardChunk, PartialEncodedChunk, ReceiptProof, ShardChunk, ShardChunkHeader,
     StateSyncInfo,
@@ -37,7 +39,6 @@ use near_primitives::utils::{
     to_timestamp,
 };
 use near_primitives::views::LightClientBlockView;
-use near_store::flat::store_helper;
 use near_store::{
     DBCol, KeyForStateChanges, ShardTries, Store, StoreUpdate, WrappedTrieChanges, CHUNK_TAIL_KEY,
     FINAL_HEAD_KEY, FORK_TAIL_KEY, HEADER_HEAD_KEY, HEAD_KEY, LARGEST_TARGET_HEIGHT_KEY,
@@ -1957,11 +1958,7 @@ impl<'a> ChainStoreUpdate<'a> {
         self.chunk_tail = Some(height);
     }
 
-    pub fn clear_chunk_data_and_headers_at_height(
-        &mut self,
-        height: BlockHeight,
-    ) -> Result<(), Error> {
-        let chunk_tail = self.chunk_tail()?;
+    fn clear_chunk_data_and_headers_at_height(&mut self, height: BlockHeight) -> Result<(), Error> {
         let chunk_hashes = self.chain_store.get_all_chunk_hashes_by_height(height)?;
         for chunk_hash in chunk_hashes {
             // 1. Delete chunk-related data
@@ -1988,16 +1985,14 @@ impl<'a> ChainStoreUpdate<'a> {
             let key: &[u8] = header_hash.as_bytes();
             store_update.delete(DBCol::BlockHeader, key);
             self.chain_store.headers.pop(key);
+            self.merge(store_update);
         }
 
         // 4. Delete chunks_tail-related data
         let key = index_to_bytes(height);
         self.gc_col(DBCol::ChunkHashesByHeight, &key);
         self.gc_col(DBCol::HeaderHashesByHeight, &key);
 
-        if chunk_tail > height {
-            self.update_chunk_tail(height);
-        }
         Ok(())
     }
 
@@ -2251,13 +2246,10 @@ impl<'a> ChainStoreUpdate<'a> {
     }
 
     // Delete all data in rocksdb that are partially or wholly indexed and can be looked up by hash of the current head of the chain
-    // Delete all data in rocksdb that indicates a link between current head and its prev block
-    pub fn clear_block_data_undo_block(&mut self, block_hash: CryptoHash) -> Result<(), Error> {
-        if self.head().unwrap().last_block_hash != block_hash {
-            return Err(Error::Other(format!("block_hash is not the hash of current head of the chain.")))
-        }
+    // and that indicates a link between current head and its prev block
+    pub fn clear_head_block_data(&mut self, runtime_adapter: &dyn RuntimeWithEpochManagerAdapter) -> Result<(), Error> {
+        let block_hash = self.head().unwrap().last_block_hash;
 
-        let mut store_update = self.store().store_update();
         let block =
             self.get_block(&block_hash).expect("block data is not expected to be already cleaned");
 
@@ -2289,13 +2281,12 @@ impl<'a> ChainStoreUpdate<'a> {
             }
 
             // delete flat storage columns: FlatStateChanges and FlatStateDeltaMetadata
-            #[cfg(feature = "protocol_feature_flat_state")]
-            let shard_uid_flat_storage = ShardUId::from_shard_id_and_layout(
-                shard_id as u64,
-                &ShardLayout::get_simple_nightshade_layout(),
-            );
-            #[cfg(feature = "protocol_feature_flat_state")]
-            store_helper::remove_delta(&mut store_update, shard_uid_flat_storage, block_hash);
+            #[cfg(feature = "protocol_feature_flat_state")] {
+                let mut store_update = self.store().store_update();
+                let shard_uid= runtime_adapter.shard_id_to_uid(shard_id);
+                store_helper::remove_delta(&mut store_update, shard_uid_flat_storage, block_hash);
+                self.merge(store_update);
+            }
         }
 
         // 3. Delete block_hash-indexed data
@@ -2328,7 +2319,6 @@ impl<'a> ChainStoreUpdate<'a> {
 
         self.clear_chunk_data_and_headers_at_height(height)?;
 
-        self.merge(store_update);
         Ok(())
     }
 

integration-tests/Cargo.toml

@@ -48,6 +48,7 @@ near-o11y.workspace = true
 near-telemetry.workspace = true
 near-test-contracts.workspace = true
 near-performance-metrics.workspace = true
+near-undo-block.workspace = true
 near-vm-errors.workspace = true
 near-vm-runner.workspace = true
 nearcore.workspace = true

integration-tests/src/tests/client/mod.rs

@@ -9,3 +9,4 @@ mod runtimes;
 #[cfg(feature = "sandbox")]
 mod sandbox;
 mod sharding_upgrade;
+mod undo_block;

integration-tests/src/tests/client/undo_block.rs

@@ -0,0 +1,74 @@
+use near_chain::{
+    ChainGenesis, ChainStore, ChainStoreAccess, Provenance, RuntimeWithEpochManagerAdapter,
+};
+use near_chain_configs::Genesis;
+use near_client::test_utils::TestEnv;
+use near_o11y::testonly::init_test_logger;
+use near_store::test_utils::create_test_store;
+use near_store::Store;
+use near_undo_block::undo_block;
+use nearcore::config::GenesisExt;
+use std::path::Path;
+use std::sync::Arc;
+
+/// Setup environment with one Near client for testing.
+fn setup_env(genesis: &Genesis, store: Store) -> (TestEnv, Arc<dyn RuntimeWithEpochManagerAdapter>) {
+    let chain_genesis = ChainGenesis::new(genesis);
+    let runtimes: Vec<Arc<dyn RuntimeWithEpochManagerAdapter>> =
+        vec![nearcore::NightshadeRuntime::test(Path::new("../../../.."), store, genesis)];
+    (TestEnv::builder(chain_genesis).runtime_adapters(runtimes.clone()).build(), runtimes[0].clone())
+}
+
+fn test_undo_block(epoch_length: u64, stop_height: u64) {
+    init_test_logger();
+
+    let save_trie_changes = true;
+
+    let mut genesis = Genesis::test(vec!["test0".parse().unwrap(), "test1".parse().unwrap()], 1);
+    genesis.config.epoch_length = epoch_length;
+
+    let store = create_test_store();
+    let (mut env, runtime) = setup_env(&genesis, store.clone());
+
+    for i in 1..stop_height + 1 {
+        let block = env.clients[0].produce_block(i).unwrap().unwrap();
+        // blocks.push(block.clone());
+        env.process_block(0, block, Provenance::PRODUCED);
+    }
+
+    let mut chain_store =
+        ChainStore::new(store.clone(), genesis.config.genesis_height, save_trie_changes);
+
+    let current_head = chain_store.head().unwrap();
+    let prev_block_hash = current_head.prev_block_hash;
+
+    undo_block(&mut chain_store, &*runtime).unwrap();
+
+    // after undo, the current head should be the prev_block_hash
+    assert_eq!(chain_store.head().unwrap().last_block_hash.as_bytes(), prev_block_hash.as_bytes());
+    assert_eq!(chain_store.head().unwrap().height, stop_height - 1);
+
+    // set up an environment again with the same store
+    let (mut env, _) = setup_env(&genesis, store.clone());
+    // the new env should be able to produce block normally
+    let block = env.clients[0].produce_block(stop_height).unwrap().unwrap();
+    env.process_block(0, block, Provenance::PRODUCED);
+
+    // after processing the new block, the head should now be at stop_height
+    assert_eq!(chain_store.head().unwrap().height, stop_height);
+}
+
+#[test]
+fn test_undo_block_middle_of_epoch() {
+    test_undo_block(5, 3)
+}
+
+#[test]
+fn test_undo_block_end_of_epoch() {
+    test_undo_block(5, 5)
+}
+
+#[test]
+fn test_undo_block_start_of_epoch() {
+    test_undo_block(5, 6)
+}

neard/src/cli.rs

@@ -18,11 +18,11 @@ use near_ping::PingCommand;
 use near_primitives::hash::CryptoHash;
 use near_primitives::merkle::compute_root_from_path;
 use near_primitives::types::{Gas, NumSeats, NumShards};
-use near_undo_block::UndoBlockCommand;
 use near_state_parts::cli::StatePartsCommand;
 use near_state_viewer::StateViewerSubCommand;
 use near_store::db::RocksDB;
 use near_store::Mode;
+use near_undo_block::cli::UndoBlockCommand;
 use serde_json::Value;
 use std::fs::File;
 use std::io::BufReader;

tools/undo-block/Cargo.toml

@@ -11,8 +11,8 @@ clap.workspace = true
 tracing.workspace = true
 chrono.workspace = true
 
-near-chain = { path = "../../chain/chain" }
-near-chain-configs = { path = "../../core/chain-configs" }
-near-store = { path = "../../core/store" }
-nearcore = { path = "../../nearcore" }
-near-primitives = { path = "../../core/primitives" }
+near-chain.workspace = true
+near-chain-configs.workspace = true
+near-store.workspace = true
+nearcore.workspace = true
+near-primitives.workspace = true

tools/undo-block/src/cli.rs

@@ -0,0 +1,40 @@
+use nearcore::NightshadeRuntime;
+use near_chain::ChainStore;
+use near_chain_configs::GenesisValidationMode;
+use near_store::{Mode, NodeStorage};
+use nearcore::load_config;
+use std::path::Path;
+
+#[derive(clap::Parser)]
+pub struct UndoBlockCommand {}
+
+impl UndoBlockCommand {
+    pub fn run(
+        self,
+        home_dir: &Path,
+        genesis_validation: GenesisValidationMode,
+    ) -> anyhow::Result<()> {
+        let near_config = load_config(home_dir, genesis_validation)
+            .unwrap_or_else(|e| panic!("Error loading config: {:#}", e));
+
+        let store_opener = NodeStorage::opener(
+            home_dir,
+            near_config.config.archive,
+            &near_config.config.store,
+            None,
+        );
+
+        let storage = store_opener.open_in_mode(Mode::ReadWrite).unwrap();
+        let store = storage.get_hot_store();
+
+        let runtime = NightshadeRuntime::from_config(home_dir, store.clone(), &near_config);
+
+        let mut chain_store = ChainStore::new(
+            store,
+            near_config.genesis.config.genesis_height,
+            near_config.client_config.save_trie_changes,
+        );
+
+        crate::undo_block(&mut chain_store, &*runtime)
+    }
+}

tools/undo-block/src/lib.rs

@@ -1,78 +1,46 @@
 use chrono::Utc;
+use near_chain::RuntimeWithEpochManagerAdapter;
 use near_chain::types::LatestKnown;
 use near_chain::{ChainStore, ChainStoreAccess, ChainStoreUpdate};
-use near_chain_configs::GenesisValidationMode;
 use near_primitives::block::Tip;
 use near_primitives::utils::to_timestamp;
-use near_store::{Mode, NodeStorage};
-use nearcore::load_config;
-use std::path::Path;
 
-#[derive(clap::Parser)]
-pub struct UndoBlockCommand {}
+pub mod cli;
 
-impl UndoBlockCommand {
-    pub fn run(
-        self,
-        home_dir: &Path,
-        genesis_validation: GenesisValidationMode,
-    ) -> anyhow::Result<()> {
-        let near_config = load_config(home_dir, genesis_validation)
-            .unwrap_or_else(|e| panic!("Error loading config: {:#}", e));
+pub fn undo_block(chain_store: &mut ChainStore, runtime: &dyn RuntimeWithEpochManagerAdapter) -> anyhow::Result<()> {
+    let current_head = chain_store.head()?;
+    let current_head_hash = current_head.last_block_hash;
+    let prev_block_hash = current_head.prev_block_hash;
+    let prev_header = chain_store.get_block_header(&prev_block_hash)?;
+    let prev_tip = Tip::from_header(&prev_header);
+    let current_head_height = current_head.height;
+    let prev_block_height = prev_tip.height;
 
-        let store_opener = NodeStorage::opener(
-            home_dir,
-            near_config.config.archive,
-            &near_config.config.store,
-            None,
-        );
+    tracing::info!(target: "neard", ?prev_block_hash, ?current_head_hash, ?prev_block_height, ?current_head_height, "Trying to update head");
 
-        let storage = store_opener.open_in_mode(Mode::ReadWrite).unwrap();
-        let store = storage.get_hot_store();
-
-        let mut chain_store = ChainStore::new(
-            store,
-            near_config.genesis.config.genesis_height,
-            near_config.client_config.save_trie_changes,
-        );
-
-        let current_head = chain_store.head()?;
-        let current_head_hash = current_head.last_block_hash;
-        let prev_block_hash = current_head.prev_block_hash;
-        let prev_header = chain_store.get_block_header(&prev_block_hash)?;
-        let prev_tip = Tip::from_header(&prev_header);
-        let current_head_height = current_head.height;
-        let prev_block_height = prev_tip.height;
-
-        tracing::info!(target: "neard", ?prev_block_hash, ?current_head_hash, "Trying to update head");
-        tracing::info!(target: "neard", ?prev_block_height, ?current_head_height, "Trying to update head");
-
-        // stop if it's already the final block
-        if chain_store.final_head()?.height >= current_head.height {
-            return Err(anyhow::anyhow!("Cannot revert past final block"));
-        }
+    // stop if it's already the final block
+    if chain_store.final_head()?.height >= current_head.height {
+        return Err(anyhow::anyhow!("Cannot revert past final block"));
+    }
 
-        let mut chain_store_update = ChainStoreUpdate::new(&mut chain_store);
+    let mut chain_store_update = ChainStoreUpdate::new(chain_store);
 
-        // clear block data for current head
-        chain_store_update.clear_block_data_undo_block(current_head_hash)?;
+    chain_store_update.clear_head_block_data(runtime)?;
 
-        chain_store_update.save_head(&prev_tip)?;
+    chain_store_update.save_head(&prev_tip)?;
 
-        chain_store_update.commit()?;
+    chain_store_update.commit()?;
 
-        chain_store.save_latest_known(LatestKnown {
-            height: prev_tip.height,
-            seen: to_timestamp(Utc::now()),
-        })?;
+    chain_store.save_latest_known(LatestKnown {
+        height: prev_tip.height,
+        seen: to_timestamp(Utc::now()),
+    })?;
 
-        let new_chain_store_head = chain_store.head()?;
-        let new_chain_store_header_head = chain_store.header_head()?;
-        let new_head_height = new_chain_store_head.height;
-        let new_header_height = new_chain_store_header_head.height;
+    let new_chain_store_head = chain_store.head()?;
+    let new_chain_store_header_head = chain_store.header_head()?;
+    let new_head_height = new_chain_store_head.height;
+    let new_header_height = new_chain_store_header_head.height;
 
-        tracing::info!(target: "neard", ?new_head_height, ?new_header_height, "The current chain store shows");
-        tracing::info!(target: "neard", ?new_header_height, "The current chain store shows header head height");
-        Ok(())
-    }
+    tracing::info!(target: "neard", ?new_head_height, ?new_header_height, "The current chain store shows");
+    Ok(())
 }