diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index c54691a5a8f..2576158bab6 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -1882,7 +1882,10 @@ async fn restart_without_region_close( /// Builds a bulk insert request for a region created by [CreateRequestBuilder] /// with its default columns `(tag_0, field_0, ts)`. -fn build_bulk_insert_request(region_id: RegionId, num_rows: usize) -> RegionBulkInsertsRequest { +pub(crate) fn build_bulk_insert_request( + region_id: RegionId, + num_rows: usize, +) -> RegionBulkInsertsRequest { let schema = Arc::new(Schema::new(vec![ Field::new("tag_0", DataType::Utf8, true), Field::new("field_0", DataType::Float64, true), diff --git a/src/mito2/src/engine/object_store_wal_test.rs b/src/mito2/src/engine/object_store_wal_test.rs index cd27c63b48d..5a61d8b6733 100644 --- a/src/mito2/src/engine/object_store_wal_test.rs +++ b/src/mito2/src/engine/object_store_wal_test.rs @@ -79,13 +79,13 @@ fn open_request(table_dir: String) -> RegionOpenRequest { } } -async fn new_engine(env: &mut TestEnv) -> (MitoEngine, Arc) { +pub(crate) async fn new_engine(env: &mut TestEnv) -> (MitoEngine, Arc) { env.create_engine_with_object_store_wal(MitoConfig::default(), &wal_config()) .await } /// Creates a region with object store WAL options and returns its table dir and row schema. -async fn create_region( +pub(crate) async fn create_region( engine: &MitoEngine, region_id: RegionId, ) -> (String, Vec) { diff --git a/src/mito2/src/engine/set_role_state_test.rs b/src/mito2/src/engine/set_role_state_test.rs index 53aa3511459..94a52c4f3c6 100644 --- a/src/mito2/src/engine/set_role_state_test.rs +++ b/src/mito2/src/engine/set_role_state_test.rs @@ -12,23 +12,36 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; use std::time::Duration; -use api::v1::Rows; -use common_error::ext::ErrorExt; +use api::v1::{Rows, SemanticType}; +use common_error::ext::{ErrorExt, RetryHint}; use common_error::status_code::StatusCode; +use common_recordbatch::RecordBatches; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::ColumnSchema; +use store_api::metadata::ColumnMetadata; use store_api::region_engine::{ RegionEngine, RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState, }; use store_api::region_request::{ - EnterStagingRequest, RegionFlushRequest, RegionPutRequest, RegionRequest, + AddColumn, AffectedRows, AlterKind, EnterStagingRequest, PathType, RegionAlterRequest, + RegionCatchupRequest, RegionFlushRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, StagingPartitionDirective, }; -use store_api::storage::RegionId; +use store_api::storage::{RegionId, ScanRequest}; +use tokio::sync::oneshot; use crate::config::MitoConfig; -use crate::region::{RegionLeaderState, RegionRoleState}; +use crate::engine::MitoEngine; +use crate::engine::basic_test::build_bulk_insert_request; +use crate::engine::listener::NotifyRegionChangeResultListener; +use crate::engine::object_store_wal_test; +use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState}; +use crate::region_write_ctx::test_hooks::arm_bulk_install_barrier; +use crate::request::WorkerRequest; use crate::test_util::{ CheckpointTaskBlocker, CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema, }; @@ -1001,3 +1014,520 @@ async fn test_restricted_state_transitions_with_format(flat_format: bool) { .unwrap(); assert_success_response(&result, 0); } + +const WAIT: Duration = Duration::from_secs(10); + +type WriteReceiver = oneshot::Receiver>; + +/// Submits a graceful role change straight to the region's worker. Requests +/// submitted after this returns reach the worker after it. +async fn submit_role_change( + engine: &MitoEngine, + region_id: RegionId, + state: SettableRegionRoleState, +) -> oneshot::Receiver { + let (request, receiver) = WorkerRequest::new_set_readonly_gracefully(region_id, state); + engine + .inner + .workers + .submit_to_worker(region_id, request) + .await + .unwrap(); + receiver +} + +/// Submits a region request straight to the region's worker. +async fn submit_request( + engine: &MitoEngine, + region_id: RegionId, + request: RegionRequest, +) -> WriteReceiver { + let metadata = engine.inner.get_metadata(region_id).ok(); + let (request, receiver) = + WorkerRequest::try_from_region_request(region_id, request, metadata).unwrap(); + engine + .inner + .workers + .submit_to_worker(region_id, request) + .await + .unwrap(); + receiver +} + +fn put_request(rows: Rows) -> RegionRequest { + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows, + hint: None, + partition_expr_version: None, + }) +} + +fn reported_last_entry_id(response: &SetRegionRoleStateResponse) -> u64 { + match response { + SetRegionRoleStateResponse::Success(success) => success.last_entry_id().unwrap(), + _ => panic!("Expected success response, got: {:?}", response), + } +} + +/// Lets the role change task waiting behind `manifest_guard` run and returns +/// the `last_entry_id` it reports. +async fn finish_role_change( + manifest_guard: G, + role_change: oneshot::Receiver, +) -> u64 { + drop(manifest_guard); + let response = tokio::time::timeout(WAIT, role_change) + .await + .expect("role change did not finish") + .unwrap(); + reported_last_entry_id(&response) +} + +/// Waits for the write's result if it is still in flight. +async fn write_result( + early: Option>, + write: WriteReceiver, +) -> crate::error::Result { + match early { + Some(result) => result, + None => tokio::time::timeout(WAIT, write) + .await + .expect("write did not finish") + .unwrap(), + } +} + +/// Checks the promise of `set_region_role_state_gracefully`: an acknowledged +/// write must be covered by the `last_entry_id` the call returned. +fn assert_write_covered( + region: &MitoRegionRef, + result: crate::error::Result, + reported: u64, +) { + let last_entry_id = region.version_control.current().last_entry_id; + assert!( + result.is_err() || last_entry_id <= reported, + "a write was acknowledged ({result:?}) after the role change returned last_entry_id {reported}, \ + but the region's last_entry_id is now {last_entry_id}" + ); +} + +/// The role change task waits for the manifest lock before it changes the +/// role, while the worker keeps handling requests. A bulk write handled after +/// the role change request must be rejected, or be covered by the returned +/// `last_entry_id` if it was accepted. +#[tokio::test(flavor = "multi_thread")] +async fn test_set_follower_gracefully_covers_inflight_bulk_write() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("follower-inflight-bulk-write").await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 3), + }, + ) + .await; + + let region = engine.get_region(region_id).unwrap(); + let manifest_guard = region.manifest_ctx.manifest_manager.read().await; + let role_change = + submit_role_change(&engine, region_id, SettableRegionRoleState::Follower).await; + // An accepted write pauses after its WAL append, before installing its rows. + let mut barrier = arm_bulk_install_barrier(region_id, region.version_control.clone()); + let mut write = submit_request( + &engine, + region_id, + RegionRequest::BulkInserts(build_bulk_insert_request(region_id, 2)), + ) + .await; + let early = tokio::select! { + _ = barrier.wait_until_reached() => None, + result = &mut write => Some(result.unwrap()), + _ = tokio::time::sleep(WAIT) => panic!("bulk write never reached the install barrier"), + }; + let reported = finish_role_change(manifest_guard, role_change).await; + assert_eq!(RegionRoleState::Follower, region.state()); + + barrier.release(); + let result = write_result(early, write).await; + assert_write_covered(®ion, result, reported); +} + +/// Same as [test_set_follower_gracefully_covers_inflight_bulk_write] for a row +/// write. An accepted put waits in its WAL append: the object store WAL +/// acknowledges an append once its object exists, and the test parks that +/// object. +#[tokio::test(flavor = "multi_thread")] +async fn test_set_follower_gracefully_covers_inflight_row_write() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("follower-inflight-row-write").await; + let (engine, store) = object_store_wal_test::new_engine(&mut env).await; + let region_id = RegionId::new(1, 1); + let (_, schema) = object_store_wal_test::create_region(&engine, region_id).await; + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows(0, 3), + }, + ) + .await; + + let region = engine.get_region(region_id).unwrap(); + let manifest_guard = region.manifest_ctx.manifest_manager.read().await; + let role_change = + submit_role_change(&engine, region_id, SettableRegionRoleState::Follower).await; + store.hold_creates(); + let mut write = submit_request( + &engine, + region_id, + put_request(Rows { + schema, + rows: build_rows(3, 5), + }), + ) + .await; + let early = tokio::select! { + result = store.wait_for_parked_creates(1) => { + result.unwrap(); + None + } + result = &mut write => Some(result.unwrap()), + _ = tokio::time::sleep(WAIT) => panic!("put never reached the WAL"), + }; + let reported = finish_role_change(manifest_guard, role_change).await; + assert_eq!(RegionRoleState::Follower, region.state()); + + store.release_creates(); + let result = write_result(early, write).await; + assert_write_covered(®ion, result, reported); +} + +/// A lease renewal sets Downgrading outside the worker, so a write that passed +/// the state check before it can still be in flight. The datanode's downgrade +/// steps (DowngradingLeader, Flush, Follower) must still flush that write and +/// report its entry. +#[tokio::test(flavor = "multi_thread")] +async fn test_downgrade_flush_follower_covers_inflight_write() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("downgrade-flush-follower").await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + let mut barrier = arm_bulk_install_barrier(region_id, region.version_control.clone()); + let write = submit_request( + &engine, + region_id, + RegionRequest::BulkInserts(build_bulk_insert_request(region_id, 2)), + ) + .await; + tokio::time::timeout(WAIT, barrier.wait_until_reached()) + .await + .expect("bulk write never reached the install barrier"); + engine + .set_region_role(region_id, RegionRole::DowngradingLeader) + .unwrap(); + + let downgrade = submit_role_change( + &engine, + region_id, + SettableRegionRoleState::DowngradingLeader, + ) + .await; + barrier.release(); + assert_eq!(2, write.await.unwrap().unwrap()); + reported_last_entry_id(&downgrade.await.unwrap()); + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + + let response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower) + .await + .unwrap(); + let reported = reported_last_entry_id(&response); + let current = region.version_control.current(); + assert_eq!(current.last_entry_id, reported); + assert_eq!(current.version.flushed_entry_id, reported); + assert!(current.version.memtables.is_empty()); +} + +fn add_nullable_field1() -> RegionAlterRequest { + RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: vec![AddColumn { + column_metadata: ColumnMetadata { + column_schema: ColumnSchema::new( + "field_1", + ConcreteDataType::float64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 3, + }, + location: None, + }], + }, + } +} + +async fn wait_for_state(region: &MitoRegionRef, state: RegionRoleState) { + tokio::time::timeout(WAIT, async { + while region.state() != state { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("region never reached {state:?}, now {:?}", region.state())); +} + +async fn count_rows(engine: &MitoEngine, region_id: RegionId) -> usize { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + batches.iter().map(|batch| batch.num_rows()).sum() +} + +/// Creates a region with rows `[0, 3)` in an SST and starts an ALTER whose +/// result the worker does not handle until `listener` is woken, so the region +/// stays in Altering. Returns the ALTER task. +async fn create_region_and_start_alter( + env: &TestEnv, + engine: &MitoEngine, + region_id: RegionId, + request: store_api::region_request::RegionCreateRequest, +) -> tokio::task::JoinHandle<()> { + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 3), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + + let cloned_engine = engine.clone(); + let alter = tokio::spawn(async move { + cloned_engine + .handle_request(region_id, RegionRequest::Alter(add_nullable_field1())) + .await + .unwrap(); + }); + let region = engine.get_region(region_id).unwrap(); + wait_for_state( + ®ion, + RegionRoleState::Leader(RegionLeaderState::Altering), + ) + .await; + alter +} + +/// DowngradingLeader while an ALTER keeps the region in Altering must fail +/// with a retryable error, because the region becomes writable again once the +/// ALTER finishes. A retry after the ALTER fences writes. +#[tokio::test(flavor = "multi_thread")] +async fn test_downgrading_leader_during_alter_stops_writes() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("downgrading-during-alter").await; + let listener = Arc::new(NotifyRegionChangeResultListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + let alter = create_region_and_start_alter(&env, &engine, region_id, request).await; + let region = engine.get_region(region_id).unwrap(); + + let response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader) + .await + .unwrap(); + let state_after_downgrade = region.state(); + listener.wake_notify(); + tokio::time::timeout(WAIT, alter).await.unwrap().unwrap(); + match &response { + SetRegionRoleStateResponse::InvalidTransition(err) => { + assert_eq!(RetryHint::Retryable, err.retry_hint(), "{err:?}") + } + _ => panic!( + "DowngradingLeader returned {response:?} with the region in {state_after_downgrade:?}" + ), + } + + let response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader) + .await + .unwrap(); + reported_last_entry_id(&response); + let err = engine + .handle_request( + region_id, + put_request(Rows { + schema: column_schemas, + rows: build_rows(3, 5), + }), + ) + .await + .unwrap_err(); + assert_eq!(StatusCode::RegionNotReady, err.status_code()); +} + +/// The datanode's downgrade steps (DowngradingLeader, Flush, Follower) with +/// the DowngradingLeader step overlapping an ALTER. The region accepts writes +/// until a downgrade takes effect. A migration candidate opens the region from +/// the object store with its own local WAL, so the final flush must carry +/// every write the old leader acknowledged. +#[tokio::test(flavor = "multi_thread")] +async fn test_downgrade_during_alter_keeps_acknowledged_writes_with_local_wal() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("downgrade-during-alter-local-wal").await; + let listener = Arc::new(NotifyRegionChangeResultListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + let table_dir = request.table_dir.clone(); + let options = request.options.clone(); + let alter = create_region_and_start_alter(&env, &engine, region_id, request).await; + + // Mirrors `DowngradeRegionsHandler::handle_downgrade_region` in the + // datanode. Metasrv retries the downgrade on a retryable error. + let mut response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader) + .await + .unwrap(); + listener.wake_notify(); + tokio::time::timeout(WAIT, alter).await.unwrap().unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(3, 5), + }, + ) + .await; + let mut acknowledged_rows = 5; + if let SetRegionRoleStateResponse::InvalidTransition(err) = &response { + assert_eq!(RetryHint::Retryable, err.retry_hint(), "{err:?}"); + response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader) + .await + .unwrap(); + } + reported_last_entry_id(&response); + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + if let Ok(output) = engine + .handle_request( + region_id, + put_request(Rows { + schema: column_schemas, + rows: build_rows(5, 8), + }), + ) + .await + { + acknowledged_rows += output.affected_rows; + } + let response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower) + .await + .unwrap(); + let reported = reported_last_entry_id(&response); + + // The candidate shares the object store but not the local WAL. + let mut candidate_env = TestEnv::with_prefix("downgrade-during-alter-candidate").await; + candidate_env.object_store_manager = env.get_object_store_manager(); + candidate_env.log_store = Some(candidate_env.create_log_store().await); + let candidate = candidate_env.new_mito_engine(MitoConfig::default()).await; + candidate + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + options, + skip_wal_replay: true, + path_type: PathType::Bare, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + candidate + .handle_request( + region_id, + RegionRequest::Catchup(RegionCatchupRequest { + set_writable: true, + entry_id: Some(reported), + metadata_entry_id: None, + checkpoint: None, + }), + ) + .await + .unwrap(); + assert_eq!( + acknowledged_rows, + count_rows(&candidate, region_id).await, + "the candidate misses rows the old leader acknowledged; Follower reported last_entry_id {reported}" + ); +} diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index d4753fe14bf..ba32f7aff56 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -599,12 +599,24 @@ impl MitoRegion { info!("Region {} already in downgrading mode", self.region_id); wait_for_checkpoint = true; } - _ => { + RegionRoleState::Follower => { warn!( "Cannot start downgrade for region {} from state {:?}", self.region_id, current_state ); } + RegionRoleState::Leader(_) => { + // The region may become writable again when the + // current operation, such as an ALTER, finishes. Fail + // with a retryable error instead of reporting a + // downgrade that does not fence writes. + return Err(RegionStateSnafu { + region_id: self.region_id, + state: current_state, + expect: RegionRoleState::Leader(RegionLeaderState::Writable), + } + .build()); + } } None } diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 4bf9d4b244a..763a068e106 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -52,7 +52,7 @@ use rand::{Rng, rng}; use snafu::{ResultExt, ensure}; use store_api::logstore::LogStore; use store_api::region_engine::{ - SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState, + RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState, }; use store_api::storage::{FileId, RegionId}; use tokio::sync::mpsc::{Receiver, Sender}; @@ -1417,6 +1417,15 @@ impl RegionWorkerLoop { sender: oneshot::Sender, ) { if let Some(region) = self.regions.get_region(region_id) { + if matches!( + region_role_state, + SettableRegionRoleState::Follower | SettableRegionRoleState::DowngradingLeader + ) { + // Fence writes here: writes the worker handles later see + // Downgrading, and earlier ones have finished. This only + // changes a writable or staging leader. + region.set_role(RegionRole::DowngradingLeader); + } // We need to do this in background as we need the manifest lock. common_runtime::spawn_global(async move { match region.set_role_state_gracefully(region_role_state).await {