feat(resharding): flat storage resharding mvp

chain/chain-primitives/src/error.rs

@@ -232,6 +232,9 @@ pub enum Error {
     /// GC error.
     #[error("GC Error: {0}")]
     GCError(String),
+    /// Resharding error.
+    #[error("Resharding Error: {0}")]
+    ReshardingError(String),
     /// Anything else
     #[error("Other Error: {0}")]
     Other(String),
@@ -269,6 +272,7 @@ impl Error {
             | Error::CannotBeFinalized
             | Error::StorageError(_)
             | Error::GCError(_)
+            | Error::ReshardingError(_)
             | Error::DBNotFoundErr(_) => false,
             Error::InvalidBlockPastTime(_, _)
             | Error::InvalidBlockFutureTime(_)
@@ -392,6 +396,7 @@ impl Error {
             Error::NotAValidator(_) => "not_a_validator",
             Error::NotAChunkValidator => "not_a_chunk_validator",
             Error::InvalidChallengeRoot => "invalid_challenge_root",
+            Error::ReshardingError(_) => "resharding_error",
         }
     }
 }

chain/chain/src/flat_storage_creator.rs

@@ -9,6 +9,7 @@
 //! `CatchingUp`: moves flat storage head forward, so it may reach chain final head.
 //! `Ready`: flat storage is created and it is up-to-date.
 
+use crate::flat_storage_resharder::FlatStorageResharder;
 use crate::types::RuntimeAdapter;
 use crate::{ChainStore, ChainStoreAccess};
 use assert_matches::assert_matches;
@@ -388,6 +389,12 @@ impl FlatStorageShardCreator {
             FlatStorageStatus::Disabled => {
                 panic!("initiated flat storage creation for shard {shard_id} while it is disabled");
             }
+            // If the flat storage is undergoing resharding it means it was previously created successfully,
+            // but resharding itself hasn't been finished.
+            // This case is a no-op because the flat storage resharder has already been created in `create_flat_storage_for_current_epoch`.
+            FlatStorageStatus::Resharding(_) => {
+                return Ok(true);
+            }
         };
         Ok(false)
     }
@@ -403,10 +410,13 @@ pub struct FlatStorageCreator {
 impl FlatStorageCreator {
     /// For each of tracked shards, either creates flat storage if it is already stored on DB,
     /// or starts migration to flat storage which updates DB in background and creates flat storage afterwards.
+    ///
+    /// Also resumes any resharding operation which was already in progress.
     pub fn new(
         epoch_manager: Arc<dyn EpochManagerAdapter>,
         runtime: Arc<dyn RuntimeAdapter>,
         chain_store: &ChainStore,
+        flat_storage_resharder: &FlatStorageResharder,
         num_threads: usize,
     ) -> Result<Option<Self>, Error> {
         let flat_storage_manager = runtime.get_flat_storage_manager();
@@ -420,6 +430,7 @@ impl FlatStorageCreator {
             &epoch_manager,
             &flat_storage_manager,
             &runtime,
+            &flat_storage_resharder,
         )?;
 
         // Create flat storage for the shards in the next epoch. This only
@@ -447,6 +458,7 @@ impl FlatStorageCreator {
         epoch_manager: &Arc<dyn EpochManagerAdapter>,
         flat_storage_manager: &FlatStorageManager,
         runtime: &Arc<dyn RuntimeAdapter>,
+        _flat_storage_resharder: &FlatStorageResharder,
     ) -> Result<HashMap<ShardUId, FlatStorageShardCreator>, Error> {
         let epoch_id = &chain_head.epoch_id;
         tracing::debug!(target: "store", ?epoch_id, "creating flat storage for the current epoch");
@@ -473,6 +485,10 @@ impl FlatStorageCreator {
                     );
                 }
                 FlatStorageStatus::Disabled => {}
+                FlatStorageStatus::Resharding(_status) => {
+                    // TODO(Trisfald): call resume
+                    // flat_storage_resharder.resume(shard_uid, &status, ...)?;
+                }
             }
         }
 
@@ -502,7 +518,8 @@ impl FlatStorageCreator {
                 }
                 FlatStorageStatus::Empty
                 | FlatStorageStatus::Creation(_)
-                | FlatStorageStatus::Disabled => {
+                | FlatStorageStatus::Disabled
+                | FlatStorageStatus::Resharding(_) => {
                     // The flat storage for children shards will be created
                     // separately in the resharding process.
                 }