diff --git a/src/datanode/src/alive_keeper.rs b/src/datanode/src/alive_keeper.rs index dbfb699036c..bef59d715c5 100644 --- a/src/datanode/src/alive_keeper.rs +++ b/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(®ion_lease.regions, deadline) .await; self.close_staled_regions(®ion_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) { + 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(); diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index df12cf4f2a8..c304aaa6e00 100644 --- a/src/datanode/src/datanode.rs +++ b/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()); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 3acf188ad19..06b9d42fd00 100644 --- a/src/datanode/src/region_server.rs +++ b/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(®ion_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]: diff --git a/src/datanode/src/region_server/flow_write.rs b/src/datanode/src/region_server/flow_write.rs index 6bcb112bc78..bc9bf65d69d 100644 --- a/src/datanode/src/region_server/flow_write.rs +++ b/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, diff --git a/src/datanode/src/tests.rs b/src/datanode/src/tests.rs index 793ae057076..4fa247d07ca 100644 --- a/src/datanode/src/tests.rs +++ b/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, diff --git a/src/file-engine/src/engine.rs b/src/file-engine/src/engine.rs index 05be2937703..be1d141bcbd 100644 --- a/src/file-engine/src/engine.rs +++ b/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, diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index af15ea087c9..0c1b1114448 100644 --- a/src/metric-engine/src/engine.rs +++ b/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, diff --git a/src/metric-engine/src/engine/alter.rs b/src/metric-engine/src/engine/alter.rs index 2bf59edda0f..e940fa4cbbb 100644 --- a/src/metric-engine/src/engine/alter.rs +++ b/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) } diff --git a/src/metric-engine/src/engine/create.rs b/src/metric-engine/src/engine/create.rs index 2f27a405b36..1061565a80e 100644 --- a/src/metric-engine/src/engine/create.rs +++ b/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) diff --git a/src/metric-engine/src/engine/drop.rs b/src/metric-engine/src/engine/drop.rs index 7e98effdeee..f7db294478c 100644 --- a/src/metric-engine/src/engine/drop.rs +++ b/src/metric-engine/src/engine/drop.rs @@ -157,10 +157,24 @@ impl MetricEngineInner { logical_region_id: RegionId, physical_region_id: RegionId, ) -> Result { - // 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 diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index 27ee3fba56c..e8098131612 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/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(); + } } diff --git a/src/metric-engine/src/metadata_region.rs b/src/metric-engine/src/metadata_region.rs index 18fa51f3597..7f797ff9727 100644 --- a/src/metric-engine/src/metadata_region.rs +++ b/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(()) } diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 92d848f5677..0147fb71e2a 100644 --- a/src/mito2/src/engine.rs +++ b/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, diff --git a/src/mito2/src/engine/lease_test.rs b/src/mito2/src/engine/lease_test.rs new file mode 100644 index 00000000000..a0e68f7e94b --- /dev/null +++ b/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, +) -> (Arc, MitoEngine, Vec) { + 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> { + 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> { + 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, 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::().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) { + 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(); +} diff --git a/src/mito2/src/engine/listener.rs b/src/mito2/src/engine/listener.rs index d5e4b8e8854..f754d4faf00 100644 --- a/src/mito2/src/engine/listener.rs +++ b/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; + } + } +} diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 9bb908720fd..a336fb89673 100644 --- a/src/mito2/src/error.rs +++ b/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, .. diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 613b1ce428e..e6e3de2df1e 100644 --- a/src/mito2/src/region.rs +++ b/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; +/// The instant lease deadlines count from, so that a deadline fits in an [AtomicU64]. +static LEASE_EPOCH: LazyLock = 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(), }; diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index d10545198f7..947844c8f20 100644 --- a/src/mito2/src/region/opener.rs +++ b/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(), }; diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index b0b781eb43c..e4e1e505326 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/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>, + /// 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) { + 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) { // 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. diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 789b0f539b1..50fb5bc0eb1 100644 --- a/src/mito2/src/worker.rs +++ b/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, /// 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 { /// Limiter for the temporary file sweeps, or `None` if the datanode allows /// no concurrent GC job. sweep_limiter: Option>, + node_lease: Arc, } impl WorkerStarter { @@ -754,6 +766,7 @@ impl WorkerStarter { 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 { /// Limiter for the temporary file sweeps, or `None` if the datanode allows /// no concurrent GC job. sweep_limiter: Option>, + node_lease: Arc, } impl RegionWorkerLoop { @@ -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 { diff --git a/src/mito2/src/worker/handle_catchup.rs b/src/mito2/src/worker/handle_catchup.rs index f817c1309d8..92874672789 100644 --- a/src/mito2/src/worker/handle_catchup.rs +++ b/src/mito2/src/worker/handle_catchup.rs @@ -75,6 +75,7 @@ impl RegionWorkerLoop { 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 RegionWorkerLoop { 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(®ion.leader_deadline); region.set_role(RegionRole::Leader); // Finalize leadership: persist backfilled metadata. match region diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index c31e9d4ac7c..15641385e59 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/src/mito2/src/worker/handle_create.rs @@ -106,6 +106,8 @@ impl RegionWorkerLoop { hook.on_region_opened(region_id, ®ion.metadata()).await; } + // The region takes writes under the node lease until its first Leader grant. + self.node_lease.seed(®ion.leader_deadline); // Insert the MitoRegion into the RegionMap. self.register_region(region); if let Some(state) = &self.series_index_task_state { diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 5de0439e26e..7e290a768fb 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/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 RegionWorkerLoop { 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 RegionWorkerLoop { ) -> HashMap { // 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 RegionWorkerLoop { &mut self, region_ctxs: &mut HashMap, write_requests: &mut Vec, + lease_now: Option, ) { for mut sender_req in write_requests.drain(..) { let region_id = sender_req.request.region_id; @@ -334,6 +340,13 @@ impl RegionWorkerLoop { 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 RegionWorkerLoop { continue; } - let region_ctx = RegionWriteCtx::new( + let mut region_ctx = RegionWriteCtx::new( region.region_id, ®ion.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 RegionWorkerLoop { &mut self, region_ctxs: &mut HashMap, requests: &mut Vec, + lease_now: Option, ) { let _timer = metrics::REGION_WORKER_HANDLE_WRITE_ELAPSED .with_label_values(&["prepare_bulk_request"]) @@ -482,19 +497,27 @@ impl RegionWorkerLoop { 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, ®ion.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 RegionWorkerLoop { } } + /// 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) { + 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. diff --git a/src/query/src/optimizer/test_util.rs b/src/query/src/optimizer/test_util.rs index 6542539570c..a6b95ffdeba 100644 --- a/src/query/src/optimizer/test_util.rs +++ b/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, diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index c95bf6d3b21..6dffdd03c3b 100644 --- a/src/store-api/src/region_engine.rs +++ b/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,