chore: apply suggestions from CR

Signed-off-by: WenyXu <[email protected]>

Author: WenyXu

src/mito2/src/engine/staging_test.rs

@@ -468,7 +468,7 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool)
     let normal_count_before = normal_files_before.len();
     assert_eq!(
         // One file for table creation
-        // One files for flush operation
+        // One file for flush operation
         normal_count_before,
         2,
         "Normal manifest directory should initially contain 2 file"

src/mito2/src/region.rs

@@ -344,11 +344,19 @@ impl MitoRegion {
         )
     }
 
+    /// Exits the entering staging state back to writable.
+    pub(crate) fn exit_entering_staging(&self) -> Result<()> {
+        self.compare_exchange_state(
+            RegionLeaderState::EnteringStaging,
+            RegionRoleState::Leader(RegionLeaderState::Writable),
+        )
+    }
+
     /// Exits the staging state back to writable.
     ///
     /// You should call this method in the worker loop.
     /// Transitions from Staging to Writable state.
-    fn exit_staging(&self) -> Result<()> {
+    pub fn exit_staging(&self) -> Result<()> {
         self.compare_exchange_state(
             RegionLeaderState::Staging,
             RegionRoleState::Leader(RegionLeaderState::Writable),

src/mito2/src/worker/handle_enter_staging.rs

@@ -61,7 +61,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
 
         let version = region.version();
         if !version.memtables.is_empty() {
-            // If memtable is empty, we can't enter staging directly and need to flush
+            // If memtable is not empty, we can't enter staging directly and need to flush
             // all memtables first.
             info!("Flush region: {} before entering staging", region_id);
             debug_assert!(!region.is_staging());
@@ -127,6 +127,9 @@ impl<S: LogStore> RegionWorkerLoop<S> {
                         "Failed to clear staging manifest files for region {}",
                         region.region_id
                     );
+                    if let Err(e) = region.exit_entering_staging() {
+                        error!(e; "Failed to exit entering staging after failed to clear staging manifest files");
+                    }
                     sender.send(Err(e));
                     return;
                 }
@@ -202,13 +205,13 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             common_runtime::spawn_global(async move {
                 let mut manager = region.manifest_ctx.manifest_manager.write().await;
                 if let Err(e) = manager.clear_staging_manifests().await {
-                    error!(e; "Failed to clear staging manifests after failed to switch region state to entering staging");
+                    error!(e; "Failed to clear staging manifests after failed to switch region state to staging");
                 }
             });
         };
 
         if let Err(e) = region.switch_state_to_staging(RegionLeaderState::EnteringStaging) {
-            error!(e; "Failed to switch region state to entering staging");
+            error!(e; "Failed to switch region state to staging");
             enter_staging_result.sender.send(Err(e));
             clean_staging_manifests(region);
             return;