fix(mito2): keep rows written during staging in the WAL (#10253)

* fix(mito2): keep rows written during staging in the WAL

A flush during staging committed its SSTs to the staging manifest, then
dropped the memtable and released the WAL. If staging did not end
successfully those rows were gone. A flush that committed after staging
had ended lost its rows even though staging succeeded.

Staging regions no longer flush. They are skipped by memory-pressure,
periodic and region-limit flushes, and their writes are rejected instead
of queued at twice the region limit or when the node's staging memtables
reach half of the engine's flush limit. A manual or repartition flush on
a staging region with data gets a retryable error; the migration
pre-flush and the WAL prune flush reply success without flushing. Closing
a staging region leaves staging first, so skip-WAL rows are flushed to
the normal manifest.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(mito2): keep accounting and closing staging regions during staged edits

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(mito2): track staged edits without relying on the partition info

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* test(mito2): synchronize staging close tests with worker admission

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* test(mito2): observe the manifest lock wait when closing staging regions

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(mito2): detect staged edits through the staging partition info

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* docs(mito2): record that staging regions must not flush

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(mito2): do not restore staging after a close abandons the edit

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(mito2): reject region edits during a close-time flush

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* style(mito2): format the staging tests

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(mito2): reject staging edits once staging is abandoned

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(mito2): bind apply-staging edits to their partition expression

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* docs(mito2): note that staging regions pin their topic's prune point

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

---------

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

src/mito2/AGENTS.md

@@ -126,6 +126,9 @@ Tests live next to the code as `*_test.rs` (e.g. `src/mito2/src/engine/flush_tes
 - In flat batches, equal `__primary_key` dictionary keys imply the same series, but
   not the reverse: `wrap_pk_binary_to_dict` gives every row its own key. Compare key
   bytes when grouping rows by series.
+- Staging regions never flush: a flush would move rows into the staging manifest and
+  release their WAL, losing them if staging does not end successfully. Every flush
+  trigger in `worker/handle_flush.rs` must skip them, and `handle_close` leaves staging first.
 
 ## Maintenance contract
 

src/mito2/src/engine/partition_filter_test.rs

@@ -20,13 +20,15 @@ use datatypes::arrow::array::AsArray;
 use datatypes::value::Value;
 use partition::expr::{Operand, col};
 use partition::function::IcebergHash;
-use store_api::region_engine::RegionEngine;
+use store_api::region_engine::{RegionEngine, RemapManifestsRequest};
 use store_api::region_request::{
-    EnterStagingRequest, RegionFlushRequest, RegionRequest, StagingPartitionDirective,
+    ApplyStagingManifestRequest, EnterStagingRequest, RegionFlushRequest, RegionRequest,
+    StagingPartitionDirective,
 };
 use store_api::storage::{RegionId, ScanRequest};
 
 use crate::config::MitoConfig;
+use crate::engine::MitoEngine;
 use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema};
 
 /// Helper to create a partition expression for testing.
@@ -39,6 +41,30 @@ fn range_expr_string(col_name: &str, start: f64, end: f64) -> String {
         .unwrap()
 }
 
+/// Ends staging the way repartition does, importing the region's own files.
+async fn apply_staging(engine: &MitoEngine, region_id: RegionId, partition_expr: String) {
+    let result = engine
+        .remap_manifests(RemapManifestsRequest {
+            region_id,
+            input_regions: vec![region_id],
+            region_mapping: [(region_id, vec![region_id])].into_iter().collect(),
+            new_partition_exprs: [(region_id, partition_expr.clone())].into_iter().collect(),
+        })
+        .await
+        .unwrap();
+    engine
+        .handle_request(
+            region_id,
+            RegionRequest::ApplyStagingManifest(ApplyStagingManifestRequest {
+                partition_expr,
+                central_region_id: region_id,
+                manifest_path: result.manifest_paths[&region_id].clone(),
+            }),
+        )
+        .await
+        .unwrap();
+}
+
 #[tokio::test]
 async fn test_partition_filter_basic() {
     test_partition_filter_basic_with_format(false).await;
@@ -51,7 +77,7 @@ async fn test_partition_filter_basic() {
 /// 1. Creates a region with partition expr `tag_0 >= "0"` (covers all data)
 /// 2. Writes data (tag_0 = "0"..="9") and flushes - SST has old partition expr
 /// 3. Enters staging mode with partition expr `tag_0 >= "5"` (narrower)
-/// 4. Writes more data and flushes in staging mode
+/// 4. Writes more data in staging mode
 /// 5. Exits staging mode - region now has new partition expr
 /// 6. Scans all data - old SST data not matching new expr should be filtered
 async fn test_partition_filter_basic_with_format(flat_format: bool) {
@@ -118,16 +144,7 @@ async fn test_partition_filter_basic_with_format(flat_format: bool) {
     };
     put_rows(&engine, region_id, rows_data).await;
 
-    // Flush in staging mode
-    engine
-        .handle_request(
-            region_id,
-            RegionRequest::Flush(RegionFlushRequest::default()),
-        )
-        .await
-        .unwrap();
-
-    // Scan data in staging mode - should only see initial 5 rows (staging SST not visible)
+    // Scan data in staging mode - rows written in staging stay readable in the memtable
     let projection = Some(vec![1]);
     let request = ScanRequest {
         projection,
@@ -138,20 +155,16 @@ async fn test_partition_filter_basic_with_format(flat_format: bool) {
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     let total_rows: usize = batches.iter().map(|rb| rb.num_rows()).sum();
     assert_eq!(
-        total_rows, 5,
-        "Should see 5 rows before exiting staging mode"
+        total_rows, 11,
+        "Should see 11 rows before exiting staging mode"
     );
 
     // Exit staging mode
-    use store_api::region_engine::SettableRegionRoleState;
-    engine
-        .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader)
-        .await
-        .unwrap();
+    apply_staging(&engine, region_id, new_partition_expr).await;
 
     // Scan after exiting staging - the old SST (tag_0 = "0".."4") should have
     // rows filtered by partition expr (tag_0 >= "5"), which means none of them pass.
-    // But the staging SST (tag_0 = "5".."10") satisfies the partition expr.
+    // But the memtable rows (tag_0 = "5".."10") satisfy the partition expr.
     let projection = Some(vec![1]);
     let request = ScanRequest {
         projection,
@@ -163,10 +176,10 @@ async fn test_partition_filter_basic_with_format(flat_format: bool) {
     let total_rows: usize = batches.iter().map(|rb| rb.num_rows()).sum();
     // After exit:
     // - Old SST (rows 0-4): partition expr is "tag_0 >= 5", so these are filtered out
-    // - Staging SST (rows 5-10): These satisfy partition expr, so they're visible
+    // - Memtable (rows 5-10): These satisfy partition expr, so they're visible
     assert_eq!(
         total_rows, 6,
-        "Should see 6 rows after exiting staging (rows 5-10 from staging SST), flat_format: {}",
+        "Should see 6 rows after exiting staging (rows 5-10 from the memtable), flat_format: {}",
         flat_format
     );
 }
@@ -220,13 +233,16 @@ async fn test_partition_filter_iceberg_hash_with_format(flat_format: bool) {
         IcebergHash::evaluate_value(&Value::String(i.to_string().as_str().into())).unwrap()
             < Value::Int32(0)
     };
-    let hash_expr = Operand::IcebergHash(IcebergHash::new("tag_0")).lt(Value::Int32(0));
+    let hash_expr = Operand::IcebergHash(IcebergHash::new("tag_0"))
+        .lt(Value::Int32(0))
+        .as_json_str()
+        .unwrap();
     engine
         .handle_request(
             region_id,
             RegionRequest::EnterStaging(EnterStagingRequest {
                 partition_directive: StagingPartitionDirective::UpdatePartitionExpr(
-                    hash_expr.as_json_str().unwrap(),
+                    hash_expr.clone(),
                 ),
             }),
         )
@@ -244,12 +260,7 @@ async fn test_partition_filter_iceberg_hash_with_format(flat_format: bool) {
         },
     )
     .await;
-    flush().await;
-    use store_api::region_engine::SettableRegionRoleState;
-    engine
-        .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader)
-        .await
-        .unwrap();
+    apply_staging(&engine, region_id, hash_expr).await;
 
     let request = ScanRequest {
         projection: Some(vec![0]),

src/mito2/src/flush.rs

@@ -60,7 +60,7 @@ use crate::read::flat_dedup::{FlatDedupIterator, FlatLastNonNull, FlatLastRow};
 use crate::read::flat_merge::FlatMergeIterator;
 use crate::region::options::{IndexOptions, MergeMode, RegionOptions};
 use crate::region::version::{VersionControlData, VersionControlRef, VersionRef};
-use crate::region::{ManifestContextRef, RegionLeaderState, RegionRoleState, parse_partition_expr};
+use crate::region::{ManifestContextRef, RegionLeaderState, parse_partition_expr};
 use crate::request::{
     BackgroundNotify, DdlRequest, FlushFailed, FlushFinished, OnFailure, OptionOutputTx, OutputTx,
     SenderBulkRequest, SenderDdlRequest, SenderWriteRequest, WorkerRequest, WorkerRequestWithTime,
@@ -281,8 +281,6 @@ pub(crate) struct RegionFlushTask {
     pub(crate) index_options: IndexOptions,
     /// Semaphore to control flush concurrency.
     pub(crate) flush_semaphore: Arc<Semaphore>,
-    /// Whether the region is in staging mode.
-    pub(crate) is_staging: bool,
     /// Partition expression of the region.
     ///
     /// This is used to generate the file meta.
@@ -412,7 +410,6 @@ impl RegionFlushTask {
                     _timer: timer,
                     edit,
                     memtables_to_remove,
-                    is_staging: self.is_staging,
                 };
                 WorkerRequest::Background {
                     region_id: self.region_id,
@@ -537,10 +534,7 @@ impl RegionFlushTask {
             flushed_sequence: Some(version_data.committed_sequence),
             committed_sequence: None,
         };
-        info!(
-            "Applying {edit:?} to region {}, is_staging: {}",
-            self.region_id, self.is_staging
-        );
+        info!("Applying {edit:?} to region {}", self.region_id);
 
         let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone()));
 
@@ -553,17 +547,11 @@ impl RegionFlushTask {
         let expected_state = if matches!(self.reason, FlushReason::Downgrading) {
             RegionLeaderState::Downgrading
         } else {
-            // Check if region is in staging mode
-            let current_state = self.manifest_ctx.current_state();
-            if current_state == RegionRoleState::Leader(RegionLeaderState::Staging) {
-                RegionLeaderState::Staging
-            } else {
-                RegionLeaderState::Writable
-            }
+            RegionLeaderState::Writable
         };
         let manifest_version = match self
             .manifest_ctx
-            .update_manifest(expected_state, action_list, self.is_staging)
+            .update_manifest(expected_state, action_list, false)
             .await
         {
             Ok(manifest_version) => {
@@ -584,9 +572,8 @@ impl RegionFlushTask {
             }
         };
         info!(
-            "Successfully update manifest version to {manifest_version}, region: {}, is_staging: {}, reason: {}",
+            "Successfully update manifest version to {manifest_version}, region: {}, reason: {}",
             self.region_id,
-            self.is_staging,
             self.reason.as_str()
         );
 
@@ -1507,6 +1494,13 @@ impl FlushScheduler {
         })
     }
 
+    /// Returns true if the region has a close-time flush.
+    pub(crate) fn has_close_time_flush(&self, region_id: RegionId) -> bool {
+        self.region_status
+            .get(&region_id)
+            .is_some_and(|status| status.closing)
+    }
+
     /// Returns true if the region has pending DDLs or a close-time flush.
     pub(crate) fn has_pending_ddls(&self, region_id: RegionId) -> bool {
         self.region_status
@@ -1713,7 +1707,6 @@ mod tests {
             manifest_ctx,
             index_options: IndexOptions::default(),
             flush_semaphore: Arc::new(Semaphore::new(2)),
-            is_staging: false,
             partition_expr: None,
             durability_barrier: DurabilityBarrier::noop(),
         }
@@ -1864,7 +1857,6 @@ mod tests {
                 .await,
             index_options: IndexOptions::default(),
             flush_semaphore: Arc::new(Semaphore::new(2)),
-            is_staging: false,
             partition_expr: None,
             durability_barrier: DurabilityBarrier::noop(),
         };
@@ -1950,7 +1942,6 @@ mod tests {
                     committed_sequence: None,
                 },
                 memtables_to_remove: smallvec![],
-                is_staging: false,
             }),
         };
 
@@ -2166,7 +2157,6 @@ mod tests {
                 manifest_ctx: manifest_ctx.clone(),
                 index_options: IndexOptions::default(),
                 flush_semaphore: Arc::new(Semaphore::new(2)),
-                is_staging: false,
                 partition_expr: None,
                 durability_barrier: DurabilityBarrier::noop(),
             })
@@ -2473,7 +2463,6 @@ mod tests {
                 manifest_ctx: manifest_ctx.clone(),
                 index_options: IndexOptions::default(),
                 flush_semaphore: Arc::new(Semaphore::new(2)),
-                is_staging: false,
                 partition_expr: None,
                 durability_barrier: DurabilityBarrier::noop(),
             })

src/mito2/src/region.rs

@@ -96,7 +96,7 @@ impl RegionUsage {
 pub enum RegionLeaderState {
     /// The region is opened and is writable.
     Writable,
-    /// The region is in staging mode - writable but no checkpoint/compaction.
+    /// The region is in staging mode - writable but no checkpoint/compaction/flush.
     Staging,
     /// The region is entering staging mode. - write requests will be stalled.
     EnteringStaging,
@@ -407,6 +407,17 @@ impl MitoRegion {
         self.manifest_ctx.state.load() == RegionRoleState::Leader(RegionLeaderState::Staging)
     }
 
+    /// Returns whether the region is in staging mode, counting a staged edit in progress.
+    pub(crate) fn is_staging_or_staged_edit(&self) -> bool {
+        match self.manifest_ctx.state.load() {
+            RegionRoleState::Leader(RegionLeaderState::Staging) => true,
+            RegionRoleState::Leader(RegionLeaderState::Editing) => {
+                self.manifest_ctx.staging_partition_info().is_some()
+            }
+            _ => false,
+        }
+    }
+
     /// Returns whether the region is entering staging mode.
     pub(crate) fn is_enter_staging(&self) -> bool {
         self.manifest_ctx.state.load()
@@ -527,6 +538,29 @@ impl MitoRegion {
         )
     }
 
+    /// Leaves staging without committing it, abandoning a staged edit in progress.
+    ///
+    /// You should call this method in the worker loop, which is what restores the state of an
+    /// edit, and only when [Self::is_staging_or_staged_edit] holds.
+    pub(crate) fn abort_staging(&self) {
+        let aborted = self
+            .manifest_ctx
+            .state
+            .fetch_update(|state| {
+                matches!(
+                    state,
+                    RegionRoleState::Leader(
+                        RegionLeaderState::Staging | RegionLeaderState::Editing
+                    )
+                )
+                .then_some(RegionRoleState::Leader(RegionLeaderState::Writable))
+            })
+            .is_ok();
+        if aborted {
+            self.manifest_ctx.clear_staging_partition_info();
+        }
+    }
+
     /// Sets the region role state gracefully. This acquires the manifest write lock.
     pub(crate) async fn set_role_state_gracefully(
         &self,
@@ -1142,30 +1176,6 @@ impl MitoRegion {
         Ok(Some(pending))
     }
 
-    /// Returns the partition expression string for this region.
-    ///
-    /// If the region is currently in staging state, this returns the partition expression held in
-    /// the staging partition field. Otherwise, it returns the partition expression from the primary
-    /// region metadata (current committed version).
-    pub fn maybe_staging_partition_expr_str(&self) -> Option<String> {
-        let is_staging = self.is_staging();
-        if is_staging {
-            let staging_partition_info = self.manifest_ctx.staging_partition_info();
-            if staging_partition_info.is_none() {
-                warn!(
-                    "Staging partition expr is none for region {} in staging state",
-                    self.region_id
-                );
-            }
-            staging_partition_info
-                .as_ref()
-                .and_then(|info| info.partition_expr().map(ToString::to_string))
-        } else {
-            let version = self.version();
-            version.metadata.partition_expr.clone()
-        }
-    }
-
     pub fn expected_partition_expr_version(&self) -> u64 {
         if self.is_staging() {
             self.manifest_ctx

src/mito2/src/request.rs

@@ -1152,8 +1152,6 @@ pub(crate) struct FlushFinished {
     pub(crate) edit: RegionEdit,
     /// Memtables to remove.
     pub(crate) memtables_to_remove: SmallVec<[MemtableId; 2]>,
-    /// Whether the region is in staging mode.
-    pub(crate) is_staging: bool,
 }
 
 impl FlushFinished {
@@ -1372,6 +1370,9 @@ pub(crate) struct RegionEditRequest {
     pub(crate) preload_sst_cache: bool,
     /// The waiters that are waiting for this region edit's result.
     pub(crate) waiters: Waiters,
+    /// The partition expression of the staging this edit belongs to, if any.
+    /// The edit must not run unless the region is staging for that expression.
+    pub(crate) staging_partition_expr: Option<String>,
 }
 
 impl RegionEditRequest {
@@ -1386,8 +1387,15 @@ impl RegionEditRequest {
             edit,
             preload_sst_cache,
             waiters: Waiters::one(waiter),
+            staging_partition_expr: None,
         }
     }
+
+    /// Binds the edit to the staging for `partition_expr`.
+    pub(crate) fn for_staging(mut self, partition_expr: String) -> Self {
+        self.staging_partition_expr = Some(partition_expr);
+        self
+    }
 }
 
 /// Notifies the regin the result of editing region.

src/mito2/src/worker.rs

@@ -93,6 +93,7 @@ use crate::sst::index::puffin_manager::PuffinManagerFactory;
 use crate::temp_file_sweeper;
 use crate::time_provider::{StdTimeProvider, TimeProviderRef};
 use crate::wal::Wal;
+use crate::worker::handle_flush::StagingRegionsRef;
 use crate::worker::handle_manifest::RegionEditQueues;
 
 /// Identifier for a worker.
@@ -258,6 +259,7 @@ impl WorkerGroup {
         let sweep_limiter = (config.gc.max_concurrent_gc_job > 0)
             .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job)));
         let node_lease = Arc::new(NodeLease::default());
+        let staging_regions = StagingRegionsRef::default();
 
         let workers = (0..config.num_workers)
             .map(|id| {
@@ -290,6 +292,7 @@ impl WorkerGroup {
                     #[cfg(feature = "hdfs-object-store")]
                     sweep_limiter: sweep_limiter.clone(),
                     node_lease: node_lease.clone(),
+                    staging_regions: staging_regions.clone(),
                 }
                 .start()
             })
@@ -529,6 +532,7 @@ impl WorkerGroup {
         let sweep_limiter = (config.gc.max_concurrent_gc_job > 0)
             .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job)));
         let node_lease = Arc::new(NodeLease::default());
+        let staging_regions = StagingRegionsRef::default();
         let workers = (0..config.num_workers)
             .map(|id| {
                 WorkerStarter {
@@ -560,6 +564,7 @@ impl WorkerGroup {
                     #[cfg(feature = "hdfs-object-store")]
                     sweep_limiter: sweep_limiter.clone(),
                     node_lease: node_lease.clone(),
+                    staging_regions: staging_regions.clone(),
                 }
                 .start()
             })
@@ -677,6 +682,7 @@ struct WorkerStarter<S> {
     /// no concurrent GC job.
     sweep_limiter: Option<Arc<Semaphore>>,
     node_lease: Arc<NodeLease>,
+    staging_regions: StagingRegionsRef,
 }
 
 impl<S: LogStore> WorkerStarter<S> {
@@ -782,6 +788,7 @@ impl<S: LogStore> WorkerStarter<S> {
             #[cfg(feature = "hdfs-object-store")]
             sweep_limiter: self.sweep_limiter,
             node_lease: self.node_lease,
+            staging_regions: self.staging_regions,
         };
         let handle = common_runtime::spawn_global(async move {
             worker_thread.run().await;
@@ -1122,6 +1129,8 @@ struct RegionWorkerLoop<S> {
     /// no concurrent GC job.
     sweep_limiter: Option<Arc<Semaphore>>,
     node_lease: Arc<NodeLease>,
+    /// Staging regions of this node that have received writes.
+    staging_regions: StagingRegionsRef,
 }
 
 impl<S: LogStore> RegionWorkerLoop<S> {

src/mito2/src/worker/handle_apply_staging.rs

@@ -140,7 +140,8 @@ impl<S: LogStore> RegionWorkerLoop<S> {
                         region_id, edit,
                         // we don't need to preload sst cache during repartition, as it may cause extra network overhead.
                         false, tx,
-                    ),
+                    )
+                    .for_staging(request.partition_expr),
                 )))
                 .await;
 

src/mito2/src/worker/handle_close.rs

@@ -36,6 +36,16 @@ impl<S: LogStore> RegionWorkerLoop<S> {
 
         info!("Try to close region {}, worker: {}", region_id, self.id);
 
+        if region.is_staging_or_staged_edit() {
+            // A staging region cannot flush, so it leaves staging before the close-time flush
+            // below. The manifest lock lets a successful end of staging in progress finish first,
+            // after which the region is already out of staging.
+            let _manager = region.manifest_ctx.manifest_manager.write().await;
+            region.abort_staging();
+            // Edits queued behind a staged edit belong to the staging that has just ended.
+            self.reject_region_edit_queue_as_not_found(region_id);
+        }
+
         // Memtables containing skip-WAL writes must be flushed before closing.
         // Followers cannot flush; explicit flush-on-close and region WAL options
         // retain their existing behavior.

src/mito2/src/worker/handle_flush.rs

@@ -14,9 +14,9 @@
 
 //! Handling flush related requests.
 
-use std::collections::HashSet;
-use std::sync::Arc;
+use std::collections::{HashMap, HashSet};
 use std::sync::atomic::Ordering;
+use std::sync::{Arc, Mutex, Weak};
 
 use common_base::readable_size::ReadableSize;
 use common_telemetry::{debug, error, info};
@@ -25,10 +25,10 @@ use store_api::region_request::{RegionFlushReason, RegionFlushRequest};
 use store_api::storage::RegionId;
 
 use crate::config::{IndexBuildMode, MitoConfig};
-use crate::error::{RegionNotFoundSnafu, Result};
+use crate::error::{RegionBusySnafu, RegionNotFoundSnafu, Result};
 use crate::flush::{FlushReason, RegionFlushTask};
-use crate::region::MitoRegionRef;
 use crate::region::version::VersionRef;
+use crate::region::{MitoRegion, MitoRegionRef};
 use crate::request::{BuildIndexRequest, FlushFailed, FlushFinished, OnFailure, OptionOutputTx};
 use crate::sst::index::IndexBuildType;
 use crate::worker::RegionWorkerLoop;
@@ -39,6 +39,33 @@ pub(crate) struct RegionWriteBufferPressure {
     pub(crate) rejected_region_ids: HashSet<RegionId>,
 }
 
+/// Staging regions of a node that have received writes.
+///
+/// Workers share it to bound the memtables that stay unflushed until staging ends.
+#[derive(Debug, Default)]
+pub(crate) struct StagingRegions {
+    regions: Mutex<HashMap<RegionId, Weak<MitoRegion>>>,
+}
+
+pub(crate) type StagingRegionsRef = Arc<StagingRegions>;
+
+impl StagingRegions {
+    /// Records `region` and returns the memtable usage of the regions still in staging.
+    fn record_and_memtable_usage(&self, region: &MitoRegionRef) -> usize {
+        let mut regions = self.regions.lock().unwrap();
+        regions.insert(region.region_id, Arc::downgrade(region));
+        let mut usage = 0;
+        regions.retain(|_, region| match region.upgrade() {
+            Some(region) if region.is_staging_or_staged_edit() => {
+                usage += region_memtable_usage(&region.version());
+                true
+            }
+            _ => false,
+        });
+        usage
+    }
+}
+
 #[derive(Debug, Default, PartialEq, Eq)]
 struct RegionWriteBufferStatus {
     should_flush: bool,
@@ -100,8 +127,11 @@ impl<S: LogStore> RegionWorkerLoop<S> {
         let mut pending_regions = vec![];
 
         for region in &regions {
-            if self.flush_scheduler.is_flush_requested(region.region_id) || !region.is_writable() {
-                // Already flushing or not writable.
+            if self.flush_scheduler.is_flush_requested(region.region_id)
+                || !region.is_writable()
+                || region.is_staging()
+            {
+                // Already flushing, not writable, or in staging, whose rows must stay in the WAL.
                 continue;
             }
 
@@ -193,6 +223,19 @@ impl<S: LogStore> RegionWorkerLoop<S> {
                 self.stalled_requests.estimated_size(&region_id),
             );
 
+            if region.is_staging() {
+                // Staging regions don't flush, so waiting would not release their memory. Writes
+                // are rejected at twice the region limit or at half of the engine's flush limit,
+                // which reserves memory headroom for other regions.
+                let staging_usage = self.staging_regions.record_and_memtable_usage(&region);
+                if status.should_reject
+                    || staging_usage >= self.write_buffer_manager.flush_limit() / 2
+                {
+                    pressure.rejected_region_ids.insert(region_id);
+                }
+                continue;
+            }
+
             if status.should_flush && !self.flush_scheduler.is_flush_requested(region.region_id) {
                 let task = self.new_flush_task(
                     &region,
@@ -242,8 +285,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             manifest_ctx: region.manifest_ctx.clone(),
             index_options: region.version().options.index_options.clone(),
             flush_semaphore: self.flush_semaphore.clone(),
-            is_staging: region.is_staging(),
-            partition_expr: region.maybe_staging_partition_expr_str(),
+            partition_expr: region.metadata().partition_expr.clone(),
             durability_barrier: self
                 .wal
                 .durability_barrier(region.region_id, &region.provider),
@@ -327,6 +369,18 @@ impl<S: LogStore> RegionWorkerLoop<S> {
         // when handling flush request instead of in `schedule_flush` or `flush_finished`.
         self.update_topic_latest_entry_id(&region);
 
+        // A staging region must keep its WAL until staging ends, so it flushes nothing.
+        if region.is_staging() && !region.version().memtables.is_empty() {
+            sender.send(match request.reason {
+                // These only ask to advance the flushed position, which can wait.
+                Some(RegionFlushReason::RegionMigration | RegionFlushReason::RemoteWalPrune) => {
+                    Ok(0)
+                }
+                _ => RegionBusySnafu { region_id }.fail(),
+            });
+            return;
+        }
+
         let reason = resolve_flush_reason(request.reason, region.is_downgrading());
         let mut task =
             self.new_flush_task(&region, reason, request.row_group_size, self.config.clone());
@@ -350,6 +404,12 @@ impl<S: LogStore> RegionWorkerLoop<S> {
                 continue;
             }
             self.update_topic_latest_entry_id(region);
+            // The rows of a staging region must stay in the WAL until staging ends, so a staging
+            // region with rows pins its topic's prune point: the update above only keeps the id
+            // fresh while the memtable is empty.
+            if region.is_staging() {
+                continue;
+            }
 
             let auto_flush_interval = region
                 .version()
@@ -399,27 +459,14 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             }
         };
 
-        if request.is_staging {
-            // Skip the region metadata update.
-            info!(
-                "Skipping region metadata update for region {} in staging mode",
-                region_id
-            );
-            region.version_control.apply_edit(
-                None,
-                &request.memtables_to_remove,
-                region.file_purger.clone(),
-            );
-        } else {
-            region.version_control.apply_edit(
-                Some(request.edit.clone()),
-                &request.memtables_to_remove,
-                region.file_purger.clone(),
-            );
-        }
+        region.version_control.apply_edit(
+            Some(request.edit.clone()),
+            &request.memtables_to_remove,
+            region.file_purger.clone(),
+        );
 
-        // The flush has committed its manifest, including in staging mode. Count only
-        // its Parquet output, before WAL cleanup (which can fail after the commit).
+        // The flush has committed its manifest. Count only its Parquet output, before
+        // WAL cleanup (which can fail after the commit).
         let ingested_parquet_bytes = request
             .edit
             .files_to_add

src/mito2/src/worker/handle_manifest.rs

@@ -30,7 +30,9 @@ use store_api::storage::{RegionId, SequenceNumber};
 use crate::cache::CacheManagerRef;
 use crate::cache::file_cache::{FileType, IndexKey};
 use crate::config::IndexBuildMode;
-use crate::error::{EditRegionSnafu, RegionBusySnafu, RegionNotFoundSnafu, Result};
+use crate::error::{
+    EditRegionSnafu, RegionBusySnafu, RegionNotFoundSnafu, RegionStateSnafu, Result,
+};
 use crate::manifest::action::{
     RegionChange, RegionEdit, RegionMetaAction, RegionMetaActionList, RegionTruncate, TruncateKind,
 };
@@ -310,6 +312,36 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             return RegionEditDispatch::Rejected;
         };
 
+        // A close-time flush ends with the region removed. A close that left staging keeps the
+        // region writable until then, so an edit meant for staging would run as a normal edit.
+        if self.flush_scheduler.has_close_time_flush(region_id) {
+            request
+                .waiters
+                .reply_with(|| RegionBusySnafu { region_id }.fail());
+            return RegionEditDispatch::Rejected;
+        }
+
+        // A close may have abandoned the staging this edit belongs to, even if its flush failed.
+        if let Some(expected) = &request.staging_partition_expr
+            && region
+                .manifest_ctx
+                .staging_partition_info()
+                .as_ref()
+                .and_then(|info| info.partition_expr())
+                != Some(expected.as_str())
+        {
+            let state = region.state();
+            request.waiters.reply_with(|| {
+                RegionStateSnafu {
+                    region_id,
+                    state,
+                    expect: RegionRoleState::Leader(RegionLeaderState::Staging),
+                }
+                .fail()
+            });
+            return RegionEditDispatch::Rejected;
+        }
+
         if !region.is_writable() || self.flush_scheduler.has_pending_alter(region_id) {
             if matches!(
                 region.state(),
@@ -334,6 +366,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             mut edit,
             waiters,
             preload_sst_cache,
+            staging_partition_expr: _,
         } = request;
         let file_sequence = region.version_control.current().last_sequence + 1;
         edit.committed_sequence = Some(file_sequence);
@@ -427,7 +460,10 @@ impl<S: LogStore> RegionWorkerLoop<S> {
         };
 
         let need_compaction = if edit_result.is_staging {
-            if edit_result.update_region_state {
+            // A close may have abandoned this edit and left staging already.
+            if edit_result.update_region_state
+                && region.state() == RegionRoleState::Leader(RegionLeaderState::Editing)
+            {
                 // For staging regions, edits are not applied immediately,
                 // as they remain invisible until the region exits the staging state.
                 region.switch_state_to_staging(RegionLeaderState::Editing);

src/store-api/src/region_engine.rs

@@ -50,7 +50,7 @@ pub enum SettableRegionRoleState {
     DowngradingLeader,
     /// Exit staging mode and return to normal leader state. Only allowed from staging state.
     Leader,
-    /// Enter staging mode. Region remains writable but disables checkpoint and compaction. Only allowed from normal leader state.
+    /// Enter staging mode. Region remains writable but disables checkpoint, compaction and flush. Only allowed from normal leader state.
     StagingLeader,
 }
 
@@ -606,7 +606,7 @@ pub struct RegionStatistic {
     /// Total ingested Parquet bytes since the region opened.
     /// Currently counts only committed flush output.
     ///
-    /// Includes staging flushes, but excludes index files, compaction and imported SSTs.
+    /// Excludes index files, compaction and imported SSTs.
     /// Resets when the region instance is recreated; manifest replay does not increment it.
     /// Accepts the legacy `flushed_parquet_bytes` key; upgrade metasrv before datanodes.
     #[serde(default, alias = "flushed_parquet_bytes")]