diff --git a/src/metric-engine/src/data_region.rs b/src/metric-engine/src/data_region.rs index a11b3e57b54..8847307ff14 100644 --- a/src/metric-engine/src/data_region.rs +++ b/src/metric-engine/src/data_region.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; + use api::v1::SemanticType; use common_query::native_histogram::is_native_histogram_value_type; use common_telemetry::{debug, info}; @@ -26,7 +28,7 @@ use store_api::region_request::{ use store_api::storage::consts::ReservedColumnId; use store_api::storage::{ConcreteDataType, RegionId}; -use crate::engine::IndexOptions; +use crate::engine::{IndexOptions, validate_shared_column}; use crate::error::{ AddingFieldColumnSnafu, ColumnTypeMismatchSnafu, ForbiddenPhysicalAlterSnafu, MitoReadOperationSnafu, MitoWriteOperationSnafu, Result, SetSkippingIndexOptionSnafu, @@ -44,7 +46,7 @@ pub struct DataRegion { } impl DataRegion { - pub fn new(mito: MitoEngine) -> Self { + pub(crate) fn new(mito: MitoEngine) -> Self { Self { mito } } @@ -55,10 +57,16 @@ impl DataRegion { /// logical regions. /// /// Invoker don't need to set up or verify the column id. This method will adjust - /// it using underlying schema. + /// it using underlying schema, and reuse the persisted id of a column that the + /// physical region already has. /// /// This method will also set the nullable marker to true. All of those change are applies /// to `columns` in-place. + /// + /// Callers must hold the column-add lock of the physical region, otherwise two + /// concurrent requests could read the same metadata and assign one column id to + /// different columns. Use + /// [MetricEngineInner::add_physical_columns](crate::engine::MetricEngineInner::add_physical_columns). pub async fn add_columns( &self, region_id: RegionId, @@ -72,10 +80,13 @@ impl DataRegion { let region_id = utils::to_data_region_id(region_id); - let num_columns = columns.len(); - let request = self + let Some((new_column_count, request)) = self .assemble_alter_request(region_id, columns, index_options) - .await?; + .await? + else { + // Every requested column already exists. + return Ok(()); + }; let _timer = MITO_DDL_DURATION.start_timer(); @@ -85,26 +96,55 @@ impl DataRegion { .await .context(MitoWriteOperationSnafu)?; - PHYSICAL_COLUMN_COUNT.add(num_columns as _); + PHYSICAL_COLUMN_COUNT.add(new_column_count as _); Ok(()) } /// Generate wrapped [RegionAlterRequest] with given [ColumnMetadata]. - /// This method will modify `columns` in-place. + /// + /// Columns that the physical region already has are not added again: their + /// definitions are validated and their persisted column ids are reused. + /// Returns `None` if every given column already exists. async fn assemble_alter_request( &self, region_id: RegionId, columns: Vec, index_options: IndexOptions, - ) -> Result { - // retrieve underlying version + ) -> Result> { let region_metadata = self .mito .get_metadata(region_id) .await .context(MitoReadOperationSnafu)?; + // The caller derives the columns to add from a cached schema that a cancelled or + // concurrent alter may have left behind, so the names are checked against the + // real metadata instead of being trusted. + let mut seen: HashMap = HashMap::with_capacity(columns.len()); + let mut columns_to_add = Vec::with_capacity(columns.len()); + for column in columns { + let name = column.column_schema.name.clone(); + if let Some(existing) = region_metadata.column_by_name(&name) { + validate_shared_column(region_id, existing, &column)?; + continue; + } + // The batch itself may repeat a column that is not installed yet: the duplicate + // must agree with the first occurrence, and only one of them is added. + if let Some(&index) = seen.get(&name) { + validate_shared_column(region_id, &columns_to_add[index], &column)?; + continue; + } + seen.insert(name, columns_to_add.len()); + columns_to_add.push(column); + } + + if columns_to_add.is_empty() { + debug!("No new columns to add to physical region {region_id:?}"); + return Ok(None); + } + let new_column_count = columns_to_add.len(); + // find the max column id let new_column_id_start = 1 + region_metadata .column_metadatas @@ -120,7 +160,7 @@ impl DataRegion { .unwrap_or(0); // overwrite semantic type - let new_columns = columns + let new_columns = columns_to_add .into_iter() .enumerate() .map(|(delta, mut c)| { @@ -188,7 +228,7 @@ impl DataRegion { }, }); - Ok(alter_request) + Ok(Some((new_column_count, alter_request))) } pub async fn write_data( @@ -259,6 +299,14 @@ mod test { use super::*; use crate::test_util::TestEnv; + fn tag_column(name: &str) -> ColumnMetadata { + ColumnMetadata { + column_id: 0, + semantic_type: SemanticType::Tag, + column_schema: ColumnSchema::new(name, ConcreteDataType::string_datatype(), false), + } + } + #[tokio::test] async fn test_add_columns() { let env = TestEnv::new().await; @@ -345,4 +393,122 @@ mod test { .await; assert!(result.is_err()); } + + // A duplicate column inside one batch must agree with the first occurrence instead of + // being dropped silently. + #[tokio::test] + async fn test_add_columns_rejects_conflicting_duplicates() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let physical_region_id = env.default_physical_region_id(); + let data_region_id = utils::to_data_region_id(physical_region_id); + + let conflicting = ColumnMetadata { + column_id: 0, + semantic_type: SemanticType::Field, + column_schema: ColumnSchema::new("dup", ConcreteDataType::float64_datatype(), false), + }; + let result = env + .data_region() + .add_columns( + physical_region_id, + vec![tag_column("dup"), conflicting], + IndexOptions::Inverted, + ) + .await; + + assert!(result.is_err()); + assert!( + env.mito() + .get_metadata(data_region_id) + .await + .unwrap() + .column_by_name("dup") + .is_none() + ); + } + + // A retry passes the column it added before, because the caller derives the columns to + // add from a cache that a cancelled request may not have refreshed. The column keeps the + // id it was installed with. + #[tokio::test] + async fn test_add_columns_reuses_the_existing_column_id() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let physical_region_id = env.default_physical_region_id(); + let data_region_id = utils::to_data_region_id(physical_region_id); + + env.data_region() + .add_columns( + physical_region_id, + vec![tag_column("tag2")], + IndexOptions::Inverted, + ) + .await + .unwrap(); + let tag2_id = env + .mito() + .get_metadata(data_region_id) + .await + .unwrap() + .column_by_name("tag2") + .unwrap() + .column_id; + + env.data_region() + .add_columns( + physical_region_id, + vec![tag_column("tag2"), tag_column("tag3")], + IndexOptions::Inverted, + ) + .await + .unwrap(); + + let metadata = env.mito().get_metadata(data_region_id).await.unwrap(); + assert_eq!(metadata.column_by_name("tag2").unwrap().column_id, tag2_id); + assert_eq!( + metadata.column_by_name("tag3").unwrap().column_id, + tag2_id + 1 + ); + } + + // Adding only columns that already exist must not submit an ALTER at all. + #[tokio::test] + async fn test_add_columns_skips_the_alter_when_all_columns_exist() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let physical_region_id = env.default_physical_region_id(); + let data_region_id = utils::to_data_region_id(physical_region_id); + + env.data_region() + .add_columns( + physical_region_id, + vec![tag_column("tag2")], + IndexOptions::Inverted, + ) + .await + .unwrap(); + let metadata = env.mito().get_metadata(data_region_id).await.unwrap(); + let version = metadata.schema_version; + + env.data_region() + .add_columns( + physical_region_id, + vec![tag_column("tag2")], + IndexOptions::Inverted, + ) + .await + .unwrap(); + + let metadata = env.mito().get_metadata(data_region_id).await.unwrap(); + assert_eq!(metadata.schema_version, version); + assert_eq!( + metadata + .column_metadatas + .iter() + .filter(|column| column.column_schema.name == "tag2") + .count(), + 1 + ); + } } diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index 376d9f4b564..3d3494ec7e7 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -25,7 +25,7 @@ mod put; mod read; mod region_metadata; mod staging; -mod state; +pub(crate) mod state; mod sync; use std::any::Any; @@ -59,8 +59,8 @@ use store_api::storage::{RegionId, ScanRequest, SequenceNumber}; use crate::config::EngineConfig; use crate::data_region::DataRegion; use crate::error::{ - self, Error, Result, StartRepeatedTaskSnafu, UnsupportedRegionRequestSnafu, - UnsupportedRemapManifestsRequestSnafu, + self, Error, PhysicalRegionNotFoundSnafu, Result, StartRepeatedTaskSnafu, + UnsupportedRegionRequestSnafu, UnsupportedRemapManifestsRequestSnafu, }; use crate::metadata_region::MetadataRegion; use crate::repeated_task::FlushMetadataRegionTask; @@ -608,7 +608,7 @@ impl MetricEngine { } /// Shared definitions must agree before CREATE/ALTER mutates either region. -fn validate_shared_column( +pub(crate) fn validate_shared_column( region_id: store_api::storage::RegionId, existing: &store_api::metadata::ColumnMetadata, proposed: &store_api::metadata::ColumnMetadata, @@ -647,6 +647,42 @@ struct MetricEngineInner { flush_task: RepeatedTask, } +impl MetricEngineInner { + /// Adds columns to the physical region while holding its column-add lock. + /// + /// Column ids are derived from the metadata of the physical region, so requests + /// on logical regions that share one physical region must not allocate ids at + /// the same time, otherwise two different columns could get one id. + async fn add_physical_columns( + &self, + physical_region_id: RegionId, + columns: Vec, + index_options: IndexOptions, + ) -> Result<()> { + if columns.is_empty() { + return Ok(()); + } + + let data_region_id = utils::to_data_region_id(physical_region_id); + let column_add_lock = { + let state = self.state.read().unwrap(); + state + .physical_region_states() + .get(&data_region_id) + .with_context(|| PhysicalRegionNotFoundSnafu { + region_id: data_region_id, + })? + .column_add_lock() + .clone() + }; + + let _guard = column_add_lock.lock().await; + self.data_region + .add_columns(physical_region_id, columns, index_options) + .await + } +} + #[cfg(test)] mod test { use std::assert_matches; diff --git a/src/metric-engine/src/engine/alter.rs b/src/metric-engine/src/engine/alter.rs index e940fa4cbbb..a5f3ec4a4d4 100644 --- a/src/metric-engine/src/engine/alter.rs +++ b/src/metric-engine/src/engine/alter.rs @@ -29,8 +29,8 @@ use validate::validate_alter_region_requests; use crate::engine::MetricEngineInner; use crate::error::{ - AddingFieldColumnSnafu, LogicalRegionNotFoundSnafu, PhysicalRegionNotFoundSnafu, Result, - SerializeColumnMetadataSnafu, UnexpectedRequestSnafu, + AddingFieldColumnSnafu, LogicalRegionNotFoundSnafu, MissingPhysicalColumnSnafu, + PhysicalRegionNotFoundSnafu, Result, SerializeColumnMetadataSnafu, UnexpectedRequestSnafu, }; use crate::utils::{append_manifest_info, encode_manifest_info_to_extensions, to_data_region_id}; @@ -163,8 +163,7 @@ impl MetricEngineInner { ); } - self.data_region - .add_columns(data_region_id, new_columns_to_add, index_options) + self.add_physical_columns(physical_region_id, new_columns_to_add, index_options) .await?; let physical_columns = self.data_region.physical_columns(data_region_id).await?; @@ -173,28 +172,39 @@ impl MetricEngineInner { .map(|metadata| (metadata.column_schema.name.as_str(), metadata)) .collect::>(); - let logical_region_columns = requests.iter().map(|(region_id, request)| { + // The physical region now has every requested column, either because this call + // added it or because it was already there with the same definition. + let mut logical_region_columns = Vec::with_capacity(requests.len()); + for (region_id, request) in &requests { let AlterKind::AddColumns { columns } = &request.kind else { unreachable!() }; - ( - *region_id, - columns - .iter() - .map(|col| { - let column_name = col.column_metadata.column_schema.name.as_str(); - let column_metadata = *physical_schema_map.get(column_name).unwrap(); - (column_name, column_metadata) - }) - .collect::>(), - ) - }); + let mut column_metadata = HashMap::with_capacity(columns.len()); + for col in columns { + let column_name = col.column_metadata.column_schema.name.as_str(); + let metadata = + physical_schema_map + .get(column_name) + .copied() + .with_context(|| MissingPhysicalColumnSnafu { + name: column_name, + region_id: physical_region_id, + })?; + column_metadata.insert(column_name, metadata); + } + logical_region_columns.push((*region_id, column_metadata)); + } - let new_add_columns = new_column_indices.keys().map(|name| { - // Safety: previous steps ensure the physical region exist - let column_metadata = *physical_schema_map.get(name).unwrap(); - (name.to_string(), column_metadata.clone()) - }); + let mut new_add_columns = Vec::with_capacity(new_column_indices.len()); + for name in new_column_indices.keys() { + let column_metadata = physical_schema_map.get(name).copied().with_context(|| { + MissingPhysicalColumnSnafu { + name: name.to_string(), + region_id: physical_region_id, + } + })?; + new_add_columns.push((name.to_string(), column_metadata.clone())); + } // The data region has the new columns, whatever happens to the metadata write. self.state @@ -205,7 +215,11 @@ impl MetricEngineInner { // Writes logical regions metadata to metadata region let result = self .metadata_region - .add_logical_regions(physical_region_id, false, logical_region_columns) + .add_logical_regions( + physical_region_id, + false, + logical_region_columns.into_iter(), + ) .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. @@ -284,6 +298,9 @@ mod test { use common_meta::ddl::test_util::assert_column_name_and_id; use common_meta::ddl::utils::{parse_column_metadatas, parse_manifest_infos_from_extensions}; use common_query::prelude::{greptime_timestamp, greptime_value}; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::ColumnSchema; + use store_api::metadata::ColumnMetadata; use store_api::metric_engine_consts::ALTER_PHYSICAL_EXTENSION_KEY; use store_api::region_engine::RegionEngine; use store_api::region_request::{ @@ -292,7 +309,9 @@ mod test { use store_api::storage::RegionId; use store_api::storage::consts::ReservedColumnId; + use crate::engine::IndexOptions; use crate::test_util::{TestEnv, alter_logical_region_request, create_logical_region_request}; + use crate::utils::to_data_region_id; async fn check_alter_physical_skip_wal( engine: &super::MetricEngineInner, @@ -442,4 +461,109 @@ mod test { ], ); } + + // Column ids are derived from the current metadata, so concurrent requests on + // logical regions sharing one physical region must not hand the same id to + // different columns. + #[tokio::test] + async fn test_add_physical_columns_concurrently_assigns_distinct_ids() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let physical_region_id = env.default_physical_region_id(); + let data_region_id = to_data_region_id(physical_region_id); + let tag = |name: &str| ColumnMetadata { + column_id: 0, + semantic_type: SemanticType::Tag, + column_schema: ColumnSchema::new(name, ConcreteDataType::string_datatype(), false), + }; + + let inner = env.metric().inner.clone(); + let (first, second) = tokio::join!( + inner.add_physical_columns( + physical_region_id, + vec![tag("tag4")], + IndexOptions::Inverted, + ), + inner.add_physical_columns( + physical_region_id, + vec![tag("tag5")], + IndexOptions::Inverted, + ), + ); + first.unwrap(); + second.unwrap(); + + let metadata = env.mito().get_metadata(data_region_id).await.unwrap(); + assert_ne!( + metadata.column_by_name("tag4").unwrap().column_id, + metadata.column_by_name("tag5").unwrap().column_id + ); + } + + // A request that Mito accepted but that never wrote the logical metadata leaves the + // physical column installed while the cache still misses it. The retry must reuse the + // installed column id and repair both the cache and the logical metadata. + #[tokio::test] + async fn test_alter_logical_region_after_cancelled_physical_alter() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let physical_region_id = env.default_physical_region_id(); + let logical_region_id = env.default_logical_region_id(); + let data_region_id = to_data_region_id(physical_region_id); + + env.data_region() + .add_columns( + physical_region_id, + vec![ColumnMetadata { + column_id: 0, + semantic_type: SemanticType::Tag, + column_schema: ColumnSchema::new( + "tag1", + ConcreteDataType::string_datatype(), + false, + ), + }], + IndexOptions::Inverted, + ) + .await + .unwrap(); + let installed_id = env + .mito() + .get_metadata(data_region_id) + .await + .unwrap() + .column_by_name("tag1") + .unwrap() + .column_id; + + env.metric() + .handle_batch_ddl_requests(BatchRegionDdlRequest::Alter(vec![( + logical_region_id, + alter_logical_region_request(&["tag1"]), + )])) + .await + .unwrap(); + + let metadata = env.mito().get_metadata(data_region_id).await.unwrap(); + assert_eq!( + metadata.column_by_name("tag1").unwrap().column_id, + installed_id + ); + assert_eq!( + env.metadata_region() + .column_semantic_type(physical_region_id, logical_region_id, "tag1") + .await + .unwrap(), + Some(SemanticType::Tag) + ); + + let engine = env.metric(); + let state = engine.inner.state.read().unwrap(); + let physical_columns = state + .physical_region_states() + .get(&data_region_id) + .unwrap() + .physical_columns(); + assert!(physical_columns.contains_key("tag1")); + } } diff --git a/src/metric-engine/src/engine/create.rs b/src/metric-engine/src/engine/create.rs index 1061565a80e..91121cb457f 100644 --- a/src/metric-engine/src/engine/create.rs +++ b/src/metric-engine/src/engine/create.rs @@ -45,9 +45,9 @@ use crate::engine::create::extract_new_columns::extract_new_columns; use crate::engine::options::{PhysicalRegionOptions, set_data_region_options}; use crate::error::{ ColumnTypeMismatchSnafu, ConflictRegionOptionSnafu, CreateMitoRegionSnafu, - InternalColumnOccupiedSnafu, InvalidMetadataSnafu, MissingRegionOptionSnafu, - MultipleFieldColumnSnafu, NoFieldColumnSnafu, ParseRegionIdSnafu, PhysicalRegionNotFoundSnafu, - Result, SerializeColumnMetadataSnafu, UnexpectedRequestSnafu, + InternalColumnOccupiedSnafu, InvalidMetadataSnafu, MissingPhysicalColumnSnafu, + MissingRegionOptionSnafu, MultipleFieldColumnSnafu, NoFieldColumnSnafu, ParseRegionIdSnafu, + PhysicalRegionNotFoundSnafu, Result, SerializeColumnMetadataSnafu, UnexpectedRequestSnafu, }; use crate::metrics::PHYSICAL_REGION_COUNT; use crate::utils::{ @@ -289,8 +289,7 @@ impl MetricEngineInner { }; // TODO(weny): we dont need to pass a mutable new_columns here. - self.data_region - .add_columns(data_region_id, new_columns, index_option) + self.add_physical_columns(physical_region_id, new_columns, index_option) .await?; let physical_columns = self.data_region.physical_columns(data_region_id).await?; @@ -302,28 +301,35 @@ impl MetricEngineInner { .iter() .map(|(region_id, _)| *region_id) .collect::>(); - let logical_region_columns = requests.iter().map(|(region_id, request)| { - ( - *region_id, - request - .column_metadatas - .iter() - .map(|metadata| { - // Safety: previous steps ensure the physical region exist - let column_metadata = *physical_schema_map - .get(metadata.column_schema.name.as_str()) - .unwrap(); - (metadata.column_schema.name.as_str(), column_metadata) - }) - .collect::>(), - ) - }); + // The physical region now has every requested column, either because this call + // added it or because it was already there with the same definition. + let mut logical_region_columns = Vec::with_capacity(requests.len()); + for (region_id, request) in &requests { + let mut column_metadata = HashMap::with_capacity(request.column_metadatas.len()); + for metadata in &request.column_metadatas { + let name = metadata.column_schema.name.as_str(); + let physical_metadata = + physical_schema_map.get(name).copied().with_context(|| { + MissingPhysicalColumnSnafu { + name, + region_id: physical_region_id, + } + })?; + column_metadata.insert(name, physical_metadata); + } + logical_region_columns.push((*region_id, column_metadata)); + } - let new_add_columns = new_column_indices.keys().map(|name| { - // Safety: previous steps ensure the physical region exist - let column_metadata = *physical_schema_map.get(name).unwrap(); - (name.to_string(), column_metadata.clone()) - }); + let mut new_add_columns = Vec::with_capacity(new_column_indices.len()); + for name in new_column_indices.keys() { + let column_metadata = physical_schema_map.get(name).copied().with_context(|| { + MissingPhysicalColumnSnafu { + name: name.to_string(), + region_id: physical_region_id, + } + })?; + new_add_columns.push((name.to_string(), column_metadata.clone())); + } // The data region has the new columns, whatever happens to the metadata write. self.state .write() @@ -337,7 +343,7 @@ impl MetricEngineInner { // Writes logical regions metadata to metadata region self.metadata_region - .add_logical_regions(physical_region_id, true, logical_region_columns) + .add_logical_regions(physical_region_id, true, logical_region_columns.into_iter()) .await?; self.state diff --git a/src/metric-engine/src/engine/state.rs b/src/metric-engine/src/engine/state.rs index e0f407b919b..cdc78fc3d00 100644 --- a/src/metric-engine/src/engine/state.rs +++ b/src/metric-engine/src/engine/state.rs @@ -43,6 +43,15 @@ pub struct PhysicalRegionState { primary_key_encoding: PrimaryKeyEncoding, options: PhysicalRegionOptions, time_index_unit: TimeUnit, + /// Serializes adding columns to this physical region: the column existence + /// check, the column id allocation and the physical ALTER that installs the + /// columns. Column ids are derived from the current metadata, so two + /// concurrent add-column requests on logical regions sharing this physical + /// region could otherwise assign the same id to different columns. + /// + /// Callers clone the [`Arc`] out of the state and hold it across the + /// physical alter; see [`crate::data_region::DataRegion::add_columns`]. + column_add_lock: Arc>, } impl PhysicalRegionState { @@ -67,9 +76,15 @@ impl PhysicalRegionState { primary_key_encoding, options, time_index_unit, + column_add_lock: Arc::new(tokio::sync::Mutex::new(())), } } + /// Returns the lock that serializes adding columns to this physical region. + pub fn column_add_lock(&self) -> &Arc> { + &self.column_add_lock + } + /// Returns a reference to the logical region ids. pub fn logical_regions(&self) -> &HashSet { &self.logical_regions diff --git a/src/metric-engine/src/error.rs b/src/metric-engine/src/error.rs index f9d61d6af31..f326429e2e0 100644 --- a/src/metric-engine/src/error.rs +++ b/src/metric-engine/src/error.rs @@ -255,6 +255,18 @@ pub enum Error { location: Location, }, + #[snafu(display( + "Column {} is missing from physical region {} after adding columns", + name, + region_id + ))] + MissingPhysicalColumn { + name: String, + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Table id count mismatch. Expect {}, got {}", expected, actual))] TableIdCountMismatch { expected: usize, @@ -451,6 +463,10 @@ impl ErrorExt for Error { ColumnNotFound { .. } => StatusCode::TableColumnNotFound, + // The physical region was read twice around the alter and lost a column in + // between, so the failure is an inconsistent server state, not a bad request. + MissingPhysicalColumn { .. } => StatusCode::Unexpected, + CreateMitoRegion { source, .. } | OpenMitoRegion { source, .. } | CleanUpMitoRegion { source, .. }