feat: introduce flat storage status (near#8689)

Part of near#8577

This PR introduces `FlatStorageStatus`:
* it generalises `FlatStorageCreationStatus`
* `FlatStorageCreationStatus::DontCreate` is moved to `::Disabled` status
* flat head is now stored as a part of `::Ready` status.
* `::Empty` status is added to explicitly express that flat storage does not exist
* we no longer need `FlatStateMisc` column, so it is removed 

nayduck run: https://nayduck.near.org/#/run/2905

Author: pugachAG

chain/chain/src/chain.rs

@@ -70,8 +70,8 @@ use near_primitives::views::{
     LightClientBlockView, SignedTransactionView,
 };
 use near_store::flat::{
-    store_helper, FlatStateChanges, FlatStateDelta, FlatStateDeltaMetadata,
-    FlatStorageCreationStatus, FlatStorageError,
+    store_helper, FlatStateChanges, FlatStateDelta, FlatStateDeltaMetadata, FlatStorageError,
+    FlatStorageReadyStatus, FlatStorageStatus,
 };
 use near_store::StorageError;
 use near_store::{DBCol, ShardTries, StoreUpdate, WrappedTrieChanges};
@@ -633,6 +633,7 @@ impl Chain {
                 if cfg!(feature = "protocol_feature_flat_state") {
                     let tmp_store_update = runtime_adapter.set_flat_storage_for_genesis(
                         genesis.hash(),
+                        genesis.header().height(),
                         genesis.header().epoch_id(),
                     )?;
                     store_update.merge(tmp_store_update);
@@ -3178,23 +3179,32 @@ impl Chain {
         // flat storage. Now we can set flat head to hash of this block and create flat storage.
         // TODO (#7327): ensure that no flat storage work is done for `KeyValueRuntime`.
         if cfg!(feature = "protocol_feature_flat_state") {
-            let mut store_update = self.runtime_adapter.store().store_update();
-            store_helper::set_flat_head(&mut store_update, shard_id, block_hash);
-            store_update.commit()?;
-        }
-
-        if self.runtime_adapter.get_flat_storage_creation_status(shard_id)
-            == FlatStorageCreationStatus::Ready
-        {
             // If block_hash is equal to default - this means that we're all the way back at genesis.
             // So we don't have to add the storage state for shard in such case.
             // TODO(8438) - add additional test scenarios for this case.
             if *block_hash != CryptoHash::default() {
                 let block_header = self.get_block_header(block_hash)?;
                 let epoch_id = block_header.epoch_id();
                 let shard_uid = self.runtime_adapter.shard_id_to_uid(shard_id, epoch_id)?;
-
-                self.runtime_adapter.create_flat_storage_for_shard(shard_uid);
+                if !matches!(
+                    self.runtime_adapter.get_flat_storage_status(shard_uid),
+                    FlatStorageStatus::Disabled
+                ) {
+                    let mut store_update = self.runtime_adapter.store().store_update();
+                    store_helper::set_flat_storage_status(
+                        &mut store_update,
+                        shard_uid,
+                        FlatStorageStatus::Ready(FlatStorageReadyStatus {
+                            flat_head: near_store::flat::BlockInfo {
+                                hash: *block_hash,
+                                prev_hash: *block_header.prev_hash(),
+                                height: block_header.height(),
+                            },
+                        }),
+                    );
+                    store_update.commit()?;
+                    self.runtime_adapter.create_flat_storage_for_shard(shard_uid);
+                }
             }
         }
 

chain/chain/src/flat_storage_creator.rs

@@ -19,8 +19,8 @@ use near_primitives::state::ValueRef;
 use near_primitives::state_part::PartId;
 use near_primitives::types::{AccountId, BlockHeight, StateRoot};
 use near_store::flat::{
-    store_helper, FetchingStateStatus, FlatStateChanges, FlatStorageCreationStatus,
-    NUM_PARTS_IN_ONE_STEP, STATE_PART_MEMORY_LIMIT,
+    store_helper, BlockInfo, FetchingStateStatus, FlatStateChanges, FlatStorageCreationStatus,
+    FlatStorageReadyStatus, FlatStorageStatus, NUM_PARTS_IN_ONE_STEP, STATE_PART_MEMORY_LIMIT,
 };
 use near_store::{Store, FLAT_STORAGE_HEAD_HEIGHT};
 use near_store::{Trie, TrieDBStorage, TrieTraversalItem};
@@ -165,10 +165,19 @@ impl FlatStorageShardCreator {
     ) -> Result<bool, Error> {
         let shard_id = self.shard_uid.shard_id();
         let current_status =
-            store_helper::get_flat_storage_creation_status(chain_store.store(), shard_id);
+            store_helper::get_flat_storage_status(chain_store.store(), self.shard_uid);
         self.metrics.status.set((&current_status).into());
         match &current_status {
-            FlatStorageCreationStatus::SavingDeltas => {
+            FlatStorageStatus::Empty => {
+                let mut store_update = chain_store.store().store_update();
+                store_helper::set_flat_storage_status(
+                    &mut store_update,
+                    self.shard_uid,
+                    FlatStorageStatus::Creation(FlatStorageCreationStatus::SavingDeltas),
+                );
+                store_update.commit()?;
+            }
+            FlatStorageStatus::Creation(FlatStorageCreationStatus::SavingDeltas) => {
                 let final_head = chain_store.final_head()?;
                 let final_height = final_head.height;
 
@@ -220,15 +229,19 @@ impl FlatStorageShardCreator {
 
                     let mut store_update = chain_store.store().store_update();
                     self.metrics.flat_head_height.set(final_head.height as i64);
-                    store_helper::set_flat_storage_creation_status(
+                    store_helper::set_flat_storage_status(
                         &mut store_update,
-                        shard_id,
-                        FlatStorageCreationStatus::FetchingState(status),
+                        self.shard_uid,
+                        FlatStorageStatus::Creation(FlatStorageCreationStatus::FetchingState(
+                            status,
+                        )),
                     );
                     store_update.commit()?;
                 }
             }
-            FlatStorageCreationStatus::FetchingState(fetching_state_status) => {
+            FlatStorageStatus::Creation(FlatStorageCreationStatus::FetchingState(
+                fetching_state_status,
+            )) => {
                 let store = self.runtime_adapter.store().clone();
                 let block_hash = fetching_state_status.block_hash;
                 let start_part_id = fetching_state_status.part_id;
@@ -297,26 +310,30 @@ impl FlatStorageShardCreator {
                                 num_parts,
                             };
                             debug!(target: "chain", %shard_id, %block_hash, ?new_status);
-                            store_helper::set_flat_storage_creation_status(
+                            store_helper::set_flat_storage_status(
                                 &mut store_update,
-                                shard_id,
-                                FlatStorageCreationStatus::FetchingState(new_status),
+                                self.shard_uid,
+                                FlatStorageStatus::Creation(
+                                    FlatStorageCreationStatus::FetchingState(new_status),
+                                ),
                             );
                         } else {
                             // If all parts were fetched, we can start catchup.
                             info!(target: "chain", %shard_id, %block_hash, "Finished fetching state");
                             self.metrics.remaining_state_parts.set(0);
-                            store_helper::set_flat_storage_creation_status(
+                            store_helper::set_flat_storage_status(
                                 &mut store_update,
-                                shard_id,
-                                FlatStorageCreationStatus::CatchingUp(block_hash),
+                                self.shard_uid,
+                                FlatStorageStatus::Creation(FlatStorageCreationStatus::CatchingUp(
+                                    block_hash,
+                                )),
                             );
                         }
                         store_update.commit()?;
                     }
                 }
             }
-            FlatStorageCreationStatus::CatchingUp(old_flat_head) => {
+            FlatStorageStatus::Creation(FlatStorageCreationStatus::CatchingUp(old_flat_head)) => {
                 let store = self.runtime_adapter.store();
                 let mut flat_head = *old_flat_head;
                 let chain_final_head = chain_store.final_head()?;
@@ -346,37 +363,46 @@ impl FlatStorageShardCreator {
                     let epoch_id = self.runtime_adapter.get_epoch_id(&flat_head)?;
                     let shard_uid = self.runtime_adapter.shard_id_to_uid(shard_id, &epoch_id)?;
                     let old_height = chain_store.get_block_height(&old_flat_head).unwrap();
-                    let height = chain_store.get_block_height(&flat_head).unwrap();
+                    let flat_head_block_header = chain_store.get_block_header(&flat_head).unwrap();
+                    let height = flat_head_block_header.height();
                     debug!(target: "chain", %shard_id, %old_flat_head, %old_height, %flat_head, %height, "Catching up flat head");
                     self.metrics.flat_head_height.set(height as i64);
                     let mut store_update = self.runtime_adapter.store().store_update();
                     merged_changes.apply_to_flat_state(&mut store_update, shard_uid);
                     if flat_head == chain_final_head.last_block_hash {
                         // If we reached chain final head, we can finish catchup and finally create flat storage.
-                        store_helper::remove_flat_storage_creation_status(
+                        store_helper::set_flat_storage_status(
                             &mut store_update,
-                            shard_id,
+                            self.shard_uid,
+                            FlatStorageStatus::Ready(FlatStorageReadyStatus {
+                                flat_head: BlockInfo {
+                                    hash: flat_head,
+                                    prev_hash: *flat_head_block_header.prev_hash(),
+                                    height,
+                                },
+                            }),
                         );
-                        store_helper::set_flat_head(&mut store_update, shard_id, &flat_head);
                         store_update.commit()?;
                         self.runtime_adapter.create_flat_storage_for_shard(shard_uid);
                         info!(target: "chain", %shard_id, %flat_head, %height, "Flat storage creation done");
                     } else {
-                        store_helper::set_flat_storage_creation_status(
+                        store_helper::set_flat_storage_status(
                             &mut store_update,
-                            shard_id,
-                            FlatStorageCreationStatus::CatchingUp(flat_head),
+                            self.shard_uid,
+                            FlatStorageStatus::Creation(FlatStorageCreationStatus::CatchingUp(
+                                flat_head,
+                            )),
                         );
                         store_update.commit()?;
                     }
                 }
             }
-            FlatStorageCreationStatus::Ready => {}
-            FlatStorageCreationStatus::DontCreate => {
-                panic!("We initiated flat storage creation for shard {shard_id} but according to flat storage state status in db it cannot be created");
+            FlatStorageStatus::Ready(_) => return Ok(true),
+            FlatStorageStatus::Disabled => {
+                panic!("initiated flat storage creation for shard {shard_id} while it is disabled");
             }
         };
-        Ok(current_status == FlatStorageCreationStatus::Ready)
+        Ok(false)
     }
 }
 
@@ -402,15 +428,14 @@ impl FlatStorageCreator {
         let mut creation_needed = false;
         for shard_id in 0..num_shards {
             if runtime_adapter.cares_about_shard(me, &chain_head.prev_block_hash, shard_id, true) {
-                let status = runtime_adapter.get_flat_storage_creation_status(shard_id);
                 let shard_uid = runtime_adapter.shard_id_to_uid(shard_id, &chain_head.epoch_id)?;
+                let status = runtime_adapter.get_flat_storage_status(shard_uid);
 
                 match status {
-                    FlatStorageCreationStatus::Ready => {
+                    FlatStorageStatus::Ready(_) => {
                         runtime_adapter.create_flat_storage_for_shard(shard_uid);
                     }
-                    FlatStorageCreationStatus::DontCreate => {}
-                    _ => {
+                    FlatStorageStatus::Empty | FlatStorageStatus::Creation(_) => {
                         creation_needed = true;
                         shard_creators.insert(
                             shard_uid,
@@ -421,6 +446,7 @@ impl FlatStorageCreator {
                             ),
                         );
                     }
+                    FlatStorageStatus::Disabled => {}
                 }
             }
         }

chain/chain/src/store.rs

@@ -2438,7 +2438,7 @@ impl<'a> ChainStoreUpdate<'a> {
             DBCol::FlatState
             | DBCol::FlatStateChanges
             | DBCol::FlatStateDeltaMetadata
-            | DBCol::FlatStateMisc => {
+            | DBCol::FlatStorageStatus => {
                 unreachable!();
             }
         }

chain/chain/src/test_utils/kv_runtime.rs

@@ -52,7 +52,7 @@ use crate::{BlockHeader, RuntimeWithEpochManagerAdapter};
 
 use near_primitives::epoch_manager::ShardConfig;
 
-use near_store::flat::{FlatStorage, FlatStorageCreationStatus};
+use near_store::flat::{FlatStorage, FlatStorageStatus};
 
 use super::ValidatorSchedule;
 
@@ -833,8 +833,8 @@ impl RuntimeAdapter for KeyValueRuntime {
         None
     }
 
-    fn get_flat_storage_creation_status(&self, _shard_id: ShardId) -> FlatStorageCreationStatus {
-        FlatStorageCreationStatus::DontCreate
+    fn get_flat_storage_status(&self, _shard_uid: ShardUId) -> FlatStorageStatus {
+        FlatStorageStatus::Disabled
     }
 
     fn create_flat_storage_for_shard(&self, shard_uid: ShardUId) {
@@ -852,6 +852,7 @@ impl RuntimeAdapter for KeyValueRuntime {
     fn set_flat_storage_for_genesis(
         &self,
         _genesis_block: &CryptoHash,
+        _genesis_block_height: BlockHeight,
         _genesis_epoch_id: &EpochId,
     ) -> Result<StoreUpdate, Error> {
         Ok(self.store.store_update())

chain/chain/src/types.rs

@@ -31,7 +31,7 @@ use near_primitives::version::{
     MIN_PROTOCOL_VERSION_NEP_92_FIX,
 };
 use near_primitives::views::{QueryRequest, QueryResponse};
-use near_store::flat::{FlatStorage, FlatStorageCreationStatus};
+use near_store::flat::{FlatStorage, FlatStorageStatus};
 use near_store::{PartialStorage, ShardTries, Store, StoreUpdate, Trie, WrappedTrieChanges};
 
 pub use near_epoch_manager::EpochManagerAdapter;
@@ -299,8 +299,7 @@ pub trait RuntimeAdapter: Send + Sync {
 
     fn get_flat_storage_for_shard(&self, shard_id: ShardId) -> Option<FlatStorage>;
 
-    /// Gets status of flat storage state background creation.
-    fn get_flat_storage_creation_status(&self, shard_id: ShardId) -> FlatStorageCreationStatus;
+    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.
@@ -318,6 +317,7 @@ pub trait RuntimeAdapter: Send + Sync {
     fn set_flat_storage_for_genesis(
         &self,
         genesis_block: &CryptoHash,
+        genesis_block_height: BlockHeight,
         genesis_epoch_id: &EpochId,
     ) -> Result<StoreUpdate, Error>;
 

core/store/src/columns.rs

@@ -258,27 +258,25 @@ pub enum DBCol {
     /// *Column type*: ExecutionOutcomeWithProof
     TransactionResultForBlock,
     /// Flat state contents. Used to get `ValueRef` by trie key faster than doing a trie lookup.
-    /// - *Rows*: trie key (Vec<u8>)
+    /// - *Rows*: `shard_uid` + trie key (Vec<u8>)
     /// - *Column type*: ValueRef
     #[cfg(feature = "protocol_feature_flat_state")]
     FlatState,
     /// Changes for flat state delta. Stores how flat state should be updated for the given shard and block.
-    /// - *Rows*: `KeyForFlatStateDelta { shard_id, block_hash }`
+    /// - *Rows*: `KeyForFlatStateDelta { shard_uid, block_hash }`
     /// - *Column type*: `FlatStateChanges`
     #[cfg(feature = "protocol_feature_flat_state")]
     FlatStateChanges,
     /// Metadata for flat state delta.
-    /// - *Rows*: `KeyForFlatStateDelta { shard_id, block_hash }`
+    /// - *Rows*: `KeyForFlatStateDelta { shard_uid, block_hash }`
     /// - *Column type*: `FlatStateDeltaMetadata`
     #[cfg(feature = "protocol_feature_flat_state")]
     FlatStateDeltaMetadata,
-    /// Miscellaneous data for flat state. Stores intermediate flat storage creation statuses and flat
-    /// state heads for each shard.
-    /// - *Rows*: Unique key prefix (e.g. `FLAT_STATE_HEAD_KEY_PREFIX`) + ShardId
-    /// - *Column type*: FetchingStateStatus || flat storage catchup status (bool) || flat storage head (CryptoHash)
-    // TODO (#7327): use only during testing, come up with proper format.
+    /// Flat storage status for the corresponding shard.
+    /// - *Rows*: `shard_uid`
+    /// - *Column type*: `FlatStorageStatus`
     #[cfg(feature = "protocol_feature_flat_state")]
-    FlatStateMisc,
+    FlatStorageStatus,
 }
 
 /// Defines different logical parts of a db key.
@@ -483,7 +481,7 @@ impl DBCol {
             #[cfg(feature = "protocol_feature_flat_state")]
             DBCol::FlatStateDeltaMetadata => &[DBKeyType::ShardId, DBKeyType::BlockHash],
             #[cfg(feature = "protocol_feature_flat_state")]
-            DBCol::FlatStateMisc => &[DBKeyType::ShardId],
+            DBCol::FlatStorageStatus => &[DBKeyType::ShardUId],
         }
     }
 }

core/store/src/flat/manager.rs

@@ -1,8 +1,10 @@
-use crate::flat::{store_helper, POISONED_LOCK_ERR};
+use crate::flat::{
+    store_helper, BlockInfo, FlatStorageReadyStatus, FlatStorageStatus, POISONED_LOCK_ERR,
+};
 use near_primitives::errors::StorageError;
 use near_primitives::hash::CryptoHash;
-use near_primitives::shard_layout::ShardLayout;
-use near_primitives::types::ShardId;
+use near_primitives::shard_layout::{ShardLayout, ShardUId};
+use near_primitives::types::{BlockHeight, ShardId};
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
 use tracing::debug;
@@ -44,12 +46,20 @@ impl FlatStorageManager {
     pub fn set_flat_storage_for_genesis(
         &self,
         store_update: &mut StoreUpdate,
-        shard_id: ShardId,
+        shard_uid: ShardUId,
         genesis_block: &CryptoHash,
+        genesis_height: BlockHeight,
     ) {
+        let shard_id = shard_uid.shard_id();
         let flat_storages = self.0.flat_storages.lock().expect(POISONED_LOCK_ERR);
         assert!(!flat_storages.contains_key(&shard_id));
-        store_helper::set_flat_head(store_update, shard_id, genesis_block);
+        store_helper::set_flat_storage_status(
+            store_update,
+            shard_uid,
+            FlatStorageStatus::Ready(FlatStorageReadyStatus {
+                flat_head: BlockInfo::genesis(*genesis_block, genesis_height),
+            }),
+        );
     }
 
     /// Add a flat storage state for shard `shard_id`. The function also checks that

core/store/src/flat/mod.rs

@@ -36,7 +36,10 @@ pub use chunk_view::FlatStorageChunkView;
 pub use delta::{FlatStateChanges, FlatStateDelta, FlatStateDeltaMetadata};
 pub use manager::FlatStorageManager;
 pub use storage::FlatStorage;
-pub use types::{BlockInfo, FetchingStateStatus, FlatStorageCreationStatus, FlatStorageError};
+pub use types::{
+    BlockInfo, FetchingStateStatus, FlatStorageCreationStatus, FlatStorageError,
+    FlatStorageReadyStatus, FlatStorageStatus,
+};
 
 pub(crate) const POISONED_LOCK_ERR: &str = "The lock was poisoned.";