From 182cce4cc2d8b7dc0619144f841761dd5935b8d8 Mon Sep 17 00:00:00 2001
From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com>
Date: Tue, 11 Nov 2025 14:16:23 +0800
Subject: [PATCH] fix(mito): allow region edit in writable state (#7201)

* fix/region-expire-state:
 Refactor region state handling in compaction task and manifest updates

 - Introduce a variable to hold the current region state for clarity in compaction task updates.
 - Add an expected_region_state field to RegionEditResult to manage region state expectations during manifest handling.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

* fix/region-expire-state:
 Refactor region state handling in compaction task

 - Replace direct assignment of `RegionLeaderState::Writable` with dynamic state retrieval and conditional check for leader state.
 - Modify `RegionEditResult` to include a flag `update_region_state` instead of `expected_region_state` to indicate if the region state should be updated to writable.
 - Adjust handling of `RegionEditResult` in `handle_manifest` to conditionally update region state based on the new flag.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

---------

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>
---
 src/mito2/src/compaction/task.rs        | 16 +++++++++++++---
 src/mito2/src/request.rs                |  2 ++
 src/mito2/src/worker/handle_manifest.rs |  8 ++++++--
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/src/mito2/src/compaction/task.rs b/src/mito2/src/compaction/task.rs
index b54d4c7c62..8488c9af9e 100644
--- a/src/mito2/src/compaction/task.rs
+++ b/src/mito2/src/compaction/task.rs
@@ -26,7 +26,7 @@ use crate::compaction::picker::{CompactionTask, PickerOutput};
 use crate::error::CompactRegionSnafu;
 use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList};
 use crate::metrics::{COMPACTION_FAILURE_COUNT, COMPACTION_STAGE_ELAPSED};
-use crate::region::RegionLeaderState;
+use crate::region::RegionRoleState;
 use crate::request::{
     BackgroundNotify, CompactionFailed, CompactionFinished, OutputTx, RegionEditResult,
     WorkerRequest, WorkerRequestWithTime,
@@ -106,12 +106,21 @@ impl CompactionTaskImpl {
 
         // 1. Update manifest
         let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone()));
+        let RegionRoleState::Leader(current_region_state) =
+            compaction_region.manifest_ctx.current_state()
+        else {
+            warn!(
+                "Region {} not in leader state, skip removing expired files",
+                region_id
+            );
+            return;
+        };
         if let Err(e) = compaction_region
             .manifest_ctx
-            .update_manifest(RegionLeaderState::Writable, action_list)
+            .update_manifest(current_region_state, action_list)
             .await
         {
-            error!(
+            warn!(
                 e;
                 "Failed to update manifest for expired files removal, region: {region_id}, files: [{expired_files_str}]. Compaction will continue."
             );
@@ -126,6 +135,7 @@ impl CompactionTaskImpl {
                 sender: expire_delete_sender,
                 edit,
                 result: Ok(()),
+                update_region_state: false,
             }),
         })
         .await;
diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs
index 606fbbb7ea..c0a086177d 100644
--- a/src/mito2/src/request.rs
+++ b/src/mito2/src/request.rs
@@ -969,6 +969,8 @@ pub(crate) struct RegionEditResult {
     pub(crate) edit: RegionEdit,
     /// Result from the manifest manager.
     pub(crate) result: Result<()>,
+    /// Whether region state need to be set to Writable after handling this request.
+    pub(crate) update_region_state: bool,
 }
 
 #[derive(Debug)]
diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs
index be4bde6583..88c4b314b1 100644
--- a/src/mito2/src/worker/handle_manifest.rs
+++ b/src/mito2/src/worker/handle_manifest.rs
@@ -249,6 +249,8 @@ impl<S> RegionWorkerLoop<S> {
                     sender,
                     edit,
                     result,
+                    // we always need to restore region state after region edit
+                    update_region_state: true,
                 }),
             };
 
@@ -292,8 +294,10 @@ impl<S> RegionWorkerLoop<S> {
             );
         }
 
-        // Sets the region as writable.
-        region.switch_state_to_writable(RegionLeaderState::Editing);
+        if edit_result.update_region_state {
+            // Sets the region as writable.
+            region.switch_state_to_writable(RegionLeaderState::Editing);
+        }
 
         let _ = edit_result.sender.send(edit_result.result);