fix: acknowledge writes only within the region's leader lease (#9904)

* fix: acknowledge writes only within the region's leader lease

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

* fix: drop cached logical columns after a failed logical drop

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

* fix: keep the prepared region for the lease check after the WAL append

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

* fix: serialize a logical drop with fills of the logical column cache

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

* test: check the drop lock and grant long leases deterministically

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

* test: cover a bulk insert that loses its lease during the WAL append

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

---------

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

src/datanode/src/alive_keeper.rs

@@ -278,6 +278,9 @@ impl HeartbeatResponseHandler for RegionAliveKeeper {
         let start_instant = self.epoch + Duration::from_millis(region_lease.duration_since_epoch);
         let deadline = start_instant + Duration::from_secs(region_lease.lease_seconds);
 
+        if let Some(mito) = self.region_server.mito_engine() {
+            mito.renew_node_lease(deadline.into_std());
+        }
         self.renew_region_leases(&region_lease.regions, deadline)
             .await;
         self.close_staled_regions(&region_lease.closeable_region_ids)
@@ -435,7 +438,7 @@ impl CountdownTask {
                         },
                         Some(CountdownCommand::Reset((role, deadline, extension_info))) => {
                             let prev_role_leader = self.region_server.is_region_leader(self.region_id).unwrap_or(false);
-                            if let Err(err) = self.region_server.set_region_role(self.region_id, role) {
+                            if let Err(err) = self.region_server.set_region_role_with_lease(self.region_id, role, deadline.into_std()) {
                                 if err.status_code() == StatusCode::RegionNotFound {
                                     // Table metadata in metasrv is deleted after its regions are dropped.
                                     // The datanode may still receive lease renewal responses that depend on the metadata
@@ -500,14 +503,92 @@ impl CountdownTask {
 #[cfg(test)]
 mod test {
 
+    use api::v1::Rows;
+    use api::v1::meta::{HeartbeatResponse, RegionLease};
+    use common_meta::heartbeat::mailbox::HeartbeatMailbox;
     use mito2::config::MitoConfig;
-    use mito2::test_util::{CreateRequestBuilder, TestEnv};
+    use mito2::test_util::{CreateRequestBuilder, TestEnv, build_rows, rows_schema};
     use store_api::region_engine::RegionEngine;
-    use store_api::region_request::{EnterStagingRequest, StagingPartitionDirective};
+    use store_api::region_request::{
+        EnterStagingRequest, RegionPutRequest, StagingPartitionDirective,
+    };
 
     use super::*;
     use crate::tests::mock_region_server;
 
+    /// Handles a heartbeat response that grants `regions` for a heartbeat sent now.
+    async fn handle_lease_response(alive_keeper: &RegionAliveKeeper, regions: Vec<GrantedRegion>) {
+        let (tx, _rx) = mpsc::channel(1);
+        let response = HeartbeatResponse {
+            region_lease: Some(RegionLease {
+                regions,
+                duration_since_epoch: alive_keeper.epoch().elapsed().as_millis() as u64,
+                // The countdown of a granted region does not expire during the test.
+                lease_seconds: 86400,
+                closeable_region_ids: vec![],
+            }),
+            ..Default::default()
+        };
+        let mut ctx =
+            HeartbeatResponseHandlerContext::new(Arc::new(HeartbeatMailbox::new(tx)), response);
+        alive_keeper.handle(&mut ctx).await.unwrap();
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn lease_responses_let_regions_write() {
+        common_telemetry::init_default_ut_logging();
+        let mut region_server = mock_region_server();
+        let mut engine_env = TestEnv::with_prefix("region-alive-keeper-write-lease").await;
+        let engine = engine_env.create_engine(MitoConfig::default()).await;
+        engine.require_write_lease();
+        // The clock of the lease checks stands still, so only the lease propagation matters.
+        engine.freeze_lease_clock(std::time::Instant::now());
+        region_server.register_engine(Arc::new(engine));
+        let alive_keeper =
+            RegionAliveKeeper::new(region_server.clone(), None, Duration::from_millis(100));
+        let create = |region_id| {
+            let request = RegionRequest::Create(CreateRequestBuilder::new().build());
+            region_server.handle_request(region_id, request)
+        };
+        let schema = rows_schema(&CreateRequestBuilder::new().build());
+        let put = |region_id| {
+            let rows = Rows {
+                schema: schema.clone(),
+                rows: build_rows(0, 1),
+            };
+            let request = RegionRequest::Put(RegionPutRequest {
+                rows,
+                hint: None,
+                partition_expr_version: None,
+                skip_wal: false,
+            });
+            region_server.handle_request(region_id, request)
+        };
+
+        // Created before any heartbeat response, the region has no lease.
+        let region_id = RegionId::new(1024, 1);
+        create(region_id).await.unwrap();
+        assert!(put(region_id).await.unwrap_err().is_retryable());
+
+        // A response renews the lease of the node, under which regions created later write.
+        handle_lease_response(&alive_keeper, vec![]).await;
+        let created_region_id = RegionId::new(1024, 2);
+        create(created_region_id).await.unwrap();
+        put(created_region_id).await.unwrap();
+
+        // A Leader grant renews the lease of the region.
+        alive_keeper.register_region(region_id).await;
+        let grant = GrantedRegion {
+            region_id: region_id.as_u64(),
+            role: api::v1::meta::RegionRole::Leader.into(),
+            extensions: HashMap::new(),
+        };
+        handle_lease_response(&alive_keeper, vec![grant]).await;
+        // The countdown task has applied the grant once it answers.
+        alive_keeper.deadline(region_id).await.unwrap();
+        put(region_id).await.unwrap();
+    }
+
     #[tokio::test(flavor = "multi_thread")]
     async fn region_alive_keeper() {
         common_telemetry::init_default_ut_logging();

src/datanode/src/datanode.rs

@@ -370,6 +370,11 @@ impl DatanodeBuilder {
                 cache_registry.clone(),
             )
             .await?;
+        // A region acknowledges only writes whose WAL append completed within its lease from
+        // the metasrv.
+        if controlled_by_metasrv && let Some(mito) = region_server.mito_engine() {
+            mito.require_write_lease();
+        }
 
         // TODO(weny): Considering introducing a readonly kv_backend trait.
         let runtime_switch_manager = RuntimeSwitchManager::new(self.kv_backend.clone());

src/datanode/src/region_server.rs

@@ -471,6 +471,23 @@ impl RegionServer {
             .with_context(|_| HandleRegionRequestSnafu { region_id })
     }
 
+    /// Sets the role of a region from a lease the metasrv grants until `deadline`.
+    pub fn set_region_role_with_lease(
+        &self,
+        region_id: RegionId,
+        role: RegionRole,
+        deadline: std::time::Instant,
+    ) -> Result<()> {
+        let engine = self
+            .inner
+            .region_map
+            .get(&region_id)
+            .with_context(|| RegionNotFoundSnafu { region_id })?;
+        engine
+            .set_region_role_with_lease(region_id, role, deadline)
+            .with_context(|_| HandleRegionRequestSnafu { region_id })
+    }
+
     /// Set region role state gracefully.
     ///
     /// For [SettableRegionRoleState::Follower]:

src/datanode/src/region_server/flow_write.rs

@@ -430,6 +430,15 @@ mod tests {
         ) -> Result<(), common_error::ext::BoxedError> {
             self.inner.set_region_role(id, role)
         }
+
+        fn set_region_role_with_lease(
+            &self,
+            id: RegionId,
+            role: store_api::region_engine::RegionRole,
+            deadline: std::time::Instant,
+        ) -> Result<(), common_error::ext::BoxedError> {
+            self.inner.set_region_role_with_lease(id, role, deadline)
+        }
         async fn sync_region(
             &self,
             id: RegionId,

src/datanode/src/tests.rs

@@ -289,6 +289,15 @@ impl RegionEngine for MockRegionEngine {
         Ok(())
     }
 
+    fn set_region_role_with_lease(
+        &self,
+        _region_id: RegionId,
+        _role: RegionRole,
+        _deadline: std::time::Instant,
+    ) -> Result<(), BoxedError> {
+        Ok(())
+    }
+
     async fn set_region_role_state_gracefully(
         &self,
         region_id: RegionId,

src/file-engine/src/engine.rs

@@ -15,6 +15,7 @@
 use std::any::Any;
 use std::collections::HashMap;
 use std::sync::{Arc, RwLock};
+use std::time::Instant;
 
 use api::region::RegionResponse;
 use async_trait::async_trait;
@@ -131,6 +132,16 @@ impl RegionEngine for FileRegionEngine {
             .map_err(BoxedError::new)
     }
 
+    /// Regions of this engine take no writes, so the lease is not needed.
+    fn set_region_role_with_lease(
+        &self,
+        region_id: RegionId,
+        role: RegionRole,
+        _deadline: Instant,
+    ) -> Result<(), BoxedError> {
+        self.set_region_role(region_id, role)
+    }
+
     async fn set_region_role_state_gracefully(
         &self,
         region_id: RegionId,

src/metric-engine/src/engine.rs

@@ -31,6 +31,7 @@ mod sync;
 use std::any::Any;
 use std::collections::HashMap;
 use std::sync::{Arc, RwLock};
+use std::time::Instant;
 
 use api::region::RegionResponse;
 use async_trait::async_trait;
@@ -383,6 +384,29 @@ impl RegionEngine for MetricEngine {
         Ok(())
     }
 
+    fn set_region_role_with_lease(
+        &self,
+        region_id: RegionId,
+        role: RegionRole,
+        deadline: Instant,
+    ) -> Result<(), BoxedError> {
+        // ignore the region not found error
+        for x in [
+            utils::to_metadata_region_id(region_id),
+            utils::to_data_region_id(region_id),
+        ] {
+            if let Err(e) = self
+                .inner
+                .mito
+                .set_region_role_with_lease(x, role, deadline)
+                && e.status_code() != StatusCode::RegionNotFound
+            {
+                return Err(e);
+            }
+        }
+        Ok(())
+    }
+
     async fn sync_region(
         &self,
         region_id: RegionId,

src/metric-engine/src/engine/alter.rs

@@ -196,20 +196,30 @@ impl MetricEngineInner {
             (name.to_string(), column_metadata.clone())
         });
 
+        // The data region has the new columns, whatever happens to the metadata write.
+        self.state
+            .write()
+            .unwrap()
+            .add_physical_columns(data_region_id, new_add_columns);
+
         // Writes logical regions metadata to metadata region
-        self.metadata_region
+        let result = self
+            .metadata_region
             .add_logical_regions(physical_region_id, false, logical_region_columns)
-            .await?;
+            .await;
+        // A failed write may still be applied, e.g. when the region lease expires after the WAL
+        // append, so the cached logical columns are dropped either way.
+        self.state
+            .write()
+            .unwrap()
+            .invalid_logical_regions_cache(requests.iter().map(|(region_id, _)| *region_id));
+        result?;
 
         extension_return_value.insert(
             ALTER_PHYSICAL_EXTENSION_KEY.to_string(),
             ColumnMetadata::encode_list(&physical_columns).context(SerializeColumnMetadataSnafu)?,
         );
 
-        let mut state = self.state.write().unwrap();
-        state.add_physical_columns(data_region_id, new_add_columns);
-        state.invalid_logical_regions_cache(requests.iter().map(|(region_id, _)| *region_id));
-
         Ok(0)
     }
 

src/metric-engine/src/engine/create.rs

@@ -324,6 +324,11 @@ impl MetricEngineInner {
             let column_metadata = *physical_schema_map.get(name).unwrap();
             (name.to_string(), column_metadata.clone())
         });
+        // The data region has the new columns, whatever happens to the metadata write.
+        self.state
+            .write()
+            .unwrap()
+            .add_physical_columns(data_region_id, new_add_columns);
 
         extension_return_value.insert(
             ALTER_PHYSICAL_EXTENSION_KEY.to_string(),
@@ -335,11 +340,10 @@ impl MetricEngineInner {
             .add_logical_regions(physical_region_id, true, logical_region_columns)
             .await?;
 
-        {
-            let mut state = self.state.write().unwrap();
-            state.add_physical_columns(data_region_id, new_add_columns);
-            state.add_logical_regions(physical_region_id, logical_regions.clone());
-        }
+        self.state
+            .write()
+            .unwrap()
+            .add_logical_regions(physical_region_id, logical_regions.clone());
         for logical_region_id in logical_regions {
             self.metadata_region
                 .open_logical_region(logical_region_id)

src/metric-engine/src/engine/drop.rs

@@ -157,10 +157,24 @@ impl MetricEngineInner {
         logical_region_id: RegionId,
         physical_region_id: RegionId,
     ) -> Result<AffectedRows> {
-        // Update metadata
-        self.metadata_region
-            .remove_logical_region(physical_region_id, logical_region_id)
+        // A fill of the cached logical columns holds the read lock until it publishes them, so
+        // none publishes columns read before the delete after the invalidation below.
+        let _write_guard = self
+            .metadata_region
+            .write_lock_logical_region(logical_region_id)
             .await?;
+        // Update metadata
+        let result = self
+            .metadata_region
+            .remove_logical_region(physical_region_id, logical_region_id)
+            .await;
+        // A failed delete may still be applied, e.g. when the region lease expires after the WAL
+        // append, so the cached logical columns are dropped either way.
+        self.state
+            .write()
+            .unwrap()
+            .invalid_logical_regions_cache([logical_region_id]);
+        result?;
 
         // Update engine state
         self.state

src/metric-engine/src/engine/put.rs

@@ -785,6 +785,8 @@ impl MetricEngineInner {
 #[cfg(test)]
 mod tests {
     use std::collections::HashSet;
+    use std::sync::Arc;
+    use std::time::{Duration, Instant};
 
     use api::v1::value::ValueData;
     use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema};
@@ -799,6 +801,10 @@ mod tests {
     use datatypes::prelude::ConcreteDataType;
     use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema};
     use datatypes::value::Value as PartitionValue;
+    use futures_util::FutureExt;
+    use mito2::config::MitoConfig;
+    use mito2::engine::listener::WriteWalListener;
+    use mito2::test_util::TestEnv as MitoTestEnv;
     use partition::expr::col;
     use store_api::metadata::ColumnMetadata;
     use store_api::metric_engine_consts::{
@@ -806,10 +812,10 @@ mod tests {
         PHYSICAL_TABLE_METADATA_KEY, PRIMARY_KEY_ENCODING,
     };
     use store_api::path_utils::table_dir;
-    use store_api::region_engine::RegionEngine;
+    use store_api::region_engine::{RegionEngine, RegionRole};
     use store_api::region_request::{
-        EnterStagingRequest, PathType, RegionCloseRequest, RegionOpenRequest, RegionRequest,
-        StagingPartitionDirective,
+        EnterStagingRequest, PathType, RegionCloseRequest, RegionDropRequest, RegionOpenRequest,
+        RegionRequest, StagingPartitionDirective,
     };
     use store_api::storage::ScanRequest;
     use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
@@ -2350,4 +2356,128 @@ mod tests {
             "expected impure-default rejection, got: {err}"
         );
     }
+
+    #[tokio::test]
+    async fn test_logical_region_ddl_needs_leader_lease() {
+        common_telemetry::init_default_ut_logging();
+        let mut mito_env = MitoTestEnv::with_prefix("logical-region-lease").await;
+        let listener = Arc::new(WriteWalListener::default());
+        let mito = mito_env
+            .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
+            .await;
+        mito.require_write_lease();
+        let start = Instant::now();
+        mito.freeze_lease_clock(start);
+        let env = TestEnv::with_mito_engine(mito_env, mito.clone());
+        let engine = env.metric();
+        // Created before the node has a lease, the physical region has none.
+        let physical_region_id = env.default_physical_region_id();
+        env.create_physical_region(physical_region_id, &TestEnv::default_table_dir(), vec![])
+            .await;
+        // The clock of the lease checks stands at `start` until `expire(n)` moves it to `at(n)`,
+        // where a lease from `grant(n)` expires.
+        let at = |n| start + n * Duration::from_secs(10);
+        let grant = |n| {
+            engine
+                .set_region_role_with_lease(physical_region_id, RegionRole::Leader, at(n))
+                .unwrap()
+        };
+        let expire = |n| mito.freeze_lease_clock(at(n));
+        let logical_region_id = env.default_logical_region_id();
+        let alter = |tag| RegionRequest::Alter(test_util::alter_logical_region_request(&[tag]));
+        let drop_request = || {
+            RegionRequest::Drop(RegionDropRequest {
+                fast_path: false,
+                force: false,
+                partial_drop: false,
+            })
+        };
+        let put = |tags: &[&str]| {
+            let request = RegionRequest::Put(RegionPutRequest {
+                rows: Rows {
+                    schema: test_util::row_schema_with_tags(tags),
+                    rows: test_util::build_rows(tags.len(), 5),
+                },
+                hint: None,
+                partition_expr_version: None,
+                skip_wal: false,
+            });
+            engine.handle_request(logical_region_id, request)
+        };
+        // Holds `request` after its metadata region append while the lease from `grant(n)`
+        // expires. With `blocks_fills`, a fill of the cached logical columns, which takes the read
+        // lock of the logical region, cannot start until the request finishes.
+        let expire_after_append = async |request, n, blocks_fills: bool| {
+            listener.pause_next();
+            let mut task = {
+                let engine = engine.clone();
+                tokio::spawn(async move { engine.handle_request(logical_region_id, request).await })
+            };
+            tokio::select! {
+                () = listener.wait_appended() => {}
+                result = &mut task => panic!("the request finished before its append: {result:?}"),
+            }
+            if blocks_fills {
+                let metadata_region = &engine.inner.metadata_region;
+                let fill = metadata_region.read_lock_logical_region(logical_region_id);
+                assert!(fill.now_or_never().is_none());
+            }
+            expire(n);
+            listener.resume();
+            task.await.unwrap()
+        };
+
+        // Each DDL below gives the data region a new column before the metadata region rejects
+        // its write; the retry must not add the column again.
+        let create = || {
+            RegionRequest::Create(test_util::create_logical_region_request(
+                &["job"],
+                physical_region_id,
+                &table_dir("test", logical_region_id.table_id()),
+            ))
+        };
+        let err = engine
+            .handle_request(logical_region_id, create())
+            .await
+            .unwrap_err();
+        assert!(err.is_retryable(), "{err:?}");
+        grant(1);
+        engine
+            .handle_request(logical_region_id, create())
+            .await
+            .unwrap();
+        put(&["job"]).await.unwrap();
+
+        // The metadata write of the ALTER is applied, so the cached metadata must not hide it.
+        let err = expire_after_append(alter("tag1"), 1, false)
+            .await
+            .unwrap_err();
+        assert!(err.is_retryable(), "{err:?}");
+        let metadata = engine.get_metadata(logical_region_id).await.unwrap();
+        assert!(metadata.column_by_name("tag1").is_some());
+
+        let err = engine
+            .handle_request(logical_region_id, alter("tag2"))
+            .await
+            .unwrap_err();
+        assert!(err.is_retryable(), "{err:?}");
+        grant(2);
+        engine
+            .handle_request(logical_region_id, alter("tag2"))
+            .await
+            .unwrap();
+        put(&["job", "tag1", "tag2"]).await.unwrap();
+
+        // The metadata delete of the DROP is applied, so writes must not use the cached columns.
+        let err = expire_after_append(drop_request(), 2, true)
+            .await
+            .unwrap_err();
+        assert!(err.is_retryable(), "{err:?}");
+        grant(3);
+        put(&["job"]).await.unwrap_err();
+        engine
+            .handle_request(logical_region_id, drop_request())
+            .await
+            .unwrap();
+    }
 }

src/metric-engine/src/metadata_region.rs

@@ -538,11 +538,15 @@ impl MetadataRegion {
     ) -> Result<()> {
         let cache_access_lock = self.cache_access_lock(metadata_region_id);
         let _cache_guard = cache_access_lock.write().await;
-        self.mito
+        let result = self
+            .mito
             .handle_request(metadata_region_id, request)
             .await
-            .context(MitoWriteOperationSnafu)?;
+            .context(MitoWriteOperationSnafu);
+        // A failed write may still be applied, e.g. when the region lease expires after the
+        // WAL append.
         self.cache.invalidate(&metadata_region_id).await;
+        result?;
 
         Ok(())
     }

src/mito2/src/engine.rs

@@ -48,6 +48,8 @@ mod filter_deleted_test;
 mod flush_test;
 #[cfg(test)]
 mod index_build_test;
+#[cfg(test)]
+mod lease_test;
 #[cfg(any(test, feature = "test"))]
 pub mod listener;
 #[cfg(test)]
@@ -327,6 +329,22 @@ impl MitoEngine {
         self.inner.workers.gc_limiter()
     }
 
+    /// Makes regions acknowledge only writes whose WAL append completes within their lease.
+    pub fn require_write_lease(&self) {
+        self.inner.workers.node_lease().enable();
+    }
+
+    /// Renews the lease of this node from a heartbeat response that carries region leases.
+    pub fn renew_node_lease(&self, deadline: Instant) {
+        self.inner.workers.node_lease().renew(deadline);
+    }
+
+    /// Freezes the clock of the lease checks at `at`.
+    #[cfg(any(test, feature = "test"))]
+    pub fn freeze_lease_clock(&self, at: Instant) {
+        self.inner.workers.node_lease().freeze_clock(at);
+    }
+
     pub fn object_store_manager(&self) -> &ObjectStoreManagerRef {
         self.inner.workers.object_store_manager()
     }
@@ -1581,6 +1599,21 @@ impl RegionEngine for MitoEngine {
             .map_err(BoxedError::new)
     }
 
+    fn set_region_role_with_lease(
+        &self,
+        region_id: RegionId,
+        role: RegionRole,
+        deadline: Instant,
+    ) -> Result<(), BoxedError> {
+        let region = self.inner.find_region(region_id).map_err(BoxedError::new)?;
+        // Before the role, and even if the role leaves the region's state unchanged.
+        if role.writable() {
+            region.leader_deadline.extend(deadline);
+        }
+        region.set_role(role);
+        Ok(())
+    }
+
     async fn set_region_role_state_gracefully(
         &self,
         region_id: RegionId,

src/mito2/src/engine/lease_test.rs

@@ -0,0 +1,284 @@
+// 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.
+
+//! Tests that a region acknowledges only writes whose WAL append completed within its leader
+//! lease.
+
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use api::region::RegionResponse;
+use api::v1::{ColumnSchema, Rows};
+use common_error::ext::{BoxedError, ErrorExt};
+use common_recordbatch::RecordBatches;
+use rstest::rstest;
+use rstest_reuse::apply;
+use store_api::region_engine::{RegionEngine, RegionRole};
+use store_api::region_request::{RegionCatchupRequest, RegionPutRequest, RegionRequest};
+use store_api::storage::{RegionId, ScanRequest};
+use tokio::task::JoinHandle;
+
+use crate::config::MitoConfig;
+use crate::engine::MitoEngine;
+use crate::engine::basic_test::build_bulk_insert_request;
+use crate::engine::listener::WriteWalListener;
+use crate::error::Error;
+use crate::test_util::{
+    CreateRequestBuilder, LogStoreFactory, TestEnv, build_rows, kafka_log_store_factory,
+    multiple_log_store_factories, prepare_test_for_kafka_log_store, raft_engine_log_store_factory,
+    rows_schema,
+};
+
+const LEASE: Duration = Duration::from_secs(10);
+const REGION_ID: RegionId = RegionId::new(1, 1);
+
+/// Creates [REGION_ID] on an engine whose writes need a lease. The node has no lease yet, and
+/// the clock of the lease checks stands still until a test moves it.
+async fn create_region(
+    env: &mut TestEnv,
+    topic: Option<String>,
+) -> (Arc<WriteWalListener>, MitoEngine, Vec<ColumnSchema>) {
+    let listener = Arc::new(WriteWalListener::default());
+    let engine = env
+        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
+        .await;
+    engine.require_write_lease();
+    engine.freeze_lease_clock(Instant::now());
+    let request = CreateRequestBuilder::new().kafka_topic(topic).build();
+    let schema = rows_schema(&request);
+    engine
+        .handle_request(REGION_ID, RegionRequest::Create(request))
+        .await
+        .unwrap();
+    (listener, engine, schema)
+}
+
+fn grant(engine: &MitoEngine, role: RegionRole, deadline: Instant) {
+    engine
+        .set_region_role_with_lease(REGION_ID, role, deadline)
+        .unwrap();
+}
+
+/// Returns a request that writes rows `start..start + 3`.
+fn put_request(schema: &[ColumnSchema], start: usize) -> RegionRequest {
+    let rows = Rows {
+        schema: schema.to_vec(),
+        rows: build_rows(start, start + 3),
+    };
+    RegionRequest::Put(RegionPutRequest {
+        rows,
+        hint: None,
+        partition_expr_version: None,
+        skip_wal: false,
+    })
+}
+
+fn put(
+    engine: &MitoEngine,
+    region_id: RegionId,
+    schema: &[ColumnSchema],
+    start: usize,
+) -> JoinHandle<Result<RegionResponse, BoxedError>> {
+    let engine = engine.clone();
+    let request = put_request(schema, start);
+    tokio::spawn(async move { engine.handle_request(region_id, request).await })
+}
+
+/// Sends `request` to [REGION_ID] and waits until it is held after its WAL append.
+async fn hold(
+    listener: &WriteWalListener,
+    engine: &MitoEngine,
+    request: RegionRequest,
+) -> JoinHandle<Result<RegionResponse, BoxedError>> {
+    listener.pause_next();
+    let engine = engine.clone();
+    let mut write = tokio::spawn(async move { engine.handle_request(REGION_ID, request).await });
+    tokio::select! {
+        () = listener.wait_appended() => {}
+        result = &mut write => panic!("the write finished before its WAL append: {result:?}"),
+    }
+    write
+}
+
+async fn count_rows(engine: &MitoEngine) -> 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()
+}
+
+fn assert_retryable(result: Result<RegionResponse, BoxedError>, expected: fn(&Error) -> bool) {
+    let err = result.unwrap_err();
+    assert!(err.retry_hint().is_retryable(), "{err:?}");
+    let err = err.into_inner();
+    let err = match err.as_any().downcast_ref::<Error>().unwrap() {
+        Error::WriteGroup { source } => source.as_ref(),
+        err => err,
+    };
+    assert!(expected(err), "{err:?}");
+}
+
+fn lease_expired(err: &Error) -> bool {
+    matches!(err, Error::RegionLeaseExpired { .. })
+}
+
+fn unacknowledged(err: &Error) -> bool {
+    matches!(err, Error::UnacknowledgedWrite { .. })
+}
+
+#[apply(multiple_log_store_factories)]
+async fn test_lease_expires_while_writing_wal(factory: Option<LogStoreFactory>) {
+    common_telemetry::init_default_ut_logging();
+    let Some(factory) = factory else {
+        return;
+    };
+    let mut env = TestEnv::with_prefix("lease-expires-while-writing-wal")
+        .await
+        .with_log_store_factory(factory.clone());
+    let topic = prepare_test_for_kafka_log_store(&factory).await;
+    let (listener, engine, schema) = create_region(&mut env, topic).await;
+    let deadline = Instant::now() + LEASE;
+    grant(&engine, RegionRole::Leader, deadline);
+
+    let write = hold(&listener, &engine, put_request(&schema, 0)).await;
+    engine.freeze_lease_clock(deadline);
+    listener.resume();
+    assert_retryable(write.await.unwrap(), unacknowledged);
+    // The write is in the WAL, so it is applied rather than undone.
+    assert_eq!(3, count_rows(&engine).await);
+
+    // A new write fails before it reaches the WAL.
+    assert_retryable(
+        put(&engine, REGION_ID, &schema, 3).await.unwrap(),
+        lease_expired,
+    );
+    assert_eq!(3, count_rows(&engine).await);
+
+    // The same holds for a bulk insert.
+    let deadline = deadline + LEASE;
+    grant(&engine, RegionRole::Leader, deadline);
+    let bulk_insert = RegionRequest::BulkInserts(build_bulk_insert_request(REGION_ID, 3));
+    let write = hold(&listener, &engine, bulk_insert).await;
+    let committed_sequence = || {
+        let region = engine.get_region(REGION_ID).unwrap();
+        region.version_control.current().committed_sequence
+    };
+    let sequence = committed_sequence();
+    engine.freeze_lease_clock(deadline);
+    listener.resume();
+    assert_retryable(write.await.unwrap(), unacknowledged);
+    assert_eq!(6, count_rows(&engine).await);
+    assert!(committed_sequence() > sequence);
+}
+
+#[tokio::test]
+async fn test_lease_renewed_while_writing_wal() {
+    common_telemetry::init_default_ut_logging();
+    let mut env = TestEnv::with_prefix("lease-renewed-while-writing-wal").await;
+    let (listener, engine, schema) = create_region(&mut env, None).await;
+    let deadline = Instant::now() + LEASE;
+    grant(&engine, RegionRole::Leader, deadline);
+
+    let write = hold(&listener, &engine, put_request(&schema, 0)).await;
+    // The grant keeps the region writable rather than staging, and still extends the deadline.
+    grant(&engine, RegionRole::StagingLeader, deadline + LEASE);
+    engine.freeze_lease_clock(deadline);
+    listener.resume();
+    assert_eq!(3, write.await.unwrap().unwrap().affected_rows);
+}
+
+#[rstest]
+#[case::within_leader_lease(false)]
+#[case::after_leader_lease(true)]
+#[tokio::test]
+async fn test_downgrade_while_writing_wal(#[case] lease_expires: bool) {
+    common_telemetry::init_default_ut_logging();
+    let mut env = TestEnv::with_prefix("downgrade-while-writing-wal").await;
+    let (listener, engine, schema) = create_region(&mut env, None).await;
+    let deadline = Instant::now() + LEASE;
+    grant(&engine, RegionRole::Leader, deadline);
+
+    let write = hold(&listener, &engine, put_request(&schema, 0)).await;
+    // The metasrv renews the lease as DowngradingLeader once the route marks the region as
+    // downgrading. A write already in the WAL is acknowledged, but only within the Leader
+    // deadline, which such grants do not extend.
+    grant(&engine, RegionRole::DowngradingLeader, deadline + LEASE);
+    if lease_expires {
+        engine.freeze_lease_clock(deadline);
+    }
+    listener.resume();
+    let result = write.await.unwrap();
+    if lease_expires {
+        assert_retryable(result, unacknowledged);
+    } else {
+        assert_eq!(3, result.unwrap().affected_rows);
+    }
+}
+
+#[tokio::test]
+async fn test_created_region_writes_under_node_lease() {
+    common_telemetry::init_default_ut_logging();
+    let mut env = TestEnv::with_prefix("created-region-node-lease").await;
+    let (_listener, engine, schema) = create_region(&mut env, None).await;
+    let bulk_insert = |region_id| {
+        engine.handle_request(
+            region_id,
+            RegionRequest::BulkInserts(build_bulk_insert_request(region_id, 3)),
+        )
+    };
+    // Created before the node has a lease, the region has no deadline.
+    assert_retryable(
+        put(&engine, REGION_ID, &schema, 0).await.unwrap(),
+        lease_expired,
+    );
+    assert_retryable(bulk_insert(REGION_ID).await, lease_expired);
+
+    // Renewing the node lease does not extend the deadline of a region created before.
+    engine.renew_node_lease(Instant::now() + LEASE);
+    assert!(put(&engine, REGION_ID, &schema, 0).await.unwrap().is_err());
+
+    // A region created under the node lease takes writes before its first Leader grant.
+    let region_id = RegionId::new(1, 2);
+    engine
+        .handle_request(
+            region_id,
+            RegionRequest::Create(CreateRequestBuilder::new().build()),
+        )
+        .await
+        .unwrap();
+    put(&engine, region_id, &schema, 0).await.unwrap().unwrap();
+    bulk_insert(region_id).await.unwrap();
+}
+
+#[tokio::test]
+async fn test_upgraded_region_writes_before_leader_grant() {
+    common_telemetry::init_default_ut_logging();
+    let mut env = TestEnv::with_prefix("upgraded-region-node-lease").await;
+    let (_listener, engine, schema) = create_region(&mut env, None).await;
+    // A migration candidate is a follower until the upgrade.
+    grant(&engine, RegionRole::Follower, Instant::now() + LEASE);
+    engine.renew_node_lease(Instant::now() + LEASE);
+
+    let request = RegionCatchupRequest {
+        set_writable: true,
+        ..Default::default()
+    };
+    engine
+        .handle_request(REGION_ID, RegionRequest::Catchup(request))
+        .await
+        .unwrap();
+    put(&engine, REGION_ID, &schema, 0).await.unwrap().unwrap();
+}

src/mito2/src/engine/listener.rs

@@ -15,7 +15,7 @@
 //! Engine event listener for tests.
 
 use std::sync::Arc;
-use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::time::Duration;
 
 use async_trait::async_trait;
@@ -39,6 +39,9 @@ pub trait EventListener: Send + Sync {
     /// Notifies the listener that the engine is stalled.
     fn on_write_stall(&self) {}
 
+    /// Notifies the listener after a batch of writes is appended to the WAL.
+    async fn on_write_wal_end(&self) {}
+
     /// Notifies the listener that the region starts to do flush.
     async fn on_flush_begin(&self, region_id: RegionId) {
         let _ = region_id;
@@ -617,3 +620,38 @@ impl EventListener for GateIndexBuildListener {
         self.stop_notify.notify_one();
     }
 }
+
+/// Listener to pause a batch of writes after it is appended to the WAL.
+#[derive(Default)]
+pub struct WriteWalListener {
+    pause: AtomicBool,
+    appended_notify: Notify,
+    resume_notify: Notify,
+}
+
+impl WriteWalListener {
+    /// Pauses the next batch of writes after it is appended to the WAL.
+    pub fn pause_next(&self) {
+        self.pause.store(true, Ordering::Relaxed);
+    }
+
+    /// Waits until the paused batch is appended to the WAL.
+    pub async fn wait_appended(&self) {
+        self.appended_notify.notified().await;
+    }
+
+    /// Resumes the paused batch.
+    pub fn resume(&self) {
+        self.resume_notify.notify_one();
+    }
+}
+
+#[async_trait]
+impl EventListener for WriteWalListener {
+    async fn on_write_wal_end(&self) {
+        if self.pause.swap(false, Ordering::Relaxed) {
+            self.appended_notify.notify_one();
+            self.resume_notify.notified().await;
+        }
+    }
+}

src/mito2/src/error.rs

@@ -656,6 +656,23 @@ pub enum Error {
         location: Location,
     },
 
+    #[snafu(display("Region {} does not hold a valid leader lease", region_id))]
+    RegionLeaseExpired {
+        region_id: RegionId,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
+    #[snafu(display(
+        "Region {} lost its leader lease while writing to the WAL, the write may have been applied",
+        region_id
+    ))]
+    UnacknowledgedWrite {
+        region_id: RegionId,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
     #[snafu(display("Failed to compact region {}", region_id))]
     CompactRegion {
         region_id: RegionId,
@@ -1584,6 +1601,7 @@ impl ErrorExt for Error {
             CompatReader { .. } => StatusCode::Unexpected,
             InvalidRegionRequest { source, .. } => source.status_code(),
             RegionState { .. } | UpdateManifest { .. } => StatusCode::RegionNotReady,
+            RegionLeaseExpired { .. } | UnacknowledgedWrite { .. } => StatusCode::RegionNotReady,
             JsonOptions { .. } => StatusCode::InvalidArguments,
             EmptyRegionDir { .. } | EmptyManifestDir { .. } => StatusCode::RegionNotFound,
             ConvertValue { source, .. } => source.status_code(),
@@ -1681,6 +1699,8 @@ impl ErrorExt for Error {
             | RegionNotFound { .. }
             | RegionClosed { .. }
             | RejectWriteInStaging { .. }
+            | RegionLeaseExpired { .. }
+            | UnacknowledgedWrite { .. }
             | RegionState {
                 state: RegionRoleState::Follower,
                 ..

src/mito2/src/region.rs

@@ -23,8 +23,9 @@ pub(crate) mod version;
 use std::cell::OnceCell;
 use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet};
-use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
-use std::sync::{Arc, Mutex, RwLock};
+use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering};
+use std::sync::{Arc, LazyLock, Mutex, RwLock};
+use std::time::Instant;
 
 use common_base::hash::partition_expr_version;
 use common_recordbatch::adapter::RegionQueryStatCounters;
@@ -185,6 +186,8 @@ pub struct MitoRegion {
     /// There are no WAL entries in range [flushed_entry_id, topic_latest_entry_id] for current region,
     /// which means these WAL entries maybe able to be pruned up to `topic_latest_entry_id`.
     pub(crate) topic_latest_entry_id: AtomicU64,
+    /// Deadline of the region's latest Leader or StagingLeader lease from the metasrv.
+    pub(crate) leader_deadline: Deadline,
     /// Region stats.
     pub(crate) region_stats: RegionStats,
     /// manifest stats
@@ -221,6 +224,80 @@ impl RegionStats {
 
 pub type MitoRegionRef = Arc<MitoRegion>;
 
+/// The instant lease deadlines count from, so that a deadline fits in an [AtomicU64].
+static LEASE_EPOCH: LazyLock<Instant> = LazyLock::new(Instant::now);
+
+fn lease_nanos(instant: Instant) -> u64 {
+    instant.saturating_duration_since(*LEASE_EPOCH).as_nanos() as u64
+}
+
+/// A lease deadline in nanoseconds since [LEASE_EPOCH]. It only moves forward; 0 means none.
+#[derive(Debug, Default)]
+pub(crate) struct Deadline(AtomicU64);
+
+impl Deadline {
+    pub(crate) fn extend(&self, deadline: Instant) {
+        self.0.fetch_max(lease_nanos(deadline), Ordering::Relaxed);
+    }
+
+    pub(crate) fn get(&self) -> u64 {
+        self.0.load(Ordering::Relaxed)
+    }
+}
+
+/// Lease of this node: the latest deadline of the heartbeat responses that carry region leases.
+#[derive(Debug, Default)]
+pub(crate) struct NodeLease {
+    /// Whether writes need a lease. Only a datanode the metasrv controls sets it.
+    enabled: AtomicBool,
+    deadline: Deadline,
+    /// The time of the lease checks while a test freezes it, in the unit of [Deadline]; 0 when
+    /// it is not frozen.
+    #[cfg(any(test, feature = "test"))]
+    frozen_now: AtomicU64,
+}
+
+impl NodeLease {
+    pub(crate) fn enable(&self) {
+        self.enabled.store(true, Ordering::Relaxed);
+    }
+
+    pub(crate) fn is_enabled(&self) -> bool {
+        self.enabled.load(Ordering::Relaxed)
+    }
+
+    pub(crate) fn renew(&self, deadline: Instant) {
+        self.deadline.extend(deadline);
+    }
+
+    /// Returns the current time in the unit of [Deadline].
+    pub(crate) fn now(&self) -> u64 {
+        #[cfg(any(test, feature = "test"))]
+        {
+            let frozen_now = self.frozen_now.load(Ordering::Relaxed);
+            if frozen_now != 0 {
+                return frozen_now;
+            }
+        }
+        lease_nanos(Instant::now())
+    }
+
+    #[cfg(any(test, feature = "test"))]
+    pub(crate) fn freeze_clock(&self, at: Instant) {
+        self.frozen_now
+            .store(lease_nanos(at).max(1), Ordering::Relaxed);
+    }
+
+    /// Lets a region created or upgraded on this node write until the node's deadline. The
+    /// metasrv granted that deadline before the creation or upgrade, so a later migration of the
+    /// region waits past it. Later responses do not extend it; the region's Leader grants do.
+    pub(crate) fn seed(&self, leader_deadline: &Deadline) {
+        leader_deadline
+            .0
+            .fetch_max(self.deadline.get(), Ordering::Relaxed);
+    }
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct StagingPartitionInfo {
     pub(crate) partition_directive: StagingPartitionDirective,
@@ -2264,6 +2341,7 @@ mod tests {
             last_schedule_compaction_millis: Default::default(),
             time_provider: Arc::new(StdTimeProvider),
             topic_latest_entry_id: Default::default(),
+            leader_deadline: Default::default(),
             region_stats: RegionStats::new(),
             stats: ManifestStats::default(),
         }
@@ -2869,6 +2947,7 @@ mod tests {
             last_schedule_compaction_millis: Default::default(),
             time_provider: Arc::new(StdTimeProvider),
             topic_latest_entry_id: Default::default(),
+            leader_deadline: Default::default(),
             region_stats: RegionStats::new(),
             stats: ManifestStats::default(),
         };

src/mito2/src/region/opener.rs

@@ -459,6 +459,7 @@ impl RegionOpener {
             last_schedule_compaction_millis: AtomicI64::new(now),
             time_provider: self.time_provider.clone(),
             topic_latest_entry_id: AtomicU64::new(flushed_entry_id),
+            leader_deadline: Default::default(),
             region_stats: RegionStats::new(),
             stats: self.stats,
         }))
@@ -697,6 +698,7 @@ impl RegionOpener {
             last_schedule_compaction_millis: AtomicI64::new(now),
             time_provider: self.time_provider.clone(),
             topic_latest_entry_id: AtomicU64::new(topic_latest_entry_id),
+            leader_deadline: Default::default(),
             region_stats: RegionStats::new(),
             stats: self.stats.clone(),
         };

src/mito2/src/region_write_ctx.rs

@@ -27,6 +27,7 @@ use crate::error::{Error, Result, UnexpectedSnafu, WriteGroupSnafu};
 use crate::memtable::KeyValues;
 use crate::memtable::bulk::part::BulkPart;
 use crate::metrics;
+use crate::region::MitoRegionRef;
 use crate::region::version::{VersionControlData, VersionControlRef, VersionRef};
 use crate::request::OptionOutputTx;
 use crate::wal::{EntryId, WalWriter};
@@ -131,6 +132,9 @@ pub(crate) struct RegionWriteCtx {
     pub(crate) delete_num: usize,
     /// The total bytes written to the region.
     pub(crate) written_bytes: Option<Arc<AtomicU64>>,
+    /// The region and its Leader deadline when the write was prepared, if the write needs a
+    /// lease.
+    pub(crate) lease: Option<(MitoRegionRef, u64)>,
 }
 
 impl RegionWriteCtx {
@@ -167,6 +171,7 @@ impl RegionWriteCtx {
             delete_num: 0,
             bulk_parts: vec![],
             written_bytes,
+            lease: None,
         }
     }
 
@@ -247,6 +252,14 @@ impl RegionWriteCtx {
 
     /// Sets error and marks all write operations are failed.
     pub(crate) fn set_error(&mut self, err: Arc<Error>) {
+        self.fail_waiters(err);
+
+        // Fail the whole write operation.
+        self.failed = true;
+    }
+
+    /// Sets error for all notifiers. The writes still go to the memtables.
+    pub(crate) fn fail_waiters(&mut self, err: Arc<Error>) {
         // Set error for all notifiers.
         for notify in self
             .wal_notifiers
@@ -258,9 +271,6 @@ impl RegionWriteCtx {
         for notify in &mut self.bulk_notifiers {
             notify.err = Some(err.clone());
         }
-
-        // Fail the whole write operation.
-        self.failed = true;
     }
 
     /// Returns whether the write operation is already marked as failed.

src/mito2/src/worker.rs

@@ -71,8 +71,8 @@ use crate::memtable::MemtableBuilderProvider;
 use crate::metrics::{REGION_COUNT, REQUEST_WAIT_TIME, WRITE_STALLING};
 use crate::region::opener::PartitionExprFetcherRef;
 use crate::region::{
-    CatchupRegions, CatchupRegionsRef, MitoRegionRef, OpeningRegions, OpeningRegionsRef, RegionMap,
-    RegionMapRef,
+    CatchupRegions, CatchupRegionsRef, MitoRegionRef, NodeLease, OpeningRegions, OpeningRegionsRef,
+    RegionMap, RegionMapRef,
 };
 use crate::request::{
     BackgroundNotify, BulkInsertRequest, DdlRequest, OptionOutputTx, SenderBulkRequest,
@@ -157,6 +157,7 @@ pub(crate) struct WorkerGroup {
     file_ref_manager: FileReferenceManagerRef,
     /// Gc limiter to limit concurrent gc jobs.
     gc_limiter: GcLimiterRef,
+    node_lease: Arc<NodeLease>,
     /// Object store manager.
     object_store_manager: ObjectStoreManagerRef,
     /// Puffin manager factory.
@@ -250,6 +251,7 @@ impl WorkerGroup {
         #[cfg(feature = "hdfs-object-store")]
         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 workers = (0..config.num_workers)
             .map(|id| {
@@ -280,6 +282,7 @@ impl WorkerGroup {
                     flush_semaphore: flush_semaphore.clone(),
                     #[cfg(feature = "hdfs-object-store")]
                     sweep_limiter: sweep_limiter.clone(),
+                    node_lease: node_lease.clone(),
                 }
                 .start()
             })
@@ -294,6 +297,7 @@ impl WorkerGroup {
             cache_manager,
             file_ref_manager,
             gc_limiter,
+            node_lease,
             object_store_manager,
             puffin_manager_factory,
             intermediate_manager,
@@ -398,6 +402,10 @@ impl WorkerGroup {
         self.gc_limiter.clone()
     }
 
+    pub(crate) fn node_lease(&self) -> &NodeLease {
+        &self.node_lease
+    }
+
     /// Get worker for specific `region_id`.
     pub(crate) fn worker(&self, region_id: RegionId) -> &RegionWorker {
         let index = region_id_to_index(region_id, self.workers.len());
@@ -509,6 +517,7 @@ impl WorkerGroup {
         #[cfg(feature = "hdfs-object-store")]
         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 workers = (0..config.num_workers)
             .map(|id| {
                 WorkerStarter {
@@ -538,6 +547,7 @@ impl WorkerGroup {
                     flush_semaphore: flush_semaphore.clone(),
                     #[cfg(feature = "hdfs-object-store")]
                     sweep_limiter: sweep_limiter.clone(),
+                    node_lease: node_lease.clone(),
                 }
                 .start()
             })
@@ -552,6 +562,7 @@ impl WorkerGroup {
             cache_manager,
             file_ref_manager,
             gc_limiter,
+            node_lease,
             object_store_manager,
             puffin_manager_factory,
             intermediate_manager,
@@ -652,6 +663,7 @@ struct WorkerStarter<S> {
     /// Limiter for the temporary file sweeps, or `None` if the datanode allows
     /// no concurrent GC job.
     sweep_limiter: Option<Arc<Semaphore>>,
+    node_lease: Arc<NodeLease>,
 }
 
 impl<S: LogStore> WorkerStarter<S> {
@@ -754,6 +766,7 @@ impl<S: LogStore> WorkerStarter<S> {
             plugins: self.plugins,
             #[cfg(feature = "hdfs-object-store")]
             sweep_limiter: self.sweep_limiter,
+            node_lease: self.node_lease,
         };
         let handle = common_runtime::spawn_global(async move {
             worker_thread.run().await;
@@ -1085,6 +1098,7 @@ struct RegionWorkerLoop<S> {
     /// Limiter for the temporary file sweeps, or `None` if the datanode allows
     /// no concurrent GC job.
     sweep_limiter: Option<Arc<Semaphore>>,
+    node_lease: Arc<NodeLease>,
 }
 
 impl<S: LogStore> RegionWorkerLoop<S> {
@@ -1606,6 +1620,14 @@ impl WorkerListener {
         }
     }
 
+    /// A batch of writes is appended to the WAL.
+    pub(crate) async fn on_write_wal_end(&self) {
+        #[cfg(any(test, feature = "test"))]
+        if let Some(listener) = &self.listener {
+            listener.on_write_wal_end().await;
+        }
+    }
+
     pub(crate) async fn on_flush_begin(&self, region_id: RegionId) {
         #[cfg(any(test, feature = "test"))]
         if let Some(listener) = &self.listener {

src/mito2/src/worker/handle_catchup.rs

@@ -75,6 +75,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
         let sweep_older_than = self.config.gc.unknown_file_lingering_time;
         #[cfg(feature = "hdfs-object-store")]
         let sweep_limiter = self.sweep_limiter.clone();
+        let node_lease = self.node_lease.clone();
         common_runtime::spawn_global(async move {
             let mut task = RegionCatchupTask::new(region.clone(), wal, allow_stale_entries)
                 .with_entry_receiver(entry_receiver)
@@ -84,6 +85,9 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             match task.run().await {
                 Ok(_) => {
                     if request.set_writable {
+                        // The region takes writes under the node lease until its next Leader
+                        // grant.
+                        node_lease.seed(&region.leader_deadline);
                         region.set_role(RegionRole::Leader);
                         // Finalize leadership: persist backfilled metadata.
                         match region

src/mito2/src/worker/handle_create.rs

@@ -106,6 +106,8 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             hook.on_region_opened(region_id, &region.metadata()).await;
         }
 
+        // The region takes writes under the node lease until its first Leader grant.
+        self.node_lease.seed(&region.leader_deadline);
         // Insert the MitoRegion into the RegionMap.
         self.register_region(region);
         if let Some(state) = &self.series_index_task_state {

src/mito2/src/worker/handle_write.rs

@@ -25,8 +25,9 @@ use store_api::logstore::LogStore;
 use store_api::storage::RegionId;
 
 use crate::error::{
-    InvalidRequestSnafu, PartitionExprVersionMismatchSnafu, RegionNotFoundSnafu, RegionStateSnafu,
-    RejectWriteInStagingSnafu, RejectWriteSnafu, Result,
+    InvalidRequestSnafu, PartitionExprVersionMismatchSnafu, RegionLeaseExpiredSnafu,
+    RegionNotFoundSnafu, RegionStateSnafu, RejectWriteInStagingSnafu, RejectWriteSnafu, Result,
+    UnacknowledgedWriteSnafu,
 };
 use crate::metrics;
 use crate::metrics::{
@@ -118,6 +119,8 @@ impl<S: LogStore> RegionWorkerLoop<S> {
                 return;
             }
         }
+        self.listener.on_write_wal_end().await;
+        self.fail_writes_after_lease_expiry(&mut region_ctxs);
 
         let (mut put_rows, mut delete_rows) = (0, 0);
         // Write to memtables.
@@ -294,8 +297,10 @@ impl<S> RegionWorkerLoop<S> {
     ) -> HashMap<RegionId, RegionWriteCtx> {
         // Initialize region write context map.
         let mut region_ctxs = HashMap::new();
-        self.process_write_requests(&mut region_ctxs, write_requests);
-        self.process_bulk_requests(&mut region_ctxs, bulk_requests);
+        // One reading serves the batch, as preparing it does not wait.
+        let lease_now = self.node_lease.is_enabled().then(|| self.node_lease.now());
+        self.process_write_requests(&mut region_ctxs, write_requests, lease_now);
+        self.process_bulk_requests(&mut region_ctxs, bulk_requests, lease_now);
         region_ctxs
     }
 
@@ -303,6 +308,7 @@ impl<S> RegionWorkerLoop<S> {
         &mut self,
         region_ctxs: &mut HashMap<RegionId, RegionWriteCtx>,
         write_requests: &mut Vec<SenderWriteRequest>,
+        lease_now: Option<u64>,
     ) {
         for mut sender_req in write_requests.drain(..) {
             let region_id = sender_req.request.region_id;
@@ -334,6 +340,13 @@ impl<S> RegionWorkerLoop<S> {
                 match region.state() {
                     RegionRoleState::Leader(RegionLeaderState::Writable)
                     | RegionRoleState::Leader(RegionLeaderState::Staging) => {
+                        let leader_deadline = region.leader_deadline.get();
+                        if lease_now.is_some_and(|now| now >= leader_deadline) {
+                            sender_req
+                                .sender
+                                .send(RegionLeaseExpiredSnafu { region_id }.fail());
+                            continue;
+                        }
                         if region.reject_all_writes_in_staging() {
                             sender_req
                                 .sender
@@ -341,13 +354,14 @@ impl<S> RegionWorkerLoop<S> {
                             continue;
                         }
 
-                        let region_ctx = RegionWriteCtx::new(
+                        let mut region_ctx = RegionWriteCtx::new(
                             region.region_id,
                             &region.version_control,
                             region.provider.clone(),
                             Some(region.region_stats.written_bytes.clone()),
                             WriteSource::Request,
                         );
+                        region_ctx.lease = lease_now.map(|_| (region, leader_deadline));
 
                         e.insert(region_ctx);
                     }
@@ -460,6 +474,7 @@ impl<S> RegionWorkerLoop<S> {
         &mut self,
         region_ctxs: &mut HashMap<RegionId, RegionWriteCtx>,
         requests: &mut Vec<SenderBulkRequest>,
+        lease_now: Option<u64>,
     ) {
         let _timer = metrics::REGION_WORKER_HANDLE_WRITE_ELAPSED
             .with_label_values(&["prepare_bulk_request"])
@@ -482,19 +497,27 @@ impl<S> RegionWorkerLoop<S> {
                 match region.state() {
                     RegionRoleState::Leader(RegionLeaderState::Writable)
                     | RegionRoleState::Leader(RegionLeaderState::Staging) => {
+                        let leader_deadline = region.leader_deadline.get();
+                        if lease_now.is_some_and(|now| now >= leader_deadline) {
+                            bulk_req
+                                .sender
+                                .send(RegionLeaseExpiredSnafu { region_id }.fail());
+                            continue;
+                        }
                         if region.reject_all_writes_in_staging() {
                             bulk_req
                                 .sender
                                 .send(RejectWriteInStagingSnafu { region_id }.fail());
                             continue;
                         }
-                        let region_ctx = RegionWriteCtx::new(
+                        let mut region_ctx = RegionWriteCtx::new(
                             region.region_id,
                             &region.version_control,
                             region.provider.clone(),
                             Some(region.region_stats.written_bytes.clone()),
                             WriteSource::Request,
                         );
+                        region_ctx.lease = lease_now.map(|_| (region, leader_deadline));
 
                         e.insert(region_ctx);
                     }
@@ -573,6 +596,30 @@ impl<S> RegionWorkerLoop<S> {
         }
     }
 
+    /// Fails the waiters of writes whose WAL append completed at or after their region's Leader
+    /// deadline. The writes still go to the memtables, as their entries are in the WAL.
+    fn fail_writes_after_lease_expiry(&self, region_ctxs: &mut HashMap<RegionId, RegionWriteCtx>) {
+        if !self.node_lease.is_enabled() {
+            return;
+        }
+        let now = self.node_lease.now();
+        for (region_id, region_ctx) in region_ctxs.iter_mut() {
+            // A Leader grant may have extended the deadline during the append.
+            let expired = region_ctx.lease.as_ref().is_some_and(|(region, deadline)| {
+                now >= *deadline && now >= region.leader_deadline.get()
+            });
+            if !expired || region_ctx.is_failed() {
+                continue;
+            }
+            region_ctx.fail_waiters(Arc::new(
+                UnacknowledgedWriteSnafu {
+                    region_id: *region_id,
+                }
+                .build(),
+            ));
+        }
+    }
+
     /// Returns true if the engine needs to reject some write requests.
     pub(crate) fn should_reject_write(&self) -> bool {
         // If memory usage reaches high threshold (we should also consider stalled requests) returns true.

src/query/src/optimizer/test_util.rs

@@ -103,6 +103,15 @@ impl RegionEngine for MetaRegionEngine {
         unimplemented!()
     }
 
+    fn set_region_role_with_lease(
+        &self,
+        _region_id: RegionId,
+        _role: RegionRole,
+        _deadline: std::time::Instant,
+    ) -> Result<(), BoxedError> {
+        unimplemented!()
+    }
+
     async fn set_region_role_state_gracefully(
         &self,
         _region_id: RegionId,

src/store-api/src/region_engine.rs

@@ -18,6 +18,7 @@ use std::any::Any;
 use std::collections::HashMap;
 use std::fmt::{Debug, Display};
 use std::sync::{Arc, Mutex};
+use std::time::Instant;
 
 use api::greptime_proto::v1::meta::{GrantedRegion as PbGrantedRegion, RegionRole as PbRegionRole};
 use api::region::RegionResponse;
@@ -1081,6 +1082,18 @@ pub trait RegionEngine: Send + Sync {
     /// take effect.
     fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError>;
 
+    /// Sets [RegionRole] for a region from a lease the metasrv grants until `deadline`.
+    ///
+    /// The region must not acknowledge a write whose WAL append completes at or after the
+    /// deadline of its latest Leader or StagingLeader lease. Leases of other roles do not extend
+    /// that deadline.
+    fn set_region_role_with_lease(
+        &self,
+        region_id: RegionId,
+        role: RegionRole,
+        deadline: Instant,
+    ) -> Result<(), BoxedError>;
+
     /// Syncs the region manifest to the given manifest version.
     async fn sync_region(
         &self,