impl

chain/chain/src/chain.rs

@@ -2136,7 +2136,11 @@ impl Chain {
         block: &Block,
         shard_uid: ShardUId,
     ) -> Result<(), Error> {
-        if let Some(flat_storage) = self.runtime_adapter.get_flat_storage_for_shard(shard_uid) {
+        if let Some(flat_storage) = self
+            .runtime_adapter
+            .get_flat_storage_manager()
+            .and_then(|manager| manager.get_flat_storage_for_shard(shard_uid))
+        {
             let mut new_flat_head = *block.header().last_final_block();
             if new_flat_head == CryptoHash::default() {
                 new_flat_head = *self.genesis.hash();
@@ -3245,14 +3249,9 @@ impl Chain {
             let epoch_id = block_header.epoch_id();
             let shard_uid = self.epoch_manager.shard_id_to_uid(shard_id, epoch_id)?;
 
-            // Check if flat storage is disabled, which may be the case when runtime is implemented with
-            // `KeyValueRuntime`.
-            if !matches!(
-                self.runtime_adapter.get_flat_storage_status(shard_uid),
-                FlatStorageStatus::Disabled
-            ) {
+            if let Some(flat_storage_manager) = self.runtime_adapter.get_flat_storage_manager() {
                 // Flat storage must not exist at this point because leftover keys corrupt its state.
-                assert!(self.runtime_adapter.get_flat_storage_for_shard(shard_uid).is_none());
+                assert!(flat_storage_manager.get_flat_storage_for_shard(shard_uid).is_none());
 
                 let mut store_update = self.runtime_adapter.store().store_update();
                 store_helper::set_flat_storage_status(
@@ -4995,7 +4994,10 @@ impl<'a> ChainUpdate<'a> {
             },
         };
 
-        if let Some(chain_flat_storage) = self.runtime_adapter.get_flat_storage_for_shard(shard_uid)
+        if let Some(chain_flat_storage) = self
+            .runtime_adapter
+            .get_flat_storage_manager()
+            .and_then(|manager| manager.get_flat_storage_for_shard(shard_uid))
         {
             // If flat storage exists, we add a block to it.
             let store_update =

chain/chain/src/flat_storage_creator.rs

@@ -441,7 +441,7 @@ impl FlatStorageCreator {
         for shard_id in 0..num_shards {
             if shard_tracker.care_about_shard(me, &chain_head.prev_block_hash, shard_id, true) {
                 let shard_uid = epoch_manager.shard_id_to_uid(shard_id, &chain_head.epoch_id)?;
-                let status = runtime.get_flat_storage_status(shard_uid);
+                let status = store_helper::get_flat_storage_status(chain_store.store(), shard_uid);
 
                 match status {
                     FlatStorageStatus::Ready(_) => {

chain/chain/src/test_utils/kv_runtime.rs

@@ -52,8 +52,6 @@ use crate::BlockHeader;
 
 use near_primitives::epoch_manager::ShardConfig;
 
-use near_store::flat::{FlatStorage, FlatStorageStatus};
-
 use super::ValidatorSchedule;
 
 /// Simple key value runtime for tests.
@@ -976,14 +974,6 @@ impl RuntimeAdapter for KeyValueRuntime {
         ))
     }
 
-    fn get_flat_storage_for_shard(&self, _shard_uid: ShardUId) -> Option<FlatStorage> {
-        None
-    }
-
-    fn get_flat_storage_status(&self, _shard_uid: ShardUId) -> FlatStorageStatus {
-        FlatStorageStatus::Disabled
-    }
-
     fn create_flat_storage_for_shard(&self, shard_uid: ShardUId) {
         panic!("Flat storage state can't be created for shard {shard_uid} because KeyValueRuntime doesn't support this");
     }

chain/chain/src/types.rs

@@ -32,7 +32,6 @@ use near_primitives::version::{
     MIN_PROTOCOL_VERSION_NEP_92_FIX,
 };
 use near_primitives::views::{QueryRequest, QueryResponse};
-use near_store::flat::{FlatStorage, FlatStorageStatus};
 use near_store::{PartialStorage, ShardTries, Store, StoreUpdate, Trie, WrappedTrieChanges};
 
 pub use near_epoch_manager::EpochManagerAdapter;
@@ -302,10 +301,6 @@ pub trait RuntimeAdapter: Send + Sync {
 
     fn get_flat_storage_manager(&self) -> Option<FlatStorageManager>;
 
-    fn get_flat_storage_for_shard(&self, shard_uid: ShardUId) -> Option<FlatStorage>;
-
-    fn get_flat_storage_status(&self, shard_uid: ShardUId) -> FlatStorageStatus;
-
     /// Creates flat storage state for given shard, assuming that all flat storage data
     /// is already stored in DB.
     /// TODO (#7327): consider returning flat storage creation errors here

core/store/src/flat/manager.rs

@@ -91,6 +91,10 @@ impl FlatStorageManager {
         assert!(original_value.is_none());
     }
 
+    pub fn get_flat_storage_status(&self, shard_uid: ShardUId) -> FlatStorageStatus {
+        store_helper::get_flat_storage_status(&self.0.store, shard_uid)
+    }
+
     /// Creates `FlatStorageChunkView` to access state for `shard_uid` and block `block_hash`.
     /// Note that:
     /// * the state includes changes by the block `block_hash`;

nearcore/src/runtime/mod.rs

@@ -47,7 +47,7 @@ use near_primitives::views::{
     AccessKeyInfoView, CallResult, QueryRequest, QueryResponse, QueryResponseKind, ViewApplyState,
     ViewStateResult,
 };
-use near_store::flat::{store_helper, FlatStorage, FlatStorageManager, FlatStorageStatus};
+use near_store::flat::{FlatStorage, FlatStorageManager};
 use near_store::metadata::DbKind;
 use near_store::split_state::get_delayed_receipts;
 use near_store::{
@@ -739,14 +739,6 @@ impl RuntimeAdapter for NightshadeRuntime {
         Some(self.flat_storage_manager.clone())
     }
 
-    fn get_flat_storage_for_shard(&self, shard_uid: ShardUId) -> Option<FlatStorage> {
-        self.flat_storage_manager.get_flat_storage_for_shard(shard_uid)
-    }
-
-    fn get_flat_storage_status(&self, shard_uid: ShardUId) -> FlatStorageStatus {
-        store_helper::get_flat_storage_status(&self.store, shard_uid)
-    }
-
     // TODO (#7327): consider passing flat storage errors here to handle them gracefully
     fn create_flat_storage_for_shard(&self, shard_uid: ShardUId) {
         let flat_storage = FlatStorage::new(self.store.clone(), shard_uid);