fix: Save genesis bootstrap congestion info into DB and reuse it at node restarts

chain/chain/src/chain.rs

@@ -3943,54 +3943,67 @@ fn get_genesis_congestion_infos_impl(
     runtime: &dyn RuntimeAdapter,
     state_roots: &Vec<CryptoHash>,
 ) -> Result<Vec<Option<CongestionInfo>>, Error> {
-    let prev_hash = CryptoHash::default();
-    let epoch_id = epoch_manager.get_epoch_id_from_prev_block(&prev_hash)?;
-    let protocol_version = epoch_manager.get_epoch_protocol_version(&epoch_id)?;
+    let genesis_prev_hash = CryptoHash::default();
+    let genesis_epoch_id = epoch_manager.get_epoch_id_from_prev_block(&genesis_prev_hash)?;
+    let genesis_protocol_version = epoch_manager.get_epoch_protocol_version(&genesis_epoch_id)?;
+    // If congestion control is not enabled at the genesis block, we return None (congestion info) for each shard.
+    if !ProtocolFeature::CongestionControl.enabled(genesis_protocol_version) {
+        return Ok(std::iter::repeat(None).take(state_roots.len()).collect());
+    }
+
+    // Since the congestion info is already bootstrapped in statelessnet, skip another bootstrap.
+    // TODO: This is temporary mitigation for the failing genesis congestion info due to garbage
+    // collected genesis state roots. It can be removed after the statelessnet network is turned down.
+    if let Ok(protocol_config) = runtime.get_protocol_config(&genesis_epoch_id) {
+        if protocol_config.genesis_config.chain_id == near_primitives::chains::STATELESSNET {
+            return Ok(std::iter::repeat(None).take(state_roots.len()).collect());
+        }
+    }
+
+    // Check we had already computed the congestion infos from the genesis state roots.
+    if let Some(saved_infos) = near_store::get_genesis_congestion_infos(runtime.store())? {
+        tracing::debug!(target: "chain", "Reading genesis congestion infos from database.");
+        return Ok(saved_infos.into_iter().map(Option::Some).collect());
+    }
 
-    let mut result = vec![];
+    let mut new_infos = vec![];
     for (shard_id, &state_root) in state_roots.iter().enumerate() {
         let shard_id = shard_id as ShardId;
-        let congestion_info = get_congestion_info(
+        let congestion_info = get_genesis_congestion_info(
             runtime,
-            protocol_version,
-            &prev_hash,
+            genesis_protocol_version,
+            &genesis_prev_hash,
             shard_id,
             state_root,
-            &epoch_id,
         )?;
-        result.push(congestion_info);
+        new_infos.push(congestion_info);
     }
-    Ok(result)
+
+    // Store it in DB so that we can read it later, instead of recomputing from genesis state roots.
+    // Note that this is necessary because genesis state roots will be garbage-collected and will not
+    // be available, for example, when the node restarts later.
+    tracing::debug!(target: "chain", "Saving genesis congestion infos to database.");
+    let mut store_update = runtime.store().store_update();
+    near_store::set_genesis_congestion_infos(&mut store_update, &new_infos);
+    store_update.commit()?;
+
+    Ok(new_infos.into_iter().map(Option::Some).collect())
 }
 
-fn get_congestion_info(
+fn get_genesis_congestion_info(
     runtime: &dyn RuntimeAdapter,
     protocol_version: ProtocolVersion,
     prev_hash: &CryptoHash,
     shard_id: ShardId,
     state_root: StateRoot,
-    epoch_id: &EpochId,
-) -> Result<Option<CongestionInfo>, Error> {
-    if !ProtocolFeature::CongestionControl.enabled(protocol_version) {
-        return Ok(None);
-    }
-
-    // Since the congestion info is already bootstrapped in statelessnet, skip another bootstrap.
-    // TODO: This is temporary mitigation for the failing genesis congestion info due to garbage
-    // collected genesis state roots. It can be removed after the statelessnet network is turned down.
-    if let Ok(protocol_config) = runtime.get_protocol_config(&epoch_id) {
-        if protocol_config.genesis_config.chain_id == near_primitives::chains::STATELESSNET {
-            return Ok(None);
-        }
-    }
-
+) -> Result<CongestionInfo, Error> {
     // Get the view trie because it's possible that the chain is ahead of
     // genesis and doesn't have this block in flat state and memtrie.
     let trie = runtime.get_view_trie_for_shard(shard_id, prev_hash, state_root)?;
     let runtime_config = runtime.get_runtime_config(protocol_version)?;
     let congestion_info = bootstrap_congestion_info(&trie, &runtime_config, shard_id)?;
     tracing::debug!(target: "chain", ?shard_id, ?state_root, ?congestion_info, "Computed genesis congestion info.");
-    Ok(Some(congestion_info))
+    Ok(congestion_info)
 }
 
 fn shard_id_out_of_bounds(shard_id: ShardId) -> Error {

core/store/src/db.rs

@@ -34,6 +34,7 @@ pub const LATEST_KNOWN_KEY: &[u8; 12] = b"LATEST_KNOWN";
 pub const LARGEST_TARGET_HEIGHT_KEY: &[u8; 21] = b"LARGEST_TARGET_HEIGHT";
 pub const GENESIS_JSON_HASH_KEY: &[u8; 17] = b"GENESIS_JSON_HASH";
 pub const GENESIS_STATE_ROOTS_KEY: &[u8; 19] = b"GENESIS_STATE_ROOTS";
+pub const GENESIS_CONGESTION_INFO_KEY: &[u8] = b"GENESIS_CONGESTION_INFO_KEY";
 pub const COLD_HEAD_KEY: &[u8; 9] = b"COLD_HEAD";
 pub const STATE_SYNC_DUMP_KEY: &[u8; 15] = b"STATE_SYNC_DUMP";
 pub const STATE_SNAPSHOT_KEY: &[u8; 18] = b"STATE_SNAPSHOT_KEY";

core/store/src/lib.rs

@@ -12,6 +12,7 @@ pub use crate::trie::{
 };
 use borsh::{BorshDeserialize, BorshSerialize};
 pub use columns::DBCol;
+use db::GENESIS_CONGESTION_INFO_KEY;
 pub use db::{
     CHUNK_TAIL_KEY, COLD_HEAD_KEY, FINAL_HEAD_KEY, FORK_TAIL_KEY, GENESIS_JSON_HASH_KEY,
     GENESIS_STATE_ROOTS_KEY, HEADER_HEAD_KEY, HEAD_KEY, LARGEST_TARGET_HEIGHT_KEY,
@@ -21,6 +22,7 @@ use metadata::{DbKind, DbVersion, KIND_KEY, VERSION_KEY};
 use near_crypto::PublicKey;
 use near_fmt::{AbbrBytes, StorageKey};
 use near_primitives::account::{AccessKey, Account};
+use near_primitives::congestion_info::CongestionInfo;
 pub use near_primitives::errors::{MissingTrieValueContext, StorageError};
 use near_primitives::hash::CryptoHash;
 use near_primitives::receipt::{
@@ -1011,6 +1013,10 @@ pub fn get_genesis_state_roots(store: &Store) -> io::Result<Option<Vec<StateRoot
     store.get_ser::<Vec<StateRoot>>(DBCol::BlockMisc, GENESIS_STATE_ROOTS_KEY)
 }
 
+pub fn get_genesis_congestion_infos(store: &Store) -> io::Result<Option<Vec<CongestionInfo>>> {
+    store.get_ser::<Vec<CongestionInfo>>(DBCol::BlockMisc, GENESIS_CONGESTION_INFO_KEY)
+}
+
 pub fn get_genesis_hash(store: &Store) -> io::Result<Option<CryptoHash>> {
     store.get_ser::<CryptoHash>(DBCol::BlockMisc, GENESIS_JSON_HASH_KEY)
 }
@@ -1027,6 +1033,15 @@ pub fn set_genesis_state_roots(store_update: &mut StoreUpdate, genesis_roots: &[
         .expect("Borsh cannot fail");
 }
 
+pub fn set_genesis_congestion_infos(
+    store_update: &mut StoreUpdate,
+    congestion_infos: &[CongestionInfo],
+) {
+    store_update
+        .set_ser(DBCol::BlockMisc, GENESIS_CONGESTION_INFO_KEY, &congestion_infos)
+        .expect("Borsh cannot fail");
+}
+
 fn option_to_not_found<T, F>(res: io::Result<Option<T>>, field_name: F) -> io::Result<T>
 where
     F: std::string::ToString,

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

@@ -0,0 +1,72 @@
+use near_async::time::Duration;
+use near_chain::ChainStoreAccess;
+use near_chain_configs::test_genesis::TestGenesisBuilder;
+use near_client::Client;
+use near_o11y::testonly::init_test_logger;
+use near_primitives::types::AccountId;
+use near_primitives::version::{ProtocolFeature, PROTOCOL_VERSION};
+
+use crate::test_loop::builder::TestLoopBuilder;
+use crate::test_loop::env::TestLoopEnv;
+use crate::test_loop::utils::ONE_NEAR;
+
+const NUM_SHARDS: usize = 4;
+
+/// This test checks that the genesis congestion control info is saved into DB and not cleaned during GC,
+/// so that client can use it to bootstrap the genesis congestion control info after restarting.
+/// Restarting is the node is not checked here but in python/nayduck tests.
+#[test]
+fn test_congestion_control_genesis_bootstrap() {
+    if !ProtocolFeature::CongestionControl.enabled(PROTOCOL_VERSION) {
+        return;
+    }
+
+    init_test_logger();
+
+    let builder = TestLoopBuilder::new();
+
+    let initial_balance = 10000 * ONE_NEAR;
+    let accounts = ["test0", "test1"];
+    let clients: Vec<AccountId> = accounts.iter().map(|account| account.parse().unwrap()).collect();
+
+    let mut genesis_builder = TestGenesisBuilder::new();
+    genesis_builder
+        .genesis_time_from_clock(&builder.clock())
+        .protocol_version_latest()
+        .shard_layout_simple_v1(&["account3", "account5", "account7"])
+        .validators_desired_roles(&accounts[0..1], &accounts[1..2])
+        .minimum_validators_per_shard(1);
+
+    for i in 0..clients.len() {
+        genesis_builder.add_user_account_simple(clients[i].clone(), initial_balance);
+    }
+
+    let TestLoopEnv { mut test_loop, datas: node_datas, tempdir } =
+        builder.genesis(genesis_builder.build()).clients(clients.clone()).build();
+
+    test_loop.run_for(Duration::seconds(5));
+
+    for i in 0..clients.len() {
+        check_genesis_congestion_info_in_store(
+            &mut test_loop.data.get_mut(&node_datas[i].client_sender.actor_handle()).client,
+        );
+    }
+
+    TestLoopEnv { test_loop, datas: node_datas, tempdir }
+        .shutdown_and_drain_remaining_events(Duration::seconds(20));
+}
+
+fn check_genesis_congestion_info_in_store(client: &mut Client) {
+    let gc_config = client.config.gc.clone();
+    client.chain.clear_data(&gc_config).unwrap();
+
+    let infos = near_store::get_genesis_congestion_infos(client.chain.chain_store().store())
+        .unwrap()
+        .unwrap();
+    assert_eq!(infos.len(), NUM_SHARDS);
+    for i in 0..NUM_SHARDS {
+        assert_eq!(infos[i].buffered_receipts_gas(), 0);
+        assert_eq!(infos[i].delayed_receipts_gas(), 0);
+        assert_eq!(infos[i].receipt_bytes(), 0);
+    }
+}

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

@@ -1,5 +1,6 @@
 mod chunk_validator_kickout;
 pub mod congestion_control;
+pub mod congestion_control_genesis_bootstrap;
 pub mod in_memory_tries;
 pub mod multinode_stateless_validators;
 pub mod multinode_test_loop_example;

nightly/pytest-sanity.txt

@@ -188,9 +188,8 @@ pytest sanity/slow_chunk.py --features nightly
 # TODO(congestion_control) - enable pytest on stabilization
 # pytest sanity/congestion_control.py
 pytest sanity/congestion_control.py --features nightly
-# TODO(#11702) Enable these after fixing the issue and stabilization.
-# pytest sanity/congestion_control_genesis_bootstrap.py
-# pytest sanity/congestion_control_genesis_bootstrap.py --features nightly
+pytest sanity/congestion_control_genesis_bootstrap.py
+pytest sanity/congestion_control_genesis_bootstrap.py --features nightly
 
 # Tests the correct operation of the view client without using memtries (#11312).
 pytest sanity/rpc_view_history.py

pytest/tests/sanity/memtrie_disktrie_switch.py

@@ -119,11 +119,12 @@ def test(self):
         self.__restart_nodes(enable_memtries=False)
         self.__random_workload_until(target_height)
 
-        # TODO(#11675): Fix MissingTrieValue error and re-enable this step of the test.
-        # target_height = self.__next_target_height(num_epochs=1)
-        # logger.info(f"Step 3: Restarting nodes with memtries enabled until height {target_height}")
-        # self.__restart_nodes(enable_memtries=True)
-        # self.__random_workload_until(target_height)
+        target_height = self.__next_target_height(num_epochs=1)
+        logger.info(
+            f"Step 3: Restarting nodes with memtries enabled until height {target_height}"
+        )
+        self.__restart_nodes(enable_memtries=True)
+        self.__random_workload_until(target_height)
 
         self.__wait_for_txs(self.txs, assert_all_accepted=False)
         logger.info("Test ended")