fix: create FlatStorageState inside RuntimeAdapter (#7710)

We create FlatStorageState on Chain creation, but it currently happens for any RuntimeAdapter impl. Because KeyValueRuntime doesn't and shouldn't support flat storage, it leads to tests failures.

Here we move FSS creation inside RuntimeAdapter impl.

## Testing

https://buildkite.com/nearprotocol/nearcore-flat-state/builds/60#018383c5-7838-4e5b-86f1-da2baa381789 - number of failing tests drops from 77 to 4 - which are failing due to other reasons

chain/chain/src/chain.rs

@@ -81,7 +81,7 @@ use near_primitives::shard_layout::{
 };
 use near_primitives::version::PROTOCOL_VERSION;
 #[cfg(feature = "protocol_feature_flat_state")]
-use near_store::flat_state::{FlatStateDelta, FlatStorageState};
+use near_store::flat_state::FlatStateDelta;
 use once_cell::sync::OnceCell;
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
 
@@ -629,13 +629,11 @@ impl Chain {
         // set up flat storage
         #[cfg(feature = "protocol_feature_flat_state")]
         for shard_id in 0..runtime_adapter.num_shards(&block_head.epoch_id)? {
-            let flat_storage_state = FlatStorageState::new(
-                store.store().clone(),
+            runtime_adapter.create_flat_storage_state_for_shard(
                 shard_id,
                 store.head()?.height,
                 &store,
             );
-            runtime_adapter.add_flat_storage_state_for_shard(shard_id, flat_storage_state);
         }
 
         info!(target: "chain", "Init: header head @ #{} {}; block head @ #{} {}",

chain/chain/src/test_utils.rs

@@ -61,6 +61,8 @@ use crate::{BlockProcessingArtifact, Provenance};
 use near_primitives::epoch_manager::ShardConfig;
 use near_primitives::time::Clock;
 use near_primitives::utils::MaybeValidated;
+#[cfg(feature = "protocol_feature_flat_state")]
+use near_store::flat_state::ChainAccessForFlatStorage;
 use near_store::flat_state::FlatStorageState;
 
 pub use self::validator_schedule::ValidatorSchedule;
@@ -709,10 +711,13 @@ impl RuntimeAdapter for KeyValueRuntime {
     fn get_flat_storage_state_for_shard(&self, _shard_id: ShardId) -> Option<FlatStorageState> {
         None
     }
-    fn add_flat_storage_state_for_shard(
+
+    #[cfg(feature = "protocol_feature_flat_state")]
+    fn create_flat_storage_state_for_shard(
         &self,
         _shard_id: ShardId,
-        _flat_storage_state: FlatStorageState,
+        _latest_block_height: BlockHeight,
+        _chain_access: &dyn ChainAccessForFlatStorage,
     ) {
     }
 

chain/chain/src/types.rs

@@ -34,6 +34,8 @@ use near_primitives::version::{
     MIN_PROTOCOL_VERSION_NEP_92_FIX,
 };
 use near_primitives::views::{EpochValidatorInfo, QueryRequest, QueryResponse};
+#[cfg(feature = "protocol_feature_flat_state")]
+use near_store::flat_state::ChainAccessForFlatStorage;
 use near_store::flat_state::FlatStorageState;
 use near_store::{PartialStorage, ShardTries, Store, StoreUpdate, Trie, WrappedTrieChanges};
 
@@ -286,10 +288,12 @@ pub trait RuntimeAdapter: EpochManagerAdapter + Send + Sync {
 
     fn get_flat_storage_state_for_shard(&self, shard_id: ShardId) -> Option<FlatStorageState>;
 
-    fn add_flat_storage_state_for_shard(
+    #[cfg(feature = "protocol_feature_flat_state")]
+    fn create_flat_storage_state_for_shard(
         &self,
         shard_id: ShardId,
-        flat_storage_state: FlatStorageState,
+        latest_block_height: BlockHeight,
+        chain_access: &dyn ChainAccessForFlatStorage,
     );
 
     fn set_flat_storage_state_for_genesis(

core/store/src/flat_state.rs

@@ -147,7 +147,7 @@ mod imp {
             }))
         }
 
-        /// When a node starts from an empty database, this function must be called  to ensure
+        /// When a node starts from an empty database, this function must be called to ensure
         /// information such as flat head is set up correctly in the database.
         /// Note that this function is different from `add_flat_storage_state_for_shard`,
         /// it must be called before `add_flat_storage_state_for_shard` if the node starts from

nearcore/src/runtime/mod.rs

@@ -46,6 +46,8 @@ use near_primitives::views::{
     AccessKeyInfoView, CallResult, EpochValidatorInfo, QueryRequest, QueryResponse,
     QueryResponseKind, ViewApplyState, ViewStateResult,
 };
+#[cfg(feature = "protocol_feature_flat_state")]
+use near_store::flat_state::ChainAccessForFlatStorage;
 use near_store::flat_state::{FlatStateFactory, FlatStorageState};
 use near_store::split_state::get_delayed_receipts;
 use near_store::{
@@ -699,11 +701,15 @@ impl RuntimeAdapter for NightshadeRuntime {
         self.flat_state_factory.get_flat_storage_state_for_shard(shard_id)
     }
 
-    fn add_flat_storage_state_for_shard(
+    #[cfg(feature = "protocol_feature_flat_state")]
+    fn create_flat_storage_state_for_shard(
         &self,
         shard_id: ShardId,
-        flat_storage_state: FlatStorageState,
+        latest_block_height: BlockHeight,
+        chain_access: &dyn ChainAccessForFlatStorage,
     ) {
+        let flat_storage_state =
+            FlatStorageState::new(self.store.clone(), shard_id, latest_block_height, chain_access);
         self.flat_state_factory.add_flat_storage_state_for_shard(shard_id, flat_storage_state)
     }