fix(mito2): fence writes before a graceful downgrade (#9716)

* test(mito2): cover graceful role changes racing with writes

A graceful role change flips the region state in a background task while
the worker keeps handling writes. Add bulk and row write tests that check a
write handled after the role change request is either rejected or covered
by the returned last_entry_id.

Add tests for DowngradingLeader while an ALTER keeps the region in
Altering, alone and in the downgrade sequence the datanode runs
(DowngradingLeader, Flush, Follower) with a migration candidate on its own
local WAL. Add a test that the downgrade sequence covers a write in flight
when a lease renewal sets Downgrading.

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

* fix(mito2): fence writes before a graceful downgrade

set_region_role_state_gracefully hands the role change to a background task
that waits for the manifest lock. The worker keeps handling writes until the
task flips the state, so a write could be acknowledged with an entry id above
the returned last_entry_id. Set Downgrading in the worker before spawning the
task for Follower and DowngradingLeader requests, as lease renewal already
does, so every write the worker handles afterwards is rejected.

DowngradingLeader also did nothing and returned success when the region was
in another leader state such as Altering. The region became writable again
when the operation finished, and the datanode then flushed and turned it
into a follower in one step. Return a retryable error instead, so metasrv
retries the downgrade.

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

* test(mito2): assert the retryable downgrade error during ALTER

Check that DowngradingLeader during an ALTER returns a retryable error, and
release the paused ALTER on that path. The local WAL test now retries the
downgrade after the ALTER, as metasrv does.

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

* test(mito2): write before the retried downgrade in the local WAL test

Rows written before an ALTER were already flushed, so the candidate check
did not depend on the downgrade's final flush. Write rows after the ALTER
and before the downgrade is retried, which only that flush can carry to the
candidate.

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

---------

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

src/mito2/src/engine/basic_test.rs

@@ -1882,7 +1882,10 @@ async fn restart_without_region_close(
 
 /// Builds a bulk insert request for a region created by [CreateRequestBuilder]
 /// with its default columns `(tag_0, field_0, ts)`.
-fn build_bulk_insert_request(region_id: RegionId, num_rows: usize) -> RegionBulkInsertsRequest {
+pub(crate) fn build_bulk_insert_request(
+    region_id: RegionId,
+    num_rows: usize,
+) -> RegionBulkInsertsRequest {
     let schema = Arc::new(Schema::new(vec![
         Field::new("tag_0", DataType::Utf8, true),
         Field::new("field_0", DataType::Float64, true),

src/mito2/src/engine/object_store_wal_test.rs

@@ -79,13 +79,13 @@ fn open_request(table_dir: String) -> RegionOpenRequest {
     }
 }
 
-async fn new_engine(env: &mut TestEnv) -> (MitoEngine, Arc<ObjectStoreLogStore>) {
+pub(crate) async fn new_engine(env: &mut TestEnv) -> (MitoEngine, Arc<ObjectStoreLogStore>) {
     env.create_engine_with_object_store_wal(MitoConfig::default(), &wal_config())
         .await
 }
 
 /// Creates a region with object store WAL options and returns its table dir and row schema.
-async fn create_region(
+pub(crate) async fn create_region(
     engine: &MitoEngine,
     region_id: RegionId,
 ) -> (String, Vec<api::v1::ColumnSchema>) {

src/mito2/src/engine/set_role_state_test.rs

@@ -12,23 +12,36 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::sync::Arc;
 use std::time::Duration;
 
-use api::v1::Rows;
-use common_error::ext::ErrorExt;
+use api::v1::{Rows, SemanticType};
+use common_error::ext::{ErrorExt, RetryHint};
 use common_error::status_code::StatusCode;
+use common_recordbatch::RecordBatches;
+use datatypes::prelude::ConcreteDataType;
+use datatypes::schema::ColumnSchema;
+use store_api::metadata::ColumnMetadata;
 use store_api::region_engine::{
     RegionEngine, RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess,
     SettableRegionRoleState,
 };
 use store_api::region_request::{
-    EnterStagingRequest, RegionFlushRequest, RegionPutRequest, RegionRequest,
+    AddColumn, AffectedRows, AlterKind, EnterStagingRequest, PathType, RegionAlterRequest,
+    RegionCatchupRequest, RegionFlushRequest, RegionOpenRequest, RegionPutRequest, RegionRequest,
     StagingPartitionDirective,
 };
-use store_api::storage::RegionId;
+use store_api::storage::{RegionId, ScanRequest};
+use tokio::sync::oneshot;
 
 use crate::config::MitoConfig;
-use crate::region::{RegionLeaderState, RegionRoleState};
+use crate::engine::MitoEngine;
+use crate::engine::basic_test::build_bulk_insert_request;
+use crate::engine::listener::NotifyRegionChangeResultListener;
+use crate::engine::object_store_wal_test;
+use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState};
+use crate::region_write_ctx::test_hooks::arm_bulk_install_barrier;
+use crate::request::WorkerRequest;
 use crate::test_util::{
     CheckpointTaskBlocker, CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema,
 };
@@ -1001,3 +1014,520 @@ async fn test_restricted_state_transitions_with_format(flat_format: bool) {
         .unwrap();
     assert_success_response(&result, 0);
 }
+
+const WAIT: Duration = Duration::from_secs(10);
+
+type WriteReceiver = oneshot::Receiver<crate::error::Result<AffectedRows>>;
+
+/// Submits a graceful role change straight to the region's worker. Requests
+/// submitted after this returns reach the worker after it.
+async fn submit_role_change(
+    engine: &MitoEngine,
+    region_id: RegionId,
+    state: SettableRegionRoleState,
+) -> oneshot::Receiver<SetRegionRoleStateResponse> {
+    let (request, receiver) = WorkerRequest::new_set_readonly_gracefully(region_id, state);
+    engine
+        .inner
+        .workers
+        .submit_to_worker(region_id, request)
+        .await
+        .unwrap();
+    receiver
+}
+
+/// Submits a region request straight to the region's worker.
+async fn submit_request(
+    engine: &MitoEngine,
+    region_id: RegionId,
+    request: RegionRequest,
+) -> WriteReceiver {
+    let metadata = engine.inner.get_metadata(region_id).ok();
+    let (request, receiver) =
+        WorkerRequest::try_from_region_request(region_id, request, metadata).unwrap();
+    engine
+        .inner
+        .workers
+        .submit_to_worker(region_id, request)
+        .await
+        .unwrap();
+    receiver
+}
+
+fn put_request(rows: Rows) -> RegionRequest {
+    RegionRequest::Put(RegionPutRequest {
+        skip_wal: false,
+        rows,
+        hint: None,
+        partition_expr_version: None,
+    })
+}
+
+fn reported_last_entry_id(response: &SetRegionRoleStateResponse) -> u64 {
+    match response {
+        SetRegionRoleStateResponse::Success(success) => success.last_entry_id().unwrap(),
+        _ => panic!("Expected success response, got: {:?}", response),
+    }
+}
+
+/// Lets the role change task waiting behind `manifest_guard` run and returns
+/// the `last_entry_id` it reports.
+async fn finish_role_change<G>(
+    manifest_guard: G,
+    role_change: oneshot::Receiver<SetRegionRoleStateResponse>,
+) -> u64 {
+    drop(manifest_guard);
+    let response = tokio::time::timeout(WAIT, role_change)
+        .await
+        .expect("role change did not finish")
+        .unwrap();
+    reported_last_entry_id(&response)
+}
+
+/// Waits for the write's result if it is still in flight.
+async fn write_result(
+    early: Option<crate::error::Result<AffectedRows>>,
+    write: WriteReceiver,
+) -> crate::error::Result<AffectedRows> {
+    match early {
+        Some(result) => result,
+        None => tokio::time::timeout(WAIT, write)
+            .await
+            .expect("write did not finish")
+            .unwrap(),
+    }
+}
+
+/// Checks the promise of `set_region_role_state_gracefully`: an acknowledged
+/// write must be covered by the `last_entry_id` the call returned.
+fn assert_write_covered(
+    region: &MitoRegionRef,
+    result: crate::error::Result<AffectedRows>,
+    reported: u64,
+) {
+    let last_entry_id = region.version_control.current().last_entry_id;
+    assert!(
+        result.is_err() || last_entry_id <= reported,
+        "a write was acknowledged ({result:?}) after the role change returned last_entry_id {reported}, \
+         but the region's last_entry_id is now {last_entry_id}"
+    );
+}
+
+/// The role change task waits for the manifest lock before it changes the
+/// role, while the worker keeps handling requests. A bulk write handled after
+/// the role change request must be rejected, or be covered by the returned
+/// `last_entry_id` if it was accepted.
+#[tokio::test(flavor = "multi_thread")]
+async fn test_set_follower_gracefully_covers_inflight_bulk_write() {
+    common_telemetry::init_default_ut_logging();
+    let mut env = TestEnv::with_prefix("follower-inflight-bulk-write").await;
+    let engine = env.create_engine(MitoConfig::default()).await;
+    let region_id = RegionId::new(1, 1);
+    let request = CreateRequestBuilder::new().build();
+    let column_schemas = rows_schema(&request);
+    engine
+        .handle_request(region_id, RegionRequest::Create(request))
+        .await
+        .unwrap();
+    put_rows(
+        &engine,
+        region_id,
+        Rows {
+            schema: column_schemas,
+            rows: build_rows(0, 3),
+        },
+    )
+    .await;
+
+    let region = engine.get_region(region_id).unwrap();
+    let manifest_guard = region.manifest_ctx.manifest_manager.read().await;
+    let role_change =
+        submit_role_change(&engine, region_id, SettableRegionRoleState::Follower).await;
+    // An accepted write pauses after its WAL append, before installing its rows.
+    let mut barrier = arm_bulk_install_barrier(region_id, region.version_control.clone());
+    let mut write = submit_request(
+        &engine,
+        region_id,
+        RegionRequest::BulkInserts(build_bulk_insert_request(region_id, 2)),
+    )
+    .await;
+    let early = tokio::select! {
+        _ = barrier.wait_until_reached() => None,
+        result = &mut write => Some(result.unwrap()),
+        _ = tokio::time::sleep(WAIT) => panic!("bulk write never reached the install barrier"),
+    };
+    let reported = finish_role_change(manifest_guard, role_change).await;
+    assert_eq!(RegionRoleState::Follower, region.state());
+
+    barrier.release();
+    let result = write_result(early, write).await;
+    assert_write_covered(&region, result, reported);
+}
+
+/// Same as [test_set_follower_gracefully_covers_inflight_bulk_write] for a row
+/// write. An accepted put waits in its WAL append: the object store WAL
+/// acknowledges an append once its object exists, and the test parks that
+/// object.
+#[tokio::test(flavor = "multi_thread")]
+async fn test_set_follower_gracefully_covers_inflight_row_write() {
+    common_telemetry::init_default_ut_logging();
+    let mut env = TestEnv::with_prefix("follower-inflight-row-write").await;
+    let (engine, store) = object_store_wal_test::new_engine(&mut env).await;
+    let region_id = RegionId::new(1, 1);
+    let (_, schema) = object_store_wal_test::create_region(&engine, region_id).await;
+    put_rows(
+        &engine,
+        region_id,
+        Rows {
+            schema: schema.clone(),
+            rows: build_rows(0, 3),
+        },
+    )
+    .await;
+
+    let region = engine.get_region(region_id).unwrap();
+    let manifest_guard = region.manifest_ctx.manifest_manager.read().await;
+    let role_change =
+        submit_role_change(&engine, region_id, SettableRegionRoleState::Follower).await;
+    store.hold_creates();
+    let mut write = submit_request(
+        &engine,
+        region_id,
+        put_request(Rows {
+            schema,
+            rows: build_rows(3, 5),
+        }),
+    )
+    .await;
+    let early = tokio::select! {
+        result = store.wait_for_parked_creates(1) => {
+            result.unwrap();
+            None
+        }
+        result = &mut write => Some(result.unwrap()),
+        _ = tokio::time::sleep(WAIT) => panic!("put never reached the WAL"),
+    };
+    let reported = finish_role_change(manifest_guard, role_change).await;
+    assert_eq!(RegionRoleState::Follower, region.state());
+
+    store.release_creates();
+    let result = write_result(early, write).await;
+    assert_write_covered(&region, result, reported);
+}
+
+/// A lease renewal sets Downgrading outside the worker, so a write that passed
+/// the state check before it can still be in flight. The datanode's downgrade
+/// steps (DowngradingLeader, Flush, Follower) must still flush that write and
+/// report its entry.
+#[tokio::test(flavor = "multi_thread")]
+async fn test_downgrade_flush_follower_covers_inflight_write() {
+    common_telemetry::init_default_ut_logging();
+    let mut env = TestEnv::with_prefix("downgrade-flush-follower").await;
+    let engine = env.create_engine(MitoConfig::default()).await;
+    let region_id = RegionId::new(1, 1);
+    engine
+        .handle_request(
+            region_id,
+            RegionRequest::Create(CreateRequestBuilder::new().build()),
+        )
+        .await
+        .unwrap();
+
+    let region = engine.get_region(region_id).unwrap();
+    let mut barrier = arm_bulk_install_barrier(region_id, region.version_control.clone());
+    let write = submit_request(
+        &engine,
+        region_id,
+        RegionRequest::BulkInserts(build_bulk_insert_request(region_id, 2)),
+    )
+    .await;
+    tokio::time::timeout(WAIT, barrier.wait_until_reached())
+        .await
+        .expect("bulk write never reached the install barrier");
+    engine
+        .set_region_role(region_id, RegionRole::DowngradingLeader)
+        .unwrap();
+
+    let downgrade = submit_role_change(
+        &engine,
+        region_id,
+        SettableRegionRoleState::DowngradingLeader,
+    )
+    .await;
+    barrier.release();
+    assert_eq!(2, write.await.unwrap().unwrap());
+    reported_last_entry_id(&downgrade.await.unwrap());
+    engine
+        .handle_request(
+            region_id,
+            RegionRequest::Flush(RegionFlushRequest::default()),
+        )
+        .await
+        .unwrap();
+
+    let response = engine
+        .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower)
+        .await
+        .unwrap();
+    let reported = reported_last_entry_id(&response);
+    let current = region.version_control.current();
+    assert_eq!(current.last_entry_id, reported);
+    assert_eq!(current.version.flushed_entry_id, reported);
+    assert!(current.version.memtables.is_empty());
+}
+
+fn add_nullable_field1() -> RegionAlterRequest {
+    RegionAlterRequest {
+        kind: AlterKind::AddColumns {
+            columns: vec![AddColumn {
+                column_metadata: ColumnMetadata {
+                    column_schema: ColumnSchema::new(
+                        "field_1",
+                        ConcreteDataType::float64_datatype(),
+                        true,
+                    ),
+                    semantic_type: SemanticType::Field,
+                    column_id: 3,
+                },
+                location: None,
+            }],
+        },
+    }
+}
+
+async fn wait_for_state(region: &MitoRegionRef, state: RegionRoleState) {
+    tokio::time::timeout(WAIT, async {
+        while region.state() != state {
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+    })
+    .await
+    .unwrap_or_else(|_| panic!("region never reached {state:?}, now {:?}", region.state()));
+}
+
+async fn count_rows(engine: &MitoEngine, region_id: RegionId) -> usize {
+    let stream = engine
+        .scan_to_stream(region_id, ScanRequest::default())
+        .await
+        .unwrap();
+    let batches = RecordBatches::try_collect(stream).await.unwrap();
+    batches.iter().map(|batch| batch.num_rows()).sum()
+}
+
+/// Creates a region with rows `[0, 3)` in an SST and starts an ALTER whose
+/// result the worker does not handle until `listener` is woken, so the region
+/// stays in Altering. Returns the ALTER task.
+async fn create_region_and_start_alter(
+    env: &TestEnv,
+    engine: &MitoEngine,
+    region_id: RegionId,
+    request: store_api::region_request::RegionCreateRequest,
+) -> tokio::task::JoinHandle<()> {
+    env.get_schema_metadata_manager()
+        .register_region_table_info(
+            region_id.table_id(),
+            "test_table",
+            "test_catalog",
+            "test_schema",
+            None,
+            env.get_kv_backend(),
+        )
+        .await;
+    let column_schemas = rows_schema(&request);
+    engine
+        .handle_request(region_id, RegionRequest::Create(request))
+        .await
+        .unwrap();
+    put_rows(
+        engine,
+        region_id,
+        Rows {
+            schema: column_schemas,
+            rows: build_rows(0, 3),
+        },
+    )
+    .await;
+    engine
+        .handle_request(
+            region_id,
+            RegionRequest::Flush(RegionFlushRequest::default()),
+        )
+        .await
+        .unwrap();
+
+    let cloned_engine = engine.clone();
+    let alter = tokio::spawn(async move {
+        cloned_engine
+            .handle_request(region_id, RegionRequest::Alter(add_nullable_field1()))
+            .await
+            .unwrap();
+    });
+    let region = engine.get_region(region_id).unwrap();
+    wait_for_state(
+        &region,
+        RegionRoleState::Leader(RegionLeaderState::Altering),
+    )
+    .await;
+    alter
+}
+
+/// DowngradingLeader while an ALTER keeps the region in Altering must fail
+/// with a retryable error, because the region becomes writable again once the
+/// ALTER finishes. A retry after the ALTER fences writes.
+#[tokio::test(flavor = "multi_thread")]
+async fn test_downgrading_leader_during_alter_stops_writes() {
+    common_telemetry::init_default_ut_logging();
+    let mut env = TestEnv::with_prefix("downgrading-during-alter").await;
+    let listener = Arc::new(NotifyRegionChangeResultListener::default());
+    let engine = env
+        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
+        .await;
+    let region_id = RegionId::new(1, 1);
+    let request = CreateRequestBuilder::new().build();
+    let column_schemas = rows_schema(&request);
+    let alter = create_region_and_start_alter(&env, &engine, region_id, request).await;
+    let region = engine.get_region(region_id).unwrap();
+
+    let response = engine
+        .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader)
+        .await
+        .unwrap();
+    let state_after_downgrade = region.state();
+    listener.wake_notify();
+    tokio::time::timeout(WAIT, alter).await.unwrap().unwrap();
+    match &response {
+        SetRegionRoleStateResponse::InvalidTransition(err) => {
+            assert_eq!(RetryHint::Retryable, err.retry_hint(), "{err:?}")
+        }
+        _ => panic!(
+            "DowngradingLeader returned {response:?} with the region in {state_after_downgrade:?}"
+        ),
+    }
+
+    let response = engine
+        .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader)
+        .await
+        .unwrap();
+    reported_last_entry_id(&response);
+    let err = engine
+        .handle_request(
+            region_id,
+            put_request(Rows {
+                schema: column_schemas,
+                rows: build_rows(3, 5),
+            }),
+        )
+        .await
+        .unwrap_err();
+    assert_eq!(StatusCode::RegionNotReady, err.status_code());
+}
+
+/// The datanode's downgrade steps (DowngradingLeader, Flush, Follower) with
+/// the DowngradingLeader step overlapping an ALTER. The region accepts writes
+/// until a downgrade takes effect. A migration candidate opens the region from
+/// the object store with its own local WAL, so the final flush must carry
+/// every write the old leader acknowledged.
+#[tokio::test(flavor = "multi_thread")]
+async fn test_downgrade_during_alter_keeps_acknowledged_writes_with_local_wal() {
+    common_telemetry::init_default_ut_logging();
+    let mut env = TestEnv::with_prefix("downgrade-during-alter-local-wal").await;
+    let listener = Arc::new(NotifyRegionChangeResultListener::default());
+    let engine = env
+        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
+        .await;
+    let region_id = RegionId::new(1, 1);
+    let request = CreateRequestBuilder::new().build();
+    let column_schemas = rows_schema(&request);
+    let table_dir = request.table_dir.clone();
+    let options = request.options.clone();
+    let alter = create_region_and_start_alter(&env, &engine, region_id, request).await;
+
+    // Mirrors `DowngradeRegionsHandler::handle_downgrade_region` in the
+    // datanode. Metasrv retries the downgrade on a retryable error.
+    let mut response = engine
+        .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader)
+        .await
+        .unwrap();
+    listener.wake_notify();
+    tokio::time::timeout(WAIT, alter).await.unwrap().unwrap();
+    put_rows(
+        &engine,
+        region_id,
+        Rows {
+            schema: column_schemas.clone(),
+            rows: build_rows(3, 5),
+        },
+    )
+    .await;
+    let mut acknowledged_rows = 5;
+    if let SetRegionRoleStateResponse::InvalidTransition(err) = &response {
+        assert_eq!(RetryHint::Retryable, err.retry_hint(), "{err:?}");
+        response = engine
+            .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader)
+            .await
+            .unwrap();
+    }
+    reported_last_entry_id(&response);
+    engine
+        .handle_request(
+            region_id,
+            RegionRequest::Flush(RegionFlushRequest::default()),
+        )
+        .await
+        .unwrap();
+    if let Ok(output) = engine
+        .handle_request(
+            region_id,
+            put_request(Rows {
+                schema: column_schemas,
+                rows: build_rows(5, 8),
+            }),
+        )
+        .await
+    {
+        acknowledged_rows += output.affected_rows;
+    }
+    let response = engine
+        .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower)
+        .await
+        .unwrap();
+    let reported = reported_last_entry_id(&response);
+
+    // The candidate shares the object store but not the local WAL.
+    let mut candidate_env = TestEnv::with_prefix("downgrade-during-alter-candidate").await;
+    candidate_env.object_store_manager = env.get_object_store_manager();
+    candidate_env.log_store = Some(candidate_env.create_log_store().await);
+    let candidate = candidate_env.new_mito_engine(MitoConfig::default()).await;
+    candidate
+        .handle_request(
+            region_id,
+            RegionRequest::Open(RegionOpenRequest {
+                engine: String::new(),
+                table_dir,
+                options,
+                skip_wal_replay: true,
+                path_type: PathType::Bare,
+                checkpoint: None,
+                requirements: Default::default(),
+            }),
+        )
+        .await
+        .unwrap();
+    candidate
+        .handle_request(
+            region_id,
+            RegionRequest::Catchup(RegionCatchupRequest {
+                set_writable: true,
+                entry_id: Some(reported),
+                metadata_entry_id: None,
+                checkpoint: None,
+            }),
+        )
+        .await
+        .unwrap();
+    assert_eq!(
+        acknowledged_rows,
+        count_rows(&candidate, region_id).await,
+        "the candidate misses rows the old leader acknowledged; Follower reported last_entry_id {reported}"
+    );
+}

src/mito2/src/region.rs

@@ -599,12 +599,24 @@ impl MitoRegion {
                         info!("Region {} already in downgrading mode", self.region_id);
                         wait_for_checkpoint = true;
                     }
-                    _ => {
+                    RegionRoleState::Follower => {
                         warn!(
                             "Cannot start downgrade for region {} from state {:?}",
                             self.region_id, current_state
                         );
                     }
+                    RegionRoleState::Leader(_) => {
+                        // The region may become writable again when the
+                        // current operation, such as an ALTER, finishes. Fail
+                        // with a retryable error instead of reporting a
+                        // downgrade that does not fence writes.
+                        return Err(RegionStateSnafu {
+                            region_id: self.region_id,
+                            state: current_state,
+                            expect: RegionRoleState::Leader(RegionLeaderState::Writable),
+                        }
+                        .build());
+                    }
                 }
                 None
             }

src/mito2/src/worker.rs

@@ -52,7 +52,7 @@ use rand::{Rng, rng};
 use snafu::{ResultExt, ensure};
 use store_api::logstore::LogStore;
 use store_api::region_engine::{
-    SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState,
+    RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState,
 };
 use store_api::storage::{FileId, RegionId};
 use tokio::sync::mpsc::{Receiver, Sender};
@@ -1417,6 +1417,15 @@ impl<S: LogStore> RegionWorkerLoop<S> {
         sender: oneshot::Sender<SetRegionRoleStateResponse>,
     ) {
         if let Some(region) = self.regions.get_region(region_id) {
+            if matches!(
+                region_role_state,
+                SettableRegionRoleState::Follower | SettableRegionRoleState::DowngradingLeader
+            ) {
+                // Fence writes here: writes the worker handles later see
+                // Downgrading, and earlier ones have finished. This only
+                // changes a writable or staging leader.
+                region.set_role(RegionRole::DowngradingLeader);
+            }
             // We need to do this in background as we need the manifest lock.
             common_runtime::spawn_global(async move {
                 match region.set_role_state_gracefully(region_role_state).await {