chore: apply suggestions from CR

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

Author: WenyXu

src/metric-engine/src/engine.rs

@@ -23,6 +23,7 @@ mod options;
 mod put;
 mod read;
 mod region_metadata;
+mod staging;
 mod state;
 mod sync;
 
@@ -210,12 +211,7 @@ impl RegionEngine for MetricEngine {
         let result = match request {
             RegionRequest::EnterStaging(_) => {
                 if self.inner.is_physical_region(region_id) {
-                    self.inner
-                        .mito
-                        .handle_request(region_id, request)
-                        .await
-                        .context(error::MitoEnterStagingOperationSnafu)
-                        .map(|response| response.affected_rows)
+                    self.handle_enter_staging_request(region_id, request).await
                 } else {
                     UnsupportedRegionRequestSnafu { request }.fail()
                 }

src/metric-engine/src/engine/staging.rs

@@ -0,0 +1,54 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use common_base::AffectedRows;
+use snafu::ResultExt;
+use store_api::region_engine::RegionEngine;
+use store_api::region_request::{EnterStagingRequest, RegionRequest};
+use store_api::storage::RegionId;
+
+use crate::engine::MetricEngine;
+use crate::error::{MitoEnterStagingOperationSnafu, Result};
+use crate::utils;
+
+impl MetricEngine {
+    /// Handles the enter staging request for the given region.
+    pub(crate) async fn handle_enter_staging_request(
+        &self,
+        region_id: RegionId,
+        request: RegionRequest,
+    ) -> Result<AffectedRows> {
+        let metadata_region_id = utils::to_metadata_region_id(region_id);
+        let data_region_id = utils::to_data_region_id(region_id);
+
+        // For metadata region, it doesn't care about the partition expr, so we can just pass an empty string.
+        self.inner
+            .mito
+            .handle_request(
+                metadata_region_id,
+                RegionRequest::EnterStaging(EnterStagingRequest {
+                    partition_expr: String::new(),
+                }),
+            )
+            .await
+            .context(MitoEnterStagingOperationSnafu)?;
+
+        self.inner
+            .mito
+            .handle_request(data_region_id, request)
+            .await
+            .context(MitoEnterStagingOperationSnafu)
+            .map(|response| response.affected_rows)
+    }
+}

src/mito2/src/engine/staging_test.rs

@@ -471,7 +471,7 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool)
         // One file for flush operation
         normal_count_before,
         2,
-        "Normal manifest directory should initially contain 2 file"
+        "Normal manifest directory should initially contain 2 files"
     );
 
     // Try read data before exiting staging, SST files should be invisible

src/mito2/src/manifest/manager.rs

@@ -151,6 +151,9 @@ pub struct RegionManifestManager {
     last_version: Arc<AtomicU64>,
     checkpointer: Checkpointer,
     manifest: Arc<RegionManifest>,
+    // Staging manifest is used to store the manifest of the staging region before it becomes available.
+    // It is initially inherited from the previous manifest(i.e., `self.manifest`).
+    // When the staging manifest becomes available, it will be used to construct the new manifest.
     staging_manifest: Option<Arc<RegionManifest>>,
     stats: ManifestStats,
     stopped: bool,
@@ -557,7 +560,6 @@ impl RegionManifestManager {
 
         if is_staging {
             let new_manifest = manifest_builder.try_build()?;
-            debug!("Built staging manifest: {:?}", new_manifest);
             self.staging_manifest = Some(Arc::new(new_manifest));
 
             info!(

src/mito2/src/region.rs

@@ -516,8 +516,13 @@ impl MitoRegion {
     }
 
     /// Switches the region state to `RegionRoleState::Leader(RegionLeaderState::Staging)` if the current state is `expect`.
-    pub(crate) fn switch_state_to_staging(&self, expect: RegionLeaderState) -> Result<()> {
-        self.compare_exchange_state(expect, RegionRoleState::Leader(RegionLeaderState::Staging))
+    /// Otherwise, logs an error.
+    pub(crate) fn switch_state_to_staging(&self, expect: RegionLeaderState) {
+        if let Err(e) =
+            self.compare_exchange_state(expect, RegionRoleState::Leader(RegionLeaderState::Staging))
+        {
+            error!(e; "failed to switch region state to staging, expect state is {:?}", expect);
+        }
     }
 
     /// Returns the region statistic.
@@ -610,15 +615,10 @@ impl MitoRegion {
             .await
             .map(|m| m.files.clone())
             .unwrap_or_default();
-        debug!(
-            "manifest files: {:?}, staging files: {:?}",
-            manifest_files, staging_files
-        );
         let files = manifest_files
             .into_iter()
             .chain(staging_files.into_iter())
             .collect::<HashMap<_, _>>();
-        debug!("files: {:?}", files);
 
         files
             .values()

src/mito2/src/worker/handle_enter_staging.rs

@@ -201,22 +201,6 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             }
         };
 
-        let clean_staging_manifests = |region: MitoRegionRef| {
-            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 staging");
-                }
-            });
-        };
-
-        if let Err(e) = region.switch_state_to_staging(RegionLeaderState::EnteringStaging) {
-            error!(e; "Failed to switch region state to staging");
-            enter_staging_result.sender.send(Err(e));
-            clean_staging_manifests(region);
-            return;
-        }
-
         if enter_staging_result.result.is_ok() {
             info!(
                 "Updating region {} staging partition expr to {}",
@@ -226,8 +210,9 @@ impl<S: LogStore> RegionWorkerLoop<S> {
                 &region,
                 enter_staging_result.partition_expr,
             );
+            region.switch_state_to_staging(RegionLeaderState::EnteringStaging);
         } else {
-            clean_staging_manifests(region);
+            region.switch_state_to_writable(RegionLeaderState::EnteringStaging);
         }
         enter_staging_result
             .sender

src/store-api/src/region_request.rs

@@ -1410,7 +1410,7 @@ impl RegionBulkInsertsRequest {
 /// Request to stage a region with a new region rule(partition expression).
 ///
 /// This request transitions a region into the staging mode.
-/// It first flushes the memtable for the old region rule,
+/// It first flushes the memtable for the old region rule if it is not empty,
 /// then enters the staging mode with the new region rule.
 #[derive(Debug, Clone)]
 pub struct EnterStagingRequest {