diff --git a/src/datanode/src/alive_keeper.rs b/src/datanode/src/alive_keeper.rs index 06d82c932e1..fd6d7da8c43 100644 --- a/src/datanode/src/alive_keeper.rs +++ b/src/datanode/src/alive_keeper.rs @@ -26,7 +26,7 @@ use common_meta::heartbeat::handler::{ }; use common_telemetry::{debug, error, info, trace, warn}; use snafu::OptionExt; -use store_api::region_engine::RegionRole; +use store_api::region_engine::{RegionRole, SettableRegionRoleState}; use store_api::region_request::{RegionCloseRequest, RegionRequest}; use store_api::storage::RegionId; #[cfg(test)] @@ -423,17 +423,27 @@ 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 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 // during the short period before it is removed. warn!(err; "Failed to set region role to {role} for region {region_id}"); - }else{ + } else { error!(err; "Failed to set region role to {role} for region {region_id}"); } + // Finalize leadership: persist backfilled metadata. + } else if !prev_role_leader && role == RegionRole::Leader + && let Err(err) = self + .region_server + .set_region_role_state_gracefully(self.region_id, SettableRegionRoleState::Leader) + .await + { + error!(err; "Failed to set region role state gracefully to {role} for region {region_id}"); } + if let Some(ext_handler) = self.handler_ext.as_ref() { ext_handler.reset_deadline( &self.region_server, diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index 97aff454486..5c9e86848f2 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -39,6 +39,7 @@ use meta_client::MetaClientRef; use metric_engine::engine::MetricEngine; use mito2::config::MitoConfig; use mito2::engine::{MitoEngine, MitoEngineBuilder}; +use mito2::region::opener::PartitionExprFetcherRef; use mito2::sst::file_ref::{FileReferenceManager, FileReferenceManagerRef}; use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef}; use object_store::util::normalize_dir; @@ -48,7 +49,9 @@ use servers::export_metrics::ExportMetricsTask; use servers::server::ServerHandlers; use snafu::{OptionExt, ResultExt, ensure}; use store_api::path_utils::WAL_DIR; -use store_api::region_engine::{RegionEngineRef, RegionRole}; +use store_api::region_engine::{ + RegionEngineRef, RegionRole, SetRegionRoleStateResponse, SettableRegionRoleState, +}; use tokio::fs; use tokio::sync::Notify; @@ -64,6 +67,7 @@ use crate::event_listener::{ }; use crate::greptimedb_telemetry::get_greptimedb_telemetry_task; use crate::heartbeat::HeartbeatTask; +use crate::partition_expr_fetcher::MetaPartitionExprFetcher; use crate::region_server::{DummyTableProviderFactory, RegionServer}; use crate::store::{self, new_object_store_without_cache}; use crate::utils::{RegionOpenRequests, build_region_open_requests}; @@ -448,12 +452,15 @@ impl DatanodeBuilder { } } + // Build a fetcher to backfill partition_expr on open. + let fetcher = Arc::new(MetaPartitionExprFetcher::new(self.kv_backend.clone())); let mito_engine = self .build_mito_engine( object_store_manager.clone(), mito_engine_config, schema_metadata_manager.clone(), file_ref_manager.clone(), + fetcher.clone(), plugins.clone(), ) .await?; @@ -480,6 +487,7 @@ impl DatanodeBuilder { mut config: MitoConfig, schema_metadata_manager: SchemaMetadataManagerRef, file_ref_manager: FileReferenceManagerRef, + partition_expr_fetcher: PartitionExprFetcherRef, plugins: Plugins, ) -> Result { let opts = &self.opts; @@ -502,6 +510,7 @@ impl DatanodeBuilder { object_store_manager, schema_metadata_manager, file_ref_manager, + partition_expr_fetcher.clone(), plugins, ); @@ -543,6 +552,7 @@ impl DatanodeBuilder { object_store_manager, schema_metadata_manager, file_ref_manager, + partition_expr_fetcher, plugins, ); @@ -655,12 +665,25 @@ async fn open_all_regions( } for region_id in open_regions { - if open_with_writable - && let Err(e) = region_server.set_region_role(region_id, RegionRole::Leader) - { - error!( - e; "failed to convert region {region_id} to leader" - ); + if open_with_writable { + let res = region_server.set_region_role(region_id, RegionRole::Leader); + match res { + Ok(_) => { + // Finalize leadership: persist backfilled metadata. + if let SetRegionRoleStateResponse::InvalidTransition(err) = region_server + .set_region_role_state_gracefully( + region_id, + SettableRegionRoleState::Leader, + ) + .await? + { + error!(err; "failed to convert region {region_id} to leader"); + } + } + Err(e) => { + error!(e; "failed to convert region {region_id} to leader"); + } + } } } diff --git a/src/datanode/src/lib.rs b/src/datanode/src/lib.rs index bc63d71c8a0..8ca9a27c913 100644 --- a/src/datanode/src/lib.rs +++ b/src/datanode/src/lib.rs @@ -23,6 +23,7 @@ pub mod event_listener; mod greptimedb_telemetry; pub mod heartbeat; pub mod metrics; +mod partition_expr_fetcher; pub mod region_server; pub mod service; pub mod store; diff --git a/src/datanode/src/partition_expr_fetcher.rs b/src/datanode/src/partition_expr_fetcher.rs new file mode 100644 index 00000000000..1d0b81e32ce --- /dev/null +++ b/src/datanode/src/partition_expr_fetcher.rs @@ -0,0 +1,91 @@ +// 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. + +use common_meta::key::table_route::TableRouteManager; +use common_meta::kv_backend::KvBackendRef; +use mito2::region::opener::PartitionExprFetcher; +use store_api::storage::RegionId; + +/// A fetcher to fetch partition exprs from the table route. +pub struct MetaPartitionExprFetcher { + table_route_manager: TableRouteManager, +} + +impl MetaPartitionExprFetcher { + /// Creates a new [MetaPartitionExprFetcher]. + pub fn new(kv: KvBackendRef) -> Self { + Self { + table_route_manager: TableRouteManager::new(kv), + } + } +} + +#[async_trait::async_trait] +impl PartitionExprFetcher for MetaPartitionExprFetcher { + async fn fetch_expr(&self, region_id: RegionId) -> Option { + let table_id = region_id.table_id(); + let Ok((_, route)) = self + .table_route_manager + .get_physical_table_route(table_id) + .await + else { + return None; + }; + let region_number = region_id.region_number(); + let rr = route + .region_routes + .iter() + .find(|r| r.region.id.region_number() == region_number)?; + Some(rr.region.partition_expr()) + } +} + +#[cfg(test)] +mod tests { + use common_meta::key::table_route::{TableRouteStorage, TableRouteValue}; + use common_meta::kv_backend::TxnService; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::rpc::router::{Region, RegionRoute}; + + use super::*; + + #[tokio::test] + async fn test_fetch_expr_json() { + let kv = std::sync::Arc::new(MemoryKvBackend::new()); + let storage = TableRouteStorage::new(kv.clone()); + let table_id_u32 = 42; + let region_number = 7; + let region_id = RegionId::new(table_id_u32, region_number); + let expr_json = + r#"{"Expr":{"lhs":{"Column":"a"},"op":"GtEq","rhs":{"Value":{"UInt32":10}}}}"#; + + let region_route = RegionRoute { + region: Region { + id: region_id, + name: "r".to_string(), + attrs: Default::default(), + partition: None, + partition_expr: expr_json.to_string(), + }, + ..Default::default() + }; + let trv = TableRouteValue::physical(vec![region_route]); + let (txn, _dec) = storage.build_create_txn(table_id_u32, &trv).unwrap(); + kv.txn(txn).await.unwrap(); + + let fetcher = MetaPartitionExprFetcher::new(kv); + let fetched = fetcher.fetch_expr(region_id).await; + assert_eq!(fetched.as_deref(), Some(expr_json)); + } +} diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 5616a808ddd..d50c1a28bd2 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -111,6 +111,7 @@ use crate::metrics::HANDLE_REQUEST_ELAPSED; use crate::read::scan_region::{ScanRegion, Scanner}; use crate::read::stream::ScanBatchStream; use crate::region::MitoRegionRef; +use crate::region::opener::PartitionExprFetcherRef; use crate::request::{RegionEditRequest, WorkerRequest}; use crate::sst::file::FileMeta; use crate::sst::file_ref::FileReferenceManagerRef; @@ -129,12 +130,14 @@ pub struct MitoEngineBuilder<'a, S: LogStore> { object_store_manager: ObjectStoreManagerRef, schema_metadata_manager: SchemaMetadataManagerRef, file_ref_manager: FileReferenceManagerRef, + partition_expr_fetcher: PartitionExprFetcherRef, plugins: Plugins, #[cfg(feature = "enterprise")] extension_range_provider_factory: Option, } impl<'a, S: LogStore> MitoEngineBuilder<'a, S> { + #[allow(clippy::too_many_arguments)] pub fn new( data_home: &'a str, config: MitoConfig, @@ -142,6 +145,7 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> { object_store_manager: ObjectStoreManagerRef, schema_metadata_manager: SchemaMetadataManagerRef, file_ref_manager: FileReferenceManagerRef, + partition_expr_fetcher: PartitionExprFetcherRef, plugins: Plugins, ) -> Self { Self { @@ -152,6 +156,7 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> { schema_metadata_manager, file_ref_manager, plugins, + partition_expr_fetcher, #[cfg(feature = "enterprise")] extension_range_provider_factory: None, } @@ -179,6 +184,7 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> { self.object_store_manager, self.schema_metadata_manager, self.file_ref_manager, + self.partition_expr_fetcher.clone(), self.plugins, ) .await?; @@ -209,6 +215,7 @@ pub struct MitoEngine { impl MitoEngine { /// Returns a new [MitoEngine] with specific `config`, `log_store` and `object_store`. + #[allow(clippy::too_many_arguments)] pub async fn new( data_home: &str, config: MitoConfig, @@ -216,6 +223,7 @@ impl MitoEngine { object_store_manager: ObjectStoreManagerRef, schema_metadata_manager: SchemaMetadataManagerRef, file_ref_manager: FileReferenceManagerRef, + partition_expr_fetcher: PartitionExprFetcherRef, plugins: Plugins, ) -> Result { let builder = MitoEngineBuilder::new( @@ -225,6 +233,7 @@ impl MitoEngine { object_store_manager, schema_metadata_manager, file_ref_manager, + partition_expr_fetcher, plugins, ); builder.try_build().await @@ -946,6 +955,7 @@ impl MitoEngine { time_provider: crate::time_provider::TimeProviderRef, schema_metadata_manager: SchemaMetadataManagerRef, file_ref_manager: FileReferenceManagerRef, + partition_expr_fetcher: PartitionExprFetcherRef, ) -> Result { config.sanitize(data_home)?; @@ -962,6 +972,7 @@ impl MitoEngine { schema_metadata_manager, file_ref_manager, time_provider, + partition_expr_fetcher, ) .await?, config, diff --git a/src/mito2/src/engine/alter_test.rs b/src/mito2/src/engine/alter_test.rs index 383ac48fe6e..268c1f4855f 100644 --- a/src/mito2/src/engine/alter_test.rs +++ b/src/mito2/src/engine/alter_test.rs @@ -402,7 +402,7 @@ async fn test_alter_on_flushing() { let mut env = TestEnv::new().await; let listener = Arc::new(AlterFlushListener::default()); let engine = env - .create_engine_with(MitoConfig::default(), None, Some(listener.clone())) + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) .await; let region_id = RegionId::new(1, 1); @@ -506,7 +506,7 @@ async fn test_alter_column_fulltext_options() { let mut env = TestEnv::new().await; let listener = Arc::new(AlterFlushListener::default()); let engine = env - .create_engine_with(MitoConfig::default(), None, Some(listener.clone())) + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) .await; let region_id = RegionId::new(1, 1); @@ -629,7 +629,7 @@ async fn test_alter_column_set_inverted_index() { let mut env = TestEnv::new().await; let listener = Arc::new(AlterFlushListener::default()); let engine = env - .create_engine_with(MitoConfig::default(), None, Some(listener.clone())) + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) .await; let region_id = RegionId::new(1, 1); @@ -743,7 +743,7 @@ async fn test_alter_region_ttl_options() { let mut env = TestEnv::new().await; let listener = Arc::new(AlterFlushListener::default()); let engine = env - .create_engine_with(MitoConfig::default(), None, Some(listener.clone())) + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) .await; let region_id = RegionId::new(1, 1); @@ -793,7 +793,7 @@ async fn test_write_stall_on_altering() { let mut env = TestEnv::new().await; let listener = Arc::new(NotifyRegionChangeResultListener::default()); let engine = env - .create_engine_with(MitoConfig::default(), None, Some(listener.clone())) + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) .await; let region_id = RegionId::new(1, 1); diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index 731d4e40138..86cd89a2784 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -519,6 +519,7 @@ async fn test_readonly_during_compaction() { }, None, Some(listener.clone()), + None, ) .await; diff --git a/src/mito2/src/engine/drop_test.rs b/src/mito2/src/engine/drop_test.rs index 8c03f3b8488..6cc38822af3 100644 --- a/src/mito2/src/engine/drop_test.rs +++ b/src/mito2/src/engine/drop_test.rs @@ -38,7 +38,7 @@ async fn test_engine_drop_region() { let mut env = TestEnv::with_prefix("drop").await; let listener = Arc::new(DropListener::new(Duration::from_millis(100))); let engine = env - .create_engine_with(MitoConfig::default(), None, Some(listener.clone())) + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) .await; let region_id = RegionId::new(1, 1); diff --git a/src/mito2/src/engine/edit_region_test.rs b/src/mito2/src/engine/edit_region_test.rs index 35a29db4e35..20f68fe5d9b 100644 --- a/src/mito2/src/engine/edit_region_test.rs +++ b/src/mito2/src/engine/edit_region_test.rs @@ -148,6 +148,7 @@ async fn test_edit_region_fill_cache() { Some(Arc::new(EditRegionListener { tx: Mutex::new(Some(tx)), })), + None, ) .await; diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index 61cbfa96f27..f9cc2ada0ad 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -99,6 +99,7 @@ async fn test_flush_engine() { MitoConfig::default(), Some(write_buffer_manager.clone()), Some(listener.clone()), + None, ) .await; @@ -169,6 +170,7 @@ async fn test_write_stall() { MitoConfig::default(), Some(write_buffer_manager.clone()), Some(listener.clone()), + None, ) .await; @@ -243,6 +245,7 @@ async fn test_flush_empty() { MitoConfig::default(), Some(write_buffer_manager.clone()), None, + None, ) .await; @@ -478,6 +481,7 @@ async fn test_flush_workers() { }, Some(write_buffer_manager.clone()), Some(listener.clone()), + None, ) .await; @@ -560,6 +564,7 @@ async fn test_update_topic_latest_entry_id(factory: Option) { MitoConfig::default(), Some(write_buffer_manager.clone()), Some(listener.clone()), + None, ) .await; let region_id = RegionId::new(1, 1); diff --git a/src/mito2/src/engine/open_test.rs b/src/mito2/src/engine/open_test.rs index d20bc1157aa..d0a73e72091 100644 --- a/src/mito2/src/engine/open_test.rs +++ b/src/mito2/src/engine/open_test.rs @@ -13,6 +13,8 @@ // limitations under the License. use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use api::v1::Rows; @@ -20,7 +22,7 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; use either::Either; -use store_api::region_engine::{RegionEngine, RegionRole}; +use store_api::region_engine::{RegionEngine, RegionRole, SettableRegionRoleState}; use store_api::region_request::{ PathType, RegionCloseRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, }; @@ -30,6 +32,7 @@ use tokio::sync::oneshot; use crate::compaction::compactor::{OpenCompactionRegionRequest, open_compaction_region}; use crate::config::MitoConfig; use crate::error; +use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef}; use crate::region::options::RegionOptions; use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, reopen_region, rows_schema, @@ -503,3 +506,153 @@ async fn test_open_compaction_region() { assert_eq!(region_id, compaction_region.region_id); } + +/// Returns a dummy fetcher that returns the expr_json only once. +fn new_dummy_fetcher(region_id: RegionId, expr_json: String) -> PartitionExprFetcherRef { + struct DummyFetcher { + region_id: RegionId, + expr_json: String, + fetch_count: AtomicUsize, + } + #[async_trait::async_trait] + impl PartitionExprFetcher for DummyFetcher { + async fn fetch_expr(&self, region_id: RegionId) -> Option { + if region_id == self.region_id { + let count = self.fetch_count.fetch_add(1, Ordering::Relaxed); + (count == 0).then(|| self.expr_json.clone()) + } else { + None + } + } + } + + Arc::new(DummyFetcher { + region_id, + expr_json, + fetch_count: AtomicUsize::new(0), + }) +} + +#[tokio::test] +async fn test_open_backfills_partition_expr_with_fetcher() { + let mut env = TestEnv::with_prefix("open-backfill-fetcher").await; + + // prepare plugins with dummy fetcher + let region_id = RegionId::new(100, 1); + let expr_json = r#"{\"Expr\":{\"lhs\":{\"Column\":\"a\"},\"op\":\"GtEq\",\"rhs\":{\"Value\":{\"UInt32\":10}}}}"#.to_string(); + let fetcher = new_dummy_fetcher(region_id, expr_json.clone()); + let engine = env + .create_engine_with(MitoConfig::default(), None, None, Some(fetcher)) + .await; + + // create region with partition_expr=None (legacy style) + let request = CreateRequestBuilder::new() + .table_dir("t") + .partition_expr_json(None) + .build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + + let meta = engine.get_region(region_id).unwrap().metadata(); + assert!(meta.partition_expr.is_none()); + + // close and reopen to trigger backfill in opener + engine + .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir: request.table_dir.clone(), + path_type: PathType::Bare, + options: HashMap::default(), + skip_wal_replay: false, + checkpoint: None, + }), + ) + .await + .unwrap(); + + // verify partition_expr is backfilled and persisted + let meta = engine.get_region(region_id).unwrap().metadata(); + assert_eq!(meta.partition_expr.as_deref(), Some(expr_json.as_str())); + + // Set leader and trigger backfill + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader) + .await + .unwrap(); + + // reopen again to ensure no further changes and still Some + engine + .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .await + .unwrap(); + let engine = env.reopen_engine(engine, MitoConfig::default()).await; + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir: request.table_dir.clone(), + path_type: PathType::Bare, + options: HashMap::default(), + skip_wal_replay: false, + checkpoint: None, + }), + ) + .await + .unwrap(); + let meta = engine.get_region(region_id).unwrap().metadata(); + assert_eq!(meta.partition_expr.as_deref(), Some(expr_json.as_str())); +} + +#[tokio::test] +async fn test_open_keeps_none_without_fetcher() { + let mut env = TestEnv::with_prefix("open-no-fetcher").await; + // engine without fetcher + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(200, 1); + let request = CreateRequestBuilder::new() + .table_dir("t2") + .partition_expr_json(None) + .build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + let meta = engine.get_region(region_id).unwrap().metadata(); + assert!(meta.partition_expr.is_none()); + + engine + .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .await + .unwrap(); + let engine = env.reopen_engine(engine, MitoConfig::default()).await; + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir: request.table_dir.clone(), + path_type: PathType::Bare, + options: HashMap::default(), + skip_wal_replay: false, + checkpoint: None, + }), + ) + .await + .unwrap(); + + let meta = engine.get_region(region_id).unwrap().metadata(); + assert!(meta.partition_expr.is_none()); +} diff --git a/src/mito2/src/engine/truncate_test.rs b/src/mito2/src/engine/truncate_test.rs index 557a988487f..0686db95c18 100644 --- a/src/mito2/src/engine/truncate_test.rs +++ b/src/mito2/src/engine/truncate_test.rs @@ -304,6 +304,7 @@ async fn test_engine_truncate_during_flush() { MitoConfig::default(), Some(write_buffer_manager), Some(listener.clone()), + None, ) .await; diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index da47cf7126a..d3fa4072790 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -41,7 +41,9 @@ use crate::error::{ FlushableRegionStateSnafu, RegionNotFoundSnafu, RegionStateSnafu, RegionTruncatedSnafu, Result, UpdateManifestSnafu, }; -use crate::manifest::action::{RegionManifest, RegionMetaAction, RegionMetaActionList}; +use crate::manifest::action::{ + RegionChange, RegionManifest, RegionMetaAction, RegionMetaActionList, +}; use crate::manifest::manager::RegionManifestManager; use crate::memtable::MemtableBuilderRef; use crate::region::version::{VersionControlRef, VersionRef}; @@ -321,7 +323,7 @@ impl MitoRegion { &self, state: SettableRegionRoleState, ) -> Result<()> { - let _manager = self.manifest_ctx.manifest_manager.write().await; + let mut manager = self.manifest_ctx.manifest_manager.write().await; let current_state = self.state(); match state { @@ -423,6 +425,36 @@ impl MitoRegion { } } + // Hack(zhongzc): If we have just become leader (writable), persist any backfilled metadata. + if self.state() == RegionRoleState::Leader(RegionLeaderState::Writable) { + // Persist backfilled metadata if manifest is missing fields (e.g., partition_expr) + let manifest_meta = &manager.manifest().metadata; + let current_meta = &self.version().metadata; + if manifest_meta.partition_expr.is_none() && current_meta.partition_expr.is_some() { + let action = RegionMetaAction::Change(RegionChange { + metadata: current_meta.clone(), + }); + let result = manager + .update( + RegionMetaActionList::with_action(action), + RegionRoleState::Leader(RegionLeaderState::Writable), + ) + .await; + + match result { + Ok(version) => { + info!( + "Successfully persisted backfilled metadata for region {}, version: {}", + self.region_id, version + ); + } + Err(e) => { + warn!(e; "Failed to persist backfilled metadata for region {}", self.region_id); + } + } + } + } + Ok(()) } diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index d8af3e3209f..c06ca890c67 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -69,6 +69,18 @@ use crate::time_provider::TimeProviderRef; use crate::wal::entry_reader::WalEntryReader; use crate::wal::{EntryId, Wal}; +/// A fetcher to retrieve partition expr for a region. +/// +/// Compatibility: older regions didn't persist `partition_expr` in engine metadata, +/// while newer ones do. On open, we backfill it via this fetcher and persist it +/// to the manifest so future opens don't need refetching. +#[async_trait::async_trait] +pub trait PartitionExprFetcher { + async fn fetch_expr(&self, region_id: RegionId) -> Option; +} + +pub type PartitionExprFetcherRef = Arc; + /// Builder to create a new [MitoRegion] or open an existing one. pub(crate) struct RegionOpener { region_id: RegionId, @@ -88,6 +100,7 @@ pub(crate) struct RegionOpener { wal_entry_reader: Option>, replay_checkpoint: Option, file_ref_manager: FileReferenceManagerRef, + partition_expr_fetcher: PartitionExprFetcherRef, } impl RegionOpener { @@ -105,6 +118,7 @@ impl RegionOpener { intermediate_manager: IntermediateManager, time_provider: TimeProviderRef, file_ref_manager: FileReferenceManagerRef, + partition_expr_fetcher: PartitionExprFetcherRef, ) -> RegionOpener { RegionOpener { region_id, @@ -124,6 +138,7 @@ impl RegionOpener { wal_entry_reader: None, replay_checkpoint: None, file_ref_manager, + partition_expr_fetcher, } } @@ -390,8 +405,18 @@ impl RegionOpener { return Ok(None); }; + // Backfill `partition_expr` if missing. Use the backfilled metadata in-memory during this open. let manifest = manifest_manager.manifest(); - let metadata = manifest.metadata.clone(); + let metadata = if manifest.metadata.partition_expr.is_none() + && let Some(expr_json) = self.partition_expr_fetcher.fetch_expr(self.region_id).await + { + let metadata = manifest.metadata.as_ref().clone(); + let mut builder = RegionMetadataBuilder::from_existing(metadata); + builder.partition_expr_json(Some(expr_json)); + Arc::new(builder.build().context(InvalidMetadataSnafu)?) + } else { + manifest.metadata.clone() + }; let region_id = self.region_id; let provider = self.provider::(®ion_options.wal_options)?; diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index 13a91d19680..7c1e1dbac62 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -73,6 +73,7 @@ use crate::error::Result; use crate::flush::{WriteBufferManager, WriteBufferManagerRef}; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; use crate::read::{Batch, BatchBuilder, BatchReader}; +use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef}; use crate::sst::file_purger::{FilePurgerRef, NoopFilePurger}; use crate::sst::file_ref::{FileReferenceManager, FileReferenceManagerRef}; use crate::sst::index::intermediate::IntermediateManager; @@ -105,6 +106,19 @@ pub(crate) fn kafka_log_store_factory() -> Option { })) } +pub(crate) fn noop_partition_expr_fetcher() -> PartitionExprFetcherRef { + struct NoopPartitionExprFetcher; + + #[async_trait::async_trait] + impl PartitionExprFetcher for NoopPartitionExprFetcher { + async fn fetch_expr(&self, _region_id: RegionId) -> Option { + None + } + } + + Arc::new(NoopPartitionExprFetcher) +} + #[template] #[rstest] #[case::with_raft_engine(raft_engine_log_store_factory())] @@ -212,6 +226,7 @@ pub struct TestEnv { schema_metadata_manager: SchemaMetadataManagerRef, file_ref_manager: FileReferenceManagerRef, kv_backend: KvBackendRef, + partition_expr_fetcher: PartitionExprFetcherRef, } impl TestEnv { @@ -247,6 +262,7 @@ impl TestEnv { schema_metadata_manager, file_ref_manager: Arc::new(FileReferenceManager::new(None)), kv_backend, + partition_expr_fetcher: noop_partition_expr_fetcher(), } } @@ -284,6 +300,7 @@ impl TestEnv { zelf.object_store_manager.as_ref().unwrap().clone(), zelf.schema_metadata_manager.clone(), zelf.file_ref_manager.clone(), + zelf.partition_expr_fetcher.clone(), Plugins::new(), ) .await @@ -318,6 +335,7 @@ impl TestEnv { config: MitoConfig, manager: Option, listener: Option, + partition_expr_fetcher: Option, ) -> MitoEngine { let (log_store, object_store_manager) = self.create_log_and_object_store_manager().await; @@ -327,6 +345,10 @@ impl TestEnv { let data_home = self.data_home().display().to_string(); + let partition_expr_fetcher = + partition_expr_fetcher.unwrap_or_else(noop_partition_expr_fetcher); + self.partition_expr_fetcher = partition_expr_fetcher; + match log_store { LogStoreImpl::RaftEngine(log_store) => MitoEngine::new_for_test( &data_home, @@ -338,6 +360,7 @@ impl TestEnv { Arc::new(StdTimeProvider), self.schema_metadata_manager.clone(), self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), ) .await .unwrap(), @@ -351,6 +374,7 @@ impl TestEnv { Arc::new(StdTimeProvider), self.schema_metadata_manager.clone(), self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), ) .await .unwrap(), @@ -395,6 +419,7 @@ impl TestEnv { Arc::new(StdTimeProvider), self.schema_metadata_manager.clone(), self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), ) .await .unwrap(), @@ -408,6 +433,7 @@ impl TestEnv { Arc::new(StdTimeProvider), self.schema_metadata_manager.clone(), self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), ) .await .unwrap(), @@ -441,6 +467,7 @@ impl TestEnv { time_provider.clone(), self.schema_metadata_manager.clone(), self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), ) .await .unwrap(), @@ -454,6 +481,7 @@ impl TestEnv { time_provider.clone(), self.schema_metadata_manager.clone(), self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), ) .await .unwrap(), @@ -491,6 +519,7 @@ impl TestEnv { Arc::new(object_store_manager), self.schema_metadata_manager.clone(), self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), Plugins::new(), ) .await @@ -501,6 +530,7 @@ impl TestEnv { Arc::new(object_store_manager), self.schema_metadata_manager.clone(), self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), Plugins::new(), ) .await diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 3b110e20640..6e60de5c28b 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -59,6 +59,7 @@ use crate::error::{self, CreateDirSnafu, JoinSnafu, Result, WorkerStoppedSnafu}; use crate::flush::{FlushScheduler, WriteBufferManagerImpl, WriteBufferManagerRef}; use crate::memtable::MemtableBuilderProvider; use crate::metrics::{REGION_COUNT, REQUEST_WAIT_TIME, WRITE_STALLING}; +use crate::region::opener::PartitionExprFetcherRef; use crate::region::{MitoRegionRef, OpeningRegions, OpeningRegionsRef, RegionMap, RegionMapRef}; use crate::request::{ BackgroundNotify, DdlRequest, SenderBulkRequest, SenderDdlRequest, SenderWriteRequest, @@ -145,6 +146,7 @@ impl WorkerGroup { object_store_manager: ObjectStoreManagerRef, schema_metadata_manager: SchemaMetadataManagerRef, file_ref_manager: FileReferenceManagerRef, + partition_expr_fetcher: PartitionExprFetcherRef, plugins: Plugins, ) -> Result { let (flush_sender, flush_receiver) = watch::channel(()); @@ -209,6 +211,7 @@ impl WorkerGroup { plugins: plugins.clone(), schema_metadata_manager: schema_metadata_manager.clone(), file_ref_manager: file_ref_manager.clone(), + partition_expr_fetcher: partition_expr_fetcher.clone(), } .start() }) @@ -306,6 +309,7 @@ impl WorkerGroup { schema_metadata_manager: SchemaMetadataManagerRef, file_ref_manager: FileReferenceManagerRef, time_provider: TimeProviderRef, + partition_expr_fetcher: PartitionExprFetcherRef, ) -> Result { let (flush_sender, flush_receiver) = watch::channel(()); let write_buffer_manager = write_buffer_manager.unwrap_or_else(|| { @@ -363,6 +367,7 @@ impl WorkerGroup { plugins: Plugins::new(), schema_metadata_manager: schema_metadata_manager.clone(), file_ref_manager: file_ref_manager.clone(), + partition_expr_fetcher: partition_expr_fetcher.clone(), } .start() }) @@ -443,6 +448,7 @@ struct WorkerStarter { plugins: Plugins, schema_metadata_manager: SchemaMetadataManagerRef, file_ref_manager: FileReferenceManagerRef, + partition_expr_fetcher: PartitionExprFetcherRef, } impl WorkerStarter { @@ -496,6 +502,7 @@ impl WorkerStarter { region_edit_queues: RegionEditQueues::default(), schema_metadata_manager: self.schema_metadata_manager, file_ref_manager: self.file_ref_manager.clone(), + partition_expr_fetcher: self.partition_expr_fetcher, }; let handle = common_runtime::spawn_global(async move { worker_thread.run().await; @@ -747,6 +754,8 @@ struct RegionWorkerLoop { schema_metadata_manager: SchemaMetadataManagerRef, /// Datanode level file references manager. file_ref_manager: FileReferenceManagerRef, + /// Partition expr fetcher used to backfill partition expr on open for compatibility. + partition_expr_fetcher: PartitionExprFetcherRef, } impl RegionWorkerLoop { diff --git a/src/mito2/src/worker/handle_catchup.rs b/src/mito2/src/worker/handle_catchup.rs index 4924591e69c..caabb6ae550 100644 --- a/src/mito2/src/worker/handle_catchup.rs +++ b/src/mito2/src/worker/handle_catchup.rs @@ -20,7 +20,7 @@ use common_telemetry::tracing::warn; use common_telemetry::{debug, info}; use snafu::ensure; use store_api::logstore::LogStore; -use store_api::region_engine::RegionRole; +use store_api::region_engine::{RegionRole, SettableRegionRoleState}; use store_api::region_request::{AffectedRows, RegionCatchupRequest}; use store_api::storage::RegionId; use tokio::time::Instant; @@ -136,6 +136,10 @@ impl RegionWorkerLoop { if request.set_writable { region.set_role(RegionRole::Leader); + // Finalize leadership: persist backfilled metadata. + region + .set_role_state_gracefully(SettableRegionRoleState::Leader) + .await?; } Ok(0) @@ -164,6 +168,7 @@ impl RegionWorkerLoop { self.intermediate_manager.clone(), self.time_provider.clone(), self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), ) .cache(Some(self.cache_manager.clone())) .options(region.version().options.clone())? diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index 492f4d61860..3c5f091a1a7 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/src/mito2/src/worker/handle_create.rs @@ -67,6 +67,7 @@ impl RegionWorkerLoop { self.intermediate_manager.clone(), self.time_provider.clone(), self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), ) .metadata_builder(builder) .parse_options(request.options)? diff --git a/src/mito2/src/worker/handle_open.rs b/src/mito2/src/worker/handle_open.rs index b0c93919c83..aa9423b526c 100644 --- a/src/mito2/src/worker/handle_open.rs +++ b/src/mito2/src/worker/handle_open.rs @@ -104,6 +104,7 @@ impl RegionWorkerLoop { self.intermediate_manager.clone(), self.time_provider.clone(), self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), ) .skip_wal_replay(request.skip_wal_replay) .cache(Some(self.cache_manager.clone()))