From 4e14ac86769c24f9185e9df41c54058a16567a98 Mon Sep 17 00:00:00 2001 From: QuakeWang <45645138+QuakeWang@users.noreply.github.com> Date: Fri, 10 Jul 2026 10:15:08 +0800 Subject: [PATCH 01/48] feat: add strict CSV header validation (#8426) * feat: add strict CSV header validation Signed-off-by: QuakeWang * style: simplify strict header parse error Signed-off-by: QuakeWang --------- Signed-off-by: QuakeWang --- src/common/datasource/src/file_format.rs | 1 + src/common/datasource/src/file_format/csv.rs | 39 +++- src/operator/src/error.rs | 17 ++ src/operator/src/statement/copy_table_from.rs | 178 +++++++++++++++++- tests-integration/src/tests/instance_test.rs | 128 +++++++++++++ 5 files changed, 361 insertions(+), 2 deletions(-) diff --git a/src/common/datasource/src/file_format.rs b/src/common/datasource/src/file_format.rs index 028aeef14d6..61f838899ef 100644 --- a/src/common/datasource/src/file_format.rs +++ b/src/common/datasource/src/file_format.rs @@ -63,6 +63,7 @@ pub const FORMAT_SCHEMA_INFER_MAX_RECORD: &str = "schema_infer_max_record"; pub const FORMAT_HEADERS: &str = "headers"; pub const FORMAT_HAS_HEADER: &str = "has_header"; pub const FORMAT_SKIP_BAD_RECORDS: &str = "skip_bad_records"; +pub const FORMAT_STRICT_HEADERS: &str = "strict_headers"; pub const FORMAT_TYPE: &str = "format"; pub const FILE_PATTERN: &str = "pattern"; pub const TIMESTAMP_FORMAT: &str = "timestamp_format"; diff --git a/src/common/datasource/src/file_format/csv.rs b/src/common/datasource/src/file_format/csv.rs index 24b61340f5b..d016d4eabe9 100644 --- a/src/common/datasource/src/file_format/csv.rs +++ b/src/common/datasource/src/file_format/csv.rs @@ -49,6 +49,7 @@ const SKIP_BAD_RECORDS_BATCH_SIZE: usize = 1; pub struct CsvFormat { pub has_header: bool, pub skip_bad_records: bool, + pub strict_headers: bool, pub delimiter: u8, pub schema_infer_max_record: Option, pub compression_type: CompressionType, @@ -111,6 +112,16 @@ impl TryFrom<&HashMap> for CsvFormat { format.skip_bad_records = parse_bool(file_format::FORMAT_SKIP_BAD_RECORDS, skip_bad_records)?; }; + if let Some(strict_headers) = value.get(file_format::FORMAT_STRICT_HEADERS) { + format.strict_headers = parse_bool(file_format::FORMAT_STRICT_HEADERS, strict_headers)?; + }; + if format.strict_headers && !format.has_header { + return error::ParseFormatSnafu { + key: file_format::FORMAT_STRICT_HEADERS, + value: "strict_headers=true requires headers=true", + } + .fail(); + } if let Some(timestamp_format) = value.get(file_format::TIMESTAMP_FORMAT) { format.timestamp_format = Some(timestamp_format.clone()); } @@ -135,6 +146,7 @@ impl Default for CsvFormat { Self { has_header: true, skip_bad_records: false, + strict_headers: false, delimiter: b',', schema_infer_max_record: Some(file_format::DEFAULT_SCHEMA_INFER_MAX_RECORD), compression_type: CompressionType::Uncompressed, @@ -365,7 +377,8 @@ mod tests { use super::*; use crate::file_format::{ FORMAT_COMPRESSION_TYPE, FORMAT_DELIMITER, FORMAT_HAS_HEADER, FORMAT_HEADERS, - FORMAT_SCHEMA_INFER_MAX_RECORD, FORMAT_SKIP_BAD_RECORDS, FileFormat, file_to_stream, + FORMAT_SCHEMA_INFER_MAX_RECORD, FORMAT_SKIP_BAD_RECORDS, FORMAT_STRICT_HEADERS, FileFormat, + file_to_stream, }; use crate::test_util::{format_schema, test_store}; @@ -492,6 +505,7 @@ mod tests { delimiter: b'\t', has_header: false, skip_bad_records: false, + strict_headers: false, timestamp_format: None, time_format: None, date_format: None @@ -509,6 +523,17 @@ mod tests { } ); + let map = HashMap::from([(FORMAT_STRICT_HEADERS.to_string(), "true".to_string())]); + let format = CsvFormat::try_from(&map).unwrap(); + + assert_eq!( + format, + CsvFormat { + strict_headers: true, + ..CsvFormat::default() + } + ); + let map = HashMap::from([(FORMAT_HEADERS.to_string(), "true".to_string())]); let format = CsvFormat::try_from(&map).unwrap(); assert_eq!(format, CsvFormat::default()); @@ -544,6 +569,9 @@ mod tests { let map = HashMap::from([(FORMAT_HEADERS.to_string(), "yes".to_string())]); assert!(CsvFormat::try_from(&map).is_err()); + + let map = HashMap::from([(FORMAT_STRICT_HEADERS.to_string(), "yes".to_string())]); + assert!(CsvFormat::try_from(&map).is_err()); } #[test] @@ -555,6 +583,15 @@ mod tests { assert!(CsvFormat::try_from(&map).is_err()); } + #[test] + fn test_try_from_rejects_strict_headerless_csv() { + let map = HashMap::from([ + (FORMAT_HEADERS.to_string(), "false".to_string()), + (FORMAT_STRICT_HEADERS.to_string(), "true".to_string()), + ]); + assert!(CsvFormat::try_from(&map).is_err()); + } + #[tokio::test] async fn test_compressed_csv() { // Create test data diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index 13b86140aa1..66fb66d7527 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -589,6 +589,22 @@ pub enum Error { location: Location, }, + #[snafu(display( + "CSV header mismatch in path: {}, unknown columns: {:?}, missing columns: {:?}, duplicate columns: {:?}", + path, + unknown_columns, + missing_columns, + duplicate_columns + ))] + CsvHeaderMismatch { + path: String, + unknown_columns: Vec, + missing_columns: Vec, + duplicate_columns: Vec, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to project schema"))] ProjectSchema { #[snafu(source)] @@ -936,6 +952,7 @@ impl ErrorExt for Error { | Error::ColumnNotFound { .. } | Error::BuildRegex { .. } | Error::InvalidSchema { .. } + | Error::CsvHeaderMismatch { .. } | Error::ProjectSchema { .. } | Error::UnsupportedFormat { .. } | Error::ColumnNoneDefaultValue { .. } diff --git a/src/operator/src/statement/copy_table_from.rs b/src/operator/src/statement/copy_table_from.rs index b64c654923c..014a2d32e9b 100644 --- a/src/operator/src/statement/copy_table_from.rs +++ b/src/operator/src/statement/copy_table_from.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::future::Future; use std::path::Path; use std::pin::Pin; @@ -392,6 +392,7 @@ impl StatementExecutor { .collect_metadata(&object_store, format.clone(), path.to_string()) .await?; + validate_csv_headers_if_required(&file_metadata, &table_schema)?; let schema_mapping = copy_from_schema_mapping(&file_metadata, &table_schema); let projected_file_schema = Arc::new( file_metadata @@ -642,6 +643,62 @@ fn ensure_schema_compatible(from: &SchemaRef, to: &SchemaRef) -> Result<()> { } } +fn validate_csv_headers_if_required(file_metadata: &FileMetadata, table: &SchemaRef) -> Result<()> { + let FileMetadata::Csv { + schema, + format, + path, + } = file_metadata + else { + return Ok(()); + }; + + if !format.strict_headers { + return Ok(()); + } + + let mut seen_file_columns = HashSet::with_capacity(schema.fields().len()); + let duplicate_columns = schema + .fields() + .iter() + .filter_map(|field| { + if seen_file_columns.insert(field.name().clone()) { + None + } else { + Some(field.name().clone()) + } + }) + .collect::>() + .into_iter() + .collect::>(); + let file_columns = seen_file_columns.into_iter().collect::>(); + let table_columns = table + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + let unknown_columns = file_columns + .difference(&table_columns) + .cloned() + .collect::>(); + let missing_columns = table_columns + .difference(&file_columns) + .cloned() + .collect::>(); + + ensure!( + unknown_columns.is_empty() && missing_columns.is_empty() && duplicate_columns.is_empty(), + error::CsvHeaderMismatchSnafu { + path, + unknown_columns, + missing_columns, + duplicate_columns, + } + ); + + Ok(()) +} + /// Generates a maybe compatible schema of the file schema. /// /// If there is a field is found in table schema, @@ -998,12 +1055,131 @@ mod tests { } } + fn make_strict_csv_metadata(schema: Arc) -> FileMetadata { + FileMetadata::Csv { + schema, + format: CsvFormat { + strict_headers: true, + ..CsvFormat::default() + }, + path: "test.csv".to_string(), + } + } + fn assert_field(schema: &Schema, idx: usize, name: &str, data_type: &DataType) { let field = schema.field(idx); assert_eq!(field.name(), name); assert_eq!(field.data_type(), data_type); } + #[test] + fn test_strict_csv_headers_allows_reordered_columns() { + let file_schema = make_test_schema(&[ + Field::new("ts", DataType::Utf8, true), + Field::new("host_id", DataType::UInt8, true), + Field::new("reading_value", DataType::Float64, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap(); + } + + #[test] + fn test_strict_csv_headers_rejects_unknown_columns() { + let file_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt8, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + Field::new("extra", DataType::Utf8, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + let err = + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap_err(); + + assert!(matches!( + err, + error::Error::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + } if unknown_columns == vec!["extra".to_string()] + && missing_columns.is_empty() + && duplicate_columns.is_empty() + )); + } + + #[test] + fn test_strict_csv_headers_rejects_missing_columns() { + let file_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt8, true), + Field::new("ts", DataType::Utf8, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + let err = + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap_err(); + + assert!(matches!( + err, + error::Error::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + } if unknown_columns.is_empty() + && missing_columns == vec!["reading_value".to_string()] + && duplicate_columns.is_empty() + )); + } + + #[test] + fn test_strict_csv_headers_rejects_duplicate_columns() { + let file_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt8, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + Field::new("host_id", DataType::UInt16, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + let err = + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap_err(); + + assert!(matches!( + err, + error::Error::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + } if unknown_columns.is_empty() + && missing_columns.is_empty() + && duplicate_columns == vec!["host_id".to_string()] + )); + } + #[test] fn test_headerless_csv_schema_projection_is_positional() { let file_schema = make_test_schema(&[ diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index fa9bad0fac0..628f97de96a 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -1090,6 +1090,134 @@ async fn test_execute_copy_from_headerless_csv(instance: Arc) check_output_stream(output, expect).await; } +#[apply(both_instances_cases)] +async fn test_execute_copy_from_csv_strict_headers(instance: Arc) { + let instance = instance.frontend(); + let tmp_dir = temp_dir::create_temp_dir("test_execute_copy_from_csv_strict_headers"); + let matching_path = tmp_dir.path().join("matching.csv"); + let unknown_path = tmp_dir.path().join("unknown.csv"); + let missing_path = tmp_dir.path().join("missing.csv"); + let duplicate_path = tmp_dir.path().join("duplicate.csv"); + std::fs::write( + &matching_path, + "ts,host_id,reading_value\n2024-01-01T00:00:00,1,10.5\n", + ) + .unwrap(); + std::fs::write( + &unknown_path, + "host_id,reading_value,ts,extra\n1,10.5,2024-01-01T00:00:00,ignored\n", + ) + .unwrap(); + std::fs::write(&missing_path, "host_id,ts\n1,2024-01-01T00:00:00\n").unwrap(); + std::fs::write( + &duplicate_path, + "host_id,reading_value,ts,host_id\n1,10.5,2024-01-01T00:00:00,2\n", + ) + .unwrap(); + + let output = execute_sql( + &instance, + "CREATE TABLE csv_strict_headers(host_id INT, reading_value DOUBLE, ts TIMESTAMP TIME INDEX);", + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(0))); + + let output = execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + matching_path.display() + ), + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(1))); + + let output = execute_sql(&instance, "SELECT * FROM csv_strict_headers;") + .await + .data; + let expect = "\ ++---------+---------------+---------------------+ +| host_id | reading_value | ts | ++---------+---------------+---------------------+ +| 1 | 10.5 | 2024-01-01T00:00:00 | ++---------+---------------+---------------------+"; + check_output_stream(output, expect).await; + + let err = try_execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + unknown_path.display() + ), + ) + .await + .unwrap_err(); + assert!(matches!( + unwrap_frontend_error(&err), + Error::TableOperation { + source: OperatorError::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + }, + .. + } if unknown_columns == &vec!["extra".to_string()] + && missing_columns.is_empty() + && duplicate_columns.is_empty() + )); + + let err = try_execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + missing_path.display() + ), + ) + .await + .unwrap_err(); + assert!(matches!( + unwrap_frontend_error(&err), + Error::TableOperation { + source: OperatorError::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + }, + .. + } if unknown_columns.is_empty() + && missing_columns == &vec!["reading_value".to_string()] + && duplicate_columns.is_empty() + )); + + let err = try_execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + duplicate_path.display() + ), + ) + .await + .unwrap_err(); + assert!(matches!( + unwrap_frontend_error(&err), + Error::TableOperation { + source: OperatorError::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + }, + .. + } if unknown_columns.is_empty() + && missing_columns.is_empty() + && duplicate_columns == &vec!["host_id".to_string()] + )); +} + #[apply(both_instances_cases)] async fn test_execute_query_external_table_json(instance: Arc) { unsafe { From 74bdb1f2a88bec7979b0029bc37408446c4f57ea Mon Sep 17 00:00:00 2001 From: raphaelroshan <49832307+raphaelroshan@users.noreply.github.com> Date: Fri, 10 Jul 2026 13:46:20 +0800 Subject: [PATCH 02/48] fix: compare all LoggingOptions fields in PartialEq (#8449) * fix: compare all LoggingOptions fields in PartialEq Signed-off-by: raphaelroshan * test: set otlp_export_protocol in load_config_test expected configs The example configs set otlp_export_protocol = "http", but the expected LoggingOptions blocks for datanode/frontend/metasrv/standalone left it at the default None. Now that PartialEq compares all fields, add the field to match the parsed configs (as the flownode case already does). Signed-off-by: raphaelroshan --------- Signed-off-by: raphaelroshan --- src/cmd/tests/load_config_test.rs | 4 +++ src/common/telemetry/src/logging.rs | 39 ++++++++++++++++++----------- 2 files changed, 28 insertions(+), 15 deletions(-) diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 74e4127de83..3570e88cc1e 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -108,6 +108,7 @@ fn test_load_datanode_example_config() { level: Some("info".to_string()), dir: format!("{}/{}", DEFAULT_DATA_HOME, DEFAULT_LOGGING_DIR), otlp_endpoint: Some(DEFAULT_OTLP_HTTP_ENDPOINT.to_string()), + otlp_export_protocol: Some(common_telemetry::logging::OtlpExportProtocol::Http), tracing_sample_ratio: Some(Default::default()), ..Default::default() }, @@ -147,6 +148,7 @@ fn test_load_frontend_example_config() { level: Some("info".to_string()), dir: format!("{}/{}", DEFAULT_DATA_HOME, DEFAULT_LOGGING_DIR), otlp_endpoint: Some(DEFAULT_OTLP_HTTP_ENDPOINT.to_string()), + otlp_export_protocol: Some(common_telemetry::logging::OtlpExportProtocol::Http), tracing_sample_ratio: Some(Default::default()), ..Default::default() }, @@ -197,6 +199,7 @@ fn test_load_metasrv_example_config() { dir: format!("{}/{}", DEFAULT_DATA_HOME, DEFAULT_LOGGING_DIR), level: Some("info".to_string()), otlp_endpoint: Some(DEFAULT_OTLP_HTTP_ENDPOINT.to_string()), + otlp_export_protocol: Some(common_telemetry::logging::OtlpExportProtocol::Http), tracing_sample_ratio: Some(Default::default()), ..Default::default() }, @@ -314,6 +317,7 @@ fn test_load_standalone_example_config() { level: Some("info".to_string()), dir: format!("{}/{}", DEFAULT_DATA_HOME, DEFAULT_LOGGING_DIR), otlp_endpoint: Some(DEFAULT_OTLP_HTTP_ENDPOINT.to_string()), + otlp_export_protocol: Some(common_telemetry::logging::OtlpExportProtocol::Http), tracing_sample_ratio: Some(Default::default()), ..Default::default() }, diff --git a/src/common/telemetry/src/logging.rs b/src/common/telemetry/src/logging.rs index 953dbbc2df9..56513531632 100644 --- a/src/common/telemetry/src/logging.rs +++ b/src/common/telemetry/src/logging.rs @@ -237,7 +237,7 @@ enum TraceState { } /// The logging options that used to initialize the logger. -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(default)] pub struct LoggingOptions { /// The directory to store log files. If not set, logs will be written to stdout. @@ -341,20 +341,6 @@ pub enum LogFormat { Text, } -impl PartialEq for LoggingOptions { - fn eq(&self, other: &Self) -> bool { - self.dir == other.dir - && self.level == other.level - && self.enable_otlp_tracing == other.enable_otlp_tracing - && self.otlp_endpoint == other.otlp_endpoint - && self.tracing_sample_ratio == other.tracing_sample_ratio - && self.append_stdout == other.append_stdout - && self.enable_per_region_metrics == other.enable_per_region_metrics - } -} - -impl Eq for LoggingOptions {} - #[derive(Clone, Debug)] struct TraceContext { app_name: String, @@ -874,4 +860,27 @@ mod tests { let protocol: OtlpExportProtocol = serde_json::from_str(http_json).unwrap(); assert_eq!(protocol, OtlpExportProtocol::Http); } + + #[test] + fn test_logging_options_partial_eq_all_fields() { + let base = LoggingOptions::default(); + + let mut log_format = base.clone(); + log_format.log_format = LogFormat::Json; + assert_ne!(base, log_format); + + let mut max_log_files = base.clone(); + max_log_files.max_log_files += 1; + assert_ne!(base, max_log_files); + + let mut otlp_export_protocol = base.clone(); + otlp_export_protocol.otlp_export_protocol = Some(OtlpExportProtocol::Http); + assert_ne!(base, otlp_export_protocol); + + let mut otlp_headers = base.clone(); + otlp_headers + .otlp_headers + .insert("key".to_string(), "value".to_string()); + assert_ne!(base, otlp_headers); + } } From ff2fa71d50af62815d23ea83ac34d0d4bf043e41 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 10 Jul 2026 14:10:28 +0800 Subject: [PATCH 03/48] fix: repartition subset partition key joins (#8460) Signed-off-by: discord9 --- src/query/src/dist_plan/merge_scan.rs | 18 +++ src/query/src/optimizer/pass_distribution.rs | 40 +++++- ..._distribution_partition_subset_join.result | 118 ++++++++++++++++++ ...ass_distribution_partition_subset_join.sql | 52 ++++++++ 4 files changed, 223 insertions(+), 5 deletions(-) create mode 100644 tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.result create mode 100644 tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.sql diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 27a0d9196bb..39bae738c65 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -554,6 +554,24 @@ impl MergeScanExec { return None; } + let hash_expr_col_names: HashSet<_> = hash_exprs + .iter() + .filter_map(|expr| { + expr.as_any() + .downcast_ref::() + .map(|col_expr| col_expr.name()) + }) + .collect(); + + let covers_all_partition_cols = self.partition_cols.values().all(|aliases| { + aliases + .iter() + .any(|col| hash_expr_col_names.contains(col.name())) + }); + if !covers_all_partition_cols { + return None; + } + let all_partition_col_aliases: HashSet<_> = self .partition_cols .values() diff --git a/src/query/src/optimizer/pass_distribution.rs b/src/query/src/optimizer/pass_distribution.rs index 22e2b903593..8ff21e046d4 100644 --- a/src/query/src/optimizer/pass_distribution.rs +++ b/src/query/src/optimizer/pass_distribution.rs @@ -16,8 +16,9 @@ use std::sync::Arc; use datafusion::config::ConfigOptions; use datafusion::physical_optimizer::PhysicalOptimizerRule; -use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::{ExecutionPlan, Partitioning}; use datafusion_common::Result as DfResult; use datafusion_physical_expr::Distribution; use datafusion_physical_expr::utils::map_columns_before_projection; @@ -68,11 +69,26 @@ impl PassDistribution { ) -> DfResult> { // If this is a MergeScanExec, try to apply the current requirement. if let Some(merge_scan) = plan.as_any().downcast_ref::() - && let Some(distribution) = current_req.as_ref() - && let Some(new_plan) = merge_scan.try_with_new_distribution(distribution.clone()) + && let Some(Distribution::HashPartitioned(hash_exprs)) = current_req.as_ref() { - // Leaf node; no children to process - return Ok(Arc::new(new_plan) as _); + if let Partitioning::Hash(current_hash_exprs, _) = &merge_scan.properties().partitioning + && *current_hash_exprs == *hash_exprs + { + return Ok(plan); + } + + if let Some(new_plan) = merge_scan + .try_with_new_distribution(Distribution::HashPartitioned(hash_exprs.clone())) + { + // Leaf node; no children to process + return Ok(Arc::new(new_plan) as _); + } + + let partitioning = Partitioning::Hash( + hash_exprs.clone(), + merge_scan.properties().partitioning.partition_count(), + ); + return Ok(Arc::new(RepartitionExec::try_new(plan, partitioning)?) as _); } // Compute per-child requirements from the current node. @@ -266,6 +282,20 @@ mod tests { ); } + #[test] + fn merge_scan_rejects_hash_requirement_on_partition_key_subset() { + let merge_scan = test_merge_scan_exec(test_schema()); + + let new_plan = merge_scan.try_with_new_distribution(Distribution::HashPartitioned(vec![ + partition_column(DATA_SCHEMA_TSID_COLUMN_NAME, 1), + ])); + + assert!( + new_plan.is_none(), + "partitioning by a subset of multi-column partition keys is not sufficient" + ); + } + fn test_merge_scan_exec(schema: SchemaRef) -> MergeScanExec { let session_state = SessionStateBuilder::new().with_default_features().build(); let partition_cols = BTreeMap::from([ diff --git a/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.result b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.result new file mode 100644 index 00000000000..91c0e3e851d --- /dev/null +++ b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.result @@ -0,0 +1,118 @@ +CREATE TABLE pd_sub_l( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + lv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +Affected Rows: 0 + +CREATE TABLE pd_sub_r( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + rv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +Affected Rows: 0 + +INSERT INTO pd_sub_l VALUES + (1, 1, 'a', 10), + (2, 1, 'z', 20); + +Affected Rows: 2 + +INSERT INTO pd_sub_r VALUES + (1, 1, 'a', 100), + (2, 1, 'z', 200); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('pd_sub_l'); + ++-------------------------------+ +| ADMIN FLUSH_TABLE('pd_sub_l') | ++-------------------------------+ +| 0 | ++-------------------------------+ + +ADMIN FLUSH_TABLE('pd_sub_r'); + ++-------------------------------+ +| ADMIN FLUSH_TABLE('pd_sub_r') | ++-------------------------------+ +| 0 | ++-------------------------------+ + +-- Joining on only a subset of a multi-column partition key must still compare +-- rows across all matching partitions for that subset key. +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \d+\(\d+,\s+\d+\) REDACTED_REGION +-- SQLNESS REPLACE Hash\(\[device_id@0\],\s+\d+\) Hash([device_id@0],REDACTED_PARTITIONS) +-- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED +-- SQLNESS REPLACE input_partitions=REDACTED\s+\| input_partitions=REDACTED| +EXPLAIN SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + ++---------------+-------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: l.area ASC NULLS LAST, r.area ASC NULLS LAST | +| | Projection: l.area, r.area | +| | Inner Join: l.device_id = r.device_id | +| | Projection: l.device_id, l.area | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: l | +| | Filter: pd_sub_l.device_id IS NOT NULL | +| | TableScan: pd_sub_l, partial_filters=[pd_sub_l.device_id IS NOT NULL] | +| | ]] | +| | Projection: r.device_id, r.area | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: r | +| | Filter: pd_sub_r.device_id IS NOT NULL | +| | TableScan: pd_sub_r, partial_filters=[pd_sub_r.device_id IS NOT NULL] | +| | ]] | +| physical_plan | SortPreservingMergeExec: [area@0 ASC NULLS LAST, area@1 ASC NULLS LAST] | +| | SortExec: expr=[area@0 ASC NULLS LAST, area@1 ASC NULLS LAST], preserve_partitioning=[true] | +| | HashJoinExec: mode=Partitioned, join_type=Inner, on=[(device_id@0, device_id@0)], projection=[area@1, area@3] | +| | RepartitionExec: partitioning=Hash([device_id@0],REDACTED_PARTITIONS), input_partitions=REDACTED| +| | ProjectionExec: expr=[device_id@1 as device_id, area@2 as area] | +| | MergeScanExec: peers=[REDACTED_REGION, REDACTED_REGION, REDACTED_REGION, ] | +| | RepartitionExec: partitioning=Hash([device_id@0],REDACTED_PARTITIONS), input_partitions=REDACTED| +| | ProjectionExec: expr=[device_id@1 as device_id, area@2 as area] | +| | MergeScanExec: peers=[REDACTED_REGION, REDACTED_REGION, REDACTED_REGION, ] | +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------+ + +SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + ++------+------+ +| area | area | ++------+------+ +| a | a | +| a | z | +| z | a | +| z | z | ++------+------+ + +DROP TABLE pd_sub_l; + +Affected Rows: 0 + +DROP TABLE pd_sub_r; + +Affected Rows: 0 + diff --git a/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.sql b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.sql new file mode 100644 index 00000000000..41337e74a0f --- /dev/null +++ b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.sql @@ -0,0 +1,52 @@ +CREATE TABLE pd_sub_l( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + lv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +CREATE TABLE pd_sub_r( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + rv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +INSERT INTO pd_sub_l VALUES + (1, 1, 'a', 10), + (2, 1, 'z', 20); + +INSERT INTO pd_sub_r VALUES + (1, 1, 'a', 100), + (2, 1, 'z', 200); + +ADMIN FLUSH_TABLE('pd_sub_l'); +ADMIN FLUSH_TABLE('pd_sub_r'); + +-- Joining on only a subset of a multi-column partition key must still compare +-- rows across all matching partitions for that subset key. +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \d+\(\d+,\s+\d+\) REDACTED_REGION +-- SQLNESS REPLACE Hash\(\[device_id@0\],\s+\d+\) Hash([device_id@0],REDACTED_PARTITIONS) +-- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED +-- SQLNESS REPLACE input_partitions=REDACTED\s+\| input_partitions=REDACTED| +EXPLAIN SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + +SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + +DROP TABLE pd_sub_l; +DROP TABLE pd_sub_r; From 29423f8ad865d8835aec381e58b3ebf8415d2b16 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 10 Jul 2026 14:40:16 +0800 Subject: [PATCH 04/48] ci: check kafka readiness before fuzz tests (#8461) Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Signed-off-by: WenyXu Co-authored-by: Sisyphus --- .github/scripts/wait-kafka-ready.sh | 68 +++++++++++++++++++++++++++++ .github/workflows/develop.yml | 10 +++++ 2 files changed, 78 insertions(+) create mode 100755 .github/scripts/wait-kafka-ready.sh diff --git a/.github/scripts/wait-kafka-ready.sh b/.github/scripts/wait-kafka-ready.sh new file mode 100755 index 00000000000..7617bdfa1f2 --- /dev/null +++ b/.github/scripts/wait-kafka-ready.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash + +set -euo pipefail + +GT_KAFKA_NAMESPACE="${GT_KAFKA_NAMESPACE:-kafka-cluster}" +GT_KAFKA_CLIENT_NAMESPACE="${GT_KAFKA_CLIENT_NAMESPACE:-my-greptimedb}" +GT_KAFKA_BOOTSTRAP_SERVERS="${GT_KAFKA_BOOTSTRAP_SERVERS:-kafka.kafka-cluster.svc.cluster.local:9092}" +GT_KAFKA_IMAGE="${GT_KAFKA_IMAGE:-docker.io/greptime/kafka:3.9.0-debian-12-r1}" +GT_KAFKA_READY_TIMEOUT_SECS="${GT_KAFKA_READY_TIMEOUT_SECS:-120}" +GT_KAFKA_READY_INTERVAL_SECS="${GT_KAFKA_READY_INTERVAL_SECS:-5}" +GT_KAFKA_POD_LABEL="${GT_KAFKA_POD_LABEL:-app.kubernetes.io/instance=kafka}" + +log() { + printf '[wait-kafka-ready] %s\n' "$*" +} + +run_client_probe() { + local pod_name="kafka-readiness-${RANDOM}-${RANDOM}" + + kubectl run "${pod_name}" \ + --namespace "${GT_KAFKA_CLIENT_NAMESPACE}" \ + --image "${GT_KAFKA_IMAGE}" \ + --restart=Never \ + --rm \ + --attach \ + --quiet \ + --env "GT_KAFKA_BOOTSTRAP_SERVERS=${GT_KAFKA_BOOTSTRAP_SERVERS}" \ + --command -- bash -euo pipefail -c ' + if command -v kafka-broker-api-versions.sh >/dev/null 2>&1; then + kafka_tool="kafka-broker-api-versions.sh" + elif [ -x /opt/bitnami/kafka/bin/kafka-broker-api-versions.sh ]; then + kafka_tool="/opt/bitnami/kafka/bin/kafka-broker-api-versions.sh" + else + echo "kafka-broker-api-versions.sh not found" >&2 + exit 1 + fi + + "${kafka_tool}" --bootstrap-server "${GT_KAFKA_BOOTSTRAP_SERVERS}" >/tmp/kafka-ready.out + cat /tmp/kafka-ready.out + ' +} + +log "wait for Kafka pods in namespace ${GT_KAFKA_NAMESPACE}" +kubectl wait \ + --for=condition=Ready \ + pod -l "${GT_KAFKA_POD_LABEL}" \ + --timeout="${GT_KAFKA_READY_TIMEOUT_SECS}s" \ + -n "${GT_KAFKA_NAMESPACE}" + +deadline=$((SECONDS + GT_KAFKA_READY_TIMEOUT_SECS)) +attempt=1 + +while [ "${SECONDS}" -le "${deadline}" ]; do + log "probe Kafka from namespace ${GT_KAFKA_CLIENT_NAMESPACE}, attempt ${attempt}" + if run_client_probe; then + log "Kafka is ready at ${GT_KAFKA_BOOTSTRAP_SERVERS}" + exit 0 + fi + + log "Kafka probe failed; retrying in ${GT_KAFKA_READY_INTERVAL_SECS}s" + sleep "${GT_KAFKA_READY_INTERVAL_SECS}" + attempt=$((attempt + 1)) +done + +log "timed out waiting for Kafka bootstrap ${GT_KAFKA_BOOTSTRAP_SERVERS}" +kubectl get pods -o wide -n "${GT_KAFKA_NAMESPACE}" || true +kubectl get pods -o wide -n "${GT_KAFKA_CLIENT_NAMESPACE}" || true +exit 1 diff --git a/.github/workflows/develop.yml b/.github/workflows/develop.yml index 08c44a826b2..0d9eceacb8e 100644 --- a/.github/workflows/develop.yml +++ b/.github/workflows/develop.yml @@ -457,6 +457,11 @@ jobs: name: Port forward Kafka WAL helper run: | kubectl port-forward service/kafka-wal-helper 8080:8080 -n kafka-cluster& + - if: matrix.mode.kafka + name: Check Kafka readiness + shell: bash + run: | + .github/scripts/wait-kafka-ready.sh - name: Fuzz Test uses: ./.github/actions/fuzz-test env: @@ -648,6 +653,11 @@ jobs: - name: Port forward (mysql) run: | kubectl port-forward service/my-greptimedb-frontend 4002:4002 -n my-greptimedb& + - if: matrix.mode.kafka + name: Check Kafka readiness + shell: bash + run: | + .github/scripts/wait-kafka-ready.sh - name: Fuzz Test uses: ./.github/actions/fuzz-test env: From 03b304522f41adba6f12af871419279a9ce7b79a Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Fri, 10 Jul 2026 16:53:23 +0800 Subject: [PATCH 05/48] refactor: model distributed inspect fan-in as exec (#8447) * refactor: model distributed inspect fan-in as exec Signed-off-by: jeremyhi * fix: harden distributed inspect exec fan-in Signed-off-by: jeremyhi * fix: show distributed inspect scan in plans Signed-off-by: jeremyhi * fix: stabilize distributed inspect sqlness Signed-off-by: jeremyhi * fix: apply sqlness replacement to inspect explain Signed-off-by: jeremyhi * fix: redact inspect repartition plan Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/catalog/src/information_extension.rs | 227 +++++++++++++++--- src/catalog/src/system_schema.rs | 16 ++ .../src/system_schema/information_schema.rs | 19 ++ .../information_schema/region_info.rs | 17 ++ .../system_schema/information_schema/ssts.rs | 173 +++++++++++-- src/store-api/src/data_source.rs | 10 + src/table/src/table.rs | 7 + src/table/src/table/adapter.rs | 5 + .../information_schema/ssts_limit.result | 31 +++ .../information_schema/ssts_limit.sql | 9 + 10 files changed, 463 insertions(+), 51 deletions(-) diff --git a/src/catalog/src/information_extension.rs b/src/catalog/src/information_extension.rs index 08d81ac1f03..c41c816c6d7 100644 --- a/src/catalog/src/information_extension.rs +++ b/src/catalog/src/information_extension.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use api::v1::meta::ProcedureStatus; use common_error::ext::BoxedError; use common_meta::cluster::{ClusterInfo, NodeInfo, Role}; @@ -22,8 +24,16 @@ use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor}; use common_meta::rpc::procedure; use common_procedure::{ProcedureInfo, ProcedureState}; use common_query::request::QueryRequest; -use common_recordbatch::SendableRecordBatchStream; +use common_recordbatch::adapter::{AsyncRecordBatchStreamAdapter, DfRecordBatchStreamAdapter}; use common_recordbatch::util::{ChainedRecordBatchStream, LimitedRecordBatchStream}; +use common_recordbatch::{DfSendableRecordBatchStream, SendableRecordBatchStream}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; +use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datatypes::schema::SchemaRef; +use futures_util::future::try_join_all; use meta_client::MetaClientRef; use snafu::ResultExt; use store_api::storage::RegionId; @@ -43,6 +53,168 @@ impl DistributedInformationExtension { datanode_manager, } } + + async fn inspect_datanode_stream( + meta_client: MetaClientRef, + datanode_manager: DatanodeManagerRef, + request: DatanodeInspectRequest, + ) -> std::result::Result { + let nodes = meta_client + .list_nodes(Some(Role::Datanode)) + .await + .map_err(BoxedError::new) + .context(crate::error::ListNodesSnafu)?; + + let limit = request.scan.limit; + let plan = request + .build_plan() + .context(crate::error::DatafusionSnafu)?; + + let streams = try_join_all(nodes.into_iter().map(|node| { + let datanode_manager = datanode_manager.clone(); + let plan = plan.clone(); + async move { + let client = datanode_manager.datanode(&node.peer).await; + client + .handle_query(QueryRequest { + plan, + region_id: RegionId::default(), + header: None, + }) + .await + .context(crate::error::HandleQuerySnafu) + } + })) + .await?; + + let chained = + ChainedRecordBatchStream::new(streams).context(crate::error::CreateRecordBatchSnafu)?; + match limit { + Some(limit) => Ok(Box::pin(LimitedRecordBatchStream::new( + Box::pin(chained), + limit, + ))), + None => Ok(Box::pin(chained)), + } + } +} + +struct DistributedInspectExec { + request: DatanodeInspectRequest, + meta_client: MetaClientRef, + datanode_manager: DatanodeManagerRef, + schema: SchemaRef, + arrow_schema: ArrowSchemaRef, + properties: Arc, +} + +impl std::fmt::Debug for DistributedInspectExec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DistributedInspectExec") + .field("request", &self.request) + .field("schema", &self.schema) + .finish_non_exhaustive() + } +} + +impl DistributedInspectExec { + fn new( + request: DatanodeInspectRequest, + schema: SchemaRef, + meta_client: MetaClientRef, + datanode_manager: DatanodeManagerRef, + ) -> Self { + let arrow_schema = schema.arrow_schema().clone(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(arrow_schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + request, + meta_client, + datanode_manager, + schema, + arrow_schema, + properties, + } + } +} + +impl DisplayAs for DistributedInspectExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "DistributedInspectExec: kind={:?}, scan={}, schema={:?}", + self.request.kind, self.request.scan, self.arrow_schema + ) + } +} + +impl ExecutionPlan for DistributedInspectExec { + fn name(&self) -> &str { + "DistributedInspectExec" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn schema(&self) -> ArrowSchemaRef { + self.arrow_schema.clone() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> datafusion::error::Result> { + if children.is_empty() { + Ok(self) + } else { + Err(datafusion::error::DataFusionError::Internal( + "DistributedInspectExec is a leaf execution plan and cannot accept children" + .to_string(), + )) + } + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> datafusion::error::Result { + if partition != 0 { + return Err(datafusion::error::DataFusionError::Execution(format!( + "DistributedInspectExec invalid partition. Expected 0, got {partition}" + ))); + } + + let meta_client = self.meta_client.clone(); + let datanode_manager = self.datanode_manager.clone(); + let request = self.request.clone(); + let schema = self.schema.clone(); + let future = async move { + DistributedInformationExtension::inspect_datanode_stream( + meta_client, + datanode_manager, + request, + ) + .await + .map_err(BoxedError::new) + .context(common_recordbatch::error::ExternalSnafu) + }; + let stream = AsyncRecordBatchStreamAdapter::new(schema, Box::pin(future)); + Ok(Box::pin(DfRecordBatchStreamAdapter::new(Box::pin(stream)))) + } } #[async_trait::async_trait] @@ -112,41 +284,24 @@ impl InformationExtension for DistributedInformationExtension { &self, request: DatanodeInspectRequest, ) -> std::result::Result { - // Aggregate results from all datanodes - let nodes = self - .meta_client - .list_nodes(Some(Role::Datanode)) - .await - .map_err(BoxedError::new) - .context(crate::error::ListNodesSnafu)?; + Self::inspect_datanode_stream( + self.meta_client.clone(), + self.datanode_manager.clone(), + request, + ) + .await + } - let limit = request.scan.limit; - let plan = request - .build_plan() - .context(crate::error::DatafusionSnafu)?; - - let mut streams = Vec::with_capacity(nodes.len()); - for node in nodes { - let client = self.datanode_manager.datanode(&node.peer).await; - let stream = client - .handle_query(QueryRequest { - plan: plan.clone(), - region_id: RegionId::default(), - header: None, - }) - .await - .context(crate::error::HandleQuerySnafu)?; - streams.push(stream); - } - - let chained = - ChainedRecordBatchStream::new(streams).context(crate::error::CreateRecordBatchSnafu)?; - match limit { - Some(limit) => Ok(Box::pin(LimitedRecordBatchStream::new( - Box::pin(chained), - limit, - ))), - None => Ok(Box::pin(chained)), - } + fn inspect_datanode_plan( + &self, + request: DatanodeInspectRequest, + schema: SchemaRef, + ) -> std::result::Result>, Self::Error> { + Ok(Some(Arc::new(DistributedInspectExec::new( + request, + schema, + self.meta_client.clone(), + self.datanode_manager.clone(), + )))) } } diff --git a/src/catalog/src/system_schema.rs b/src/catalog/src/system_schema.rs index 6c6a5fc2028..fe926d893c9 100644 --- a/src/catalog/src/system_schema.rs +++ b/src/catalog/src/system_schema.rs @@ -25,6 +25,7 @@ use std::sync::Arc; use common_error::ext::BoxedError; use common_recordbatch::{RecordBatchStreamWrapper, SendableRecordBatchStream}; use common_telemetry::tracing::Span; +use datafusion::physical_plan::ExecutionPlan; use datatypes::schema::SchemaRef; use futures_util::StreamExt; use snafu::ResultExt; @@ -107,6 +108,10 @@ pub trait SystemTable { fn to_stream(&self, request: ScanRequest) -> Result; + fn scan_plan(&self, _request: ScanRequest) -> Result>> { + Ok(None) + } + fn table_type(&self) -> TableType { TableType::Temporary } @@ -173,4 +178,15 @@ impl DataSource for SystemTableDataSource { Ok(Box::pin(stream)) } + + fn get_physical_plan( + &self, + request: ScanRequest, + ) -> std::result::Result>, BoxedError> { + self.table + .scan_plan(request) + .map_err(BoxedError::new) + .context(TablesRecordBatchSnafu) + .map_err(BoxedError::new) + } } diff --git a/src/catalog/src/system_schema/information_schema.rs b/src/catalog/src/system_schema/information_schema.rs index 01e366da886..0a0864aed85 100644 --- a/src/catalog/src/system_schema/information_schema.rs +++ b/src/catalog/src/system_schema/information_schema.rs @@ -46,6 +46,7 @@ use common_procedure::ProcedureInfo; use common_recordbatch::SendableRecordBatchStream; use datafusion::error::DataFusionError; use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_plan::ExecutionPlan; use datatypes::schema::SchemaRef; use lazy_static::lazy_static; use paste::paste; @@ -420,6 +421,10 @@ pub trait InformationTable { fn to_stream(&self, request: ScanRequest) -> Result; + fn scan_plan(&self, _request: ScanRequest) -> Result>> { + Ok(None) + } + fn table_type(&self) -> TableType { TableType::Temporary } @@ -449,6 +454,10 @@ where fn to_stream(&self, request: ScanRequest) -> Result { InformationTable::to_stream(self, request) } + + fn scan_plan(&self, request: ScanRequest) -> Result>> { + InformationTable::scan_plan(self, request) + } } pub type InformationExtensionRef = Arc + Send + Sync>; @@ -475,6 +484,16 @@ pub trait InformationExtension { &self, request: DatanodeInspectRequest, ) -> std::result::Result; + + /// Builds a physical plan for datanode inspect if the extension can expose + /// the distributed fan-in semantics to DataFusion. + fn inspect_datanode_plan( + &self, + _request: DatanodeInspectRequest, + _schema: SchemaRef, + ) -> std::result::Result>, Self::Error> { + Ok(None) + } } /// The request to inspect the datanode. diff --git a/src/catalog/src/system_schema/information_schema/region_info.rs b/src/catalog/src/system_schema/information_schema/region_info.rs index ffc9dfc7ae7..56dc04b85c8 100644 --- a/src/catalog/src/system_schema/information_schema/region_info.rs +++ b/src/catalog/src/system_schema/information_schema/region_info.rs @@ -18,6 +18,7 @@ use common_catalog::consts::INFORMATION_SCHEMA_REGION_INFO_TABLE_ID; use common_error::ext::BoxedError; use common_recordbatch::SendableRecordBatchStream; use common_recordbatch::adapter::AsyncRecordBatchStreamAdapter; +use datafusion::physical_plan::ExecutionPlan; use datatypes::schema::SchemaRef; use snafu::ResultExt; use store_api::region_info::RegionInfoEntry; @@ -83,4 +84,20 @@ impl InformationTable for InformationSchemaRegionInfo { Box::pin(future), ))) } + + fn scan_plan(&self, request: ScanRequest) -> Result>> { + let schema = if let Some(p) = request.projection_indices() { + Arc::new(self.schema.try_project(p).context(ProjectSchemaSnafu)?) + } else { + self.schema.clone() + }; + + let info_ext = utils::information_extension(&self.catalog_manager)?; + let req = DatanodeInspectRequest { + kind: DatanodeInspectKind::RegionInfo, + scan: request, + }; + + info_ext.inspect_datanode_plan(req, schema) + } } diff --git a/src/catalog/src/system_schema/information_schema/ssts.rs b/src/catalog/src/system_schema/information_schema/ssts.rs index 520fc6d485d..19c44aee343 100644 --- a/src/catalog/src/system_schema/information_schema/ssts.rs +++ b/src/catalog/src/system_schema/information_schema/ssts.rs @@ -21,6 +21,7 @@ use common_catalog::consts::{ use common_error::ext::BoxedError; use common_recordbatch::SendableRecordBatchStream; use common_recordbatch::adapter::AsyncRecordBatchStreamAdapter; +use datafusion::physical_plan::ExecutionPlan; use datatypes::schema::SchemaRef; use snafu::ResultExt; use store_api::sst_entry::{ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry}; @@ -34,6 +35,30 @@ use crate::information_schema::{ }; use crate::system_schema::utils; +fn projected_schema(schema: &SchemaRef, request: &ScanRequest) -> Result { + if let Some(p) = request.projection_indices() { + Ok(Arc::new(schema.try_project(p).context(ProjectSchemaSnafu)?)) + } else { + Ok(schema.clone()) + } +} + +fn inspect_plan( + catalog_manager: &Weak, + kind: DatanodeInspectKind, + schema: &SchemaRef, + request: ScanRequest, +) -> Result>> { + let projected_schema = projected_schema(schema, &request)?; + let info_ext = utils::information_extension(catalog_manager)?; + let req = DatanodeInspectRequest { + kind, + scan: request, + }; + + info_ext.inspect_datanode_plan(req, projected_schema) +} + /// Information schema table for sst manifest. pub struct InformationSchemaSstsManifest { schema: SchemaRef, @@ -63,11 +88,7 @@ impl InformationTable for InformationSchemaSstsManifest { } fn to_stream(&self, request: ScanRequest) -> Result { - let schema = if let Some(p) = request.projection_indices() { - Arc::new(self.schema.try_project(p).context(ProjectSchemaSnafu)?) - } else { - self.schema.clone() - }; + let schema = projected_schema(&self.schema, &request)?; let info_ext = utils::information_extension(&self.catalog_manager)?; let req = DatanodeInspectRequest { kind: DatanodeInspectKind::SstManifest, @@ -86,6 +107,15 @@ impl InformationTable for InformationSchemaSstsManifest { Box::pin(future), ))) } + + fn scan_plan(&self, request: ScanRequest) -> Result>> { + inspect_plan( + &self.catalog_manager, + DatanodeInspectKind::SstManifest, + &self.schema, + request, + ) + } } /// Information schema table for sst storage. @@ -117,11 +147,7 @@ impl InformationTable for InformationSchemaSstsStorage { } fn to_stream(&self, request: ScanRequest) -> Result { - let schema = if let Some(p) = request.projection_indices() { - Arc::new(self.schema.try_project(p).context(ProjectSchemaSnafu)?) - } else { - self.schema.clone() - }; + let schema = projected_schema(&self.schema, &request)?; let info_ext = utils::information_extension(&self.catalog_manager)?; let req = DatanodeInspectRequest { @@ -141,6 +167,15 @@ impl InformationTable for InformationSchemaSstsStorage { Box::pin(future), ))) } + + fn scan_plan(&self, request: ScanRequest) -> Result>> { + inspect_plan( + &self.catalog_manager, + DatanodeInspectKind::SstStorage, + &self.schema, + request, + ) + } } /// Information schema table for index metadata. @@ -172,11 +207,7 @@ impl InformationTable for InformationSchemaSstsIndexMeta { } fn to_stream(&self, request: ScanRequest) -> Result { - let schema = if let Some(p) = request.projection_indices() { - Arc::new(self.schema.try_project(p).context(ProjectSchemaSnafu)?) - } else { - self.schema.clone() - }; + let schema = projected_schema(&self.schema, &request)?; let info_ext = utils::information_extension(&self.catalog_manager)?; let req = DatanodeInspectRequest { @@ -196,4 +227,116 @@ impl InformationTable for InformationSchemaSstsIndexMeta { Box::pin(future), ))) } + + fn scan_plan(&self, request: ScanRequest) -> Result>> { + inspect_plan( + &self.catalog_manager, + DatanodeInspectKind::SstIndexMeta, + &self.schema, + request, + ) + } +} + +#[cfg(test)] +mod tests { + use cache::{build_fundamental_cache_registry, with_default_composite_cache_registry}; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, INFORMATION_SCHEMA_NAME}; + use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; + use common_meta::cluster::NodeInfo; + use common_meta::datanode::RegionStat; + use common_meta::key::flow::flow_state::FlowStat; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_procedure::ProcedureInfo; + use common_recordbatch::{RecordBatches, SendableRecordBatchStream}; + use datafusion::physical_plan::empty::EmptyExec; + + use super::*; + use crate::error::Error; + use crate::information_schema::InformationExtension; + use crate::kvbackend::KvBackendCatalogManagerBuilder; + + struct PlanInformationExtension; + + #[async_trait::async_trait] + impl InformationExtension for PlanInformationExtension { + type Error = Error; + + async fn nodes(&self) -> std::result::Result, Self::Error> { + Ok(vec![]) + } + + async fn procedures( + &self, + ) -> std::result::Result, Self::Error> { + Ok(vec![]) + } + + async fn region_stats(&self) -> std::result::Result, Self::Error> { + Ok(vec![]) + } + + async fn flow_stats(&self) -> std::result::Result, Self::Error> { + Ok(None) + } + + async fn inspect_datanode( + &self, + _request: DatanodeInspectRequest, + ) -> std::result::Result { + Ok(RecordBatches::empty().as_stream()) + } + + fn inspect_datanode_plan( + &self, + _request: DatanodeInspectRequest, + schema: SchemaRef, + ) -> std::result::Result>, Self::Error> { + Ok(Some(Arc::new(EmptyExec::new( + schema.arrow_schema().clone(), + )))) + } + } + + #[tokio::test] + async fn ssts_manifest_uses_information_extension_scan_plan() { + let backend = Arc::new(MemoryKvBackend::default()); + let layered_cache_builder = LayeredCacheRegistryBuilder::default() + .add_cache_registry(CacheRegistryBuilder::default().build()); + let fundamental_cache_registry = build_fundamental_cache_registry(backend.clone()); + let layered_cache_registry = Arc::new( + with_default_composite_cache_registry( + layered_cache_builder.add_cache_registry(fundamental_cache_registry), + ) + .unwrap() + .build(), + ); + + let catalog_manager = KvBackendCatalogManagerBuilder::new( + Arc::new(PlanInformationExtension), + backend, + layered_cache_registry, + ) + .build(); + + let table = catalog_manager + .table( + DEFAULT_CATALOG_NAME, + INFORMATION_SCHEMA_NAME, + SSTS_MANIFEST, + None, + ) + .await + .unwrap() + .unwrap(); + + let request = ScanRequest { + projection_input: Some(vec![0].into()), + ..Default::default() + }; + let plan = table.scan_to_plan(request).unwrap().unwrap(); + + assert!(plan.as_any().is::()); + assert_eq!(1, plan.schema().fields().len()); + } } diff --git a/src/store-api/src/data_source.rs b/src/store-api/src/data_source.rs index 3fbef8e08c8..2e047cad6c0 100644 --- a/src/store-api/src/data_source.rs +++ b/src/store-api/src/data_source.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use common_error::ext::BoxedError; use common_recordbatch::SendableRecordBatchStream; +use datafusion_physical_plan::ExecutionPlan; use crate::storage::ScanRequest; @@ -24,6 +25,15 @@ use crate::storage::ScanRequest; pub trait DataSource { /// Retrieves a stream of record batches based on the provided scan request. fn get_stream(&self, request: ScanRequest) -> Result; + + /// Builds a physical plan for the scan request if the data source needs to + /// expose its execution semantics to DataFusion. + fn get_physical_plan( + &self, + _request: ScanRequest, + ) -> Result>, BoxedError> { + Ok(None) + } } pub type DataSourceRef = Arc; diff --git a/src/table/src/table.rs b/src/table/src/table.rs index 0ae7d580d8a..542a5066771 100644 --- a/src/table/src/table.rs +++ b/src/table/src/table.rs @@ -19,6 +19,7 @@ use common_recordbatch::SendableRecordBatchStream; use datafusion::execution::FunctionRegistry; use datafusion::logical_expr::Cast; use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionContext; use datafusion_expr::expr::Expr; use datatypes::data_type::DataType; @@ -142,6 +143,12 @@ impl Table { .context(TablesRecordBatchSnafu) } + pub fn scan_to_plan(&self, request: ScanRequest) -> Result>> { + self.data_source + .get_physical_plan(request) + .context(TablesRecordBatchSnafu) + } + /// Tests whether the table provider can make use of any or all filter expressions /// to optimise data retrieval. pub fn supports_filters_pushdown(&self, filters: &[&Expr]) -> Result> { diff --git a/src/table/src/table/adapter.rs b/src/table/src/table/adapter.rs index 339c82dddeb..8907140c312 100644 --- a/src/table/src/table/adapter.rs +++ b/src/table/src/table/adapter.rs @@ -113,6 +113,11 @@ impl TableProvider for DfTableProviderAdapter { request.limit = limit; request.clone() }; + + if let Some(plan) = self.table.scan_to_plan(request.clone())? { + return Ok(plan); + } + let stream = self.table.scan_to_stream(request).await?; // build sort physical expr diff --git a/tests/cases/distributed/information_schema/ssts_limit.result b/tests/cases/distributed/information_schema/ssts_limit.result index 98a19842995..0b64a3c0fa0 100644 --- a/tests/cases/distributed/information_schema/ssts_limit.result +++ b/tests/cases/distributed/information_schema/ssts_limit.result @@ -63,6 +63,37 @@ FROM ( | 1 | +-----------------------+ +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +EXPLAIN SELECT COUNT(*) AS filtered_limited_rows +FROM ( + SELECT region_id + FROM information_schema.ssts_manifest + WHERE table_id > 0 + LIMIT 1 +); + ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: count(Int64(1)) AS count(*) AS filtered_limited_rows | +| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] | +| | Limit: skip=0, fetch=1 | +| | Projection: information_schema.ssts_manifest.region_id | +| | Filter: information_schema.ssts_manifest.table_id > UInt32(0) | +| | TableScan: information_schema.ssts_manifest, partial_filters=[information_schema.ssts_manifest.table_id > UInt32(0)] | +| | ]] | +| physical_plan | ProjectionExec: expr=[count(Int64(1))@0 as filtered_limited_rows] | +| | AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] | +| | CoalescePartitionsExec | +| | AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] | +| | CoalescePartitionsExec: fetch=1 | +| | FilterExec: table_id@0 > 0, projection=[], fetch=1 | +| | RepartitionExec: REDACTED +| | DistributedInspectExec: kind=SstManifest, scan=ScanRequest { projection: ProjectionInput { projection: [2], nested_paths: [] }, filters: [table_id > UInt32(0)] }, schema=Schema { fields: [Field { name: "table_id", data_type: UInt32 }], metadata: {"greptime:version": "0"} } | +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + DROP TABLE ssts_limit_case; Affected Rows: 0 diff --git a/tests/cases/distributed/information_schema/ssts_limit.sql b/tests/cases/distributed/information_schema/ssts_limit.sql index c02f31d63d7..d548bd283a5 100644 --- a/tests/cases/distributed/information_schema/ssts_limit.sql +++ b/tests/cases/distributed/information_schema/ssts_limit.sql @@ -35,4 +35,13 @@ FROM ( LIMIT 1 ); +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +EXPLAIN SELECT COUNT(*) AS filtered_limited_rows +FROM ( + SELECT region_id + FROM information_schema.ssts_manifest + WHERE table_id > 0 + LIMIT 1 +); + DROP TABLE ssts_limit_case; From 02283b6ba0493206b1723d749ab3820ec1f19f0b Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 10 Jul 2026 20:27:27 +0800 Subject: [PATCH 06/48] test(perf): add remote write storage inspection (#8444) * feat: add remote write value distributions Signed-off-by: discord9 * test: extend remote write value distributions Signed-off-by: discord9 * test: inspect remote write parquet storage Signed-off-by: discord9 * test: normalize remote write perf fixtures Signed-off-by: discord9 * test: add heavy remote write perf case Signed-off-by: discord9 * test: use head greptime for read bench Signed-off-by: discord9 * test: keep heavy remote write case local Signed-off-by: discord9 * test: tune remote write perf smoke case Signed-off-by: discord9 * test: fix query fixture import style Signed-off-by: discord9 * test: cover remote write value distributions Signed-off-by: discord9 * test: add integer counter perf case Signed-off-by: discord9 --------- Signed-off-by: discord9 --- .github/scripts/query-regression-run.py | 14 +- .github/workflows/query-regression.yml | 16 +- src/cmd/Cargo.toml | 7 +- src/cmd/src/bin/greptime.rs | 1 - src/cmd/src/bin/prom_remote_write_fixture.rs | 144 ---- src/cmd/src/bin/query_perf_fixture.rs | 652 +----------------- src/cmd/src/bin/query_perf_fixture/case.rs | 470 +++++++++++++ .../src/bin/query_perf_fixture/direct_sst.rs | 529 ++++++++++++++ .../bin/query_perf_fixture/inspect_footer.rs | 181 +++++ src/cmd/src/bin/query_perf_fixture/mod.rs | 133 ++++ .../query_perf_fixture/prom_remote_write.rs | 333 +++++++++ src/cmd/src/bin/query_perf_fixture/util.rs | 23 + src/cmd/src/datanode.rs | 9 +- tests/perf/AGENTS.md | 8 +- tests/perf/README.md | 166 ++++- .../case.toml | 48 ++ .../prom_remote_write_mixed_every/case.toml | 49 ++ .../prom_remote_write_run_heavy/case.toml | 50 ++ .../prom_remote_write_seeded_random/case.toml | 48 ++ .../prom_remote_write_smoke/case.toml | 31 - tests/perf/query_regression_runner.py | 377 ++++++++-- ...ery_regression_runner_compaction_toctou.py | 158 +++++ 22 files changed, 2520 insertions(+), 927 deletions(-) delete mode 100644 src/cmd/src/bin/prom_remote_write_fixture.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/case.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/direct_sst.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/inspect_footer.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/mod.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/util.rs create mode 100644 tests/perf/query_cases/prom_remote_write_integer_counter/case.toml create mode 100644 tests/perf/query_cases/prom_remote_write_mixed_every/case.toml create mode 100644 tests/perf/query_cases/prom_remote_write_run_heavy/case.toml create mode 100644 tests/perf/query_cases/prom_remote_write_seeded_random/case.toml delete mode 100644 tests/perf/query_cases/prom_remote_write_smoke/case.toml create mode 100644 tests/perf/test_query_regression_runner_compaction_toctou.py diff --git a/.github/scripts/query-regression-run.py b/.github/scripts/query-regression-run.py index b89857be81f..62d11f44203 100644 --- a/.github/scripts/query-regression-run.py +++ b/.github/scripts/query-regression-run.py @@ -26,6 +26,10 @@ from pathlib import Path DEFAULT_CASES = [ "tests/perf/query_cases/smoke_direct_sst/case.toml", + "tests/perf/query_cases/prom_remote_write_seeded_random/case.toml", + "tests/perf/query_cases/prom_remote_write_run_heavy/case.toml", + "tests/perf/query_cases/prom_remote_write_mixed_every/case.toml", + "tests/perf/query_cases/prom_remote_write_integer_counter/case.toml", "tests/perf/query_cases/promql_pushdown_7913/case.toml", "tests/perf/query_cases/sql_topk_order_by/case.toml", "tests/perf/query_cases/sql_aggregate_order_by/case.toml", @@ -106,8 +110,6 @@ def run_case(args: argparse.Namespace, case_path: Path, work_dir: Path) -> int: str(candidate_bin), "--fixture-generator", str(fixture_generator), - "--remote-write-generator", - str(args.remote_write_generator), "--work-dir", str(work_dir), "--http-timeout", @@ -160,11 +162,6 @@ def main() -> int: parser.add_argument("--work-dir", default=Path("query-regression-work"), type=Path) parser.add_argument("--http-timeout", default=os.environ.get("HTTP_TIMEOUT", "300")) parser.add_argument("--allow-large-fixture", default=os.environ.get("ALLOW_LARGE_FIXTURE", "false")) - parser.add_argument( - "--remote-write-generator", - type=Path, - default=configured_path(os.environ.get("REMOTE_WRITE_GENERATOR")), - ) parser.add_argument( "--summary-script", type=Path, @@ -177,9 +174,6 @@ def main() -> int: parser.add_argument("--candidate-ref", default=os.environ.get("CANDIDATE_REF", "")) parser.add_argument("--github-output", default=os.environ.get("GITHUB_OUTPUT")) args = parser.parse_args() - if args.remote_write_generator is None: - args.remote_write_generator = args.candidate_src / "target" / profile_dir(args.cargo_profile) / "prom_remote_write_fixture" - try: cases = split_cases(args.cases or [os.environ.get("CASE_PATHS", "all")]) except ValueError as err: diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index a898be6ad29..619442e2f95 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -119,10 +119,15 @@ jobs: git reset --hard FETCH_HEAD git clean -ffdx + - name: Test query regression runner lifecycle + working-directory: src + run: uv run --no-project python tests/perf/test_query_regression_runner_compaction_toctou.py + - name: Build candidate greptime and fixture generators working-directory: src run: | - cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime --bin query_perf_fixture --bin prom_remote_write_fixture + cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime + cargo build --profile "${CARGO_PROFILE}" -p cmd --bin query_perf_fixture --features dev-tools target_dir="${CARGO_PROFILE}" if [[ "${CARGO_PROFILE}" == "dev" ]]; then target_dir="debug" @@ -132,8 +137,6 @@ jobs: "${GITHUB_WORKSPACE}/query-regression-bins/candidate/greptime" cp "${CARGO_TARGET_DIR}/${target_dir}/query_perf_fixture" \ "${GITHUB_WORKSPACE}/query-regression-bins/candidate/query_perf_fixture" - cp "${CARGO_TARGET_DIR}/${target_dir}/prom_remote_write_fixture" \ - "${GITHUB_WORKSPACE}/query-regression-bins/candidate/prom_remote_write_fixture" - name: Run query regression id: run @@ -146,7 +149,6 @@ jobs: BASE_BIN: ${{ github.workspace }}/query-regression-bins/base/greptime CANDIDATE_BIN: ${{ github.workspace }}/query-regression-bins/candidate/greptime FIXTURE_GENERATOR: ${{ github.workspace }}/query-regression-bins/candidate/query_perf_fixture - REMOTE_WRITE_GENERATOR: ${{ github.workspace }}/query-regression-bins/candidate/prom_remote_write_fixture run: >- uv run --no-project python src/.github/scripts/query-regression-run.py --base-src src @@ -171,9 +173,13 @@ jobs: with: name: query-regression-report path: | - query-regression-work/** + query-regression-work/**/query-regression-report.json + query-regression-work/**/base/report.json + query-regression-work/**/candidate/report.json + query-regression-work/**/logs/** query-regression-summary.md if-no-files-found: warn + retention-days: 7 - name: Upload trusted comment artifact if: ${{ always() && github.event_name == 'pull_request' }} diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index a8adfa01a48..b06da6e517a 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -12,10 +12,7 @@ path = "src/bin/greptime.rs" [[bin]] name = "query_perf_fixture" path = "src/bin/query_perf_fixture.rs" - -[[bin]] -name = "prom_remote_write_fixture" -path = "src/bin/prom_remote_write_fixture.rs" +required-features = ["dev-tools"] [features] default = [ @@ -25,7 +22,7 @@ default = [ "meta-srv/mysql_kvbackend", ] enterprise = ["common-meta/enterprise", "frontend/enterprise", "meta-srv/enterprise"] -# Developer-only tools (benchmarks, diagnostics) gating the `datanode parquetbench` subcommand. +# Developer-only helper binary gate for `query_perf_fixture`. # Kept out of `default` so normal/release builds don't compile them. dev-tools = [] tokio-console = ["common-telemetry/tokio-console"] diff --git a/src/cmd/src/bin/greptime.rs b/src/cmd/src/bin/greptime.rs index e0f326908a2..957e227355a 100644 --- a/src/cmd/src/bin/greptime.rs +++ b/src/cmd/src/bin/greptime.rs @@ -117,7 +117,6 @@ async fn start(cli: Command) -> Result<()> { } datanode::SubCommand::Objbench(ref bench) => bench.run().await, datanode::SubCommand::Scanbench(ref bench) => bench.run().await, - #[cfg(feature = "dev-tools")] datanode::SubCommand::Parquetbench(ref bench) => bench.run().await, }, SubCommand::Flownode(cmd) => { diff --git a/src/cmd/src/bin/prom_remote_write_fixture.rs b/src/cmd/src/bin/prom_remote_write_fixture.rs deleted file mode 100644 index 0bfdf954823..00000000000 --- a/src/cmd/src/bin/prom_remote_write_fixture.rs +++ /dev/null @@ -1,144 +0,0 @@ -// 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. - -#![allow(clippy::print_stdout)] - -use std::time::{Duration, Instant}; - -use api::prom_store::remote::{Label, Sample, TimeSeries, WriteRequest}; -use clap::Parser; -use prost::Message; -use serde_json::json; -use servers::prom_store::snappy_compress; - -#[derive(Debug, Parser)] -#[command(about = "Generate a deterministic Prometheus remote-write fixture")] -struct Args { - #[arg(long, default_value = "http://127.0.0.1:4000/v1/prometheus/write")] - endpoint: String, - #[arg(long, default_value = "public")] - database: String, - #[arg(long)] - metric: String, - #[arg(long, default_value = "greptime_physical_table")] - physical_table: String, - #[arg(long, default_value_t = 8)] - series_count: u64, - #[arg(long, default_value_t = 30)] - samples_per_series: u64, - #[arg(long, default_value_t = 1_704_067_200_000)] - start_unix_millis: i64, - #[arg(long, default_value_t = 15_000)] - step_millis: i64, - #[arg(long, alias = "batch-size", default_value_t = 8)] - chunk_series_count: u64, - #[arg(long, default_value_t = 60)] - timeout_seconds: u64, -} - -#[tokio::main] -async fn main() -> Result<(), Box> { - let args = Args::parse(); - let started = Instant::now(); - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(args.timeout_seconds)) - .build()?; - let mut batches = 0_u64; - let mut rows = 0_u64; - let mut http_statuses = Vec::new(); - - let chunk = args.chunk_series_count.max(1); - for first in (0..args.series_count).step_by(chunk as usize) { - let last = (first + chunk).min(args.series_count); - let timeseries = (first..last) - .map(|series_idx| make_series(&args, series_idx)) - .collect::>(); - rows += timeseries - .iter() - .map(|ts| ts.samples.len() as u64) - .sum::(); - let body = snappy_compress( - &WriteRequest { - timeseries, - ..Default::default() - } - .encode_to_vec(), - )?; - let response = client - .post(&args.endpoint) - .header("Content-Encoding", "snappy") - .body(body) - .send() - .await?; - let status = response.status(); - http_statuses.push(status.as_u16()); - if !status.is_success() { - let body_text = response.text().await.unwrap_or_default(); - return Err(format!("remote-write failed with status {status}: {body_text}").into()); - } - batches += 1; - } - - println!( - "{}", - json!({ - "status": "ok", - "endpoint": args.endpoint, - "database": args.database, - "metric": args.metric, - "physical_table": args.physical_table, - "series_count": args.series_count, - "samples_per_series": args.samples_per_series, - "rows": rows, - "samples_written": rows, - "batches": batches, - "elapsed_seconds": started.elapsed().as_secs_f64(), - "http_statuses": http_statuses, - }) - ); - Ok(()) -} - -fn make_series(args: &Args, series_idx: u64) -> TimeSeries { - let mut labels = vec![ - label("__name__", &args.metric), - label("x_greptime_database", &args.database), - label("x_greptime_physical_table", &args.physical_table), - label("host", &format!("host{:04}", series_idx % 1024)), - label("instance", &format!("instance{:06}", series_idx)), - ]; - labels.sort_unstable_by(|left, right| left.name.cmp(&right.name)); - let samples = (0..args.samples_per_series) - .map(|sample_idx| Sample { - value: deterministic_value(series_idx, sample_idx), - timestamp: args.start_unix_millis + (sample_idx as i64 * args.step_millis), - }) - .collect(); - TimeSeries { - labels, - samples, - ..Default::default() - } -} - -fn label(name: &str, value: &str) -> Label { - Label { - name: name.to_string(), - value: value.to_string(), - } -} - -fn deterministic_value(series_idx: u64, sample_idx: u64) -> f64 { - ((series_idx % 97) as f64) + (sample_idx as f64 * 0.125) -} diff --git a/src/cmd/src/bin/query_perf_fixture.rs b/src/cmd/src/bin/query_perf_fixture.rs index dfc3a664a03..d0e185038fa 100644 --- a/src/cmd/src/bin/query_perf_fixture.rs +++ b/src/cmd/src/bin/query_perf_fixture.rs @@ -14,656 +14,10 @@ #![allow(clippy::print_stderr, clippy::print_stdout)] -//! Generic direct-readable-SST fixture generator for query performance labs. -//! -//! This is a bounded MVP: it reads one table/one region case TOML, builds -//! synthetic region metadata, writes real Mito readable SST parquet files, and -//! emits a manifest checkpoint plus validation metadata. The first case uses -//! `sst_format = "flat"`; primary-key output is also accepted for early lab use -//! because the writer can convert the same flat input batch to primary-key SSTs. -//! -//! TODO(query-perf): add multi-region generation and seed-manifest/catalog DB -//! integration so the generated fixture can be mounted directly by a full -//! GreptimeDB open/query path without manual metadata wiring. - -use std::collections::HashMap; -use std::fs; -use std::io::{BufWriter, Write}; -use std::num::{NonZeroU64, NonZeroUsize}; -use std::path::{Path, PathBuf}; -use std::sync::Arc; - -use api::v1::{OpType, SemanticType}; -use async_trait::async_trait; -use clap::Parser; -use datatypes::arrow::array::{ - ArrayRef, BinaryDictionaryBuilder, Float64Array, RecordBatch, StringDictionaryBuilder, - TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt64Array, -}; -use datatypes::arrow::datatypes::UInt32Type; -use datatypes::prelude::ConcreteDataType; -use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; -use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; -use mito2::access_layer::{FilePathProvider, Metrics, WriteType}; -use mito2::config::IndexConfig; -use mito2::manifest::action::{RegionCheckpoint, RegionManifest, RemovedFilesRecord}; -use mito2::read::FlatSource; -use mito2::sst::file::{FileMeta, RegionFileId}; -use mito2::sst::index::{Indexer, IndexerBuilder}; -use mito2::sst::parquet::writer::ParquetWriter; -use mito2::sst::parquet::{SstInfo, WriteOptions}; -use mito2::sst::{ - DEFAULT_WRITE_BUFFER_SIZE, FlatSchemaOptions, FormatType, to_flat_sst_arrow_schema, -}; -use object_store::ObjectStore; -use object_store::services::Fs as FsBuilder; -use serde::Deserialize; -use store_api::codec::PrimaryKeyEncoding; -use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataRef}; -use store_api::path_utils::region_name; -use store_api::region_request::PathType; -use store_api::storage::{FileId, RegionId}; - -#[derive(Parser, Debug)] -#[command(name = "query_perf_fixture")] -#[command(about = "Generate direct-readable SST fixtures for query performance cases")] -struct Args { - /// TOML case file, for example tests/perf/query_cases/.../case.toml. - #[arg(long, value_name = "PATH")] - case: PathBuf, - - /// Output directory (creates object-store/, manifest/, files.jsonl, summary.json). - #[arg(long, value_name = "DIR")] - out_dir: PathBuf, - - /// Region ID to synthesize. Defaults to table id 1024, region 0. - #[arg(long, default_value = "4398046511104")] - region_id: u64, - - /// Table directory relative to object-store root. Defaults to data/{database}/{table}/. - #[arg(long, value_name = "DIR")] - table_dir: Option, - - /// Table name to generate when the case contains multiple [[tables]]. - #[arg(long, value_name = "NAME")] - table: Option, - - /// Manifest/checkpoint version. - #[arg(long, default_value = "1000000")] - checkpoint_version: u64, - - /// Safety flag when sst_count exceeds 1000. - #[arg(long)] - allow_large: bool, - - /// Print plan only. - #[arg(long)] - dry_run: bool, -} - -#[derive(Debug, Deserialize)] -struct CaseFile { - scenario: Scenario, -} - -#[derive(Debug, Deserialize)] -#[serde(tag = "kind")] -enum Scenario { - #[serde(rename = "direct_readable_sst")] - DirectReadableSst(DirectReadableSstScenario), -} - -#[derive(Debug, Deserialize)] -struct DirectReadableSstScenario { - #[serde(default)] - seed: Option, - tables: Vec, - layout: LayoutConfig, -} - -impl Scenario { - fn kind(&self) -> &'static str { - match self { - Scenario::DirectReadableSst(_) => "direct_readable_sst", - } - } - - fn direct_readable_sst(&self) -> &DirectReadableSstScenario { - match self { - Scenario::DirectReadableSst(scenario) => scenario, - } - } -} - -#[derive(Debug, Deserialize)] -struct TableConfig { - database: String, - name: String, - engine: String, - #[serde(default)] - append_mode: Option, - #[serde(default)] - sst_format: Option, - primary_key: Vec, - time_index: String, - columns: Vec, -} - -#[derive(Debug, Deserialize)] -struct ColumnConfig { - name: String, - #[serde(rename = "type")] - ty: String, - semantic: String, - distribution: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(tag = "kind")] -enum Distribution { - #[serde(rename = "cardinality")] - Cardinality { - values: NonZeroUsize, - prefix: String, - }, - #[serde(rename = "deterministic_wave")] - DeterministicWave { min: f64, max: f64 }, -} - -#[derive(Debug, Deserialize)] -struct LayoutConfig { - regions: usize, - sst_count: usize, - rows_per_sst: usize, - row_group_size: usize, - series_count: NonZeroUsize, - start_unix_nanos: i64, - step_nanos: i64, - time_range_layout: String, - series_layout: String, -} - -struct NoopIndexBuilder; - -#[async_trait] -impl IndexerBuilder for NoopIndexBuilder { - async fn build(&self, _file_id: FileId, _index_version: u64) -> Indexer { - Indexer::default() - } -} - -#[derive(Clone)] -struct FixedPathProvider { - table_dir: String, -} - -impl FilePathProvider for FixedPathProvider { - fn build_index_file_path(&self, file_id: RegionFileId) -> String { - mito2::sst::location::index_file_path_legacy(&self.table_dir, file_id, PathType::Bare) - } - fn build_index_file_path_with_version( - &self, - index_id: mito2::sst::file::RegionIndexId, - ) -> String { - mito2::sst::location::index_file_path(&self.table_dir, index_id, PathType::Bare) - } - fn build_sst_file_path(&self, file_id: RegionFileId) -> String { - mito2::sst::location::sst_file_path(&self.table_dir, file_id, PathType::Bare) - } -} - -fn semantic_type(s: &str) -> SemanticType { - match s.to_ascii_lowercase().as_str() { - "tag" => SemanticType::Tag, - "field" => SemanticType::Field, - "timestamp" => SemanticType::Timestamp, - other => panic!("unsupported semantic {other}"), - } -} - -fn concrete_type(ty: &str) -> ConcreteDataType { - match ty.to_ascii_uppercase().as_str() { - "STRING" => ConcreteDataType::string_datatype(), - "DOUBLE" => ConcreteDataType::float64_datatype(), - "UINT64" => ConcreteDataType::uint64_datatype(), - "TIMESTAMP(9)" => ConcreteDataType::timestamp_nanosecond_datatype(), - "TIMESTAMP(3)" => ConcreteDataType::timestamp_millisecond_datatype(), - other => panic!("unsupported column type {other}"), - } -} - -fn build_region_metadata(table: &TableConfig, region_id: RegionId) -> RegionMetadata { - let mut builder = store_api::metadata::RegionMetadataBuilder::new(region_id); - for (idx, col) in table.columns.iter().enumerate() { - let mut schema = ColumnSchema::new( - &col.name, - concrete_type(&col.ty), - col.semantic != "timestamp", - ); - if col.semantic == "tag" { - schema = schema.with_inverted_index(true); - } - if col.semantic == "tag" || col.semantic == "field" { - schema = schema - .with_skipping_options(SkippingIndexOptions { - granularity: 1, - ..Default::default() - }) - .expect("valid skipping index options for query perf fixture"); - } - if col.name == table.time_index { - schema = schema.with_time_index(true); - } - builder.push_column_metadata(ColumnMetadata { - column_schema: schema, - semantic_type: semantic_type(&col.semantic), - column_id: idx as u32, - }); - } - let pk_ids = table - .primary_key - .iter() - .map(|name| { - table - .columns - .iter() - .position(|c| c.name == *name) - .unwrap_or_else(|| panic!("primary key column {name} not found")) as u32 - }) - .collect(); - builder.primary_key(pk_ids); - builder.primary_key_encoding(PrimaryKeyEncoding::Dense); - builder - .build() - .expect("region metadata should be valid for fixture table config") -} - -fn encode_dense_primary_key(table: &TableConfig, values: &HashMap) -> Vec { - let fields = table - .primary_key - .iter() - .enumerate() - .map(|(idx, _)| { - ( - idx as u32, - SortField::new(ConcreteDataType::string_datatype()), - ) - }) - .collect(); - let converter = DensePrimaryKeyCodec::with_fields(fields); - converter - .encode( - table - .primary_key - .iter() - .map(|name| datatypes::value::ValueRef::String(values[name].as_str())), - ) - .expect("dense primary key encoding should match fixture primary key fields") -} - -fn tag_value(col: &ColumnConfig, series: usize) -> String { - match col - .distribution - .as_ref() - .expect("tag distribution is required") - { - Distribution::Cardinality { values, prefix } => { - format!("{}{}", prefix, series % values.get()) - } - _ => panic!("tag column {} requires cardinality distribution", col.name), - } -} - -fn series_for_row(layout: &LayoutConfig, sst_idx: usize, base_row: usize, row: usize) -> usize { - match layout.series_layout.as_str() { - "round_robin" | "timestamp_major" => (base_row + row) % layout.series_count.get(), - "per_sst" => sst_idx % layout.series_count.get(), - other => panic!("unsupported series_layout {other}"), - } -} - -fn timestamp_for_row(layout: &LayoutConfig, base_row: usize, row: usize) -> i64 { - let logical_row = base_row + row; - let timestamp_step = if layout.series_layout == "timestamp_major" { - logical_row / layout.series_count.get() - } else { - logical_row - }; - - layout.start_unix_nanos + (timestamp_step as i64 * layout.step_nanos) -} - -fn wave_value(min: f64, max: f64, row: usize) -> f64 { - let span = max - min; - let phase = (row % 1024) as f64 / 1023.0; - min + span * (0.5 - 0.5 * (std::f64::consts::TAU * phase).cos()) -} - -fn generate_record_batch( - table: &TableConfig, - metadata: &RegionMetadataRef, - layout: &LayoutConfig, - sst_idx: usize, - sequence: u64, -) -> RecordBatch { - let flat_schema = to_flat_sst_arrow_schema(metadata, &FlatSchemaOptions::default()); - let rows = layout.rows_per_sst; - let mut columns = Vec::new(); - let base_row = sst_idx * rows; - - for col in &table.columns { - match (col.semantic.as_str(), col.ty.to_ascii_uppercase().as_str()) { - ("tag", "STRING") => { - let mut b = StringDictionaryBuilder::::new(); - for row in 0..rows { - let series = series_for_row(layout, sst_idx, base_row, row); - let value = tag_value(col, series); - b.append_value(value); - } - columns.push(Arc::new(b.finish()) as ArrayRef); - } - ("field", "DOUBLE") => { - let (min, max) = match col.distribution.as_ref() { - Some(Distribution::DeterministicWave { min, max }) => (*min, *max), - _ => (0.0, 1.0), - }; - columns.push(Arc::new(Float64Array::from( - (0..rows) - .map(|r| wave_value(min, max, base_row + r)) - .collect::>(), - )) as ArrayRef); - } - ("field", "UINT64") => columns.push(Arc::new(UInt64Array::from( - (0..rows).map(|r| (base_row + r) as u64).collect::>(), - )) as ArrayRef), - ("timestamp", "TIMESTAMP(9)") => columns.push(Arc::new(TimestampNanosecondArray::from( - (0..rows) - .map(|r| timestamp_for_row(layout, base_row, r)) - .collect::>(), - )) as ArrayRef), - ("timestamp", "TIMESTAMP(3)") => { - columns.push(Arc::new(TimestampMillisecondArray::from( - (0..rows) - .map(|r| timestamp_for_row(layout, base_row, r) / 1_000_000) - .collect::>(), - )) as ArrayRef) - } - other => panic!("unsupported column combination {other:?}"), - } - } - - let mut pk_builder = BinaryDictionaryBuilder::::new(); - for row in 0..rows { - let series = series_for_row(layout, sst_idx, base_row, row); - let values = table - .columns - .iter() - .filter(|c| c.semantic == "tag") - .map(|c| (c.name.clone(), tag_value(c, series))) - .collect(); - pk_builder - .append(encode_dense_primary_key(table, &values)) - .expect("append generated dense primary key to Arrow dictionary builder"); - } - columns.push(Arc::new(pk_builder.finish())); - columns.push(Arc::new(UInt64Array::from_value(sequence, rows))); - columns.push(Arc::new(UInt8Array::from_value(OpType::Put as u8, rows))); - RecordBatch::try_new(flat_schema, columns) - .expect("generated fixture columns should match flat SST Arrow schema") -} - -fn file_meta_from_sst_info( - info: &SstInfo, - region_id: RegionId, - file_id: FileId, - sequence: u64, -) -> FileMeta { - FileMeta { - region_id, - file_id, - time_range: info.time_range, - level: 0, - file_size: info.file_size, - max_row_group_uncompressed_size: info.max_row_group_uncompressed_size, - available_indexes: Default::default(), - indexes: Default::default(), - index_file_size: 0, - index_version: 0, - num_rows: info.num_rows as u64, - num_row_groups: info.num_row_groups, - sequence: NonZeroU64::new(sequence), - partition_expr: None, - num_series: info.num_series, - primary_key_min: None, - primary_key_max: None, - } -} - -fn deterministic_file_id(seed: u64, index: usize) -> FileId { - let seed_hi = ((seed >> 16) & 0xffff) as u16; - let seed_lo = (seed & 0xffff) as u16; - FileId::parse_str(&format!( - "00000000-0000-{seed_hi:04x}-{seed_lo:04x}-{index:012x}" - )) - .expect("deterministic file id must be a valid UUID") -} - -fn remap_sst_file( - object_store_dir: &Path, - table_dir: &str, - region_id: RegionId, - old_file_id: FileId, - new_file_id: FileId, -) { - if old_file_id == new_file_id { - return; - } - let old_path = object_store_dir.join(mito2::sst::location::sst_file_path( - table_dir, - RegionFileId::new(region_id, old_file_id), - PathType::Bare, - )); - let new_path = object_store_dir.join(mito2::sst::location::sst_file_path( - table_dir, - RegionFileId::new(region_id, new_file_id), - PathType::Bare, - )); - if let Some(parent) = new_path.parent() { - fs::create_dir_all(parent).expect("failed to create remapped SST parent directory"); - } - fs::rename(&old_path, &new_path).unwrap_or_else(|err| { - panic!( - "failed to remap SST file {} -> {}: {err}", - old_path.display(), - new_path.display() - ) - }); -} - -fn case_name_from_path(path: &Path) -> String { - path.parent() - .and_then(Path::file_name) - .and_then(|name| name.to_str()) - .unwrap_or("query_perf_case") - .to_string() -} +#[path = "query_perf_fixture/mod.rs"] +mod query_perf_fixture; #[tokio::main] async fn main() { - let args = Args::parse(); - let case_text = fs::read_to_string(&args.case).expect("failed to read query perf case TOML"); - let case: CaseFile = toml::from_str(&case_text).expect("failed to parse query perf case TOML"); - let case_name = case_name_from_path(&args.case); - let scenario = case.scenario.direct_readable_sst(); - if scenario.tables.is_empty() || scenario.layout.regions != 1 { - panic!("fixture generator supports one or more tables and exactly one region per table"); - } - if scenario.layout.time_range_layout != "non_overlapping_per_sst" { - panic!("MVP supports time_range_layout=non_overlapping_per_sst"); - } - if scenario.layout.series_layout == "timestamp_major" - && !scenario - .layout - .rows_per_sst - .is_multiple_of(scenario.layout.series_count.get()) - { - panic!( - "series_layout=timestamp_major requires rows_per_sst to be divisible by series_count" - ); - } - if scenario.layout.sst_count > 1000 && !args.allow_large { - panic!("sst_count exceeds 1000; pass --allow-large"); - } - let seed = scenario.seed.unwrap_or(0); - let table_index = match args.table.as_deref() { - Some(name) => scenario - .tables - .iter() - .position(|table| table.name == name) - .unwrap_or_else(|| panic!("--table {name} was not found in case {case_name}")), - None if scenario.tables.len() == 1 => 0, - None => panic!( - "case {} contains {} tables; pass --table to choose one", - case_name, - scenario.tables.len() - ), - }; - let table = &scenario.tables[table_index]; - if table.engine != "mito" { - panic!("MVP supports engine=mito"); - } - let region_id = RegionId::from(args.region_id); - let table_dir = args - .table_dir - .unwrap_or_else(|| format!("data/{}/{}/", table.database, table.name)); - if Path::new(&table_dir).is_absolute() { - panic!("--table-dir must be relative"); - } - if table_dir.trim_end_matches('/').rsplit('/').next() - == Some(region_name(region_id.table_id(), region_id.region_sequence()).as_str()) - { - panic!("--table-dir must be table-level directory"); - } - let region_dir = - mito2::sst::location::region_dir_from_table_dir(&table_dir, region_id, PathType::Bare); - let format = match table.sst_format.as_deref().unwrap_or("flat") { - "flat" => FormatType::Flat, - "primary_key" | "primary-key" => FormatType::PrimaryKey, - other => panic!("unsupported sst_format {other}"), - }; - println!( - "query_perf_fixture case={} scenario={} table={}.{} ssts={} rows_per_sst={} format={format:?}", - case_name, - case.scenario.kind(), - table.database, - table.name, - scenario.layout.sst_count, - scenario.layout.rows_per_sst - ); - println!( - "out_dir={} table_dir={} region_dir={}", - args.out_dir.display(), - table_dir, - region_dir - ); - if args.dry_run { - return; - } - - let obj_store_dir = args.out_dir.join("object-store"); - let manifest_dir = args.out_dir.join("manifest"); - fs::create_dir_all(&obj_store_dir).expect("failed to create fixture object-store directory"); - fs::create_dir_all(&manifest_dir).expect("failed to create fixture manifest directory"); - let ostorage = ObjectStore::new(FsBuilder::default().root(&obj_store_dir.to_string_lossy())) - .expect("failed to create filesystem object store for fixture output") - .finish(); - let metadata: RegionMetadataRef = Arc::new(build_region_metadata(table, region_id)); - let mut files = HashMap::with_capacity(scenario.layout.sst_count); - let mut next_file_index = 1; - for i in 0..scenario.layout.sst_count { - let sequence = 1000 + i as u64; - let batch = generate_record_batch(table, &metadata, &scenario.layout, i, sequence); - let source = - FlatSource::new_iter(batch.schema(), Box::new(vec![batch].into_iter().map(Ok))); - let mut metrics = Metrics::new(WriteType::Flush); - let mut writer = ParquetWriter::new_with_object_store( - ostorage.clone(), - metadata.clone(), - IndexConfig::default(), - NoopIndexBuilder, - FixedPathProvider { - table_dir: table_dir.clone(), - }, - &mut metrics, - ) - .await; - let opts = WriteOptions { - write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, - row_group_size: scenario.layout.row_group_size, - max_file_size: None, - }; - let infos = match format { - FormatType::Flat => writer.write_all_flat(source, Some(sequence), &opts).await, - FormatType::PrimaryKey => { - writer - .write_all_flat_as_primary_key(source, Some(sequence), &opts) - .await - } - } - .expect("failed to write fixture SST through Mito parquet writer"); - for info in infos { - let file_id = deterministic_file_id( - seed.wrapping_add((table_index as u64) << 16), - next_file_index, - ); - next_file_index += 1; - remap_sst_file(&obj_store_dir, &table_dir, region_id, info.file_id, file_id); - files.insert( - file_id, - file_meta_from_sst_info(&info, region_id, file_id, sequence), - ); - } - } - let overall_seq = files - .values() - .filter_map(|m| m.sequence) - .map(|s| s.get()) - .max() - .unwrap_or(0); - let manifest = RegionManifest { - metadata, - files, - removed_files: RemovedFilesRecord::default(), - flushed_entry_id: overall_seq, - flushed_sequence: overall_seq, - committed_sequence: Some(overall_seq), - manifest_version: args.checkpoint_version, - truncated_entry_id: None, - compaction_time_window: None, - sst_format: format, - append_mode: table.append_mode, - }; - let checkpoint = RegionCheckpoint { - last_version: args.checkpoint_version, - compacted_actions: manifest.files.len(), - checkpoint: Some(manifest.clone()), - }; - let checkpoint_bytes = checkpoint - .encode() - .expect("failed to encode fixture region checkpoint"); - let checkpoint_path = manifest_dir.join(format!("{:020}.checkpoint", args.checkpoint_version)); - fs::write(&checkpoint_path, &checkpoint_bytes).expect("failed to write checkpoint file"); - fs::write(manifest_dir.join("_last_checkpoint"), serde_json::to_vec_pretty(&serde_json::json!({ "size": checkpoint_bytes.len(), "version": args.checkpoint_version, "checksum": null, "extend_metadata": {} })).expect("failed to serialize _last_checkpoint metadata")).expect("failed to write _last_checkpoint file"); - let files_jsonl_path = args.out_dir.join("files.jsonl"); - let mut jsonl = BufWriter::new( - fs::File::create(&files_jsonl_path).expect("failed to create fixture files.jsonl"), - ); - let mut entries: Vec<_> = manifest.files.iter().collect(); - entries.sort_by_key(|(id, _)| id.to_string()); - for (file_id, meta) in entries { - writeln!(jsonl, "{}", serde_json::to_string(&serde_json::json!({ "file_id": file_id.to_string(), "region_id": meta.region_id.as_u64(), "object_path": mito2::sst::location::sst_file_path(&table_dir, RegionFileId::new(meta.region_id, *file_id), PathType::Bare), "time_range_start": meta.time_range.0.value(), "time_range_end": meta.time_range.1.value(), "num_rows": meta.num_rows, "num_row_groups": meta.num_row_groups, "file_size": meta.file_size, "num_series": meta.num_series, "sequence": meta.sequence.map(|s| s.get()) })).expect("failed to serialize fixture SST metadata JSONL entry")).expect("failed to write fixture SST metadata JSONL entry"); - } - jsonl.flush().expect("failed to flush fixture files.jsonl"); - fs::write(args.out_dir.join("summary.json"), serde_json::to_vec_pretty(&serde_json::json!({ "case": case_name, "seed": seed, "table_index": table_index, "table": table.name, "database": table.database, "region_id": region_id.as_u64(), "table_dir": table_dir, "region_dir": region_dir, "sst_format": format!("{format:?}"), "sst_count": scenario.layout.sst_count, "rows_per_sst": scenario.layout.rows_per_sst, "row_group_size": scenario.layout.row_group_size, "total_rows": scenario.layout.sst_count * scenario.layout.rows_per_sst, "checkpoint_path": checkpoint_path, "files_jsonl_path": files_jsonl_path, "readback_validated": false, "metadata_source": "synthetic" })).expect("failed to serialize fixture summary")).expect("failed to write fixture summary.json"); - println!("Done. wrote {} SST file entries", manifest.files.len()); + query_perf_fixture::run().await; } diff --git a/src/cmd/src/bin/query_perf_fixture/case.rs b/src/cmd/src/bin/query_perf_fixture/case.rs new file mode 100644 index 00000000000..7dce735d7b9 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/case.rs @@ -0,0 +1,470 @@ +// 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 std::collections::HashMap; +use std::num::NonZeroUsize; + +use clap::ValueEnum; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct CaseFile { + pub(super) scenario: Scenario, +} + +#[derive(Debug, Deserialize, Serialize)] +#[serde(tag = "kind")] +pub(super) enum Scenario { + #[serde(rename = "direct_readable_sst")] + DirectReadableSst(DirectReadableSstScenario), + #[serde(rename = "prom_remote_write_then_query")] + PromRemoteWriteThenQuery(PromRemoteWriteThenQueryScenario), +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct DirectReadableSstScenario { + #[serde(default)] + pub(super) seed: Option, + pub(super) tables: Vec, + pub(super) layout: LayoutConfig, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct PromRemoteWriteThenQueryScenario { + #[serde(default)] + pub(super) queries: Vec, + pub(super) remote_write: PromRemoteWritePlan, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct PromRemoteWritePlan { + #[serde(default = "default_database")] + pub(super) database: String, + #[serde(alias = "metric_name")] + pub(super) metric: String, + #[serde(default = "default_physical_table")] + pub(super) physical_table: String, + #[serde(default = "default_series_count")] + pub(super) series_count: u64, + #[serde(default = "default_samples_per_series")] + pub(super) samples_per_series: u64, + #[serde(default = "default_start_unix_millis")] + pub(super) start_unix_millis: i64, + #[serde(default = "default_step_millis")] + pub(super) step_millis: i64, + #[serde(default = "default_chunk_series_count", alias = "batch_size")] + pub(super) chunk_series_count: u64, + #[serde(default = "default_timeout_seconds")] + pub(super) timeout_seconds: u64, + #[serde(default)] + pub(super) sample_chunk_size: Option, + #[serde(default = "default_flush_every_sample_chunks")] + pub(super) flush_every_sample_chunks: u64, + #[serde(default = "default_visibility_timeout_seconds")] + pub(super) visibility_timeout_seconds: u64, + #[serde(default)] + pub(super) prom_store: PromStoreConfig, + #[serde(default)] + pub(super) value: ValueConfig, + #[serde(default)] + pub(super) storage: Option, + #[serde(default)] + pub(super) read_bench: Option, +} + +pub(super) fn default_database() -> String { + "public".to_string() +} +pub(super) fn default_physical_table() -> String { + "greptime_physical_table".to_string() +} +pub(super) fn default_series_count() -> u64 { + 8 +} +pub(super) fn default_samples_per_series() -> u64 { + 30 +} +pub(super) fn default_start_unix_millis() -> i64 { + 1_704_067_200_000 +} +pub(super) fn default_step_millis() -> i64 { + 15_000 +} +pub(super) fn default_chunk_series_count() -> u64 { + 8 +} +pub(super) fn default_timeout_seconds() -> u64 { + 60 +} +pub(super) fn default_flush_every_sample_chunks() -> u64 { + 1 +} +pub(super) fn default_visibility_timeout_seconds() -> u64 { + 30 +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct PromStoreConfig { + #[serde(default = "default_pending_rows_flush_interval")] + pub(super) pending_rows_flush_interval: String, + #[serde(default = "default_max_batch_rows")] + pub(super) max_batch_rows: u64, + #[serde(default = "default_max_concurrent_flushes")] + pub(super) max_concurrent_flushes: u64, + #[serde(default = "default_worker_channel_capacity")] + pub(super) worker_channel_capacity: u64, + #[serde(default = "default_max_inflight_requests")] + pub(super) max_inflight_requests: u64, +} + +impl Default for PromStoreConfig { + fn default() -> Self { + Self { + pending_rows_flush_interval: default_pending_rows_flush_interval(), + max_batch_rows: default_max_batch_rows(), + max_concurrent_flushes: default_max_concurrent_flushes(), + worker_channel_capacity: default_worker_channel_capacity(), + max_inflight_requests: default_max_inflight_requests(), + } + } +} + +pub(super) fn default_pending_rows_flush_interval() -> String { + "1s".to_string() +} +pub(super) fn default_max_batch_rows() -> u64 { + 100000 +} +pub(super) fn default_max_concurrent_flushes() -> u64 { + 256 +} +pub(super) fn default_worker_channel_capacity() -> u64 { + 65526 +} +pub(super) fn default_max_inflight_requests() -> u64 { + 3000 +} + +#[derive(Debug, Clone, Copy, Deserialize, Serialize, ValueEnum, Eq, PartialEq)] +#[serde(rename_all = "snake_case")] +#[value(rename_all = "snake_case")] +pub(super) enum ValuePattern { + Linear, + Constant, + Modulo, + Unique, + SeededRandom, + RunLength, + QuantizedSignal, + SignalWithSporadicStalls, + MixedSignalRepeated, +} + +impl std::fmt::Display for ValuePattern { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}", + serde_json::to_value(self).unwrap().as_str().unwrap() + ) + } +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct ValueConfig { + #[serde(default = "default_value_pattern")] + pub(super) pattern: ValuePattern, + #[serde(default)] + pub(super) base: f64, + #[serde(default = "default_value_step")] + pub(super) step: f64, + #[serde(default = "default_value_cardinality")] + pub(super) cardinality: u64, + #[serde(default)] + pub(super) seed: u64, + #[serde(default = "default_value_run_length")] + pub(super) run_length: u64, + #[serde(default = "default_value_stall_every")] + pub(super) stall_every: u64, + #[serde(default = "default_value_stall_length")] + pub(super) stall_length: u64, + #[serde(default = "default_value_mixed_every")] + pub(super) mixed_every: u64, +} +impl Default for ValueConfig { + fn default() -> Self { + Self { + pattern: default_value_pattern(), + base: 0.0, + step: default_value_step(), + cardinality: default_value_cardinality(), + seed: 0, + run_length: default_value_run_length(), + stall_every: default_value_stall_every(), + stall_length: default_value_stall_length(), + mixed_every: default_value_mixed_every(), + } + } +} +pub(super) fn default_value_pattern() -> ValuePattern { + ValuePattern::Linear +} +pub(super) fn default_value_step() -> f64 { + 0.125 +} +pub(super) fn default_value_cardinality() -> u64 { + 97 +} +pub(super) fn default_value_run_length() -> u64 { + 8 +} +pub(super) fn default_value_stall_every() -> u64 { + 100 +} +pub(super) fn default_value_stall_length() -> u64 { + 16 +} +pub(super) fn default_value_mixed_every() -> u64 { + 5 +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct StorageConfig { + #[serde(default = "default_true")] + pub(super) inspect: bool, + #[serde(default = "default_storage_column")] + pub(super) column: String, + #[serde(default)] + pub(super) root_suffix: Option, + #[serde(default)] + pub(super) include_metadata_files: bool, + #[serde(default = "default_min_files")] + pub(super) min_files: u64, + #[serde(default = "default_min_files")] + pub(super) min_files_with_column: u64, + #[serde(default)] + pub(super) require_encodings: Vec, + #[serde(default)] + pub(super) forbid_encodings: Vec, + #[serde(default)] + pub(super) max_total_file_size_bytes: Option, + #[serde(default)] + pub(super) max_column_compressed_size_bytes: Option, + #[serde(default)] + pub(super) max_column_uncompressed_size_bytes: Option, + #[serde(default)] + pub(super) max_candidate_total_file_size_regression_pct: Option, + #[serde(default)] + pub(super) max_candidate_column_compressed_size_regression_pct: Option, + #[serde(default)] + pub(super) max_candidate_column_uncompressed_size_regression_pct: Option, + #[serde(skip_deserializing, default)] + pub(super) planned_thresholds: Vec, +} + +#[derive(Debug, Serialize)] +pub(super) struct StorageThresholdPlan { + pub(super) threshold: String, + pub(super) status: String, + pub(super) value: serde_json::Value, +} + +impl StorageConfig { + pub(super) fn populate_planned_thresholds(&mut self) { + let mut planned = Vec::new(); + planned.push(StorageThresholdPlan::new("min_files", self.min_files)); + planned.push(StorageThresholdPlan::new( + "min_files_with_column", + self.min_files_with_column, + )); + if !self.require_encodings.is_empty() { + planned.push(StorageThresholdPlan::new( + "require_encodings", + &self.require_encodings, + )); + } + if !self.forbid_encodings.is_empty() { + planned.push(StorageThresholdPlan::new( + "forbid_encodings", + &self.forbid_encodings, + )); + } + macro_rules! push_optional { + ($name:literal, $value:expr) => { + if let Some(value) = $value { + planned.push(StorageThresholdPlan::new($name, value)); + } + }; + } + push_optional!("max_total_file_size_bytes", self.max_total_file_size_bytes); + push_optional!( + "max_column_compressed_size_bytes", + self.max_column_compressed_size_bytes + ); + push_optional!( + "max_column_uncompressed_size_bytes", + self.max_column_uncompressed_size_bytes + ); + push_optional!( + "max_candidate_total_file_size_regression_pct", + self.max_candidate_total_file_size_regression_pct + ); + push_optional!( + "max_candidate_column_compressed_size_regression_pct", + self.max_candidate_column_compressed_size_regression_pct + ); + push_optional!( + "max_candidate_column_uncompressed_size_regression_pct", + self.max_candidate_column_uncompressed_size_regression_pct + ); + self.planned_thresholds = planned; + } +} + +impl StorageThresholdPlan { + fn new(threshold: &str, value: T) -> Self { + Self { + threshold: threshold.to_string(), + status: "planned".to_string(), + value: serde_json::to_value(value).expect("storage threshold value must serialize"), + } + } +} +pub(super) fn default_true() -> bool { + true +} +pub(super) fn default_storage_column() -> String { + "greptime_value".to_string() +} +pub(super) fn default_min_files() -> u64 { + 1 +} + +#[derive(Debug, Deserialize, Serialize, Clone)] +pub(super) struct ReadBenchConfig { + #[serde(default = "default_true")] + pub(super) enabled: bool, + #[serde(default = "default_true")] + pub(super) parquetbench: bool, + #[serde(default = "default_true")] + pub(super) scanbench: bool, + #[serde(default = "default_iterations")] + pub(super) iterations: u64, + #[serde(default)] + pub(super) projection: Vec, + #[serde(default = "default_parquet_reader")] + pub(super) parquet_reader: String, + #[serde(default = "default_scan_scanner")] + pub(super) scan_scanner: String, + #[serde(default = "default_parallelism")] + pub(super) parallelism: u64, + #[serde(default)] + pub(super) max_files: Option, + #[serde(flatten)] + pub(super) thresholds: HashMap, +} + +impl Default for ReadBenchConfig { + fn default() -> Self { + Self { + enabled: true, + parquetbench: true, + scanbench: true, + iterations: default_iterations(), + projection: vec![], + parquet_reader: default_parquet_reader(), + scan_scanner: default_scan_scanner(), + parallelism: default_parallelism(), + max_files: None, + thresholds: HashMap::new(), + } + } +} + +pub(super) fn default_iterations() -> u64 { + 7 +} +pub(super) fn default_parquet_reader() -> String { + "direct".to_string() +} +pub(super) fn default_scan_scanner() -> String { + "seq".to_string() +} +pub(super) fn default_parallelism() -> u64 { + 1 +} + +impl Scenario { + pub(super) fn kind(&self) -> &'static str { + match self { + Scenario::DirectReadableSst(_) => "direct_readable_sst", + Scenario::PromRemoteWriteThenQuery(_) => "prom_remote_write_then_query", + } + } + + pub(super) fn direct_readable_sst(&self) -> &DirectReadableSstScenario { + match self { + Scenario::DirectReadableSst(scenario) => scenario, + _ => panic!("scenario is not direct_readable_sst"), + } + } +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct TableConfig { + pub(super) database: String, + pub(super) name: String, + pub(super) engine: String, + #[serde(default)] + pub(super) append_mode: Option, + #[serde(default)] + pub(super) sst_format: Option, + pub(super) primary_key: Vec, + pub(super) time_index: String, + pub(super) columns: Vec, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct ColumnConfig { + pub(super) name: String, + #[serde(rename = "type")] + pub(super) ty: String, + pub(super) semantic: String, + pub(super) distribution: Option, +} + +#[derive(Debug, Deserialize, Serialize)] +#[serde(tag = "kind")] +pub(super) enum Distribution { + #[serde(rename = "cardinality")] + Cardinality { + values: NonZeroUsize, + prefix: String, + }, + #[serde(rename = "deterministic_wave")] + DeterministicWave { min: f64, max: f64 }, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct LayoutConfig { + pub(super) regions: usize, + pub(super) sst_count: usize, + pub(super) rows_per_sst: usize, + pub(super) row_group_size: usize, + pub(super) series_count: NonZeroUsize, + pub(super) start_unix_nanos: i64, + pub(super) step_nanos: i64, + pub(super) time_range_layout: String, + pub(super) series_layout: String, +} diff --git a/src/cmd/src/bin/query_perf_fixture/direct_sst.rs b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs new file mode 100644 index 00000000000..0f2e3b34876 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs @@ -0,0 +1,529 @@ +// 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 std::collections::HashMap; +use std::fs::{self}; +use std::io::{BufWriter, Write}; +use std::num::NonZeroU64; +use std::path::Path; +use std::sync::Arc; + +use api::v1::{OpType, SemanticType}; +use async_trait::async_trait; +use datatypes::arrow::array::{ + ArrayRef, BinaryDictionaryBuilder, Float64Array, RecordBatch, StringDictionaryBuilder, + TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt64Array, +}; +use datatypes::arrow::datatypes::UInt32Type; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; +use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; +use mito2::access_layer::{FilePathProvider, Metrics, WriteType}; +use mito2::config::IndexConfig; +use mito2::manifest::action::{RegionCheckpoint, RegionManifest, RemovedFilesRecord}; +use mito2::read::FlatSource; +use mito2::sst::file::{FileMeta, RegionFileId}; +use mito2::sst::index::{Indexer, IndexerBuilder}; +use mito2::sst::parquet::writer::ParquetWriter; +use mito2::sst::parquet::{SstInfo, WriteOptions}; +use mito2::sst::{ + DEFAULT_WRITE_BUFFER_SIZE, FlatSchemaOptions, FormatType, to_flat_sst_arrow_schema, +}; +use object_store::ObjectStore; +use object_store::services::Fs as FsBuilder; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataRef}; +use store_api::path_utils::region_name; +use store_api::region_request::PathType; +use store_api::storage::{FileId, RegionId}; + +use crate::query_perf_fixture::DirectArgs; +use crate::query_perf_fixture::case::*; +use crate::query_perf_fixture::util::case_name_from_path; + +struct NoopIndexBuilder; + +#[async_trait] +impl IndexerBuilder for NoopIndexBuilder { + async fn build(&self, _file_id: FileId, _index_version: u64) -> Indexer { + Indexer::default() + } +} + +#[derive(Clone)] +struct FixedPathProvider { + table_dir: String, +} + +impl FilePathProvider for FixedPathProvider { + fn build_index_file_path(&self, file_id: RegionFileId) -> String { + mito2::sst::location::index_file_path_legacy(&self.table_dir, file_id, PathType::Bare) + } + fn build_index_file_path_with_version( + &self, + index_id: mito2::sst::file::RegionIndexId, + ) -> String { + mito2::sst::location::index_file_path(&self.table_dir, index_id, PathType::Bare) + } + fn build_sst_file_path(&self, file_id: RegionFileId) -> String { + mito2::sst::location::sst_file_path(&self.table_dir, file_id, PathType::Bare) + } +} + +fn semantic_type(s: &str) -> SemanticType { + match s.to_ascii_lowercase().as_str() { + "tag" => SemanticType::Tag, + "field" => SemanticType::Field, + "timestamp" => SemanticType::Timestamp, + other => panic!("unsupported semantic {other}"), + } +} + +fn concrete_type(ty: &str) -> ConcreteDataType { + match ty.to_ascii_uppercase().as_str() { + "STRING" => ConcreteDataType::string_datatype(), + "DOUBLE" => ConcreteDataType::float64_datatype(), + "UINT64" => ConcreteDataType::uint64_datatype(), + "TIMESTAMP(9)" => ConcreteDataType::timestamp_nanosecond_datatype(), + "TIMESTAMP(3)" => ConcreteDataType::timestamp_millisecond_datatype(), + other => panic!("unsupported column type {other}"), + } +} + +fn build_region_metadata(table: &TableConfig, region_id: RegionId) -> RegionMetadata { + let mut builder = store_api::metadata::RegionMetadataBuilder::new(region_id); + for (idx, col) in table.columns.iter().enumerate() { + let mut schema = ColumnSchema::new( + &col.name, + concrete_type(&col.ty), + col.semantic != "timestamp", + ); + if col.semantic == "tag" { + schema = schema.with_inverted_index(true); + } + if col.semantic == "tag" || col.semantic == "field" { + schema = schema + .with_skipping_options(SkippingIndexOptions { + granularity: 1, + ..Default::default() + }) + .expect("valid skipping index options for query perf fixture"); + } + if col.name == table.time_index { + schema = schema.with_time_index(true); + } + builder.push_column_metadata(ColumnMetadata { + column_schema: schema, + semantic_type: semantic_type(&col.semantic), + column_id: idx as u32, + }); + } + let pk_ids = table + .primary_key + .iter() + .map(|name| { + table + .columns + .iter() + .position(|c| c.name == *name) + .unwrap_or_else(|| panic!("primary key column {name} not found")) as u32 + }) + .collect(); + builder.primary_key(pk_ids); + builder.primary_key_encoding(PrimaryKeyEncoding::Dense); + builder + .build() + .expect("region metadata should be valid for fixture table config") +} + +fn encode_dense_primary_key(table: &TableConfig, values: &HashMap) -> Vec { + let fields = table + .primary_key + .iter() + .enumerate() + .map(|(idx, _)| { + ( + idx as u32, + SortField::new(ConcreteDataType::string_datatype()), + ) + }) + .collect(); + let converter = DensePrimaryKeyCodec::with_fields(fields); + converter + .encode( + table + .primary_key + .iter() + .map(|name| datatypes::value::ValueRef::String(values[name].as_str())), + ) + .expect("dense primary key encoding should match fixture primary key fields") +} + +fn tag_value(col: &ColumnConfig, series: usize) -> String { + match col + .distribution + .as_ref() + .expect("tag distribution is required") + { + Distribution::Cardinality { values, prefix } => { + format!("{}{}", prefix, series % values.get()) + } + _ => panic!("tag column {} requires cardinality distribution", col.name), + } +} + +fn series_for_row(layout: &LayoutConfig, sst_idx: usize, base_row: usize, row: usize) -> usize { + match layout.series_layout.as_str() { + "round_robin" | "timestamp_major" => (base_row + row) % layout.series_count.get(), + "per_sst" => sst_idx % layout.series_count.get(), + other => panic!("unsupported series_layout {other}"), + } +} + +fn timestamp_for_row(layout: &LayoutConfig, base_row: usize, row: usize) -> i64 { + let logical_row = base_row + row; + let timestamp_step = if layout.series_layout == "timestamp_major" { + logical_row / layout.series_count.get() + } else { + logical_row + }; + + layout.start_unix_nanos + (timestamp_step as i64 * layout.step_nanos) +} + +fn wave_value(min: f64, max: f64, row: usize) -> f64 { + let span = max - min; + let phase = (row % 1024) as f64 / 1023.0; + min + span * (0.5 - 0.5 * (std::f64::consts::TAU * phase).cos()) +} + +fn generate_record_batch( + table: &TableConfig, + metadata: &RegionMetadataRef, + layout: &LayoutConfig, + sst_idx: usize, + sequence: u64, +) -> RecordBatch { + let flat_schema = to_flat_sst_arrow_schema(metadata, &FlatSchemaOptions::default()); + let rows = layout.rows_per_sst; + let mut columns = Vec::new(); + let base_row = sst_idx * rows; + + for col in &table.columns { + match (col.semantic.as_str(), col.ty.to_ascii_uppercase().as_str()) { + ("tag", "STRING") => { + let mut b = StringDictionaryBuilder::::new(); + for row in 0..rows { + let series = series_for_row(layout, sst_idx, base_row, row); + let value = tag_value(col, series); + b.append_value(value); + } + columns.push(Arc::new(b.finish()) as ArrayRef); + } + ("field", "DOUBLE") => { + let (min, max) = match col.distribution.as_ref() { + Some(Distribution::DeterministicWave { min, max }) => (*min, *max), + _ => (0.0, 1.0), + }; + columns.push(Arc::new(Float64Array::from( + (0..rows) + .map(|r| wave_value(min, max, base_row + r)) + .collect::>(), + )) as ArrayRef); + } + ("field", "UINT64") => columns.push(Arc::new(UInt64Array::from( + (0..rows).map(|r| (base_row + r) as u64).collect::>(), + )) as ArrayRef), + ("timestamp", "TIMESTAMP(9)") => columns.push(Arc::new(TimestampNanosecondArray::from( + (0..rows) + .map(|r| timestamp_for_row(layout, base_row, r)) + .collect::>(), + )) as ArrayRef), + ("timestamp", "TIMESTAMP(3)") => { + columns.push(Arc::new(TimestampMillisecondArray::from( + (0..rows) + .map(|r| timestamp_for_row(layout, base_row, r) / 1_000_000) + .collect::>(), + )) as ArrayRef) + } + other => panic!("unsupported column combination {other:?}"), + } + } + + let mut pk_builder = BinaryDictionaryBuilder::::new(); + for row in 0..rows { + let series = series_for_row(layout, sst_idx, base_row, row); + let values = table + .columns + .iter() + .filter(|c| c.semantic == "tag") + .map(|c| (c.name.clone(), tag_value(c, series))) + .collect(); + pk_builder + .append(encode_dense_primary_key(table, &values)) + .expect("append generated dense primary key to Arrow dictionary builder"); + } + columns.push(Arc::new(pk_builder.finish())); + columns.push(Arc::new(UInt64Array::from_value(sequence, rows))); + columns.push(Arc::new(UInt8Array::from_value(OpType::Put as u8, rows))); + RecordBatch::try_new(flat_schema, columns) + .expect("generated fixture columns should match flat SST Arrow schema") +} + +fn file_meta_from_sst_info( + info: &SstInfo, + region_id: RegionId, + file_id: FileId, + sequence: u64, +) -> FileMeta { + FileMeta { + region_id, + file_id, + time_range: info.time_range, + level: 0, + file_size: info.file_size, + max_row_group_uncompressed_size: info.max_row_group_uncompressed_size, + available_indexes: Default::default(), + indexes: Default::default(), + index_file_size: 0, + index_version: 0, + num_rows: info.num_rows as u64, + num_row_groups: info.num_row_groups, + sequence: NonZeroU64::new(sequence), + partition_expr: None, + num_series: info.num_series, + primary_key_min: None, + primary_key_max: None, + } +} + +fn deterministic_file_id(seed: u64, index: usize) -> FileId { + let seed_hi = ((seed >> 16) & 0xffff) as u16; + let seed_lo = (seed & 0xffff) as u16; + FileId::parse_str(&format!( + "00000000-0000-{seed_hi:04x}-{seed_lo:04x}-{index:012x}" + )) + .expect("deterministic file id must be a valid UUID") +} + +fn remap_sst_file( + object_store_dir: &Path, + table_dir: &str, + region_id: RegionId, + old_file_id: FileId, + new_file_id: FileId, +) { + if old_file_id == new_file_id { + return; + } + let old_path = object_store_dir.join(mito2::sst::location::sst_file_path( + table_dir, + RegionFileId::new(region_id, old_file_id), + PathType::Bare, + )); + let new_path = object_store_dir.join(mito2::sst::location::sst_file_path( + table_dir, + RegionFileId::new(region_id, new_file_id), + PathType::Bare, + )); + if let Some(parent) = new_path.parent() { + fs::create_dir_all(parent).expect("failed to create remapped SST parent directory"); + } + fs::rename(&old_path, &new_path).unwrap_or_else(|err| { + panic!( + "failed to remap SST file {} -> {}: {err}", + old_path.display(), + new_path.display() + ) + }); +} +pub(super) async fn run_direct_sst(args: DirectArgs) { + let case_path = args.case.as_ref().expect("--case is required"); + let out_dir = args.out_dir.as_ref().expect("--out-dir is required"); + let case_text = fs::read_to_string(case_path).expect("failed to read query perf case TOML"); + let case: CaseFile = toml::from_str(&case_text).expect("failed to parse query perf case TOML"); + let case_name = case_name_from_path(case_path); + let scenario = case.scenario.direct_readable_sst(); + if scenario.tables.is_empty() || scenario.layout.regions != 1 { + panic!("fixture generator supports one or more tables and exactly one region per table"); + } + if scenario.layout.time_range_layout != "non_overlapping_per_sst" { + panic!("MVP supports time_range_layout=non_overlapping_per_sst"); + } + if scenario.layout.series_layout == "timestamp_major" + && !scenario + .layout + .rows_per_sst + .is_multiple_of(scenario.layout.series_count.get()) + { + panic!( + "series_layout=timestamp_major requires rows_per_sst to be divisible by series_count" + ); + } + if scenario.layout.sst_count > 1000 && !args.allow_large { + panic!("sst_count exceeds 1000; pass --allow-large"); + } + let seed = scenario.seed.unwrap_or(0); + let table_index = match args.table.as_deref() { + Some(name) => scenario + .tables + .iter() + .position(|table| table.name == name) + .unwrap_or_else(|| panic!("--table {name} was not found in case {case_name}")), + None if scenario.tables.len() == 1 => 0, + None => panic!( + "case {} contains {} tables; pass --table to choose one", + case_name, + scenario.tables.len() + ), + }; + let table = &scenario.tables[table_index]; + if table.engine != "mito" { + panic!("MVP supports engine=mito"); + } + let region_id = RegionId::from(args.region_id); + let table_dir = args + .table_dir + .unwrap_or_else(|| format!("data/{}/{}/", table.database, table.name)); + if Path::new(&table_dir).is_absolute() { + panic!("--table-dir must be relative"); + } + if table_dir.trim_end_matches('/').rsplit('/').next() + == Some(region_name(region_id.table_id(), region_id.region_sequence()).as_str()) + { + panic!("--table-dir must be table-level directory"); + } + let region_dir = + mito2::sst::location::region_dir_from_table_dir(&table_dir, region_id, PathType::Bare); + let format = match table.sst_format.as_deref().unwrap_or("flat") { + "flat" => FormatType::Flat, + "primary_key" | "primary-key" => FormatType::PrimaryKey, + other => panic!("unsupported sst_format {other}"), + }; + println!( + "query_perf_fixture case={} scenario={} table={}.{} ssts={} rows_per_sst={} format={format:?}", + case_name, + case.scenario.kind(), + table.database, + table.name, + scenario.layout.sst_count, + scenario.layout.rows_per_sst + ); + println!( + "out_dir={} table_dir={} region_dir={}", + out_dir.display(), + table_dir, + region_dir + ); + if args.dry_run { + return; + } + + let obj_store_dir = out_dir.join("object-store"); + let manifest_dir = out_dir.join("manifest"); + fs::create_dir_all(&obj_store_dir).expect("failed to create fixture object-store directory"); + fs::create_dir_all(&manifest_dir).expect("failed to create fixture manifest directory"); + let ostorage = ObjectStore::new(FsBuilder::default().root(&obj_store_dir.to_string_lossy())) + .expect("failed to create filesystem object store for fixture output") + .finish(); + let metadata: RegionMetadataRef = Arc::new(build_region_metadata(table, region_id)); + let mut files = HashMap::with_capacity(scenario.layout.sst_count); + let mut next_file_index = 1; + for i in 0..scenario.layout.sst_count { + let sequence = 1000 + i as u64; + let batch = generate_record_batch(table, &metadata, &scenario.layout, i, sequence); + let source = + FlatSource::new_iter(batch.schema(), Box::new(vec![batch].into_iter().map(Ok))); + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + ostorage.clone(), + metadata.clone(), + IndexConfig::default(), + NoopIndexBuilder, + FixedPathProvider { + table_dir: table_dir.clone(), + }, + &mut metrics, + ) + .await; + let opts = WriteOptions { + write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, + row_group_size: scenario.layout.row_group_size, + max_file_size: None, + }; + let infos = match format { + FormatType::Flat => writer.write_all_flat(source, Some(sequence), &opts).await, + FormatType::PrimaryKey => { + writer + .write_all_flat_as_primary_key(source, Some(sequence), &opts) + .await + } + } + .expect("failed to write fixture SST through Mito parquet writer"); + for info in infos { + let file_id = deterministic_file_id( + seed.wrapping_add((table_index as u64) << 16), + next_file_index, + ); + next_file_index += 1; + remap_sst_file(&obj_store_dir, &table_dir, region_id, info.file_id, file_id); + files.insert( + file_id, + file_meta_from_sst_info(&info, region_id, file_id, sequence), + ); + } + } + let overall_seq = files + .values() + .filter_map(|m| m.sequence) + .map(|s| s.get()) + .max() + .unwrap_or(0); + let manifest = RegionManifest { + metadata, + files, + removed_files: RemovedFilesRecord::default(), + flushed_entry_id: overall_seq, + flushed_sequence: overall_seq, + committed_sequence: Some(overall_seq), + manifest_version: args.checkpoint_version, + truncated_entry_id: None, + compaction_time_window: None, + sst_format: format, + append_mode: table.append_mode, + }; + let checkpoint = RegionCheckpoint { + last_version: args.checkpoint_version, + compacted_actions: manifest.files.len(), + checkpoint: Some(manifest.clone()), + }; + let checkpoint_bytes = checkpoint + .encode() + .expect("failed to encode fixture region checkpoint"); + let checkpoint_path = manifest_dir.join(format!("{:020}.checkpoint", args.checkpoint_version)); + fs::write(&checkpoint_path, &checkpoint_bytes).expect("failed to write checkpoint file"); + fs::write(manifest_dir.join("_last_checkpoint"), serde_json::to_vec_pretty(&serde_json::json!({ "size": checkpoint_bytes.len(), "version": args.checkpoint_version, "checksum": null, "extend_metadata": {} })).expect("failed to serialize _last_checkpoint metadata")).expect("failed to write _last_checkpoint file"); + let files_jsonl_path = out_dir.join("files.jsonl"); + let mut jsonl = BufWriter::new( + fs::File::create(&files_jsonl_path).expect("failed to create fixture files.jsonl"), + ); + let mut entries: Vec<_> = manifest.files.iter().collect(); + entries.sort_by_key(|(id, _)| id.to_string()); + for (file_id, meta) in entries { + writeln!(jsonl, "{}", serde_json::to_string(&serde_json::json!({ "file_id": file_id.to_string(), "region_id": meta.region_id.as_u64(), "object_path": mito2::sst::location::sst_file_path(&table_dir, RegionFileId::new(meta.region_id, *file_id), PathType::Bare), "time_range_start": meta.time_range.0.value(), "time_range_end": meta.time_range.1.value(), "num_rows": meta.num_rows, "num_row_groups": meta.num_row_groups, "file_size": meta.file_size, "num_series": meta.num_series, "sequence": meta.sequence.map(|s| s.get()) })).expect("failed to serialize fixture SST metadata JSONL entry")).expect("failed to write fixture SST metadata JSONL entry"); + } + jsonl.flush().expect("failed to flush fixture files.jsonl"); + fs::write(out_dir.join("summary.json"), serde_json::to_vec_pretty(&serde_json::json!({ "case": case_name, "seed": seed, "table_index": table_index, "table": table.name, "database": table.database, "region_id": region_id.as_u64(), "table_dir": table_dir, "region_dir": region_dir, "sst_format": format!("{format:?}"), "sst_count": scenario.layout.sst_count, "rows_per_sst": scenario.layout.rows_per_sst, "row_group_size": scenario.layout.row_group_size, "total_rows": scenario.layout.sst_count * scenario.layout.rows_per_sst, "checkpoint_path": checkpoint_path, "files_jsonl_path": files_jsonl_path, "readback_validated": false, "metadata_source": "synthetic" })).expect("failed to serialize fixture summary")).expect("failed to write fixture summary.json"); + println!("Done. wrote {} SST file entries", manifest.files.len()); +} diff --git a/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs b/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs new file mode 100644 index 00000000000..a52005c4e30 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs @@ -0,0 +1,181 @@ +// 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 std::collections::BTreeSet; +use std::fs::{self, File}; +use std::path::{Path, PathBuf}; + +use clap::Args as ClapArgs; +use parquet::file::reader::{FileReader, SerializedFileReader}; +use serde::Serialize; + +#[derive(Debug, ClapArgs)] +pub(super) struct InspectFooterArgs { + #[arg(long)] + root: PathBuf, + #[arg(long, default_value = "greptime_value")] + column: String, + #[arg(long)] + include_metadata_files: bool, +} + +#[derive(Debug, Serialize)] +struct FooterReport { + root: String, + summary: FooterSummary, + files: Vec, +} +#[derive(Debug, Default, Serialize)] +struct FooterSummary { + column: String, + file_count: usize, + files_with_column: usize, + total_file_size: u64, + total_rows: i64, + column_compressed_size: i64, + column_uncompressed_size: i64, + unique_encodings: Vec, +} +#[derive(Debug, Serialize)] +struct FooterFileReport { + path: String, + relative_path: String, + file_size: u64, + num_rows: i64, + num_row_groups: usize, + columns: Vec, +} +#[derive(Debug, Serialize)] +struct FooterColumnChunkReport { + row_group_index: usize, + column_index: usize, + column_path: String, + encodings: Vec, + compression: String, + compressed_size: i64, + uncompressed_size: i64, + num_values: i64, +} + +pub(super) fn run_inspect_footer( + args: InspectFooterArgs, +) -> Result<(), Box> { + let files = collect_parquet_files(&args.root, args.include_metadata_files)?; + let reports = files + .iter() + .map(|p| inspect_file(&args.root, p, &args.column)) + .collect::, _>>()?; + println!( + "{}", + serde_json::to_string_pretty(&FooterReport { + root: args.root.display().to_string(), + summary: summarize_footer(&args.column, &reports), + files: reports + })? + ); + Ok(()) +} +fn collect_parquet_files( + root: &Path, + include_metadata_files: bool, +) -> Result, Box> { + let mut files = Vec::new(); + collect_parquet_files_inner(root, root, include_metadata_files, &mut files)?; + files.sort(); + Ok(files) +} +fn collect_parquet_files_inner( + root: &Path, + path: &Path, + include_metadata_files: bool, + files: &mut Vec, +) -> Result<(), Box> { + if path.is_dir() { + for entry in fs::read_dir(path)? { + collect_parquet_files_inner(root, &entry?.path(), include_metadata_files, files)?; + } + } else if path.extension().is_some_and(|ext| ext == "parquet") + && (include_metadata_files + || !path + .strip_prefix(root) + .unwrap_or(path) + .components() + .any(|c| c.as_os_str() == "metadata")) + { + files.push(path.to_path_buf()); + } + Ok(()) +} +fn inspect_file( + root: &Path, + path: &Path, + column: &str, +) -> Result> { + let file_size = fs::metadata(path)?.len(); + let reader = SerializedFileReader::new(File::open(path)?)?; + let metadata = reader.metadata().file_metadata(); + let row_groups = reader.metadata().row_groups(); + let mut columns = Vec::new(); + for (row_group_index, rg) in row_groups.iter().enumerate() { + for (column_index, chunk) in rg.columns().iter().enumerate() { + let column_path = chunk.column_path().string(); + if column_path == column { + columns.push(FooterColumnChunkReport { + row_group_index, + column_index, + column_path, + encodings: chunk.encodings().map(|e| format!("{e:?}")).collect(), + compression: format!("{:?}", chunk.compression()), + compressed_size: chunk.compressed_size(), + uncompressed_size: chunk.uncompressed_size(), + num_values: chunk.num_values(), + }); + } + } + } + Ok(FooterFileReport { + path: path.display().to_string(), + relative_path: path + .strip_prefix(root) + .unwrap_or(path) + .display() + .to_string(), + file_size, + num_rows: metadata.num_rows(), + num_row_groups: row_groups.len(), + columns, + }) +} +fn summarize_footer(column: &str, files: &[FooterFileReport]) -> FooterSummary { + let mut encodings = BTreeSet::new(); + let mut s = FooterSummary { + column: column.to_string(), + file_count: files.len(), + ..Default::default() + }; + for file in files { + s.total_file_size += file.file_size; + s.total_rows += file.num_rows; + if !file.columns.is_empty() { + s.files_with_column += 1; + } + for c in &file.columns { + s.column_compressed_size += c.compressed_size; + s.column_uncompressed_size += c.uncompressed_size; + encodings.extend(c.encodings.iter().cloned()); + } + } + s.unique_encodings = encodings.into_iter().collect(); + s +} diff --git a/src/cmd/src/bin/query_perf_fixture/mod.rs b/src/cmd/src/bin/query_perf_fixture/mod.rs new file mode 100644 index 00000000000..00cef327c78 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/mod.rs @@ -0,0 +1,133 @@ +// 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. + +mod case; +mod direct_sst; +mod inspect_footer; +mod prom_remote_write; +mod util; + +use std::fs; +use std::path::PathBuf; + +use case::*; +use clap::{Args as ClapArgs, Parser, Subcommand}; +use direct_sst::run_direct_sst; +use inspect_footer::{InspectFooterArgs, run_inspect_footer}; +use prom_remote_write::{PromRemoteWriteArgs, run_prom_remote_write}; +use serde_json::json; + +#[derive(Parser, Debug)] +#[command(name = "query_perf_fixture")] +#[command(about = "Query performance fixture helper")] +struct Args { + #[command(subcommand)] + command: Option, + + #[command(flatten)] + legacy: DirectArgs, +} + +#[derive(Subcommand, Debug)] +enum Command { + DirectSst(DirectArgs), + PromRemoteWrite(PromRemoteWriteArgs), + InspectFooter(InspectFooterArgs), + Plan(PlanArgs), +} + +#[derive(ClapArgs, Debug, Clone)] +struct DirectArgs { + /// TOML case file, for example tests/perf/query_cases/.../case.toml. + #[arg(long, value_name = "PATH")] + case: Option, + + /// Output directory (creates object-store/, manifest/, files.jsonl, summary.json). + #[arg(long, value_name = "DIR")] + out_dir: Option, + + /// Region ID to synthesize. Defaults to table id 1024, region 0. + #[arg(long, default_value = "4398046511104")] + region_id: u64, + + /// Table directory relative to object-store root. Defaults to data/{database}/{table}/. + #[arg(long, value_name = "DIR")] + table_dir: Option, + + /// Table name to generate when the case contains multiple [[tables]]. + #[arg(long, value_name = "NAME")] + table: Option, + + /// Manifest/checkpoint version. + #[arg(long, default_value = "1000000")] + checkpoint_version: u64, + + /// Safety flag when sst_count exceeds 1000. + #[arg(long)] + allow_large: bool, + + /// Print plan only. + #[arg(long)] + dry_run: bool, +} + +#[derive(ClapArgs, Debug)] +struct PlanArgs { + #[arg(long, value_name = "PATH")] + case: PathBuf, +} + +pub async fn run() { + let args = Args::parse(); + match args.command { + Some(Command::DirectSst(direct)) => run_direct_sst(direct).await, + Some(Command::PromRemoteWrite(rw)) => run_prom_remote_write(rw) + .await + .expect("prom remote write failed"), + Some(Command::InspectFooter(inspect)) => { + run_inspect_footer(inspect).expect("inspect footer failed") + } + Some(Command::Plan(plan)) => run_plan(plan).expect("plan failed"), + None => run_direct_sst(args.legacy).await, + } +} + +fn run_plan(args: PlanArgs) -> Result<(), Box> { + let case_text = fs::read_to_string(&args.case)?; + let mut case: CaseFile = toml::from_str(&case_text)?; + if let Scenario::PromRemoteWriteThenQuery(s) = &mut case.scenario + && let Some(storage) = &mut s.remote_write.storage + { + storage.populate_planned_thresholds(); + if storage.inspect { + let mut rb = s.remote_write.read_bench.clone().unwrap_or_default(); + if rb.enabled && rb.projection.is_empty() { + rb.projection = vec![storage.column.clone()]; + } + s.remote_write.read_bench = Some(rb); + } else if s + .remote_write + .read_bench + .as_ref() + .is_some_and(|rb| rb.enabled) + { + return Err("scenario.remote_write.read_bench requires scenario.remote_write.storage.inspect = true".into()); + } + } + println!( + "{}", + serde_json::to_string_pretty(&json!({"schema_version": 1, "scenario": case.scenario}))? + ); + Ok(()) +} diff --git a/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs b/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs new file mode 100644 index 00000000000..1f36b0c2d23 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs @@ -0,0 +1,333 @@ +// 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 clap::Args as ClapArgs; +use prost::Message; +use serde_json::json; +use servers::prom_store::snappy_compress; + +use crate::query_perf_fixture::case::ValuePattern; + +#[derive(Debug, ClapArgs)] +pub(super) struct PromRemoteWriteArgs { + #[arg(long, default_value = "http://127.0.0.1:4000/v1/prometheus/write")] + endpoint: String, + #[arg(long, default_value = "public")] + database: String, + #[arg(long)] + metric: String, + #[arg(long, default_value = "greptime_physical_table")] + physical_table: String, + #[arg(long, default_value_t = 8)] + series_count: u64, + #[arg(long, default_value_t = 30)] + samples_per_series: u64, + #[arg(long, default_value_t = 1_704_067_200_000)] + start_unix_millis: i64, + #[arg(long, default_value_t = 15_000)] + step_millis: i64, + #[arg(long, alias = "batch-size", default_value_t = 8)] + chunk_series_count: u64, + #[arg(long, default_value_t = 60)] + timeout_seconds: u64, + #[arg(long, default_value_t = ValuePattern::Linear)] + value_pattern: ValuePattern, + #[arg(long, default_value_t = 0.0)] + value_base: f64, + #[arg(long, default_value_t = 0.125)] + value_step: f64, + #[arg(long, default_value_t = 97)] + value_cardinality: u64, + #[arg(long, default_value_t = 0)] + value_seed: u64, + #[arg(long, default_value_t = 8)] + value_run_length: u64, + #[arg(long, default_value_t = 100)] + value_stall_every: u64, + #[arg(long, default_value_t = 16)] + value_stall_length: u64, + #[arg(long, default_value_t = 5)] + value_mixed_every: u64, + #[arg(long, default_value_t = 0)] + value_sample_offset: u64, + #[arg(long)] + value_total_samples_per_series: Option, +} + +pub(super) async fn run_prom_remote_write( + args: PromRemoteWriteArgs, +) -> Result<(), Box> { + use api::prom_store::remote::WriteRequest; + let started = std::time::Instant::now(); + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(args.timeout_seconds)) + .build()?; + let mut batches = 0_u64; + let mut rows = 0_u64; + let mut http_statuses = Vec::new(); + for first in (0..args.series_count).step_by(args.chunk_series_count.max(1) as usize) { + let last = (first + args.chunk_series_count.max(1)).min(args.series_count); + let timeseries = (first..last) + .map(|series_idx| make_prom_series(&args, series_idx)) + .collect::>(); + rows += timeseries + .iter() + .map(|ts| ts.samples.len() as u64) + .sum::(); + let body = snappy_compress( + &WriteRequest { + timeseries, + ..Default::default() + } + .encode_to_vec(), + )?; + let response = client + .post(&args.endpoint) + .header("Content-Encoding", "snappy") + .body(body) + .send() + .await?; + let status = response.status(); + http_statuses.push(status.as_u16()); + if !status.is_success() { + return Err(format!( + "remote-write failed with status {status}: {}", + response.text().await.unwrap_or_default() + ) + .into()); + } + batches += 1; + } + println!( + "{}", + json!({"status":"ok","endpoint":args.endpoint,"database":args.database,"metric":args.metric,"physical_table":args.physical_table,"series_count":args.series_count,"samples_per_series":args.samples_per_series,"rows":rows,"samples_written":rows,"batches":batches,"elapsed_seconds":started.elapsed().as_secs_f64(),"http_statuses":http_statuses,"value":{"pattern":args.value_pattern.to_string(),"base":args.value_base,"step":args.value_step,"cardinality":args.value_cardinality,"seed":args.value_seed,"run_length":args.value_run_length,"stall_every":args.value_stall_every,"stall_length":args.value_stall_length,"mixed_every":args.value_mixed_every,"sample_offset":args.value_sample_offset,"total_samples_per_series":args.value_total_samples_per_series.unwrap_or(args.samples_per_series)}}) + ); + Ok(()) +} + +fn make_prom_series( + args: &PromRemoteWriteArgs, + series_idx: u64, +) -> api::prom_store::remote::TimeSeries { + let mut labels = vec![ + prom_label("__name__", &args.metric), + prom_label("x_greptime_database", &args.database), + prom_label("x_greptime_physical_table", &args.physical_table), + prom_label("host", &format!("host{:04}", series_idx % 1024)), + prom_label("instance", &format!("instance{:06}", series_idx)), + ]; + labels.sort_unstable_by(|l, r| l.name.cmp(&r.name)); + let samples = (0..args.samples_per_series) + .map(|sample_idx| api::prom_store::remote::Sample { + value: deterministic_prom_value(args, series_idx, sample_idx), + timestamp: args.start_unix_millis + (sample_idx as i64 * args.step_millis), + }) + .collect(); + api::prom_store::remote::TimeSeries { + labels, + samples, + ..Default::default() + } +} +fn prom_label(name: &str, value: &str) -> api::prom_store::remote::Label { + api::prom_store::remote::Label { + name: name.to_string(), + value: value.to_string(), + } +} +fn deterministic_prom_value(args: &PromRemoteWriteArgs, series_idx: u64, sample_idx: u64) -> f64 { + let ordinal = series_idx + * args + .value_total_samples_per_series + .unwrap_or(args.samples_per_series) + + args.value_sample_offset + + sample_idx; + let local = args.value_sample_offset + sample_idx; + match args.value_pattern { + ValuePattern::Linear => { + args.value_base + (series_idx % 97) as f64 + local as f64 * args.value_step + } + ValuePattern::Constant => args.value_base, + ValuePattern::Modulo => { + args.value_base + (ordinal % args.value_cardinality.max(1)) as f64 * args.value_step + } + ValuePattern::Unique => args.value_base + ordinal as f64 * args.value_step, + ValuePattern::SeededRandom => { + args.value_base + + (splitmix64(ordinal ^ args.value_seed) % args.value_cardinality.max(1)) as f64 + * args.value_step + } + ValuePattern::RunLength => { + args.value_base + + ((ordinal / args.value_run_length.max(1)) % args.value_cardinality.max(1)) as f64 + * args.value_step + } + ValuePattern::QuantizedSignal => { + args.value_base + + (((series_idx % args.value_cardinality.max(1)) + + (local / args.value_run_length.max(1))) + % args.value_cardinality.max(1)) as f64 + * args.value_step + } + ValuePattern::SignalWithSporadicStalls => { + let every = args.value_stall_every.max(1); + let phase = local % every; + let effective = if phase < args.value_stall_length.min(every) { + local - phase + } else { + local + }; + args.value_base + (series_idx % 97) as f64 + effective as f64 * args.value_step + } + ValuePattern::MixedSignalRepeated => { + if local.is_multiple_of(args.value_mixed_every.max(1)) { + args.value_base + } else { + args.value_base + (series_idx % 97) as f64 + local as f64 * args.value_step + } + } + } +} +fn splitmix64(mut value: u64) -> u64 { + value = value.wrapping_add(0x9e3779b97f4a7c15); + let mut mixed = value; + mixed = (mixed ^ (mixed >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); + mixed = (mixed ^ (mixed >> 27)).wrapping_mul(0x94d049bb133111eb); + mixed ^ (mixed >> 31) +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use super::*; + + fn args(pattern: ValuePattern) -> PromRemoteWriteArgs { + PromRemoteWriteArgs { + endpoint: "http://127.0.0.1:4000/v1/prometheus/write".to_string(), + database: "public".to_string(), + metric: "test_metric".to_string(), + physical_table: "greptime_physical_table".to_string(), + series_count: 2, + samples_per_series: 64, + start_unix_millis: 1_704_067_200_000, + step_millis: 60_000, + chunk_series_count: 2, + timeout_seconds: 60, + value_pattern: pattern, + value_base: 0.0, + value_step: 1.0, + value_cardinality: 1_024, + value_seed: 8_444, + value_run_length: 16, + value_stall_every: 100, + value_stall_length: 16, + value_mixed_every: 5, + value_sample_offset: 0, + value_total_samples_per_series: None, + } + } + + #[test] + fn seeded_random_is_deterministic_in_range_and_highly_distinct() { + let mut args = args(ValuePattern::SeededRandom); + args.samples_per_series = 4_096; + args.value_cardinality = 29_491_200; + + let values = (0..args.samples_per_series) + .map(|sample_idx| deterministic_prom_value(&args, 0, sample_idx)) + .collect::>(); + let repeated = (0..args.samples_per_series) + .map(|sample_idx| deterministic_prom_value(&args, 0, sample_idx)) + .collect::>(); + + assert_eq!(values, repeated); + assert!(values.iter().all(|value| { + (args.value_base..=args.value_base + (args.value_cardinality - 1) as f64) + .contains(value) + })); + assert!( + values + .iter() + .map(|value| value.to_bits()) + .collect::>() + .len() + > 4_000 + ); + } + + #[test] + fn run_length_values_have_exact_sixteen_sample_runs() { + let args = args(ValuePattern::RunLength); + + for sample_idx in 0..64 { + assert_eq!( + deterministic_prom_value(&args, 0, sample_idx), + (sample_idx / 16) as f64 + ); + } + } + + #[test] + fn mixed_every_five_repeats_only_the_expected_positions() { + let mut args = args(ValuePattern::MixedSignalRepeated); + args.value_step = 0.125; + + for sample_idx in 0..20 { + let expected = if sample_idx % 5 == 0 { + 0.0 + } else { + 1.0 + sample_idx as f64 * 0.125 + }; + assert_eq!(deterministic_prom_value(&args, 1, sample_idx), expected); + } + } + + #[test] + fn chunked_values_match_monolithic_values_through_offsets() { + let mut monolithic = args(ValuePattern::SeededRandom); + monolithic.samples_per_series = 14_400; + monolithic.value_total_samples_per_series = Some(14_400); + + for series_idx in 0..2 { + let expected = (0..monolithic.samples_per_series) + .map(|sample_idx| deterministic_prom_value(&monolithic, series_idx, sample_idx)) + .collect::>(); + for chunk_idx in 0..10 { + let mut chunk = args(ValuePattern::SeededRandom); + chunk.samples_per_series = 1_440; + chunk.value_sample_offset = chunk_idx * 1_440; + chunk.value_total_samples_per_series = Some(14_400); + for sample_idx in 0..chunk.samples_per_series { + assert_eq!( + deterministic_prom_value(&chunk, series_idx, sample_idx), + expected[(chunk.value_sample_offset + sample_idx) as usize] + ); + } + } + } + } + + #[test] + fn values_continue_across_series_boundaries() { + let mut args = args(ValuePattern::Unique); + args.samples_per_series = 16; + args.value_total_samples_per_series = Some(16); + args.value_base = 1.0; + args.value_step = 0.5; + + assert_eq!(deterministic_prom_value(&args, 0, 15), 8.5); + assert_eq!(deterministic_prom_value(&args, 1, 0), 9.0); + } +} diff --git a/src/cmd/src/bin/query_perf_fixture/util.rs b/src/cmd/src/bin/query_perf_fixture/util.rs new file mode 100644 index 00000000000..2c86baa6839 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/util.rs @@ -0,0 +1,23 @@ +// 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 std::path::Path; + +pub(super) fn case_name_from_path(path: &Path) -> String { + path.parent() + .and_then(Path::file_name) + .and_then(|name| name.to_str()) + .unwrap_or("query_perf_case") + .to_string() +} diff --git a/src/cmd/src/datanode.rs b/src/cmd/src/datanode.rs index 744624a5084..c33be6c6f33 100644 --- a/src/cmd/src/datanode.rs +++ b/src/cmd/src/datanode.rs @@ -15,11 +15,10 @@ pub mod builder; #[allow(clippy::print_stdout)] pub(crate) mod objbench; -#[cfg(feature = "dev-tools")] #[allow(clippy::print_stdout)] -mod parquetbench; +pub mod parquetbench; #[allow(clippy::print_stdout)] -mod scanbench; +pub mod scanbench; use std::path::Path; use std::time::Duration; @@ -40,7 +39,6 @@ use tracing_appender::non_blocking::WorkerGuard; use crate::App; use crate::datanode::builder::InstanceBuilder; use crate::datanode::objbench::ObjbenchCommand; -#[cfg(feature = "dev-tools")] use crate::datanode::parquetbench::ParquetbenchCommand; use crate::datanode::scanbench::ScanbenchCommand; use crate::error::{ @@ -115,7 +113,6 @@ impl Command { SubCommand::Start(cmd) => cmd.load_options(global_options), // Bench commands are standalone utilities and don't need to load DatanodeOptions. SubCommand::Objbench(_) | SubCommand::Scanbench(_) => Self::default_bench_options(), - #[cfg(feature = "dev-tools")] SubCommand::Parquetbench(_) => Self::default_bench_options(), } } @@ -141,7 +138,6 @@ pub enum SubCommand { /// Scan benchmark tool - benchmarks scanning a region directly from storage Scanbench(ScanbenchCommand), /// Benchmark scanning a single parquet SST. - #[cfg(feature = "dev-tools")] Parquetbench(ParquetbenchCommand), } @@ -160,7 +156,6 @@ impl SubCommand { cmd.run().await?; std::process::exit(0); } - #[cfg(feature = "dev-tools")] SubCommand::Parquetbench(cmd) => { cmd.run().await?; std::process::exit(0); diff --git a/tests/perf/AGENTS.md b/tests/perf/AGENTS.md index eb95544d22d..7f64b283bdc 100644 --- a/tests/perf/AGENTS.md +++ b/tests/perf/AGENTS.md @@ -11,14 +11,14 @@ docs together. - `[case]` is report metadata only. `[scenario]` is the executable regression configuration and must include `kind`, data layout, tables, queries, and - thresholds. Keep scenario parsing explicit in both - `tests/perf/query_regression_runner.py` and - `src/cmd/src/bin/query_perf_fixture.rs`. + thresholds. Rust owns case schema, defaults, validation, and normalized plan + output through `query_perf_fixture plan`; `tests/perf/query_regression_runner.py` + should only orchestrate from that normalized JSON. - Keep the direct-SST generator generic. Issue-specific behavior belongs in case files and thresholds, not in Rust generator logic. - Before pushing perf harness changes, run at least: - `uv run --no-project python -m py_compile .github/scripts/query-regression-run.py .github/scripts/query-regression-summary.py .github/scripts/query-regression-pr-metadata.py tests/perf/query_regression_runner.py` - `cargo fmt --all -- --check` - - `cargo build -p cmd --bin query_perf_fixture` + - `cargo build -p cmd --bin query_perf_fixture --features dev-tools` - dry-run the Python runner and Rust fixture generator against all built-in cases when the DSL or workflow case selection changes. diff --git a/tests/perf/README.md b/tests/perf/README.md index a59081b5681..91bc6c1637e 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -4,6 +4,11 @@ This directory is for query performance cases that compare a base build with a candidate build. It is not a replacement for sqlness: the goal is to measure the effect of optimizer/query-engine changes on realistic scan work. +The Prometheus remote-write cases are end-to-end regression coverage for the +automatic write, flush, SST, and query path. They do not replace controlled +encoding experiments that explicitly compare `plain`/dictionary, +`no_dictionary`, BYTE_STREAM_SPLIT, or Auto policies. + ## Phase 1: direct readable SST fixtures Phase 1 should generate data by writing readable Mito SST files and matching @@ -40,8 +45,10 @@ The runner also supports `scenario.kind = "prom_remote_write_then_query"` for a bounded write-path smoke/regression flow. This path is explicit and separate from the direct-SST fixture path: it starts the base and candidate distributed clusters, writes deterministic Prometheus remote-write v1 samples through -`/v1/prometheus/write`, flushes the configured physical metric table, then runs -the configured SQL/TQL queries against the logical metric table. +`/v1/prometheus/write`, flushes the configured physical metric table, checks +visibility, runs the configured SQL/TQL queries, stops and awaits the datanode, +then optionally inspects the landed SST/Parquet footer, encoding, and size and +runs the read-bench against the quiescent data directory. Remote-write cases configure one database, one logical metric, and one physical table under `[scenario.remote_write]`: @@ -52,34 +59,121 @@ kind = "prom_remote_write_then_query" [scenario.remote_write] database = "public" -metric = "prom_remote_write_smoke" +metric = "prom_remote_write_seeded_random" physical_table = "greptime_physical_table" -series_count = 8 -samples_per_series = 30 +series_count = 512 +samples_per_series = 1440 +sample_chunk_size = 480 +flush_every_sample_chunks = 1 start_unix_millis = 1_704_067_200_000 -step_millis = 15_000 -chunk_series_count = 4 -visibility_timeout_seconds = 30 +step_millis = 60_000 +chunk_series_count = 128 +timeout_seconds = 180 +visibility_timeout_seconds = 120 # Optional: split by time so each helper invocation writes only this many samples # per series, then periodically flush the physical metric table to produce # multiple time-interval SSTs. -# sample_chunk_size = 300 -# flush_every_sample_chunks = 1 [scenario.remote_write.prom_store] pending_rows_flush_interval = "1s" +max_batch_rows = 100000 ``` +Cases can optionally add `[scenario.remote_write.value]` to control the generated +sample values without changing label cardinality or the runner lifecycle: + +```toml +[scenario.remote_write.value] +pattern = "quantized_signal" # linear, constant, modulo, unique, seeded_random, + # run_length, quantized_signal, + # signal_with_sporadic_stalls, mixed_signal_repeated +base = 0.0 +step = 0.125 +cardinality = 4096 # buckets for modulo/seeded_random/quantized_signal/run_length +seed = 12345 # deterministic seeded_random input +run_length = 8 # adjacent samples per bucket for run_length/quantized_signal +stall_every = 100 # interval for signal_with_sporadic_stalls +stall_length = 16 # held samples inside each stall interval +mixed_every = 5 # every Nth sample becomes the repeated base value +``` + +The default `linear` pattern preserves the helper's historical formula. Use +`constant` or low-cardinality `modulo`/`seeded_random` values for repeated-value +data shapes, `run_length` for run-heavy low-cardinality series, `quantized_signal` +for signal-like values collapsed into a finite bucket set, `signal_with_sporadic_stalls` +for mostly continuous signals with periodic flat spots, and `mixed_signal_repeated` +for signal-plus-periodic-default mixtures. `unique` or high-cardinality buckets +still work for broad sample-value distributions. This is a generic sample-value +control for query/ingestion cases; it does not inspect or assert storage +encoding, Parquet footers, or storage policy choices. For chunked remote-write +ingestion, the runner passes the sample offset and total sample count to the +helper so non-linear value patterns use a stable global/per-series ordinal across +chunks. + +Case schema, value distribution defaults, storage defaults, and read-bench +defaults are owned by Rust. The Python runner calls +`query_perf_fixture plan --case ` and orchestrates the normalized JSON. +The same helper exposes `direct-sst`, `prom-remote-write`, and `inspect-footer` +subcommands; the old direct invocation (`query_perf_fixture --case ... --out-dir +...`) remains compatible. + +Remote-write cases that need to validate storage output can add +`[scenario.remote_write.storage]`. When present, the scenario body becomes: +remote-write → flush → visibility check → query measurement → stop and await +datanode → Parquet footer/size/encoding inspection → optional read-bench. Footer +inspection is handled by `query_perf_fixture inspect-footer`, which reads local +Parquet footers directly; it does not require Python `pyarrow`. + +```toml +[scenario.remote_write.storage] +inspect = true +column = "greptime_value" +include_metadata_files = false +# Optional: inspect below datanode data home instead of the whole fresh data dir. +# root_suffix = "greptime/public/" +min_files = 1 +min_files_with_column = 1 +require_encodings = ["BYTE_STREAM_SPLIT"] +forbid_encodings = ["PLAIN_DICTIONARY"] +max_total_file_size_bytes = 104857600 +max_column_compressed_size_bytes = 52428800 +max_column_uncompressed_size_bytes = 209715200 +max_candidate_total_file_size_regression_pct = 10.0 +max_candidate_column_compressed_size_regression_pct = 10.0 +max_candidate_column_uncompressed_size_regression_pct = 10.0 +``` + +The `*_pct` storage thresholds are percent regression limits comparing candidate +against base; `pct` means percent, not percentile. Per-target storage checks +(`min_files`, `min_files_with_column`, required/forbidden encodings, and absolute +`max_*_bytes`) run for both base and candidate. Comparative +`max_candidate_*_regression_pct` checks remain base-vs-candidate. These checks +are generic footer and byte-size assertions and do not encode product-specific +heuristics. When storage inspection is enabled, `min_files` and +`min_files_with_column` default to `1` even if omitted, so dry-runs show these +planned checks and an empty flush or missing target column fails the scenario. +By default the inspector root is the target datanode data home, which is suitable +for a fresh single-case data directory. For reused or more complex data homes, set +`root_suffix` to a path relative to the datanode data home to narrow inspection. + +When storage inspection is enabled, `[scenario.remote_write.read_bench]` defaults +to enabled with both datanode `parquetbench` and `scanbench`. Disable it with +`enabled = false`. `parquetbench` measures per-SST reader cost, `scanbench` +measures region scan cost, and query measurements still exercise the SQL/TQL +frontend path. Treat all performance conclusions as release-only; debug builds +are suitable only for command wiring and correctness checks. + The runner creates the configured database if needed, writes a per-target frontend config enabling `[prom_store]` with metric engine storage and a non-zero `pending_rows_flush_interval`, and validates that the logical metric table reaches `series_count * samples_per_series` rows before trusting the query measurements. -Use `--remote-write-generator /path/to/prom_remote_write_fixture` to provide the -Rust payload helper. `--fixture-only` is rejected for remote-write cases; use -`--dry-run` for planning. +Use `--fixture-generator /path/to/query_perf_fixture` to provide the Rust helper. +Deprecated `--remote-write-generator` and `--storage-inspector` options are kept +only for CLI compatibility. `--fixture-only` is rejected for remote-write cases; +use `--dry-run` for planning. Large manual remote-write cases can set `sample_chunk_size` to split ingestion by -time. For each chunk, the runner invokes `prom_remote_write_fixture` with the +time. For each chunk, the runner invokes `query_perf_fixture prom-remote-write` with the same series cardinality but a shorter `--samples-per-series` and an advanced `--start-unix-millis`. `flush_every_sample_chunks` controls periodic `ADMIN FLUSH_TABLE('')` calls; with `flush_every_sample_chunks = 1`, @@ -89,6 +183,25 @@ need multi-window file distribution should span multiple compaction windows. If `sample_chunk_size` is omitted, the runner keeps the older single-helper-invocation behavior and flushes once at the end. +The default remote-write coverage set contains four cases with the same 2048 +series × 14,400 samples (29,491,200 rows per target) shape. Each writes ten +one-day chunks, flushes every chunk, requires at least two visible SSTs with the +value column, runs seven read-bench iterations with `parquetbench` capped at +four SSTs while `scanbench` covers the landed region, and measures one ten-day +TQL selector with two warmups and 15 iterations: + +- `prom_remote_write_seeded_random`: high-distinct seeded-random values with + cardinality 29,491,200 and seed 8444. +- `prom_remote_write_run_heavy`: low-cardinality values in exact 16-sample runs. +- `prom_remote_write_mixed_every`: continuous signal values with every fifth + sample repeated at the base value. +- `prom_remote_write_integer_counter`: strictly increasing integer counter + values, with disjoint ranges for each series. + +These cases record normal storage/footer and read-bench results but do not gate +specific value encodings or storage-size outcomes. Use the controlled encoding +experiment matrix for those policy comparisons. + `tests/perf/query_cases/prom_remote_write_7913/case.toml` is a larger manual case for issue #7913. It writes 8192 series × 20160 samples through remote-write in 1440-sample daily time chunks, flushing after each chunk before running 1d/7d/14d @@ -246,7 +359,8 @@ as a fallback. Manifest files are reflinked or copied, not hardlinked. Fixture generator smoke test: ```bash -cargo run -p cmd --bin query_perf_fixture -- \ +cargo run -p cmd --bin query_perf_fixture --features dev-tools -- \ + direct-sst \ --case tests/perf/query_cases/smoke_direct_sst/case.toml \ --out-dir /tmp/query-perf-smoke ``` @@ -271,10 +385,10 @@ Remote-write runner dry-run: ```bash uv run --no-project python tests/perf/query_regression_runner.py \ - --case tests/perf/query_cases/prom_remote_write_smoke/case.toml \ + --case tests/perf/query_cases/prom_remote_write_seeded_random/case.toml \ --base-bin /path/to/base/greptime \ --candidate-bin /path/to/candidate/greptime \ - --remote-write-generator /path/to/prom_remote_write_fixture \ + --fixture-generator /path/to/query_perf_fixture \ --work-dir /tmp/query-perf-remote-write \ --dry-run ``` @@ -289,14 +403,22 @@ query regression runs. It builds its own binaries for now: candidate checkout - runner and summary formatter from the candidate checkout -The workflow runs automatically only for non-draft PRs labeled -`query-regression` (on label/ready-for-review/reopen events, not every push). +The workflow builds base and candidate `greptime` as normal release-equivalent +binaries. Candidate `query_perf_fixture` is the extra head-side helper binary; +the runner uses candidate `greptime datanode parquetbench/scanbench` as the +read-bench tool against each target's data directory. + +The workflow runs automatically only when `query-regression` is added to a +non-draft PR; it does not rerun on pushes, ready-for-review, or reopen events. PR runs build base/candidate once and then run the default case set with `--allow-large-fixture`. Manual `workflow_dispatch` runs can pass `all`, one case path, or a comma/whitespace-separated list of case paths, and can override refs. -It always uploads `query-regression-work/**` and `query-regression-summary.md`, -writes the Markdown summary to the workflow step summary, and updates a sticky PR -comment through the trusted follow-up workflow. +The main report artifact uploads only aggregate/per-target JSON reports, +component logs, and `query-regression-summary.md` with seven-day retention; +fixture data, SSTs, and cluster state are excluded. PR runs also upload a +separate trusted-comment artifact containing PR metadata and aggregate reports. +The workflow writes the Markdown summary to the workflow step summary and +updates a sticky PR comment through the trusted follow-up workflow. ## Built-in cases diff --git a/tests/perf/query_cases/prom_remote_write_integer_counter/case.toml b/tests/perf/query_cases/prom_remote_write_integer_counter/case.toml new file mode 100644 index 00000000000..3ac64a94b9d --- /dev/null +++ b/tests/perf/query_cases/prom_remote_write_integer_counter/case.toml @@ -0,0 +1,48 @@ +[case] +name = "prom_remote_write_integer_counter" +description = "End-to-end Prometheus remote-write regression with strictly increasing integer counter values" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "prom_remote_write_integer_counter" +physical_table = "greptime_physical_table" +series_count = 2048 +samples_per_series = 14400 +sample_chunk_size = 1440 +flush_every_sample_chunks = 1 +start_unix_millis = 1704067200000 +step_millis = 60000 +chunk_series_count = 256 +timeout_seconds = 600 +visibility_timeout_seconds = 300 + +[scenario.remote_write.value] +pattern = "unique" +base = 0 +step = 1 + +[scenario.remote_write.prom_store] +pending_rows_flush_interval = "1s" +max_batch_rows = 100000 + +[scenario.remote_write.storage] +column = "greptime_value" +min_files = 2 +min_files_with_column = 2 + +[scenario.remote_write.read_bench] +iterations = 7 +max_files = 4 + +[[scenario.queries]] +name = "selector_10d_integer_counter" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704931200, '1h') prom_remote_write_integer_counter{host=~'host.*'}" +warmup = 2 +iterations = 15 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 75 diff --git a/tests/perf/query_cases/prom_remote_write_mixed_every/case.toml b/tests/perf/query_cases/prom_remote_write_mixed_every/case.toml new file mode 100644 index 00000000000..9ba39f2fbed --- /dev/null +++ b/tests/perf/query_cases/prom_remote_write_mixed_every/case.toml @@ -0,0 +1,49 @@ +[case] +name = "prom_remote_write_mixed_every" +description = "End-to-end Prometheus remote-write regression with periodic repeated signal values" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "prom_remote_write_mixed_every" +physical_table = "greptime_physical_table" +series_count = 2048 +samples_per_series = 14400 +sample_chunk_size = 1440 +flush_every_sample_chunks = 1 +start_unix_millis = 1704067200000 +step_millis = 60000 +chunk_series_count = 256 +timeout_seconds = 600 +visibility_timeout_seconds = 300 + +[scenario.remote_write.value] +pattern = "mixed_signal_repeated" +base = 0 +step = 0.125 +mixed_every = 5 + +[scenario.remote_write.prom_store] +pending_rows_flush_interval = "1s" +max_batch_rows = 100000 + +[scenario.remote_write.storage] +column = "greptime_value" +min_files = 2 +min_files_with_column = 2 + +[scenario.remote_write.read_bench] +iterations = 7 +max_files = 4 + +[[scenario.queries]] +name = "selector_10d_mixed_every" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704931200, '1h') prom_remote_write_mixed_every{host=~'host.*'}" +warmup = 2 +iterations = 15 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 75 diff --git a/tests/perf/query_cases/prom_remote_write_run_heavy/case.toml b/tests/perf/query_cases/prom_remote_write_run_heavy/case.toml new file mode 100644 index 00000000000..243892e42b1 --- /dev/null +++ b/tests/perf/query_cases/prom_remote_write_run_heavy/case.toml @@ -0,0 +1,50 @@ +[case] +name = "prom_remote_write_run_heavy" +description = "End-to-end Prometheus remote-write regression with repeated 16-sample value runs" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "prom_remote_write_run_heavy" +physical_table = "greptime_physical_table" +series_count = 2048 +samples_per_series = 14400 +sample_chunk_size = 1440 +flush_every_sample_chunks = 1 +start_unix_millis = 1704067200000 +step_millis = 60000 +chunk_series_count = 256 +timeout_seconds = 600 +visibility_timeout_seconds = 300 + +[scenario.remote_write.value] +pattern = "run_length" +base = 0 +step = 1 +cardinality = 1024 +run_length = 16 + +[scenario.remote_write.prom_store] +pending_rows_flush_interval = "1s" +max_batch_rows = 100000 + +[scenario.remote_write.storage] +column = "greptime_value" +min_files = 2 +min_files_with_column = 2 + +[scenario.remote_write.read_bench] +iterations = 7 +max_files = 4 + +[[scenario.queries]] +name = "selector_10d_run_heavy" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704931200, '1h') prom_remote_write_run_heavy{host=~'host.*'}" +warmup = 2 +iterations = 15 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 75 diff --git a/tests/perf/query_cases/prom_remote_write_seeded_random/case.toml b/tests/perf/query_cases/prom_remote_write_seeded_random/case.toml new file mode 100644 index 00000000000..6a63edbcbbf --- /dev/null +++ b/tests/perf/query_cases/prom_remote_write_seeded_random/case.toml @@ -0,0 +1,48 @@ +[case] +name = "prom_remote_write_seeded_random" +description = "End-to-end Prometheus remote-write regression with high-distinct seeded-random values" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "prom_remote_write_seeded_random" +physical_table = "greptime_physical_table" +series_count = 2048 +samples_per_series = 14400 +sample_chunk_size = 1440 +flush_every_sample_chunks = 1 +start_unix_millis = 1704067200000 +step_millis = 60000 +chunk_series_count = 256 +timeout_seconds = 600 +visibility_timeout_seconds = 300 + +[scenario.remote_write.value] +pattern = "seeded_random" +cardinality = 29_491_200 +seed = 8444 + +[scenario.remote_write.prom_store] +pending_rows_flush_interval = "1s" +max_batch_rows = 100000 + +[scenario.remote_write.storage] +column = "greptime_value" +min_files = 2 +min_files_with_column = 2 + +[scenario.remote_write.read_bench] +iterations = 7 +max_files = 4 + +[[scenario.queries]] +name = "selector_10d_seeded_random" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704931200, '1h') prom_remote_write_seeded_random{host=~'host.*'}" +warmup = 2 +iterations = 15 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 75 diff --git a/tests/perf/query_cases/prom_remote_write_smoke/case.toml b/tests/perf/query_cases/prom_remote_write_smoke/case.toml deleted file mode 100644 index b946fbc8f51..00000000000 --- a/tests/perf/query_cases/prom_remote_write_smoke/case.toml +++ /dev/null @@ -1,31 +0,0 @@ -[case] -name = "prom_remote_write_smoke" -description = "Tiny Prometheus remote-write ingestion then TQL query smoke case" - -[scenario] -kind = "prom_remote_write_then_query" - -[scenario.remote_write] -database = "public" -metric = "prom_remote_write_smoke" -physical_table = "greptime_physical_table" -series_count = 8 -samples_per_series = 30 -start_unix_millis = 1_704_067_200_000 -step_millis = 15_000 -chunk_series_count = 4 -timeout_seconds = 30 - -[scenario.remote_write.prom_store] -pending_rows_flush_interval = "1s" -max_batch_rows = 10000 - -[[scenario.queries]] -name = "selector_5m_smoke" -kind = "tql" -query = "TQL ANALYZE VERBOSE (1704067200, 1704067650, '15s') prom_remote_write_smoke{host=~'host.*'}" -warmup = 0 -iterations = 2 - -[scenario.queries.thresholds] -max_candidate_latency_regression_pct = 100 diff --git a/tests/perf/query_regression_runner.py b/tests/perf/query_regression_runner.py index a7f2646ea8b..f5c9628b83d 100644 --- a/tests/perf/query_regression_runner.py +++ b/tests/perf/query_regression_runner.py @@ -66,7 +66,8 @@ def parse_args() -> argparse.Namespace: p.add_argument("--base-bin", required=True, type=Path) p.add_argument("--candidate-bin", required=True, type=Path) p.add_argument("--fixture-generator", type=Path) - p.add_argument("--remote-write-generator", type=Path, help="prom_remote_write_fixture helper binary") + p.add_argument("--remote-write-generator", type=Path, help="deprecated: use --fixture-generator query_perf_fixture") + p.add_argument("--storage-inspector", type=Path, help="deprecated: use --fixture-generator query_perf_fixture") p.add_argument("--work-dir", required=True, type=Path) p.add_argument("--fixture-cache-dir", type=Path, help="persistent directory for generated fixtures, keyed by case content") p.add_argument("--reuse-fixture", action="store_true") @@ -83,6 +84,16 @@ def load_case(path: Path) -> dict[str, Any]: return tomllib.load(f) +def load_normalized_case(case_path: Path, fixture_generator: Path | None) -> dict[str, Any]: + if fixture_generator is None: + raise ValueError("--fixture-generator query_perf_fixture is required for Rust-owned case planning") + result = run_command([str(fixture_generator), "plan", "--case", str(case_path)]) + if result["returncode"] != 0: + raise RuntimeError(f"query_perf_fixture plan failed: {result['stderr'][:2000]}") + plan = json.loads(result["stdout"]) + return {"case": load_case(case_path).get("case", {}), "scenario": plan["scenario"], "schema_version": plan.get("schema_version")} + + def require_binary(path: Path, name: str, *, dry_run: bool) -> None: if dry_run: return @@ -157,11 +168,9 @@ def scenario(case: dict[str, Any]) -> dict[str, Any]: def case_tables(case: dict[str, Any]) -> list[dict[str, Any]]: value = scenario(case) if value.get("kind") == "prom_remote_write_then_query": - remote = value.get("remote_write") or {} - metric = remote.get("metric") or remote.get("metric_name") - database = remote.get("database", "public") - if not metric: - raise ValueError("remote-write scenario requires scenario.remote_write.metric") + remote = value["remote_write"] + metric = remote["metric"] + database = remote["database"] return [{"database": database, "name": metric, "engine": "metric", "validate_show_create_engine": False}] tables = value.get("tables") or [] if not tables or value.get("layout", {}).get("regions") != 1: @@ -548,7 +557,7 @@ def generate_fixture(generator: Path | None, case_path: Path, fixture_dir: Path, if fixture_dir.exists() and not reuse_fixture: shutil.rmtree(fixture_dir) fixture_dir.mkdir(parents=True, exist_ok=True) - cmd = [str(generator), "--case", str(case_path), "--out-dir", str(fixture_dir)] + cmd = [str(generator), "direct-sst", "--case", str(case_path), "--out-dir", str(fixture_dir)] if table_name is not None: cmd += ["--table", table_name] if region_id is not None: @@ -777,6 +786,243 @@ def enforce_thresholds(case: dict[str, Any], base: dict[str, Any], candidate: di return results +def storage_config(remote: dict[str, Any]) -> dict[str, Any] | None: + value = remote["storage"] + if not isinstance(value, dict): + return None + if value["inspect"] is False: + return None + return value + + +def run_storage_inspection(helper: Path | None, target: RunTarget, storage: dict[str, Any], *, dry_run: bool) -> dict[str, Any]: + if helper is None: + if not dry_run: + raise ValueError("--fixture-generator query_perf_fixture is required when storage inspection is enabled") + helper = Path("query_perf_fixture") + root = target.datanode_data_dir + if storage.get("root_suffix"): + root = root / str(storage["root_suffix"]) + cmd = [ + str(helper), + "inspect-footer", + "--root", str(root), + "--column", str(storage["column"]), + ] + if storage["include_metadata_files"]: + cmd.append("--include-metadata-files") + if dry_run: + return {"status": "dry-run", "cmd": cmd, "root": str(root)} + result = run_command(cmd) + result["status"] = "ok" if result["returncode"] == 0 else "failed" + if result["returncode"] != 0: + raise RuntimeError(f"storage inspector failed for {target.name}: {result['stderr'][:2000]}") + try: + result["summary"] = json.loads(result["stdout"]) + except json.JSONDecodeError: + result["summary_parse_error"] = result["stdout"] + result["status"] = "failed" + return result + + +def storage_summary(inspection: dict[str, Any]) -> dict[str, Any]: + summary = inspection.get("summary") + if isinstance(summary, dict): + nested = summary.get("summary") + if isinstance(nested, dict): + return nested + return {} + + +def enforce_storage_thresholds(storage: dict[str, Any] | None, base: dict[str, Any] | None, candidate: dict[str, Any] | None) -> list[dict[str, Any]]: + if not storage: + return [] + results: list[dict[str, Any]] = [] + target_summaries = [("base", storage_summary(base or {})), ("candidate", storage_summary(candidate or {}))] + cand = target_summaries[1][1] + base_summary = storage_summary(base or {}) + + def add_abs(name: str, field: str) -> None: + limit = storage.get(name) + if limit is None: + return + for target_name, summary in target_summaries: + actual = summary.get(field) + ok = actual is not None and float(actual) <= float(limit) + results.append({"target": target_name, "threshold": name, "status": "passed" if ok else "failed", "actual": actual, "limit": limit}) + + def add_min(name: str, field: str) -> None: + limit = storage[name] + if limit is None: + return + for target_name, summary in target_summaries: + actual = summary.get(field) + ok = actual is not None and int(actual) >= int(limit) + results.append({"target": target_name, "threshold": name, "status": "passed" if ok else "failed", "actual": actual, "limit": limit}) + + def add_cmp(name: str, field: str) -> None: + limit = storage.get(name) + if limit is None: + return + base_value = base_summary.get(field) + candidate_value = cand.get(field) + if base_value in (None, 0) or candidate_value is None: + results.append({"threshold": name, "status": "failed", "reason": "missing or zero base/candidate value", "base": base_value, "candidate": candidate_value}) + return + actual = (float(candidate_value) - float(base_value)) / float(base_value) * 100.0 + results.append({"threshold": name, "status": "passed" if actual <= float(limit) else "failed", "actual_pct": actual, "limit_pct": limit, "base": base_value, "candidate": candidate_value}) + + add_min("min_files", "file_count") + add_min("min_files_with_column", "files_with_column") + add_abs("max_total_file_size_bytes", "total_file_size") + add_abs("max_column_compressed_size_bytes", "column_compressed_size") + add_abs("max_column_uncompressed_size_bytes", "column_uncompressed_size") + for target_name, summary in target_summaries: + encodings = set(summary.get("unique_encodings") or []) + for required in storage["require_encodings"]: + results.append({"target": target_name, "threshold": "require_encodings", "encoding": required, "status": "passed" if required in encodings else "failed"}) + for forbidden in storage["forbid_encodings"]: + results.append({"target": target_name, "threshold": "forbid_encodings", "encoding": forbidden, "status": "passed" if forbidden not in encodings else "failed"}) + add_cmp("max_candidate_total_file_size_regression_pct", "total_file_size") + add_cmp("max_candidate_column_compressed_size_regression_pct", "column_compressed_size") + add_cmp("max_candidate_column_uncompressed_size_regression_pct", "column_uncompressed_size") + return results + + +def planned_storage_thresholds(storage: dict[str, Any] | None) -> list[dict[str, Any]]: + if not storage: + return [] + return list(storage["planned_thresholds"]) + + +REGION_DIR_RE = re.compile(r"^(\d+)_(\d{10})$") + + +def inspection_report(inspection: dict[str, Any] | None) -> dict[str, Any]: + if not inspection: + return {} + report = inspection.get("summary") + return report if isinstance(report, dict) else {} + + +def inspected_relative_path(target: RunTarget, inspection: dict[str, Any], relative_path: str) -> Path: + report = inspection_report(inspection) + root_text = report.get("root") or inspection.get("root") + if not root_text: + return Path(relative_path) + root = Path(root_text).resolve(strict=False) + data_home = target.datanode_data_dir.resolve(strict=False) + try: + root_suffix = root.relative_to(data_home) + except ValueError as e: + raise RuntimeError(f"storage inspection root {root} is not under datanode data home {data_home}") from e + return root_suffix / relative_path + + +def parse_sst_bench_target(target: RunTarget, inspection: dict[str, Any], file: dict[str, Any], *, dry_run: bool) -> dict[str, str] | None: + relative_path = str(file.get("relative_path") or "") + if dry_run and (not relative_path or "<" in relative_path): + return { + "relative_path": relative_path or "/_/data/.parquet", + "table_dir": "/", + "region_id": ":", + "path_type": "data", + "file_id": "", + } + if not relative_path: + return None + full_relative = inspected_relative_path(target, inspection, relative_path) + parts = full_relative.parts + if not parts or not parts[-1].endswith(".parquet"): + return None + region_index = next((idx for idx, part in enumerate(parts) if REGION_DIR_RE.match(part)), None) + if region_index is None or region_index == 0: + if dry_run: + return { + "relative_path": relative_path, + "table_dir": "/", + "region_id": ":", + "path_type": "data", + "file_id": Path(parts[-1]).stem, + } + return None + region_match = REGION_DIR_RE.match(parts[region_index]) + assert region_match is not None + file_index = len(parts) - 1 + path_type = "bare" + if region_index + 1 < file_index and parts[region_index + 1] in ("data", "metadata"): + path_type = parts[region_index + 1] + if path_type == "metadata": + return None + table_dir = "/".join(parts[:region_index]).rstrip("/") + "/" + return { + "relative_path": str(full_relative), + "table_dir": table_dir, + "region_id": f"{int(region_match.group(1))}:{int(region_match.group(2))}", + "path_type": path_type, + "file_id": Path(parts[-1]).stem, + } + + +def run_read_bench(bench_binary: Path, target: RunTarget, read_bench: dict[str, Any] | None, storage_inspection: dict[str, Any] | None, *, dry_run: bool) -> dict[str, Any]: + if not read_bench or read_bench["enabled"] is False: + return {"status": "skipped", "reason": "read_bench disabled"} + run_parquetbench = read_bench["parquetbench"] + run_scanbench = read_bench["scanbench"] + if not run_parquetbench and not run_scanbench: + return {"status": "skipped", "reason": "parquetbench and scanbench disabled"} + if storage_inspection is None: + raise ValueError("read_bench requires storage inspection") + bench_dir = target.work_dir / "read_bench" + config_toml = bench_dir / "bench.toml" + scan_json = bench_dir / "scan.json" + files = inspection_report(storage_inspection).get("files") or [] + ssts = [f for f in files if f.get("relative_path", "").endswith(".parquet") and f.get("columns")] + if dry_run and not ssts: + ssts = [{"relative_path": ".parquet"}] + if read_bench["max_files"] is not None: + ssts = ssts[: int(read_bench["max_files"])] + bench_targets = [parsed for f in ssts if (parsed := parse_sst_bench_target(target, storage_inspection, f, dry_run=dry_run)) is not None] + if not bench_targets: + return {"status": "failed", "reason": "no inspected data SST files available for read_bench"} + config_text = f'[storage]\ndata_home = "{target.datanode_data_dir}"\ntype = "File"\n\n[[region_engine]]\n[region_engine.mito]\n' + commands = [] + parquet_runs = [] + scan_runs = [] + iterations = str(int(read_bench["iterations"])) + if run_parquetbench: + for bench_target in bench_targets: + cmd = [str(bench_binary), "datanode", "parquetbench", "--config", str(config_toml), "--region-id", bench_target["region_id"], "--table-dir", bench_target["table_dir"], "--file-id", bench_target["file_id"], "--scan-config", str(scan_json), "--path-type", bench_target["path_type"], "--iterations", iterations, "--reader", str(read_bench["parquet_reader"])] + commands.append(cmd) + parquet_runs.append({**bench_target, "cmd": cmd, "status": "dry-run" if dry_run else "planned"}) + regions: dict[tuple[str, str, str], list[str]] = {} + if run_scanbench: + for bench_target in bench_targets: + regions.setdefault((bench_target["table_dir"], bench_target["region_id"], bench_target["path_type"]), []).append(bench_target["relative_path"]) + for (table_dir, region_id, path_type), region_files in regions.items(): + cmd = [str(bench_binary), "datanode", "scanbench", "--config", str(config_toml), "--region-id", region_id, "--table-dir", table_dir, "--scan-config", str(scan_json), "--path-type", path_type, "--scanner", str(read_bench["scan_scanner"]), "--parallelism", str(int(read_bench["parallelism"])), "--iterations", iterations] + commands.append(cmd) + scan_runs.append({"table_dir": table_dir, "region_id": region_id, "path_type": path_type, "files": region_files, "cmd": cmd, "status": "dry-run" if dry_run else "planned"}) + if dry_run: + return {"status": "dry-run", "config_path": str(config_toml), "scan_config_path": str(scan_json), "commands": commands, "parquetbench": parquet_runs, "scanbench": scan_runs} + bench_dir.mkdir(parents=True, exist_ok=True) + config_toml.write_text(config_text) + scan_json.write_text(json.dumps({"projection_names": read_bench["projection"]}, indent=2) + "\n") + avg_re = re.compile(r"Average duration[^0-9]*([0-9.]+)\s*ms", re.I) + for runs in (parquet_runs, scan_runs): + for run in runs: + result = run_command(run["cmd"]) + run.update(result) + run["status"] = "ok" if result["returncode"] == 0 else "failed" + m = avg_re.search(result.get("stdout", "")) + if m: + run["average_ms"] = float(m.group(1)) + def median(runs: list[dict[str, Any]]) -> float | None: + vals = [r["average_ms"] for r in runs if "average_ms" in r] + return statistics.median(vals) if vals else None + return {"status": "ok" if all(r.get("status") == "ok" for r in parquet_runs + scan_runs) else "failed", "config_path": str(config_toml), "scan_config_path": str(scan_json), "parquetbench": parquet_runs, "scanbench": scan_runs, "aggregate": {"parquetbench_median_average_ms": median(parquet_runs), "scanbench_median_average_ms": median(scan_runs)}} + + def write_json(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") @@ -791,18 +1037,18 @@ def require_fresh_work_dirs(targets: list[RunTarget], *, reuse_work_dir: bool, d def write_frontend_prom_config(target: RunTarget, remote: dict[str, Any]) -> Path: - prom = remote.get("prom_store") or {} + prom = remote["prom_store"] path = target.work_dir / "frontend-prom-store.toml" path.parent.mkdir(parents=True, exist_ok=True) path.write_text( "[prom_store]\n" "enable = true\n" "with_metric_engine = true\n" - f"pending_rows_flush_interval = {json.dumps(str(prom.get('pending_rows_flush_interval', '1s')))}\n" - f"max_batch_rows = {int(prom.get('max_batch_rows', 100000))}\n" - f"max_concurrent_flushes = {int(prom.get('max_concurrent_flushes', 256))}\n" - f"worker_channel_capacity = {int(prom.get('worker_channel_capacity', 65526))}\n" - f"max_inflight_requests = {int(prom.get('max_inflight_requests', 3000))}\n" + f"pending_rows_flush_interval = {json.dumps(str(prom['pending_rows_flush_interval']))}\n" + f"max_batch_rows = {int(prom['max_batch_rows'])}\n" + f"max_concurrent_flushes = {int(prom['max_concurrent_flushes'])}\n" + f"worker_channel_capacity = {int(prom['worker_channel_capacity'])}\n" + f"max_inflight_requests = {int(prom['max_inflight_requests'])}\n" ) return path @@ -810,23 +1056,38 @@ def write_frontend_prom_config(target: RunTarget, remote: dict[str, Any]) -> Pat def run_remote_write(generator: Path | None, target: RunTarget, remote: dict[str, Any], *, dry_run: bool) -> dict[str, Any]: if generator is None: if not dry_run: - raise ValueError("--remote-write-generator is required for prom_remote_write_then_query unless --dry-run is set") - generator = Path("prom_remote_write_fixture") - metric = remote.get("metric") or remote.get("metric_name") - physical_table = remote.get("physical_table", "greptime_physical_table") + raise ValueError("--fixture-generator query_perf_fixture is required for prom_remote_write_then_query") + generator = Path("query_perf_fixture") + metric = remote["metric"] + physical_table = remote["physical_table"] cmd = [ str(generator), + "prom-remote-write", "--endpoint", f"http://127.0.0.1:{target.http_port}/v1/prometheus/write", - "--database", remote.get("database", "public"), + "--database", remote["database"], "--metric", metric, "--physical-table", physical_table, - "--series-count", str(int(remote.get("series_count", 8))), - "--samples-per-series", str(int(remote.get("samples_per_series", 30))), - "--start-unix-millis", str(int(remote.get("start_unix_millis", 1_704_067_200_000))), - "--step-millis", str(int(remote.get("step_millis", 15_000))), - "--chunk-series-count", str(int(remote.get("chunk_series_count", remote.get("batch_size", 8)))), - "--timeout-seconds", str(int(remote.get("timeout_seconds", 60))), + "--series-count", str(int(remote["series_count"])), + "--samples-per-series", str(int(remote["samples_per_series"])), + "--start-unix-millis", str(int(remote["start_unix_millis"])), + "--step-millis", str(int(remote["step_millis"])), + "--chunk-series-count", str(int(remote["chunk_series_count"])), + "--timeout-seconds", str(int(remote["timeout_seconds"])), ] + value = remote["value"] + cmd.extend(["--value-pattern", str(value["pattern"])]) + cmd.extend(["--value-base", str(value["base"])]) + cmd.extend(["--value-step", str(value["step"])]) + cmd.extend(["--value-cardinality", str(int(value["cardinality"]))]) + cmd.extend(["--value-seed", str(int(value["seed"]))]) + cmd.extend(["--value-run-length", str(int(value["run_length"]))]) + cmd.extend(["--value-stall-every", str(int(value["stall_every"]))]) + cmd.extend(["--value-stall-length", str(int(value["stall_length"]))]) + cmd.extend(["--value-mixed-every", str(int(value["mixed_every"]))]) + if "sample_offset" in remote: + cmd.extend(["--value-sample-offset", str(int(remote["sample_offset"]))]) + if "total_samples_per_series" in remote: + cmd.extend(["--value-total-samples-per-series", str(int(remote["total_samples_per_series"]))]) if dry_run: return {"status": "dry-run", "cmd": cmd} result = run_command(cmd) @@ -862,23 +1123,23 @@ def flush_remote_physical_table(target: RunTarget, db: str, physical_table: str, def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote: dict[str, Any], args: argparse.Namespace, *, dry_run: bool) -> tuple[dict[str, Any], list[dict[str, Any]]]: - sample_chunk_size = remote.get("sample_chunk_size") - db = remote.get("database", "public") - physical_table = remote.get("physical_table", "greptime_physical_table") + sample_chunk_size = remote["sample_chunk_size"] + db = remote["database"] + physical_table = remote["physical_table"] if sample_chunk_size is None: rw = run_remote_write(generator, target, remote, dry_run=dry_run) flush = flush_remote_physical_table(target, db, physical_table, args, dry_run=dry_run, reason="final") return rw, [flush] - total_samples = int(remote.get("samples_per_series", 30)) + total_samples = int(remote["samples_per_series"]) chunk_samples = int(sample_chunk_size) if chunk_samples <= 0: raise ValueError("scenario.remote_write.sample_chunk_size must be positive") - flush_every = int(remote.get("flush_every_sample_chunks", 1)) + flush_every = int(remote["flush_every_sample_chunks"]) if flush_every <= 0: raise ValueError("scenario.remote_write.flush_every_sample_chunks must be positive") - start = int(remote.get("start_unix_millis", 1_704_067_200_000)) - step = int(remote.get("step_millis", 15_000)) + start = int(remote["start_unix_millis"]) + step = int(remote["step_millis"]) chunks: list[dict[str, Any]] = [] flushes: list[dict[str, Any]] = [] chunk_index = 0 @@ -889,6 +1150,8 @@ def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote chunk_remote = dict(remote) chunk_remote["samples_per_series"] = current chunk_remote["start_unix_millis"] = start + offset * step + chunk_remote["sample_offset"] = offset + chunk_remote["total_samples_per_series"] = total_samples result = run_remote_write(generator, target, chunk_remote, dry_run=dry_run) result["sample_offset"] = offset result["samples_per_series"] = current @@ -901,8 +1164,8 @@ def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote flushes.append(flush_remote_physical_table(target, db, physical_table, args, dry_run=dry_run, reason="final", chunk_index=chunk_index)) aggregate = summarize_remote_write_chunks(chunks) if dry_run: - series_count = int(remote.get("series_count", 8)) - chunk_series_count = int(remote.get("chunk_series_count", remote.get("batch_size", 8))) + series_count = int(remote["series_count"]) + chunk_series_count = int(remote["chunk_series_count"]) aggregate["rows"] = series_count * total_samples aggregate["samples_written"] = aggregate["rows"] aggregate["batches"] = chunk_index * ((series_count + chunk_series_count - 1) // chunk_series_count) @@ -917,7 +1180,7 @@ def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote def expected_remote_write_rows(remote: dict[str, Any]) -> int: - return int(remote.get("series_count", 8)) * int(remote.get("samples_per_series", 30)) + return int(remote["series_count"]) * int(remote["samples_per_series"]) def extract_count_value(result: dict[str, Any]) -> int | None: @@ -966,16 +1229,21 @@ def poll_expected_count(target: RunTarget, table_name: str, db: str, expected_ro def run_remote_write_scenario(args: argparse.Namespace, case: dict[str, Any], case_path: Path, targets: list[RunTarget], report: dict[str, Any]) -> None: - remote = scenario(case).get("remote_write") or {} + remote = scenario(case)["remote_write"] tables = case_tables(case) if args.fixture_only: raise ValueError("--fixture-only is not supported for prom_remote_write_then_query; use --dry-run for planning") - if args.remote_write_generator is not None: - require_binary(args.remote_write_generator, "remote-write generator", dry_run=args.dry_run) + helper = args.fixture_generator or args.remote_write_generator + if helper is not None: + require_binary(helper, "query_perf_fixture", dry_run=args.dry_run) elif not args.dry_run: - raise ValueError("--remote-write-generator is required for prom_remote_write_then_query") + raise ValueError("--fixture-generator is required for prom_remote_write_then_query") + storage = storage_config(remote) + if storage and helper is None and not args.dry_run: + raise ValueError("--fixture-generator is required when scenario.remote_write.storage.inspect is enabled") clusters: list[DistributedCluster] = [] target_results = [] + storage_results = [] try: for target in targets: target.work_dir.mkdir(parents=True, exist_ok=True) @@ -984,28 +1252,40 @@ def run_remote_write_scenario(args: argparse.Namespace, case: dict[str, Any], ca clusters.append(cluster) if not args.dry_run: cluster.start_all(config_path) - db = remote.get("database", "public") + db = remote["database"] create_database = {"status": "dry-run", "database": db} if not args.dry_run: create_database = http_post_sql(target.http_port, f"CREATE DATABASE IF NOT EXISTS {sql_ident(db)}", "public", args.http_timeout) - rw, flushes = run_remote_write_ingestion(args.remote_write_generator, target, remote, args, dry_run=args.dry_run) - physical = remote.get("physical_table", "greptime_physical_table") + rw, flushes = run_remote_write_ingestion(helper, target, remote, args, dry_run=args.dry_run) visibility = {"status": "dry-run"} query_result = {"validation": [], "validation_errors": [], "measurements": [], "status": "planned"} if not args.dry_run: - visibility = poll_expected_count(target, tables[0]["name"], db, expected_remote_write_rows(remote), float(remote.get("visibility_timeout_seconds", 30)), args.http_timeout) + visibility = poll_expected_count(target, tables[0]["name"], db, expected_remote_write_rows(remote), float(remote["visibility_timeout_seconds"]), args.http_timeout) query_result = run_queries(target, case, tables, args.http_timeout) + cluster.stop_component("datanode") + storage_inspection = None + if storage: + storage_inspection = run_storage_inspection(helper or args.storage_inspector, target, storage, dry_run=args.dry_run) + storage_results.append(storage_inspection) + read_bench_result = run_read_bench(args.candidate_bin, target, remote["read_bench"], storage_inspection, dry_run=args.dry_run) tr = {"name": target.name, "binary": str(target.binary), "work_dir": str(target.work_dir), "components": cluster.component_report(), "frontend_config": str(config_path), "create_database": create_database, "remote_write": rw, "flushes": flushes, "flush": flushes[-1] if flushes else None, "visibility": visibility, **query_result} + if storage_inspection is not None: + tr["storage_inspection"] = storage_inspection + tr["read_bench"] = read_bench_result flushes_ok = all(flush.get("ok") for flush in flushes) - remote_checks_ok = args.dry_run or (create_database.get("ok") and flushes_ok and visibility.get("ok") and visibility.get("row_count_ok")) + storage_ok = storage_inspection is None or args.dry_run or storage_inspection.get("status") == "ok" + read_bench_ok = args.dry_run or read_bench_result.get("status") in ("ok", "skipped") + remote_checks_ok = args.dry_run or (create_database.get("ok") and flushes_ok and visibility.get("ok") and visibility.get("row_count_ok") and storage_ok and read_bench_ok) if not remote_checks_ok: - tr.setdefault("validation_errors", []).append({"phase": "remote_write_visibility", "create_database_ok": create_database.get("ok"), "flushes_ok": flushes_ok, "visibility_ok": visibility.get("ok"), "row_count_ok": visibility.get("row_count_ok"), "create_database": create_database, "flushes": flushes, "visibility": visibility}) + tr.setdefault("validation_errors", []).append({"phase": "remote_write_visibility_storage", "create_database_ok": create_database.get("ok"), "flushes_ok": flushes_ok, "visibility_ok": visibility.get("ok"), "row_count_ok": visibility.get("row_count_ok"), "storage_ok": storage_ok, "read_bench_ok": read_bench_ok, "create_database": create_database, "flushes": flushes, "visibility": visibility, "read_bench": read_bench_result}) tr["status"] = "planned" if args.dry_run else ("measured" if remote_checks_ok and query_result["status"] == "ok" else "failed") write_json(target.report_path, tr) report["targets"].append(tr) target_results.append(query_result) if not args.dry_run: - report["thresholds"] = enforce_thresholds(case, target_results[0], target_results[1]) + report["thresholds"] = enforce_thresholds(case, target_results[0], target_results[1]) + enforce_storage_thresholds(storage, storage_results[0] if storage_results else None, storage_results[1] if len(storage_results) > 1 else None) + elif storage: + report["thresholds"] = planned_storage_thresholds(storage) report["status"] = "planned" if args.dry_run else ("failed" if any(t["status"] == "failed" for t in report["thresholds"]) or any(t.get("status") == "failed" for t in report["targets"]) else "ok") finally: for cluster in reversed(clusters): @@ -1015,7 +1295,9 @@ def run_remote_write_scenario(args: argparse.Namespace, case: dict[str, Any], ca def main() -> int: args = parse_args() case_path = args.case.resolve() - case = load_case(case_path) + if args.fixture_generator is not None: + require_binary(args.fixture_generator, "fixture generator", dry_run=False) + case = load_normalized_case(case_path, args.fixture_generator) scenario_config = scenario(case) scenario_kind = scenario_config.get("kind") tables = case_tables(case) @@ -1025,9 +1307,6 @@ def main() -> int: require_binary(args.candidate_bin, "candidate", dry_run=args.dry_run or args.fixture_only) if scenario_kind == "direct_readable_sst" and not args.dry_run and args.fixture_generator is None: raise ValueError("--fixture-generator is required unless --dry-run is set") - if args.fixture_generator is not None: - require_binary(args.fixture_generator, "fixture generator", dry_run=args.dry_run) - ports = allocate_ports(16) targets = [make_target("base", args.base_bin.resolve(), work_root, ports[:8]), make_target("candidate", args.candidate_bin.resolve(), work_root, ports[8:])] require_fresh_work_dirs(targets, reuse_work_dir=args.reuse_work_dir, dry_run=args.dry_run, fixture_only=args.fixture_only) diff --git a/tests/perf/test_query_regression_runner_compaction_toctou.py b/tests/perf/test_query_regression_runner_compaction_toctou.py new file mode 100644 index 00000000000..675f451a10e --- /dev/null +++ b/tests/perf/test_query_regression_runner_compaction_toctou.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +# 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. + +"""Regression coverage for quiescing remote-write storage before read benches.""" + +import argparse +import importlib.util +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + + +RUNNER_PATH = Path(__file__).with_name("query_regression_runner.py") +SPEC = importlib.util.spec_from_file_location("query_regression_runner_under_test", RUNNER_PATH) +assert SPEC is not None and SPEC.loader is not None +runner = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = runner +SPEC.loader.exec_module(runner) + + +class RemoteWriteCompactionToctouTest(unittest.TestCase): + def test_quiesces_each_target_before_storage_inspection_and_read_bench(self) -> None: + lifecycle: dict[str, list[str]] = {"base": [], "candidate": []} + clusters = [] + + class RecordingCluster: + def __init__(self, target): + self.target = target + self.active = set() + self.stop_all_calls = 0 + clusters.append(self) + + def start_all(self, _config_path): + self.active.update(("metasrv", "datanode", "frontend")) + + def stop_component(self, name): + if name in self.active: + lifecycle[self.target.name].append(f'cluster.stop_component("{name}")') + self.active.remove(name) + + def stop_all(self): + self.stop_all_calls += 1 + for name in ("frontend", "datanode", "metasrv"): + self.stop_component(name) + + def component_report(self): + return {} + + def run_queries(target, _case, _tables, _http_timeout): + lifecycle[target.name].append("query measurement completes") + return {"validation": [], "validation_errors": [], "measurements": [], "status": "ok"} + + def inspect_storage(_helper, target, _storage, *, dry_run): + self.assertFalse(dry_run) + lifecycle[target.name].append("storage inspection") + return {"status": "ok", "summary": {"summary": {}}} + + def read_bench(_binary, target, _read_bench, _inspection, *, dry_run): + self.assertFalse(dry_run) + lifecycle[target.name].append("read bench") + return {"status": "ok"} + + remote_write = { + "database": "public", + "metric": "cpu_usage", + "physical_table": "cpu_usage_physical", + "series_count": 1, + "samples_per_series": 1, + "visibility_timeout_seconds": 1, + "storage": { + "inspect": True, + "min_files": None, + "min_files_with_column": None, + "max_total_file_size_bytes": None, + "max_column_compressed_size_bytes": None, + "max_column_uncompressed_size_bytes": None, + "require_encodings": [], + "forbid_encodings": [], + }, + "read_bench": {"enabled": True}, + "prom_store": { + "pending_rows_flush_interval": "1s", + "max_batch_rows": 1, + "max_concurrent_flushes": 1, + "worker_channel_capacity": 1, + "max_inflight_requests": 1, + }, + } + case = { + "scenario": { + "kind": "prom_remote_write_then_query", + "remote_write": remote_write, + "queries": [], + } + } + args = argparse.Namespace( + fixture_only=False, + fixture_generator=Path("/bin/true"), + remote_write_generator=None, + storage_inspector=None, + candidate_bin=Path("/bin/true"), + dry_run=False, + http_timeout=1.0, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + root = Path(tmpdir) + targets = [ + runner.make_target("base", Path("/bin/true"), root, list(range(10000, 10008))), + runner.make_target("candidate", Path("/bin/true"), root, list(range(10008, 10016))), + ] + report = {"targets": [], "thresholds": []} + with ( + patch.object(runner, "DistributedCluster", RecordingCluster), + patch.object(runner, "run_remote_write_ingestion", return_value=({"status": "ok"}, [{"ok": True}])), + patch.object(runner, "http_post_sql", return_value={"ok": True}), + patch.object(runner, "poll_expected_count", return_value={"ok": True, "row_count_ok": True}), + patch.object(runner, "run_queries", side_effect=run_queries), + patch.object(runner, "run_storage_inspection", side_effect=inspect_storage), + patch.object(runner, "run_read_bench", side_effect=read_bench), + ): + runner.run_remote_write_scenario(args, case, root / "case.toml", targets, report) + + for cluster in clusters: + self.assertEqual(cluster.stop_all_calls, 1) + self.assertEqual(cluster.active, set()) + cluster.stop_all() + self.assertEqual(cluster.stop_all_calls, 2) + self.assertEqual(cluster.active, set()) + + expected = [ + "query measurement completes", + 'cluster.stop_component("datanode")', + "storage inspection", + "read bench", + ] + relevant = {name: [event for event in events if event in expected] for name, events in lifecycle.items()} + for target_name in ("base", "candidate"): + with self.subTest(target=target_name): + self.assertEqual(relevant[target_name], expected) + + +if __name__ == "__main__": + unittest.main() From c12f40cec232dda23429a0995d70bb4a230a562c Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sat, 11 Jul 2026 21:13:36 +0800 Subject: [PATCH 07/48] chore: update star history link (#8483) Signed-off-by: Dennis Zhuang --- README.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 3ba52bc5266..99c78814609 100644 --- a/README.md +++ b/README.md @@ -209,7 +209,13 @@ Read the [v1.0 highlights](https://greptime.com/blogs/2025-11-05-greptimedb-v1-h If GreptimeDB is useful to you, please star the repo. -[![Star History Chart](https://api.star-history.com/svg?repos=GreptimeTeam/GreptimeDB&type=Date)](https://www.star-history.com/#GreptimeTeam/GreptimeDB&Date) + + + + + Star History Chart + + Known Users From b411c6265c19debefddac5a38f9365cd59122398 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 13 Jul 2026 10:46:33 +0800 Subject: [PATCH 08/48] test: add json2 variant payload compat case (#8459) * test: add json2 variant payload compat case Signed-off-by: discord9 * test: refine json2 variant compat case Signed-off-by: discord9 --------- Signed-off-by: discord9 --- .../legacy_json2_variant_payload/case.toml | 8 +++ .../legacy_json2_variant_payload/setup.sql | 29 +++++++++ .../verify.result | 60 +++++++++++++++++++ .../legacy_json2_variant_payload/verify.sql | 16 +++++ 4 files changed, 113 insertions(+) create mode 100644 tests/compatibility/cases/legacy_json2_variant_payload/case.toml create mode 100644 tests/compatibility/cases/legacy_json2_variant_payload/setup.sql create mode 100644 tests/compatibility/cases/legacy_json2_variant_payload/verify.result create mode 100644 tests/compatibility/cases/legacy_json2_variant_payload/verify.sql diff --git a/tests/compatibility/cases/legacy_json2_variant_payload/case.toml b/tests/compatibility/cases/legacy_json2_variant_payload/case.toml new file mode 100644 index 00000000000..dd46e28ba4d --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_variant_payload/case.toml @@ -0,0 +1,8 @@ +name = "legacy_json2_variant_payload" +reason = "Verify current binaries can read JSON2 variant payloads written by old binaries after the JSON2 encoding change." +introduced_by = "PR #8435" +topologies = ["distributed"] +from_range = ["=v1.1.0", "=v1.1.1", "=v1.1.2"] +to_range = [">v1.1.1"] +features = ["json", "json2", "table", "read"] +owner = "query" diff --git a/tests/compatibility/cases/legacy_json2_variant_payload/setup.sql b/tests/compatibility/cases/legacy_json2_variant_payload/setup.sql new file mode 100644 index 00000000000..cef3a913c93 --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_variant_payload/setup.sql @@ -0,0 +1,29 @@ +CREATE TABLE t_legacy_json2_variant_payload ( + ts TIMESTAMP TIME INDEX, + j JSON2 +) WITH ( + 'append_mode' = 'true' +); + +INSERT INTO t_legacy_json2_variant_payload (ts, j) VALUES + ('2026-07-08 00:00:00+0000', '{"a": {"b": 1}, "c": "s1", "d": [{"e": {"f": 0.1}}]}'), + ('2026-07-08 00:01:00+0000', '{"a": {"b": -2}, "c": [1], "d": [{"e": {"f": 0.2}}]}'), + ('2026-07-08 00:02:00+0000', '{"a": {"b": "s3"}, "c": "s3", "d": [{"e": {"g": -0.3}}]}'), + ('2026-07-08 00:03:00+0000', '{"a": {"b": null}, "c": null}'), + ('2026-07-08 00:04:00+0000', '{"c": "s5", "d": {"e": true}}'), + ('2026-07-08 00:05:00+0000', '{"a": {"b": 6}}'); + +ADMIN FLUSH_TABLE('t_legacy_json2_variant_payload'); + +CREATE TABLE t_legacy_json2_variant_payload_object_d ( + ts TIMESTAMP TIME INDEX, + j JSON2 +) WITH ( + 'append_mode' = 'true' +); + +INSERT INTO t_legacy_json2_variant_payload_object_d (ts, j) VALUES + ('2026-07-08 01:00:00+0000', '{"d": {"e": true}}'), + ('2026-07-08 01:01:00+0000', '{"d": {"e": false}}'); + +ADMIN FLUSH_TABLE('t_legacy_json2_variant_payload_object_d'); diff --git a/tests/compatibility/cases/legacy_json2_variant_payload/verify.result b/tests/compatibility/cases/legacy_json2_variant_payload/verify.result new file mode 100644 index 00000000000..a474cc84345 --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_variant_payload/verify.result @@ -0,0 +1,60 @@ +SELECT ts, j.a.b AS a_b, j.c AS c +FROM t_legacy_json2_variant_payload +ORDER BY ts; + ++---------------------+-----+-----+ +| ts | a_b | c | ++---------------------+-----+-----+ +| 2026-07-08T00:00:00 | 1 | s1 | +| 2026-07-08T00:01:00 | -2 | [1] | +| 2026-07-08T00:02:00 | s3 | s3 | +| 2026-07-08T00:03:00 | | | +| 2026-07-08T00:04:00 | | s5 | +| 2026-07-08T00:05:00 | 6 | | ++---------------------+-----+-----+ + +SELECT + ts, + j.d.e AS d_e +FROM t_legacy_json2_variant_payload_object_d +ORDER BY ts; + ++---------------------+-------+ +| ts | d_e | ++---------------------+-------+ +| 2026-07-08T01:00:00 | true | +| 2026-07-08T01:01:00 | false | ++---------------------+-------+ + +-- JSON2 projection rendering is a user-visible compatibility contract. +SELECT ts, j.d AS d +FROM t_legacy_json2_variant_payload +ORDER BY ts; + ++---------------------+--------------------+ +| ts | d | ++---------------------+--------------------+ +| 2026-07-08T00:00:00 | [{"e":{"f":0.1}}] | +| 2026-07-08T00:01:00 | [{"e":{"f":0.2}}] | +| 2026-07-08T00:02:00 | [{"e":{"g":-0.3}}] | +| 2026-07-08T00:03:00 | | +| 2026-07-08T00:04:00 | {"e":true} | +| 2026-07-08T00:05:00 | | ++---------------------+--------------------+ + +SHOW CREATE TABLE t_legacy_json2_variant_payload; + ++--------------------------------+---------------------------------------------------------------+ +| Table | Create Table | ++--------------------------------+---------------------------------------------------------------+ +| t_legacy_json2_variant_payload | CREATE TABLE IF NOT EXISTS "t_legacy_json2_variant_payload" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2 NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++--------------------------------+---------------------------------------------------------------+ diff --git a/tests/compatibility/cases/legacy_json2_variant_payload/verify.sql b/tests/compatibility/cases/legacy_json2_variant_payload/verify.sql new file mode 100644 index 00000000000..a3c2cab1b92 --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_variant_payload/verify.sql @@ -0,0 +1,16 @@ +SELECT ts, j.a.b AS a_b, j.c AS c +FROM t_legacy_json2_variant_payload +ORDER BY ts; + +SELECT + ts, + j.d.e AS d_e +FROM t_legacy_json2_variant_payload_object_d +ORDER BY ts; + +-- JSON2 projection rendering is a user-visible compatibility contract. +SELECT ts, j.d AS d +FROM t_legacy_json2_variant_payload +ORDER BY ts; + +SHOW CREATE TABLE t_legacy_json2_variant_payload; From d667728fdeb5a55e63665f070453b4649801af60 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 13 Jul 2026 11:27:06 +0800 Subject: [PATCH 09/48] chore!: update promql-parser to v0.10.0, remove `holt_winters` (#8457) * chore: update promql parser and fix compile Signed-off-by: shuiyisong * fix: sqlness Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- Cargo.lock | 8 +-- Cargo.toml | 2 +- src/query/src/parser.rs | 49 ++++++++++++------- src/query/src/promql/planner.rs | 41 ++++++++++++++++ .../standalone/common/promql/functions.result | 10 ---- .../standalone/common/promql/functions.sql | 4 -- 6 files changed, 76 insertions(+), 38 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e00e30e28a6..75ba5ce410b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11038,9 +11038,9 @@ dependencies = [ [[package]] name = "promql-parser" -version = "0.7.3" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f514ab521bb2b4ce26aca29bf133498b2a2651b785d03a262bbf5772329c279f" +checksum = "9e8c90c956be4237a93be6f2298302b56974ef862ab11e0f98869d1f27c69f88" dependencies = [ "cfgrammar", "chrono", @@ -11183,7 +11183,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -11196,7 +11196,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", diff --git a/Cargo.toml b/Cargo.toml index 5ee60c450e8..fd878c7c0a1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -198,7 +198,7 @@ paste = "1.0" pin-project = "1.0" pretty_assertions = "1.4.0" prometheus = { version = "0.14", features = ["process"] } -promql-parser = { version = "0.7.3", features = ["ser"] } +promql-parser = { version = "0.10.0", features = ["ser"] } prost = { version = "0.14", features = ["no-recursion-limit"] } prost-types = "0.14" raft-engine = { version = "0.4.1", default-features = false } diff --git a/src/query/src/parser.rs b/src/query/src/parser.rs index 53f24557c73..87b827e0cc1 100644 --- a/src/query/src/parser.rs +++ b/src/query/src/parser.rs @@ -52,45 +52,37 @@ pub enum QueryStatement { } impl QueryStatement { - pub fn post_process(&self, params: HashMap) -> Result { + pub fn post_process(self, params: HashMap) -> Result { match self { QueryStatement::Sql(_) => UnimplementedSnafu { operation: "sql post process", } .fail(), - QueryStatement::Promql(eval_stmt, alias) => { + QueryStatement::Promql(mut eval_stmt, alias) => { let node_name = match params.get("name") { Some(name) => name.as_str(), None => "", }; - let extension_node = Self::create_extension_node(node_name, &eval_stmt.expr); - Ok(QueryStatement::Promql( - EvalStmt { - expr: Extension(extension_node.unwrap()), - start: eval_stmt.start, - end: eval_stmt.end, - interval: eval_stmt.interval, - lookback_delta: eval_stmt.lookback_delta, - }, - alias.clone(), - )) + let extension_node = Self::create_extension_node(node_name, eval_stmt.expr); + eval_stmt.expr = Extension(extension_node.unwrap()); + Ok(QueryStatement::Promql(eval_stmt, alias)) } } } - fn create_extension_node(node_name: &str, expr: &Expr) -> Option { + fn create_extension_node(node_name: &str, expr: Expr) -> Option { match node_name { ANALYZE_NODE_NAME => Some(NodeExtension { - expr: Arc::new(AnalyzeExpr { expr: expr.clone() }), + expr: Arc::new(AnalyzeExpr { expr }), }), ANALYZE_VERBOSE_NODE_NAME => Some(NodeExtension { - expr: Arc::new(AnalyzeVerboseExpr { expr: expr.clone() }), + expr: Arc::new(AnalyzeVerboseExpr { expr }), }), EXPLAIN_NODE_NAME => Some(NodeExtension { - expr: Arc::new(ExplainExpr { expr: expr.clone() }), + expr: Arc::new(ExplainExpr { expr }), }), EXPLAIN_VERBOSE_NODE_NAME => Some(NodeExtension { - expr: Arc::new(ExplainVerboseExpr { expr: expr.clone() }), + expr: Arc::new(ExplainVerboseExpr { expr }), }), _ => None, } @@ -204,9 +196,10 @@ impl QueryLanguageParser { } pub(crate) fn apply_alias_extension(mut eval_stmt: EvalStmt, alias: &str) -> EvalStmt { + let expr = eval_stmt.expr; eval_stmt.expr = Extension(NodeExtension { expr: Arc::new(AliasExpr { - expr: eval_stmt.expr.clone(), + expr, alias: alias.to_string(), }), }); @@ -263,6 +256,14 @@ macro_rules! define_node_ast_extension { fn children(&self) -> &[Expr] { std::slice::from_ref(&self.expr) } + + fn with_new_children(&self, children: Vec) -> Arc { + let mut iter = children.into_iter(); + match (iter.next(), iter.next()) { + (Some(expr), None) => Arc::new($name_expr { expr }), + _ => Arc::new(self.clone()), + } + } } #[allow(rustdoc::broken_intra_doc_links)] @@ -313,6 +314,16 @@ impl ExtensionExpr for AliasExpr { fn children(&self) -> &[Expr] { std::slice::from_ref(&self.expr) } + fn with_new_children(&self, children: Vec) -> Arc { + let mut iter = children.into_iter(); + match (iter.next(), iter.next()) { + (Some(expr), None) => Arc::new(Self { + expr, + alias: self.alias.clone(), + }), + _ => Arc::new(self.clone()), + } + } } #[derive(Debug, Clone)] pub struct Alias { diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index c2cd6316ada..61913156cf8 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -229,6 +229,8 @@ impl IslandExpr { !modifier.return_bool && modifier.matching.is_none() && matches!(modifier.card, VectorMatchCardinality::OneToOne) + && modifier.fill_values.lhs.is_none() + && modifier.fill_values.rhs.is_none() }) => { let lhs = Self::try_new(lhs, env)?; @@ -1093,6 +1095,18 @@ impl PromPlanner { query_engine_state: &QueryEngineState, binary_expr: &PromBinaryExpr, ) -> Result { + // promql-parser accepts fill modifiers, but Greptime does not implement the + // required outer joins and missing-value substitution. Reject them before the + // binary-island fast path so they cannot silently behave like normal inner joins. + if let Some(modifier) = &binary_expr.modifier { + ensure!( + modifier.fill_values.lhs.is_none() && modifier.fill_values.rhs.is_none(), + UnsupportedExprSnafu { + name: "PromQL fill modifiers" + } + ); + } + if let Some(plan) = self.try_plan_binary_island(binary_expr).await? { return Ok(plan); } @@ -5542,6 +5556,33 @@ mod test { ); } + #[tokio::test] + async fn reject_binary_fill_modifiers() { + let state = build_query_engine_state(); + + for query in [ + "some_metric + fill(0) some_alt_metric", + "some_metric + fill_left(0) some_alt_metric", + "some_metric + fill_right(0) some_alt_metric", + "(some_metric + fill(0) some_alt_metric) + some_metric", + ] { + let eval_stmt = build_eval_stmt(query); + let table_provider = build_test_table_provider(&[], 0, 0).await; + let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state) + .await + .unwrap_err(); + + assert!( + matches!( + &err, + crate::promql::error::Error::UnsupportedExpr { name, .. } + if name == "PromQL fill modifiers" + ), + "{err}" + ); + } + } + #[tokio::test] async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() { for query in [ diff --git a/tests/cases/standalone/common/promql/functions.result b/tests/cases/standalone/common/promql/functions.result index 02ff599ced9..325fbccd17d 100644 --- a/tests/cases/standalone/common/promql/functions.result +++ b/tests/cases/standalone/common/promql/functions.result @@ -62,16 +62,6 @@ tql eval (10, 10, '1s') double_exponential_smoothing(prom_series[10s], 0.4 + 0.1 | 1970-01-01T00:00:10 | 47.0806953125 | p | +---------------------+------------------------------------------------------------------------------------------+------+ --- holt_winters (backward compatibility) --- SQLNESS SORT_RESULT 3 1 -tql eval (10, 10, '1s') holt_winters(prom_series[10s], 0.4 + 0.1, 0.1); - -+---------------------+------------------------------------------------------------------------------------------+------+ -| ts | prom_double_exponential_smoothing(ts_range,val,Float64(0.4) + Float64(0.1),Float64(0.1)) | host | -+---------------------+------------------------------------------------------------------------------------------+------+ -| 1970-01-01T00:00:10 | 47.0806953125 | p | -+---------------------+------------------------------------------------------------------------------------------+------+ - DROP TABLE prom_series; Affected Rows: 0 diff --git a/tests/cases/standalone/common/promql/functions.sql b/tests/cases/standalone/common/promql/functions.sql index 6657fc49e81..0a257f29da8 100644 --- a/tests/cases/standalone/common/promql/functions.sql +++ b/tests/cases/standalone/common/promql/functions.sql @@ -34,10 +34,6 @@ tql eval (3, 3, '1s') predict_linear(prom_series[3s], 40 + 2); -- SQLNESS SORT_RESULT 3 1 tql eval (10, 10, '1s') double_exponential_smoothing(prom_series[10s], 0.4 + 0.1, 0.1); --- holt_winters (backward compatibility) --- SQLNESS SORT_RESULT 3 1 -tql eval (10, 10, '1s') holt_winters(prom_series[10s], 0.4 + 0.1, 0.1); - DROP TABLE prom_series; CREATE TABLE From 29b345fe2d07a04937bf5e64d37ea6c69ab4dbc0 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 13 Jul 2026 11:56:34 +0800 Subject: [PATCH 10/48] ci: increase Kafka resource requests (#8482) * ci: increase Kafka resource requests Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus Signed-off-by: WenyXu * ci: enable rskafka debug logs in remote WAL fuzz Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus Signed-off-by: WenyXu --------- Signed-off-by: WenyXu Co-authored-by: Sisyphus --- .../with-remote-wal-logical-prune.yaml | 1 + .github/actions/setup-greptimedb-cluster/with-remote-wal.yaml | 1 + .github/actions/setup-kafka-cluster/action.yml | 4 ++-- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/actions/setup-greptimedb-cluster/with-remote-wal-logical-prune.yaml b/.github/actions/setup-greptimedb-cluster/with-remote-wal-logical-prune.yaml index 60c23def6e2..3b25e80a8f2 100644 --- a/.github/actions/setup-greptimedb-cluster/with-remote-wal-logical-prune.yaml +++ b/.github/actions/setup-greptimedb-cluster/with-remote-wal-logical-prune.yaml @@ -3,6 +3,7 @@ logging: format: "json" filters: - log_store=debug + - rskafka=debug meta: configData: |- [runtime] diff --git a/.github/actions/setup-greptimedb-cluster/with-remote-wal.yaml b/.github/actions/setup-greptimedb-cluster/with-remote-wal.yaml index ec595b9e9b8..6a15b25b422 100644 --- a/.github/actions/setup-greptimedb-cluster/with-remote-wal.yaml +++ b/.github/actions/setup-greptimedb-cluster/with-remote-wal.yaml @@ -3,6 +3,7 @@ logging: format: "json" filters: - log_store=debug + - rskafka=debug meta: configData: |- [runtime] diff --git a/.github/actions/setup-kafka-cluster/action.yml b/.github/actions/setup-kafka-cluster/action.yml index fca420da709..3c30ca285e3 100644 --- a/.github/actions/setup-kafka-cluster/action.yml +++ b/.github/actions/setup-kafka-cluster/action.yml @@ -16,8 +16,8 @@ runs: helm upgrade \ --install kafka oci://registry-1.docker.io/bitnamicharts/kafka \ --set controller.replicaCount=${{ inputs.controller-replicas }} \ - --set controller.resources.requests.cpu=50m \ - --set controller.resources.requests.memory=128Mi \ + --set controller.resources.requests.cpu=500m \ + --set controller.resources.requests.memory=512Mi \ --set controller.resources.limits.cpu=2000m \ --set controller.resources.limits.memory=2Gi \ --set listeners.controller.protocol=PLAINTEXT \ From c70cbaa132d45f69f0ad5548e98e5eb208fa5ac0 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 13 Jul 2026 13:21:21 +0800 Subject: [PATCH 11/48] feat: more region lifecycle hooks (#8467) * feat: hook extension for region/file close/drop * feat: cover region lifecycle with open/close/drop hooks * docs(mito2): clarify on_region_opened runs in spawned task on open path Address review feedback: the module/trait docs claimed on_region_opened runs inline in the region worker loop. That holds for the create path but not for the open path, where it fires inside the spawned open task (common_runtime::spawn_global) after WAL replay and before registration. Correct both the Notes block and the trait method doc so hook authors don't assume worker-loop-thread affinity or strict ordering on open. Signed-off-by: Ning Sun * docs(mito2): fix region_hook inventory and lifecycle wording Address shuiyisong's review feedback: - Overview no longer hardcodes 'two methods'; the lifecycle bullet list now includes on_region_opened. - on_region_opened is described as firing after open/create succeeds but before registration (it runs before insert_region on both paths). - on_region_closed drops the inaccurate 'follower/catchup regions' exclusion: remove_region fires it for any role, which is consistent with on_region_opened firing for followers too. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- src/mito2/src/engine/close_test.rs | 43 +++++++++ src/mito2/src/engine/drop_test.rs | 76 ++++++++++++++++ src/mito2/src/engine/flush_test.rs | 59 +++++++++++- src/mito2/src/engine/open_test.rs | 35 +++++++ src/mito2/src/engine/region_hook.rs | 126 +++++++++++++++++++++++++- src/mito2/src/worker/handle_close.rs | 8 ++ src/mito2/src/worker/handle_create.rs | 6 ++ src/mito2/src/worker/handle_drop.rs | 47 ++++++++++ src/mito2/src/worker/handle_open.rs | 7 ++ 9 files changed, 399 insertions(+), 8 deletions(-) diff --git a/src/mito2/src/engine/close_test.rs b/src/mito2/src/engine/close_test.rs index 4c06583b0b8..0092634ec74 100644 --- a/src/mito2/src/engine/close_test.rs +++ b/src/mito2/src/engine/close_test.rs @@ -12,11 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; +use std::sync::atomic::Ordering; + +use common_base::Plugins; use store_api::region_engine::RegionEngine; use store_api::region_request::{RegionCloseRequest, RegionRequest}; use store_api::storage::RegionId; use crate::config::MitoConfig; +use crate::engine::flush_test::MockRegionHook; +use crate::engine::region_hook::RegionHookRef; use crate::test_util::{CreateRequestBuilder, TestEnv}; #[tokio::test] @@ -60,3 +66,40 @@ async fn test_engine_close_region_with_format(flat_format: bool) { .await .unwrap(); } + +#[tokio::test] +async fn test_region_hook_on_close() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + + let hook = Arc::new(MockRegionHook::new()); + let plugins = Plugins::new(); + plugins.insert(hook.clone() as RegionHookRef); + + let engine = env + .create_engine_with_plugins(MitoConfig::default(), plugins) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Sanity: no lifecycle events before closing besides the create open. + assert_eq!(hook.opened_count.load(Ordering::Relaxed), 1); + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0); + assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0); + + engine + .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .await + .unwrap(); + + // Closing fires on_region_closed exactly once. + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 1); + // Close must not be confused with drop. + assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0); + assert_eq!(hook.files_removed_count.load(Ordering::Relaxed), 0); +} diff --git a/src/mito2/src/engine/drop_test.rs b/src/mito2/src/engine/drop_test.rs index a34a5d11727..90afd66ba32 100644 --- a/src/mito2/src/engine/drop_test.rs +++ b/src/mito2/src/engine/drop_test.rs @@ -13,9 +13,11 @@ // limitations under the License. use std::sync::Arc; +use std::sync::atomic::Ordering; use std::time::Duration; use api::v1::Rows; +use common_base::Plugins; use common_meta::key::SchemaMetadataManager; use common_meta::kv_backend::KvBackendRef; use object_store::util::join_path; @@ -25,7 +27,9 @@ use store_api::storage::RegionId; use crate::config::MitoConfig; use crate::engine::MitoEngine; +use crate::engine::flush_test::MockRegionHook; use crate::engine::listener::DropListener; +use crate::engine::region_hook::RegionHookRef; use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows_for_key, flush_region, put_rows, rows_schema, }; @@ -302,3 +306,75 @@ async fn test_engine_drop_region_for_custom_store_with_format(flat_format: bool) .unwrap() ); } + +#[tokio::test] +async fn test_region_hook_on_drop() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("drop_hook").await; + + let hook = Arc::new(MockRegionHook::new()); + let plugins = Plugins::new(); + plugins.insert(hook.clone() as RegionHookRef); + + let engine = env + .create_engine_with_plugins(MitoConfig::default(), plugins) + .await; + + let region_id = RegionId::new(1, 1); + 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 request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Put some rows so the region has active data (kept in memtable; we + // intentionally do not flush so the drop GC worker finds no parquet files + // and removes the directory immediately, avoiding the multi-minute GC wait). + let rows = Rows { + schema: column_schemas, + rows: build_rows_for_key("a", 0, 2, 0), + }; + put_rows(&engine, region_id, rows).await; + + // Sanity: nothing fired yet besides the open from create. + assert_eq!(hook.opened_count.load(Ordering::Relaxed), 1); + assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0); + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0); + + // full drop (partial_drop = false) so the GC worker physically deletes the + // directory and fires on_region_files_removed. + engine + .handle_request( + region_id, + RegionRequest::Drop(RegionDropRequest { + fast_path: false, + force: false, + partial_drop: false, + }), + ) + .await + .unwrap(); + + // on_region_dropped fires inline while handling the drop request. + hook.wait_for_dropped().await; + assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 1); + // Drop must not be confused with close. + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0); + + // on_region_files_removed fires from the background GC worker once the + // directory is physically deleted. + hook.wait_for_files_removed().await; + assert_eq!(hook.files_removed_count.load(Ordering::Relaxed), 1); +} diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index 4870e8c2bbe..dc0063a7e63 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -895,24 +895,46 @@ async fn test_update_topic_latest_entry_id(factory: Option) { } #[derive(Debug)] -struct MockRegionHook { - sst_written_count: AtomicUsize, - manifest_updated_count: AtomicUsize, +pub(super) struct MockRegionHook { + pub(super) sst_written_count: AtomicUsize, + pub(super) manifest_updated_count: AtomicUsize, + pub(super) opened_count: AtomicUsize, + pub(super) closed_count: AtomicUsize, + pub(super) dropped_count: AtomicUsize, + pub(super) files_removed_count: AtomicUsize, notify: Notify, + dropped_notify: Notify, + files_removed_notify: Notify, } impl MockRegionHook { - fn new() -> Self { + pub(super) fn new() -> Self { Self { sst_written_count: AtomicUsize::new(0), manifest_updated_count: AtomicUsize::new(0), + opened_count: AtomicUsize::new(0), + closed_count: AtomicUsize::new(0), + dropped_count: AtomicUsize::new(0), + files_removed_count: AtomicUsize::new(0), notify: Notify::new(), + dropped_notify: Notify::new(), + files_removed_notify: Notify::new(), } } async fn wait_for_manifest_update(&self) { self.notify.notified().await; } + + /// Waits until `on_region_dropped` has fired at least once. + pub(super) async fn wait_for_dropped(&self) { + self.dropped_notify.notified().await; + } + + /// Waits until `on_region_files_removed` has fired at least once. + pub(super) async fn wait_for_files_removed(&self) { + self.files_removed_notify.notified().await; + } } #[async_trait] @@ -966,6 +988,35 @@ impl RegionHook for MockRegionHook { ); self.notify.notify_one(); } + + async fn on_region_opened(&self, region_id: RegionId, _region_metadata: &RegionMetadataRef) { + self.opened_count.fetch_add(1, Ordering::Relaxed); + common_telemetry::info!("MockRegionHook::on_region_opened: region={}", region_id); + } + + async fn on_region_closed(&self, region_id: RegionId, _region_metadata: &RegionMetadataRef) { + self.closed_count.fetch_add(1, Ordering::Relaxed); + common_telemetry::info!("MockRegionHook::on_region_closed: region={}", region_id); + } + + async fn on_region_dropped(&self, region_id: RegionId, _region_metadata: &RegionMetadataRef) { + self.dropped_count.fetch_add(1, Ordering::Relaxed); + common_telemetry::info!("MockRegionHook::on_region_dropped: region={}", region_id); + self.dropped_notify.notify_one(); + } + + async fn on_region_files_removed( + &self, + region_id: RegionId, + _region_metadata: &RegionMetadataRef, + ) { + self.files_removed_count.fetch_add(1, Ordering::Relaxed); + common_telemetry::info!( + "MockRegionHook::on_region_files_removed: region={}", + region_id + ); + self.files_removed_notify.notify_one(); + } } #[tokio::test] diff --git a/src/mito2/src/engine/open_test.rs b/src/mito2/src/engine/open_test.rs index cf4bfb1ff1d..ee0d9560093 100644 --- a/src/mito2/src/engine/open_test.rs +++ b/src/mito2/src/engine/open_test.rs @@ -32,6 +32,8 @@ use tokio::sync::oneshot; use crate::compaction::compactor::{OpenCompactionRegionRequest, open_compaction_region}; use crate::config::MitoConfig; +use crate::engine::flush_test::MockRegionHook; +use crate::engine::region_hook::RegionHookRef; use crate::error; use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef}; use crate::region::options::RegionOptions; @@ -805,3 +807,36 @@ async fn test_open_keeps_none_without_fetcher() { let meta = engine.get_region(region_id).unwrap().metadata(); assert!(meta.partition_expr.is_none()); } + +#[tokio::test] +async fn test_region_hook_on_open() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("open_hook").await; + + let hook = Arc::new(MockRegionHook::new()); + let plugins = Plugins::new(); + plugins.insert(hook.clone() as RegionHookRef); + + let engine = env + .create_engine_with_plugins(MitoConfig::default(), plugins) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Creating a region fires on_region_opened exactly once. + assert_eq!(hook.opened_count.load(Ordering::Relaxed), 1); + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0); + + // `reopen_region` closes then opens; the open must fire the hook again. + reopen_region(&engine, region_id, table_dir, false, Default::default()).await; + + assert_eq!(hook.opened_count.load(Ordering::Relaxed), 2); + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 1); + assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0); +} diff --git a/src/mito2/src/engine/region_hook.rs b/src/mito2/src/engine/region_hook.rs index d1b28f6ebae..c4b683ca831 100644 --- a/src/mito2/src/engine/region_hook.rs +++ b/src/mito2/src/engine/region_hook.rs @@ -16,7 +16,8 @@ //! //! ## Design //! -//! The [`RegionHook`] trait provides two methods with clear separation of concerns: +//! The [`RegionHook`] trait observes region activity through two categories of +//! callbacks — manifest/file observation and region lifecycle: //! //! - [`on_sst_files_written`]: Fires when mito2 physically writes SST **data files**. //! Provides per-file [`SstInfo`] + [`FileMeta`]; metadata richness varies by path @@ -26,6 +27,10 @@ //! Receives the full [`RegionMetaActionList`] so consumers can inspect what changed //! (file additions/removals, schema changes, truncation, partition expression changes, etc.). //! +//! - [`on_region_opened`] / [`on_region_closed`] / [`on_region_dropped`] / [`on_region_files_removed`]: +//! Region **lifecycle** callbacks for open, close, logical drop, and physical file removal. +//! See [Region lifecycle](#region-lifecycle) below. +//! //! Hook implementations are registered via the [`Plugins`](common_base::Plugins) system: //! ```ignore //! plugins.insert(Arc::new(MyHook) as RegionHookRef); @@ -56,7 +61,39 @@ //! //! The following paths do **not** trigger any hook: //! - Follower region sync / catchup (manifest read-only; followers don't author changes) -//! - GC / checkpoint / drop / remap (internal bookkeeping, not logical state changes) +//! - GC / checkpoint / remap (internal bookkeeping, not logical state changes) +//! +//! An explicit region **drop** does fire lifecycle hooks — see +//! [Region lifecycle](#region-lifecycle). +//! +//! ## Region lifecycle +//! +//! Beyond manifest/SST observation, the hook observes the high-level lifecycle of an +//! active region: +//! +//! | Event | Method | When | +//! |-------|--------|------| +//! | Open | [`on_region_opened`] | A create or open request registers the region as active (the counterpart to close/drop). Does not fire for the compactor's transient regions or the catch-up reopen. | +//! | Close | [`on_region_closed`] | A close request (or a close-after-flush) removes the region from the active set. Data files, manifest and WAL state are **preserved**; the region may be reopened. | +//! | Logical drop | [`on_region_dropped`] | A drop request has been handled: the region leaves the active set and its WAL entries are marked obsolete. Data files are **not yet deleted**. | +//! | Physical file removal | [`on_region_files_removed`] | The drop GC worker has deleted the region directory. Terminal file-lifecycle event. | +//! +//! Notes: +//! - `on_region_closed` / `on_region_dropped` run **inline in the region worker loop**, +//! so implementations must be fast (same contract as `on_manifest_updated`). +//! - `on_region_opened` runs inline in the worker loop on the **create** path, but on the +//! **open** path it fires inside the spawned open task (`common_runtime::spawn_global`), +//! i.e. concurrently with the worker loop — after WAL replay, before the region is +//! registered and its open request is acknowledged. Implementations must still be fast +//! and must not assume worker-loop-thread affinity or strict ordering against concurrent +//! requests to other regions. +//! - `on_region_files_removed` runs on the background drop GC task, outside the worker loop. +//! - When global GC is enabled and a normal table region is dropped with `partial_drop`, its +//! directory is left for global reclamation and `on_region_files_removed` is **not** fired +//! by the drop worker (observe it via the global GC path instead). +//! - Logical file removal (compaction, region edit, truncate) is already observable via +//! [`on_manifest_updated`] (`Edit.files_to_remove` / `Truncate` action); only the drop +//! worker's physical directory deletion needs a dedicated file hook. //! //! ## Invocation points //! @@ -77,11 +114,18 @@ //! //! ## Future work //! -//! A future `on_files_removed` hook may be added to observe file lifecycle end -//! (GC, drop, truncate, compaction removal). This is not yet implemented. +//! `on_region_files_removed` currently covers only the **drop** GC worker's physical +//! directory removal. A broader per-file `on_files_removed` hook covering compaction +//! removal, truncate, and the global GC reclamation path is not yet implemented +//! (though logical file removal is already observable via `on_manifest_updated`). +//! Role/leadership transitions (`on_region_role_changed`) are also not hooked. //! //! [`on_sst_files_written`]: RegionHook::on_sst_files_written //! [`on_manifest_updated`]: RegionHook::on_manifest_updated +//! [`on_region_opened`]: RegionHook::on_region_opened +//! [`on_region_closed`]: RegionHook::on_region_closed +//! [`on_region_dropped`]: RegionHook::on_region_dropped +//! [`on_region_files_removed`]: RegionHook::on_region_files_removed //! [`RegionManifestManager::update`]: crate::manifest::manager::RegionManifestManager::update //! [`ManifestContext::update_locked`]: crate::region::ManifestContext::update_locked //! [`ManifestContext::update_manifest`]: crate::region::ManifestContext::update_manifest @@ -228,6 +272,80 @@ pub trait RegionHook: Send + Sync + Debug { ) { let _ = (region_id, action_list, manifest_version); } + + /// Called once a region **open** or **create** succeeds, but **before** the + /// region is registered in the engine's active set (`insert_region` runs + /// immediately afterwards). + /// + /// Fires once when a region becomes active via a create or open request — + /// the natural counterpart to [`on_region_closed`] / [`on_region_dropped`]. + /// It does **not** fire for the compactor's transient compaction regions + /// (`open_compaction_region`), nor for the internal reopen performed during + /// follower catch-up / leadership promotion. + /// + /// On the **create** path it runs inline in the region worker loop; on the + /// **open** path it runs inside the spawned open task + /// (`common_runtime::spawn_global`), concurrently with the worker loop + /// (after WAL replay, before the region is registered/acknowledged). + /// Implementations must be fast and must **not** assume worker-loop-thread + /// affinity or strict ordering against concurrent requests to other regions. + /// + /// [`on_region_closed`]: RegionHook::on_region_closed + /// [`on_region_dropped`]: RegionHook::on_region_dropped + async fn on_region_opened(&self, region_id: RegionId, region_metadata: &RegionMetadataRef) { + let _ = (region_id, region_metadata); + } + + /// Called after a region is **closed** via a close request. + /// + /// The region is removed from the engine's active set, but its data files, + /// manifest, and WAL state are **preserved**; the region may be reopened + /// later. Fires once per successful close, after the region's background + /// tasks (flush/compaction) have been stopped. + /// + /// Fires for a region of **any** role (leader or follower) that is closed. + /// Does **not** fire when a region is dropped (see [`on_region_dropped`]). + /// + /// Runs inline in the region worker loop; implementations should be fast. + /// + /// [`on_region_dropped`]: RegionHook::on_region_dropped + async fn on_region_closed(&self, region_id: RegionId, region_metadata: &RegionMetadataRef) { + let _ = (region_id, region_metadata); + } + + /// Called after a region is **logically dropped** (a drop request has been + /// handled). + /// + /// The region is removed from the active set and its WAL entries are marked + /// obsolete. Its data files are **not yet deleted** — they are scheduled for + /// asynchronous removal by the GC worker. Observe physical deletion via + /// [`on_region_files_removed`]. + /// + /// Runs inline in the region worker loop; implementations should be fast. + /// + /// [`on_region_files_removed`]: RegionHook::on_region_files_removed + async fn on_region_dropped(&self, region_id: RegionId, region_metadata: &RegionMetadataRef) { + let _ = (region_id, region_metadata); + } + + /// Called after a dropped region's data files are **physically removed** by + /// the drop GC worker (the region directory has been deleted). + /// + /// This is the terminal event in a region's file lifecycle; no further + /// callbacks fire for this region id afterwards. Fires only when the drop + /// worker itself deletes the directory. When global GC is enabled and the + /// region is a normal table region dropped with `partial_drop`, the + /// directory is left for global reclamation and this hook is **not** fired + /// by the drop worker. + /// + /// Runs on a background task, outside the region worker loop. + async fn on_region_files_removed( + &self, + region_id: RegionId, + region_metadata: &RegionMetadataRef, + ) { + let _ = (region_id, region_metadata); + } } pub type RegionHookRef = Arc; diff --git a/src/mito2/src/worker/handle_close.rs b/src/mito2/src/worker/handle_close.rs index a616a0e32dc..b582c276927 100644 --- a/src/mito2/src/worker/handle_close.rs +++ b/src/mito2/src/worker/handle_close.rs @@ -89,5 +89,13 @@ impl RegionWorkerLoop { // clean index build status. self.index_build_scheduler.on_region_closed(region_id).await; self.region_count.dec(); + + // Notify the region hook that the region has been closed. The region is + // fully stopped and unregistered, but its files/manifest are preserved. + // Runs inline; the hook contract requires it to be fast. + if let Some(hook) = region.manifest_ctx.hook() { + let metadata = region.metadata(); + hook.on_region_closed(region_id, &metadata).await; + } } } diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index 6e43e74a0e4..bf254dc1e6b 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/src/mito2/src/worker/handle_create.rs @@ -86,6 +86,12 @@ impl RegionWorkerLoop { self.region_count.inc(); + // Notify the region hook that the region has been opened (created). + // Fires before registration; allocates nothing when no hook is registered. + if let Some(hook) = region.manifest_ctx.hook() { + hook.on_region_opened(region_id, ®ion.metadata()).await; + } + // Insert the MitoRegion into the RegionMap. self.regions.insert_region(region); diff --git a/src/mito2/src/worker/handle_drop.rs b/src/mito2/src/worker/handle_drop.rs index 00a96b0d166..3d111776b82 100644 --- a/src/mito2/src/worker/handle_drop.rs +++ b/src/mito2/src/worker/handle_drop.rs @@ -23,10 +23,12 @@ use object_store::util::join_path; use object_store::{EntryMode, ObjectStore}; use snafu::ResultExt; use store_api::logstore::LogStore; +use store_api::metadata::RegionMetadataRef; use store_api::region_request::{AffectedRows, PathType}; use store_api::storage::RegionId; use tokio::time::sleep; +use crate::engine::region_hook::RegionHookRef; use crate::error::{OpenDalSnafu, Result}; use crate::region::{RegionLeaderState, RegionMapRef}; use crate::worker::{DROPPING_MARKER_FILE, RegionWorkerLoop}; @@ -106,6 +108,22 @@ where self.region_count.dec(); + // Notify registered hooks that the region has been logically dropped, and + // prepare a payload for the background GC task to fire the terminal + // `on_region_files_removed`. When no hook is registered this allocates + // nothing — no metadata snapshot is taken and no payload is built. + let hook_payload = match region.manifest_ctx.hook() { + Some(hook) => { + let region_metadata = region.metadata(); + hook.on_region_dropped(region_id, ®ion_metadata).await; + Some(DropHookPayload { + hook, + metadata: region_metadata, + }) + } + None => None, + }; + let object_store = region.access_layer.object_store().clone(); let dropping_regions = self.dropping_regions.clone(); let listener = self.listener.clone(); @@ -122,6 +140,7 @@ where object_store, dropping_regions, partial_drop, + hook_payload.as_ref(), ) .await } else { @@ -135,6 +154,7 @@ where object_store, dropping_regions, gc_duration, + hook_payload.as_ref(), ) .await }; @@ -156,6 +176,19 @@ where } } +/// Carries the region hook and region metadata into the background GC task so +/// it can fire [`RegionHook::on_region_files_removed`] once the dropped region's +/// directory is physically deleted. +/// +/// Only constructed when a hook is registered; the task receives it as +/// `Option<&DropHookPayload>` so the no-hook path allocates nothing. +/// +/// [`RegionHook::on_region_files_removed`]: crate::engine::region_hook::RegionHook::on_region_files_removed +struct DropHookPayload { + hook: RegionHookRef, + metadata: RegionMetadataRef, +} + /// Background GC task to remove the entire region path once one of the following /// conditions is true: /// - It finds there is no parquet file left. @@ -172,6 +205,7 @@ async fn later_drop_task_without_global_gc( object_store: ObjectStore, dropping_regions: RegionMapRef, gc_duration: Duration, + hook_payload: Option<&DropHookPayload>, ) -> bool { remove_region_with_retry( region_id, @@ -180,6 +214,7 @@ async fn later_drop_task_without_global_gc( dropping_regions, Some(gc_duration), false, + hook_payload, ) .await } @@ -191,6 +226,7 @@ async fn remove_region_with_retry( dropping_regions: std::sync::Arc, gc_duration: Option, mut force: bool, + hook_payload: Option<&DropHookPayload>, ) -> bool { for _ in 0..MAX_RETRY_TIMES { let result = remove_region_dir_once(®ion_path, &object_store, force).await; @@ -204,6 +240,15 @@ async fn remove_region_with_retry( Ok(true) => { dropping_regions.remove_region(region_id); info!("Region {} is dropped, force: {}", region_path, force); + // The region directory has been physically deleted; fire the + // terminal file-lifecycle event. The partial-drop/global-GC path + // never reaches here (it does not delete the directory itself). + if let Some(hook_payload) = hook_payload { + hook_payload + .hook + .on_region_files_removed(region_id, &hook_payload.metadata) + .await; + } return true; } Ok(false) => (), @@ -230,6 +275,7 @@ async fn later_drop_task_with_global_gc( object_store: ObjectStore, dropping_regions: RegionMapRef, partial_drop: bool, + hook_payload: Option<&DropHookPayload>, ) -> bool { // For metadata regions or regions marked for full deletion (such as when dropping a table) // the region directory is forcefully removed immediately. @@ -243,6 +289,7 @@ async fn later_drop_task_with_global_gc( dropping_regions, None, true, + hook_payload, ) .await } else { diff --git a/src/mito2/src/worker/handle_open.rs b/src/mito2/src/worker/handle_open.rs index 3d3520b0cc8..48fef5ddd6f 100644 --- a/src/mito2/src/worker/handle_open.rs +++ b/src/mito2/src/worker/handle_open.rs @@ -149,6 +149,13 @@ impl RegionWorkerLoop { ); region_count.inc(); + // Notify the region hook that the region has been opened. + // Fires before registration; allocates nothing when no hook + // is registered. + if let Some(hook) = region.manifest_ctx.hook() { + hook.on_region_opened(region_id, ®ion.metadata()).await; + } + // Insert the Region into the RegionMap. regions.insert_region(region); From a12ad8df525dd25a7ae2050ad298b9a39e4ee87d Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 13 Jul 2026 14:35:54 +0800 Subject: [PATCH 12/48] refactor: suppress region hook notifications for staging-only manifest writes (#8471) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Only fire `on_manifest_updated` for writes to the live (normal) manifest directory. Writes to the staging directory (enter staging, operations during staging, the intermediate apply-staging `RegionEdit`) are suppressed — their effects accumulate and are delivered in a single notification when `exit_staging_on_success` promotes all staged actions to the live manifest. This simplifies hook implementations that only care about "live" state, since they no longer receive intermediate staging notifications they cannot act on, and no file-list information is lost (the promote step carries all accumulated files from staging operations). Key changes: - `PendingManifestHook` gains an `is_staging` field; `fire()` is a no-op when `is_staging` is `true` - `ManifestContext::update_locked` threads its `is_staging` parameter through - Updated module-level docs and coverage table Signed-off-by: Ning Sun --- src/mito2/src/engine/region_hook.rs | 62 +++++++++++++++++++++++------ src/mito2/src/region.rs | 1 + 2 files changed, 50 insertions(+), 13 deletions(-) diff --git a/src/mito2/src/engine/region_hook.rs b/src/mito2/src/engine/region_hook.rs index c4b683ca831..9c5a7dfc95d 100644 --- a/src/mito2/src/engine/region_hook.rs +++ b/src/mito2/src/engine/region_hook.rs @@ -23,9 +23,13 @@ //! Provides per-file [`SstInfo`] + [`FileMeta`]; metadata richness varies by path //! (see [`SstFileInfo`] and the coverage footnote). //! -//! - [`on_manifest_updated`]: Fires after **any** manifest write is successfully committed. -//! Receives the full [`RegionMetaActionList`] so consumers can inspect what changed -//! (file additions/removals, schema changes, truncation, partition expression changes, etc.). +//! - [`on_manifest_updated`]: Fires after a manifest write is committed to the **live** +//! (normal) manifest directory. Writes to the staging directory (enter staging, +//! operations during staging, the intermediate apply-staging edit) are suppressed — +//! their effects are accumulated and delivered in a single notification when the +//! staged actions are promoted to the live manifest. Receives the full +//! [`RegionMetaActionList`] so consumers can inspect what changed (file additions/ +//! removals, schema changes, truncation, partition expression changes, etc.). //! //! - [`on_region_opened`] / [`on_region_closed`] / [`on_region_dropped`] / [`on_region_files_removed`]: //! Region **lifecycle** callbacks for open, close, logical drop, and physical file removal. @@ -38,6 +42,12 @@ //! //! ## Coverage //! +//! Only manifest writes to the **normal** (live) manifest directory trigger +//! `on_manifest_updated`. Writes to the staging manifest directory (operations +//! that happen while the region is in staging mode) are intentionally suppressed: +//! their effects are accumulated and delivered in a single notification when +//! the staged actions are promoted to the live manifest via `exit_staging_on_success`. +//! //! | Scenario | `on_sst_files_written` | `on_manifest_updated` | //! |------------------------------|:----------------------:|:---------------------:| //! | Flush (memtable → SST) | ✅ Yes | ✅ Yes | @@ -45,19 +55,21 @@ //! | Remote compaction | ✅ (compactor node) ¹ | ✅ (compactor node) ¹ | //! | RegionEdit / bulk ingestion | ❌ (files pre-written) | ✅ Yes | //! | Copy region | ❌ (object-store copy) | ✅ Yes | -//! | Apply staging | ❌ (delegates to edit) | ✅ Yes ² | +//! | Apply staging (promote) | ❌ (delegates to edit) | ✅ Yes ² | //! | Alter (schema change) | ❌ (no SST files) | ✅ Yes | //! | Truncate | ❌ (removes files) | ✅ Yes | -//! | Enter staging | ❌ (no SST files) | ✅ Yes | +//! | Enter staging | ❌ (no SST files) | ❌ (staging dir) | +//! | Operations during staging | N/A | ❌ (staging dir) | //! | Async index build | ❌ (index files only) | ✅ Yes | //! //! ¹ Remote compaction runs on a dedicated compactor node via `open_compaction_region()`; //! pass plugins via `OpenCompactionRegionRequest` to enable hooks there. `sst_infos` is //! `#[serde(skip)]` over the wire, so each [`SstInfo`] is rebuilt from [`FileMeta`] with //! empty footer/index — see [`SstFileInfo`] for field-level detail. -//! ² Apply staging fires `on_manifest_updated` twice: once when the staging SST files are -//! committed via `RegionEdit`, and once when `exit_staging_on_success` merges all staged -//! manifest actions into the live manifest. +//! ² Apply staging fires `on_manifest_updated` once when `exit_staging_on_success` promotes +//! all staged manifest actions (including the SST file additions) into the live manifest. +//! The intermediate staging `RegionEdit` is written to the staging directory and does not +//! fire the hook — its file list is included in the promote notification. //! //! The following paths do **not** trigger any hook: //! - Follower region sync / catchup (manifest read-only; followers don't author changes) @@ -147,6 +159,13 @@ use crate::sst::parquet::SstInfo; /// /// Must be [`fire`](Self::fire)d **after** the manifest write lock is released /// (the hook may read the manifest). `#[must_use]` so a forgotten receipt warns. +/// +/// ## Staging suppression +/// +/// When `is_staging` is `true`, [`fire`](Self::fire) is a no-op — the write went to +/// the staging manifest directory. The hook only observes writes to the live (normal) +/// manifest directory. Staging actions are accumulated and delivered in a single +/// notification when `exit_staging_on_success` promotes them (`is_staging = false`). #[must_use = "the region hook must be fired after releasing the manifest write lock"] pub(crate) struct PendingManifestHook { region_id: RegionId, @@ -154,6 +173,9 @@ pub(crate) struct PendingManifestHook { action_list: Option, version: ManifestVersion, hook: Option, + /// Whether the manifest write went to the staging directory. + /// When `true`, `fire()` is suppressed — the hook only observes live manifest writes. + is_staging: bool, } impl PendingManifestHook { @@ -162,12 +184,14 @@ impl PendingManifestHook { action_list: Option, version: ManifestVersion, hook: Option, + is_staging: bool, ) -> Self { Self { region_id, action_list, version, hook, + is_staging, } } @@ -176,9 +200,13 @@ impl PendingManifestHook { self.version } - /// Fires the hook if one is registered. Safe to call unconditionally: it is - /// a no-op when no hook is registered. + /// Fires the hook if one is registered, **unless** the write went to the staging + /// manifest directory (`is_staging = true`). Safe to call unconditionally: + /// it is a no-op when no hook is registered or when the write is staging-only. pub(crate) async fn fire(self) { + if self.is_staging { + return; + } if let (Some(hook), Some(action_list)) = (self.hook, self.action_list) { hook.on_manifest_updated(self.region_id, &action_list, self.version) .await; @@ -190,6 +218,9 @@ impl PendingManifestHook { /// keeps `self`'s actions followed by `other`'s, and the *later* manifest /// version wins. Used when a sequence of writes (e.g. staging-exit followed /// by metadata backfill) should notify the hook exactly once. + /// + /// `is_staging` is `true` only if **both** sides are staging; if either side + /// is a live write, the merged result is also live. pub(crate) fn merge(self, other: PendingManifestHook) -> PendingManifestHook { debug_assert_eq!( self.region_id, other.region_id, @@ -208,6 +239,7 @@ impl PendingManifestHook { }, version: self.version.max(other.version), hook: self.hook.or(other.hook), + is_staging: self.is_staging && other.is_staging, } } } @@ -255,10 +287,14 @@ pub trait RegionHook: Send + Sync + Debug { let _ = (region_id, region_metadata, files); } - /// Called after the region manifest is successfully committed. + /// Called after the region manifest is successfully committed to the **live** + /// (normal) manifest directory. /// - /// Fires for **all** manifest write paths: flush, compaction, region edit, - /// copy region, alter, truncate, enter staging, index build, etc. + /// Fires for: flush, compaction, region edit, copy region, alter, truncate, + /// async index build, and apply-staging promote. Does **not** fire for writes + /// to the staging manifest directory (enter staging, operations during staging, + /// the intermediate apply-staging edit) — those are suppressed because their + /// effects are accumulated and delivered in a single promote notification. /// /// Does **not** fire for: /// - Manifest reads / follower sync (no write) diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 1e173ab62d2..c566dd47c04 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -1236,6 +1236,7 @@ impl ManifestContext { action_list_for_hook, version, self.hook.clone(), + is_staging, )) } From 758664eefab604ccbf3130ec40e7c38bfbcaede4 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 13 Jul 2026 18:19:57 +0800 Subject: [PATCH 13/48] fix: Use prepared file locations for CSV strict headers integration test (#8493) fix: normalize Windows COPY test paths Signed-off-by: Codex --- tests-integration/src/tests/instance_test.rs | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 628f97de96a..8572e9ce5b2 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -1098,6 +1098,10 @@ async fn test_execute_copy_from_csv_strict_headers(instance: Arc Date: Mon, 13 Jul 2026 20:32:46 +0800 Subject: [PATCH 14/48] feat: prepare soft-drop WAL retirement (#8475) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: flush soft-dropped regions on close Signed-off-by: Lei, HUANG * feat(mito2): handle flush-on-close race with concurrent in-flight flush When a region close with `flush_on_close: true` races with an already-running flush, pass the actual close request (including the flush_on_close flag) to the DDL handler instead of a default request so the pending flush is correctly awaited. Files: `src/mito2/src/worker/handle_close.rs` Also adds a test verifying that closing with flush-on-close while a flush is in progress still persists all written data correctly. Files: `src/mito2/src/engine/close_test.rs` Signed-off-by: Lei, HUANG * feat: support full WAL retirement Signed-off-by: Lei, HUANG * fix: complete close request migration Signed-off-by: Lei, HUANG * fix: finish close request callsites Signed-off-by: Lei, HUANG * refactor: guard Kafka provider setup behind index collector check Move Kafka provider initialization and `get_or_insert` inside the existing `if let Some(collector)` block so these operations are skipped when no global index collector is configured. Affected file: - `src/log-store/src/kafka/log_store.rs` Signed-off-by: Lei, HUANG * chore: avoid to_vec Signed-off-by: Lei, HUANG * refactor: replace imperative close-region loop with functional combinators Transform the region close dispatch in `DropTableExecutor` from mutable `Vec` and `push` loops to iterator chains with `join_all`, improving idiomatic Rust style and readability. - `src/common/meta/src/ddl/drop_table/executor.rs` — rewired datanode region-close logic to use `peers.map()` and nested `join_all`, moving `node_manager.datanode()` inside the closure to align with the new structure Signed-off-by: Lei, HUANG * fix: decouple Kafka client from WAL checkpoint Signed-off-by: Lei, HUANG * fix: merge Kafka WAL index checkpoints Signed-off-by: Lei, HUANG * fix: delegate Kafka WAL retirement to metasrv Signed-off-by: Lei, HUANG * chore: rebase main and resolve conflicts Signed-off-by: Lei, HUANG * fix: license header Signed-off-by: Lei, HUANG * chore: bump proto to commits on main Signed-off-by: Lei, HUANG * refactor: remove Kafka obsolete-all index changes Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/common/meta/src/ddl/drop_table.rs | 1 + .../meta/src/ddl/drop_table/executor.rs | 77 +++++--- src/common/meta/src/ddl/tests/drop_table.rs | 10 +- src/datanode/src/alive_keeper.rs | 2 +- src/datanode/src/heartbeat/handler.rs | 5 +- .../src/heartbeat/handler/close_region.rs | 6 +- .../src/heartbeat/handler/open_region.rs | 10 +- src/datanode/src/region_server.rs | 10 +- src/log-store/src/kafka/log_store.rs | 4 + src/log-store/src/noop/log_store.rs | 4 + src/log-store/src/raft_engine/log_store.rs | 38 ++++ src/metric-engine/src/engine.rs | 11 +- src/metric-engine/src/engine/close.rs | 18 +- src/metric-engine/src/engine/open.rs | 2 +- src/metric-engine/src/engine/sync/region.rs | 2 +- src/mito2/src/engine/catchup_test.rs | 5 +- src/mito2/src/engine/close_test.rs | 186 +++++++++++++++++- src/mito2/src/engine/create_test.rs | 5 +- src/mito2/src/engine/edit_region_test.rs | 5 +- src/mito2/src/engine/flush_test.rs | 10 +- src/mito2/src/engine/open_test.rs | 30 ++- src/mito2/src/engine/skip_wal_test.rs | 40 +++- src/mito2/src/flush.rs | 2 +- src/mito2/src/test_util.rs | 5 +- src/mito2/src/wal.rs | 9 + src/mito2/src/wal/raw_entry_reader.rs | 8 + src/mito2/src/worker.rs | 5 +- src/mito2/src/worker/handle_close.rs | 10 +- src/store-api/src/logstore.rs | 8 + src/store-api/src/region_request.rs | 15 +- 32 files changed, 457 insertions(+), 90 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 75ba5ce410b..7a9ec59760f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5938,7 +5938,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=4a69c9e75240a76afeb0796798085813669288f8#4a69c9e75240a76afeb0796798085813669288f8" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=6fc4d88851c5ad16846698e4dff6e74c1f7a48ba#6fc4d88851c5ad16846698e4dff6e74c1f7a48ba" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", diff --git a/Cargo.toml b/Cargo.toml index fd878c7c0a1..5fecc43fa5b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -158,7 +158,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "4a69c9e75240a76afeb0796798085813669288f8" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "6fc4d88851c5ad16846698e4dff6e74c1f7a48ba" } hex = "0.4" http = "1" humantime = "2.1" diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs index f5b5bc7a22a..8747bb111d2 100644 --- a/src/common/meta/src/ddl/drop_table.rs +++ b/src/common/meta/src/ddl/drop_table.rs @@ -179,6 +179,7 @@ impl DropTableProcedure { &self.context.node_manager, &self.context.leader_region_registry, &self.data.physical_region_routes, + true, ) .await?; self.context diff --git a/src/common/meta/src/ddl/drop_table/executor.rs b/src/common/meta/src/ddl/drop_table/executor.rs index 1e85a64ef60..74b2b2482b9 100644 --- a/src/common/meta/src/ddl/drop_table/executor.rs +++ b/src/common/meta/src/ddl/drop_table/executor.rs @@ -313,6 +313,7 @@ impl DropTableExecutor { }), body: Some(region_request::Body::Close(PbCloseRegionRequest { region_id: region_id.as_u64(), + flush_on_close: false, })), }; @@ -348,11 +349,13 @@ impl DropTableExecutor { } /// Closes all table regions on datanodes without deleting region files or metadata tombstones. + /// When `flush_leaders_on_close` is set, only leader regions are flushed before close. pub async fn on_close_regions( &self, node_manager: &NodeManagerRef, leader_region_registry: &LeaderRegionRegistryRef, region_routes: &[RegionRoute], + flush_leaders_on_close: bool, ) -> Result<()> { let table_id = self.table_id; let mut seen_peer_ids = HashSet::new(); @@ -360,39 +363,57 @@ impl DropTableExecutor { .into_iter() .chain(find_followers(region_routes)) .filter(|peer| seen_peer_ids.insert(peer.id)); - let mut close_region_tasks = Vec::new(); - - for datanode in peers { - let requester = node_manager.datanode(&datanode).await; + let close_region_tasks = peers.map(|datanode| { let region_ids = find_leader_regions(region_routes, &datanode) .into_iter() - .chain(find_follower_regions(region_routes, &datanode)) - .map(|region_number| RegionId::new(table_id, region_number)); + .map(|region_number| { + ( + RegionId::new(table_id, region_number), + flush_leaders_on_close, + ) + }) + .chain( + find_follower_regions(region_routes, &datanode) + .into_iter() + .map(|region_number| (RegionId::new(table_id, region_number), false)), + ) + .collect::>(); - for region_id in region_ids { - debug!("Closing region {region_id} on Datanode {datanode:?}"); - let request = RegionRequest { - header: Some(RegionRequestHeader { - tracing_context: TracingContext::from_current_span().to_w3c(), - ..Default::default() - }), - body: Some(region_request::Body::Close(PbCloseRegionRequest { - region_id: region_id.as_u64(), - })), - }; + async move { + let requester = node_manager.datanode(&datanode).await; + let close_region_tasks = + region_ids.into_iter().map(|(region_id, flush_on_close)| { + debug!("Closing region {region_id} on Datanode {datanode:?}"); + let request = RegionRequest { + header: Some(RegionRequestHeader { + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::Close(PbCloseRegionRequest { + region_id: region_id.as_u64(), + flush_on_close, + })), + }; - let datanode = datanode.clone(); - let requester = requester.clone(); - close_region_tasks.push(async move { - if let Err(err) = requester.handle(request).await - && err.status_code() != StatusCode::RegionNotFound - { - return Err(add_peer_context_if_needed(datanode)(err)); - } - Ok(()) - }); + let datanode = datanode.clone(); + let requester = requester.clone(); + async move { + if let Err(err) = requester.handle(request).await + && err.status_code() != StatusCode::RegionNotFound + { + return Err(add_peer_context_if_needed(datanode)(err)); + } + Ok(()) + } + }); + + join_all(close_region_tasks) + .await + .into_iter() + .collect::>>()?; + Ok(()) } - } + }); join_all(close_region_tasks) .await diff --git a/src/common/meta/src/ddl/tests/drop_table.rs b/src/common/meta/src/ddl/tests/drop_table.rs index fd41dd5126d..5c90dcbd79b 100644 --- a/src/common/meta/src/ddl/tests/drop_table.rs +++ b/src/common/meta/src/ddl/tests/drop_table.rs @@ -293,16 +293,16 @@ async fn test_soft_drop_closes_regions_and_keeps_tombstone() { let Some(region_request::Body::Close(req)) = request.body else { unreachable!(); }; - requests.push((peer.id, req.region_id)); + requests.push((peer.id, req.region_id, req.flush_on_close)); } requests.sort_unstable(); assert_eq!( requests, vec![ - (1, RegionId::new(table_id, 1).as_u64()), - (1, RegionId::new(table_id, 2).as_u64()), - (2, RegionId::new(table_id, 1).as_u64()), - (2, RegionId::new(table_id, 2).as_u64()), + (1, RegionId::new(table_id, 1).as_u64(), true), + (1, RegionId::new(table_id, 2).as_u64(), false), + (2, RegionId::new(table_id, 1).as_u64(), false), + (2, RegionId::new(table_id, 2).as_u64(), true), ] ); assert!(rx.try_recv().is_err()); diff --git a/src/datanode/src/alive_keeper.rs b/src/datanode/src/alive_keeper.rs index dbf99fdb284..dbfb699036c 100644 --- a/src/datanode/src/alive_keeper.rs +++ b/src/datanode/src/alive_keeper.rs @@ -153,7 +153,7 @@ impl RegionAliveKeeper { async fn close_staled_region(&self, region_id: RegionId) { info!("Closing staled region: {region_id}"); - let request = RegionRequest::Close(RegionCloseRequest {}); + let request = RegionRequest::Close(RegionCloseRequest::default()); if let Err(e) = self.region_server.handle_request(region_id, request).await && e.status_code() != StatusCode::RegionNotFound { diff --git a/src/datanode/src/heartbeat/handler.rs b/src/datanode/src/heartbeat/handler.rs index 79e0baaef36..71797592c3b 100644 --- a/src/datanode/src/heartbeat/handler.rs +++ b/src/datanode/src/heartbeat/handler.rs @@ -522,7 +522,10 @@ mod tests { .unwrap(); region_server - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); let mut heartbeat_env = HeartbeatResponseTestEnv::new(); diff --git a/src/datanode/src/heartbeat/handler/close_region.rs b/src/datanode/src/heartbeat/handler/close_region.rs index 581c35f3d8f..930c3f11525 100644 --- a/src/datanode/src/heartbeat/handler/close_region.rs +++ b/src/datanode/src/heartbeat/handler/close_region.rs @@ -40,8 +40,10 @@ impl InstructionHandler for CloseRegionsHandler { .collect::>(); let futs = region_ids.iter().map(|region_id| { - ctx.region_server - .handle_request(*region_id, RegionRequest::Close(RegionCloseRequest {})) + ctx.region_server.handle_request( + *region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) }); let results = join_all(futs).await; diff --git a/src/datanode/src/heartbeat/handler/open_region.rs b/src/datanode/src/heartbeat/handler/open_region.rs index 6d57d26e2a6..89ce857ce4d 100644 --- a/src/datanode/src/heartbeat/handler/open_region.rs +++ b/src/datanode/src/heartbeat/handler/open_region.rs @@ -169,11 +169,17 @@ mod tests { .await .unwrap(); region_server - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); region_server - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 25c2c4684a5..3932e307489 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -1852,7 +1852,10 @@ impl RegionServerInner { for (region_id, engine) in regions { let closed = engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await; match closed { Ok(_) => debug!("Region {region_id} is closed"), @@ -2218,7 +2221,10 @@ mod tests { ); let response = mock_region_server - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); assert_eq!(response.affected_rows, 0); diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 01484a7d902..591a3c88ea5 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -511,6 +511,10 @@ impl LogStore for KafkaLogStore { Ok(()) } + async fn obsolete_all(&self, _provider: &Provider, _region_id: RegionId) -> Result<()> { + Ok(()) + } + /// Returns the highest entry id of the specified topic in remote WAL. fn latest_entry_id(&self, provider: &Provider) -> Result { let provider = provider diff --git a/src/log-store/src/noop/log_store.rs b/src/log-store/src/noop/log_store.rs index f0f056514e8..11b77e79108 100644 --- a/src/log-store/src/noop/log_store.rs +++ b/src/log-store/src/noop/log_store.rs @@ -86,6 +86,10 @@ impl LogStore for NoopLogStore { Ok(()) } + async fn obsolete_all(&self, _provider: &Provider, _region_id: RegionId) -> Result<()> { + Ok(()) + } + fn latest_entry_id(&self, _provider: &Provider) -> Result { Ok(0) } diff --git a/src/log-store/src/raft_engine/log_store.rs b/src/log-store/src/raft_engine/log_store.rs index dbd4d0daa7d..923110858b9 100644 --- a/src/log-store/src/raft_engine/log_store.rs +++ b/src/log-store/src/raft_engine/log_store.rs @@ -484,6 +484,12 @@ impl LogStore for RaftEngineLogStore { Ok(()) } + async fn obsolete_all(&self, provider: &Provider, region_id: RegionId) -> Result<()> { + let latest_entry_id = self.latest_entry_id(provider)?; + self.obsolete(provider, region_id, latest_entry_id).await?; + self.delete_namespace(provider).await + } + fn latest_entry_id(&self, provider: &Provider) -> Result { let ns = provider .as_raft_engine_provider() @@ -576,6 +582,38 @@ mod tests { assert!(logstore.list_namespaces().await.unwrap().is_empty()); } + #[tokio::test] + async fn test_obsolete_all_removes_entries_and_namespace() { + let dir = create_temp_dir("raft-engine-logstore-test"); + let logstore = RaftEngineLogStore::try_new( + dir.path().to_str().unwrap().to_string(), + &RaftEngineConfig::default(), + ) + .await + .unwrap(); + let region_id = RegionId::new(1, 1); + let provider = Provider::raft_engine_provider(region_id.as_u64()); + logstore.create_namespace(&provider).await.unwrap(); + for entry_id in 1..=3 { + logstore + .append( + EntryImpl::create( + entry_id, + region_id.as_u64(), + entry_id.to_string().into_bytes(), + ) + .into(), + ) + .await + .unwrap(); + } + + logstore.obsolete_all(&provider, region_id).await.unwrap(); + + assert_eq!(0, logstore.latest_entry_id(&provider).unwrap()); + assert!(logstore.list_namespaces().await.unwrap().is_empty()); + } + #[tokio::test] async fn test_append_and_read() { let dir = create_temp_dir("raft-engine-logstore-test"); diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index 3cc18f5f153..f8eceda9e8a 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -604,7 +604,7 @@ mod test { engine .handle_request( physical_region_id, - RegionRequest::Close(RegionCloseRequest {}), + RegionRequest::Close(RegionCloseRequest::default()), ) .await .unwrap(); @@ -632,7 +632,7 @@ mod test { engine .handle_request( nonexistent_region_id, - RegionRequest::Close(RegionCloseRequest {}), + RegionRequest::Close(RegionCloseRequest::default()), ) .await .unwrap(); @@ -736,7 +736,7 @@ mod test { metric_engine .handle_request( physical_region_id, - RegionRequest::Close(RegionCloseRequest {}), + RegionRequest::Close(RegionCloseRequest::default()), ) .await .unwrap(); @@ -828,7 +828,10 @@ mod test { // Closes all regions for region_id in logical_region_ids.iter().chain(physical_region_ids.iter()) { metric_engine - .handle_request(*region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + *region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); } diff --git a/src/metric-engine/src/engine/close.rs b/src/metric-engine/src/engine/close.rs index 1af507cf367..fdf2685c0d1 100644 --- a/src/metric-engine/src/engine/close.rs +++ b/src/metric-engine/src/engine/close.rs @@ -29,7 +29,7 @@ impl MetricEngineInner { pub async fn close_region( &self, region_id: RegionId, - _req: RegionCloseRequest, + req: RegionCloseRequest, ) -> Result { let data_region_id = utils::to_data_region_id(region_id); if self @@ -38,7 +38,8 @@ impl MetricEngineInner { .unwrap() .exist_physical_region(data_region_id) { - self.close_physical_region(data_region_id).await?; + self.close_physical_region(data_region_id, req.flush_on_close) + .await?; self.state .write() .unwrap() @@ -59,18 +60,25 @@ impl MetricEngineInner { } } - pub(crate) async fn close_physical_region(&self, region_id: RegionId) -> Result { + pub(crate) async fn close_physical_region( + &self, + region_id: RegionId, + flush_on_close: bool, + ) -> Result { let data_region_id = utils::to_data_region_id(region_id); let metadata_region_id = utils::to_metadata_region_id(region_id); self.mito - .handle_request(data_region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + data_region_id, + RegionRequest::Close(RegionCloseRequest { flush_on_close }), + ) .await .with_context(|_| CloseMitoRegionSnafu { region_id })?; self.mito .handle_request( metadata_region_id, - RegionRequest::Close(RegionCloseRequest {}), + RegionRequest::Close(RegionCloseRequest { flush_on_close }), ) .await .with_context(|_| CloseMitoRegionSnafu { region_id })?; diff --git a/src/metric-engine/src/engine/open.rs b/src/metric-engine/src/engine/open.rs index 8fcdfcd8211..65ffcd77eb4 100644 --- a/src/metric-engine/src/engine/open.rs +++ b/src/metric-engine/src/engine/open.rs @@ -105,7 +105,7 @@ impl MetricEngineInner { utils::to_metadata_region_id(physical_region_id), utils::to_data_region_id(physical_region_id) ); - if let Err(err) = self.close_physical_region(physical_region_id).await { + if let Err(err) = self.close_physical_region(physical_region_id, false).await { error!(err; "Failed to close physical region {}", physical_region_id); } } diff --git a/src/metric-engine/src/engine/sync/region.rs b/src/metric-engine/src/engine/sync/region.rs index d1f92bef64b..c7b3ac8682c 100644 --- a/src/metric-engine/src/engine/sync/region.rs +++ b/src/metric-engine/src/engine/sync/region.rs @@ -304,7 +304,7 @@ mod tests { metric_engine .handle_request( target_physical_region_id, - RegionRequest::Close(RegionCloseRequest {}), + RegionRequest::Close(RegionCloseRequest::default()), ) .await .unwrap(); diff --git a/src/mito2/src/engine/catchup_test.rs b/src/mito2/src/engine/catchup_test.rs index 54326cecefd..5c8e357c056 100644 --- a/src/mito2/src/engine/catchup_test.rs +++ b/src/mito2/src/engine/catchup_test.rs @@ -498,7 +498,10 @@ async fn test_catchup_with_manifest_update(factory: Option) { async fn close_region(engine: &MitoEngine, region_id: RegionId) { engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); } diff --git a/src/mito2/src/engine/close_test.rs b/src/mito2/src/engine/close_test.rs index 0092634ec74..c3ace2f5c42 100644 --- a/src/mito2/src/engine/close_test.rs +++ b/src/mito2/src/engine/close_test.rs @@ -15,15 +15,20 @@ use std::sync::Arc; use std::sync::atomic::Ordering; +use api::v1::Rows; use common_base::Plugins; +use common_recordbatch::RecordBatches; use store_api::region_engine::RegionEngine; use store_api::region_request::{RegionCloseRequest, RegionRequest}; -use store_api::storage::RegionId; +use store_api::storage::{RegionId, ScanRequest}; use crate::config::MitoConfig; use crate::engine::flush_test::MockRegionHook; +use crate::engine::listener::AlterFlushListener; use crate::engine::region_hook::RegionHookRef; -use crate::test_util::{CreateRequestBuilder, TestEnv}; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, rows_schema, +}; #[tokio::test] async fn test_engine_close_region() { @@ -43,7 +48,10 @@ async fn test_engine_close_region_with_format(flat_format: bool) { let region_id = RegionId::new(1, 1); // It's okay to close a region doesn't exist. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -55,14 +63,20 @@ async fn test_engine_close_region_with_format(flat_format: bool) { // Close the created region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); assert!(!engine.is_region_exists(region_id)); // It's okay to close this region again. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); } @@ -93,7 +107,10 @@ async fn test_region_hook_on_close() { assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0); engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -103,3 +120,160 @@ async fn test_region_hook_on_close() { assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0); assert_eq!(hook.files_removed_count.load(Ordering::Relaxed), 0); } + +#[tokio::test] +async fn test_engine_close_region_flush_on_close() { + let mut env = TestEnv::with_prefix("close-flush-on-close").await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + ) + .await; + assert!( + !engine + .get_region(region_id) + .unwrap() + .version() + .memtables + .is_empty() + ); + + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest { + flush_on_close: true, + }), + ) + .await + .unwrap(); + assert!(!engine.is_region_exists(region_id)); + + engine + .handle_request( + region_id, + RegionRequest::Open(store_api::region_request::RegionOpenRequest { + engine: String::new(), + table_dir: request.table_dir.clone(), + path_type: store_api::region_request::PathType::Bare, + options: request.options.clone(), + skip_wal_replay: true, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(3, batches.iter().map(|b| b.num_rows()).sum::()); +} + +#[tokio::test] +async fn test_engine_close_region_flush_on_close_while_flushing() { + let mut env = TestEnv::with_prefix("close-flush-on-close-while-flushing").await; + let listener = Arc::new(AlterFlushListener::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(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + ) + .await; + + let flush_engine = engine.clone(); + let flush_task = tokio::spawn(async move { + flush_region(&flush_engine, region_id, None).await; + }); + listener.wait_flush_begin().await; + + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&request), + rows: build_rows(3, 6), + }, + ) + .await; + + let request_count = listener.request_count(); + let close_engine = engine.clone(); + let close_task = tokio::spawn(async move { + close_engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest { + flush_on_close: true, + }), + ) + .await + .unwrap(); + }); + listener.wait_request_count(request_count + 1).await; + + let second_flush_listener = listener.clone(); + let second_flush_task = tokio::spawn(async move { + second_flush_listener.wait_flush_begin().await; + second_flush_listener.wake_flush(); + }); + listener.wake_flush(); + + flush_task.await.unwrap(); + close_task.await.unwrap(); + second_flush_task.abort(); + assert!(!engine.is_region_exists(region_id)); + + engine + .handle_request( + region_id, + RegionRequest::Open(store_api::region_request::RegionOpenRequest { + engine: String::new(), + table_dir: request.table_dir, + path_type: store_api::region_request::PathType::Bare, + options: request.options, + skip_wal_replay: true, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(6, batches.iter().map(|b| b.num_rows()).sum::()); +} diff --git a/src/mito2/src/engine/create_test.rs b/src/mito2/src/engine/create_test.rs index 8c4fe25b727..25b6995190e 100644 --- a/src/mito2/src/engine/create_test.rs +++ b/src/mito2/src/engine/create_test.rs @@ -105,7 +105,10 @@ async fn test_engine_create_close_create_region_with_format(flat_format: bool) { .unwrap(); // Close the region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); // Create the same region id again. diff --git a/src/mito2/src/engine/edit_region_test.rs b/src/mito2/src/engine/edit_region_test.rs index 106156a42c0..5aa514de436 100644 --- a/src/mito2/src/engine/edit_region_test.rs +++ b/src/mito2/src/engine/edit_region_test.rs @@ -419,7 +419,10 @@ async fn test_stalled_write_fails_fast_if_region_closed_during_editing() { let close_engine = engine.clone(); let close_task = tokio::spawn(async move { close_engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await }); diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index dc0063a7e63..6d68a33abda 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -447,7 +447,10 @@ async fn test_skip_remote_wal_replay_sets_topic_latest_entry_id(factory: Option< assert_eq!(1, region.topic_latest_entry_id.load(Ordering::Relaxed)); engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -589,7 +592,10 @@ async fn test_remote_wal_open_without_replayed_memtable_sets_topic_latest_entry_ // The empty flush updates `topic_latest_entry_id` from KafkaLogStore's topic stats. flush_region(&engine, region_id, None).await; engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); diff --git a/src/mito2/src/engine/open_test.rs b/src/mito2/src/engine/open_test.rs index ee0d9560093..2b709a7f08b 100644 --- a/src/mito2/src/engine/open_test.rs +++ b/src/mito2/src/engine/open_test.rs @@ -245,7 +245,10 @@ async fn test_engine_region_open_with_options_with_format(flat_format: bool) { // Close the region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -306,7 +309,10 @@ async fn test_engine_region_open_with_custom_store_with_format(flat_format: bool // Close the custom region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -628,7 +634,10 @@ async fn test_open_compaction_region_with_format(flat_format: bool) { // Close the region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -708,7 +717,10 @@ async fn test_open_backfills_partition_expr_with_fetcher() { // close and reopen to trigger backfill in opener engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); engine @@ -742,7 +754,10 @@ async fn test_open_backfills_partition_expr_with_fetcher() { // reopen again to ensure no further changes and still Some engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); let engine = env.reopen_engine(engine, MitoConfig::default()).await; @@ -784,7 +799,10 @@ async fn test_open_keeps_none_without_fetcher() { assert!(meta.partition_expr.is_none()); engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); let engine = env.reopen_engine(engine, MitoConfig::default()).await; diff --git a/src/mito2/src/engine/skip_wal_test.rs b/src/mito2/src/engine/skip_wal_test.rs index eca19f22d77..2309d91569b 100644 --- a/src/mito2/src/engine/skip_wal_test.rs +++ b/src/mito2/src/engine/skip_wal_test.rs @@ -78,7 +78,10 @@ async fn test_close_region_skip_wal(insert: bool) { // Close the region. This should trigger a flush. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -146,7 +149,10 @@ async fn test_close_follower_region_skip_wal() { // Close the region. This should trigger a flush. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -214,7 +220,10 @@ async fn test_close_follower_region_skip_wal_with_pending_data() { assert!(region.is_follower()); engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -268,7 +277,10 @@ async fn test_close_region_skip_wal_while_flush_in_flight_closes_region() { let engine_cloned = engine.clone(); let close_job = tokio::spawn(async move { engine_cloned - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); }); @@ -344,7 +356,10 @@ async fn test_close_region_skip_wal_rejects_writes_queued_after_close() { let engine_cloned = engine.clone(); let close_job = tokio::spawn(async move { engine_cloned - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); }); @@ -445,7 +460,10 @@ async fn test_concurrent_close_region_skip_wal_while_flush_in_flight_succeeds() let engine_cloned = engine.clone(); let first_close = tokio::spawn(async move { engine_cloned - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); }); @@ -455,7 +473,10 @@ async fn test_concurrent_close_region_skip_wal_while_flush_in_flight_succeeds() let engine_cloned = engine.clone(); let second_close = tokio::spawn(async move { engine_cloned - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); }); @@ -518,7 +539,10 @@ async fn test_close_region_after_truncate_skip_wal() { assert!(!region.version().memtables.is_empty()); engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index f087d8a5545..e6aa15f9d48 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -2107,7 +2107,7 @@ mod tests { scheduler.add_ddl_request_to_pending(SenderDdlRequest { sender: OptionOutputTx::from(sender), region_id: builder.region_id(), - request: DdlRequest::Close(store_api::region_request::RegionCloseRequest {}), + request: DdlRequest::Close(store_api::region_request::RegionCloseRequest::default()), }); let version_data = version_control.current(); diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index c90c93df297..5e027884985 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -1307,7 +1307,10 @@ pub async fn reopen_region( ) { // Close the region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); diff --git a/src/mito2/src/wal.rs b/src/mito2/src/wal.rs index ff152a4f603..eb9d4a251bb 100644 --- a/src/mito2/src/wal.rs +++ b/src/mito2/src/wal.rs @@ -159,6 +159,15 @@ impl Wal { .map_err(BoxedError::new) .context(DeleteWalSnafu { region_id }) } + /// Marks all WAL entries of a region as obsolete and removes its dedicated namespace when + /// supported by the backend. + pub async fn obsolete_all(&self, region_id: RegionId, provider: &Provider) -> Result<()> { + self.store + .obsolete_all(provider, region_id) + .await + .map_err(BoxedError::new) + .context(DeleteWalSnafu { region_id }) + } } /// WAL batch writer. diff --git a/src/mito2/src/wal/raw_entry_reader.rs b/src/mito2/src/wal/raw_entry_reader.rs index de63f92251b..8a79d6c3b16 100644 --- a/src/mito2/src/wal/raw_entry_reader.rs +++ b/src/mito2/src/wal/raw_entry_reader.rs @@ -187,6 +187,14 @@ mod tests { unreachable!() } + async fn obsolete_all( + &self, + _provider: &Provider, + _region_id: RegionId, + ) -> Result<(), Self::Error> { + unreachable!() + } + fn entry( &self, _data: Vec, diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 3e1d42caf2c..3cd4a402c88 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -1148,8 +1148,9 @@ impl RegionWorkerLoop { .await; continue; } - DdlRequest::Close(_) => { - self.handle_close_request(ddl.region_id, ddl.sender).await; + DdlRequest::Close(req) => { + self.handle_close_request(ddl.region_id, req, ddl.sender) + .await; continue; } DdlRequest::Alter(req) => { diff --git a/src/mito2/src/worker/handle_close.rs b/src/mito2/src/worker/handle_close.rs index b582c276927..ef5702c8367 100644 --- a/src/mito2/src/worker/handle_close.rs +++ b/src/mito2/src/worker/handle_close.rs @@ -27,6 +27,7 @@ impl RegionWorkerLoop { pub(crate) async fn handle_close_request( &mut self, region_id: RegionId, + request: RegionCloseRequest, sender: OptionOutputTx, ) { let Some(region) = self.regions.get_region(region_id) else { @@ -36,9 +37,10 @@ impl RegionWorkerLoop { info!("Try to close region {}, worker: {}", region_id, self.id); - // If the region is using Noop WAL and has data in memtable and region is flushable (like, - // not in follower state), we should flush it before closing to ensure durability. - if region.provider == Provider::Noop + // If the close request asks for a flush, or the region is using Noop WAL, + // and has data in memtable and region is flushable (like, not in follower state), + // we should flush it before closing to ensure durability. + if (request.flush_on_close || region.provider == Provider::Noop) && !region .version_control .current() @@ -53,7 +55,7 @@ impl RegionWorkerLoop { .add_ddl_request_to_pending(SenderDdlRequest { region_id, sender, - request: DdlRequest::Close(RegionCloseRequest {}), + request: DdlRequest::Close(request), }); return; } diff --git a/src/store-api/src/logstore.rs b/src/store-api/src/logstore.rs index 19bbe195f0b..a2c1e3ee803 100644 --- a/src/store-api/src/logstore.rs +++ b/src/store-api/src/logstore.rs @@ -86,6 +86,14 @@ pub trait LogStore: Send + Sync + 'static + std::fmt::Debug { entry_id: EntryId, ) -> Result<(), Self::Error>; + /// Marks all entries of a region as obsolete and removes its dedicated namespace when + /// supported by the backend. + async fn obsolete_all( + &self, + provider: &Provider, + region_id: RegionId, + ) -> Result<(), Self::Error>; + /// Makes an entry instance of the associated Entry type fn entry( &self, diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index e046e9df568..61b353956b1 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -189,6 +189,10 @@ impl RegionRequest { reason: "RemoteDynFilter request should be handled separately by RegionServer", } .fail(), + region_request::Body::CleanUp(_) => UnexpectedSnafu { + reason: "CleanUp request should be handled separately by RegionServer", + } + .fail(), region_request::Body::ApplyStagingManifest(apply) => { make_region_apply_staging_manifest(apply) } @@ -329,7 +333,9 @@ fn make_region_close(close: CloseRequest) -> Result Date: Tue, 14 Jul 2026 11:25:37 +0800 Subject: [PATCH 15/48] chore: add development Docker image skill (#8477) * chore: add development Docker image skill Signed-off-by: WenyXu * fix: harden development image skill Signed-off-by: WenyXu * fix: address development image review feedback Signed-off-by: WenyXu * fix: add image tag script license header Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .../SKILL.md | 368 ++++++++++++++++++ .../assets/Dockerfile | 14 + .../scripts/binary_platform.py | 78 ++++ .../scripts/build_binary.sh | 65 ++++ .../scripts/build_image.sh | 136 +++++++ .../scripts/collect_context.py | 200 ++++++++++ .../scripts/image_config.py | 155 ++++++++ .../scripts/next_image_tag.py | 48 +++ .../scripts/prepare_context.sh | 65 ++++ .../scripts/update_image_env.py | 53 +++ .../tests/test_binary_platform.py | 100 +++++ .../tests/test_image_config.py | 100 +++++ 12 files changed, 1382 insertions(+) create mode 100644 .agents/skills/greptimedb-development-docker-image/SKILL.md create mode 100644 .agents/skills/greptimedb-development-docker-image/assets/Dockerfile create mode 100644 .agents/skills/greptimedb-development-docker-image/scripts/binary_platform.py create mode 100755 .agents/skills/greptimedb-development-docker-image/scripts/build_binary.sh create mode 100755 .agents/skills/greptimedb-development-docker-image/scripts/build_image.sh create mode 100755 .agents/skills/greptimedb-development-docker-image/scripts/collect_context.py create mode 100644 .agents/skills/greptimedb-development-docker-image/scripts/image_config.py create mode 100755 .agents/skills/greptimedb-development-docker-image/scripts/next_image_tag.py create mode 100755 .agents/skills/greptimedb-development-docker-image/scripts/prepare_context.sh create mode 100755 .agents/skills/greptimedb-development-docker-image/scripts/update_image_env.py create mode 100644 .agents/skills/greptimedb-development-docker-image/tests/test_binary_platform.py create mode 100644 .agents/skills/greptimedb-development-docker-image/tests/test_image_config.py diff --git a/.agents/skills/greptimedb-development-docker-image/SKILL.md b/.agents/skills/greptimedb-development-docker-image/SKILL.md new file mode 100644 index 00000000000..02a9a61d871 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/SKILL.md @@ -0,0 +1,368 @@ +--- +name: greptimedb-development-docker-image +description: Builds a development-only GreptimeDB Docker image from a local debug binary for local-cluster testing, and optionally pushes it to a development registry. Use when the user asks to package, build, tag, publish, or cross-build a non-release GreptimeDB or GreptimeDB Enterprise image for debugging. +compatibility: Requires Docker with Buildx for cross-platform builds; Podman is supported for native-platform builds. +metadata: + protocols: docker buildx podman + platforms: linux-amd64 linux-arm64 +--- + +# GreptimeDB Development Docker Image + +## Goal + +Package a locally built GreptimeDB binary into a **development-only** Docker +image for debugging and local-cluster testing, optionally push it to a +development registry, and retain the non-secret image settings in `.env` for +the next build. It is **not** a release-image workflow and must not be used to +publish a production or release artifact. + +This skill follows the documented development-image build procedure: + +- Build `greptime` with Cargo's `nightly` profile. +- Copy the target binary into the Docker build context as `./greptime`. +- Build the supplied `Dockerfile`, which uses Ubuntu 24.04 only as the runtime + base image and exposes the `greptime` binary as its entrypoint. + +## Inputs + +Collect or discover the following before any build or push: + +| Input | Discovery and rule | +| --- | --- | +| Source repository | Treat the directory in which this skill is invoked as the default workspace. Include it as an editable field in the batch configuration; do not ask a separate confirmation. Do not assume open-source versus Enterprise. | +| Edition and binary | Default to `greptime`; use the binary name the user requests for Enterprise builds. The copied Docker-context filename must always be `greptime`. | +| Cargo profile and binary source | Include this in the batch configuration. Default to rebuilding `nightly`; reuse an existing binary only after the user explicitly accepts that its freshness is unverified. | +| Target platform | Preselect `linux/amd64` unless Docker's server is `linux/arm64`, then preselect `linux/arm64`. Let the user override it in the batch configuration. Ubuntu 24.04 is the runtime base image, not a target-platform choice. Each image has exactly one target platform, but it may differ from the host platform. | +| Build mode | Ask whether the user wants a locally loadable debug image or a registry push. | +| Registry/repository | Read `IMAGE_REGISTRY` and `IMAGE_REPOSITORY` only from the selected workspace's `.env` and prefill them in the batch configuration. Example: `registry.example.com/team` + `greptimedb-dev`. | +| Tag | Read `IMAGE_TAG` only from the selected workspace's `.env` and prefill it in the batch configuration. When the tag exists in the selected registry, preselect an incremented version. | + +Use the image reference `${IMAGE_REGISTRY}/${IMAGE_REPOSITORY}:${IMAGE_TAG}`. +If the registry is intentionally empty, omit its slash rather than producing a +leading slash. + +## Interactive Workflow + +Use the platform's interactive prompt component for every question, selection, +and confirmation. Do not ask an open-ended text question when a single-choice, +multi-select, or confirmation component can represent the decision. If the +platform does not provide an interactive component, use the equivalent numbered +or lettered prompt below and wait for input before continuing. + +### Batch configuration + +Minimize user round trips: run the collector first against the invocation +directory, then use one interactive form or batched prompt to collect workspace, +edition/binary, profile/reuse-or-rebuild choice, target platform, build mode, +registry/repository, tag, and whether to inspect the configured registry tag. +Use the invocation directory as the editable workspace default. Prefill +collector values and mark recommended defaults. If the user changes workspace, +rerun the collector for that workspace without asking another configuration +question. Only ask a follow-up if information is missing or invalid, or a safety +gate is required. Treat unchanged prefilled values as accepted. + +Keep separate confirmation components only for elevated privileges: `sudo`, +package installation, and privileged QEMU setup. Use one final confirmation for +all non-privileged selected work. + +Always offer **Cancel** for an action that can write state, build, push, or +require elevated privileges. Display the relevant preview before the user +confirms it: source/binary path, `file` architecture result, platform(s), image +reference, build mode, and non-secret `.env` values. + +### Batch configuration fields + +Use single-choice controls inside the one batch form for mutually exclusive +choices: + +| Decision | Required choices | +| --- | --- | +| Cargo output | Rebuild `nightly` (recommended), reuse an existing binary only with a freshness-unverified acknowledgement, choose a profile/binary | +| Target platform | `linux/amd64` (default unless Docker server is arm64), `linux/arm64` (default when Docker server is arm64). The selected image may be cross-built with Buildx. | +| Build mode | Load locally (recommended), push to development registry | +| Registry/repository | Use `.env` defaults, enter new values | +| Tag | Use proposed increment when the configured tag exists, enter a version | + +When push is selected and the image configuration is complete, automatically +run the read-only registry-tag check. It may return `unknown`; do not treat that +as a missing tag. + +Use one final confirmation for all selected non-privileged operations: + +- persist changed non-secret image settings to `.env`; +- run the selected Cargo build or reuse the resolved binary; +- prepare the Docker context; +- build, verify, and, when selected, push the image. + +The final confirmation must clearly state: **“This creates a development and +local-cluster test image, not a release or production artifact.”** + +For a real multi-select only, use lettered choices (`[A]`, `[B]`) and accept +`all`, `none`, or `cancel`; otherwise use single-choice components. + +### Preflight context collection + +Before asking the profile, registry, or tag questions, run the bundled, +read-only collector. It works on macOS and Linux, does not invoke `cargo build`, +and emits JSON that can be shown or summarized to the user: + +```bash +python3 /scripts/collect_context.py \ + --source \ + --bin greptime \ + --profile nightly \ + [--target ] +``` + +Use its report to identify: + +- missing `IMAGE_REGISTRY`, `IMAGE_REPOSITORY`, or `IMAGE_TAG` values in the + selected workspace's `.env`; +- Cargo binary targets available in the workspace and whether the requested + `--bin` exists; +- the resolved Cargo target directory, expected profile output path, and whether + that binary already exists; and +- macOS/Linux host architecture plus Docker's server platform when Docker is + available. + +When push is selected and image configuration is complete, rerun the collector +with `--check-registry-tag` and all selected image values. It uses the selected +engine's read-only manifest inspection and returns `exists` or `unknown`; +`unknown` includes an unavailable registry, an absent tag, or missing +authentication and must not be treated as an absent tag. + +```bash +python3 /scripts/collect_context.py \ + --source \ + --bin \ + --profile \ + --engine \ + --registry \ + --repository \ + --tag \ + --check-registry-tag +``` + +When the status is `exists` and the tag ends in a number, preselect the +collector's `candidate_tag` as the next development tag. The user must still +confirm it before `.env` is updated or an image is built. Never use automatic +incrementing as permission to overwrite or push an existing image. + +If Cargo metadata or the requested binary target is unavailable, stop before +building and ask the user to correct the source path or binary target. + +The profile/binary choice is a field in the batch configuration, not a separate +question. Its default is the existing `nightly` output when present: + +```text +Question: Which Cargo build output should package this image? +Options: +- Rebuild nightly (Recommended): run Cargo with `--profile nightly`, then use its output. +- Reuse nightly output: use the existing output only after showing its path, modification time, size, and matching platform; freshness is unverified. +- Choose profile or binary: provide a different Cargo profile or an explicit binary path. +``` + +Do not infer that a binary under `target/debug`, `target/release`, or another +profile is suitable. For an existing binary, show its resolved path and run +`file` before asking for final build confirmation. + +## Environment State + +The image configuration lives only in the selected workspace's `.env`. The +collector reads only the following keys and never reads or displays other `.env` +values: + +```dotenv +IMAGE_REGISTRY=registry.example.com/team +IMAGE_REPOSITORY=greptimedb-dev +IMAGE_TAG=dev-001 +``` + +1. Read `.env` if it exists. Ignore blank lines and `#` comments. +2. Show the discovered registry, repository, and proposed tag to the user. +3. Include registry/repository and tag in the batch configuration even when + values already exist; saved values are defaults, not authorization to reuse + them. +4. Compare the selected configuration with the three current managed values. If + all match, skip `.env` confirmation and do not update the file. Otherwise, + preview only the changed `IMAGE_*` values, request confirmation, and update + the file while preserving unrelated entries and comments. +5. Never put registry credentials, access tokens, passwords, or `docker login` + output in `.env`, build commands, logs, or responses. Ask the user to log in + themselves if a push needs authentication. + +For a repeated build, offer an incremented tag before asking. Use +`scripts/next_image_tag.py --tag ` to calculate it. It increments +the final numeric component while preserving zero padding: + +```text +weny-2025-0715-01 -> weny-2025-0715-02 +v0.1.4 -> v0.1.5 +debug-009 -> debug-010 +``` + +If the saved tag has no trailing number, ask the user for the next tag rather +than inventing a versioning convention. Do not overwrite an existing image tag +without explicit user confirmation. + +Only if a managed value is missing or differs, persist the values with the +bundled helper after confirmation instead of hand-editing the file: + +```bash +python3 /scripts/update_image_env.py \ + --env /.env \ + --registry \ + --repository \ + --tag +``` + +## Build Procedure + +### 1. Inspect the host and container tooling + +Run these checks and report the selected path: + +```bash +uname -s +uname -m +docker version --format '{{.Server.Os}}/{{.Server.Arch}}' +docker buildx version +docker buildx inspect --bootstrap +``` + +For a native-platform build, Podman can be used if Docker is unavailable. For +any non-native request, require Docker Buildx. Do not silently fall back to a +native build when the requested image platform differs. + +If the requested target is unavailable, stop and ask the user to configure a +Buildx builder and cross-compilation toolchain externally. Do not automate +privileged QEMU or binfmt setup, and never run an unpinned privileged image. + +### 2. Build the executable for the image platform + +From the source checkout, build the requested binary. For open-source amd64: + +```bash +/scripts/build_binary.sh \ + --source --package cmd --bin greptime --profile nightly +``` + +For open-source arm64: + +```bash +/scripts/build_binary.sh \ + --source \ + --package cmd \ + --bin greptime \ + --profile nightly \ + --target aarch64-unknown-linux-gnu +``` + +For Enterprise, use the user-provided executable target, for example: + +```bash +/scripts/build_binary.sh \ + --source --bin greptime-ent-cloud --profile nightly +``` + +For a user-selected profile, replace `nightly` in the helper call and resolve +the binary under `target//` (or +`target///` for cross-compilation). Reuse the +existing target file only when the interactive profile question selected reuse; +otherwise invoke the helper to rebuild it. + +Do not use a binary compiled for the host architecture in an image intended for +another architecture. Determine the binary path from the Cargo target and +profile, then copy it into the Docker context as `greptime`: + +```bash +cp /greptime +``` + +Verify it before building the image with `file /greptime`; its +reported architecture must match the requested target platform. + +Create a fresh isolated build context (never use the repository root) and use +the bundled preparation script. The helper rejects an existing Dockerfile or +binary output to prevent accidental context reuse: + +```bash +/scripts/prepare_context.sh \ + --context "$(mktemp -d)" \ + --binary \ + --dockerfile /assets/Dockerfile \ + --platform +``` + +### 3. Build or push the image + +Run the command only after the user confirms the final image reference and +whether it should be pushed. + +Use the bundled build helper rather than spelling out individual Docker or +Podman commands. It follows the existing scripts' Docker-first/Podman-fallback +behavior, uses Buildx for a non-native Docker request, and accepts exactly one +target platform per image: + +```bash +/scripts/build_image.sh \ + --context \ + --image \ + --platform \ + --mode local|push +``` + +`--mode local` builds a native image locally or uses `docker buildx --load` for +a non-native image. `--mode push` pushes the selected single-platform image. +The helper never pushes in `local` mode. + +This is intentionally a runtime image built from a precompiled binary, not a +multi-stage Dockerfile. Ubuntu 24.04 is the runtime base image only; select one +target platform separately as `linux/amd64` or `linux/arm64`. Only introduce a +multi-stage Dockerfile if the user asks to compile within Docker or the local +Rust toolchain cannot build the required target. Keep the final Ubuntu 24.04 +runtime stage and existing entrypoint unless the user asks to change runtime +behavior. + +At the final confirmation, repeat that the selected reference is a +**development/local-cluster test image**, not a release artifact. If the user +requests a release or production image, stop and direct them to the release +process instead. + +### 4. Verify the result + +For a local image, use the selected engine to inspect its architecture and +start it with `--help`: + +```bash +docker image inspect --format '{{.Os}}/{{.Architecture}}' +docker run --rm --help +``` + +For Podman, use `podman image inspect ` and +`podman run --rm --help`. For a pushed Podman image, use +`podman manifest inspect `. + +For a pushed image, inspect its remote manifest: + +```bash +docker buildx imagetools inspect +``` + +Report the final image reference, target platform, source binary path, and +whether the image was loaded locally or pushed. + +## Safety Rules + +- This skill is only for development and local-cluster testing. Never present + its image as a production, release, or officially published artifact. +- Never push by default. Local debug builds should use a local tag and load the + image unless the user explicitly requests a push. +- Do not run `docker login`, privileged QEMU setup, package installation, or + overwrite a tag without confirmation. +- Do not delete existing images, builders, binaries, or `.env` entries as part + of this workflow. +- Stop and explain if Docker's server is unavailable, Buildx cannot support the + requested platform, or the copied binary architecture does not match the + image platform. diff --git a/.agents/skills/greptimedb-development-docker-image/assets/Dockerfile b/.agents/skills/greptimedb-development-docker-image/assets/Dockerfile new file mode 100644 index 00000000000..bc74c830310 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/assets/Dockerfile @@ -0,0 +1,14 @@ +FROM ubuntu:24.04 + +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + curl \ + unzip \ + mysql-client \ + wget \ + && rm -rf /var/lib/apt/lists/* + +COPY ./greptime /greptime/bin/ + +ENV PATH=/greptime/bin/:$PATH + +ENTRYPOINT ["greptime"] diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/binary_platform.py b/.agents/skills/greptimedb-development-docker-image/scripts/binary_platform.py new file mode 100644 index 00000000000..86f855a11cf --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/binary_platform.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +# 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. + +"""Validate that a binary is an executable ELF for a requested Linux platform.""" + +import argparse +import stat +import struct +import sys +from pathlib import Path +from typing import Optional + +PLATFORM_MACHINES = {"linux/amd64": 0x3E, "linux/arm64": 0xB7} + + +def validate_platform(platform: str) -> int: + try: + return PLATFORM_MACHINES[platform] + except KeyError as error: + raise ValueError("platform must be linux/amd64 or linux/arm64") from error + + +def inspect_elf(path: Path) -> int: + executable_bits = stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH + if path.is_symlink() or not path.is_file() or not (path.stat().st_mode & executable_bits): + raise ValueError("binary must be an executable regular file") + try: + with path.open("rb") as handle: + header = handle.read(64) + except OSError as error: + raise ValueError(f"cannot read binary: {error}") from error + if len(header) != 64 or header[:4] != b"\x7fELF": + raise ValueError("binary is not ELF") + if header[4] != 2 or header[5] != 1 or header[6] != 1: + raise ValueError("binary must be a 64-bit little-endian ELF") + if header[7] not in (0, 3): + raise ValueError("binary must use the System V or Linux ELF ABI") + elf_type, machine, version = struct.unpack_from(" None: + expected = validate_platform(platform) + actual = inspect_elf(path) + if actual != expected: + raise ValueError("binary architecture does not match requested platform") + + +def main(argv: Optional[list[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", required=True, type=Path) + parser.add_argument("--platform", required=True) + args = parser.parse_args(argv) + try: + validate_binary(args.binary, args.platform) + except ValueError as error: + print(f"error: {error}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/build_binary.sh b/.agents/skills/greptimedb-development-docker-image/scripts/build_binary.sh new file mode 100755 index 00000000000..605685fa7c7 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/build_binary.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash + +set -euo pipefail + +usage() { + printf 'Usage: %s --source PATH --bin NAME --profile PROFILE [--target RUST_TARGET] [--package PACKAGE] [--features FEATURES]\n' "$0" >&2 +} + +source_dir="" +binary_name="" +profile="" +target="" +package="" +features="" + +while [[ $# -gt 0 ]]; do + case "$1" in + --source|--bin|--profile|--target|--package|--features) + if [[ $# -lt 2 ]]; then + usage + exit 2 + fi + case "$1" in + --source) source_dir="$2" ;; + --bin) binary_name="$2" ;; + --profile) profile="$2" ;; + --target) target="$2" ;; + --package) package="$2" ;; + --features) features="$2" ;; + esac + shift 2 + ;; + *) + usage + exit 2 + ;; + esac +done + +if [[ -z "$source_dir" || -z "$binary_name" || -z "$profile" ]]; then + usage + exit 2 +fi + +if [[ ! -f "$source_dir/Cargo.toml" ]]; then + printf 'error: not a Cargo workspace: %s\n' "$source_dir" >&2 + exit 1 +fi + +command=(cargo build --locked --bin "$binary_name" --profile "$profile") +if [[ -n "$package" ]]; then + command+=(--package "$package") +fi +if [[ -n "$features" ]]; then + command+=(--features "$features") +fi +if [[ -n "$target" ]]; then + command+=(--target "$target") +fi + +printf 'building %s with profile %s\n' "$binary_name" "$profile" +( + cd "$source_dir" + exec "${command[@]}" +) diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/build_image.sh b/.agents/skills/greptimedb-development-docker-image/scripts/build_image.sh new file mode 100755 index 00000000000..3b480334671 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/build_image.sh @@ -0,0 +1,136 @@ +#!/usr/bin/env bash + +set -euo pipefail + +usage() { + printf 'Usage: %s --context PATH --image REFERENCE --platform PLATFORM --mode local|push [--dockerfile PATH] [--engine auto|docker|podman] [--allow-local-tag-overwrite]\n' "$0" >&2 +} + +context="" +image="" +platform="" +mode="" +dockerfile="" +engine="auto" +allow_local_tag_overwrite=false + +while [[ $# -gt 0 ]]; do + case "$1" in + --context|--image|--platform|--mode|--dockerfile|--engine) + if [[ $# -lt 2 ]]; then + usage + exit 2 + fi + case "$1" in + --context) context="$2" ;; + --image) image="$2" ;; + --platform) platform="$2" ;; + --mode) mode="$2" ;; + --dockerfile) dockerfile="$2" ;; + --engine) engine="$2" ;; + esac + shift 2 + ;; + --allow-local-tag-overwrite) + allow_local_tag_overwrite=true + shift + ;; + *) + usage + exit 2 + ;; + esac +done + +if [[ -z "$context" || -z "$image" || -z "$platform" || -z "$mode" ]]; then + usage + exit 2 +fi + +if [[ "$mode" != "local" && "$mode" != "push" ]]; then + printf 'error: --mode must be local or push\n' >&2 + exit 2 +fi + +if [[ "$engine" != "auto" && "$engine" != "docker" && "$engine" != "podman" ]]; then + printf 'error: --engine must be auto, docker, or podman\n' >&2 + exit 2 +fi + +dockerfile="${dockerfile:-$context/Dockerfile}" +if [[ ! -f "$dockerfile" ]]; then + printf 'error: Dockerfile does not exist: %s\n' "$dockerfile" >&2 + exit 1 +fi + +if [[ "$platform" == *,* ]]; then + printf 'error: only one target platform is supported per image build\n' >&2 + exit 2 +fi + +case "$platform" in + linux/amd64|linux/arm64) ;; + *) + printf 'error: unsupported platform: %s\n' "$platform" >&2 + exit 2 + ;; +esac + +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +python3 "$script_dir/binary_platform.py" --binary "$context/greptime" --platform "$platform" + +check_collision() { + local runtime="$1" + if [[ "$allow_local_tag_overwrite" == true ]]; then + return 0 + fi + if [[ "$runtime" == docker ]] && docker image inspect "$image" >/dev/null 2>&1; then + printf 'error: local image tag exists: %s (pass --allow-local-tag-overwrite to replace it)\n' "$image" >&2 + exit 3 + fi + if [[ "$runtime" == podman ]] && podman image exists "$image"; then + printf 'error: local image tag exists: %s (pass --allow-local-tag-overwrite to replace it)\n' "$image" >&2 + exit 3 + fi +} + +if [[ "$engine" != podman ]] && command -v docker >/dev/null 2>&1 && server_platform="$(docker version --format '{{.Server.Os}}/{{.Server.Arch}}' 2>/dev/null)"; then + if [[ "$platform" == "$server_platform" ]]; then + check_collision docker + docker build --platform "$platform" -f "$dockerfile" -t "$image" "$context" + if [[ "$mode" == "push" ]]; then + exec docker push "$image" + fi + exit 0 + fi + + if [[ "$mode" == "push" ]]; then + exec docker buildx build --platform "$platform" --push -f "$dockerfile" -t "$image" "$context" + fi + check_collision docker + exec docker buildx build --platform "$platform" --load -f "$dockerfile" -t "$image" "$context" +fi + +if [[ "$engine" != docker ]] && command -v podman >/dev/null 2>&1; then + native_platform="$(podman info --format '{{.Host.Os}}/{{.Host.Arch}}' 2>/dev/null || true)" + case "$native_platform" in + linux/amd64|linux/arm64) ;; + *) + printf 'error: Podman is unavailable or has unsupported platform: %s\n' "$native_platform" >&2 + exit 1 + ;; + esac + if [[ "$platform" != "$native_platform" ]]; then + printf 'error: Docker Buildx is required for non-native platform %s\n' "$platform" >&2 + exit 1 + fi + check_collision podman + podman build -f "$dockerfile" -t "$image" "$context" + if [[ "$mode" == "push" ]]; then + exec podman push "$image" + fi + exit 0 +fi + +printf 'error: Docker or Podman is required\n' >&2 +exit 1 diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/collect_context.py b/.agents/skills/greptimedb-development-docker-image/scripts/collect_context.py new file mode 100755 index 00000000000..214d56ee789 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/collect_context.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +# 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. + +"""Collect read-only build context for a GreptimeDB development image.""" + +import argparse +import json +import os +import platform +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Optional, cast + +from next_image_tag import increment_tag +from binary_platform import PLATFORM_MACHINES, inspect_elf +from image_config import image_reference, parse_env_file, validate_values + + +ENV_KEYS = ("IMAGE_REGISTRY", "IMAGE_REPOSITORY", "IMAGE_TAG") +PROFILE_DIRECTORIES = {"dev": "debug", "test": "debug", "bench": "release"} +PLATFORM_ARCHITECTURES = { + "x86_64": "linux/amd64", + "amd64": "linux/amd64", + "aarch64": "linux/arm64", + "arm64": "linux/arm64", +} + + +def parse_arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", required=True, type=Path) + parser.add_argument("--bin", dest="binary", default="greptime") + parser.add_argument("--profile", default="nightly") + parser.add_argument("--target", help="Rust target triple for cross-compilation") + parser.add_argument("--package", help="Cargo package containing the selected binary") + parser.add_argument("--platform", choices=("linux/amd64", "linux/arm64")) + parser.add_argument("--engine", choices=("auto", "docker", "podman"), default="auto") + parser.add_argument("--registry") + parser.add_argument("--repository") + parser.add_argument("--tag") + parser.add_argument( + "--check-registry-tag", + action="store_true", + help="inspect the configured image tag with Docker without modifying it", + ) + return parser.parse_args() + + +def native_platform(architecture: str) -> Optional[str]: + return PLATFORM_ARCHITECTURES.get(architecture.lower()) + + +def parse_env(env_path: Path) -> dict[str, object]: + values = parse_env_file(env_path) + missing = [key for key in ("IMAGE_REPOSITORY", "IMAGE_TAG") if not values.get(key)] + try: + if not missing: + validate_values(values) + error = None + except ValueError as exc: + error = str(exc) + return {"path": str(env_path), "exists": env_path.is_file(), "values": values, "missing": missing, "error": error} + + +def cargo_metadata(source: Path) -> dict[str, Any]: + command = ["cargo", "metadata", "--locked", "--format-version=1", "--no-deps"] + try: + result = subprocess.run(command, cwd=source, text=True, capture_output=True, check=True) + except FileNotFoundError as error: + raise RuntimeError("cargo is not installed or is not on PATH") from error + except subprocess.CalledProcessError as error: + detail = error.stderr.strip() or error.stdout.strip() + raise RuntimeError(f"cargo metadata failed: {detail}") from error + return json.loads(result.stdout) + + +def binary_targets(metadata: dict[str, Any]) -> list[dict[str, Any]]: + targets = [] + for package in metadata["packages"]: + for target in package["targets"]: + if "bin" in target["kind"]: + targets.append({ + "package": package["name"], + "name": target["name"], + "path": target["src_path"], + "required_features": target.get("required-features", []), + }) + return sorted(targets, key=lambda target: (target["name"], target["package"])) + + +def expected_binary(target_dir: Path, binary: str, profile: str, rust_target: Optional[str]) -> Path: + profile_directory = PROFILE_DIRECTORIES.get(profile, profile) + output_dir = target_dir / rust_target if rust_target else target_dir + return output_dir / profile_directory / binary + + +def engine_status(engine: str) -> dict[str, object]: + if engine == "docker": + command = ["docker", "version", "--format", "{{.Server.Os}}/{{.Server.Arch}}"] + else: + command = ["podman", "info", "--format", "{{.Host.Os}}/{{.Host.Arch}}"] + try: + result = subprocess.run(command, text=True, capture_output=True, check=True) + except (FileNotFoundError, subprocess.CalledProcessError): + return {"available": False, "server_platform": None} + return {"available": True, "server_platform": result.stdout.strip()} + + +def registry_tag_status(values: dict[str, str], enabled: bool, engine: str, source: str) -> dict[str, object]: + if not enabled: + return {"checked": False, "status": "not_checked", "candidate_tag": None} + + try: + reference = image_reference(values) + except ValueError: + return {"checked": False, "status": "not_configured", "candidate_tag": None, "configuration_source": source} + command = ["docker", "buildx", "imagetools", "inspect", reference] if engine == "docker" else ["podman", "manifest", "inspect", reference] + try: + subprocess.run(command, text=True, capture_output=True, check=True) + except FileNotFoundError: + return {"checked": True, "status": "unknown", "candidate_tag": None, "reference": reference, "engine": engine, "configuration_source": source} + except subprocess.CalledProcessError: + return {"checked": True, "status": "unknown", "candidate_tag": None, "reference": reference, "engine": engine, "configuration_source": source} + + try: + candidate_tag = increment_tag(values["IMAGE_TAG"]) + except ValueError: + candidate_tag = None + return {"checked": True, "status": "exists", "candidate_tag": candidate_tag, "reference": reference, "engine": engine, "configuration_source": source} + + +def collect_report(args: argparse.Namespace) -> dict[str, object]: + source = args.source.expanduser().resolve() + if not (source / "Cargo.toml").is_file(): + raise RuntimeError(f"not a Cargo workspace: {source}") + + metadata = cargo_metadata(source) + targets = binary_targets(metadata) + target_dir = Path(metadata["target_directory"]) + expected = expected_binary(target_dir, args.binary, args.profile, args.target) + host_architecture = platform.machine() + selected_target = next((target for target in targets if target["name"] == args.binary), None) + environment = parse_env(source / ".env") + supplied_values = (args.registry, args.repository, args.tag) + if any(value is not None for value in supplied_values) and not all(value is not None for value in supplied_values): + raise RuntimeError("--registry, --repository, and --tag must be supplied together") + values: dict[str, str] = {"IMAGE_REGISTRY": args.registry or "", "IMAGE_REPOSITORY": args.repository or "", "IMAGE_TAG": args.tag or ""} if all(value is not None for value in supplied_values) else cast(dict[str, str], environment["values"]) + configuration_source = "arguments" if all(value is not None for value in supplied_values) else "environment" + docker = engine_status("docker") + podman = engine_status("podman") + engine = args.engine if args.engine != "auto" else "docker" if docker["available"] else "podman" + binary_details: dict[str, object] = {"platform": None, "error": None} + if expected.is_file(): + try: + machine = inspect_elf(expected) + binary_details["platform"] = next( + (name for name, expected_machine in PLATFORM_MACHINES.items() if expected_machine == machine), + None, + ) + if binary_details["platform"] is None: + binary_details["error"] = f"binary has unsupported ELF machine: {machine}" + except ValueError as exc: + binary_details["error"] = str(exc) + return { + "host": {"os": platform.system(), "arch": host_architecture, "native_platform": native_platform(host_architecture)}, + "docker": docker, + "podman": podman, + "environment": environment, + "registry_tag": registry_tag_status(values, args.check_registry_tag, engine, configuration_source), + "cargo": {"workspace_root": metadata["workspace_root"], "target_dir": str(target_dir), "binaries": targets}, + "selected_binary": {"name": args.binary, "package": args.package, "cargo_target_exists": selected_target is not None, "target": selected_target}, + "expected_binary": {"profile": args.profile, "rust_target": args.target, "path": str(expected), "exists": expected.is_file(), "executable": os.access(expected, os.X_OK), "platform": binary_details["platform"], "requested_platform": args.platform, "platform_matches": binary_details["platform"] == args.platform if args.platform and binary_details["platform"] else None, "validation_error": binary_details["error"], "mtime": datetime.fromtimestamp(expected.stat().st_mtime, timezone.utc).isoformat() if expected.exists() else None, "freshness": "unverified"}, + } + + +def main() -> int: + try: + print(json.dumps(collect_report(parse_arguments()), indent=2, sort_keys=True)) + except RuntimeError as error: + print(f"error: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/image_config.py b/.agents/skills/greptimedb-development-docker-image/scripts/image_config.py new file mode 100644 index 00000000000..cb3260d8c8f --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/image_config.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +# 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. + +"""Bounded parsing and safe updating of the image settings in a .env file.""" + +import os +import re +import tempfile +from pathlib import Path +from typing import Mapping + +MANAGED_KEYS = ("IMAGE_REGISTRY", "IMAGE_REPOSITORY", "IMAGE_TAG") +_REGISTRY = re.compile(r"[a-z0-9][a-z0-9.-]*(?::[0-9]+)?(?:/[a-z0-9][a-z0-9._-]*)*") +_REPOSITORY = re.compile(r"[a-z0-9][a-z0-9._-]*(?:/[a-z0-9][a-z0-9._-]*)*") +_TAG = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}") + + +def _value(text: str) -> str: + text = text.strip() + if not text: + return "" + quote = text[0] if text[0] in "'\"" else None + if quote: + end = text.find(quote, 1) + if end < 0 or text[end + 1 :].strip() and not text[end + 1 :].lstrip().startswith("#"): + raise ValueError("unmatched quote or trailing value syntax") + return text[1:end] + # A comment is inline only when introduced by whitespace. This keeps + # valid values such as an image tag containing '#'. + match = re.search(r"\s+#", text) + if match: + text = text[: match.start()].rstrip() + if any(character in text for character in "'\";$`"): + raise ValueError("unsupported shell syntax in .env value") + return text + + +def parse_env(text: str) -> dict[str, str]: + """Parse the three managed keys, without validating their values.""" + lines = text.splitlines() + result: dict[str, str] = {} + for number, original in enumerate(lines, 1): + line = original.strip() + if not line or line.startswith("#"): + continue + export = re.match(r"export[ \t]+", line) + if export: + line = line[export.end() :].lstrip() + key, separator, raw = line.partition("=") + if not separator: + continue + key = key.strip() + if key in MANAGED_KEYS: + try: + result[key] = _value(raw) + except ValueError as error: + raise ValueError(f"line {number}: {error}") from error + return result + + +def parse_env_file(path: Path) -> dict[str, str]: + return parse_env(path.read_text() if path.exists() else "") + + +def normalize_values(values: Mapping[str, str]) -> dict[str, str]: + return { + "IMAGE_REGISTRY": values.get("IMAGE_REGISTRY", ""), + "IMAGE_REPOSITORY": values.get("IMAGE_REPOSITORY", ""), + "IMAGE_TAG": values.get("IMAGE_TAG", ""), + } + + +def validate_values(values: Mapping[str, str]) -> None: + values = normalize_values(values) + missing = [key for key in ("IMAGE_REPOSITORY", "IMAGE_TAG") if not values.get(key)] + if missing: + raise ValueError("missing managed values: " + ", ".join(missing)) + registry = values.get("IMAGE_REGISTRY", "") + if registry and ( + not _REGISTRY.fullmatch(registry) + ): + raise ValueError("invalid registry") + if registry: + host = registry.split("/", 1)[0] + if host != "localhost" and "." not in host and ":" not in host: + raise ValueError("ambiguous registry; use an FQDN, explicit port, or localhost") + if not _REPOSITORY.fullmatch(values["IMAGE_REPOSITORY"]): + raise ValueError("invalid repository") + if not _TAG.fullmatch(values["IMAGE_TAG"]): + raise ValueError("invalid tag") + if any(any(character.isspace() for character in value) for value in values.values()): + raise ValueError("managed values cannot contain whitespace") + + +def image_reference(values: Mapping[str, str]) -> str: + values = normalize_values(values) + validate_values(values) + prefix = values.get("IMAGE_REGISTRY", "") + repository = values["IMAGE_REPOSITORY"] + return f"{prefix}/{repository}:{values['IMAGE_TAG']}" if prefix else f"{repository}:{values['IMAGE_TAG']}" + + +def update_env(path: Path, values: Mapping[str, str]) -> bool: + """Atomically update managed entries; return whether the file changed.""" + values = normalize_values(values) + validate_values(values) + if path.is_symlink(): + raise ValueError("refusing to update symlink .env") + old = path.read_text() if path.exists() else "" + parse_env(old) # detect malformed managed syntax before changing anything + lines = old.splitlines() + output: list[str] = [] + written: set[str] = set() + for line in lines: + stripped = line.strip() + export = re.match(r"export[ \t]+", stripped) + candidate = stripped[export.end() :].lstrip() if export else stripped + key = candidate.split("=", 1)[0].strip() if "=" in candidate else "" + if key in MANAGED_KEYS: + if key in written: + continue + output.append(f"{key}={values[key]}") + written.add(key) + else: + output.append(line) + output.extend(f"{key}={values.get(key, '')}" for key in MANAGED_KEYS if key not in written) + new = "\n".join(output) + "\n" + if old == new: + return False + path.parent.mkdir(parents=True, exist_ok=True) + mode = os.stat(path).st_mode & 0o777 if path.exists() else 0o600 + fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + try: + os.fchmod(fd, mode) + with os.fdopen(fd, "w") as handle: + handle.write(new) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + return True diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/next_image_tag.py b/.agents/skills/greptimedb-development-docker-image/scripts/next_image_tag.py new file mode 100755 index 00000000000..ad261ebd657 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/next_image_tag.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python3 +# 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. + +"""Print a tag with its final numeric component incremented.""" + +import argparse +import re +import sys + + +def increment_tag(tag: str) -> str: + match = re.search(r"(\d+)$", tag) + if match is None: + raise ValueError("tag must end with a numeric version component") + + number = match.group(1) + return f"{tag[:match.start()]}{int(number) + 1:0{len(number)}d}" + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Increment the final numeric component of a Docker image tag." + ) + parser.add_argument("--tag", required=True, help="existing Docker image tag") + args = parser.parse_args() + + try: + print(increment_tag(args.tag)) + except ValueError as error: + print(f"error: {error}: {args.tag!r}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/prepare_context.sh b/.agents/skills/greptimedb-development-docker-image/scripts/prepare_context.sh new file mode 100755 index 00000000000..f0ea66fdf34 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/prepare_context.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash + +set -euo pipefail + +usage() { + printf 'Usage: %s --context PATH --binary PATH --dockerfile PATH --platform PLATFORM\n' "$0" >&2 +} + +context="" +binary="" +dockerfile="" +platform="" + +while [[ $# -gt 0 ]]; do + case "$1" in + --context) + context="$2" + shift 2 + ;; + --binary) + binary="$2" + shift 2 + ;; + --dockerfile) + dockerfile="$2" + shift 2 + ;; + --platform) + platform="$2" + shift 2 + ;; + *) + usage + exit 2 + ;; + esac +done + +if [[ -z "$context" || -z "$binary" || -z "$dockerfile" || -z "$platform" ]]; then + usage + exit 2 +fi + +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +python3 "$script_dir/binary_platform.py" --binary "$binary" --platform "$platform" + +if [[ ! -f "$dockerfile" ]]; then + printf 'error: Dockerfile template does not exist: %s\n' "$dockerfile" >&2 + exit 1 +fi + +if [[ -e "$context" && ( -L "$context" || -n "$(ls -A "$context")" ) ]]; then + printf 'error: build context must be a new or empty directory: %s\n' "$context" >&2 + exit 1 +fi +mkdir -p "$context" +if [[ -e "$context/greptime" || -e "$context/Dockerfile" || -L "$context/greptime" || -L "$context/Dockerfile" ]]; then + printf 'error: build context already contains output files: %s\n' "$context" >&2 + exit 1 +fi +cp "$binary" "$context/greptime" +cp "$dockerfile" "$context/Dockerfile" +printf 'created Dockerfile: %s\n' "$context/Dockerfile" + +file "$context/greptime" diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/update_image_env.py b/.agents/skills/greptimedb-development-docker-image/scripts/update_image_env.py new file mode 100755 index 00000000000..9bf8cbeea04 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/update_image_env.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +# 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. + +"""Create or update non-secret GreptimeDB image settings in a .env file.""" + +import argparse +import sys +from pathlib import Path + +from image_config import image_reference, normalize_values, update_env, validate_values + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--env", type=Path) + parser.add_argument("--registry", required=True) + parser.add_argument("--repository", required=True) + parser.add_argument("--tag", required=True) + parser.add_argument("--validate-only", action="store_true") + args = parser.parse_args() + + try: + values = normalize_values({"IMAGE_REGISTRY": args.registry, "IMAGE_REPOSITORY": args.repository, "IMAGE_TAG": args.tag}) + validate_values(values) + except ValueError as error: + print(f"error: {error}", file=sys.stderr) + return 2 + + if args.validate_only: + print(image_reference(values)) + return 0 + if args.env is None: + parser.error("--env is required unless --validate-only is set") + if update_env(args.env, values): + print(f"updated {args.env}") + else: + print(f"unchanged {args.env}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.agents/skills/greptimedb-development-docker-image/tests/test_binary_platform.py b/.agents/skills/greptimedb-development-docker-image/tests/test_binary_platform.py new file mode 100644 index 00000000000..e89bafa0696 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/tests/test_binary_platform.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +# 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. + +import importlib +import struct +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +SCRIPTS = Path(__file__).resolve().parents[1] / "scripts" +sys.path.insert(0, str(SCRIPTS)) + +binary_platform = importlib.import_module("binary_platform") + + +def elf(machine: int, *, bits: int = 2, endian: int = 1, abi: int = 3, elf_type: int = 3) -> bytes: + header = bytearray(64) + header[:4] = b"\x7fELF" + header[4:8] = bytes((bits, endian, 1, abi)) + struct.pack_into("H", header, 18, machine) + struct.pack_into(" Path: + path = Path(self.directory.name) / "greptime" + path.write_bytes(content) + path.chmod(0o755) + return path + + def setUp(self): + self.directory = tempfile.TemporaryDirectory() + + def tearDown(self): + self.directory.cleanup() + + def test_matching_architectures(self): + binary_platform.validate_binary(self.make_binary(elf(0x3E)), "linux/amd64") + binary_platform.validate_binary(self.make_binary(elf(0xB7)), "linux/arm64") + + def test_rejects_bad_executable_elf_and_platform(self): + path = self.make_binary(elf(0x3E)) + path.chmod(0o644) + with self.assertRaises(ValueError): + binary_platform.validate_binary(path, "linux/amd64") + path = self.make_binary(b"not an elf") + with self.assertRaises(ValueError): + binary_platform.validate_binary(path, "linux/amd64") + path = self.make_binary(elf(0x3E, bits=1)) + with self.assertRaises(ValueError): + binary_platform.validate_binary(path, "linux/amd64") + with self.assertRaises(ValueError): + binary_platform.validate_binary(self.make_binary(elf(0x3E, abi=9)), "linux/amd64") + with self.assertRaises(ValueError): + binary_platform.validate_binary(self.make_binary(elf(0x3E, elf_type=1)), "linux/amd64") + with self.assertRaises(ValueError): + binary_platform.validate_binary(self.make_binary(b"\x7fELF"), "linux/amd64") + with self.assertRaises(ValueError): + binary_platform.validate_binary(self.make_binary(elf(0x3E)), "darwin/amd64") + + def test_rejects_wrong_architecture_and_symlink(self): + with self.assertRaises(ValueError): + binary_platform.validate_binary(self.make_binary(elf(0x3E)), "linux/arm64") + target = self.make_binary(elf(0x3E)) + link = Path(self.directory.name) / "link" + link.symlink_to(target) + with self.assertRaises(ValueError): + binary_platform.validate_binary(link, "linux/amd64") + + def test_cli_success_and_failures(self): + path = self.make_binary(elf(0x3E)) + assert binary_platform.__file__ is not None + script = Path(binary_platform.__file__) + ok = subprocess.run([sys.executable, str(script), "--binary", str(path), "--platform", "linux/amd64"], capture_output=True, text=True) + self.assertEqual(ok.returncode, 0, ok.stderr) + bad = subprocess.run([sys.executable, str(script), "--binary", str(path), "--platform", "linux/arm64"], capture_output=True, text=True) + self.assertEqual(bad.returncode, 2) + missing = subprocess.run([sys.executable, str(script), "--binary", str(path) + ".missing", "--platform", "linux/amd64"], capture_output=True, text=True) + self.assertEqual(missing.returncode, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/.agents/skills/greptimedb-development-docker-image/tests/test_image_config.py b/.agents/skills/greptimedb-development-docker-image/tests/test_image_config.py new file mode 100644 index 00000000000..95dc913ea5d --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/tests/test_image_config.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +# 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. + +import importlib +import stat +import sys +import tempfile +import unittest +from unittest import mock +from pathlib import Path + +SCRIPTS = Path(__file__).resolve().parents[1] / "scripts" +sys.path.insert(0, str(SCRIPTS)) + +image_config = importlib.import_module("image_config") + + +VALUES = {"IMAGE_REGISTRY": "registry.example.com/team", "IMAGE_REPOSITORY": "greptime-dev", "IMAGE_TAG": "dev-001"} + + +class ImageConfigTest(unittest.TestCase): + def test_parse_supported_bounded_dotenv(self): + text = """\n# comment\n export IMAGE_REGISTRY = \"registry.example.com/team\" # note\nIMAGE_REPOSITORY='greptime-dev'\nIMAGE_TAG=dev-001 # safe comment\nOTHER=$HOME\n""" + self.assertEqual(image_config.parse_env(text), VALUES) + + def test_parser_does_not_validate_values(self): + self.assertEqual(image_config.parse_env("IMAGE_TAG=not valid\n"), {"IMAGE_TAG": "not valid"}) + + def test_parser_rejects_unmatched_and_unsupported_syntax(self): + with self.assertRaises(ValueError): + image_config.parse_env('IMAGE_TAG="unterminated\n') + with self.assertRaises(ValueError): + image_config.parse_env("IMAGE_TAG=$(date)\n") + + def test_validation_and_reference(self): + image_config.validate_values(VALUES) + self.assertEqual(image_config.image_reference(VALUES), "registry.example.com/team/greptime-dev:dev-001") + local = dict(VALUES, IMAGE_REGISTRY="") + self.assertEqual(image_config.image_reference(local), "greptime-dev:dev-001") + image_config.validate_values(dict(VALUES, IMAGE_REGISTRY="localhost")) + image_config.validate_values(dict(VALUES, IMAGE_REGISTRY="example.com")) + with self.assertRaises(ValueError): + image_config.validate_values(dict(VALUES, IMAGE_REGISTRY="registry/team")) + + def test_atomic_update_preserves_mode_and_canonicalizes_duplicates(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / ".env" + path.write_text("# keep\nIMAGE_TAG=old\nIMAGE_TAG=duplicate\nOTHER=value\n") + path.chmod(0o640) + self.assertTrue(image_config.update_env(path, VALUES)) + self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o640) + self.assertEqual(path.read_text(), "# keep\nIMAGE_TAG=dev-001\nOTHER=value\nIMAGE_REGISTRY=registry.example.com/team\nIMAGE_REPOSITORY=greptime-dev\n") + self.assertFalse(image_config.update_env(path, VALUES)) + + def test_update_rejects_symlink(self): + with tempfile.TemporaryDirectory() as directory: + target = Path(directory) / "target" + target.write_text("IMAGE_TAG=old\n") + link = Path(directory) / ".env" + link.symlink_to(target) + with self.assertRaises(ValueError): + image_config.update_env(link, VALUES) + self.assertEqual(target.read_text(), "IMAGE_TAG=old\n") + + def test_update_clears_omitted_registry_and_uses_private_new_file_mode(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / ".env" + path.write_text("IMAGE_REGISTRY=registry.example.com\nIMAGE_REPOSITORY=old\nIMAGE_TAG=old\n") + values = {"IMAGE_REPOSITORY": "greptime-dev", "IMAGE_TAG": "dev-001"} + image_config.update_env(path, values) + self.assertIn("IMAGE_REGISTRY=\n", path.read_text()) + new_path = Path(directory) / "new.env" + image_config.update_env(new_path, VALUES) + self.assertEqual(stat.S_IMODE(new_path.stat().st_mode), 0o600) + + def test_failed_replace_keeps_original_and_removes_temporary_file(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / ".env" + path.write_text("OTHER=value\n") + with mock.patch("image_config.os.replace", side_effect=OSError("no space")): + with self.assertRaises(OSError): + image_config.update_env(path, VALUES) + self.assertEqual(path.read_text(), "OTHER=value\n") + self.assertEqual(list(Path(directory).glob("..env.*")), []) + + +if __name__ == "__main__": + unittest.main() From 6688ca13f6009ee39e2610726dbf258d672a0141 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 14 Jul 2026 12:02:14 +0800 Subject: [PATCH 16/48] fix: disable WAL index creation by default (#8505) Signed-off-by: WenyXu --- config/config.md | 2 +- config/datanode.example.toml | 2 +- src/common/wal/src/config/kafka/datanode.rs | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/config/config.md b/config/config.md index b4fb40c4778..a717b0d6134 100644 --- a/config/config.md +++ b/config/config.md @@ -514,7 +514,7 @@ | `wal.timeout` | String | `3s` | The timeout for kafka client.
**It's only used when the provider is `kafka`**. | | `wal.max_batch_bytes` | String | `1MB` | The max size of a single producer batch.
Warning: Kafka has a default limit of 1MB per message in a topic.
**It's only used when the provider is `kafka`**. | | `wal.consumer_wait_timeout` | String | `100ms` | The consumer wait timeout.
**It's only used when the provider is `kafka`**. | -| `wal.create_index` | Bool | `true` | Whether to enable WAL index creation.
**It's only used when the provider is `kafka`**. | +| `wal.create_index` | Bool | `false` | Whether to enable WAL index creation.
**It's only used when the provider is `kafka`**. | | `wal.dump_index_interval` | String | `60s` | The interval for dumping WAL indexes.
**It's only used when the provider is `kafka`**. | | `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.
**It's only used when the provider is `kafka`**.

This option ensures that when Kafka messages are deleted, the system
can still successfully replay memtable data without throwing an
out-of-range error.
However, enabling this option might lead to unexpected data loss,
as the system will skip over missing entries instead of treating
them as critical errors. | | `query` | -- | -- | The query engine options. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 678af443768..7a5a486bbeb 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -185,7 +185,7 @@ consumer_wait_timeout = "100ms" ## Whether to enable WAL index creation. ## **It's only used when the provider is `kafka`**. -create_index = true +create_index = false ## The interval for dumping WAL indexes. ## **It's only used when the provider is `kafka`**. diff --git a/src/common/wal/src/config/kafka/datanode.rs b/src/common/wal/src/config/kafka/datanode.rs index 374a7229d4a..4003b52e4bc 100644 --- a/src/common/wal/src/config/kafka/datanode.rs +++ b/src/common/wal/src/config/kafka/datanode.rs @@ -60,7 +60,7 @@ impl Default for DatanodeKafkaConfig { consumer_wait_timeout: Duration::from_millis(100), kafka_topic: KafkaTopicConfig::default(), auto_create_topics: true, - create_index: true, + create_index: false, dump_index_interval: Duration::from_secs(60), topic_latest_offset_fetch_interval: DEFAULT_TOPIC_LATEST_OFFSET_FETCH_INTERVAL, overwrite_entry_start_id: false, From 5d9dddffa77b0584d23dd6890fe911b8a1b4e1ab Mon Sep 17 00:00:00 2001 From: QuakeWang <45645138+QuakeWang@users.noreply.github.com> Date: Tue, 14 Jul 2026 12:42:23 +0800 Subject: [PATCH 17/48] fix(flow): convert streaming expiration to milliseconds (#8481) * fix(flow): convert streaming expiration to milliseconds Flow EXPIRE AFTER values are stored in seconds, while streaming flow timestamps and durations use milliseconds. Passing the value through unchanged expires state and limits refill scans 1000 times too early. Convert the value at streaming create and refill boundaries, reject overflow, and document the CreateFlowArgs unit. Signed-off-by: QuakeWang * fix(flow): reject negative streaming expiration Reject negative EXPIRE AFTER values at the streaming adapter boundary so refill ranges cannot move into the future. Keep zero valid and retain checked seconds-to-milliseconds conversion. Signed-off-by: QuakeWang --------- Signed-off-by: QuakeWang --- src/flow/src/adapter.rs | 22 +++++++++++++++++++++- src/flow/src/adapter/refill.rs | 5 ++++- src/flow/src/adapter/tests.rs | 22 ++++++++++++++++++++++ src/flow/src/engine.rs | 1 + 4 files changed, 48 insertions(+), 2 deletions(-) diff --git a/src/flow/src/adapter.rs b/src/flow/src/adapter.rs index 1e663582a28..e8dd10ad7d8 100644 --- a/src/flow/src/adapter.rs +++ b/src/flow/src/adapter.rs @@ -76,6 +76,23 @@ pub(crate) mod table_source; use crate::FrontendInvoker; use crate::error::Error; +fn expire_after_secs_to_millis(expire_after_secs: i64) -> Result { + ensure!( + expire_after_secs >= 0, + InvalidQuerySnafu { + reason: format!("EXPIRE AFTER must be non-negative, got {expire_after_secs} seconds"), + } + ); + + expire_after_secs + .checked_mul(1_000) + .with_context(|| InvalidQuerySnafu { + reason: format!( + "EXPIRE AFTER value {expire_after_secs} seconds cannot be represented in milliseconds" + ), + }) +} + // `GREPTIME_TIMESTAMP` is not used to distinguish when table is created automatically by flow pub const AUTO_CREATED_PLACEHOLDER_TS_COL: &str = "__ts_placeholder"; @@ -740,7 +757,7 @@ impl StreamingEngine { source_table_ids, create_if_not_exists, or_replace, - expire_after, + expire_after: expire_after_secs, eval_interval: _, comment, sql, @@ -748,6 +765,9 @@ impl StreamingEngine { query_ctx, .. } = args; + let expire_after = expire_after_secs + .map(expire_after_secs_to_millis) + .transpose()?; let mut node_ctx = self.node_context.write().await; // assign global id to source and sink table diff --git a/src/flow/src/adapter/refill.rs b/src/flow/src/adapter/refill.rs index 6d66505e89e..3672c0427e1 100644 --- a/src/flow/src/adapter/refill.rs +++ b/src/flow/src/adapter/refill.rs @@ -85,7 +85,10 @@ impl StreamingEngine { } } - let expire_after = info.expire_after(); + let expire_after = info + .expire_after() + .map(super::expire_after_secs_to_millis) + .transpose()?; // TODO(discord9): better way to get last point let now = self.tick_manager.tick(); let plan = self diff --git a/src/flow/src/adapter/tests.rs b/src/flow/src/adapter/tests.rs index 9ed7e676c29..df6a58c46ce 100644 --- a/src/flow/src/adapter/tests.rs +++ b/src/flow/src/adapter/tests.rs @@ -61,3 +61,25 @@ pub fn new_test_table_info_with_name>( /// /// containing several default table info and schema fn mock_harness_flow_node_manager() {} + +#[test] +fn test_expire_after_secs_to_millis() { + assert_eq!(expire_after_secs_to_millis(300).unwrap(), 300_000); + assert_eq!(expire_after_secs_to_millis(0).unwrap(), 0); + + let max_secs = i64::MAX / 1_000; + assert_eq!( + expire_after_secs_to_millis(max_secs).unwrap(), + max_secs * 1_000 + ); +} + +#[test] +fn test_expire_after_secs_to_millis_invalid() { + let invalid_values = [i64::MAX / 1_000 + 1, -1]; + + for invalid_secs in invalid_values { + let error = expire_after_secs_to_millis(invalid_secs).unwrap_err(); + assert!(matches!(error, Error::InvalidQuery { .. })); + } +} diff --git a/src/flow/src/engine.rs b/src/flow/src/engine.rs index 386610ad716..3a139a898a9 100644 --- a/src/flow/src/engine.rs +++ b/src/flow/src/engine.rs @@ -34,6 +34,7 @@ pub struct CreateFlowArgs { pub source_table_ids: Vec, pub create_if_not_exists: bool, pub or_replace: bool, + /// Duration in seconds. pub expire_after: Option, pub eval_interval: Option, pub comment: Option, From 99235e4b2dc53accc66ea4e82052ec46f9061326 Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 14 Jul 2026 14:19:17 +0800 Subject: [PATCH 18/48] fix: honor exclusive end in log queries (#8495) * fix: exclude logs end boundary Signed-off-by: discord9 * fix: normalize log query date bounds Signed-off-by: discord9 --------- Signed-off-by: discord9 --- src/log-query/src/log_query.rs | 110 +++++++++++++++++++---------- src/query/src/log_query/planner.rs | 27 ++++--- tests-integration/tests/http.rs | 41 +++++++++-- 3 files changed, 123 insertions(+), 55 deletions(-) diff --git a/src/log-query/src/log_query.rs b/src/log-query/src/log_query.rs index 2955b7bdedf..78e622200d4 100644 --- a/src/log-query/src/log_query.rs +++ b/src/log-query/src/log_query.rs @@ -214,9 +214,9 @@ impl TimeFilter { (Some(start), Some(end), _) => { // Both 'start' and 'end' are provided let (start, _) = Self::parse_datetime(start)?; - let (end, _) = Self::parse_datetime(end)?; + let (end, inferred_end) = Self::parse_datetime(end)?; start_dt = Some(start); - end_dt = Some(end); + end_dt = Some(inferred_end.unwrap_or(end)); } (Some(start), None, Some(span)) => { let (start, _) = Self::parse_datetime(start)?; @@ -276,41 +276,39 @@ impl TimeFilter { if let Ok(dt) = DateTime::parse_from_rfc3339(s) { Ok((dt.with_timezone(&Utc), None)) } else { - let formats = ["%Y-%m-%d", "%Y-%m", "%Y"]; - for format in &formats { - if let Ok(naive_date) = NaiveDate::parse_from_str(s, format) { - let start = Utc.from_utc_datetime( - &naive_date.and_time(NaiveTime::from_hms_opt(0, 0, 0).unwrap()), - ); - let end = match *format { - "%Y-%m-%d" => start + Duration::days(1), - "%Y-%m" => { - let next_month = if naive_date.month() == 12 { - NaiveDate::from_ymd_opt(naive_date.year() + 1, 1, 1).unwrap() - } else { - NaiveDate::from_ymd_opt( - naive_date.year(), - naive_date.month() + 1, - 1, - ) - .unwrap() - }; - Utc.from_utc_datetime(&next_month.and_hms_opt(0, 0, 0).unwrap()) - } - "%Y" => { - let next_year = - NaiveDate::from_ymd_opt(naive_date.year() + 1, 1, 1).unwrap(); - Utc.from_utc_datetime(&next_year.and_hms_opt(0, 0, 0).unwrap()) - } - _ => unreachable!(), - }; - return Ok((start, Some(end))); + let (naive_date, format) = if let Ok(date) = NaiveDate::parse_from_str(s, "%Y-%m-%d") { + (date, "%Y-%m-%d") + } else if let Ok(date) = NaiveDate::parse_from_str(&format!("{s}-01"), "%Y-%m-%d") { + (date, "%Y-%m") + } else if let Ok(date) = NaiveDate::parse_from_str(&format!("{s}-01-01"), "%Y-%m-%d") { + (date, "%Y") + } else { + return Err(InvalidDateFormatSnafu { + input: s.to_string(), } - } - Err(InvalidDateFormatSnafu { - input: s.to_string(), - } - .build()) + .build()); + }; + + let start = Utc + .from_utc_datetime(&naive_date.and_time(NaiveTime::from_hms_opt(0, 0, 0).unwrap())); + let end = match format { + "%Y-%m-%d" => start + Duration::days(1), + "%Y-%m" => { + let next_month = if naive_date.month() == 12 { + NaiveDate::from_ymd_opt(naive_date.year() + 1, 1, 1).unwrap() + } else { + NaiveDate::from_ymd_opt(naive_date.year(), naive_date.month() + 1, 1) + .unwrap() + }; + Utc.from_utc_datetime(&next_month.and_hms_opt(0, 0, 0).unwrap()) + } + "%Y" => { + let next_year = NaiveDate::from_ymd_opt(naive_date.year() + 1, 1, 1).unwrap(); + Utc.from_utc_datetime(&next_year.and_hms_opt(0, 0, 0).unwrap()) + } + _ => unreachable!(), + }; + Ok((start, Some(end))) } } @@ -545,4 +543,44 @@ mod tests { let result = tf.canonicalize(); assert!(matches!(result, Err(Error::EndBeforeStart { .. }))); } + + #[test] + fn test_canonicalize_date_end() { + let mut date_end = TimeFilter { + start: Some("2023-10-01".to_string()), + end: Some("2023-10-02".to_string()), + span: None, + }; + date_end.canonicalize().unwrap(); + assert_eq!(date_end.end.as_deref(), Some("2023-10-03T00:00:00+00:00")); + + let mut timestamp_end = TimeFilter { + start: Some("2023-10-01".to_string()), + end: Some("2023-10-02T12:34:56Z".to_string()), + span: None, + }; + timestamp_end.canonicalize().unwrap(); + assert_eq!( + timestamp_end.end.as_deref(), + Some("2023-10-02T12:34:56+00:00") + ); + + let mut month = TimeFilter { + start: Some("2023-10".to_string()), + end: Some("2023-10".to_string()), + span: None, + }; + month.canonicalize().unwrap(); + assert_eq!(month.start.as_deref(), Some("2023-10-01T00:00:00+00:00")); + assert_eq!(month.end.as_deref(), Some("2023-11-01T00:00:00+00:00")); + + let mut year = TimeFilter { + start: Some("2023".to_string()), + end: Some("2023".to_string()), + span: None, + }; + year.canonicalize().unwrap(); + assert_eq!(year.start.as_deref(), Some("2023-01-01T00:00:00+00:00")); + assert_eq!(year.end.as_deref(), Some("2024-01-01T00:00:00+00:00")); + } } diff --git a/src/query/src/log_query/planner.rs b/src/query/src/log_query/planner.rs index f9c641b23b5..53d49f98db0 100644 --- a/src/query/src/log_query/planner.rs +++ b/src/query/src/log_query/planner.rs @@ -126,15 +126,14 @@ impl LogQueryPlanner { .clone(); let start_time = ScalarValue::Utf8(time_filter.start.clone()); - let end_time = ScalarValue::Utf8( - time_filter - .end - .clone() - .or(Some("9999-12-31T23:59:59Z".to_string())), - ); let expr = col(timestamp_col.clone()) .gt_eq(lit(start_time)) - .and(col(timestamp_col).lt_eq(lit(end_time))); + .and(match &time_filter.end { + Some(end) => col(timestamp_col).lt(lit(ScalarValue::Utf8(Some(end.clone())))), + None => col(timestamp_col).lt_eq(lit(ScalarValue::Utf8(Some( + "9999-12-31T23:59:59Z".to_string(), + )))), + }); Ok(expr) } @@ -730,7 +729,7 @@ mod tests { let plan = planner.query_to_plan(log_query).await.unwrap(); let expected = "Limit: skip=0, fetch=100 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ +\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); @@ -757,7 +756,7 @@ mod tests { .gt_eq(lit(ScalarValue::Utf8(Some( "2021-01-01T00:00:00Z".to_string(), )))) - .and(col("timestamp").lt_eq(lit(ScalarValue::Utf8(Some( + .and(col("timestamp").lt(lit(ScalarValue::Utf8(Some( "2021-01-02T00:00:00Z".to_string(), ))))); @@ -851,7 +850,7 @@ mod tests { let plan = planner.query_to_plan(log_query).await.unwrap(); let expected = "Limit: skip=10, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ +\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); @@ -886,7 +885,7 @@ mod tests { let plan = planner.query_to_plan(log_query).await.unwrap(); let expected = "Limit: skip=0, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ +\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); @@ -934,7 +933,7 @@ mod tests { let plan = planner.query_to_plan(log_query).await.unwrap(); let expected = "Aggregate: groupBy=[[greptime.public.test_table.host]], aggr=[[count(greptime.public.test_table.message) AS count_result]] [host:Utf8;N, count_result:Int64]\ \n Limit: skip=0, fetch=100 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ +\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); @@ -974,7 +973,7 @@ mod tests { let plan = planner.query_to_plan(log_query).await.unwrap(); let expected = "Projection: date_trunc(Utf8(\"day\"), greptime.public.test_table.timestamp) AS time_bucket [time_bucket:Timestamp(ms)]\ \n Limit: skip=0, fetch=100 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ - \n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ + \n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); @@ -1058,7 +1057,7 @@ mod tests { let expected = "Aggregate: groupBy=[[2__date_histogram__time_bucket]], aggr=[[count(2__date_histogram__time_bucket) AS count_result]] [2__date_histogram__time_bucket:Timestamp(ns);N, count_result:Int64]\ \n Projection: date_bin(Utf8(\"30 seconds\"), greptime.public.test_table.timestamp) AS 2__date_histogram__time_bucket [2__date_histogram__time_bucket:Timestamp(ns);N]\ \n Limit: skip=0, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ +\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index c6366f48efe..a541b82b09e 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -7810,14 +7810,16 @@ pub async fn test_log_query(store_type: StorageType) { .await; assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); let res = client - .post("/v1/sql?sql=insert into logs values ('2024-11-07 10:53:50', 'hello');") + .post( + "/v1/sql?sql=insert into logs values ('2024-11-06 23:59:59', 'before-date-end'), ('2024-11-07 10:53:50', 'before-explicit-end'), ('2024-11-07 10:53:51', 'at-explicit-end'), ('2024-11-07 10:53:52', 'after-explicit-end');", + ) .header("Content-Type", "application/x-www-form-urlencoded") .send() .await; assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); // test log query - let log_query = LogQuery { + let mut log_query = LogQuery { table: TableName { catalog_name: "greptime".to_string(), schema_name: "public".to_string(), @@ -7825,12 +7827,12 @@ pub async fn test_log_query(store_type: StorageType) { }, time_filter: TimeFilter { start: Some("2024-11-07".to_string()), - end: None, + end: Some("2024-11-07T10:53:51Z".to_string()), span: None, }, limit: Limit { skip: None, - fetch: Some(1), + fetch: Some(3), }, columns: vec!["ts".to_string(), "message".to_string()], filters: Default::default(), @@ -7847,7 +7849,36 @@ pub async fn test_log_query(store_type: StorageType) { assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); let resp = res.text().await; let v = get_rows_from_output(&resp); - assert_eq!(v, "[[1730976830000,\"hello\"]]"); + assert_eq!(v, "[[1730976830000,\"before-explicit-end\"]]"); + + log_query.time_filter.start = Some("2024-11-06".to_string()); + log_query.time_filter.end = Some("2024-11-07".to_string()); + log_query.limit.fetch = Some(4); + let res = client + .post("/v1/logs") + .header("Content-Type", "application/json") + .body(serde_json::to_string(&log_query).unwrap()) + .send() + .await; + + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + let resp = res.text().await; + let output = serde_json::from_str::(&resp).unwrap(); + let rows = output["output"][0]["records"]["rows"].as_array().unwrap(); + let mut messages = rows + .iter() + .map(|row| row[1].as_str().unwrap()) + .collect::>(); + messages.sort_unstable(); + assert_eq!( + messages, + vec![ + "after-explicit-end", + "at-explicit-end", + "before-date-end", + "before-explicit-end", + ] + ); guard.remove_all().await; } From 0fcf61c3fbfbdf436640d77cd8d8bd696f5ef82f Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 14 Jul 2026 14:50:59 +0800 Subject: [PATCH 19/48] feat: skip oversized compaction tasks (#8466) fix: skip oversized compaction tasks Signed-off-by: jeremyhi --- src/mito2/src/compaction.rs | 87 ++++++++++++++++++++++++- src/mito2/src/test_util/version_util.rs | 11 +++- 2 files changed, 95 insertions(+), 3 deletions(-) diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index 10900b5540c..9c0c0520458 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -64,7 +64,9 @@ use crate::error::{ RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, RemoteCompactionSnafu, Result, TimeRangePredicateOverflowSnafu, TimeoutSnafu, }; -use crate::metrics::{COMPACTION_STAGE_ELAPSED, INFLIGHT_COMPACTION_COUNT}; +use crate::metrics::{ + COMPACTION_MEMORY_REJECTED, COMPACTION_STAGE_ELAPSED, INFLIGHT_COMPACTION_COUNT, +}; use crate::read::FlatSource; use crate::read::flat_projection::FlatProjectionMapper; use crate::read::read_columns::ReadColumns; @@ -613,8 +615,21 @@ impl CompactionScheduler { waiters }; - // Create a local compaction task. + // Check whether this local compaction can ever fit before submitting it. let estimated_bytes = estimate_compaction_bytes(&picker_output); + if let Some(limit_bytes) = self.exceeds_compaction_memory_limit(estimated_bytes) { + COMPACTION_MEMORY_REJECTED + .with_label_values(&["oversized"]) + .inc(); + warn!( + "Skip compaction for region {} because estimated memory {} bytes exceeds compaction memory limit {} bytes", + region_id, estimated_bytes, limit_bytes, + ); + for waiter in waiters { + waiter.send(Ok(0)); + } + return Ok(None); + } let cancel_handle = Arc::new(CancellationHandle::default()); let state = LocalCompactionState::new(cancel_handle.clone()); @@ -652,6 +667,15 @@ impl CompactionScheduler { ) } + fn exceeds_compaction_memory_limit(&self, estimated_bytes: u64) -> Option { + let limit_bytes = self.memory_manager.limit_bytes(); + if limit_bytes > 0 && estimated_bytes > limit_bytes { + Some(limit_bytes) + } else { + None + } + } + fn remove_region_on_failure(&mut self, region_id: RegionId, err: Arc) { // Remove this region. let Some(status) = self.region_status.remove(®ion_id) else { @@ -1551,6 +1575,65 @@ mod tests { assert!(scheduler.region_status.contains_key(®ion_id)); } + #[tokio::test] + async fn test_schedule_compaction_skips_task_exceeding_memory_limit() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + scheduler.memory_manager = Arc::new(new_compaction_memory_manager(1024 * 1024)); + + let mut builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let end = 1000 * 1000; + let version_control = Arc::new( + builder + .push_l0_file_with_max_row_group_size(0, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(10, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(50, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(80, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(90, end, 1024 * 1024) + .build(), + ); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + schema_metadata_manager + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + kv_backend, + ) + .await; + let (output_tx, output_rx) = oneshot::channel(); + let rejected = COMPACTION_MEMORY_REJECTED.with_label_values(&["oversized"]); + let rejected_before = rejected.get(); + + let scheduled = scheduler + .schedule_compaction( + region_id, + Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::from(output_tx), + &manifest_ctx, + schema_metadata_manager, + 1, + ) + .await + .unwrap(); + + assert!(!scheduled); + assert_eq!(output_rx.await.unwrap().unwrap(), 0); + assert_eq!(rejected_before + 1, rejected.get()); + assert_eq!(0, job_scheduler.num_jobs()); + assert!(!scheduler.region_status.contains_key(®ion_id)); + } + #[tokio::test] async fn test_schedule_on_finished() { common_telemetry::init_default_ut_logging(); diff --git a/src/mito2/src/test_util/version_util.rs b/src/mito2/src/test_util/version_util.rs index b14eea09581..1c0d7bfebc8 100644 --- a/src/mito2/src/test_util/version_util.rs +++ b/src/mito2/src/test_util/version_util.rs @@ -89,6 +89,15 @@ impl VersionControlBuilder { } pub(crate) fn push_l0_file(&mut self, start_ms: i64, end_ms: i64) -> &mut Self { + self.push_l0_file_with_max_row_group_size(start_ms, end_ms, 0) + } + + pub(crate) fn push_l0_file_with_max_row_group_size( + &mut self, + start_ms: i64, + end_ms: i64, + max_row_group_uncompressed_size: u64, + ) -> &mut Self { let file_id = FileId::random(); self.files.insert( file_id, @@ -101,7 +110,7 @@ impl VersionControlBuilder { ), level: 0, file_size: 0, // We don't care file size. - max_row_group_uncompressed_size: 0, + max_row_group_uncompressed_size, available_indexes: Default::default(), indexes: Default::default(), index_file_size: 0, From 0beeef7fc097b4546a7b20e8c47492420c0846b2 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Tue, 14 Jul 2026 15:09:35 +0800 Subject: [PATCH 20/48] refactor: optimize json2 write, again (#8498) refactor: optimize json2 write performance Signed-off-by: luofucong --- src/api/src/helper.rs | 110 +++++++++++---- src/datatypes/src/json/value.rs | 20 ++- src/datatypes/src/vectors/json/builder.rs | 156 +++++++++++++--------- 3 files changed, 184 insertions(+), 102 deletions(-) diff --git a/src/api/src/helper.rs b/src/api/src/helper.rs index ed8d64b57a3..1ff437f127a 100644 --- a/src/api/src/helper.rs +++ b/src/api/src/helper.rs @@ -20,7 +20,7 @@ use common_decimal::decimal128::{DECIMAL128_DEFAULT_SCALE, DECIMAL128_MAX_PRECIS use common_time::time::Time; use common_time::timestamp::TimeUnit; use common_time::{Date, IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth, Timestamp}; -use datatypes::json::value::{JsonNumber, JsonValue, JsonValueRef, JsonVariant}; +use datatypes::json::value::{JsonNumber, JsonValue, JsonValueRef, JsonVariant, JsonVariantRef}; use datatypes::prelude::{ConcreteDataType, ValueRef}; use datatypes::types::json_type::JsonNativeType; use datatypes::types::{ @@ -931,33 +931,59 @@ pub fn encode_json_value(value: JsonValue) -> v1::JsonValue { } fn decode_json_value(value: &v1::JsonValue) -> JsonValueRef<'_> { + let (variant, json_type) = decode_json_value_parts(value); + JsonValueRef::new_with_type(variant, json_type) +} + +fn decode_json_value_parts(value: &v1::JsonValue) -> (JsonVariantRef<'_>, JsonNativeType) { let Some(value) = &value.value else { - return JsonValueRef::null(); + return (JsonVariantRef::Null, JsonNativeType::Null); }; match value { - json_value::Value::Boolean(x) => (*x).into(), - json_value::Value::Int(x) => (*x).into(), - json_value::Value::Uint(x) => (*x).into(), - json_value::Value::Float(x) => (*x).into(), - json_value::Value::Str(x) => (x.as_str()).into(), - json_value::Value::Array(array) => array - .items - .iter() - .map(|x| decode_json_value(x).into_variant()) - .collect::>() - .into(), - json_value::Value::Object(x) => x - .entries - .iter() - .filter_map(|entry| { - entry - .value - .as_ref() - .map(|v| (entry.key.as_str(), decode_json_value(v).into_variant())) - }) - .collect::>() - .into(), - json_value::Value::Variant(x) => x.as_slice().into(), + json_value::Value::Boolean(x) => (JsonVariantRef::Bool(*x), JsonNativeType::Bool), + json_value::Value::Int(x) => ((*x).into(), JsonNativeType::i64()), + json_value::Value::Uint(x) => ((*x).into(), JsonNativeType::u64()), + json_value::Value::Float(x) => ((*x).into(), JsonNativeType::f64()), + json_value::Value::Str(x) => (x.as_str().into(), JsonNativeType::String), + json_value::Value::Array(array) => { + let mut variants = Vec::with_capacity(array.items.len()); + let mut item_type = JsonNativeType::Null; + for item in &array.items { + let (variant, ty) = decode_json_value_parts(item); + variants.push(variant); + if !matches!(item_type, JsonNativeType::Variant) { + item_type.merge(&ty); + } + } + ( + JsonVariantRef::Array(variants), + JsonNativeType::Array(Box::new(item_type)), + ) + } + json_value::Value::Object(object) => { + let mut variants = Vec::with_capacity(object.entries.len()); + let mut fields = Vec::with_capacity(object.entries.len()); + for entry in &object.entries { + let Some(value) = &entry.value else { + continue; + }; + let (variant, json_type) = decode_json_value_parts(value); + variants.push((entry.key.as_str(), variant)); + fields.push((entry.key.clone(), json_type)); + } + let variants = variants.into_iter().collect::>(); + let fields = fields.into_iter().collect::>(); + let json_type = if fields.is_empty() { + JsonNativeType::Null + } else { + JsonNativeType::Object(fields) + }; + (JsonVariantRef::Object(variants), json_type) + } + json_value::Value::Variant(x) => ( + JsonVariantRef::Variant(x.as_slice()), + JsonNativeType::Variant, + ), } } @@ -1854,15 +1880,40 @@ mod tests { ] })) ); + assert_eq!( + JsonNativeType::Array(Box::new(JsonNativeType::i64())), + decode_json_value_parts(&proto).1 + ); let value = decode_json_value(&proto); assert_eq!(json.as_ref(), value); + let proto = v1::JsonValue { + value: Some(json_value::Value::Array(JsonList { + items: vec![ + v1::JsonValue { + value: Some(json_value::Value::Int(1)), + }, + v1::JsonValue { + value: Some(json_value::Value::Float(2.0)), + }, + ], + })), + }; + assert_eq!( + JsonNativeType::Array(Box::new(JsonNativeType::f64())), + decode_json_value_parts(&proto).1 + ); + let json: JsonValue = [(); 0].into(); let proto = encode_json_value(json.clone()); assert_eq!( proto.value, Some(json_value::Value::Array(JsonList { items: vec![] })) ); + assert_eq!( + JsonNativeType::Array(Box::new(JsonNativeType::Null)), + decode_json_value_parts(&proto).1 + ); let value = decode_json_value(&proto); assert_eq!(json.as_ref(), value); @@ -1893,6 +1944,14 @@ mod tests { ] })) ); + assert_eq!( + JsonNativeType::Object(JsonObjectType::from([ + ("k1".to_string(), JsonNativeType::i64()), + ("k2".to_string(), JsonNativeType::i64()), + ("k3".to_string(), JsonNativeType::i64()), + ])), + decode_json_value_parts(&proto).1 + ); let value = decode_json_value(&proto); assert_eq!(json.as_ref(), value); @@ -1902,6 +1961,7 @@ mod tests { proto.value, Some(json_value::Value::Object(JsonObject { entries: vec![] })) ); + assert_eq!(JsonNativeType::Null, decode_json_value_parts(&proto).1); let value = decode_json_value(&proto); assert_eq!(json.as_ref(), value); diff --git a/src/datatypes/src/json/value.rs b/src/datatypes/src/json/value.rs index 40217f75141..b727bd8b5ab 100644 --- a/src/datatypes/src/json/value.rs +++ b/src/datatypes/src/json/value.rs @@ -317,13 +317,6 @@ impl JsonValue { } } - pub(crate) fn new_with(json_variant: JsonVariant, json_type: Arc) -> Self { - Self { - json_type: OnceLock::from(json_type), - json_variant, - } - } - pub(crate) fn data_type(&self) -> ConcreteDataType { ConcreteDataType::json2(self.json_type().clone()) } @@ -383,10 +376,6 @@ impl JsonValue { self.json_variant } - pub(crate) fn variant(&self) -> &JsonVariant { - &self.json_variant - } - pub(crate) fn into_value(self) -> Value { fn helper(v: JsonVariant) -> Value { match v { @@ -828,6 +817,15 @@ impl<'a> JsonValueRef<'a> { ().into() } + /// Creates a JSON value reference with its precomputed native type. + /// The native type must describe `json_variant` exactly. + pub fn new_with_type(json_variant: JsonVariantRef<'a>, json_type: JsonNativeType) -> Self { + Self { + json_type: OnceLock::from(Arc::new(json_type)), + json_variant, + } + } + pub(crate) fn data_type(&self) -> ConcreteDataType { ConcreteDataType::json2(self.json_type().as_ref().clone()) } diff --git a/src/datatypes/src/vectors/json/builder.rs b/src/datatypes/src/vectors/json/builder.rs index 41aff8c29c6..ebf11ab7961 100644 --- a/src/datatypes/src/vectors/json/builder.rs +++ b/src/datatypes/src/vectors/json/builder.rs @@ -19,17 +19,17 @@ use arrow_schema::DataType; use crate::data_type::ConcreteDataType; use crate::error::{Result, TryFromValueSnafu, UnexpectedSnafu, UnsupportedOperationSnafu}; -use crate::json::value::{JsonNumber, JsonValue, JsonVariant}; +use crate::json::value::{JsonNumber, JsonVariant, encode_json_variant}; use crate::prelude::{ValueRef, Vector, VectorRef}; use crate::types::StructType; use crate::types::json_type::{JsonNativeType, is_include}; -use crate::value::{ListValueRef, StructValueRef}; +use crate::value::{ListValue, StructValue, StructValueRef, Value}; use crate::vectors::{MutableVector, StructVectorBuilder}; #[derive(Clone)] pub(crate) struct JsonVectorBuilder { merged_type: JsonNativeType, - values: Vec, + values: Vec, } impl JsonVectorBuilder { @@ -56,23 +56,22 @@ impl JsonVectorBuilder { let mut builder = StructVectorBuilder::with_type_and_capacity(struct_type.clone(), self.values.len()); - for value in self.values.iter_mut() { - if value.is_null() { + for value in std::mem::take(&mut self.values) { + if matches!(&value, JsonVariant::Null) { builder.push_null(); continue; } - value.try_align(&self.merged_type)?; - let value_ref = json_variant_to_struct_value_ref(value.variant(), struct_type.clone())?; - builder.push_struct_value_ref(value_ref)?; + let value = json_variant_into_struct_value(value, struct_type.clone())?; + builder.push_struct_value_ref(StructValueRef::Ref(&value))?; } Ok(builder.to_vector()) } } -fn json_variant_to_struct_value_ref( - value: &JsonVariant, +fn json_variant_into_struct_value( + value: JsonVariant, struct_type: StructType, -) -> Result> { +) -> Result { let JsonVariant::Object(object) = value else { return TryFromValueSnafu { reason: format!("expected json object value, got {value:?}"), @@ -80,63 +79,87 @@ fn json_variant_to_struct_value_ref( .fail(); }; - let values = struct_type - .fields() - .iter() - .map(|field| { - object - .get(field.name()) - .map(|v| json_variant_to_value_ref(v, field.data_type())) - .unwrap_or(Ok(ValueRef::Null)) - }) - .collect::>>()?; + let mut entries = object.into_iter(); + let mut entry = entries.next(); + let mut values = Vec::with_capacity(struct_type.fields().len()); + for field in struct_type.fields().iter() { + let value = match entry.take() { + Some((name, value)) if name == field.name() => { + entry = entries.next(); + json_variant_into_value(value, field.data_type())? + } + Some((name, _)) if name.as_str() < field.name() => { + return TryFromValueSnafu { + reason: format!("field {name} is missing from merged JSON type"), + } + .fail(); + } + next => { + entry = next; + Value::Null + } + }; + values.push(value); + } + if let Some((name, _)) = entry { + return TryFromValueSnafu { + reason: format!("field {name} is missing from merged JSON type"), + } + .fail(); + } - Ok(StructValueRef::RefList { - val: values, - fields: struct_type, - }) + Ok(StructValue::new(values, struct_type)) } -fn json_variant_to_value_ref<'a>( - value: &'a JsonVariant, - expected_type: &ConcreteDataType, -) -> Result> { - let value = match value { - JsonVariant::Null => ValueRef::Null, - JsonVariant::Bool(x) => ValueRef::Boolean(*x), - JsonVariant::Number(x) => match x { - JsonNumber::PosInt(i) => ValueRef::UInt64(*i), - JsonNumber::NegInt(i) => ValueRef::Int64(*i), - JsonNumber::Float(f) => ValueRef::Float64(*f), - }, - JsonVariant::String(x) => ValueRef::String(x), - JsonVariant::Array(array) => { - let item_type = match expected_type { - ConcreteDataType::List(list_type) => list_type.item_type().clone(), - _ => ConcreteDataType::null_datatype(), - }; - let values = array - .iter() - .map(|v| json_variant_to_value_ref(v, &item_type)) - .collect::>>()?; - ValueRef::List(ListValueRef::RefList { - val: values, - item_datatype: Arc::new(item_type), - }) +fn json_variant_into_value(value: JsonVariant, expected_type: &ConcreteDataType) -> Result { + let value = match (value, expected_type) { + (JsonVariant::Null, _) | (_, ConcreteDataType::Null(_)) => Value::Null, + (JsonVariant::Bool(x), ConcreteDataType::Boolean(_)) => Value::Boolean(x), + (JsonVariant::Number(JsonNumber::PosInt(x)), ConcreteDataType::UInt64(_)) => { + Value::UInt64(x) } - JsonVariant::Object(_) => { - let ConcreteDataType::Struct(struct_type) = expected_type else { + (JsonVariant::Number(x), ConcreteDataType::Int64(_)) => { + let x = match x { + JsonNumber::PosInt(x) => i64::try_from(x).ok(), + JsonNumber::NegInt(x) => Some(x), + JsonNumber::Float(_) => None, + }; + let Some(x) = x else { return TryFromValueSnafu { - reason: format!("expected struct type, got {expected_type}"), + reason: format!("unable to convert {x:?} to Int64"), } .fail(); }; - ValueRef::Struct(json_variant_to_struct_value_ref( - value, - struct_type.clone(), - )?) + Value::Int64(x) + } + (JsonVariant::Number(JsonNumber::PosInt(x)), ConcreteDataType::Float64(_)) => { + Value::Float64((x as f64).into()) + } + (JsonVariant::Number(JsonNumber::NegInt(x)), ConcreteDataType::Float64(_)) => { + Value::Float64((x as f64).into()) + } + (JsonVariant::Number(JsonNumber::Float(x)), ConcreteDataType::Float64(_)) => { + Value::Float64(x) + } + (JsonVariant::String(x), ConcreteDataType::String(_)) => Value::String(x.into()), + (JsonVariant::Array(array), ConcreteDataType::List(list_type)) => { + let item_type = list_type.item_type().clone(); + let values = array + .into_iter() + .map(|v| json_variant_into_value(v, &item_type)) + .collect::>>()?; + Value::List(ListValue::new(values, Arc::new(item_type))) + } + (value @ JsonVariant::Object(_), ConcreteDataType::Struct(struct_type)) => { + Value::Struct(json_variant_into_struct_value(value, struct_type.clone())?) + } + (value, ConcreteDataType::Binary(_)) => Value::from(encode_json_variant(value)?), + (value, expected_type) => { + return TryFromValueSnafu { + reason: format!("unable to convert json value {value:?} to {expected_type}"), + } + .fail(); } - JsonVariant::Variant(x) => ValueRef::Binary(x), }; Ok(value) } @@ -185,13 +208,12 @@ impl MutableVector for JsonVectorBuilder { self.merged_type.merge(json_type); } - let value = JsonValue::new_with(JsonVariant::from(value.variant()), value.json_type()); - self.values.push(value); + self.values.push(JsonVariant::from(value.variant())); Ok(()) } fn push_null(&mut self) { - self.values.push(JsonValue::null()) + self.values.push(JsonVariant::Null) } fn extend_slice_of(&mut self, _: &dyn Vector, _: usize, _: usize) -> Result<()> { @@ -227,7 +249,7 @@ mod tests { } // Object inputs should merge into a superset schema, preserve null rows, - // and align conflicting nested values into Variant payloads. + // and project conflicting nested values into Variant payloads. let mut builder = JsonVectorBuilder::new(JsonNativeType::Object(Default::default()), 3); let first = parse_json_value(r#"{"id":1,"payload":{"name":"foo"}}"#); let second = parse_json_value(r#"{"id":2,"extra":true,"payload":"raw"}"#); @@ -332,7 +354,7 @@ mod tests { } #[test] - fn test_json_variant_to_struct_value_ref() -> Result<()> { + fn test_json_variant_into_struct_value() -> Result<()> { let item_type = ConcreteDataType::struct_datatype(StructType::new(Arc::new(vec![StructField::new( "id".to_string(), @@ -370,8 +392,10 @@ mod tests { JsonVariant::from([("name", JsonVariant::from("foo"))]), ), ]); - let value_ref = json_variant_to_struct_value_ref(&variant, struct_type.clone())?; - let value = value_ref.to_value(); + let value = Value::Struct(json_variant_into_struct_value( + variant, + struct_type.clone(), + )?); assert_eq!( value, From 81d8beee3a4a323f73be00a93821bd3316d9e4b2 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:39:36 +0800 Subject: [PATCH 21/48] feat: clean up soft-dropped regions offline (#8458) * fix(meta): skip reopening dropped tables during purge Purge soft-dropped tables by dropping stored routes directly instead of reopening tombstoned regions first. Treat legacy \`PurgeDroppedTableState::OpenRegions\` snapshots as a compatibility-only transition to \`DropRegions\`. Files: - \`src/common/meta/src/ddl/purge_dropped_table.rs\` - \`src/common/meta/src/ddl/undrop_table.rs\` - \`src/common/meta/src/ddl/tests/drop_table.rs\` Signed-off-by: Lei, HUANG * fix(meta): guard undrop restore race Serialize `UNDROP TABLE` with same-name creates by seeding the original table name before procedure submission, and clean up reopened regions when metadata restore fails. Files: - `src/common/meta/src/ddl/undrop_table.rs` - `src/common/meta/src/ddl_manager.rs` - `src/common/meta/src/ddl/tests/drop_table.rs` Signed-off-by: Lei, HUANG * feat: clean up soft-dropped regions offline Use an explicit RegionCleanUp request for purge-table cleanup so tombstoned regions can be removed without reopening them. Route cleanup through datanode, Mito, and metric-engine offline paths, including WAL obsoletion and region directory removal. Signed-off-by: Lei, HUANG * fix(datanode): reject cleanup for registered regions Return `RegionBusy` when `RegionCleanUp` targets a region already tracked by the datanode, so offline cleanup only runs for regions without a local mapping. Add coverage for `OfflineCleanup` engine selection and registered-region rejection. Files: - `src/datanode/src/region_server.rs` Signed-off-by: Lei, HUANG * fix(meta): require tombstone before undrop Signed-off-by: Lei, HUANG * fix(meta): reject file-engine soft drop Signed-off-by: Lei, HUANG * fix: harden soft-drop cleanup paths Reject `RegionCleanUp` for already-open Mito regions instead of turning cleanup into a drop. Make `UndropTableProcedure` tolerate missing persisted table names and always deregister failure detectors after restore-failure cleanup. Files: - `src/common/meta/src/ddl/undrop_table.rs` - `src/mito2/src/engine/open_test.rs` - `src/mito2/src/worker/handle_open.rs` Signed-off-by: Lei, HUANG * chore: bump cleanup proto dependency Bump \`greptime-proto\` to the reviewed cleanup RPC revision and align cleanup request parsing and dispatch with the renamed \`CleanUpRequest\` payload. Files: - \`Cargo.toml\` - \`Cargo.lock\` - \`src/store-api/src/region_request.rs\` - \`src/common/meta/src/ddl/drop_table/executor.rs\` Signed-off-by: Lei, HUANG * refactor: share region cleanup helpers Share common region cleanup helpers between normal drop and offline cleanup while keeping their preconditions separate. - Extract shared dropped-region runtime cleanup for `handle_drop_request` and `handle_offline_cleanup_request`. - Share runtime artifact and manifest cache cleanup after full deletion paths. - Make full-drop directory removal policy explicit: full drop and purge cleanup force physical deletion, while partial drop may defer to global GC. Files: - `src/mito2/src/worker/handle_drop.rs` - `src/mito2/src/worker/handle_open.rs` Signed-off-by: Lei, HUANG * chore: preserve soft-drop cleanup split state Signed-off-by: Lei, HUANG * refactor: remove obsolete CleanUp match arm from RegionRequest The `CleanUp` variant in the `region_request::Body` match is now handled exclusively by `RegionServer` via a separate path. This arm would have returned an unexpected error, so removing it eliminates dead code. Signed-off-by: Lei, HUANG * fix(meta): clean every soft-dropped region replica Signed-off-by: Lei, HUANG * fix(meta): order soft-drop replica cleanup Signed-off-by: Lei, HUANG * Revert "fix(meta): order soft-drop replica cleanup" This reverts commit e77162d3e5ebcf2817e2845a6a5177c328fb2c60. Signed-off-by: Lei, HUANG * Revert "fix(meta): clean every soft-dropped region replica" This reverts commit 2378e00cc258ca1b6a85a1aafbd68c79c666f43c. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- .../meta/src/ddl/drop_table/executor.rs | 80 ++++- .../meta/src/ddl/drop_table/metadata.rs | 7 + .../meta/src/ddl/purge_dropped_table.rs | 34 +- src/common/meta/src/ddl/tests/drop_table.rs | 329 +++++++++++------- src/common/meta/src/ddl/undrop_table.rs | 285 ++++++++++++--- src/common/meta/src/ddl_manager.rs | 69 +++- src/datanode/src/region_server.rs | 139 +++++++- src/metric-engine/src/engine.rs | 59 +++- src/metric-engine/src/engine/open.rs | 95 ++++- src/metric-engine/src/error.rs | 10 + src/mito2/src/engine/open_test.rs | 106 +++++- src/mito2/src/region/opener.rs | 57 +-- src/mito2/src/request.rs | 13 +- src/mito2/src/wal.rs | 13 + src/mito2/src/worker.rs | 4 + src/mito2/src/worker/handle_drop.rs | 71 +++- src/mito2/src/worker/handle_open.rs | 54 ++- src/store-api/src/region_request.rs | 71 +++- 18 files changed, 1229 insertions(+), 267 deletions(-) diff --git a/src/common/meta/src/ddl/drop_table/executor.rs b/src/common/meta/src/ddl/drop_table/executor.rs index 74b2b2482b9..1b0993170e6 100644 --- a/src/common/meta/src/ddl/drop_table/executor.rs +++ b/src/common/meta/src/ddl/drop_table/executor.rs @@ -15,8 +15,8 @@ use std::collections::{HashMap, HashSet}; use api::v1::region::{ - CloseRequest as PbCloseRegionRequest, DropRequest as PbDropRegionRequest, RegionRequest, - RegionRequestHeader, region_request, + CleanUpRequest as PbCleanUpRequest, CloseRequest as PbCloseRegionRequest, + DropRequest as PbDropRegionRequest, RegionRequest, RegionRequestHeader, region_request, }; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; @@ -26,12 +26,14 @@ use common_wal::options::WalOptions; use futures::future::join_all; use snafu::ensure; use store_api::storage::{RegionId, RegionNumber}; -use table::metadata::TableId; +use table::metadata::{TableId, TableInfo}; use table::table_name::TableName; use crate::cache_invalidator::Context; -use crate::ddl::DdlContext; -use crate::ddl::utils::{add_peer_context_if_needed, convert_region_routes_to_detecting_regions}; +use crate::ddl::utils::{ + add_peer_context_if_needed, convert_region_routes_to_detecting_regions, region_storage_path, +}; +use crate::ddl::{CreateRequestBuilder, DdlContext, build_template_from_raw_table_info}; use crate::error::{self, Result}; use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; @@ -348,6 +350,74 @@ impl DropTableExecutor { Ok(()) } + /// Cleans leader regions on datanodes without reopening them as live regions. + pub async fn on_cleanup_regions_offline( + &self, + node_manager: &NodeManagerRef, + leader_region_registry: &LeaderRegionRegistryRef, + table_info: &TableInfo, + region_routes: &[RegionRoute], + region_wal_options: &HashMap, + ) -> Result<()> { + let template = build_template_from_raw_table_info(table_info)?; + let builder = CreateRequestBuilder::new(template, None); + let storage_path = region_storage_path(&self.table.catalog_name, &self.table.schema_name); + + let leaders = find_leaders(region_routes); + let mut cleanup_region_tasks = Vec::with_capacity(leaders.len()); + let table_id = self.table_id; + for datanode in leaders { + let requester = node_manager.datanode(&datanode).await; + let regions = find_leader_regions(region_routes, &datanode); + let region_ids = regions + .iter() + .map(|region_number| RegionId::new(table_id, *region_number)) + .collect::>(); + + for region_id in region_ids { + debug!("Cleaning region {region_id} offline on Datanode {datanode:?}"); + let create_request = builder.build_one( + region_id, + storage_path.clone(), + region_wal_options, + &HashMap::new(), + )?; + let request = RegionRequest { + header: Some(RegionRequestHeader { + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::CleanUp(PbCleanUpRequest { + region_id: create_request.region_id, + engine: create_request.engine, + path: create_request.path, + options: create_request.options, + })), + }; + let datanode = datanode.clone(); + let requester = requester.clone(); + cleanup_region_tasks.push(async move { + if let Err(err) = requester.handle(request).await + && err.status_code() != StatusCode::RegionNotFound + { + return Err(add_peer_context_if_needed(datanode)(err)); + } + Ok(()) + }); + } + } + + join_all(cleanup_region_tasks) + .await + .into_iter() + .collect::>>()?; + + let region_ids = operating_leader_regions(region_routes); + leader_region_registry.batch_delete(region_ids.into_iter().map(|(region_id, _)| region_id)); + + Ok(()) + } + /// Closes all table regions on datanodes without deleting region files or metadata tombstones. /// When `flush_leaders_on_close` is set, only leader regions are flushed before close. pub async fn on_close_regions( diff --git a/src/common/meta/src/ddl/drop_table/metadata.rs b/src/common/meta/src/ddl/drop_table/metadata.rs index 6289b6c61ab..982182ffbf1 100644 --- a/src/common/meta/src/ddl/drop_table/metadata.rs +++ b/src/common/meta/src/ddl/drop_table/metadata.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_catalog::consts::FILE_ENGINE; use common_catalog::format_full_table_name; use snafu::{OptionExt, ensure}; use store_api::metric_engine_consts::METRIC_ENGINE_NAME; @@ -61,6 +62,12 @@ impl DropTableProcedure { if physical_table_id == self.data.table_id() { let engine = table_info_value.table_info.meta.engine; + ensure!( + !(self.context.soft_drop_enabled && engine == FILE_ENGINE), + error::UnsupportedSnafu { + operation: "soft-dropping file-engine tables".to_string() + } + ); // rollback only if dropping the metric physical table fails self.data.allow_rollback = engine.as_str() == METRIC_ENGINE_NAME; diff --git a/src/common/meta/src/ddl/purge_dropped_table.rs b/src/common/meta/src/ddl/purge_dropped_table.rs index b95e5db976d..759e908c8b9 100644 --- a/src/common/meta/src/ddl/purge_dropped_table.rs +++ b/src/common/meta/src/ddl/purge_dropped_table.rs @@ -29,7 +29,6 @@ use table::table_name::TableName; use crate::ddl::DdlContext; use crate::ddl::drop_table::executor::DropTableExecutor; -use crate::ddl::undrop_table::open_regions_ignore_region_not_found; use crate::ddl::utils::{ convert_region_routes_to_detecting_regions, is_metric_engine_logical_table, map_to_procedure_error, @@ -84,22 +83,6 @@ impl PurgeDroppedTableProcedure { self.data.table_info = Some(dropped_table.table_info_value.table_info); self.data.table_route_value = Some(dropped_table.table_route_value); self.data.region_wal_options = dropped_table.region_wal_options; - self.data.state = PurgeDroppedTableState::OpenRegions; - Ok(Status::executing(true)) - } - - async fn on_open_regions(&mut self) -> Result { - if let Some(region_routes) = self.data.physical_region_routes() { - open_regions_ignore_region_not_found( - &self.context, - self.data.table_id(), - self.data.table_name(), - self.data.table_info(), - region_routes, - &self.data.region_wal_options, - ) - .await?; - } self.data.state = PurgeDroppedTableState::DropRegions; Ok(Status::executing(true)) } @@ -107,13 +90,12 @@ impl PurgeDroppedTableProcedure { async fn on_drop_regions(&mut self) -> Result { if let Some(region_routes) = self.data.physical_region_routes() { self.executor() - .on_drop_regions( + .on_cleanup_regions_offline( &self.context.node_manager, &self.context.leader_region_registry, + self.data.table_info(), region_routes, - false, - false, - false, + &self.data.region_wal_options, ) .await?; self.context @@ -157,7 +139,6 @@ impl Procedure for PurgeDroppedTableProcedure { async fn execute(&mut self, _: &ProcedureContext) -> ProcedureResult { match self.data.state { PurgeDroppedTableState::Prepare => self.on_prepare().await, - PurgeDroppedTableState::OpenRegions => self.on_open_regions().await, PurgeDroppedTableState::DropRegions => self.on_drop_regions().await, PurgeDroppedTableState::DeleteTombstone => self.on_delete_tombstone().await, } @@ -206,14 +187,14 @@ impl PurgeDroppedTableData { self.table_name.as_ref().unwrap() } - fn table_info(&self) -> &TableInfo { - self.table_info.as_ref().unwrap() - } - fn table_route_value(&self) -> &TableRouteValue { self.table_route_value.as_ref().unwrap() } + fn table_info(&self) -> &TableInfo { + self.table_info.as_ref().unwrap() + } + fn physical_region_routes(&self) -> Option<&[RegionRoute]> { match self.table_route_value() { TableRouteValue::Physical(route) => Some(&route.region_routes), @@ -225,7 +206,6 @@ impl PurgeDroppedTableData { #[derive(Debug, Serialize, Deserialize, AsRefStr, PartialEq)] enum PurgeDroppedTableState { Prepare, - OpenRegions, DropRegions, DeleteTombstone, } diff --git a/src/common/meta/src/ddl/tests/drop_table.rs b/src/common/meta/src/ddl/tests/drop_table.rs index 5c90dcbd79b..b14fe77cc5b 100644 --- a/src/common/meta/src/ddl/tests/drop_table.rs +++ b/src/common/meta/src/ddl/tests/drop_table.rs @@ -14,19 +14,17 @@ use std::assert_matches; use std::collections::{HashMap, HashSet}; -use std::sync::{Arc, Mutex as StdMutex}; +use std::sync::Arc; -use api::region::RegionResponse; use api::v1::region::{RegionRequest, region_request}; use async_trait::async_trait; -use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; -use common_error::ext::{BoxedError, ErrorExt, StackError}; +use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, FILE_ENGINE}; +use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; -use common_procedure::Procedure; +use common_procedure::{Procedure, StringKey}; use common_procedure_test::{ execute_procedure_until, execute_procedure_until_done, new_test_procedure_context, }; -use snafu::ResultExt; use store_api::region_engine::RegionRole; use store_api::storage::RegionId; use table::metadata::TableId; @@ -44,7 +42,7 @@ use crate::ddl::test_util::{ }; use crate::ddl::undrop_table::UndropTableProcedure; use crate::ddl::{DetectingRegion, RegionFailureDetectorController, TableMetadata}; -use crate::error::{self, Error}; +use crate::error::Error; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; use crate::kv_backend::memory::MemoryKvBackend; @@ -741,6 +739,34 @@ async fn test_soft_drop_metric_logical_table_fails() { assert_eq!(err.status_code(), StatusCode::Unsupported); } +#[tokio::test] +async fn test_soft_drop_file_engine_table_fails() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let table_name = "foo"; + let mut task = test_create_table_task(table_name, table_id); + task.table_info.meta.engine = FILE_ENGINE.to_string(); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + + let mut procedure = DropTableProcedure::new( + new_drop_table_task(table_name, table_id, false), + ddl_context, + ); + let err = procedure.on_prepare().await.unwrap_err(); + + assert_eq!(err.status_code(), StatusCode::Unsupported); +} + #[tokio::test] async fn test_undrop_metric_logical_table_fails() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); @@ -872,6 +898,174 @@ async fn test_undrop_table_fails_when_live_name_is_created_after_prepare() { assert_eq!(live_table.table_id(), live_table_id); } +#[tokio::test] +async fn test_undrop_table_closes_opened_regions_when_restore_metadata_races_with_create() { + let (tx, mut rx) = mpsc::channel(8); + let datanode_handler = DatanodeWatcher::new(tx); + let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); + let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + ddl_context.region_failure_detector_controller = detector_controller.clone(); + let dropped_table_id = 1024; + let live_table_id = 1025; + let table_name = "foo"; + let task = test_create_table_task(table_name, dropped_table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(RegionId::new(dropped_table_id, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![Peer::empty(2)], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + HashMap::new(), + ) + .await + .unwrap(); + let mut drop_procedure = DropTableProcedure::new( + new_drop_table_task(table_name, dropped_table_id, false), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut drop_procedure).await; + while rx.try_recv().is_ok() {} + detector_controller.clear().await; + + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(dropped_table_id), ddl_context.clone()); + let ctx = new_test_procedure_context(); + procedure.execute(&ctx).await.unwrap(); + procedure.execute(&ctx).await.unwrap(); + + let mut opened_regions = HashSet::new(); + for _ in 0..2 { + let (peer, request) = rx.try_recv().unwrap(); + let Some(region_request::Body::Open(req)) = request.body else { + unreachable!(); + }; + opened_regions.insert((peer.id, req.region_id)); + } + assert_eq!( + opened_regions, + HashSet::from([ + (1, RegionId::new(dropped_table_id, 1).as_u64()), + (2, RegionId::new(dropped_table_id, 1).as_u64()), + ]) + ); + assert_eq!( + detector_controller.registered().await, + vec![(1, RegionId::new(dropped_table_id, 1))] + ); + + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task(table_name, live_table_id).table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + + let err = procedure.execute(&ctx).await.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); + + let mut closed_regions = HashSet::new(); + for _ in 0..2 { + let (peer, request) = rx.try_recv().unwrap(); + let Some(region_request::Body::Close(req)) = request.body else { + unreachable!(); + }; + closed_regions.insert((peer.id, req.region_id)); + } + assert_eq!( + closed_regions, + HashSet::from([ + (1, RegionId::new(dropped_table_id, 1).as_u64()), + (2, RegionId::new(dropped_table_id, 1).as_u64()), + ]) + ); + assert!(rx.try_recv().is_err()); + assert_eq!( + detector_controller.deregistered().await, + vec![(1, RegionId::new(dropped_table_id, 1))] + ); + + let live_table = ddl_context + .table_metadata_manager + .table_name_manager() + .get(TableNameKey::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + table_name, + )) + .await + .unwrap() + .unwrap(); + assert_eq!(live_table.table_id(), live_table_id); +} + +#[tokio::test] +async fn test_undrop_table_lock_key_includes_original_table_name_before_prepare() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let table_name = "foo"; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + let mut drop_procedure = DropTableProcedure::new( + new_drop_table_task(table_name, table_id, false), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut drop_procedure).await; + + let original_table_name = ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .unwrap() + .table_name; + let procedure = UndropTableProcedure::new_with_original_table_name( + new_undrop_table_task(table_id), + ddl_context, + Some(original_table_name), + ); + + let keys = procedure + .lock_key() + .keys_to_lock() + .cloned() + .collect::>(); + assert!( + keys.iter().any(|key| matches!( + key, + StringKey::Exclusive(key) if key == "__table_name_lock/greptime.public.foo" + )), + "undrop lock keys should include the original table name: {keys:?}" + ); + assert!( + keys.iter().any(|key| matches!( + key, + StringKey::Exclusive(key) if key == "__table_lock/1024" + )), + "undrop lock keys should include the dropped table id: {keys:?}" + ); +} + #[tokio::test] async fn test_undrop_table_replayed_restore_metadata_is_idempotent() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); @@ -921,7 +1115,7 @@ async fn test_undrop_table_replayed_restore_metadata_is_idempotent() { } #[tokio::test] -async fn test_purge_dropped_table_drops_regions_and_deletes_tombstone() { +async fn test_purge_dropped_table_cleans_regions_offline_and_deletes_tombstone() { let (tx, mut rx) = mpsc::channel(8); let datanode_handler = DatanodeWatcher::new(tx); let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); @@ -963,15 +1157,16 @@ async fn test_purge_dropped_table_drops_regions_and_deletes_tombstone() { execute_procedure_until_done(&mut procedure).await; let mut requests = Vec::new(); - for _ in 0..4 { + for _ in 0..1 { let (peer, request) = rx.try_recv().unwrap(); requests.push((peer.id, request.body.unwrap())); } requests.sort_unstable_by_key(|(peer_id, _)| *peer_id); - assert_matches!(requests[0].1, region_request::Body::Open(_)); - assert_matches!(requests[1].1, region_request::Body::Drop(_)); - assert_matches!(requests[2].1, region_request::Body::Open(_)); - assert_matches!(requests[3].1, region_request::Body::Close(_)); + let region_request::Body::CleanUp(req) = &requests[0].1 else { + unreachable!(); + }; + assert_eq!(requests[0].0, 1); + assert_eq!(req.region_id, RegionId::new(table_id, 1).as_u64()); assert!(rx.try_recv().is_err()); assert!( ddl_context @@ -1050,119 +1245,13 @@ async fn test_purge_dropped_table_by_id_selects_tombstone_when_live_table_exists assert_eq!(live_table.table_id(), live_table_id); let (_, request) = rx.try_recv().unwrap(); - let Some(region_request::Body::Open(req)) = request.body else { - unreachable!(); - }; - assert_eq!(req.region_id, RegionId::new(dropped_table_id, 1).as_u64()); - - let (_, request) = rx.try_recv().unwrap(); - let Some(region_request::Body::Drop(req)) = request.body else { + let Some(region_request::Body::CleanUp(req)) = request.body else { unreachable!(); }; assert_eq!(req.region_id, RegionId::new(dropped_table_id, 1).as_u64()); assert!(rx.try_recv().is_err()); } -#[tokio::test] -async fn test_purge_dropped_table_replayed_open_regions_ignores_dropped_regions() { - let (tx, mut rx) = mpsc::channel(8); - let dropped_regions = Arc::new(StdMutex::new(HashSet::new())); - let datanode_handler = DatanodeWatcher::new(tx).with_handler({ - let dropped_regions = dropped_regions.clone(); - move |_peer, request| { - let Some(body) = request.body.as_ref() else { - return Ok(RegionResponse::new(0)); - }; - match body { - region_request::Body::Open(req) - if dropped_regions.lock().unwrap().contains(&req.region_id) => - { - Err::(BoxedError::new(MockRegionNotFoundError)) - .context(error::ExternalSnafu) - } - region_request::Body::Drop(req) => { - dropped_regions.lock().unwrap().insert(req.region_id); - Ok(RegionResponse::new(0)) - } - _ => Ok(RegionResponse::new(0)), - } - } - }); - let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); - let mut ddl_context = new_ddl_context(node_manager); - ddl_context.soft_drop_enabled = true; - let table_id = 1024; - let table_name = "foo"; - let task = test_create_table_task(table_name, table_id); - ddl_context - .table_metadata_manager - .create_table_metadata( - task.table_info.clone(), - TableRouteValue::physical(vec![RegionRoute { - region: Region::new_test(RegionId::new(table_id, 1)), - leader_peer: Some(Peer::empty(1)), - follower_peers: vec![], - leader_state: None, - leader_down_since: None, - write_route_policy: None, - }]), - HashMap::new(), - ) - .await - .unwrap(); - let mut drop_procedure = DropTableProcedure::new( - new_drop_table_task(table_name, table_id, false), - ddl_context.clone(), - ); - execute_procedure_until_done(&mut drop_procedure).await; - while rx.try_recv().is_ok() {} - - let mut procedure = PurgeDroppedTableProcedure::new( - new_purge_dropped_table_task(table_id), - ddl_context.clone(), - ); - let ctx = new_test_procedure_context(); - procedure.execute(&ctx).await.unwrap(); - let open_regions_data = procedure.dump().unwrap(); - procedure.execute(&ctx).await.unwrap(); - procedure.execute(&ctx).await.unwrap(); - - let mut replayed = - PurgeDroppedTableProcedure::from_json(&open_regions_data, ddl_context.clone()).unwrap(); - execute_procedure_until_done(&mut replayed).await; - - assert!( - ddl_context - .table_metadata_manager - .get_dropped_table(&drop_procedure.data.task.table_name()) - .await - .unwrap() - .is_none() - ); -} - -#[derive(Debug, snafu::Snafu)] -#[snafu(display("mock region not found"))] -struct MockRegionNotFoundError; - -impl StackError for MockRegionNotFoundError { - fn debug_fmt(&self, _: usize, _: &mut Vec) {} - - fn next(&self) -> Option<&dyn StackError> { - None - } -} - -impl ErrorExt for MockRegionNotFoundError { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn status_code(&self) -> StatusCode { - StatusCode::RegionNotFound - } -} - #[tokio::test] async fn test_on_rollback() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs index f08d606a827..9bd3b24822c 100644 --- a/src/common/meta/src/ddl/undrop_table.rs +++ b/src/common/meta/src/ddl/undrop_table.rs @@ -18,13 +18,12 @@ use api::v1::region::{ OpenRequest as PbOpenRegionRequest, RegionRequest, RegionRequestHeader, region_request, }; use async_trait::async_trait; -use common_error::ext::ErrorExt; -use common_error::status_code::StatusCode; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ Context as ProcedureContext, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::tracing_context::TracingContext; +use common_telemetry::warn; use common_wal::options::WalOptions; use futures::future::join_all; use serde::{Deserialize, Serialize}; @@ -34,6 +33,7 @@ use strum::AsRefStr; use table::metadata::TableId; use table::table_name::TableName; +use crate::ddl::drop_table::executor::DropTableExecutor; use crate::ddl::utils::{ add_peer_context_if_needed, convert_region_routes_to_detecting_regions, is_metric_engine_logical_table, map_to_procedure_error, region_storage_path, @@ -43,7 +43,7 @@ use crate::error::{self, Result}; use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; -use crate::lock_key::TableLock; +use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use crate::rpc::ddl::UndropTableTask; use crate::rpc::router::{ RegionRoute, find_follower_regions, find_followers, find_leader_regions, find_leaders, @@ -58,10 +58,17 @@ impl UndropTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::UndropTable"; pub fn new(task: UndropTableTask, context: DdlContext) -> Self { - Self { - context, - data: UndropTableData::new(task), - } + Self::new_with_original_table_name(task, context, None) + } + + pub(crate) fn new_with_original_table_name( + task: UndropTableTask, + context: DdlContext, + table_name: Option, + ) -> Self { + let mut data = UndropTableData::new(task); + data.table_name = table_name; + Self { context, data } } pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult { @@ -108,29 +115,101 @@ impl UndropTableProcedure { } async fn on_restore_metadata(&mut self) -> Result { + self.ensure_table_name_loaded().await?; + let table_name = self.data.table_name().clone(); let table_route_value = self.data.table_route_value(); - self.context + if let Err(err) = self + .context .table_metadata_manager .restore_table_metadata( self.data.task.table_id, - self.data.table_name(), + &table_name, table_route_value, &self.data.region_wal_options, ) .await - .map_err(|err| match err { - error::Error::TombstoneTargetAlreadyExists { .. } => { - error::TableAlreadyExistsSnafu { - table_name: self.data.table_name().to_string(), - } - .build() - } - err => err, - })?; + { + let should_cleanup_opened_regions = !err.is_retry_later(); + let err = self.map_restore_metadata_error(err); + if should_cleanup_opened_regions + && let Err(cleanup_err) = self.cleanup_opened_regions_after_restore_failure().await + { + warn!( + cleanup_err; + "Failed to close opened regions after undrop metadata restore failure, table_id: {}", + self.data.task.table_id + ); + } + return Err(err); + } self.data.state = UndropTableState::InvalidateTableCache; Ok(Status::executing(true)) } + async fn ensure_table_name_loaded(&mut self) -> Result<()> { + if self.data.table_name.is_some() { + return Ok(()); + } + + let dropped_table = self + .context + .table_metadata_manager + .get_dropped_table_by_id(self.data.task.table_id) + .await? + .with_context(|| error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + })?; + self.data.table_name = Some(dropped_table.table_name); + Ok(()) + } + + fn map_restore_metadata_error(&self, err: error::Error) -> error::Error { + match err { + error::Error::TombstoneTargetAlreadyExists { .. } => { + let table_name = self + .data + .table_name + .as_ref() + .map(ToString::to_string) + .unwrap_or_else(|| self.data.task.table_id.to_string()); + error::TableAlreadyExistsSnafu { table_name }.build() + } + err => err, + } + } + + async fn cleanup_opened_regions_after_restore_failure(&self) -> Result<()> { + let Some(table_route_value) = self.data.table_route_value.as_ref() else { + return Ok(()); + }; + let TableRouteValue::Physical(route) = table_route_value else { + return Ok(()); + }; + let region_routes = route.region_routes.clone(); + let close_result = if let Some(table_name) = self.data.table_name.as_ref() { + let executor = + DropTableExecutor::new(table_name.clone(), self.data.task.table_id, false); + + executor + .on_close_regions( + &self.context.node_manager, + &self.context.leader_region_registry, + ®ion_routes, + false, + ) + .await + } else { + Ok(()) + }; + self.context + .deregister_failure_detectors(convert_region_routes_to_detecting_regions( + ®ion_routes, + )) + .await; + close_result?; + Ok(()) + } + async fn on_open_regions(&mut self) -> Result { self.ensure_live_table_not_exists().await?; let TableRouteValue::Physical(route) = self.data.table_route_value() else { @@ -213,7 +292,22 @@ impl Procedure for UndropTableProcedure { } fn lock_key(&self) -> LockKey { - LockKey::new(vec![TableLock::Write(self.data.task.table_id).into()]) + let mut lock_key = Vec::new(); + if let Some(table_name) = &self.data.table_name { + lock_key.push(CatalogLock::Read(&table_name.catalog_name).into()); + lock_key + .push(SchemaLock::read(&table_name.catalog_name, &table_name.schema_name).into()); + lock_key.push( + TableNameLock::new( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + ) + .into(), + ); + } + lock_key.push(TableLock::Write(self.data.task.table_id).into()); + LockKey::new(lock_key) } } @@ -232,27 +326,6 @@ pub(crate) async fn open_regions( table_info, region_routes, region_wal_options, - false, - ) - .await -} - -pub(crate) async fn open_regions_ignore_region_not_found( - context: &DdlContext, - table_id: TableId, - table_name: &TableName, - table_info: &table::metadata::TableInfo, - region_routes: &[RegionRoute], - region_wal_options: &HashMap, -) -> Result<()> { - open_regions_inner( - context, - table_id, - table_name, - table_info, - region_routes, - region_wal_options, - true, ) .await } @@ -264,7 +337,6 @@ async fn open_regions_inner( table_info: &table::metadata::TableInfo, region_routes: &[RegionRoute], region_wal_options: &HashMap, - ignore_region_not_found: bool, ) -> Result<()> { let template = build_template_from_raw_table_info(table_info)?; let builder = CreateRequestBuilder::new(template, None); @@ -303,9 +375,7 @@ async fn open_regions_inner( let datanode = datanode.clone(); let requester = requester.clone(); tasks.push(async move { - if let Err(err) = requester.handle(request).await - && !(ignore_region_not_found && err.status_code() == StatusCode::RegionNotFound) - { + if let Err(err) = requester.handle(request).await { return Err(add_peer_context_if_needed(datanode)(err)); } Ok(()) @@ -366,3 +436,130 @@ enum UndropTableState { OpenRegions, InvalidateTableCache, } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use api::region::RegionResponse; + use api::v1::region::region_request; + use async_trait::async_trait; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; + use store_api::storage::RegionId; + use table::table_name::TableName; + use tokio::sync::{Mutex, mpsc}; + + use super::*; + use crate::ddl::test_util::datanode_handler::DatanodeWatcher; + use crate::ddl::{DetectingRegion, RegionFailureDetectorController}; + use crate::peer::Peer; + use crate::rpc::router::{Region, RegionRoute}; + use crate::test_util::{MockDatanodeManager, new_ddl_context}; + + #[derive(Default)] + struct RecordingRegionFailureDetectorController { + deregistered: Mutex>, + } + + #[async_trait] + impl RegionFailureDetectorController for RecordingRegionFailureDetectorController { + async fn register_failure_detectors(&self, _detecting_regions: Vec) {} + + async fn reset_failure_detectors(&self, _detecting_regions: Vec) {} + + async fn deregister_failure_detectors(&self, detecting_regions: Vec) { + self.deregistered.lock().await.extend(detecting_regions); + } + } + + #[test] + fn test_map_restore_metadata_error_without_table_name() { + let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let procedure = UndropTableProcedure::new(UndropTableTask { table_id: 42 }, context); + + let err = procedure.map_restore_metadata_error( + error::TombstoneTargetAlreadyExistsSnafu { + key: "table-name-key".to_string(), + } + .build(), + ); + + assert_eq!(StatusCode::TableAlreadyExists, err.status_code()); + } + + #[tokio::test] + async fn test_cleanup_opened_regions_without_table_name_deregisters_detectors() { + let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); + let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + context.region_failure_detector_controller = detector_controller.clone(); + + let table_id = 1024; + let region_id = RegionId::new(table_id, 1); + let mut procedure = UndropTableProcedure::new(UndropTableTask { table_id }, context); + procedure.data.table_route_value = Some(TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }])); + + procedure + .cleanup_opened_regions_after_restore_failure() + .await + .unwrap(); + + assert_eq!( + detector_controller.deregistered.lock().await.as_slice(), + &[(1, region_id)] + ); + } + + #[tokio::test] + async fn test_cleanup_opened_regions_deregisters_detectors_when_close_fails() { + let (tx, _rx) = mpsc::channel(8); + let datanode_handler = DatanodeWatcher::new(tx).with_handler(|_, request| { + if matches!(request.body, Some(region_request::Body::Close(_))) { + return error::UnexpectedSnafu { + err_msg: "mock close error".to_string(), + } + .fail(); + } + Ok(RegionResponse::new(0)) + }); + let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); + let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(datanode_handler))); + context.region_failure_detector_controller = detector_controller.clone(); + + let table_id = 1024; + let region_id = RegionId::new(table_id, 1); + let mut procedure = UndropTableProcedure::new(UndropTableTask { table_id }, context); + procedure.data.table_name = Some(TableName::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + "foo", + )); + procedure.data.table_route_value = Some(TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }])); + + let err = procedure + .cleanup_opened_regions_after_restore_failure() + .await + .unwrap_err(); + + assert_eq!(StatusCode::Unexpected, err.status_code()); + assert_eq!( + detector_controller.deregistered.lock().await.as_slice(), + &[(1, region_id)] + ); + } +} diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index bfd858a1963..3f421d3b028 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -480,7 +480,19 @@ impl DdlManager { undrop_table_task: UndropTableTask, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = UndropTableProcedure::new(undrop_table_task, context); + let original_table_name = context + .table_metadata_manager + .get_dropped_table_by_id(undrop_table_task.table_id) + .await? + .with_context(|| TableNotFoundSnafu { + table_name: undrop_table_task.table_id.to_string(), + })? + .table_name; + let procedure = UndropTableProcedure::new_with_original_table_name( + undrop_table_task, + context, + Some(original_table_name), + ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); self.execute_procedure_and_wait(procedure_with_id).await @@ -1204,7 +1216,8 @@ mod tests { use std::sync::Arc; use std::time::Duration; - use common_error::ext::BoxedError; + use common_error::ext::{BoxedError, ErrorExt}; + use common_error::status_code::StatusCode; use common_procedure::local::LocalManager; use common_procedure::test_util::InMemoryPoisonStore; use common_procedure::{BoxedProcedure, ProcedureManagerRef}; @@ -1228,6 +1241,7 @@ mod tests { use crate::peer::Peer; use crate::region_keeper::MemoryRegionKeeper; use crate::region_registry::LeaderRegionRegistry; + use crate::rpc::ddl::UndropTableTask; use crate::sequence::SequenceBuilder; use crate::state_store::KvStateStore; use crate::wal_provider::WalProvider; @@ -1326,4 +1340,55 @@ mod tests { assert!(procedure_manager.contains_loader(loader)); } } + + #[tokio::test] + async fn test_submit_undrop_missing_tombstone_returns_table_not_found_directly() { + let kv_backend = Arc::new(MemoryKvBackend::new()); + let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); + let table_metadata_allocator = Arc::new(TableMetadataAllocator::new( + Arc::new(SequenceBuilder::new("test", kv_backend.clone()).build()), + Arc::new(WalProvider::default()), + )); + let flow_metadata_manager = Arc::new(FlowMetadataManager::new(kv_backend.clone())); + let flow_metadata_allocator = Arc::new(FlowMetadataAllocator::with_noop_peer_allocator( + Arc::new(SequenceBuilder::new("flow-test", kv_backend.clone()).build()), + )); + + let state_store = Arc::new(KvStateStore::new(kv_backend.clone())); + let poison_manager = Arc::new(InMemoryPoisonStore::default()); + let procedure_manager = Arc::new(LocalManager::new( + Default::default(), + state_store, + poison_manager, + None, + None, + )); + + let ddl_manager = DdlManager::try_new( + DdlContext { + node_manager: Arc::new(DummyDatanodeManager), + cache_invalidator: Arc::new(DummyCacheInvalidator), + table_metadata_manager, + table_metadata_allocator, + flow_metadata_manager, + flow_metadata_allocator, + memory_region_keeper: Arc::new(MemoryRegionKeeper::default()), + leader_region_registry: Arc::new(LeaderRegionRegistry::default()), + region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), + soft_drop_enabled: true, + }, + procedure_manager, + Arc::new(DummyRepartitionProcedureFactory), + true, + ) + .unwrap(); + + let err = ddl_manager + .submit_undrop_table_task(UndropTableTask { table_id: 1024 }) + .await + .unwrap_err(); + + assert_eq!(err.status_code(), StatusCode::TableNotFound); + assert!(matches!(err, crate::error::Error::TableNotFound { .. })); + } } diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 3932e307489..5a6ab700854 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -1214,6 +1214,24 @@ impl RegionServerInner { })? .clone(), }, + RegionChange::OfflineCleanup(attribute) => match current_region_status { + Some(status) => match status.clone() { + RegionEngineWithStatus::Registering(_) + | RegionEngineWithStatus::Deregistering(_) + | RegionEngineWithStatus::Ready(_) => { + return error::RegionBusySnafu { region_id }.fail(); + } + }, + None => self + .engines + .read() + .unwrap() + .get(attribute.engine()) + .with_context(|| RegionEngineNotFoundSnafu { + name: attribute.engine(), + })? + .clone(), + }, RegionChange::Deregisters => match current_region_status { Some(status) => match status.clone() { RegionEngineWithStatus::Registering(_) => { @@ -1559,6 +1577,10 @@ impl RegionServerInner { let attribute = parse_region_attribute(&open.engine, &open.options)?; RegionChange::Register(attribute) } + RegionRequest::CleanUp(clean_up) => { + let attribute = parse_region_attribute(&clean_up.engine, &clean_up.options)?; + RegionChange::OfflineCleanup(attribute) + } RegionRequest::Close(_) | RegionRequest::Drop(_) => RegionChange::Deregisters, RegionRequest::Put(_) | RegionRequest::Delete(_) | RegionRequest::BulkInserts(_) => { RegionChange::Ingest @@ -1678,7 +1700,7 @@ impl RegionServerInner { region_change: RegionChange, ) { match region_change { - RegionChange::None | RegionChange::Ingest => {} + RegionChange::None | RegionChange::Ingest | RegionChange::OfflineCleanup(_) => {} RegionChange::Register(_) => { self.region_map.remove(®ion_id); } @@ -1698,7 +1720,7 @@ impl RegionServerInner { ) -> Result<()> { let engine_type = engine.name(); match region_change { - RegionChange::None | RegionChange::Ingest => {} + RegionChange::None | RegionChange::Ingest | RegionChange::OfflineCleanup(_) => {} RegionChange::Register(attribute) => { info!( "Region {region_id} is registered to engine {}", @@ -1883,6 +1905,7 @@ impl RegionServerInner { enum RegionChange { None, Register(RegionAttribute), + OfflineCleanup(RegionAttribute), Deregisters, Catchup, Ingest, @@ -1947,8 +1970,8 @@ mod tests { use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::region_engine::RegionEngine; use store_api::region_request::{ - PathType, RegionCompactRequest, RegionDeleteRequest, RegionDropRequest, RegionOpenRequest, - RegionPutRequest, RegionTruncateRequest, + PathType, RegionCleanUpRequest, RegionCompactRequest, RegionDeleteRequest, + RegionDropRequest, RegionOpenRequest, RegionPutRequest, RegionTruncateRequest, }; use store_api::storage::RegionId; @@ -2119,6 +2142,77 @@ mod tests { assert!(pinned.as_ref().get_ref().metrics().is_none()); } + #[tokio::test] + async fn test_offline_cleanup_does_not_register_region() { + let mut mock_region_server = mock_region_server(); + let (engine, mut receiver) = MockRegionEngine::new(MITO_ENGINE_NAME); + mock_region_server.register_engine(engine); + + let region_id = RegionId::new(1, 1); + let response = mock_region_server + .handle_request( + region_id, + RegionRequest::CleanUp(RegionCleanUpRequest { + engine: MITO_ENGINE_NAME.to_string(), + table_dir: String::new(), + path_type: PathType::Bare, + options: HashMap::new(), + }), + ) + .await + .unwrap(); + + assert_eq!(response.affected_rows, 0); + let (handled_region_id, handled_request) = receiver.try_recv().unwrap(); + assert_eq!(handled_region_id, region_id); + assert_matches!(handled_request, RegionRequest::CleanUp(_)); + assert!( + mock_region_server + .inner + .region_map + .get(®ion_id) + .is_none() + ); + } + + #[tokio::test] + async fn test_offline_cleanup_rejects_registered_region() { + let mut mock_region_server = mock_region_server(); + let (engine, mut receiver) = MockRegionEngine::new(MITO_ENGINE_NAME); + mock_region_server.register_engine(engine.clone()); + + let region_id = RegionId::new(1, 1); + mock_region_server + .inner + .region_map + .insert(region_id, RegionEngineWithStatus::Ready(engine)); + + let err = mock_region_server + .handle_request( + region_id, + RegionRequest::CleanUp(RegionCleanUpRequest { + engine: MITO_ENGINE_NAME.to_string(), + table_dir: String::new(), + path_type: PathType::Bare, + options: HashMap::new(), + }), + ) + .await + .unwrap_err(); + + assert_eq!(err.status_code(), StatusCode::RegionBusy); + assert!(receiver.try_recv().is_err()); + assert!(matches!( + mock_region_server + .inner + .region_map + .get(®ion_id) + .unwrap() + .clone(), + RegionEngineWithStatus::Ready(_) + )); + } + #[tokio::test] async fn test_region_registering() { common_telemetry::init_default_ut_logging(); @@ -2530,6 +2624,43 @@ mod tests { assert_matches!(current_engine, CurrentEngine::Engine(_)); }), }, + // RegionChange::OfflineCleanup + CurrentEngineTest { + region_id, + current_region_status: None, + region_change: RegionChange::OfflineCleanup(RegionAttribute::Mito), + assert: Box::new(|result| { + let current_engine = result.unwrap(); + assert_matches!(current_engine, CurrentEngine::Engine(_)); + }), + }, + CurrentEngineTest { + region_id, + current_region_status: Some(RegionEngineWithStatus::Registering(engine.clone())), + region_change: RegionChange::OfflineCleanup(RegionAttribute::Mito), + assert: Box::new(|result| { + let err = result.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::RegionBusy); + }), + }, + CurrentEngineTest { + region_id, + current_region_status: Some(RegionEngineWithStatus::Deregistering(engine.clone())), + region_change: RegionChange::OfflineCleanup(RegionAttribute::Mito), + assert: Box::new(|result| { + let err = result.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::RegionBusy); + }), + }, + CurrentEngineTest { + region_id, + current_region_status: Some(RegionEngineWithStatus::Ready(engine.clone())), + region_change: RegionChange::OfflineCleanup(RegionAttribute::Mito), + assert: Box::new(|result| { + let err = result.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::RegionBusy); + }), + }, ]; for test in tests { diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index f8eceda9e8a..f5c7863e825 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -237,6 +237,9 @@ impl RegionEngine for MetricEngine { } RegionRequest::Drop(drop) => self.inner.drop_region(region_id, drop).await, RegionRequest::Open(open) => self.inner.open_region(region_id, open).await, + RegionRequest::CleanUp(clean_up) => { + self.inner.clean_up_region(region_id, clean_up).await + } RegionRequest::Close(close) => self.inner.close_region(region_id, close).await, RegionRequest::Alter(alter) => { self.inner @@ -585,8 +588,8 @@ mod test { use store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY; use store_api::mito_engine_options::WAL_OPTIONS_KEY; use store_api::region_request::{ - PathType, RegionCloseRequest, RegionDropRequest, RegionFlushRequest, RegionOpenRequest, - RegionRequest, + PathType, RegionCleanUpRequest, RegionCloseRequest, RegionDropRequest, RegionFlushRequest, + RegionOpenRequest, RegionRequest, }; use super::*; @@ -656,6 +659,58 @@ mod test { .unwrap(); } + #[tokio::test] + async fn test_offline_cleanup_physical_region() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let engine = env.metric(); + let mito = env.mito(); + let physical_region_id = env.default_physical_region_id(); + let metadata_region_id = crate::utils::to_metadata_region_id(physical_region_id); + let data_region_id = crate::utils::to_data_region_id(physical_region_id); + + engine + .handle_request( + physical_region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + + let object_store = env.get_object_store().unwrap(); + let metadata_region_dir = region_dir_from_table_dir( + &TestEnv::default_table_dir(), + metadata_region_id, + PathType::Metadata, + ); + let data_region_dir = region_dir_from_table_dir( + &TestEnv::default_table_dir(), + data_region_id, + PathType::Data, + ); + assert!(object_store.exists(&metadata_region_dir).await.unwrap()); + assert!(object_store.exists(&data_region_dir).await.unwrap()); + + let physical_region_option = [(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new())] + .into_iter() + .collect(); + let clean_up_request = RegionCleanUpRequest { + engine: METRIC_ENGINE_NAME.to_string(), + table_dir: TestEnv::default_table_dir(), + path_type: PathType::Bare, + options: physical_region_option, + }; + engine + .handle_request(physical_region_id, RegionRequest::CleanUp(clean_up_request)) + .await + .unwrap(); + + assert!(!mito.is_region_exists(metadata_region_id)); + assert!(!mito.is_region_exists(data_region_id)); + assert!(!object_store.exists(&metadata_region_dir).await.unwrap()); + assert!(!object_store.exists(&data_region_dir).await.unwrap()); + } + #[tokio::test] async fn test_role() { let env = TestEnv::new().await; diff --git a/src/metric-engine/src/engine/open.rs b/src/metric-engine/src/engine/open.rs index 65ffcd77eb4..54bbdfd7038 100644 --- a/src/metric-engine/src/engine/open.rs +++ b/src/metric-engine/src/engine/open.rs @@ -22,14 +22,17 @@ use datafusion::common::HashMap; use mito2::engine::MITO_ENGINE_NAME; use snafu::{OptionExt, ResultExt}; use store_api::region_engine::{BatchResponses, RegionEngine}; -use store_api::region_request::{AffectedRows, PathType, RegionOpenRequest, ReplayCheckpoint}; +use store_api::region_request::{ + AffectedRows, PathType, RegionCleanUpRequest, RegionOpenRequest, RegionRequest, + ReplayCheckpoint, +}; use store_api::storage::RegionId; use crate::engine::MetricEngineInner; use crate::engine::create::region_options_for_metadata_region; use crate::engine::options::{PhysicalRegionOptions, set_data_region_options}; use crate::error::{ - BatchOpenMitoRegionSnafu, NoOpenRegionResultSnafu, OpenMitoRegionSnafu, + BatchOpenMitoRegionSnafu, CleanUpMitoRegionSnafu, NoOpenRegionResultSnafu, OpenMitoRegionSnafu, PhysicalRegionNotFoundSnafu, Result, }; use crate::metrics::{LOGICAL_REGION_COUNT, PHYSICAL_REGION_COUNT}; @@ -200,6 +203,94 @@ impl MetricEngineInner { } } + pub async fn clean_up_region( + &self, + region_id: RegionId, + request: RegionCleanUpRequest, + ) -> Result { + if request.is_physical_table() { + self.cleanup_physical_region_offline(region_id, request) + .await + } else { + Ok(0) + } + } + + async fn cleanup_physical_region_offline( + &self, + region_id: RegionId, + request: RegionCleanUpRequest, + ) -> Result { + let metadata_region_id = utils::to_metadata_region_id(region_id); + let data_region_id = utils::to_data_region_id(region_id); + let (clean_up_metadata_region_request, clean_up_data_region_request) = + self.transform_clean_up_physical_region_request(request); + let _ = self + .mito + .handle_request( + metadata_region_id, + RegionRequest::CleanUp(clean_up_metadata_region_request), + ) + .await + .context(CleanUpMitoRegionSnafu { + region_type: "metadata", + })?; + let data_region_response = self + .mito + .handle_request( + data_region_id, + RegionRequest::CleanUp(clean_up_data_region_request), + ) + .await + .context(CleanUpMitoRegionSnafu { + region_type: "data", + })?; + + if self.state.read().unwrap().exist_physical_region(region_id) { + self.state + .write() + .unwrap() + .remove_physical_region(region_id)?; + PHYSICAL_REGION_COUNT.dec(); + } + + Ok(data_region_response.affected_rows) + } + + /// Transform the cleanup request to metadata region and data region cleanup requests. + /// + /// Returns: + /// - The cleanup request for metadata region. + /// - The cleanup request for data region. + fn transform_clean_up_physical_region_request( + &self, + request: RegionCleanUpRequest, + ) -> (RegionCleanUpRequest, RegionCleanUpRequest) { + let clean_up_metadata_region_request = RegionCleanUpRequest { + table_dir: request.table_dir.clone(), + path_type: PathType::Metadata, + options: region_options_for_metadata_region(&request.options), + engine: MITO_ENGINE_NAME.to_string(), + }; + + let mut data_region_options = request.options; + set_data_region_options( + &mut data_region_options, + self.config.sparse_primary_key_encoding, + ); + let clean_up_data_region_request = RegionCleanUpRequest { + table_dir: request.table_dir, + path_type: PathType::Data, + options: data_region_options, + engine: MITO_ENGINE_NAME.to_string(), + }; + + ( + clean_up_metadata_region_request, + clean_up_data_region_request, + ) + } + /// Transform the open request to open metadata region and data region. /// /// Returns: diff --git a/src/metric-engine/src/error.rs b/src/metric-engine/src/error.rs index 59aee77b67b..316588e6b31 100644 --- a/src/metric-engine/src/error.rs +++ b/src/metric-engine/src/error.rs @@ -43,6 +43,14 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to clean up mito region, region type: {}", region_type))] + CleanUpMitoRegion { + region_type: String, + source: BoxedError, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to batch open mito region"))] BatchOpenMitoRegion { source: BoxedError, @@ -438,6 +446,7 @@ impl ErrorExt for Error { CreateMitoRegion { source, .. } | OpenMitoRegion { source, .. } + | CleanUpMitoRegion { source, .. } | CloseMitoRegion { source, .. } | MitoReadOperation { source, .. } | MitoWriteOperation { source, .. } @@ -471,6 +480,7 @@ impl ErrorExt for Error { match self { CreateMitoRegion { source, .. } | OpenMitoRegion { source, .. } + | CleanUpMitoRegion { source, .. } | BatchOpenMitoRegion { source, .. } | BatchCatchupMitoRegion { source, .. } | CloseMitoRegion { source, .. } diff --git a/src/mito2/src/engine/open_test.rs b/src/mito2/src/engine/open_test.rs index 2b709a7f08b..f5ee4dd19f7 100644 --- a/src/mito2/src/engine/open_test.rs +++ b/src/mito2/src/engine/open_test.rs @@ -23,9 +23,12 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; use either::Either; +use store_api::logstore::LogStore; +use store_api::logstore::provider::Provider; use store_api::region_engine::{RegionEngine, RegionRole, SettableRegionRoleState}; use store_api::region_request::{ - PathType, RegionCloseRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, + PathType, RegionCleanUpRequest, RegionCloseRequest, RegionOpenRequest, RegionPutRequest, + RegionRequest, }; use store_api::storage::{RegionId, ScanRequest}; use tokio::sync::oneshot; @@ -37,8 +40,10 @@ use crate::engine::region_hook::RegionHookRef; use crate::error; use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef}; use crate::region::options::RegionOptions; +use crate::sst::location::region_dir_from_table_dir; use crate::test_util::{ - CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, reopen_region, rows_schema, + CreateRequestBuilder, LogStoreImpl, TestEnv, build_rows, flush_region, put_rows, reopen_region, + rows_schema, }; #[tokio::test] @@ -81,6 +86,103 @@ async fn test_engine_open_empty_with_format(flat_format: bool) { assert_eq!(role, None); } +#[tokio::test] +async fn test_engine_offline_cleanup_closed_region() { + let mut env = TestEnv::with_prefix("offline-cleanup").await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + let path_type = request.path_type; + let options = request.options.clone(); + let region_dir = region_dir_from_table_dir(&table_dir, region_id, path_type); + let object_store = env.get_object_store().unwrap(); + + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + assert!(object_store.exists(®ion_dir).await.unwrap()); + + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + ) + .await; + let Some(LogStoreImpl::RaftEngine(log_store)) = env.get_log_store() else { + unreachable!() + }; + let provider = Provider::raft_engine_provider(region_id.as_u64()); + assert!(log_store.latest_entry_id(&provider).unwrap() > 0); + + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + assert!(!engine.is_region_exists(region_id)); + + let cleanup_request = RegionCleanUpRequest { + engine: String::new(), + table_dir, + path_type, + options, + }; + for _ in 0..2 { + engine + .handle_request(region_id, RegionRequest::CleanUp(cleanup_request.clone())) + .await + .unwrap(); + } + + assert!(!engine.is_region_exists(region_id)); + assert!(!object_store.exists(®ion_dir).await.unwrap()); + assert_eq!(0, log_store.latest_entry_id(&provider).unwrap()); +} + +#[tokio::test] +async fn test_engine_offline_cleanup_rejects_opened_region() { + let mut env = TestEnv::with_prefix("offline-cleanup-opened").await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + let path_type = request.path_type; + let options = request.options.clone(); + let region_dir = region_dir_from_table_dir(&table_dir, region_id, path_type); + let object_store = env.get_object_store().unwrap(); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let err = engine + .handle_request( + region_id, + RegionRequest::CleanUp(RegionCleanUpRequest { + engine: String::new(), + table_dir, + path_type, + options, + }), + ) + .await + .unwrap_err(); + + assert_eq!(StatusCode::RegionBusy, err.status_code()); + assert!(engine.is_region_exists(region_id)); + assert!(object_store.exists(®ion_dir).await.unwrap()); +} + #[tokio::test] async fn test_engine_open_existing() { test_engine_open_existing_with_format(false).await; diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 2d5cdcb2847..ad59e4fb94f 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -436,31 +436,7 @@ impl RegionOpener { } fn provider(&self, wal_options: &WalOptions) -> Result { - match wal_options { - WalOptions::RaftEngine => { - ensure!( - TypeId::of::() == TypeId::of::() - || TypeId::of::() == TypeId::of::(), - error::IncompatibleWalProviderChangeSnafu { - global: "`kafka`", - region: "`raft_engine`", - } - ); - Ok(Provider::raft_engine_provider(self.region_id.as_u64())) - } - WalOptions::Kafka(options) => { - ensure!( - TypeId::of::() == TypeId::of::() - || TypeId::of::() == TypeId::of::(), - error::IncompatibleWalProviderChangeSnafu { - global: "`raft_engine`", - region: "`kafka`", - } - ); - Ok(Provider::kafka_provider(options.topic.clone())) - } - WalOptions::Noop => Ok(Provider::noop_provider()), - } + provider_from_wal_options::(self.region_id, wal_options) } /// Tries to open the region and returns `None` if the region directory is empty. @@ -656,6 +632,37 @@ impl RegionOpener { } } +pub(crate) fn provider_from_wal_options( + region_id: RegionId, + wal_options: &WalOptions, +) -> Result { + match wal_options { + WalOptions::RaftEngine => { + ensure!( + TypeId::of::() == TypeId::of::() + || TypeId::of::() == TypeId::of::(), + error::IncompatibleWalProviderChangeSnafu { + global: "`kafka`", + region: "`raft_engine`", + } + ); + Ok(Provider::raft_engine_provider(region_id.as_u64())) + } + WalOptions::Kafka(options) => { + ensure!( + TypeId::of::() == TypeId::of::() + || TypeId::of::() == TypeId::of::(), + error::IncompatibleWalProviderChangeSnafu { + global: "`raft_engine`", + region: "`kafka`", + } + ); + Ok(Provider::kafka_provider(options.topic.clone())) + } + WalOptions::Noop => Ok(Provider::noop_provider()), + } +} + #[cfg(not(feature = "test-shared-fs-region-migration"))] fn supports_open_region_object_storage_requirement(object_store: &ObjectStore) -> bool { is_object_storage(object_store) diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index e9df654f834..3fdd15db4f1 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -38,9 +38,10 @@ use store_api::region_engine::{ }; use store_api::region_request::{ AffectedRows, ApplyStagingManifestRequest, EnterStagingRequest, RegionAlterRequest, - RegionBuildIndexRequest, RegionBulkInsertsRequest, RegionCatchupRequest, RegionCloseRequest, - RegionCompactRequest, RegionCreateRequest, RegionDropRequest, RegionFlushRequest, - RegionOpenRequest, RegionRequest, RegionTruncateRequest, StagingPartitionDirective, + RegionBuildIndexRequest, RegionBulkInsertsRequest, RegionCatchupRequest, RegionCleanUpRequest, + RegionCloseRequest, RegionCompactRequest, RegionCreateRequest, RegionDropRequest, + RegionFlushRequest, RegionOpenRequest, RegionRequest, RegionTruncateRequest, + StagingPartitionDirective, }; use store_api::storage::{FileId, RegionId}; use tokio::sync::oneshot::{self, Receiver, Sender}; @@ -711,6 +712,11 @@ impl WorkerRequest { sender: sender.into(), request: DdlRequest::Open((v, None)), }), + RegionRequest::CleanUp(v) => WorkerRequest::Ddl(SenderDdlRequest { + region_id, + sender: sender.into(), + request: DdlRequest::OfflineCleanup(v), + }), RegionRequest::Close(v) => WorkerRequest::Ddl(SenderDdlRequest { region_id, sender: sender.into(), @@ -863,6 +869,7 @@ pub(crate) enum DdlRequest { Create(RegionCreateRequest), Drop(RegionDropRequest), Open((RegionOpenRequest, Option)), + OfflineCleanup(RegionCleanUpRequest), Close(RegionCloseRequest), Alter(RegionAlterRequest), Flush(RegionFlushRequest), diff --git a/src/mito2/src/wal.rs b/src/mito2/src/wal.rs index eb9d4a251bb..316cbbd6126 100644 --- a/src/mito2/src/wal.rs +++ b/src/mito2/src/wal.rs @@ -159,6 +159,19 @@ impl Wal { .map_err(BoxedError::new) .context(DeleteWalSnafu { region_id }) } + + /// Deletes all WAL entries in the namespace represented by `provider`. + pub async fn delete_namespace(&self, region_id: RegionId, provider: &Provider) -> Result<()> { + if let Provider::Noop = provider { + return Ok(()); + } + self.store + .delete_namespace(provider) + .await + .map_err(BoxedError::new) + .context(DeleteWalSnafu { region_id }) + } + /// Marks all WAL entries of a region as obsolete and removes its dedicated namespace when /// supported by the backend. pub async fn obsolete_all(&self, region_id: RegionId, provider: &Provider) -> Result<()> { diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 3cd4a402c88..23b6dd4d7b8 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -1148,6 +1148,10 @@ impl RegionWorkerLoop { .await; continue; } + DdlRequest::OfflineCleanup(req) => { + self.handle_offline_cleanup_request(ddl.region_id, req) + .await + } DdlRequest::Close(req) => { self.handle_close_request(ddl.region_id, req, ddl.sender) .await; diff --git a/src/mito2/src/worker/handle_drop.rs b/src/mito2/src/worker/handle_drop.rs index 3d111776b82..1efa7cb99ed 100644 --- a/src/mito2/src/worker/handle_drop.rs +++ b/src/mito2/src/worker/handle_drop.rs @@ -28,9 +28,11 @@ use store_api::region_request::{AffectedRows, PathType}; use store_api::storage::RegionId; use tokio::time::sleep; +use crate::cache::CacheManagerRef; use crate::engine::region_hook::RegionHookRef; use crate::error::{OpenDalSnafu, Result}; use crate::region::{RegionLeaderState, RegionMapRef}; +use crate::sst::index::intermediate::IntermediateManager; use crate::worker::{DROPPING_MARKER_FILE, RegionWorkerLoop}; const GC_TASK_INTERVAL_SEC: u64 = 5 * 60; // 5 minutes @@ -90,14 +92,7 @@ where ®ion.provider, ) .await?; - // Notifies flush scheduler. - self.flush_scheduler.on_region_dropped(region_id); - // Notifies compaction scheduler. - self.compaction_scheduler.on_region_dropped(region_id); - // notifies index build scheduler. - self.index_build_scheduler - .on_region_dropped(region_id) - .await; + self.cleanup_dropped_region_runtime_state(region_id).await; // Marks region version as dropped region.version_control.mark_dropped(); @@ -159,21 +154,57 @@ where .await }; - if let Err(err) = intm_manager.prune_region_dir(®ion_id).await { - warn!(err; "Failed to prune intermediate region directory, region_id: {}", region_id); - } - - if let Some(write_cache) = cache_manager.write_cache() - && let Some(manifest_cache) = write_cache.manifest_cache() - { - manifest_cache.clean_manifests(&table_dir).await; - } + cleanup_region_file_artifacts(region_id, &table_dir, &intm_manager, &cache_manager) + .await; listener.on_later_drop_end(region_id, removed); }); Ok(0) } + + /// Cleans runtime state for a region that is no longer available to serve requests. + pub(crate) async fn cleanup_dropped_region_runtime_state(&mut self, region_id: RegionId) { + // Notifies flush scheduler. + self.flush_scheduler.on_region_dropped(region_id); + // Notifies compaction scheduler. + self.compaction_scheduler.on_region_dropped(region_id); + // Notifies index build scheduler. + self.index_build_scheduler + .on_region_dropped(region_id) + .await; + } +} + +/// Cleans files and caches that are produced at runtime but are not part of the +/// primary region directory deletion. +pub(crate) async fn cleanup_region_file_artifacts( + region_id: RegionId, + table_dir: &str, + intermediate_manager: &IntermediateManager, + cache_manager: &CacheManagerRef, +) { + if let Err(err) = intermediate_manager.prune_region_dir(®ion_id).await { + warn!(err; "Failed to prune intermediate region directory, region_id: {}", region_id); + } + + if let Some(write_cache) = cache_manager.write_cache() + && let Some(manifest_cache) = write_cache.manifest_cache() + { + manifest_cache.clean_manifests(table_dir).await; + } +} + +/// Removes a region directory for full-drop style cleanup. +/// +/// Full drop and purge/offline cleanup force physical deletion. Only partial +/// drop may leave data files for global GC. +pub(crate) async fn remove_region_dir_for_full_drop( + region_path: &str, + object_store: &ObjectStore, +) -> Result<()> { + remove_region_dir_once(region_path, object_store, true).await?; + Ok(()) } /// Carries the region hook and region metadata into the background GC task so @@ -281,7 +312,7 @@ async fn later_drop_task_with_global_gc( // the region directory is forcefully removed immediately. // // TODO(discord9): Evaluate removing files instantly rather than waiting for the GC period. - if path_type == PathType::Metadata || !partial_drop { + if should_force_remove_region_dir(path_type, partial_drop) { remove_region_with_retry( region_id, region_path, @@ -299,6 +330,10 @@ async fn later_drop_task_with_global_gc( } } +fn should_force_remove_region_dir(path_type: PathType, partial_drop: bool) -> bool { + path_type == PathType::Metadata || !partial_drop +} + // TODO(ruihang): place the marker in a separate dir /// Removes region dir if there is no parquet files, returns whether the directory is removed. /// If `force = true`, always removes the dir. diff --git a/src/mito2/src/worker/handle_open.rs b/src/mito2/src/worker/handle_open.rs index 48fef5ddd6f..802ab9c1a83 100644 --- a/src/mito2/src/worker/handle_open.rs +++ b/src/mito2/src/worker/handle_open.rs @@ -18,24 +18,68 @@ use std::sync::Arc; use std::time::Instant; use common_telemetry::info; -use object_store::util::join_path; +use object_store::util::{join_path, normalize_dir}; use snafu::{OptionExt, ResultExt}; use store_api::logstore::LogStore; -use store_api::region_request::RegionOpenRequest; +use store_api::region_request::{AffectedRows, RegionCleanUpRequest, RegionOpenRequest}; use store_api::storage::RegionId; use table::requests::STORAGE_KEY; use crate::error::{ - ObjectStoreNotFoundSnafu, OpenDalSnafu, OpenRegionSnafu, RegionNotFoundSnafu, Result, + ObjectStoreNotFoundSnafu, OpenDalSnafu, OpenRegionSnafu, RegionBusySnafu, RegionNotFoundSnafu, + Result, }; -use crate::region::opener::{RegionOpener, sanitize_open_request_options}; +use crate::region::opener::{ + RegionOpener, get_object_store, provider_from_wal_options, sanitize_open_request_options, +}; +use crate::region::options::RegionOptions; use crate::request::OptionOutputTx; use crate::sst::location::region_dir_from_table_dir; use crate::wal::entry_distributor::WalEntryReceiver; -use crate::worker::handle_drop::remove_region_dir_once; +use crate::worker::handle_drop::{ + cleanup_region_file_artifacts, remove_region_dir_for_full_drop, remove_region_dir_once, +}; use crate::worker::{DROPPING_MARKER_FILE, RegionWorkerLoop}; impl RegionWorkerLoop { + pub(crate) async fn handle_offline_cleanup_request( + &mut self, + region_id: RegionId, + mut request: RegionCleanUpRequest, + ) -> Result { + info!( + "Try to clean region {} offline, worker: {}", + region_id, self.id + ); + + if self.regions.is_region_exists(region_id) { + return RegionBusySnafu { region_id }.fail(); + } + + sanitize_open_request_options(&mut request.options); + + let options = RegionOptions::try_from_options(region_id, &request.options)?; + let object_store = get_object_store(&options.storage, &self.object_store_manager)?; + let provider = provider_from_wal_options::(region_id, &options.wal_options)?; + self.wal.obsolete_all(region_id, &provider).await?; + + let table_dir = normalize_dir(&request.table_dir); + let region_dir = region_dir_from_table_dir(&table_dir, region_id, request.path_type); + remove_region_dir_for_full_drop(®ion_dir, &object_store).await?; + + self.cleanup_dropped_region_runtime_state(region_id).await; + self.dropping_regions.remove_region(region_id); + cleanup_region_file_artifacts( + region_id, + &table_dir, + &self.intermediate_manager, + &self.cache_manager, + ) + .await; + + Ok(0) + } + async fn check_and_cleanup_region( &self, region_id: RegionId, diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 61b353956b1..6aa3e3f49d0 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -23,10 +23,10 @@ use api::v1::column_def::{ }; use api::v1::region::bulk_insert_request::Body; use api::v1::region::{ - AlterRequest, AlterRequests, BuildIndexRequest, BulkInsertRequest, CloseRequest, - CompactRequest, CreateRequest, CreateRequests, DeleteRequests, DropRequest, DropRequests, - FlushRequest, InsertRequests, OpenRequest, TruncateRequest, alter_request, compact_request, - region_request, truncate_request, + AlterRequest, AlterRequests, BuildIndexRequest, BulkInsertRequest, + CleanUpRequest as PbCleanUpRequest, CloseRequest, CompactRequest, CreateRequest, + CreateRequests, DeleteRequests, DropRequest, DropRequests, FlushRequest, InsertRequests, + OpenRequest, TruncateRequest, alter_request, compact_request, region_request, truncate_request, }; use api::v1::{ self, Analyzer, ArrowIpc, FulltextBackend as PbFulltextBackend, Option as PbOption, Rows, @@ -145,6 +145,7 @@ pub enum RegionRequest { Create(RegionCreateRequest), Drop(RegionDropRequest), Open(RegionOpenRequest), + CleanUp(RegionCleanUpRequest), Close(RegionCloseRequest), Alter(RegionAlterRequest), Flush(RegionFlushRequest), @@ -167,6 +168,7 @@ impl RegionRequest { region_request::Body::Create(create) => make_region_create(create), region_request::Body::Drop(drop) => make_region_drop(drop), region_request::Body::Open(open) => make_region_open(open), + region_request::Body::CleanUp(clean_up) => make_region_clean_up(clean_up), region_request::Body::Close(close) => make_region_close(close), region_request::Body::Alter(alter) => make_region_alter(alter), region_request::Body::Flush(flush) => make_region_flush(flush), @@ -189,10 +191,6 @@ impl RegionRequest { reason: "RemoteDynFilter request should be handled separately by RegionServer", } .fail(), - region_request::Body::CleanUp(_) => UnexpectedSnafu { - reason: "CleanUp request should be handled separately by RegionServer", - } - .fail(), region_request::Body::ApplyStagingManifest(apply) => { make_region_apply_staging_manifest(apply) } @@ -329,6 +327,20 @@ fn make_region_open(open: OpenRequest) -> Result> )]) } +fn make_region_clean_up(clean_up: PbCleanUpRequest) -> Result> { + let region_id = RegionId::from(clean_up.region_id); + let table_dir = table_dir(&clean_up.path, region_id.table_id()); + Ok(vec![( + region_id, + RegionRequest::CleanUp(RegionCleanUpRequest { + engine: clean_up.engine, + table_dir, + path_type: PathType::Bare, + options: clean_up.options, + }), + )]) +} + fn make_region_close(close: CloseRequest) -> Result> { let region_id = close.region_id.into(); Ok(vec![( @@ -652,6 +664,26 @@ impl RegionOpenRequest { } } +/// Offline region cleanup request. +#[derive(Debug, Clone)] +pub struct RegionCleanUpRequest { + /// Region engine name + pub engine: String, + /// Directory for table's data home. Usually is composed by catalog and table id + pub table_dir: String, + /// Path type for generating paths + pub path_type: PathType, + /// Options of the cleaned region. + pub options: HashMap, +} + +impl RegionCleanUpRequest { + /// Returns true when the region belongs to the metric engine's physical table. + pub fn is_physical_table(&self) -> bool { + self.options.contains_key(PHYSICAL_TABLE_METADATA_KEY) + } +} + /// Close region request. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct RegionCloseRequest { @@ -1642,6 +1674,7 @@ impl fmt::Display for RegionRequest { RegionRequest::Create(_) => write!(f, "Create"), RegionRequest::Drop(_) => write!(f, "Drop"), RegionRequest::Open(_) => write!(f, "Open"), + RegionRequest::CleanUp(_) => write!(f, "CleanUp"), RegionRequest::Close(_) => write!(f, "Close"), RegionRequest::Alter(_) => write!(f, "Alter"), RegionRequest::Flush(_) => write!(f, "Flush"), @@ -2203,6 +2236,28 @@ mod tests { assert_eq!(request.requirements, RegionRequirements::object_storage()); } + #[test] + fn test_parse_region_cleanup_from_proto() { + let clean_up = api::v1::region::CleanUpRequest { + region_id: RegionId::new(42, 3).as_u64(), + engine: "mito".to_string(), + path: "test".to_string(), + options: HashMap::from([("k".to_string(), "v".to_string())]), + }; + + let requests = + RegionRequest::try_from_request_body(region_request::Body::CleanUp(clean_up)).unwrap(); + let RegionRequest::CleanUp(request) = &requests[0].1 else { + unreachable!() + }; + + assert_eq!(requests[0].0, RegionId::new(42, 3)); + assert_eq!(request.engine, "mito"); + assert_eq!(request.table_dir, "data/test/42/"); + assert_eq!(request.path_type, PathType::Bare); + assert_eq!(request.options.get("k"), Some(&"v".to_string())); + } + #[test] fn test_validate_modify_column_fulltext_options() { let kind = AlterKind::SetIndexes { From d494d763ee1820958c7c476393549984c88f04c4 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Tue, 14 Jul 2026 16:46:53 +0800 Subject: [PATCH 22/48] feat: enlarge file meta cache (#8499) Signed-off-by: Ruihang Xia --- config/config.md | 4 ++-- config/datanode.example.toml | 4 ++-- config/standalone.example.toml | 4 ++-- src/mito2/src/config.rs | 35 ++++++++++++++++++++++++++++++---- 4 files changed, 37 insertions(+), 10 deletions(-) diff --git a/config/config.md b/config/config.md index a717b0d6134..b0da5722f98 100644 --- a/config/config.md +++ b/config/config.md @@ -157,7 +157,7 @@ | `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size`. | -| `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/32 of OS memory with a max limitation of 128MB. | +| `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. | | `region_engine.mito.vector_cache_size` | String | Auto | Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. | | `region_engine.mito.page_cache_size` | String | Auto | Cache size for pages of SST row groups. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory. | | `region_engine.mito.selector_result_cache_size` | String | Auto | Cache size for time series selector (e.g. `last_value()`). Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. | @@ -560,7 +560,7 @@ | `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size` | -| `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/32 of OS memory with a max limitation of 128MB. | +| `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. | | `region_engine.mito.vector_cache_size` | String | Auto | Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. | | `region_engine.mito.page_cache_size` | String | Auto | Cache size for pages of SST row groups. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory. | | `region_engine.mito.selector_result_cache_size` | String | Auto | Cache size for time series selector (e.g. `last_value()`). Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 7a5a486bbeb..f823d3114be 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -469,9 +469,9 @@ auto_flush_interval = "1h" #+ global_write_buffer_reject_size = "2GB" ## Cache size for SST metadata. Setting it to 0 to disable the cache. -## If not set, it's default to 1/32 of OS memory with a max limitation of 128MB. +## If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. ## @toml2docs:none-default="Auto" -#+ sst_meta_cache_size = "128MB" +#+ sst_meta_cache_size = "512MB" ## Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache. ## If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 074b827a2b5..822f9ba568d 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -598,9 +598,9 @@ auto_flush_interval = "1h" #+ global_write_buffer_reject_size = "2GB" ## Cache size for SST metadata. Setting it to 0 to disable the cache. -## If not set, it's default to 1/32 of OS memory with a max limitation of 128MB. +## If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. ## @toml2docs:none-default="Auto" -#+ sst_meta_cache_size = "128MB" +#+ sst_meta_cache_size = "512MB" ## Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache. ## If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index f9aaf8bb2ee..407fedc955d 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -38,7 +38,11 @@ pub(crate) const DEFAULT_MAX_CONCURRENT_SCAN_FILES: usize = 384; // Use `1/GLOBAL_WRITE_BUFFER_SIZE_FACTOR` of OS memory as global write buffer size in default mode const GLOBAL_WRITE_BUFFER_SIZE_FACTOR: u64 = 8; /// Use `1/SST_META_CACHE_SIZE_FACTOR` of OS memory size as SST meta cache size in default mode -const SST_META_CACHE_SIZE_FACTOR: u64 = 32; +const SST_META_CACHE_SIZE_FACTOR: u64 = 8; +/// Use `1/PREFILTER_RESULT_CACHE_SIZE_FACTOR` of OS memory size as prefilter result cache size in default mode +const PREFILTER_RESULT_CACHE_SIZE_FACTOR: u64 = 32; +/// Use `1/INDEX_METADATA_CACHE_SIZE_FACTOR` of OS memory size as index metadata cache size in default mode +const INDEX_METADATA_CACHE_SIZE_FACTOR: u64 = 32; /// Use `1/MEM_CACHE_SIZE_FACTOR` of OS memory size as mem cache size in default mode const MEM_CACHE_SIZE_FACTOR: u64 = 16; /// Use `1/PAGE_CACHE_SIZE_FACTOR` of OS memory size as page cache size in default mode @@ -321,9 +325,14 @@ impl MitoConfig { ); // Use 2x of global write buffer size as global write buffer reject size. let global_write_buffer_reject_size = global_write_buffer_size * 2; - // shouldn't be greater than 128MB in default mode. + // Page-index-bearing SST metadata can be much larger than footers alone. + // Keep the auto-sized default bounded, but allow a larger warm working set. let sst_meta_cache_size = cmp::min( sys_memory / SST_META_CACHE_SIZE_FACTOR, + ReadableSize::mb(512), + ); + let prefilter_result_cache_size = cmp::min( + sys_memory / PREFILTER_RESULT_CACHE_SIZE_FACTOR, ReadableSize::mb(128), ); // shouldn't be greater than 512MB in default mode. @@ -338,7 +347,7 @@ impl MitoConfig { self.selector_result_cache_size = mem_cache_size; self.range_result_cache_size = mem_cache_size; // Use a smaller cache size because prefilter result usually should be small. - self.prefilter_result_cache_size = sst_meta_cache_size; + self.prefilter_result_cache_size = prefilter_result_cache_size; self.index.adjust_buffer_and_cache_size(sys_memory); } @@ -359,6 +368,24 @@ impl MitoConfig { } } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_adjust_sst_metadata_and_prefilter_cache_caps_independently() { + let mut config = MitoConfig::default(); + + config.adjust_buffer_and_cache_size(ReadableSize::gb(1)); + assert_eq!(ReadableSize::mb(128), config.sst_meta_cache_size); + assert_eq!(ReadableSize::mb(32), config.prefilter_result_cache_size); + + config.adjust_buffer_and_cache_size(ReadableSize::gb(64)); + assert_eq!(ReadableSize::mb(512), config.sst_meta_cache_size); + assert_eq!(ReadableSize::mb(128), config.prefilter_result_cache_size); + } +} + /// Index build mode. #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, Default)] #[serde(rename_all = "snake_case")] @@ -466,7 +493,7 @@ impl IndexConfig { self.content_cache_size = cmp::min(self.content_cache_size, cache_size); let metadata_cache_size = cmp::min( - sys_memory / SST_META_CACHE_SIZE_FACTOR, + sys_memory / INDEX_METADATA_CACHE_SIZE_FACTOR, ReadableSize::mb(64), ); self.metadata_cache_size = cmp::min(self.metadata_cache_size, metadata_cache_size); From 1e94c3adab16eac9033e4bcaf79ffa8d8f5bc913 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 14 Jul 2026 17:34:27 +0800 Subject: [PATCH 23/48] chore(skill): Add ca-certificates to Dockerfile (#8507) as title --- .../skills/greptimedb-development-docker-image/assets/Dockerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/.agents/skills/greptimedb-development-docker-image/assets/Dockerfile b/.agents/skills/greptimedb-development-docker-image/assets/Dockerfile index bc74c830310..4e1a7d2c502 100644 --- a/.agents/skills/greptimedb-development-docker-image/assets/Dockerfile +++ b/.agents/skills/greptimedb-development-docker-image/assets/Dockerfile @@ -1,6 +1,7 @@ FROM ubuntu:24.04 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + ca-certificates \ curl \ unzip \ mysql-client \ From 9cde31a647bc362b9f507040c98dc849fc700964 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 14 Jul 2026 17:35:17 +0800 Subject: [PATCH 24/48] fix: account logical record batch slice memory (#8480) * fix: account logical record batch slice memory Signed-off-by: jeremyhi * fix: account nested view payload memory Signed-off-by: jeremyhi * test: align scan memory expectation with logical size Signed-off-by: jeremyhi * refactor: name Arrow inline view limit Signed-off-by: jeremyhi * fix: account visible list child memory Signed-off-by: jeremyhi * fix: exclude null list child memory Signed-off-by: jeremyhi * fix: exclude null struct child memory Signed-off-by: jeremyhi * fix: account fixed-size list slices Signed-off-by: jeremyhi * fix: limit scan accounting to flat slices Signed-off-by: jeremyhi * docs: clarify nested view accounting scope Signed-off-by: jeremyhi * perf: reduce view memory accounting overhead Signed-off-by: jeremyhi * chore: fix benchmark license header Signed-off-by: jeremyhi * perf: optimize view slice accounting Signed-off-by: jeremyhi * test: cover sliced mixed-null view accounting Signed-off-by: jeremyhi * refactor: narrow slice accounting to arrow buffers Signed-off-by: jeremyhi * docs: note view accounting assumption Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/common/recordbatch/Cargo.toml | 4 + .../recordbatch/benches/memory_accounting.rs | 117 +++++++++++++++++ src/common/recordbatch/src/lib.rs | 33 ++++- src/common/recordbatch/src/recordbatch.rs | 120 +++++++++++++++++- src/table/src/table/scan.rs | 2 +- 5 files changed, 270 insertions(+), 6 deletions(-) create mode 100644 src/common/recordbatch/benches/memory_accounting.rs diff --git a/src/common/recordbatch/Cargo.toml b/src/common/recordbatch/Cargo.toml index efc6b6f60e9..59929ba3477 100644 --- a/src/common/recordbatch/Cargo.toml +++ b/src/common/recordbatch/Cargo.toml @@ -34,3 +34,7 @@ tokio.workspace = true [[bench]] name = "iter_record_batch_rows" harness = false + +[[bench]] +name = "memory_accounting" +harness = false diff --git a/src/common/recordbatch/benches/memory_accounting.rs b/src/common/recordbatch/benches/memory_accounting.rs new file mode 100644 index 00000000000..05fb63555a4 --- /dev/null +++ b/src/common/recordbatch/benches/memory_accounting.rs @@ -0,0 +1,117 @@ +// 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 std::hint::black_box; +use std::sync::Arc; + +use common_recordbatch::RecordBatch; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datatypes::data_type::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; +use datatypes::vectors::{Int32Vector, StringVector, UInt64Vector, VectorRef}; + +const LARGE_ROWS: usize = 8_192; + +fn batch(columns: Vec<(&str, ConcreteDataType, bool, VectorRef)>) -> RecordBatch { + let schema = Arc::new(Schema::new( + columns + .iter() + .map(|(name, data_type, nullable, _)| { + ColumnSchema::new(*name, data_type.clone(), *nullable) + }) + .collect(), + )); + let vectors: Vec = columns + .into_iter() + .map(|(_, _, _, vector)| vector) + .collect(); + RecordBatch::new(schema, vectors).unwrap() +} + +fn primitive_string_batch(rows: usize) -> RecordBatch { + let strings = (0..rows) + .map(|row| (row % 5 != 0).then(|| format!("value-{row:08}-with-payload"))) + .collect::>(); + batch(vec![ + ( + "i32", + ConcreteDataType::int32_datatype(), + false, + Arc::new(Int32Vector::from_slice( + (0..rows).map(|value| value as i32).collect::>(), + )), + ), + ( + "u64", + ConcreteDataType::uint64_datatype(), + false, + Arc::new(UInt64Vector::from_slice( + (0..rows).map(|value| value as u64).collect::>(), + )), + ), + ( + "string", + ConcreteDataType::string_datatype(), + true, + Arc::new(StringVector::from(strings)), + ), + ]) +} + +fn bench_one_call(c: &mut Criterion) { + let mut group = c.benchmark_group("memory_accounting/one_call"); + for rows in [128, LARGE_ROWS] { + let batch = primitive_string_batch(rows); + group.bench_function(BenchmarkId::new("buffer", rows), |b| { + b.iter(|| black_box(black_box(&batch).buffer_memory_size())) + }); + group.bench_function(BenchmarkId::new("logical", rows), |b| { + b.iter(|| black_box(black_box(&batch).logical_slice_memory_size())) + }); + } + group.finish(); +} + +fn bench_shared_buffer_slices(c: &mut Criterion) { + let backing = primitive_string_batch(LARGE_ROWS); + let mut group = c.benchmark_group("memory_accounting/shared_buffer_slices"); + for slice_len in [1, 16] { + let slices = (0..512) + .map(|index| backing.slice(index * slice_len, slice_len).unwrap()) + .collect::>(); + group.throughput(Throughput::Elements(slices.len() as u64)); + group.bench_function(BenchmarkId::new("buffer", slice_len), |b| { + b.iter(|| { + let total = black_box(&slices) + .iter() + .map(RecordBatch::buffer_memory_size) + .fold(0usize, usize::saturating_add); + black_box(total) + }) + }); + group.bench_function(BenchmarkId::new("logical", slice_len), |b| { + b.iter(|| { + let total = black_box(&slices) + .iter() + .map(RecordBatch::logical_slice_memory_size) + .fold(0usize, usize::saturating_add); + black_box(total) + }) + }); + } + group.finish(); +} + +criterion_group!(benches, bench_one_call, bench_shared_buffer_slices); +criterion_main!(benches); diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs index 3db3d9c08b6..46375d8fd77 100644 --- a/src/common/recordbatch/src/lib.rs +++ b/src/common/recordbatch/src/lib.rs @@ -769,7 +769,7 @@ impl MemoryTrackedStream { batch: RecordBatch, cx: &mut Context<'_>, ) -> Poll>> { - let additional = batch.buffer_memory_size(); + let additional = batch.logical_slice_memory_size(); let tracker = self.ready_tracker_mut(); if let Err(error) = tracker.try_track(additional) { @@ -860,6 +860,35 @@ mod tests { as usize } + #[tokio::test] + async fn test_memory_tracked_stream_charges_logical_slice_size() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "payload", + ConcreteDataType::string_datatype(), + false, + )])); + let payloads: Vec<_> = (0..1024) + .map(|value| format!("payload-{value:04}")) + .collect(); + let batch = RecordBatch::new(schema, vec![Arc::new(StringVector::from(payloads)) as _]) + .unwrap() + .slice(512, 1) + .unwrap(); + let expected_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size()); + assert!(expected_bytes < aligned_tracked_bytes(batch.buffer_memory_size())); + let tracker = QueryMemoryTracker::builder(MB, OnExhaustedPolicy::Fail).build(); + let mut stream = MemoryTrackedStream::new( + RecordBatches::try_new(batch.schema.clone(), vec![batch]) + .unwrap() + .as_stream(), + tracker.clone(), + ); + + stream.next().await.unwrap().unwrap(); + + assert_eq!(tracker.current(), expected_bytes); + } + #[test] fn test_recordbatches_try_from_columns() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( @@ -1052,7 +1081,7 @@ mod tests { }) .build(); let batch = large_string_batch(700 * 1024); - let expected_bytes = aligned_tracked_bytes(batch.buffer_memory_size()); + let expected_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size()); let mut stream1 = MemoryTrackedStream::new( RecordBatches::try_new(batch.schema.clone(), vec![batch.clone()]) diff --git a/src/common/recordbatch/src/recordbatch.rs b/src/common/recordbatch/src/recordbatch.rs index 7280772d05d..77a36836d6a 100644 --- a/src/common/recordbatch/src/recordbatch.rs +++ b/src/common/recordbatch/src/recordbatch.rs @@ -269,6 +269,28 @@ impl RecordBatch { .sum() } + /// Returns the logical memory size of this batch's array slices. + /// + /// This sums Arrow's logical visible slice buffers rather than the capacity of their shared + /// backing buffers. View out-of-line payloads and nested custom payloads are not separately + /// traversed or accounted. It is not an exact measure of live physical memory. If Arrow cannot + /// calculate a slice's size, the full buffer size is used conservatively. + /// + /// Mito's current scan paths do not produce top-level View arrays. If they do in the future, + /// their out-of-line payload accounting must be reassessed here. + pub fn logical_slice_memory_size(&self) -> usize { + self.df_record_batch + .columns() + .iter() + .fold(0, |total, array| { + let array_size = array + .to_data() + .get_slice_memory_size() + .unwrap_or_else(|_| array.get_buffer_memory_size()); + total.saturating_add(array_size) + }) + } + /// Iterate the values as strings in the column at index `i`. /// /// Note that if the underlying array is not a valid GreptimeDB vector, an empty iterator is @@ -382,13 +404,14 @@ fn maybe_align_json_array_with_schema( mod tests { use std::sync::Arc; - use datatypes::arrow::array::{AsArray, BinaryArray, UInt32Array}; + use datatypes::arrow::array::{ + AsArray, BinaryArray, StringArray, StringViewArray, UInt32Array, + }; use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema, UInt32Type}; - use datatypes::arrow_array::StringArray; use datatypes::data_type::ConcreteDataType; use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; use datatypes::schema::{ColumnSchema, Schema}; - use datatypes::vectors::{StringVector, UInt32Vector}; + use datatypes::vectors::{BinaryVector, StringVector, UInt32Vector}; use super::*; @@ -475,6 +498,97 @@ mod tests { assert!(recordbatch.slice(1, 5).is_err()); } + #[test] + fn test_logical_slice_memory_size_for_visible_primitive_string_binary_slices() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("numbers", ConcreteDataType::uint32_datatype(), false), + ColumnSchema::new("strings", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("binary", ConcreteDataType::binary_datatype(), true), + ])); + let numbers: Vec<_> = (0..1024).collect(); + let strings = (0..1024) + .map(|value| (value % 3 != 0).then(|| format!("value-{value}"))) + .collect::>(); + let binary = (0_u32..1024) + .map(|value| (value % 3 != 0).then(|| value.to_le_bytes().to_vec())) + .collect::>(); + let columns: Vec = vec![ + Arc::new(UInt32Vector::from_slice(numbers)), + Arc::new(StringVector::from(strings)), + Arc::new(BinaryVector::from(binary)), + ]; + let batch = RecordBatch::new(schema, columns).unwrap(); + let slice = batch.slice(511, 3).unwrap(); + + assert!(slice.columns().iter().any(|column| column.null_count() > 0)); + assert!(slice.logical_slice_memory_size() < slice.buffer_memory_size()); + assert_eq!( + slice.logical_slice_memory_size(), + slice + .columns() + .iter() + .map(|column| column.to_data().get_slice_memory_size().unwrap()) + .sum::() + ); + } + + #[test] + fn test_logical_slice_memory_size_for_many_shared_buffer_slices() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "strings", + ConcreteDataType::string_datatype(), + false, + )])); + let strings = (0..1024) + .map(|value| format!("shared-value-{value}")) + .collect::>(); + let backing = RecordBatch::new( + schema, + vec![Arc::new(StringVector::from(strings)) as VectorRef], + ) + .unwrap(); + let slices = (0..128) + .map(|index| backing.slice(index * 4, 2).unwrap()) + .collect::>(); + let logical_total = slices + .iter() + .map(RecordBatch::logical_slice_memory_size) + .sum::(); + let buffer_total = slices + .iter() + .map(RecordBatch::buffer_memory_size) + .sum::(); + + assert!(logical_total < buffer_total); + assert!(slices.iter().all(|slice| { + slice.logical_slice_memory_size() + == slice.column(0).to_data().get_slice_memory_size().unwrap() + })); + } + + #[test] + fn test_logical_slice_memory_size_uses_arrow_slice_scope_for_views() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "strings", + ConcreteDataType::utf8_view_datatype(), + false, + )])); + let columns: Vec = + vec![Arc::new(StringVector::from(StringViewArray::from(vec![ + "unrelated backing payload", + "visible string view payload", + ])))]; + let batch = RecordBatch::new(schema, columns) + .unwrap() + .slice(1, 1) + .unwrap(); + + assert_eq!( + batch.column(0).to_data().get_slice_memory_size().unwrap(), + batch.logical_slice_memory_size() + ); + } + #[test] fn test_merge_record_batch() { let column_schemas = vec![ diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 825d930a6fc..9ce752553e5 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -541,7 +541,7 @@ impl Stream for StreamWithMetricWrapper { Ok(record_batch) => { // we don't record elapsed time here // since it's calling storage api involving I/O ops - let batch_bytes = record_batch.buffer_memory_size(); + let batch_bytes = record_batch.logical_slice_memory_size(); this.metric.record_output_bytes(batch_bytes); this.metric.record_output(record_batch.num_rows()); Poll::Ready(Some(Ok(record_batch.into_df_record_batch()))) From 7039e54835adfe95992cc48901f59e74075e5094 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 14 Jul 2026 20:04:03 +0800 Subject: [PATCH 25/48] fix: require metasrv GC for repartition (#8497) * fix: require GC for repartition Signed-off-by: WenyXu * test: update sqlness result Signed-off-by: WenyXu * test: enable GC for repartition integration tests Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/meta-srv/src/metasrv/builder.rs | 23 ++++-- src/meta-srv/src/procedure/repartition.rs | 77 ++++++++++++++++++- tests-integration/tests/repartition.rs | 44 +++++++++++ .../repartition/repartition.result | 38 ++++----- .../distributed/repartition/repartition.sql | 14 ++-- 5 files changed, 162 insertions(+), 34 deletions(-) diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 90ce54142f1..8a6f3147bfe 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -28,7 +28,9 @@ use common_meta::ddl::table_meta::{TableMetadataAllocator, TableMetadataAllocato use common_meta::ddl::{ DdlContext, NoopRegionFailureDetectorControl, RegionFailureDetectorControllerRef, }; -use common_meta::ddl_manager::{DdlManager, DdlManagerConfiguratorRef}; +use common_meta::ddl_manager::{ + DdlManager, DdlManagerConfiguratorRef, RepartitionProcedureFactoryRef, +}; use common_meta::distributed_time_constants::default_distributed_time_constants; use common_meta::key::TableMetadataManager; use common_meta::key::flow::FlowMetadataManager; @@ -70,7 +72,9 @@ use crate::metasrv::{ use crate::peer::MetasrvPeerAllocator; use crate::procedure::region_migration::DefaultContextFactory; use crate::procedure::region_migration::manager::RegionMigrationManager; -use crate::procedure::repartition::DefaultRepartitionProcedureFactory; +use crate::procedure::repartition::{ + DefaultRepartitionProcedureFactory, GcDisabledRepartitionProcedureFactory, +}; use crate::procedure::wal_prune::Context as WalPruneContext; use crate::procedure::wal_prune::manager::{WalPruneManager, WalPruneTicker}; use crate::region::flush_trigger::RegionFlushTrigger; @@ -428,10 +432,17 @@ impl MetasrvBuilder { soft_drop_enabled: ddl_soft_drop_enabled(&options), }; let procedure_manager_c = procedure_manager.clone(); - let repartition_procedure_factory = Arc::new(DefaultRepartitionProcedureFactory::new( - mailbox.clone(), - options.grpc.server_addr.clone(), - )); + let repartition_procedure_factory: RepartitionProcedureFactoryRef = if options.gc.enable { + Arc::new(DefaultRepartitionProcedureFactory::new( + mailbox.clone(), + options.grpc.server_addr.clone(), + )) + } else { + Arc::new(GcDisabledRepartitionProcedureFactory::new( + mailbox.clone(), + options.grpc.server_addr.clone(), + )) + }; let ddl_manager = DdlManager::try_new( ddl_context, procedure_manager_c, diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index 53e4bacf1ec..2aa8f18799c 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -808,6 +808,50 @@ impl DefaultRepartitionProcedureFactory { } } +/// Rejects new repartition requests when metasrv GC is disabled. +/// +/// Procedure loaders are still delegated to the enabled factory so procedures +/// persisted before a metasrv restart can be recovered. +pub struct GcDisabledRepartitionProcedureFactory { + enabled_factory: DefaultRepartitionProcedureFactory, +} + +impl GcDisabledRepartitionProcedureFactory { + pub fn new(mailbox: MailboxRef, server_addr: String) -> Self { + Self { + enabled_factory: DefaultRepartitionProcedureFactory::new(mailbox, server_addr), + } + } +} + +impl RepartitionProcedureFactory for GcDisabledRepartitionProcedureFactory { + fn create( + &self, + _ddl_ctx: &DdlContext, + _table_name: TableName, + _table_id: TableId, + _source: RepartitionSource, + _to_exprs: Vec, + _timeout: Option, + ) -> std::result::Result { + Err(BoxedError::new( + error::InvalidArgumentsSnafu { + err_msg: "Repartition requires metasrv GC to be enabled".to_string(), + } + .build(), + )) + } + + fn register_loaders( + &self, + ddl_ctx: &DdlContext, + procedure_manager: &ProcedureManagerRef, + ) -> std::result::Result<(), BoxedError> { + self.enabled_factory + .register_loaders(ddl_ctx, procedure_manager) + } +} + impl RepartitionProcedureFactory for DefaultRepartitionProcedureFactory { fn create( &self, @@ -919,7 +963,7 @@ mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; - use common_error::ext::BoxedError; + use common_error::ext::{BoxedError, ErrorExt}; use common_error::mock::MockError; use common_error::status_code::StatusCode; use common_meta::ddl::test_util::datanode_handler::{ @@ -1050,6 +1094,37 @@ mod tests { ) } + #[test] + fn test_gc_disabled_factory_rejects_repartition() { + let env = TestingEnv::new(); + let node_manager = Arc::new(MockDatanodeManager::new(UnexpectedErrorDatanodeHandler)); + let ddl_ctx = env.ddl_context(node_manager); + let factory = GcDisabledRepartitionProcedureFactory::new( + env.mailbox_ctx.mailbox().clone(), + env.server_addr.clone(), + ); + + let err = factory + .create( + &ddl_ctx, + TableName::new("test_catalog", "test_schema", "test_table"), + 1024, + RepartitionSource::Unpartitioned { + partition_columns: vec![], + }, + vec![], + None, + ) + .err() + .expect("GC-disabled factory must reject repartition"); + + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert_eq!( + "Invalid arguments: Repartition requires metasrv GC to be enabled", + err.to_string() + ); + } + #[test] fn test_filter_allocated_region_routes() { let table_id = 1024; diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index 32f80a32cba..cc55cb1bfdc 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -345,6 +345,17 @@ pub async fn test_partition_unpartitioned_mito(store_type: StorageType) { .with_datanodes(datanodes as u32) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + gc_cooldown_period: Duration::from_nanos(1), + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + lingering_time: Some(Duration::from_secs(0)), + unknown_file_lingering_time: Duration::from_secs(0), + ..Default::default() + }) .build(true) .await; @@ -472,6 +483,17 @@ pub async fn test_repartition_on_columns_metadata_mito(store_type: StorageType) .with_datanodes(datanodes as u32) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + gc_cooldown_period: Duration::from_nanos(1), + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + lingering_time: Some(Duration::from_secs(0)), + unknown_file_lingering_time: Duration::from_secs(0), + ..Default::default() + }) .build(true) .await; @@ -557,6 +579,17 @@ pub async fn test_repartition_on_columns_data_correctness_mito(store_type: Stora .with_datanodes(datanodes as u32) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + gc_cooldown_period: Duration::from_nanos(1), + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + lingering_time: Some(Duration::from_secs(0)), + unknown_file_lingering_time: Duration::from_secs(0), + ..Default::default() + }) .build(true) .await; @@ -612,6 +645,17 @@ pub async fn test_partition_unpartitioned_metric(store_type: StorageType) { .with_datanodes(datanodes as u32) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + gc_cooldown_period: Duration::from_nanos(1), + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + lingering_time: Some(Duration::from_secs(0)), + unknown_file_lingering_time: Duration::from_secs(0), + ..Default::default() + }) .build(true) .await; diff --git a/tests/cases/distributed/repartition/repartition.result b/tests/cases/distributed/repartition/repartition.result index 72212699af8..7b25a13a061 100644 --- a/tests/cases/distributed/repartition/repartition.result +++ b/tests/cases/distributed/repartition/repartition.result @@ -12,6 +12,7 @@ CREATE TABLE alter_repartition_table( Affected Rows: 0 +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE alter_repartition_table REPARTITION ( device_id < 100 ) INTO ( @@ -19,7 +20,7 @@ ALTER TABLE alter_repartition_table REPARTITION ( device_id < 100 AND area >= 'South' ); -Affected Rows: 0 +Error: 1004(InvalidArguments), Invalid arguments: Repartition requires metasrv GC to be enabled SHOW CREATE TABLE alter_repartition_table; @@ -35,21 +36,21 @@ SHOW CREATE TABLE alter_repartition_table; | | PRIMARY KEY ("device_id") | | | ) | | | PARTITION ON COLUMNS ("device_id", "area") ( | -| | device_id < 100 AND area < 'South', | +| | device_id < 100, | | | device_id >= 100 AND device_id < 200, | -| | device_id >= 200, | -| | device_id < 100 AND area >= 'South' | +| | device_id >= 200 | | | ) | | | ENGINE=mito | | | | +-------------------------+--------------------------------------------------------+ +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE alter_repartition_table MERGE PARTITION ( - device_id < 100 AND area < 'South', - device_id < 100 AND area >= 'South' + device_id < 100, + device_id >= 100 AND device_id < 200, ); -Affected Rows: 0 +Error: 1004(InvalidArguments), Invalid arguments: Repartition requires metasrv GC to be enabled SHOW CREATE TABLE alter_repartition_table; @@ -124,11 +125,7 @@ ALTER TABLE alter_repartition_table_with_options REPARTITION ( WAIT = false ); -+--------------------------------------+ -| Procedure ID | -+--------------------------------------+ -| PROC_ID | -+--------------------------------------+ +Error: 1004(InvalidArguments), Invalid arguments: Repartition requires metasrv GC to be enabled DROP TABLE alter_repartition_table_with_options; @@ -177,7 +174,7 @@ WITH ( Affected Rows: 0 --- Split physical table partition +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE metric_physical_table SPLIT PARTITION ( host < 'h1' ) INTO ( @@ -185,7 +182,7 @@ ALTER TABLE metric_physical_table SPLIT PARTITION ( host >= 'h0' AND host < 'h1' ); -Affected Rows: 0 +Error: 1004(InvalidArguments), Invalid arguments: Repartition requires metasrv GC to be enabled SHOW CREATE TABLE metric_physical_table; @@ -200,10 +197,9 @@ SHOW CREATE TABLE metric_physical_table; | | PRIMARY KEY ("host") | | | ) | | | PARTITION ON COLUMNS ("host") ( | -| | host < 'h0', | +| | host < 'h1', | | | host >= 'h1' AND host < 'h2', | -| | host >= 'h2', | -| | host >= 'h0' AND host < 'h1' | +| | host >= 'h2' | | | ) | | | ENGINE=metric | | | WITH( | @@ -227,13 +223,13 @@ SELECT * FROM logical_table_v2; ++ ++ --- Merge physical table partition +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE metric_physical_table MERGE PARTITION ( - host < 'h0', - host >= 'h0' AND host < 'h1' + host < 'h1', + host >= 'h1' AND host < 'h2' ); -Affected Rows: 0 +Error: 1004(InvalidArguments), Invalid arguments: Repartition requires metasrv GC to be enabled SHOW CREATE TABLE metric_physical_table; diff --git a/tests/cases/distributed/repartition/repartition.sql b/tests/cases/distributed/repartition/repartition.sql index 92b1d6ad97b..180969f96fa 100644 --- a/tests/cases/distributed/repartition/repartition.sql +++ b/tests/cases/distributed/repartition/repartition.sql @@ -10,6 +10,7 @@ CREATE TABLE alter_repartition_table( device_id >= 200 ); +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE alter_repartition_table REPARTITION ( device_id < 100 ) INTO ( @@ -19,9 +20,10 @@ ALTER TABLE alter_repartition_table REPARTITION ( SHOW CREATE TABLE alter_repartition_table; +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE alter_repartition_table MERGE PARTITION ( - device_id < 100 AND area < 'South', - device_id < 100 AND area >= 'South' + device_id < 100, + device_id >= 100 AND device_id < 200, ); SHOW CREATE TABLE alter_repartition_table; @@ -109,7 +111,7 @@ WITH ( on_physical_table = "metric_physical_table" ); --- Split physical table partition +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE metric_physical_table SPLIT PARTITION ( host < 'h1' ) INTO ( @@ -126,10 +128,10 @@ SELECT * FROM logical_table_v1; SELECT * FROM logical_table_v2; --- Merge physical table partition +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE metric_physical_table MERGE PARTITION ( - host < 'h0', - host >= 'h0' AND host < 'h1' + host < 'h1', + host >= 'h1' AND host < 'h2' ); SHOW CREATE TABLE metric_physical_table; From b513dbaf4a3b21add12c031e5abc965251f80423 Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 14 Jul 2026 20:48:13 +0800 Subject: [PATCH 26/48] fix: reject datanode startup on GC config mismatch (#8509) * fix: reject datanode gc config mismatch Signed-off-by: discord9 * refactor: minimize datanode gc startup check Signed-off-by: discord9 * chore: update greptime-proto revision Signed-off-by: discord9 * chore: use merged greptime-proto revision Signed-off-by: discord9 --------- Signed-off-by: discord9 --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/datanode/src/error.rs | 13 +++++++++++++ src/datanode/src/heartbeat.rs | 17 +++++++++++++++++ src/frontend/src/frontend.rs | 1 + src/meta-client/src/client/heartbeat.rs | 8 +++++--- src/meta-srv/src/handler.rs | 7 ++++--- src/meta-srv/src/handler/test_utils.rs | 1 + src/meta-srv/src/metasrv.rs | 3 +++ 9 files changed, 46 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7a9ec59760f..cb6cb42c2e2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5938,7 +5938,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=6fc4d88851c5ad16846698e4dff6e74c1f7a48ba#6fc4d88851c5ad16846698e4dff6e74c1f7a48ba" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=c6f92ee9e12ee7bd5ee3c0956b1c9caa39be3f76#c6f92ee9e12ee7bd5ee3c0956b1c9caa39be3f76" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", diff --git a/Cargo.toml b/Cargo.toml index 5fecc43fa5b..31093a556c6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -158,7 +158,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "6fc4d88851c5ad16846698e4dff6e74c1f7a48ba" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "c6f92ee9e12ee7bd5ee3c0956b1c9caa39be3f76" } hex = "0.4" http = "1" humantime = "2.1" diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index 0e2e1da959c..6c176547edd 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -287,6 +287,18 @@ pub enum Error { location: Location, }, + #[snafu(display( + "GC configuration mismatch: metasrv.gc.enable={}, datanode.region_engine.mito.gc.enable={}", + metasrv_gc_enabled, + datanode_gc_enabled, + ))] + GcConfigMismatch { + metasrv_gc_enabled: bool, + datanode_gc_enabled: bool, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Unsupported output type, expected: {}", expected))] UnsupportedOutput { expected: String, @@ -456,6 +468,7 @@ impl ErrorExt for Error { | ColumnNoneDefaultValue { .. } | MissingRequiredField { .. } | RegionEngineNotFound { .. } + | GcConfigMismatch { .. } | ParseAddr { .. } | TomlFormat { .. } | BuildDatanode { .. } => StatusCode::InvalidArguments, diff --git a/src/datanode/src/heartbeat.rs b/src/datanode/src/heartbeat.rs index bfdaf32a1c6..0e0c51624e3 100644 --- a/src/datanode/src/heartbeat.rs +++ b/src/datanode/src/heartbeat.rs @@ -121,6 +121,7 @@ impl HeartbeatTask { pub async fn create_streams( meta_client: &MetaClient, + local_gc_enabled: bool, running: Arc, handler_executor: HeartbeatResponseHandlerExecutorRef, mailbox: MailboxRef, @@ -129,6 +130,13 @@ impl HeartbeatTask { ) -> Result<(HeartbeatSender, HeartbeatConfig)> { let client_id = meta_client.id(); let (tx, mut rx, config) = meta_client.heartbeat().await.context(MetaClientInitSnafu)?; + if config.gc_enabled != local_gc_enabled { + return error::GcConfigMismatchSnafu { + metasrv_gc_enabled: config.gc_enabled, + datanode_gc_enabled: local_gc_enabled, + } + .fail(); + } let mut last_received_lease = Instant::now(); @@ -220,9 +228,17 @@ impl HeartbeatTask { let mailbox = Arc::new(HeartbeatMailbox::new(outgoing_tx)); let quit_signal = Arc::new(Notify::new()); + let local_gc_enabled = self + .region_server + .mito_engine() + .context(RegionEngineNotFoundSnafu { name: "mito" })? + .mito_config() + .gc + .enable; let (mut tx, config) = Self::create_streams( &meta_client, + local_gc_enabled, running.clone(), handler_executor.clone(), mailbox.clone(), @@ -368,6 +384,7 @@ impl HeartbeatTask { error!(e; "Failed to send heartbeat to metasrv"); match Self::create_streams( &meta_client, + local_gc_enabled, running.clone(), handler_executor.clone(), mailbox.clone(), diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 918185cb8f6..c7f7e71d34c 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -237,6 +237,7 @@ mod tests { is_handshake.then_some(api::v1::meta::HeartbeatConfig { heartbeat_interval_ms, retry_interval_ms: heartbeat_interval_ms, + gc_enabled: false, }); is_handshake = false; let response = HeartbeatResponse { diff --git a/src/meta-client/src/client/heartbeat.rs b/src/meta-client/src/client/heartbeat.rs index c31369dac78..add7bd54de6 100644 --- a/src/meta-client/src/client/heartbeat.rs +++ b/src/meta-client/src/client/heartbeat.rs @@ -39,6 +39,7 @@ use crate::error::{InvalidResponseHeaderSnafu, Result}; pub struct HeartbeatConfig { pub interval: Duration, pub retry_interval: Duration, + pub gc_enabled: bool, } impl Default for HeartbeatConfig { @@ -46,6 +47,7 @@ impl Default for HeartbeatConfig { Self { interval: BASE_HEARTBEAT_INTERVAL, retry_interval: BASE_HEARTBEAT_INTERVAL, + gc_enabled: false, } } } @@ -54,8 +56,8 @@ impl fmt::Display for HeartbeatConfig { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!( f, - "interval={:?}, retry={:?}", - self.interval, self.retry_interval + "interval={:?}, retry={:?}, gc_enabled={}", + self.interval, self.retry_interval, self.gc_enabled ) } } @@ -68,6 +70,7 @@ impl HeartbeatConfig { Self { interval: Duration::from_millis(cfg.heartbeat_interval_ms), retry_interval: Duration::from_millis(cfg.retry_interval_ms), + gc_enabled: cfg.gc_enabled, } } else { let fallback = Self::default(); @@ -259,7 +262,6 @@ impl Inner { .map_err(error::Error::from)? .context(error::CreateHeartbeatStreamSnafu)?; - // Extract heartbeat configuration from handshake response let config = HeartbeatConfig::from_response(&res); info!( diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs index 9cfd4e6079d..72d5fcdab8a 100644 --- a/src/meta-srv/src/handler.rs +++ b/src/meta-srv/src/handler.rs @@ -21,8 +21,8 @@ use std::time::{Duration, Instant}; use api::v1::meta::mailbox_message::Payload; use api::v1::meta::{ - HeartbeatRequest, HeartbeatResponse, MailboxMessage, PROTOCOL_VERSION, RegionLease, - ResponseHeader, Role, + HeartbeatConfig, HeartbeatRequest, HeartbeatResponse, MailboxMessage, PROTOCOL_VERSION, + RegionLease, ResponseHeader, Role, }; use check_leader_handler::CheckLeaderHandler; use collect_cluster_info_handler::{ @@ -387,7 +387,8 @@ impl HeartbeatHandlerGroup { // Populate heartbeat_config during handshake let heartbeat_config = if is_handshake { - let config = ctx.heartbeat_options_for(role).into(); + let mut config: HeartbeatConfig = ctx.heartbeat_options_for(role).into(); + config.gc_enabled = ctx.gc_enabled; info!( "Handshake with {:?} node, sending config: {:?}", diff --git a/src/meta-srv/src/handler/test_utils.rs b/src/meta-srv/src/handler/test_utils.rs index 7bfdfbee792..d8241c785ea 100644 --- a/src/meta-srv/src/handler/test_utils.rs +++ b/src/meta-srv/src/handler/test_utils.rs @@ -92,6 +92,7 @@ impl TestEnv { leader_region_registry: self.leader_region_registry.clone(), topic_stats_registry: self.topic_stats_registry.clone(), heartbeat_interval: BASE_HEARTBEAT_INTERVAL, + gc_enabled: false, is_handshake: false, } } diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index f7f5bbf77da..7480b8cb500 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -175,6 +175,7 @@ impl From for HeartbeatConfig { Self { heartbeat_interval_ms: opts.interval.as_millis() as u64, retry_interval_ms: opts.retry_interval.as_millis() as u64, + gc_enabled: false, } } } @@ -438,6 +439,7 @@ pub struct Context { pub leader_region_registry: LeaderRegionRegistryRef, pub topic_stats_registry: TopicStatsRegistryRef, pub heartbeat_interval: Duration, + pub gc_enabled: bool, pub is_handshake: bool, } @@ -932,6 +934,7 @@ impl Metasrv { leader_region_registry, topic_stats_registry, heartbeat_interval: self.options().heartbeat_interval, + gc_enabled: self.options().gc.enable, is_handshake: false, } } From b30765e5e96388ddfefbe9297a5bef234a54f4eb Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 14 Jul 2026 20:53:07 +0800 Subject: [PATCH 27/48] fix: preserve distributed topk merge ordering (#8432) * fix: preserve distributed topk for scalar latest Signed-off-by: discord9 * fix: preserve scalar latest ordering across merge scan Signed-off-by: discord9 * fix: carry merge scan ordering from planner rewrite Signed-off-by: discord9 * test: cover merge scan ordering metadata Signed-off-by: discord9 * test: cover merge scan partition ordering gate Signed-off-by: discord9 * fix: require adjacent merge sort for ordering metadata Signed-off-by: discord9 * test: update ordering sqlness plans Signed-off-by: discord9 * fix: preserve distributed topk merge ordering Signed-off-by: discord9 * test: cover latest per series queries Signed-off-by: discord9 * test: update distributed merge sort sqlness Signed-off-by: discord9 * test: cover merge scan ordering partition gate Signed-off-by: discord9 * test: cover merge scan over-partition ordering Signed-off-by: discord9 * test: cover distributed latest with low parallelism Signed-off-by: discord9 * fix: keep distributed merge sort opaque Signed-off-by: discord9 * test: assert distributed scalar latest merge sort Signed-off-by: discord9 * test: clarify merge scan ordering helper Signed-off-by: discord9 * fix: complete merge sort exec delegation Signed-off-by: discord9 * docs: explain merge sort limit pushdown Signed-off-by: discord9 * test: cover merge sort optimizer opacity Signed-off-by: discord9 * docs: explain merge sort optimizer hooks Signed-off-by: discord9 * test: cover merge sort optimizer hooks Signed-off-by: discord9 * fix: preserve merge sort child topk Signed-off-by: discord9 * test: update order by topk plan Signed-off-by: discord9 * fix: recognize merge sort global fetch Signed-off-by: discord9 --------- Signed-off-by: discord9 --- src/query/src/dist_plan.rs | 1 + src/query/src/dist_plan/merge_scan.rs | 61 +- src/query/src/dist_plan/merge_sort.rs | 673 +++++++++++++++++- src/query/src/dist_plan/planner.rs | 83 ++- src/query/src/optimizer/global_limit.rs | 43 +- tests-integration/src/tests/instance_test.rs | 103 +++ .../explain/multi_partitions.result | 2 +- .../explain/step_aggr_advance.result | 20 +- .../distributed/explain/subqueries.result | 12 +- .../optimizer/first_value_advance.result | 8 +- .../optimizer/last_value_advance.result | 8 +- .../time_index_filter_pushdown.result | 76 ++ .../optimizer/time_index_filter_pushdown.sql | 45 +- .../standalone/common/order/order_by.result | 2 +- .../standalone/common/tql/partition.result | 2 +- .../common/window/latest_per_series.result | 187 +++++ .../common/window/latest_per_series.sql | 134 ++++ .../optimizer/first_value_advance.result | 8 +- .../optimizer/last_value_advance.result | 8 +- .../time_index_filter_pushdown.result | 25 + .../optimizer/time_index_filter_pushdown.sql | 19 + 21 files changed, 1447 insertions(+), 73 deletions(-) create mode 100644 tests/cases/standalone/common/window/latest_per_series.result create mode 100644 tests/cases/standalone/common/window/latest_per_series.sql diff --git a/src/query/src/dist_plan.rs b/src/query/src/dist_plan.rs index 1983e7824c7..969f00e7c9c 100644 --- a/src/query/src/dist_plan.rs +++ b/src/query/src/dist_plan.rs @@ -27,6 +27,7 @@ mod remote_dyn_filter_registry; pub use analyzer::{DistPlannerAnalyzer, DistPlannerOptions}; pub use filter_id::{FilterFingerprint, FilterId, ParseFilterIdError, RemoteDynFilterProducerId}; pub use merge_scan::{MergeScanExec, MergeScanLogicalPlan}; +pub(crate) use merge_sort::MergeSortExec; pub use planner::{DistExtensionPlanner, MergeSortExtensionPlanner}; pub use predicate_extractor::PredicateExtractor; pub use region_pruner::ConstraintPruner; diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 39bae738c65..126ee9cb068 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -1031,7 +1031,7 @@ mod tests { use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult; use datafusion_common::TableReference; - use datafusion_expr::{LogicalPlanBuilder, lit}; + use datafusion_expr::{LogicalPlanBuilder, col, lit}; use datafusion_physical_expr::Distribution; use datafusion_physical_expr::expressions::{ Column, DynamicFilterPhysicalExpr, lit as physical_lit, @@ -1051,6 +1051,65 @@ mod tests { QueryId::from(Uuid::from_u128(value)) } + fn merge_scan_exec_with_sorted_input( + region_count: u64, + target_partition: usize, + ) -> MergeScanExec { + let session_state = SessionStateBuilder::new().build(); + let plan = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("ts")]) + .unwrap() + .sort(vec![col("ts").sort(false, true)]) + .unwrap() + .build() + .unwrap(); + let schema = plan.schema().as_arrow().clone(); + let regions = (0..region_count) + .map(|region_number| RegionId::new(1024, region_number as u32)) + .collect(); + + MergeScanExec::new( + &session_state, + // The table name is not relevant to these ordering metadata tests; + // `MergeScanExec::new` requires one to model the production plan. + TableName::new("catalog", "schema", "table"), + regions, + plan, + &schema, + Arc::new(TestRegionQueryHandler), + QueryContext::arc(), + target_partition, + AliasMapping::new(), + None, + false, + ) + .unwrap() + } + + #[test] + fn merge_scan_does_not_advertise_ordering_when_partition_may_merge_regions() { + let exec = merge_scan_exec_with_sorted_input(3, 2); + + assert!( + exec.properties().output_ordering().is_none(), + "target_partition < region_count means one output partition may concatenate multiple sorted region streams" + ); + } + + #[test] + fn merge_scan_advertises_ordering_when_each_partition_reads_at_most_one_region() { + let exec = merge_scan_exec_with_sorted_input(3, 3); + + assert!(exec.properties().output_ordering().is_some()); + } + + #[test] + fn merge_scan_advertises_ordering_when_partitions_exceed_regions() { + let exec = merge_scan_exec_with_sorted_input(3, 4); + + assert!(exec.properties().output_ordering().is_some()); + } + #[test] fn remote_dyn_filter_region_query_context_registers_before_do_get() { let registry_manager = Arc::new(DynFilterRegistryManager::default()); diff --git a/src/query/src/dist_plan/merge_sort.rs b/src/query/src/dist_plan/merge_sort.rs index d9bc6119c08..fca81dc4a52 100644 --- a/src/query/src/dist_plan/merge_sort.rs +++ b/src/query/src/dist_plan/merge_sort.rs @@ -16,11 +16,22 @@ //! `SortPreservingMergeExec` operator in datafusion //! +use std::any::Any; use std::fmt; use std::sync::Arc; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::projection::{ProjectionExec, make_with_child, update_ordering}; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, + Statistics, +}; use datafusion_common::{DataFusionError, Result}; use datafusion_expr::{Extension, LogicalPlan, SortExpr, UserDefinedLogicalNodeCore}; +use datafusion_physical_expr::{Distribution, LexOrdering, OrderingRequirements}; /// MergeSort Logical Plan, have same field as `Sort`, but indicate it is a merge sort, /// which assume each input partition is a sorted stream, and will use `SortPreserveingMergeExec` @@ -53,14 +64,245 @@ impl MergeSortLogicalPlan { node: Arc::new(self), }) } +} - /// Convert self to a [`Sort`] logical plan with same input and expressions - pub fn into_sort(self) -> LogicalPlan { - LogicalPlan::Sort(datafusion::logical_expr::Sort { - input: self.input.clone(), - expr: self.expr, - fetch: self.fetch, - }) +/// An opaque physical execution node for [`MergeSortLogicalPlan`]. +/// +/// It delegates execution and physical properties to DataFusion's +/// [`SortPreservingMergeExec`], but intentionally does not expose itself as a +/// `SortPreservingMergeExec`. `EnforceSorting` is allowed to replace a bare +/// `SortPreservingMergeExec` with `CoalescePartitionsExec` when the parent does +/// not require ordering. `MergeSortExec` represents the distributed TopK merge +/// stage itself, so later physical optimizer rules must not rewrite it into an +/// unordered fetch. +#[derive(Debug, Clone)] +pub(crate) struct MergeSortExec { + inner: SortPreservingMergeExec, +} + +impl MergeSortExec { + pub(crate) fn new( + ordering: LexOrdering, + input: Arc, + fetch: Option, + ) -> Self { + Self { + inner: SortPreservingMergeExec::new(ordering, input).with_fetch(fetch), + } + } + + fn input_with_fetch(&self, fetch: Option) -> Arc { + let input = Arc::clone(self.inner.input()); + if let Some(sort) = input.as_any().downcast_ref::() + && sort.preserve_partitioning() + && sort.expr() == self.inner.expr() + { + // Mirror DataFusion's bare SPM plan quality for distributed TopK: + // keep the parent `MergeSortExec(fetch)` as the global merge, and + // bound the partition-preserving child sort to the same local TopK. + // Local top-K is safe because every global top-K row must be within + // the top-K rows of its own input partition. + Arc::new(sort.with_fetch(fetch)) + } else { + input + } + } +} + +impl DisplayAs for MergeSortExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "MergeSortExec: [{}]", self.inner.expr())?; + if let Some(fetch) = self.inner.fetch() { + write!(f, ", fetch={fetch}")?; + } + Ok(()) + } + DisplayFormatType::TreeRender => { + if let Some(fetch) = self.inner.fetch() { + writeln!(f, "limit={fetch}")?; + } + + for (i, expr) in self.inner.expr().iter().enumerate() { + expr.fmt_sql(f)?; + if i != self.inner.expr().len() - 1 { + write!(f, ", ")?; + } + } + + Ok(()) + } + } + } +} + +impl ExecutionPlan for MergeSortExec { + fn name(&self) -> &str { + "MergeSortExec" + } + + /// Keeps this node intentionally opaque to DataFusion's type-specialized + /// optimizer rewrites. + /// + /// `MergeSortExec` delegates most behavior to DataFusion's + /// `SortPreservingMergeExec`, but it must not expose itself as that type. + /// DataFusion's `EnforceSorting` optimizer recognizes a bare + /// `SortPreservingMergeExec` via `as_any().downcast_ref::<...>()` and may + /// replace it with an unordered `CoalescePartitionsExec(fetch)` when the + /// parent does not require sorted output. + /// + /// That rewrite is valid for an ordinary SPM used only to satisfy parent + /// ordering, but not for GreptimeDB's distributed TopK merge stage. In a + /// scalar-subquery shape like `ORDER BY ts DESC LIMIT 1`, this node is the + /// operator that merges region-local TopK streams into the global TopK. + /// Replacing it with unordered coalescing can return a partial/latest row + /// from one region instead of the global latest row. + /// + /// `required_input_ordering()` separately tells DataFusion what ordering this + /// node needs from its child, so `EnforceSorting` can insert a `SortExec` + /// below `MergeSortExec` when `MergeScanExec` cannot preserve per-partition + /// ordering. This opacity is specifically about protecting the merge stage + /// itself from the `EnforceSorting` rewrite above. + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &Arc { + self.inner.properties() + } + + /// Forwards DataFusion's order-preserving scan hint through this wrapper. + /// + /// This mirrors `SortPreservingMergeExec::with_preserve_order()`: if the + /// child can produce an order-preserving variant, rebuild the same merge + /// stage on top of that child. The returned plan must stay a + /// `MergeSortExec`, not a bare SPM, so the distributed TopK merge remains + /// opaque to `EnforceSorting`'s SPM-specific rewrite. + fn with_preserve_order(&self, preserve_order: bool) -> Option> { + self.inner + .input() + .with_preserve_order(preserve_order) + .map(|new_input| { + Arc::new(Self::new( + self.inner.expr().clone(), + new_input, + self.inner.fetch(), + )) as Arc + }) + } + + fn required_input_distribution(&self) -> Vec { + self.inner.required_input_distribution() + } + + fn benefits_from_input_partitioning(&self) -> Vec { + self.inner.benefits_from_input_partitioning() + } + + /// Tells DataFusion that `MergeSortExec` requires each input partition to be + /// ordered. This is the contract that makes `EnforceSorting` insert a + /// `SortExec` below `MergeSortExec` when the input cannot preserve ordering. + /// + /// The opacity of `MergeSortExec::as_any`, not this requirement, is what + /// prevents DataFusion from rewriting the merge stage itself as a bare + /// `SortPreservingMergeExec`. + fn required_input_ordering(&self) -> Vec> { + vec![Some(OrderingRequirements::from(self.inner.expr().clone()))] + } + + fn maintains_input_order(&self) -> Vec { + self.inner.maintains_input_order() + } + + fn children(&self) -> Vec<&Arc> { + self.inner.children() + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(DataFusionError::Internal(format!( + "MergeSortExec expects exactly one child, got {}", + children.len() + ))); + } + + Ok(Arc::new(Self::new( + self.inner.expr().clone(), + children.swap_remove(0), + self.inner.fetch(), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.execute(partition, context) + } + + fn metrics(&self) -> Option { + self.inner.metrics() + } + + fn partition_statistics(&self, partition: Option) -> Result { + self.inner.partition_statistics(partition) + } + + /// Intentionally keeps DataFusion's generic limit pushdown disabled. + /// + /// `MergeSortExec` still supports its own global fetch through + /// `with_fetch()`. What we must not allow is pushing an external limit below + /// this required distributed TopK merge. DataFusion's limit pushdown rules + /// know how to treat a bare `SortPreservingMergeExec` as a + /// partition-combining node, but `MergeSortExec` is intentionally opaque to + /// those SPM-specific downcasts. Enabling generic limit pushdown without also + /// teaching the optimizer about this wrapper could return partition-local + /// rows instead of the global TopK. + fn supports_limit_pushdown(&self) -> bool { + false + } + + fn fetch(&self) -> Option { + self.inner.fetch() + } + + fn with_fetch(&self, limit: Option) -> Option> { + Some(Arc::new(Self::new( + self.inner.expr().clone(), + self.input_with_fetch(limit), + limit, + ))) + } + + /// Lets DataFusion push a projection below this merge when it can rewrite + /// the ordering expressions safely. + /// + /// This mirrors `SortPreservingMergeExec::try_swapping_with_projection()` + /// for plan quality, but re-wraps the result as `MergeSortExec` so the + /// distributed merge stage keeps its type identity and opacity. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result>> { + if projection.expr().len() >= projection.input().schema().fields().len() { + return Ok(None); + } + + let Some(updated_exprs) = update_ordering(self.inner.expr().clone(), projection.expr())? + else { + return Ok(None); + }; + + Ok(Some(Arc::new(Self::new( + updated_exprs, + make_with_child(projection, self.inner.input())?, + self.inner.fetch(), + )))) } } @@ -127,3 +369,420 @@ pub fn merge_sort_transformer(plan: &LogicalPlan) -> Option { None } } + +#[cfg(test)] +mod tests { + use arrow_schema::{DataType, Field, Schema, SortOptions}; + use datafusion::physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{ + OrderPreservationContext, plan_with_order_breaking_variants, + }; + use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; + use datafusion::physical_plan::displayable; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_physical_expr::expressions::col as physical_col; + + use super::*; + + /// Test double that records DataFusion's preserve-order signal while + /// otherwise behaving like a transparent wrapper around its child. + #[derive(Debug, Clone)] + struct PreserveOrderProbeExec { + inner: Arc, + preserve_order: bool, + } + + impl PreserveOrderProbeExec { + fn new(inner: Arc) -> Self { + Self { + inner, + preserve_order: false, + } + } + } + + impl DisplayAs for PreserveOrderProbeExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "PreserveOrderProbeExec: preserve_order={}", + self.preserve_order + ) + } + } + + impl ExecutionPlan for PreserveOrderProbeExec { + fn name(&self) -> &str { + "PreserveOrderProbeExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &Arc { + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.inner] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(DataFusionError::Internal(format!( + "PreserveOrderProbeExec expects exactly one child, got {}", + children.len() + ))); + } + + Ok(Arc::new(Self { + inner: children.swap_remove(0), + preserve_order: self.preserve_order, + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.execute(partition, context) + } + + fn with_preserve_order(&self, preserve_order: bool) -> Option> { + Some(Arc::new(Self { + inner: Arc::clone(&self.inner), + preserve_order, + })) + } + } + + fn test_ordering(schema: &Schema) -> LexOrdering { + LexOrdering::new([PhysicalSortExpr::new( + physical_col("ts", schema).unwrap(), + SortOptions { + descending: true, + nulls_first: false, + }, + )]) + .unwrap() + } + + #[test] + fn merge_sort_exec_is_opaque_and_preserves_topk_requirements() { + let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _; + let ordering = test_ordering(schema.as_ref()); + + let merge_sort = + Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc; + + assert_eq!(merge_sort.name(), "MergeSortExec"); + assert!( + merge_sort + .as_any() + .downcast_ref::() + .is_none(), + "MergeSortExec must stay opaque to EnforceSorting's bare SortPreservingMerge rewrite" + ); + assert_eq!(merge_sort.fetch(), Some(1)); + assert!(!merge_sort.supports_limit_pushdown()); + assert!(merge_sort.required_input_ordering()[0].is_some()); + + let tree = displayable(merge_sort.as_ref()).tree_render().to_string(); + assert!(tree.contains("MergeSortExec")); + assert!(!tree.contains("SortPreservingMergeExec")); + + let fetched = merge_sort.with_fetch(Some(2)).unwrap(); + assert!(fetched.as_any().downcast_ref::().is_some()); + assert_eq!(fetched.fetch(), Some(2)); + } + + #[test] + fn merge_sort_exec_required_input_ordering_matches_spm() { + let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _; + let ordering = test_ordering(schema.as_ref()); + + let merge_sort = MergeSortExec::new(ordering.clone(), Arc::clone(&input), Some(1)); + let bare_spm = + SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)); + + assert_eq!( + merge_sort.required_input_ordering(), + vec![Some(OrderingRequirements::from(ordering))], + "MergeSortExec must require locally sorted input partitions for the merge key" + ); + assert_eq!( + merge_sort.required_input_ordering(), + bare_spm.required_input_ordering(), + "MergeSortExec's child ordering contract should mirror SortPreservingMergeExec" + ); + assert_eq!( + merge_sort.maintains_input_order(), + bare_spm.maintains_input_order() + ); + } + + #[test] + fn merge_sort_exec_with_fetch_pushes_fetch_to_child_sort() { + let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _; + let ordering = test_ordering(schema.as_ref()); + let child_sort = + Arc::new(SortExec::new(ordering.clone(), input).with_preserve_partitioning(true)) + as Arc; + let merge_sort = MergeSortExec::new(ordering, child_sort, None); + + let fetched = merge_sort.with_fetch(Some(2)).unwrap(); + + assert!(fetched.as_any().downcast_ref::().is_some()); + assert_eq!(fetched.fetch(), Some(2)); + let child_sort = fetched.children()[0] + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(child_sort.fetch(), Some(2)); + assert!(child_sort.preserve_partitioning()); + } + + #[test] + fn merge_sort_exec_with_preserve_order_matches_spm_but_keeps_wrapper() { + let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)])); + let input = Arc::new(PreserveOrderProbeExec::new(Arc::new( + EmptyExec::new(schema.clone()).with_partitions(2), + ))) as _; + let ordering = test_ordering(schema.as_ref()); + + let bare_spm = + SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)); + let preserved_spm = bare_spm.with_preserve_order(true).unwrap(); + assert!( + preserved_spm + .as_any() + .downcast_ref::() + .is_some(), + "bare SPM should rebuild as bare SPM" + ); + assert!( + preserved_spm.children()[0] + .as_any() + .downcast_ref::() + .unwrap() + .preserve_order + ); + + let merge_sort = MergeSortExec::new(ordering, input, Some(1)); + let preserved_merge_sort = merge_sort.with_preserve_order(true).unwrap(); + assert!( + preserved_merge_sort + .as_any() + .downcast_ref::() + .is_some(), + "MergeSortExec must rewrap the preserve-order child as MergeSortExec" + ); + assert!( + preserved_merge_sort + .as_any() + .downcast_ref::() + .is_none(), + "MergeSortExec must not expose a bare SPM after with_preserve_order" + ); + assert_eq!(preserved_merge_sort.fetch(), Some(1)); + assert_eq!( + preserved_merge_sort.required_input_ordering(), + preserved_spm.required_input_ordering(), + "preserve-order rewrite should keep the same SPM child-ordering contract" + ); + assert!( + preserved_merge_sort.children()[0] + .as_any() + .downcast_ref::() + .unwrap() + .preserve_order + ); + } + + #[test] + fn merge_sort_exec_projection_swap_matches_spm_but_keeps_wrapper() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("value", DataType::Int64, false), + Field::new("ts", DataType::Int64, false), + Field::new("tag", DataType::Utf8, false), + ])); + let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _; + let ordering = test_ordering(schema.as_ref()); + + let bare_spm = Arc::new( + SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)), + ) as Arc; + let spm_projection = ProjectionExec::try_new( + vec![ + (physical_col("ts", schema.as_ref())?, "ts".to_string()), + (physical_col("tag", schema.as_ref())?, "tag".to_string()), + ], + Arc::clone(&bare_spm), + )?; + let swapped_spm = bare_spm + .try_swapping_with_projection(&spm_projection)? + .expect("SPM should accept a narrowing projection that preserves the sort key"); + assert!( + swapped_spm + .as_any() + .downcast_ref::() + .is_some(), + "bare SPM should rebuild as bare SPM" + ); + + let merge_sort = + Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc; + let merge_projection = ProjectionExec::try_new( + vec![ + (physical_col("ts", schema.as_ref())?, "ts".to_string()), + (physical_col("tag", schema.as_ref())?, "tag".to_string()), + ], + Arc::clone(&merge_sort), + )?; + let swapped_merge_sort = merge_sort + .try_swapping_with_projection(&merge_projection)? + .expect("MergeSortExec should accept the same projection swap as SPM"); + + assert!( + swapped_merge_sort + .as_any() + .downcast_ref::() + .is_some(), + "MergeSortExec must rewrap projection swaps as MergeSortExec" + ); + assert!( + swapped_merge_sort + .as_any() + .downcast_ref::() + .is_none(), + "MergeSortExec must not expose a bare SPM after projection swap" + ); + assert_eq!(swapped_merge_sort.fetch(), Some(1)); + assert!( + swapped_merge_sort.children()[0] + .as_any() + .downcast_ref::() + .is_some(), + "the projection should move below MergeSortExec" + ); + let swapped_schema = swapped_merge_sort.schema(); + assert_eq!( + swapped_schema + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + vec!["ts", "tag"], + "swapped MergeSortExec should expose the projected schema" + ); + + let projected_ordering = LexOrdering::new([PhysicalSortExpr::new( + physical_col("ts", swapped_merge_sort.children()[0].schema().as_ref())?, + SortOptions { + descending: true, + nulls_first: false, + }, + )]) + .unwrap(); + assert_eq!( + swapped_merge_sort.required_input_ordering(), + vec![Some(OrderingRequirements::from(projected_ordering))], + "projection swap must rewrite the ordering to the child projection's schema" + ); + assert_eq!( + swapped_merge_sort.required_input_ordering(), + swapped_spm.required_input_ordering(), + "MergeSortExec projection swap should mirror SPM's ordering rewrite" + ); + + let spm_projection_without_sort_key = ProjectionExec::try_new( + vec![(physical_col("tag", schema.as_ref())?, "tag".to_string())], + Arc::clone(&bare_spm), + )?; + let merge_projection_without_sort_key = ProjectionExec::try_new( + vec![(physical_col("tag", schema.as_ref())?, "tag".to_string())], + Arc::clone(&merge_sort), + )?; + assert!( + bare_spm + .try_swapping_with_projection(&spm_projection_without_sort_key)? + .is_none(), + "SPM must reject projection swaps that drop the sort key" + ); + assert!( + merge_sort + .try_swapping_with_projection(&merge_projection_without_sort_key)? + .is_none(), + "MergeSortExec should reject the same projection swap as SPM" + ); + + Ok(()) + } + + #[test] + fn enforce_sorting_rewrite_keeps_merge_sort_exec_opaque() { + let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _; + let ordering = test_ordering(schema.as_ref()); + + let bare_spm = Arc::new( + SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)), + ) as Arc; + let optimized_spm = plan_with_order_breaking_variants(OrderPreservationContext::new( + bare_spm, + false, + vec![OrderPreservationContext::new( + Arc::clone(&input), + false, + vec![], + )], + )) + .unwrap() + .plan; + assert!( + optimized_spm + .as_any() + .downcast_ref::() + .is_some(), + "this regression test must exercise EnforceSorting's bare SPM -> CoalescePartitionsExec rewrite" + ); + + let merge_sort = + Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc; + let optimized_merge_sort = + plan_with_order_breaking_variants(OrderPreservationContext::new( + Arc::clone(&merge_sort), + false, + vec![OrderPreservationContext::new( + Arc::clone(merge_sort.children()[0]), + false, + vec![], + )], + )) + .unwrap() + .plan; + assert!( + optimized_merge_sort + .as_any() + .downcast_ref::() + .is_some(), + "MergeSortExec must stay opaque to the bare SPM rewrite" + ); + assert!( + optimized_merge_sort + .as_any() + .downcast_ref::() + .is_none(), + "MergeSortExec(fetch) is the required distributed TopK merge stage, not an unordered coalesce" + ); + assert_eq!(optimized_merge_sort.fetch(), Some(1)); + } +} diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index f613807b291..f3b7851fde3 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use ahash::HashMap; +use arrow_schema::SortOptions; use async_trait::async_trait; use catalog::CatalogManagerRef; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; @@ -29,6 +30,7 @@ use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner}; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}; use datafusion_common::{DataFusionError, TableReference}; use datafusion_expr::{LogicalPlan, UserDefinedLogicalNode}; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; use partition::manager::{PartitionRuleManagerRef, create_partitions_from_region_routes}; use session::context::QueryContext; use snafu::{OptionExt, ResultExt}; @@ -39,27 +41,54 @@ use table::table_name::TableName; use crate::dist_plan::PredicateExtractor; use crate::dist_plan::merge_scan::{MergeScanExec, MergeScanLogicalPlan}; -use crate::dist_plan::merge_sort::MergeSortLogicalPlan; +use crate::dist_plan::merge_sort::{MergeSortExec, MergeSortLogicalPlan}; use crate::dist_plan::region_pruner::ConstraintPruner; use crate::error::{CatalogSnafu, PartitionRuleManagerSnafu, TableNotFoundSnafu}; use crate::region_query::RegionQueryHandlerRef; -/// Planner for convert merge sort logical plan to physical plan +/// Planner for converting merge sort logical plan to physical plan. /// -/// it is currently a fallback to sort, and doesn't change the execution plan: -/// `MergeSort(MergeScan) -> Sort(MergeScan) - to physical plan -> ...` -/// It should be applied after `DistExtensionPlanner` -/// -/// (Later when actually impl this merge sort) -/// -/// We should ensure the number of partition is not smaller than the number of region at present. Otherwise this would result in incorrect output. +/// `MergeSortExec` always represents the distributed merge stage. It declares +/// the required input ordering to DataFusion, so `EnforceSorting` inserts a +/// `SortExec` below it when the input `MergeScanExec` cannot preserve per-region +/// ordering, for example when one output partition may merge multiple region +/// streams. pub struct MergeSortExtensionPlanner {} +impl MergeSortExtensionPlanner { + fn ordering( + session_state: &SessionState, + merge_sort: &MergeSortLogicalPlan, + ) -> Result { + let ordering = merge_sort + .expr + .iter() + .map(|sort_expr| { + let physical_expr = session_state + .create_physical_expr(sort_expr.expr.clone(), merge_sort.input.schema())?; + Ok(PhysicalSortExpr::new( + physical_expr, + SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }, + )) + }) + .collect::>>()?; + + LexOrdering::new(ordering).ok_or_else(|| { + DataFusionError::Internal( + "Expect MergeSort to have non-empty sort expressions".to_string(), + ) + }) + } +} + #[async_trait] impl ExtensionPlanner for MergeSortExtensionPlanner { async fn plan_extension( &self, - planner: &dyn PhysicalPlanner, + _planner: &dyn PhysicalPlanner, node: &dyn UserDefinedLogicalNode, _logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc], @@ -73,28 +102,24 @@ impl ExtensionPlanner for MergeSortExtensionPlanner { .downcast_ref::() .is_some() { - let merge_scan_exec = physical_inputs - .first() - .and_then(|p| p.as_any().downcast_ref::()) - .ok_or(DataFusionError::Internal(format!( + let input = physical_inputs.first().cloned().ok_or_else(|| { + DataFusionError::Internal( + "Expect MergeSort to have one physical input".to_string(), + ) + })?; + if input.as_any().downcast_ref::().is_none() { + return Err(DataFusionError::Internal(format!( "Expect MergeSort's input is a MergeScanExec, found {:?}", physical_inputs - )))?; - - let partition_cnt = merge_scan_exec.partition_count(); - let region_cnt = merge_scan_exec.region_count(); - // if partition >= region, we know that every partition stream of merge scan is ordered - // and we only need to do a merge sort, otherwise fallback to quick sort - let can_merge_sort = partition_cnt >= region_cnt; - if can_merge_sort { - // TODO(discord9): use `SortPreservingMergeExec here` + ))); } - // for now merge sort only exist in logical plan, and have the same effect as `Sort` - // doesn't change the execution plan, this will change in the future - let ret = planner - .create_physical_plan(&merge_sort.clone().into_sort(), session_state) - .await?; - Ok(Some(ret)) + + let ordering = Self::ordering(session_state, merge_sort)?; + Ok(Some(Arc::new(MergeSortExec::new( + ordering, + input, + merge_sort.fetch, + )))) } else { Ok(None) } diff --git a/src/query/src/optimizer/global_limit.rs b/src/query/src/optimizer/global_limit.rs index 2e5350f7797..4b2ff721f64 100644 --- a/src/query/src/optimizer/global_limit.rs +++ b/src/query/src/optimizer/global_limit.rs @@ -25,6 +25,8 @@ use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use datafusion_common::Result as DfResult; use datafusion_physical_expr::{Distribution, OrderingRequirements, Partitioning}; +use crate::dist_plan::MergeSortExec; + #[derive(Debug)] pub struct EnsureGlobalLimitForFetch; @@ -149,7 +151,8 @@ fn provided_global_fetch(plan: &Arc) -> Option { let fetch = plan.fetch()?; (plan.as_any().is::() || plan.as_any().is::() - || plan.as_any().is::()) + || plan.as_any().is::() + || plan.as_any().is::()) .then_some(fetch) } @@ -327,6 +330,36 @@ mod tests { assert_eq!(child.fetch(), Some(5)); } + #[test] + fn keeps_filter_under_parent_merge_sort_fetch() { + let (input, ordering) = ordered_input(); + let filter = filter_fetch(input, 1); + let merge = merge_sort_fetch(ordering, filter, 1); + + let optimized = + EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap(); + let child = optimized.children()[0]; + + assert!(optimized.as_any().is::()); + assert!(child.as_any().is::()); + } + + #[test] + fn adds_tighter_global_fetch_under_looser_merge_sort_fetch() { + let (input, ordering) = ordered_input(); + let filter = filter_fetch(input, 5); + let merge = merge_sort_fetch(ordering, filter, 10); + + let optimized = + EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap(); + let child = optimized.children()[0]; + + assert!(optimized.as_any().is::()); + assert!(child.as_any().is::()); + assert_eq!(child.fetch(), Some(5)); + assert!(child.children()[0].as_any().is::()); + } + #[test] fn preserves_parent_ordering_requirement() { let (input, ordering) = ordered_input(); @@ -552,6 +585,14 @@ mod tests { ) } + fn merge_sort_fetch( + ordering: LexOrdering, + input: Arc, + fetch: usize, + ) -> Arc { + Arc::new(MergeSortExec::new(ordering, input, Some(fetch))) + } + fn hash_repartition(input: Arc) -> Arc { let partitioning = Partitioning::Hash(vec![col("a", input.schema().as_ref()).unwrap()], 3); Arc::new(RepartitionExec::try_new(input, partitioning).unwrap()) diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 8572e9ce5b2..90c3708b534 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -28,6 +28,7 @@ use datatypes::arrow::array::{ use frontend::error::Error; use frontend::instance::Instance; use operator::error::Error as OperatorError; +use query::datafusion::QUERY_PARALLELISM_HINT; use rstest::rstest; use rstest_reuse::apply; use servers::error as server_error; @@ -89,6 +90,98 @@ async fn test_create_database_and_insert_query(instance: Arc) } } +#[tokio::test(flavor = "multi_thread")] +async fn test_distributed_scalar_latest_with_query_parallelism_below_regions() { + common_telemetry::init_default_ut_logging(); + + let distributed = crate::tests::create_distributed_instance( + "test_distributed_scalar_latest_with_query_parallelism_below_regions", + ) + .await; + let frontend = distributed.frontend(); + + execute_sql( + &frontend, + r#" +CREATE TABLE cpu ( + rack STRING NULL, + os STRING NULL, + usage_user BIGINT NULL, + greptime_timestamp TIMESTAMP(9) NOT NULL, + TIME INDEX (greptime_timestamp) +) +PARTITION ON COLUMNS (rack) ( + rack < '2', + rack >= '2' AND rack < '4', + rack >= '4' AND rack < '6', + rack >= '6' AND rack < '8', + rack >= '8' +) +ENGINE = mito +WITH (append_mode = 'true', sst_format = 'flat') +"#, + ) + .await; + + execute_sql( + &frontend, + r#" +INSERT INTO cpu VALUES + ('1', 'linux', 10, '2023-06-12 01:04:49'), + ('1', 'linux', 15, '2023-06-12 01:04:50'), + ('3', 'windows', 25, '2023-06-12 01:05:00'), + ('5', 'mac', 30, '2023-06-12 01:03:00'), + ('7', 'linux', 45, '2023-06-12 02:00:00'), + ('2', 'linux', 20, '2023-06-12 01:04:51'), + ('2', 'windows', 22, '2023-06-12 01:06:00'), + ('4', 'mac', 12, '2023-06-12 00:59:00'), + ('6', 'linux', 35, '2023-06-12 01:04:55'), + ('8', 'windows', 50, '2023-06-12 02:10:00') +"#, + ) + .await; + + let latest_sql = r#" +SELECT rack, os, greptime_timestamp +FROM cpu +WHERE greptime_timestamp = ( + SELECT greptime_timestamp + FROM cpu + ORDER BY greptime_timestamp DESC + LIMIT 1 +) +"#; + + let result = execute_sql_with_query_parallelism(&frontend, latest_sql, 1) + .await + .data + .pretty_print() + .await; + assert_eq!( + result, + r#"+------+---------+---------------------+ +| rack | os | greptime_timestamp | ++------+---------+---------------------+ +| 8 | windows | 2023-06-12T02:10:00 | ++------+---------+---------------------+"# + ); + + let explain = + execute_sql_with_query_parallelism(&frontend, &format!("EXPLAIN {latest_sql}"), 1) + .await + .data + .pretty_print() + .await; + assert!( + explain.contains("SortExec"), + "query_parallelism=1 with five regions should insert SortExec below MergeSortExec; explain:\n{explain}" + ); + assert!( + explain.contains("MergeSortExec"), + "query_parallelism=1 with five regions should keep the distributed merge stage stable; explain:\n{explain}" + ); +} + #[apply(both_instances_cases)] async fn test_show_create_table(instance: Arc) { let frontend = instance.frontend(); @@ -2876,6 +2969,16 @@ async fn execute_sql(instance: &Arc, sql: &str) -> Output { execute_sql_with(instance, sql, QueryContext::arc()).await } +async fn execute_sql_with_query_parallelism( + instance: &Arc, + sql: &str, + parallelism: usize, +) -> Output { + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, parallelism.to_string()); + execute_sql_with(instance, sql, Arc::new(query_ctx)).await +} + async fn try_execute_sql(instance: &Arc, sql: &str) -> server_error::Result { try_execute_sql_with(instance, sql, QueryContext::arc()).await } diff --git a/tests/cases/distributed/explain/multi_partitions.result b/tests/cases/distributed/explain/multi_partitions.result index 0c3d4fa4b70..47d807542d8 100644 --- a/tests/cases/distributed/explain/multi_partitions.result +++ b/tests/cases/distributed/explain/multi_partitions.result @@ -30,7 +30,7 @@ explain SELECT * FROM multi_partitions_test_table WHERE ts > cast(1000000000 as |_|_Filter: multi_partitions_test_table.ts > TimestampMillisecond(1000000000, None)_| |_|_TableScan: multi_partitions_test_table, partial_filters=[multi_partitions_test_table.ts > TimestampMillisecond(1000000000, None)]_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/distributed/explain/step_aggr_advance.result b/tests/cases/distributed/explain/step_aggr_advance.result index c2304df5eb9..7ea86dbccbf 100644 --- a/tests/cases/distributed/explain/step_aggr_advance.result +++ b/tests/cases/distributed/explain/step_aggr_advance.result @@ -37,7 +37,7 @@ tql explain (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr | | Filter: aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None) AND aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None) | | | TableScan: aggr_optimize_not, partial_filters=[aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None), aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)] | | | ]] | -| physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | +| physical_plan | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | | | MergeScanExec: REDACTED | | | +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -54,7 +54,7 @@ tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED @@ -244,7 +244,7 @@ tql explain (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time | | Filter: aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None) AND aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None) | | | TableScan: aggr_optimize_not, partial_filters=[aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None), aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)] | | | ]] | -| physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] | +| physical_plan | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] | | | MergeScanExec: REDACTED | | | +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -261,7 +261,7 @@ tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED @@ -485,9 +485,10 @@ tql explain (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize | | PromSeriesDivideExec: tags=["a", "b", "c", "d"] | | | SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true] | | | MergeScanExec: REDACTED -| | SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] | -| | CooperativeExec | -| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=REDACTED +| | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | +| | SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] | +| | MergeScanExec: REDACTED | | | +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -514,8 +515,9 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED |_|_|_SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true] REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED @@ -974,7 +976,7 @@ EXPLAIN SELECT pk_col_1, pk_col_2, sum(val_col_1) FROM step_aggr_extended GROUP |_|_Aggregate: groupBy=[[step_aggr_extended.pk_col_1, step_aggr_extended.pk_col_2]], aggr=[[sum(step_aggr_extended.val_col_1)]] | |_|_TableScan: step_aggr_extended_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [pk_col_1@0 ASC NULLS LAST, pk_col_2@1 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [pk_col_1@0 ASC NULLS LAST, pk_col_2@1 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/distributed/explain/subqueries.result b/tests/cases/distributed/explain/subqueries.result index 49bf5d97566..e6f7286535e 100644 --- a/tests/cases/distributed/explain/subqueries.result +++ b/tests/cases/distributed/explain/subqueries.result @@ -287,7 +287,7 @@ EXPLAIN SELECT x FROM (SELECT a AS x FROM t) sq ORDER BY x; |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -326,7 +326,7 @@ EXPLAIN SELECT x, COUNT(*) AS c FROM (SELECT a AS x FROM t) sq GROUP BY x ORDER |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -397,7 +397,7 @@ EXPLAIN SELECT sq.x FROM (SELECT a AS x FROM t) sq ORDER BY sq.x; |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -433,7 +433,7 @@ EXPLAIN SELECT y FROM (SELECT x AS y FROM (SELECT a AS x FROM t) sq1) sq2 ORDER |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [y@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [y@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -471,7 +471,7 @@ EXPLAIN SELECT x, x + 1 AS y FROM (SELECT a AS x FROM t) sq ORDER BY x; |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -625,7 +625,7 @@ EXPLAIN SELECT x FROM (SELECT a AS x FROM t) sq ORDER BY x LIMIT 2; |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST], fetch=2_| +| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST], fetch=2_| |_|_SortExec: TopK(fetch=2), expr=[x@0 ASC NULLS LAST], preserve_partitioning=[true]_| |_|_MergeScanExec: REDACTED |_|_| diff --git a/tests/cases/distributed/optimizer/first_value_advance.result b/tests/cases/distributed/optimizer/first_value_advance.result index ac977b6fdfb..08e1b8a4112 100644 --- a/tests/cases/distributed/optimizer/first_value_advance.result +++ b/tests/cases/distributed/optimizer/first_value_advance.result @@ -240,7 +240,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t.host]], aggr=[[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] | |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -267,7 +267,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 as ordered_host, first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST]@1 as first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]@2 as first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -606,7 +606,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t1.host]], aggr=[[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] | |_|_TableScan: t1_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -632,7 +632,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 as ordered_host, first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]@1 as first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/distributed/optimizer/last_value_advance.result b/tests/cases/distributed/optimizer/last_value_advance.result index d01ea184300..7b131d44f37 100644 --- a/tests/cases/distributed/optimizer/last_value_advance.result +++ b/tests/cases/distributed/optimizer/last_value_advance.result @@ -240,7 +240,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t.host]], aggr=[[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] | |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -267,7 +267,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 as ordered_host, last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST]@1 as last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]@2 as last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -606,7 +606,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t1.host]], aggr=[[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] | |_|_TableScan: t1_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -632,7 +632,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 as ordered_host, last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]@1 as last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/distributed/optimizer/time_index_filter_pushdown.result b/tests/cases/distributed/optimizer/time_index_filter_pushdown.result index 22e96793747..016c49deea0 100644 --- a/tests/cases/distributed/optimizer/time_index_filter_pushdown.result +++ b/tests/cases/distributed/optimizer/time_index_filter_pushdown.result @@ -84,6 +84,82 @@ group by | windows | 2 | +---------+----------+ +-- Query the row at the latest timestamp through a scalar subquery. +-- Keep the distributed merge sort visible, but redact the optional local SortExec +-- that appears when query parallelism is lower than the region count. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (?m)^\|_\|_SortExec:.*\n +EXPLAIN SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Projection: cpu.rack, cpu.os, cpu.greptime_timestamp_| +|_|_Inner Join: cpu.greptime_timestamp = __scalar_sq_1.greptime_timestamp_| +|_|_Projection: cpu.rack, cpu.os, cpu.greptime_timestamp_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| TableScan: cpu_| +|_| ]]_| +|_|_SubqueryAlias: __scalar_sq_1_| +|_|_Limit: skip=0, fetch=1_| +|_|_MergeSort: cpu.greptime_timestamp DESC NULLS FIRST_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Limit: skip=0, fetch=1_| +|_|_Sort: cpu.greptime_timestamp DESC NULLS FIRST_| +|_|_Projection: cpu.greptime_timestamp_| +|_|_TableScan: cpu_| +|_| ]]_| +| physical_plan | HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(greptime_timestamp@0, greptime_timestamp@2)], projection=[rack@1, os@2, greptime_timestamp@3]_| +|_|_MergeSortExec: [greptime_timestamp@0 DESC], fetch=1_| +|_|_MergeScanExec: REDACTED +|_|_ProjectionExec: expr=[rack@0 as rack, os@1 as os, greptime_timestamp@3 as greptime_timestamp]_| +|_|_CooperativeExec_| +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + ++------+---------+---------------------+ +| rack | os | greptime_timestamp | ++------+---------+---------------------+ +| 8 | windows | 2023-06-12T02:10:00 | ++------+---------+---------------------+ + drop table cpu; Affected Rows: 0 diff --git a/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql b/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql index f5fa24efe77..c9c80c24a61 100644 --- a/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql +++ b/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql @@ -59,4 +59,47 @@ where group by os; -drop table cpu; \ No newline at end of file +-- Query the row at the latest timestamp through a scalar subquery. +-- Keep the distributed merge sort visible, but redact the optional local SortExec +-- that appears when query parallelism is lower than the region count. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (?m)^\|_\|_SortExec:.*\n +EXPLAIN SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + +SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + +drop table cpu; diff --git a/tests/cases/standalone/common/order/order_by.result b/tests/cases/standalone/common/order/order_by.result index ceebb94ff5d..aed1a491c27 100644 --- a/tests/cases/standalone/common/order/order_by.result +++ b/tests/cases/standalone/common/order/order_by.result @@ -296,7 +296,7 @@ explain analyze select tag from t where num > 6 order by ts desc limit 2; | stage | node | plan_| +-+-+-+ | 0_| 0_|_ProjectionExec: expr=[tag@0 as tag] REDACTED -|_|_|_SortPreservingMergeExec: [ts@1 DESC], fetch=2 REDACTED +|_|_|_MergeSortExec: [ts@1 DESC], fetch=2 REDACTED |_|_|_SortExec: TopK(fetch=2), expr=[ts@1 DESC], preserve_partitioning=[true] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| diff --git a/tests/cases/standalone/common/tql/partition.result b/tests/cases/standalone/common/tql/partition.result index 5d0fd981238..45927f8abbd 100644 --- a/tests/cases/standalone/common/tql/partition.result +++ b/tests/cases/standalone/common/tql/partition.result @@ -133,8 +133,8 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED -|_|_|_SortExec: expr=[k@2 ASC, l@3 ASC, j@1 ASC], preserve_partitioning=[true] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeSortExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED diff --git a/tests/cases/standalone/common/window/latest_per_series.result b/tests/cases/standalone/common/window/latest_per_series.result new file mode 100644 index 00000000000..c389c8aeb8a --- /dev/null +++ b/tests/cases/standalone/common/window/latest_per_series.result @@ -0,0 +1,187 @@ +CREATE TABLE netflow_raw ( + greptime_timestamp TIMESTAMP TIME INDEX, + src STRING, + dst_port INT, + dst_addr STRING, + byte_count DOUBLE, + flow_id INT, + PRIMARY KEY (src, dst_port) +); + +Affected Rows: 0 + +INSERT INTO netflow_raw VALUES + (1000, 'source-a', 80, 'target-80-old', 100.0, 1), + (5000, 'source-a', 80, 'target-80-new', 500.0, 2), + (2000, 'source-a', 443, 'target-443-old', 200.0, 3), + (7000, 'source-a', 443, 'target-443-new', 700.0, 4), + (3000, 'source-a', 22, 'target-22-only', 300.0, 5), + (4000, 'source-a', 8080, 'target-8080-only', 400.0, 6), + (8000, 'source-b', 9999, 'target-other-src', 999.0, 7); + +Affected Rows: 7 + +-- Latest row per dst_port using a window subquery. The window result must be +-- filtered outside the subquery, not directly in the WHERE clause. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + ROW_NUMBER() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS rn + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE rn = 1 +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + ++---------------------+----------+------------------+------------+---------+ +| greptime_timestamp | dst_port | dst_addr | byte_count | flow_id | ++---------------------+----------+------------------+------------+---------+ +| 1970-01-01T00:00:07 | 443 | target-443-new | 700.0 | 4 | +| 1970-01-01T00:00:05 | 80 | target-80-new | 500.0 | 2 | +| 1970-01-01T00:00:04 | 8080 | target-8080-only | 400.0 | 6 | ++---------------------+----------+------------------+------------+---------+ + +-- DISTINCT ON keeps the first row per dst_port according to its own ORDER BY. +-- Wrap it to apply the final "latest groups first" ordering and LIMIT. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id +FROM ( + SELECT DISTINCT ON (dst_port) + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp + ORDER BY dst_port, greptime_timestamp DESC +) latest_per_port +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + ++---------------------+----------+------------------+------------+---------+ +| greptime_timestamp | dst_port | dst_addr | byte_count | flow_id | ++---------------------+----------+------------------+------------+---------+ +| 1970-01-01T00:00:07 | 443 | target-443-new | 700.0 | 4 | +| 1970-01-01T00:00:05 | 80 | target-80-new | 500.0 | 2 | +| 1970-01-01T00:00:04 | 8080 | target-8080-only | 400.0 | 6 | ++---------------------+----------+------------------+------------+---------+ + +-- RANK returns all rows tied at the latest timestamp for each dst_port. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + RANK() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS r + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE r = 1 +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + ++---------------------+----------+------------------+------------+---------+---+ +| greptime_timestamp | dst_port | dst_addr | byte_count | flow_id | r | ++---------------------+----------+------------------+------------+---------+---+ +| 1970-01-01T00:00:07 | 443 | target-443-new | 700.0 | 4 | 1 | +| 1970-01-01T00:00:05 | 80 | target-80-new | 500.0 | 2 | 1 | +| 1970-01-01T00:00:04 | 8080 | target-8080-only | 400.0 | 6 | 1 | ++---------------------+----------+------------------+------------+---------+---+ + +-- Ordered aggregate variant when the caller can enumerate the columns to fetch. +SELECT + dst_port, + last_value(greptime_timestamp ORDER BY greptime_timestamp) AS latest_ts, + last_value(dst_addr ORDER BY greptime_timestamp) AS latest_dst_addr, + last_value(byte_count ORDER BY greptime_timestamp) AS latest_byte_count, + last_value(flow_id ORDER BY greptime_timestamp) AS latest_flow_id +FROM netflow_raw +WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +GROUP BY dst_port +ORDER BY latest_ts DESC, dst_port +LIMIT 3; + ++----------+---------------------+------------------+-------------------+----------------+ +| dst_port | latest_ts | latest_dst_addr | latest_byte_count | latest_flow_id | ++----------+---------------------+------------------+-------------------+----------------+ +| 443 | 1970-01-01T00:00:07 | target-443-new | 700.0 | 4 | +| 80 | 1970-01-01T00:00:05 | target-80-new | 500.0 | 2 | +| 8080 | 1970-01-01T00:00:04 | target-8080-only | 400.0 | 6 | ++----------+---------------------+------------------+-------------------+----------------+ + +CREATE TABLE netflow_raw_ties ( + greptime_timestamp TIMESTAMP TIME INDEX, + src STRING, + dst_port INT, + dst_addr STRING, + byte_count DOUBLE, + flow_id INT, + PRIMARY KEY (src, dst_port, flow_id) +); + +Affected Rows: 0 + +INSERT INTO netflow_raw_ties VALUES + (1000, 'source-a', 443, 'target-443-old', 100.0, 1), + (9000, 'source-a', 443, 'target-443-tie-a', 900.0, 2), + (9000, 'source-a', 443, 'target-443-tie-b', 901.0, 3), + (8000, 'source-a', 80, 'target-80-new', 800.0, 4), + (9500, 'source-b', 443, 'target-other-src', 950.0, 5); + +Affected Rows: 5 + +-- RANK with r = 1 keeps all rows tied at the latest timestamp per dst_port. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + RANK() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS r + FROM netflow_raw_ties + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE r = 1 +ORDER BY greptime_timestamp DESC, dst_port, flow_id; + ++---------------------+----------+------------------+------------+---------+---+ +| greptime_timestamp | dst_port | dst_addr | byte_count | flow_id | r | ++---------------------+----------+------------------+------------+---------+---+ +| 1970-01-01T00:00:09 | 443 | target-443-tie-a | 900.0 | 2 | 1 | +| 1970-01-01T00:00:09 | 443 | target-443-tie-b | 901.0 | 3 | 1 | +| 1970-01-01T00:00:08 | 80 | target-80-new | 800.0 | 4 | 1 | ++---------------------+----------+------------------+------------+---------+---+ + +DROP TABLE netflow_raw_ties; + +Affected Rows: 0 + +DROP TABLE netflow_raw; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/window/latest_per_series.sql b/tests/cases/standalone/common/window/latest_per_series.sql new file mode 100644 index 00000000000..f809c8777ee --- /dev/null +++ b/tests/cases/standalone/common/window/latest_per_series.sql @@ -0,0 +1,134 @@ +CREATE TABLE netflow_raw ( + greptime_timestamp TIMESTAMP TIME INDEX, + src STRING, + dst_port INT, + dst_addr STRING, + byte_count DOUBLE, + flow_id INT, + PRIMARY KEY (src, dst_port) +); + +INSERT INTO netflow_raw VALUES + (1000, 'source-a', 80, 'target-80-old', 100.0, 1), + (5000, 'source-a', 80, 'target-80-new', 500.0, 2), + (2000, 'source-a', 443, 'target-443-old', 200.0, 3), + (7000, 'source-a', 443, 'target-443-new', 700.0, 4), + (3000, 'source-a', 22, 'target-22-only', 300.0, 5), + (4000, 'source-a', 8080, 'target-8080-only', 400.0, 6), + (8000, 'source-b', 9999, 'target-other-src', 999.0, 7); + +-- Latest row per dst_port using a window subquery. The window result must be +-- filtered outside the subquery, not directly in the WHERE clause. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + ROW_NUMBER() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS rn + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE rn = 1 +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + +-- DISTINCT ON keeps the first row per dst_port according to its own ORDER BY. +-- Wrap it to apply the final "latest groups first" ordering and LIMIT. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id +FROM ( + SELECT DISTINCT ON (dst_port) + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp + ORDER BY dst_port, greptime_timestamp DESC +) latest_per_port +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + +-- RANK returns all rows tied at the latest timestamp for each dst_port. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + RANK() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS r + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE r = 1 +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + +-- Ordered aggregate variant when the caller can enumerate the columns to fetch. +SELECT + dst_port, + last_value(greptime_timestamp ORDER BY greptime_timestamp) AS latest_ts, + last_value(dst_addr ORDER BY greptime_timestamp) AS latest_dst_addr, + last_value(byte_count ORDER BY greptime_timestamp) AS latest_byte_count, + last_value(flow_id ORDER BY greptime_timestamp) AS latest_flow_id +FROM netflow_raw +WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +GROUP BY dst_port +ORDER BY latest_ts DESC, dst_port +LIMIT 3; + +CREATE TABLE netflow_raw_ties ( + greptime_timestamp TIMESTAMP TIME INDEX, + src STRING, + dst_port INT, + dst_addr STRING, + byte_count DOUBLE, + flow_id INT, + PRIMARY KEY (src, dst_port, flow_id) +); + +INSERT INTO netflow_raw_ties VALUES + (1000, 'source-a', 443, 'target-443-old', 100.0, 1), + (9000, 'source-a', 443, 'target-443-tie-a', 900.0, 2), + (9000, 'source-a', 443, 'target-443-tie-b', 901.0, 3), + (8000, 'source-a', 80, 'target-80-new', 800.0, 4), + (9500, 'source-b', 443, 'target-other-src', 950.0, 5); + +-- RANK with r = 1 keeps all rows tied at the latest timestamp per dst_port. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + RANK() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS r + FROM netflow_raw_ties + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE r = 1 +ORDER BY greptime_timestamp DESC, dst_port, flow_id; + +DROP TABLE netflow_raw_ties; + +DROP TABLE netflow_raw; diff --git a/tests/cases/standalone/optimizer/first_value_advance.result b/tests/cases/standalone/optimizer/first_value_advance.result index ffc0b341f0a..afcefae4cff 100644 --- a/tests/cases/standalone/optimizer/first_value_advance.result +++ b/tests/cases/standalone/optimizer/first_value_advance.result @@ -240,7 +240,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t.host]], aggr=[[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] | |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -267,7 +267,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED @@ -603,7 +603,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t1.host]], aggr=[[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] | |_|_TableScan: t1_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -629,7 +629,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED diff --git a/tests/cases/standalone/optimizer/last_value_advance.result b/tests/cases/standalone/optimizer/last_value_advance.result index 0af561e6185..a1882a75887 100644 --- a/tests/cases/standalone/optimizer/last_value_advance.result +++ b/tests/cases/standalone/optimizer/last_value_advance.result @@ -240,7 +240,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t.host]], aggr=[[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] | |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -267,7 +267,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED @@ -603,7 +603,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t1.host]], aggr=[[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] | |_|_TableScan: t1_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -629,7 +629,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED diff --git a/tests/cases/standalone/optimizer/time_index_filter_pushdown.result b/tests/cases/standalone/optimizer/time_index_filter_pushdown.result index de6222ae28e..fd0d869a33f 100644 --- a/tests/cases/standalone/optimizer/time_index_filter_pushdown.result +++ b/tests/cases/standalone/optimizer/time_index_filter_pushdown.result @@ -197,6 +197,31 @@ group by | windows | 2 | +---------+----------+ +-- Query the row at the latest timestamp through a scalar subquery. +SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + ++------+---------+---------------------+ +| rack | os | greptime_timestamp | ++------+---------+---------------------+ +| 8 | windows | 2023-06-12T02:10:00 | ++------+---------+---------------------+ + drop table cpu; Affected Rows: 0 diff --git a/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql b/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql index c0636e26a4b..7a7030d0925 100644 --- a/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql +++ b/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql @@ -120,4 +120,23 @@ where group by os; +-- Query the row at the latest timestamp through a scalar subquery. +SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + drop table cpu; From 589ae56a6fd33434140ef12e5175c27763cd39cd Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Tue, 14 Jul 2026 22:16:08 +0800 Subject: [PATCH 28/48] fix: close database ACL gaps in permission checks (#8492) * fix: database ACL Signed-off-by: shuiyisong * chore: reduce code Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/auth/src/permission.rs | 49 ++- src/frontend/src/instance.rs | 15 +- src/frontend/src/instance/grpc.rs | 15 +- src/operator/src/statement.rs | 8 +- src/operator/src/statement/set.rs | 30 -- src/sql/src/statements/set_variables.rs | 45 +- src/sql/src/statements/statement.rs | 3 +- src/sql/src/util.rs | 520 ++++++++++++++++-------- 8 files changed, 473 insertions(+), 212 deletions(-) diff --git a/src/auth/src/permission.rs b/src/auth/src/permission.rs index 8914635290f..cff9b414777 100644 --- a/src/auth/src/permission.rs +++ b/src/auth/src/permission.rs @@ -36,7 +36,11 @@ pub enum PermissionReq<'a> { PromStoreRead, Otlp, LogWrite, - BulkInsert, + BulkInsert { + catalog: &'a str, + schema: &'a str, + table: &'a str, + }, } impl<'a> PermissionReq<'a> { @@ -57,7 +61,7 @@ impl<'a> PermissionReq<'a> { | PermissionReq::PromStoreWrite | PermissionReq::Otlp | PermissionReq::LogWrite - | PermissionReq::BulkInsert => false, + | PermissionReq::BulkInsert { .. } => false, } } @@ -79,6 +83,15 @@ pub trait PermissionChecker: Send + Sync { user_info: UserInfoRef, req: PermissionReq, ) -> Result; + + fn check_permission_with_context( + &self, + user_info: UserInfoRef, + req: PermissionReq, + _current_schema: Option<&str>, + ) -> Result { + self.check_permission(user_info, req) + } } impl PermissionChecker for Option<&PermissionCheckerRef> { @@ -86,13 +99,24 @@ impl PermissionChecker for Option<&PermissionCheckerRef> { &self, user_info: UserInfoRef, req: PermissionReq, + ) -> Result { + self.check_permission_with_context(user_info, req, None) + } + + fn check_permission_with_context( + &self, + user_info: UserInfoRef, + req: PermissionReq, + current_schema: Option<&str>, ) -> Result { match self { - Some(checker) => match checker.check_permission(user_info, req) { - Ok(PermissionResp::Reject) => PermissionDeniedSnafu.fail(), - Ok(PermissionResp::Allow) => Ok(PermissionResp::Allow), - Err(e) => Err(e), - }, + Some(checker) => { + match checker.check_permission_with_context(user_info, req, current_schema) { + Ok(PermissionResp::Reject) => PermissionDeniedSnafu.fail(), + Ok(PermissionResp::Allow) => Ok(PermissionResp::Allow), + Err(e) => Err(e), + } + } None => Ok(PermissionResp::Allow), } } @@ -209,4 +233,15 @@ mod tests { assert!(req.is_write()); } + + #[test] + fn test_bulk_insert_is_write_request() { + let req = PermissionReq::BulkInsert { + catalog: "greptime", + schema: "public", + table: "metrics", + }; + + assert!(req.is_write()); + } } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 7242d601c1b..5beb8eb827a 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -611,7 +611,11 @@ impl Instance { let checker_ref = self.plugins.get::(); checker_ref .as_ref() - .check_permission(query_ctx.current_user(), PermissionReq::SqlStatement(&stmt)) + .check_permission_with_context( + query_ctx.current_user(), + PermissionReq::SqlStatement(&stmt), + Some(&query_ctx.current_schema()), + ) .context(PermissionSnafu)?; check_permission(self.plugins.clone(), &stmt, &query_ctx)?; let catalog_name = query_ctx.current_catalog().to_string(); @@ -672,9 +676,10 @@ impl Instance { let mut results = Vec::with_capacity(stmts.len()); for stmt in stmts { if let Err(e) = checker - .check_permission( + .check_permission_with_context( query_ctx.current_user(), PermissionReq::SqlStatement(&stmt), + Some(&query_ctx.current_schema()), ) .context(PermissionSnafu) { @@ -848,7 +853,11 @@ impl Instance { self.plugins .get::() .as_ref() - .check_permission(query_ctx.current_user(), PermissionReq::SqlStatement(&stmt)) + .check_permission_with_context( + query_ctx.current_user(), + PermissionReq::SqlStatement(&stmt), + Some(&query_ctx.current_schema()), + ) .context(PermissionSnafu)?; let plan = self diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index 0ca1a2cf209..808a2bd6ad8 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -71,7 +71,11 @@ impl GrpcQueryHandler for Instance { self.plugins .get::() .as_ref() - .check_permission(ctx.current_user(), PermissionReq::GrpcRequest(&request)) + .check_permission_with_context( + ctx.current_user(), + PermissionReq::GrpcRequest(&request), + Some(&ctx.current_schema()), + ) .context(PermissionSnafu)?; let output = match request { @@ -356,7 +360,14 @@ impl Instance { plugins .get::() .as_ref() - .check_permission(ctx.current_user(), PermissionReq::BulkInsert) + .check_permission( + ctx.current_user(), + PermissionReq::BulkInsert { + catalog: &table_name.catalog_name, + schema: &table_name.schema_name, + table: &table_name.table_name, + }, + ) .context(PermissionSnafu)?; // Resolve table reference diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 34b2550acf2..5ec7b4bc3cc 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -78,8 +78,7 @@ use table::table_name::TableName; use table::table_reference::TableReference; use self::set::{ - set_bytea_output, set_datestyle, set_intervalstyle, set_search_path, set_timezone, - validate_client_encoding, + set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding, }; use crate::error::{ self, CatalogSnafu, ExecLogicalPlanSnafu, ExternalSnafu, InvalidSqlSnafu, NotSupportedSnafu, @@ -528,7 +527,10 @@ impl StatementExecutor { }, "SEARCH_PATH" => { if query_ctx.channel() == Channel::Postgres { - set_search_path(set_var.value, query_ctx)? + let search_path = set_var.search_path().context(NotSupportedSnafu { + feat: "Unsupported search path in set variable statement", + })?; + query_ctx.set_current_schema(search_path); } else { return NotSupportedSnafu { feat: format!("Unsupported set variable {}", var_name), diff --git a/src/operator/src/statement/set.rs b/src/operator/src/statement/set.rs index 38a16e1b051..77515e62e09 100644 --- a/src/operator/src/statement/set.rs +++ b/src/operator/src/statement/set.rs @@ -129,36 +129,6 @@ pub fn set_bytea_output(exprs: Vec, ctx: QueryContextRef) -> Result<()> { Ok(()) } -pub fn set_search_path(exprs: Vec, ctx: QueryContextRef) -> Result<()> { - let search_expr = exprs.first().context(NotSupportedSnafu { - feat: "No search path find in set variable statement", - })?; - match search_expr { - Expr::Value(ValueWithSpan { - value: Value::SingleQuotedString(search_path), - .. - }) - | Expr::Value(ValueWithSpan { - value: Value::DoubleQuotedString(search_path), - .. - }) => { - ctx.set_current_schema(search_path); - Ok(()) - } - Expr::Identifier(Ident { value, .. }) => { - ctx.set_current_schema(value); - Ok(()) - } - expr => NotSupportedSnafu { - feat: format!( - "Unsupported search path expr {} in set variable statement", - expr - ), - } - .fail(), - } -} - pub fn validate_client_encoding(set: SetVariables) -> Result<()> { let Some((encoding, [])) = set.value.split_first() else { return InvalidSqlSnafu { diff --git a/src/sql/src/statements/set_variables.rs b/src/sql/src/statements/set_variables.rs index 4c6eb16692f..346f1762e62 100644 --- a/src/sql/src/statements/set_variables.rs +++ b/src/sql/src/statements/set_variables.rs @@ -15,7 +15,7 @@ use std::fmt::Display; use serde::Serialize; -use sqlparser::ast::{Expr, ObjectName}; +use sqlparser::ast::{Expr, ObjectName, ObjectNamePart, Value, ValueWithSpan}; use sqlparser_derive::{Visit, VisitMut}; /// SET variables statement. @@ -25,6 +25,27 @@ pub struct SetVariables { pub value: Vec, } +impl SetVariables { + /// Returns the first supported `search_path` value. + pub fn search_path(&self) -> Option<&str> { + let [ObjectNamePart::Identifier(variable)] = self.variable.0.as_slice() else { + return None; + }; + if variable.quote_style.is_some() || !variable.value.eq_ignore_ascii_case("search_path") { + return None; + } + + match self.value.first()? { + Expr::Value(ValueWithSpan { + value: Value::SingleQuotedString(value) | Value::DoubleQuotedString(value), + .. + }) => Some(value), + Expr::Identifier(value) => Some(&value.value), + _ => None, + } + } +} + impl Display for SetVariables { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let variable = &self.variable; @@ -70,4 +91,26 @@ SET delayed_insert_timeout = 300"#, } } } + + #[test] + fn test_search_path() { + for (sql, expected) in [ + ("SET search_path TO public", Some("public")), + ("SET SEARCH_PATH TO 'private'", Some("private")), + ("SET search_path TO \"quoted\"", Some("quoted")), + ("SET \"search_path\" TO private", None), + ("SET timezone TO 'UTC'", None), + ] { + let statements = ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let [Statement::SetVariables(set)] = statements.as_slice() else { + unreachable!() + }; + assert_eq!(expected, set.search_path()); + } + } } diff --git a/src/sql/src/statements/statement.rs b/src/sql/src/statements/statement.rs index bac7e367cd9..ec7c73573b5 100644 --- a/src/sql/src/statements/statement.rs +++ b/src/sql/src/statements/statement.rs @@ -178,7 +178,6 @@ impl Statement { | Statement::ShowSearchPath(_) | Statement::ShowViews(_) | Statement::DescribeTable(_) - | Statement::Explain(_) | Statement::ShowVariables(_) | Statement::ShowProcesslist(_) | Statement::FetchCursor(_) @@ -189,6 +188,8 @@ impl Statement { #[cfg(feature = "enterprise")] Statement::ShowTriggers(_) => true, + Statement::Explain(explain) => !explain.analyze || explain.statement.is_readonly(), + // Write operations Statement::Insert(_) | Statement::Delete(_) diff --git a/src/sql/src/util.rs b/src/sql/src/util.rs index 6a7abe23bb5..5868be6aa0e 100644 --- a/src/sql/src/util.rs +++ b/src/sql/src/util.rs @@ -14,20 +14,19 @@ use std::collections::HashSet; use std::fmt::{Display, Formatter}; +use std::ops::ControlFlow; use itertools::Itertools; use promql_parser::label::{METRIC_NAME, MatchOp}; use promql_parser::parser::{ - AggregateExpr as PromAggregateExpr, BinaryExpr as PromBinaryExpr, Call as PromCall, - Expr as PromExpr, MatrixSelector as PromMatrixSelector, ParenExpr as PromParenExpr, - SubqueryExpr as PromSubqueryExpr, UnaryExpr as PromUnaryExpr, - VectorSelector as PromVectorSelector, + Expr as PromExpr, MatrixSelector as PromMatrixSelector, VectorSelector as PromVectorSelector, }; +use promql_parser::util::{ExprVisitor, walk_expr}; use serde::Serialize; use snafu::ensure; use sqlparser::ast::{ - Array, Expr, Ident, ObjectName, ObjectNamePart, SetExpr, SqlOption, TableFactor, - TableWithJoins, Value, ValueWithSpan, + Array, Expr, Ident, ObjectName, ObjectNamePart, SqlOption, Value, ValueWithSpan, + Visit as AstVisit, Visitor, }; use sqlparser_derive::{Visit, VisitMut}; @@ -184,48 +183,60 @@ pub fn parse_option_string(option: SqlOption) -> Result<(String, OptionValue)> { /// Walk through a [Query] and extract all the tables referenced in it. pub fn extract_tables_from_query(query: &SqlOrTql) -> impl Iterator { - let mut names = HashSet::new(); - - match query { - SqlOrTql::Sql(query, _) => { - extract_tables_from_sql_query(&query.inner, &mut names); - extract_tables_from_hybrid_cte_query(query, &mut names); - } - SqlOrTql::Tql(tql, _) => extract_tables_from_tql(tql, &mut names), - } - - names.into_iter() + extract_tables_from_query_inner(query).0.into_iter() } -fn extract_tables_from_hybrid_cte_query(query: &Query, sql_names: &mut HashSet) { - if let Some(hybrid_cte) = &query.hybrid_cte { - let mut cte_names: HashSet = hybrid_cte - .cte_tables - .iter() - .map(|cte| ParserContext::canonicalize_identifier(cte.name.clone()).value) - .collect(); - remove_cte_names(sql_names, &cte_names); +/// Walk through a [Query] and extract its referenced tables, returning `None` +/// if any table reference cannot be resolved statically. +pub fn extract_tables_from_query_checked( + query: &SqlOrTql, +) -> Option> { + let (names, complete) = extract_tables_from_query_inner(query); + complete.then_some(names.into_iter()) +} - cte_names.clear(); - for cte in &hybrid_cte.cte_tables { - let cte_name = ParserContext::canonicalize_identifier(cte.name.clone()).value; - let mut cte_query_names = HashSet::new(); - match &cte.content { - CteContent::Sql(cte_query) => { - extract_tables_from_sql_query(cte_query, &mut cte_query_names) - } - CteContent::Tql(tql) => extract_tables_from_tql(tql, &mut cte_query_names), - } - if hybrid_cte.recursive { - cte_names.insert(cte_name.clone()); - } - remove_cte_names(&mut cte_query_names, &cte_names); - sql_names.extend(cte_query_names); - if !hybrid_cte.recursive { - cte_names.insert(cte_name); - } +fn extract_tables_from_query_inner(query: &SqlOrTql) -> (HashSet, bool) { + let mut names = HashSet::new(); + + let complete = match query { + SqlOrTql::Sql(query, _) => { + extract_tables_from_sql_query(&query.inner, &mut names); + extract_tables_from_hybrid_cte_query(query, &mut names) } + SqlOrTql::Tql(tql, _) => extract_tables_from_tql(tql, &mut names), + }; + + (names, complete) +} + +fn extract_tables_from_hybrid_cte_query( + query: &Query, + sql_names: &mut HashSet, +) -> bool { + let Some(hybrid_cte) = &query.hybrid_cte else { + return true; + }; + + let mut complete = true; + let cte_names: HashSet = hybrid_cte + .cte_tables + .iter() + .map(|cte| ParserContext::canonicalize_identifier(cte.name.clone()).value) + .collect(); + remove_cte_names(sql_names, &cte_names); + + for cte in &hybrid_cte.cte_tables { + let mut cte_query_names = HashSet::new(); + match &cte.content { + CteContent::Sql(cte_query) => { + extract_tables_from_sql_query(cte_query, &mut cte_query_names) + } + CteContent::Tql(tql) => complete &= extract_tables_from_tql(tql, &mut cte_query_names), + } + sql_names.extend(cte_query_names); } + + complete } fn remove_cte_names(names: &mut HashSet, cte_names: &HashSet) { @@ -246,55 +257,55 @@ fn remove_cte_names(names: &mut HashSet, cte_names: &HashSet }); } -fn extract_tables_from_tql(tql: &Tql, names: &mut HashSet) { +fn extract_tables_from_tql(tql: &Tql, names: &mut HashSet) -> bool { let promql = match tql { Tql::Eval(eval) => &eval.query, Tql::Explain(explain) => &explain.query, Tql::Analyze(analyze) => &analyze.query, }; - if let Ok(expr) = promql_parser::parser::parse(promql) { - extract_tables_from_prom_expr(&expr, names); - } + let Ok(expr) = promql_parser::parser::parse(promql) else { + return false; + }; + extract_tables_from_prom_expr(&expr, names) } -fn extract_tables_from_prom_expr(expr: &PromExpr, names: &mut HashSet) { - match expr { - PromExpr::Aggregate(PromAggregateExpr { expr, .. }) => { - extract_tables_from_prom_expr(expr, names); - } - PromExpr::Unary(PromUnaryExpr { expr, .. }) => { - extract_tables_from_prom_expr(expr, names); - } - PromExpr::Binary(PromBinaryExpr { lhs, rhs, .. }) => { - extract_tables_from_prom_expr(lhs, names); - extract_tables_from_prom_expr(rhs, names); - } - PromExpr::Paren(PromParenExpr { expr }) => { - extract_tables_from_prom_expr(expr, names); - } - PromExpr::Subquery(PromSubqueryExpr { expr, .. }) => { - extract_tables_from_prom_expr(expr, names); - } - PromExpr::VectorSelector(selector) => { - extract_metric_name_from_vector_selector(selector, names); - } - PromExpr::MatrixSelector(PromMatrixSelector { vs, .. }) => { - extract_metric_name_from_vector_selector(vs, names); - } - PromExpr::Call(PromCall { args, .. }) => { - for arg in &args.args { - extract_tables_from_prom_expr(arg, names); - } - } - PromExpr::NumberLiteral(_) | PromExpr::StringLiteral(_) | PromExpr::Extension(_) => {} +fn extract_tables_from_prom_expr(expr: &PromExpr, names: &mut HashSet) -> bool { + struct TableCollector<'a> { + names: &'a mut HashSet, + complete: bool, } + + impl ExprVisitor for TableCollector<'_> { + type Error = (); + + fn pre_visit(&mut self, expr: &PromExpr) -> std::result::Result { + self.complete &= match expr { + PromExpr::VectorSelector(selector) => { + extract_metric_name_from_vector_selector(selector, self.names) + } + PromExpr::MatrixSelector(PromMatrixSelector { vs, .. }) => { + extract_metric_name_from_vector_selector(vs, self.names) + } + PromExpr::Extension(_) => false, + _ => true, + }; + Ok(true) + } + } + + let mut collector = TableCollector { + names, + complete: true, + }; + let _ = walk_expr(&mut collector, expr); + collector.complete } fn extract_metric_name_from_vector_selector( selector: &PromVectorSelector, names: &mut HashSet, -) { +) -> bool { let metric_name = selector.name.clone().or_else(|| { let mut metric_name_matchers = selector.matchers.find_matchers(METRIC_NAME); if metric_name_matchers.len() == 1 && metric_name_matchers[0].op == MatchOp::Equal { @@ -304,9 +315,16 @@ fn extract_metric_name_from_vector_selector( } }); let Some(metric_name) = metric_name else { - return; + return false; }; + if selector.matchers.matchers.iter().any(|matcher| { + (matcher.name == SCHEMA_MATCHER || matcher.name == DATABASE_MATCHER) + && matcher.op != MatchOp::Equal + }) { + return false; + } + let schema_matcher = selector.matchers.matchers.iter().rev().find(|matcher| { matcher.op == MatchOp::Equal && (matcher.name == SCHEMA_MATCHER || matcher.name == DATABASE_MATCHER) @@ -322,6 +340,7 @@ fn extract_metric_name_from_vector_selector( metric_name, ))])); } + true } /// translate the start location to the index in the sql string @@ -340,96 +359,116 @@ pub fn location_to_index(sql: &str, location: &sqlparser::tokenizer::Location) - index - 1 } +// Tracks CTE aliases as sqlparser's visitor enters their query bodies. +struct QueryScope { + cte_names: Vec, + next_cte: usize, + recursive: bool, + scope_start: usize, + add_to_parent_after: Option, +} + +struct RelationCollector<'a> { + names: &'a mut HashSet, + ctes_in_scope: Vec, + query_scopes: Vec, + #[cfg(test)] + query_visits: usize, +} + +impl<'a> RelationCollector<'a> { + fn new(names: &'a mut HashSet) -> Self { + Self { + names, + ctes_in_scope: Vec::new(), + query_scopes: Vec::new(), + #[cfg(test)] + query_visits: 0, + } + } +} + +impl Visitor for RelationCollector<'_> { + type Break = (); + + fn pre_visit_query(&mut self, query: &sqlparser::ast::Query) -> ControlFlow { + #[cfg(test)] + { + self.query_visits += 1; + } + + // Query::visit enters WITH definitions in declaration order before its body. + let parent_cte = self.query_scopes.last_mut().and_then(|scope| { + let cte_name = scope.cte_names.get(scope.next_cte)?.clone(); + scope.next_cte += 1; + Some((cte_name, scope.recursive)) + }); + + let add_to_parent_after = match parent_cte { + Some((cte_name, true)) => { + self.ctes_in_scope.push(cte_name); + None + } + Some((cte_name, false)) => Some(cte_name), + None => None, + }; + + let scope_start = self.ctes_in_scope.len(); + let (cte_names, recursive) = if let Some(with) = &query.with { + ( + with.cte_tables + .iter() + .map(|cte| ParserContext::canonicalize_identifier(cte.alias.name.clone()).value) + .collect(), + with.recursive, + ) + } else { + (Vec::new(), false) + }; + + self.query_scopes.push(QueryScope { + cte_names, + next_cte: 0, + recursive, + scope_start, + add_to_parent_after, + }); + ControlFlow::Continue(()) + } + + fn post_visit_query(&mut self, _query: &sqlparser::ast::Query) -> ControlFlow { + let Some(scope) = self.query_scopes.pop() else { + return ControlFlow::Break(()); + }; + self.ctes_in_scope.truncate(scope.scope_start); + if let Some(cte_name) = scope.add_to_parent_after { + self.ctes_in_scope.push(cte_name); + } + ControlFlow::Continue(()) + } + + fn pre_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow { + let is_cte = matches!( + relation.0.as_slice(), + [part] + if part.as_ident().is_some_and(|ident| { + self.ctes_in_scope.contains( + &ParserContext::canonicalize_identifier(ident.clone()).value, + ) + }) + ); + if !is_cte { + self.names.insert(relation.clone()); + } + ControlFlow::Continue(()) + } +} + /// Helper function for [extract_tables_from_query]. /// /// Handle [sqlparser::ast::Query]. fn extract_tables_from_sql_query(query: &sqlparser::ast::Query, names: &mut HashSet) { - let mut cte_names = HashSet::new(); - if let Some(with) = &query.with { - for cte in &with.cte_tables { - let cte_name = ParserContext::canonicalize_identifier(cte.alias.name.clone()).value; - let mut cte_query_names = HashSet::new(); - extract_tables_from_sql_query(&cte.query, &mut cte_query_names); - if with.recursive { - cte_names.insert(cte_name.clone()); - } - remove_cte_names(&mut cte_query_names, &cte_names); - names.extend(cte_query_names); - if !with.recursive { - cte_names.insert(cte_name); - } - } - } - - let mut body_names = HashSet::new(); - extract_tables_from_set_expr(&query.body, &mut body_names); - remove_cte_names(&mut body_names, &cte_names); - names.extend(body_names); -} - -/// Helper function for [extract_tables_from_query]. -/// -/// Handle [SetExpr]. -fn extract_tables_from_set_expr(set_expr: &SetExpr, names: &mut HashSet) { - match set_expr { - SetExpr::Select(select) => { - for from in &select.from { - extract_tables_from_table_with_joins(from, names); - } - } - SetExpr::Query(query) => { - extract_tables_from_sql_query(query, names); - } - SetExpr::SetOperation { left, right, .. } => { - extract_tables_from_set_expr(left, names); - extract_tables_from_set_expr(right, names); - } - _ => {} - }; -} - -/// Helper function for [extract_tables_from_query]. -/// -/// Handle [TableWithJoins]. -fn extract_tables_from_table_with_joins( - table_with_joins: &TableWithJoins, - names: &mut HashSet, -) { - table_factor_to_object_name(&table_with_joins.relation, names); - for join in &table_with_joins.joins { - table_factor_to_object_name(&join.relation, names); - } -} - -/// Helper function for [extract_tables_from_query]. -/// -/// Handle [TableFactor]. -fn table_factor_to_object_name(table_factor: &TableFactor, names: &mut HashSet) { - match table_factor { - TableFactor::Table { name, .. } => { - names.insert(name.to_owned()); - } - TableFactor::Derived { subquery, .. } => { - extract_tables_from_sql_query(subquery, names); - } - TableFactor::NestedJoin { - table_with_joins, .. - } => { - extract_tables_from_table_with_joins(table_with_joins, names); - } - TableFactor::Pivot { table, .. } - | TableFactor::Unpivot { table, .. } - | TableFactor::MatchRecognize { table, .. } => { - table_factor_to_object_name(table, names); - } - TableFactor::TableFunction { .. } - | TableFactor::Function { .. } - | TableFactor::UNNEST { .. } - | TableFactor::JsonTable { .. } - | TableFactor::OpenJsonTable { .. } - | TableFactor::XmlTable { .. } - | TableFactor::SemanticView { .. } => {} - } + let _ = query.visit(&mut RelationCollector::new(names)); } #[cfg(test)] @@ -518,6 +557,65 @@ TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", {__n } } + #[test] + fn test_extract_tables_from_tql_query_completeness() { + let testcases = [ + ("cpu", true, vec!["cpu"]), + (r#"{__name__="cpu"}"#, true, vec!["cpu"]), + (r#"cpu + {job="api"}"#, false, vec!["cpu"]), + (r#"{__name__=~"cpu.*"}"#, false, vec![]), + (r#"cpu{__schema__=~"private.*"}"#, false, vec![]), + ]; + + for (promql, expected_complete, expected_tables) in testcases { + let sql = format!("TQL EVAL (0, 10, '5s') {promql}"); + let mut stmts = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::Tql(tql) = stmts.pop().unwrap() else { + unreachable!() + }; + + let query = SqlOrTql::Tql(tql, sql); + let (tables, complete) = extract_tables_from_query_inner(&query); + let mut tables = tables + .into_iter() + .map(|table| format_raw_object_name(&table)) + .collect_vec(); + tables.sort(); + assert_eq!(expected_complete, complete, "{promql}"); + assert_eq!(expected_tables, tables, "{promql}"); + assert_eq!( + expected_complete, + extract_tables_from_query_checked(&query).is_some(), + "{promql}" + ); + } + } + + #[test] + fn test_extract_tables_from_chained_tql_ctes() { + let sql = "WITH denied AS (TQL EVAL (0, 10, '5s') allowed), \ + leak AS (TQL EVAL (0, 10, '5s') denied) \ + SELECT * FROM leak"; + let mut stmts = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + let Statement::Query(query) = stmts.pop().unwrap() else { + unreachable!() + }; + + let mut tables = extract_tables_from_query_checked(&SqlOrTql::Sql(*query, sql.to_string())) + .unwrap() + .map(|table| format_raw_object_name(&table)) + .collect_vec(); + tables.sort(); + assert_eq!(vec!["allowed".to_string(), "denied".to_string()], tables); + } + #[test] fn test_extract_tables_from_sql_query_with_derived_join() { let sql = r#" @@ -556,6 +654,37 @@ LEFT JOIN ( ); } + #[test] + fn test_extract_tables_from_sql_query_with_expression_subqueries() { + let sql = r#" +SELECT + (SELECT max(value) FROM scalar_source) +FROM outer_source +WHERE EXISTS (SELECT 1 FROM exists_source) +ORDER BY (SELECT max(value) FROM order_source); +"#; + let mut stmts = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + let Statement::Query(query) = stmts.pop().unwrap() else { + unreachable!() + }; + + let mut tables = extract_tables_from_query(&SqlOrTql::Sql(*query, sql.to_string())) + .map(|table| format_raw_object_name(&table)) + .collect_vec(); + tables.sort(); + assert_eq!( + vec![ + "exists_source".to_string(), + "order_source".to_string(), + "outer_source".to_string(), + "scalar_source".to_string(), + ], + tables + ); + } + #[test] fn test_extract_tables_from_sql_query_with_cte_scopes() { let testcases = vec![ @@ -570,6 +699,15 @@ SELECT * FROM source; ), ( r#" +WITH RECURSIVE source AS ( + SELECT * FROM source +) +SELECT * FROM source; +"#, + vec![], + ), + ( + r#" WITH first_cte AS ( SELECT * FROM physical_source ), second_cte AS ( @@ -579,6 +717,24 @@ SELECT * FROM second_cte; "#, vec!["physical_source".to_string()], ), + ( + r#" +SELECT * FROM ( + WITH nested_cte AS (SELECT * FROM nested_source) + SELECT * FROM nested_cte +); +"#, + vec!["nested_source".to_string()], + ), + ( + r#" +SELECT * FROM ( + WITH nested_cte AS (SELECT * FROM nested_source) + SELECT (SELECT * FROM nested_cte) +); +"#, + vec!["nested_source".to_string()], + ), ]; for (sql, expected_tables) in testcases { @@ -603,6 +759,40 @@ SELECT * FROM second_cte; } } + #[test] + fn test_extract_tables_from_deeply_nested_ctes_visits_each_query_once() { + let mut sql = "SELECT * FROM physical_source".to_string(); + for depth in 0..20 { + sql = format!("WITH cte_{depth} AS ({sql}) SELECT * FROM cte_{depth}"); + } + + let mut stmts = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::Query(query) = stmts.pop().unwrap() else { + unreachable!() + }; + + let mut tables = HashSet::new(); + let query_visits = { + let mut collector = RelationCollector::new(&mut tables); + let _ = query.inner.visit(&mut collector); + collector.query_visits + }; + + assert_eq!(21, query_visits); + assert_eq!( + vec!["physical_source"], + tables + .into_iter() + .map(|table| format_raw_object_name(&table)) + .collect_vec() + ); + } + #[test] fn test_extract_tables_from_tql_query_with_schema_matcher() { let sql = r#" From 16217ff567052b020a07083f07036ed8d139325d Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 15 Jul 2026 09:46:40 +0800 Subject: [PATCH 29/48] ci: add persistent query regression cache (#8474) * ci: add persistent query regression cache Signed-off-by: discord9 * ci: add sccache to query regression runner Signed-off-by: discord9 * ci: pin query regression label revision Signed-off-by: discord9 * ci: isolate query regression toolchain state Signed-off-by: discord9 --------- Signed-off-by: discord9 --- .../query-regression/Dockerfile | 51 +++ .../query-regression/README.md | 344 +++++++++----- .../query-regression/cache-pvc.yaml | 15 + .../query-regression/values-8-cores.yaml | 101 ++++- .../query-regression/values-paused.yaml | 2 + .github/scripts/query-regression-comment.cjs | 4 +- .../scripts/query-regression-pr-metadata.py | 5 +- .github/workflows/query-regression.yml | 424 +++++++++++++++++- 8 files changed, 817 insertions(+), 129 deletions(-) create mode 100644 .github/runner-scale-sets/query-regression/cache-pvc.yaml create mode 100644 .github/runner-scale-sets/query-regression/values-paused.yaml diff --git a/.github/runner-scale-sets/query-regression/Dockerfile b/.github/runner-scale-sets/query-regression/Dockerfile index 1f9f7203b9d..95161f6093b 100644 --- a/.github/runner-scale-sets/query-regression/Dockerfile +++ b/.github/runner-scale-sets/query-regression/Dockerfile @@ -3,6 +3,19 @@ FROM ghcr.io/actions/actions-runner@sha256:08c30b0a7105f64bddfc485d2487a22aa0393 USER root +ARG SCCACHE_VERSION=0.16.0 +ARG SCCACHE_SHA256=aec995a83ad3dff3d14b6314e08858b7b73d35ca85a5bcf3d3a9ec07dee35588 +ARG RUSTUP_INIT_VERSION=1.29.0 +ARG RUSTUP_INIT_TARGET=x86_64-unknown-linux-gnu +ARG RUSTUP_INIT_SHA256=4acc9acc76d5079515b46346a485974457b5a79893cfb01112423c89aeb5aa10 +ARG RUST_TOOLCHAIN=nightly-2026-03-21 + +ENV RUSTUP_HOME=/opt/rustup \ + CARGO_HOME=/opt/cargo \ + RUSTUP_TOOLCHAIN=${RUST_TOOLCHAIN} \ + RUSTUP_AUTO_INSTALL=0 \ + PATH=/opt/cargo/bin:${PATH} + RUN apt-get update \ && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ build-essential \ @@ -30,4 +43,42 @@ RUN apt-get update \ COPY --from=uv /uv /uvx /usr/local/bin/ +RUN curl --fail --location --silent --show-error \ + --output /tmp/sccache.tar.gz \ + "https://github.com/mozilla/sccache/releases/download/v${SCCACHE_VERSION}/sccache-v${SCCACHE_VERSION}-x86_64-unknown-linux-musl.tar.gz" \ + && echo "${SCCACHE_SHA256} /tmp/sccache.tar.gz" | sha256sum --check --status - \ + && mkdir --parents /tmp/sccache \ + && tar --extract --gzip --file /tmp/sccache.tar.gz \ + --directory /tmp/sccache --strip-components=1 \ + "sccache-v${SCCACHE_VERSION}-x86_64-unknown-linux-musl/sccache" \ + && install --mode=0755 /tmp/sccache/sccache /usr/local/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache + +RUN curl --fail --location --silent --show-error \ + --output /tmp/rustup-init \ + "https://static.rust-lang.org/rustup/archive/${RUSTUP_INIT_VERSION}/${RUSTUP_INIT_TARGET}/rustup-init" \ + && echo "${RUSTUP_INIT_SHA256} /tmp/rustup-init" | sha256sum --check --status - \ + && chmod 0755 /tmp/rustup-init \ + && /tmp/rustup-init -y --profile minimal --default-toolchain "${RUST_TOOLCHAIN}" --no-modify-path \ + && rm -f /tmp/rustup-init \ + && chown -R root:root /opt/rustup /opt/cargo \ + && chmod -R go-w /opt/rustup /opt/cargo \ + && find /opt/rustup /opt/cargo -type d -exec chmod 0755 {} + \ + && find /opt/rustup /opt/cargo -type f -exec chmod a+r {} + + USER runner + +RUN set -eu \ + && test "$(id -u)" = "1001" \ + && temporary_cargo_home="$(mktemp --directory)" \ + && trap 'rm -rf "${temporary_cargo_home}"' 0 \ + && export CARGO_HOME="${temporary_cargo_home}" \ + && test "${RUSTUP_AUTO_INSTALL}" = "0" \ + && rustup --version \ + && cargo --version \ + && rustc --version \ + && active_toolchain="$(rustup show active-toolchain)" \ + && printf 'Active toolchain: %s\n' "${active_toolchain}" \ + && case "${active_toolchain}" in "${RUST_TOOLCHAIN}-x86_64-unknown-linux-gnu"|"${RUST_TOOLCHAIN}-x86_64-unknown-linux-gnu "*) ;; *) exit 1;; esac \ + && test ! -w /opt/rustup \ + && test ! -w /opt/cargo/bin diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index eb1a7e12f79..fd68b4c8321 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -1,21 +1,14 @@ # Query regression self-hosted runners -The `Query Regression` workflow targets self-hosted GitHub Actions runners via -runner labels or ARC runner scale set names: +The `Query Regression` workflow uses the dedicated ARC runner scale set +`perf-regression-8-cores`. ARC runner Pods run in the target Kubernetes cluster +and connect outbound to GitHub. The live scale set is currently **paused**: +`minRunners=0`, `maxRunners=0`, and no runner Pods. Do not resume it without +explicit approval. -- `perf-regression-8-cores` +## Prerequisites and trust admission -The name intentionally avoids generic labels such as `ubuntu-22.04-8-cores`, -which may already be used by GitHub-hosted larger runners or other runner pools. - -For Kubernetes-based runners, ARC runner pods run inside the target Kubernetes -cluster and connect outbound to GitHub; GitHub can then dispatch jobs whose -`runs-on` value matches the runner scale set name or runner labels. - -## Prerequisites - -Install the ARC scale set controller in the Kubernetes cluster if it is not -already installed: +Install the ARC scale set controller if it is not already installed: ```bash helm upgrade --install arc \ @@ -25,8 +18,8 @@ helm upgrade --install arc \ --version 0.14.2 ``` -Create the GitHub App or PAT secret in the namespace that will host the runner -scale sets. Prefer a GitHub App with access limited to `GreptimeTeam/greptimedb`. +Create the GitHub App secret in the runner namespace. Prefer an App limited to +`GreptimeTeam/greptimedb`: ```bash kubectl -n arc-runners create secret generic greptimedb-arc-github-app \ @@ -35,14 +28,46 @@ kubectl -n arc-runners create secret generic greptimedb-arc-github-app \ --from-file=github_app_private_key= ``` -The values files in this directory reference that secret by name. +The values files here reference that secret by name. -## Build the runner image +A maintainer applying the `query-regression` label is **trust admission for +that exact PR revision**. The admitted job may use this scale set's dedicated, +writable persistent cache. `pull_request: labeled` is the only PR trigger: the +label event snapshots its merge, head, and base SHAs. A queued job fetches that +immutable event merge SHA directly, verifies it is a two-parent merge whose +parents include the snapshotted head exactly once, and uses its other parent as +the actual base build revision. The snapshotted event base is retained for audit +only, so a difference from the merge's non-head parent is not a failure. The job +never follows a newer mutable PR merge ref. An unavailable event merge, or one +that does not contain exactly one snapshotted head parent, fails closed. A later +PR head change does not retarget an already queued run: it may execute only its +previously trusted event revision if that revision remains fetchable. To run the +new revision, the maintainer must review it, remove the label, and re-add +`query-regression`; cancel the old run if it is no longer wanted. An existing +label does not automatically rerun the benchmark. -Use a derived ARC runner image instead of the minimal upstream runner image. The -image keeps the official `/home/runner/run.sh` entrypoint layout and adds the -tools this workflow expects to be present on a normal CI host, including `wget`, -`uv`, C/C++ build tools, OpenSSL headers, and protobuf tooling. +Admission does not relax runner hardening or GitHub permissions. Keep +service-account token mounting disabled; do not mount host paths, the Docker +socket, kubeconfig, or long-lived credentials. The runner and cache initializer +use UID/GID 1001, disallow privilege escalation, drop all capabilities, and use +the RuntimeDefault seccomp profile. Keep GitHub tokens least-privilege and +review workflow changes before admission. Where the CNI supports it, restrict +egress to required GitHub Actions, artifact/cache, Rust/crate/toolchain, DNS, +and image-registry endpoints; block unrelated cluster services, private ranges, +and metadata endpoints unless a case requires them. + +### Office routing prerequisite + +Direct split routing through `.2` is an **external office-network prerequisite**. +The gateway at `192.168.50.2` must route GitHub Actions, GitHub content, +artifact/cache, crates.io, Rust toolchain, and image-registry traffic directly +rather than through the VPN. Neither this repository nor Kubernetes configures +that route. Verify it with the responsible network operator before any canary. + +## Runner image and workflow tools + +Build and push the derived runner image; it preserves the official +`/home/runner/run.sh` entrypoint and supplies CI tools needed at runtime: ```bash docker build \ @@ -53,100 +78,219 @@ docker build \ docker push greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest ``` -Deploy the runner image by digest rather than by a mutable tag. Update -`values-8-cores.yaml` after pushing a rebuilt image. +Deploy by digest, not mutable tag, by updating `values-8-cores.yaml` after a +rebuild. If the registry is private, use a dedicated read-only pull secret only +as `imagePullSecrets`; never expose registry credentials to runner containers. +Both digest-pinned init and runner containers use `IfNotPresent`: the immutable +digest makes a cached image safe and avoids adding a registry dependency to every +runner startup. -The workflow still runs setup actions for pinned Rust and `uv` behavior. `mold` -is installed in the image and selected through `CARGO_BUILD_RUSTFLAGS`, so jobs -do not need privileged package installation at runtime. +Before builds, the workflow asserts UID/GID 1001 and exact image tool versions: +`libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.30.0`, `Python 3.12.3`, `sccache +0.16.0`, root-owned `rustup 1.29.0`, and the image-baked +`nightly-2026-03-21` Rust toolchain. Rustup, Cargo, and Rustc must resolve from +`/opt/cargo/bin`; the runner cannot write `/opt/rustup` or `/opt/cargo/bin`. +`actions-rust-lang/setup-rust-toolchain@v1` is intentionally removed. The +workflow sets its warning-denying mold `RUSTFLAGS` directly, disables automatic +Rustup installation, and performs no runtime toolchain downloads. -Decide whether the registry repository is public or private. Public pull access -avoids distributing registry credentials to the runner namespace. If the image -must be private, create a dedicated read-only image-pull secret for this image -and attach it only as `imagePullSecrets`; do not expose registry credentials to -runner containers. +The workflow no longer uses GitHub `rust-cache`, `setup-protoc`, `setup-uv`, or +runtime Rust setup: the image establishes immutable executable state and the PVC +supplies only reusable Cargo data. Do not reintroduce those actions unless the +corresponding cache or image contract changes. -## Install the query-regression scale set +## Capacity and persistent cache -Install the runner scale set. The Helm release name and `runnerScaleSetName` -should match the `runs-on` value used by the workflow. +`values-8-cores.yaml` is normal operation: `minRunners=0`, `maxRunners=1`. +`values-paused.yaml` is the mandatory pause overlay: `minRunners=0`, +`maxRunners=0`. The job uses group `query-regression-persistent-cache-v1`, +`queue: max`, and `cancel-in-progress: false`; admitted jobs queue rather than +replacing older pending jobs. During maintenance, cancel admitted queued runs as +well as pausing ARC. Runner Pods have `activeDeadlineSeconds=12600`. + +The cache claim `query-regression-build-cache` is a nominal 600Gi `local-path` +PVC in `arc-runners`. It is `ReadWriteOnce`; `local-path` uses +WaitForFirstConsumer binding and Delete reclaim behavior, produces a +node-affine local PV, is non-expandable, and the 600Gi request is not a hard +storage quota. The runner's `minipc-3` selector is its only consumer candidate. + +The initializer mounts the PVC root at `/cache`, creates and write-tests these +versioned subpaths as non-root UID/GID 1001, and the runner mounts them as: + +| Persistent state | PVC subpath | Runner mount | +| --- | --- | --- | +| Ephemeral Cargo home | `emptyDir` | `/home/runner/.cargo` | +| Cargo registry data | `cargo-registry-v1` | `/home/runner/.cargo/registry` | +| Cargo Git data | `cargo-git-v1` | `/home/runner/.cargo/git` | +| Cargo target | `query-regression-target-v1` | `/home/runner/query-regression-target` | +| Cache metadata | `meta-v1` | `/home/runner/query-regression-cache-meta` | +| sccache local disk cache | `sccache-v1` | `/home/runner/.cache/sccache` | +| Immutable Rust toolchain | image-owned | `/opt/rustup`, `/opt/cargo/bin` | + +The Pod security context uses UID/GID and `fsGroup` 1001 with +`fsGroupChangePolicy: OnRootMismatch`; no privileged `chown` or raw `hostPath` +is used. `CARGO_HOME` is a per-Pod `emptyDir`; only its nested `registry` and +`git` mounts are persistent. `RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are +image-owned immutable paths, while `CARGO_TARGET_DIR`, cache metadata, and +`SCCACHE_DIR` are persistent absolute paths. The runner sets +`RUSTC_WRAPPER=/usr/local/bin/sccache`, +`SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=40G`, and +`CARGO_INCREMENTAL=0`. sccache uses its local PVC disk backend and self-evicts +at 40G; do not add runtime downloads, object storage, or a shared backend. + +The repository's `.cargo/config.toml` remains a trusted per-revision build input. +In contrast, `$CARGO_HOME/config*`, credentials, installed bins, and Cargo +metadata outside the persistent `registry` and `git` data mounts are ephemeral +and cannot survive to another Pod. + +The local disk backend has a one-server constraint. `maxRunners=1` and the +unchanged `query-regression-persistent-cache-v1` workflow concurrency group +serialize runs; do not increase runner capacity or relax that serialization +while this backend is in use. Base and candidate builds share the target; Cargo +fingerprints invalidate source and dependency changes. The workflow records the +sccache version and relevant environment in the target ABI marker, starts and +zeros sccache after cache and toolchain checks, shows initial/base/candidate +statistics, and resets statistics between base and candidate builds. + +### Disk preflight and cleanup contract + +Before applying or unpausing, verify the backing filesystem on `minipc-3` has +at least 900GiB free. The current local-path provisioner source is +`/opt/local-path-provisioner`; measure the filesystem containing it: + +```bash +df -PB1G /opt/local-path-provisioner +``` + +The workflow reports `du`, `df -P`, human-readable free space, and inode +availability before builds and in an always-run report. Its cleanup is narrow +and non-destructive: + +- warn at target size 400GiB; at 450GiB clear only the complete target root; +- warn at Cargo registry-plus-Git data size 60GiB; at 80GiB remove only + `registry/src` and `git/checkouts`, then abort if that persistent data remains + at least 80GiB; +- below 300GiB backing free space, clear the complete target root first, + remeasure, then remove only those Cargo extracted trees and checkouts; abort + if free space is still below 300GiB; +- never automatically remove Cargo registry cache/index, Git database, the + image-owned Cargo bin or Rustup toolchain, cache metadata, the self-evicting + sccache directory, or the PVC. + +The target clear uses fixed absolute roots and removes all entries, including +dotfiles. After migration validation, remove obsolete `cargo-home-v1` and +`rustup-home-v1` only in explicit maintenance while ARC is 0/0 and no runner Pod +exists; they are not mounted by the current configuration. + +## Deploy and pause safely + +First verify external `.2` routing and the disk preflight. Apply the PVC; while +the scale set is paused, it is expected to remain `Pending` because +WaitForFirstConsumer has no scheduled runner: + +```bash +kubectl apply --dry-run=server \ + -f .github/runner-scale-sets/query-regression/cache-pvc.yaml +kubectl apply -f .github/runner-scale-sets/query-regression/cache-pvc.yaml +``` + +Render normal and paused configurations. Normal values are always first; the +pause overlay is always last: + +```bash +helm template perf-regression-8-cores \ + oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ + --namespace arc-runners --version 0.14.2 \ + --set controllerServiceAccount.name=arc-gha-rs-controller \ + --set controllerServiceAccount.namespace=arc-systems \ + -f .github/runner-scale-sets/query-regression/values-8-cores.yaml + +helm template perf-regression-8-cores \ + oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ + --namespace arc-runners --version 0.14.2 \ + --set controllerServiceAccount.name=arc-gha-rs-controller \ + --set controllerServiceAccount.namespace=arc-systems \ + -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ + -f .github/runner-scale-sets/query-regression/values-paused.yaml +``` + +The **first post-merge Helm deployment must reconcile the release in paused +mode**. Keep the pause overlay last: + +```bash +# First post-merge deployment and every return to paused mode: 0/0. +helm upgrade --install perf-regression-8-cores \ + oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ + --namespace arc-runners --create-namespace --version 0.14.2 \ + --reset-values --wait \ + -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ + -f .github/runner-scale-sets/query-regression/values-paused.yaml + +# Expect 0/0 and no runner resources before considering normal mode. +kubectl -n arc-runners get autoscalingrunnerset perf-regression-8-cores \ + -o jsonpath='{.spec.minRunners}{"/"}{.spec.maxRunners}{"\n"}' +kubectl -n arc-runners get ephemeralrunners,pods \ + -l actions.github.com/scale-set-name=perf-regression-8-cores +``` + +Only after that verification and separate explicit approval, apply normal 0/1 +operation without the pause overlay: ```bash helm upgrade --install perf-regression-8-cores \ oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners \ - --create-namespace \ - --version 0.14.2 \ + --namespace arc-runners --create-namespace --version 0.14.2 \ + --reset-values --wait \ -f .github/runner-scale-sets/query-regression/values-8-cores.yaml ``` -Check registration and pods: +Do not use bare `helm rollback`, `--atomic`, or `--reuse-values`: a stored +revision can restore nonzero runner capacity. Inspect rendered manifests for +capacity, the `minipc-3` selector, cache claim and mounts, initializer, +security context, and resources. After approved normal mode receives its first +canary, the PVC binds to `minipc-3`. -```bash -kubectl -n arc-runners get pods -kubectl -n arc-runners get autoscalingrunnersets -``` +For local-PV node loss, cache recovery is intentionally disposable: return to +0/0, recreate the PVC on a healthy node, cold-fill it, and run a new canary. -The scale sets should also appear under repository Actions runner settings. +## Canary and rollback -## Security notes for fork PRs +With explicit approval, run two identical `workflow_dispatch` canaries on +`perf-regression-8-cores`, using immutable full base and candidate commit SHAs +and `cargo_profile=nightly`. The first is the cold fill; the second verifies warm +reuse. Record the workflow's base/candidate build elapsed logs and cache +ABI-marker output, initial/base/candidate sccache statistics, and cache report. +Confirm the image tool contract (root-owned Rustup/Cargo paths, baked nightly, +and non-writable `/opt` roots) and that the ephemeral Cargo home contains only +the mounted registry/Git data before Cargo creates per-Pod state. +Obtain dependency and tool network byte counters from the +identified `.2` counter source, filtered to `minipc-3` and the dependency/tool +destinations. -Maintainer-approved fork PRs can run on self-hosted runners. Approval only lets -the workflow execute; it does not make the fork code trusted. +Accept the canary only when all of the following hold: -The `query-regression` label is the explicit trigger for PR runs. Updating a PR -does not rerun the benchmark automatically; remove and re-add the label after -reviewing the updated changes. +- exactly one runner Pod runs on `minipc-3`, and both jobs use the same bound PV; +- UID/GID 1001 cache mounts are writable; the warm run does not invalidate the + target or bulk-redownload crates or toolchains; sccache reports separate base + and candidate build statistics without server or cache-path errors; immutable + Rustup/Cargo roots remain non-writable and only registry/Git data persists; +- warm base build time is at most 50% of cold base build time; +- warm dependency/tool network bytes are at most 10% of cold fill bytes; +- cache sizes remain below soft watermarks, node free space remains at least + 300GiB, and the benchmark is correct without TLS EOFs or timeouts; +- `.2` confirms this traffic is outside VPN accounting. -These runners execute PR code and should be treated as untrusted execution -capacity: +Immediately return to 0/0 after either canary unless ongoing normal operation +has been explicitly approved; return immediately on any traffic, cache, disk, +TLS, or correctness failure. To roll back, use the paused Helm upgrade above, +or another explicit `helm upgrade --install` with known-good values followed by +`values-paused.yaml`, `--reset-values`, and `--wait`. Do not delete the PVC +automatically; preserve it for diagnosis unless intentionally discarding cache. -- keep them isolated from sensitive internal services unless explicitly required; -- do not mount host paths, Docker socket, kubeconfig, or long-lived credentials; -- use ephemeral runner pods and no shared work directory with trusted jobs; -- disable service account token mounting in runner pods unless Kubernetes API - access is intentionally required; -- use a runner image whose default user is non-root, disable privilege - escalation, drop Linux capabilities, and use the runtime-default seccomp - profile; -- keep GitHub tokens least-privilege and rely on normal `pull_request` behavior - for fork PRs, where repository secrets are withheld and `GITHUB_TOKEN` is - read-only; -- review fork workflow changes before approving the run. +## Future optional phases -If stronger isolation is required, install a separate runner group/namespace just -for query-regression PR workloads and restrict repository/workflow access to the -`Query Regression` workflow. - -Use namespace or cluster network policy to restrict runner egress where the CNI -supports the needed controls. Query-regression runners need outbound access to -GitHub Actions services, GitHub artifact/cache endpoints, Rust/crate/toolchain -endpoints, DNS, and the configured image registries. Block access to unrelated -cluster services, private network ranges, and cloud metadata endpoints unless a -case explicitly needs them. - -## Build cache - -The workflow builds base and candidate in the same job, the same source path, -and a shared `CARGO_TARGET_DIR`. It checks out the base ref into `src`, builds -and copies the base binary aside, then resets that same `src` checkout to the -candidate ref before building the candidate binary. Keeping the workspace path -stable improves Cargo incremental reuse for local workspace crates compared with -building separate `base-src` and `candidate-src` checkouts. - -The workflow also uses the GitHub Actions Rust cache for restore-only cache -reuse. PR and dispatch runs do not save cache entries. Refresh shared caches from -trusted maintenance workflows only. - -The ARC values in this directory do not configure a cross-run runner-local -compiler cache. If cross-run compile time still dominates the benchmark, -deployers may add one of the following at the runner infrastructure layer: - -- a custom runner image with Rust tooling, `mold`, and `sccache` preinstalled; -- `RUSTC_WRAPPER=sccache` plus a shared `sccache` backend such as object storage, - Redis, or a Kubernetes storage class that is safe for concurrent runner pods; -- separate cache namespaces/buckets for untrusted PR code and trusted branches to - avoid cache poisoning across trust boundaries. - -Do not mount a shared writable host path, Docker socket, kubeconfig, or other -privileged credentials into runners that execute PR code. +The current phase uses a digest-pinned image with sccache 0.16.0 and no shared +cache service. Optional follow-ups are an image additionally seeded with the +exact Rust toolchain and `cargo fetch --locked`; or an internal read/write +sccache backend or Cargo/Git mirror. Evaluate them only if persistent PVC reuse +is insufficient. diff --git a/.github/runner-scale-sets/query-regression/cache-pvc.yaml b/.github/runner-scale-sets/query-regression/cache-pvc.yaml new file mode 100644 index 00000000000..692f3a2d6e2 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/cache-pvc.yaml @@ -0,0 +1,15 @@ +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: query-regression-build-cache + namespace: arc-runners + labels: + app.kubernetes.io/part-of: perf-regression-8-cores +spec: + volumeMode: Filesystem + accessModes: + - ReadWriteOnce + storageClassName: local-path + resources: + requests: + storage: 600Gi diff --git a/.github/runner-scale-sets/query-regression/values-8-cores.yaml b/.github/runner-scale-sets/query-regression/values-8-cores.yaml index 54ade99bad1..77ec4055e5b 100644 --- a/.github/runner-scale-sets/query-regression/values-8-cores.yaml +++ b/.github/runner-scale-sets/query-regression/values-8-cores.yaml @@ -3,21 +3,96 @@ githubConfigSecret: greptimedb-arc-github-app runnerScaleSetName: "perf-regression-8-cores" minRunners: 0 -maxRunners: 4 +maxRunners: 1 template: spec: automountServiceAccountToken: false - activeDeadlineSeconds: 11400 + activeDeadlineSeconds: 12600 + nodeSelector: + kubernetes.io/hostname: minipc-3 securityContext: + runAsNonRoot: true + runAsUser: 1001 + runAsGroup: 1001 + fsGroup: 1001 + fsGroupChangePolicy: OnRootMismatch seccompProfile: type: RuntimeDefault + volumes: + - name: cargo-home + emptyDir: {} + - name: build-cache + persistentVolumeClaim: + claimName: query-regression-build-cache + initContainers: + - name: initialize-build-cache + image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:7022b551771dc82ec4c493f0c3db6ee3050ececd336598690a2165240ebe953e + imagePullPolicy: IfNotPresent + command: + - /bin/sh + - -ec + - | + umask 0002 + for directory in cargo-registry-v1 cargo-git-v1 query-regression-target-v1 meta-v1 sccache-v1; do + cache_directory="/cache/${directory}" + mkdir -p "${cache_directory}" + test -w "${cache_directory}" + probe_file="${cache_directory}/.write-probe" + : > "${probe_file}" + test -f "${probe_file}" + rm "${probe_file}" + done + securityContext: + runAsNonRoot: true + runAsUser: 1001 + runAsGroup: 1001 + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + resources: + requests: + cpu: 100m + memory: 128Mi + ephemeral-storage: 1Gi + limits: + cpu: 500m + memory: 512Mi + ephemeral-storage: 1Gi + volumeMounts: + - name: build-cache + mountPath: /cache containers: - name: runner - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:108b7357532e5cf3329853373f561b0c674e954d6307b53d37075df5a189b0b2 - imagePullPolicy: Always + image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:7022b551771dc82ec4c493f0c3db6ee3050ececd336598690a2165240ebe953e + imagePullPolicy: IfNotPresent command: ["/home/runner/run.sh"] + env: + - name: CARGO_HOME + value: /home/runner/.cargo + - name: RUSTUP_HOME + value: /opt/rustup + - name: RUSTUP_TOOLCHAIN + value: nightly-2026-03-21 + - name: RUSTUP_AUTO_INSTALL + value: "0" + - name: CARGO_TARGET_DIR + value: /home/runner/query-regression-target + - name: QUERY_REGRESSION_CACHE_META + value: /home/runner/query-regression-cache-meta + - name: RUSTC_WRAPPER + value: /usr/local/bin/sccache + - name: SCCACHE_DIR + value: /home/runner/.cache/sccache + - name: SCCACHE_CACHE_SIZE + value: 40G + - name: CARGO_INCREMENTAL + value: "0" securityContext: + runAsNonRoot: true + runAsUser: 1001 + runAsGroup: 1001 allowPrivilegeEscalation: false capabilities: drop: @@ -31,3 +106,21 @@ template: cpu: "8" memory: 24Gi ephemeral-storage: 80Gi + volumeMounts: + - name: cargo-home + mountPath: /home/runner/.cargo + - name: build-cache + mountPath: /home/runner/.cargo/registry + subPath: cargo-registry-v1 + - name: build-cache + mountPath: /home/runner/.cargo/git + subPath: cargo-git-v1 + - name: build-cache + mountPath: /home/runner/query-regression-target + subPath: query-regression-target-v1 + - name: build-cache + mountPath: /home/runner/query-regression-cache-meta + subPath: meta-v1 + - name: build-cache + mountPath: /home/runner/.cache/sccache + subPath: sccache-v1 diff --git a/.github/runner-scale-sets/query-regression/values-paused.yaml b/.github/runner-scale-sets/query-regression/values-paused.yaml new file mode 100644 index 00000000000..fe734ef7615 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/values-paused.yaml @@ -0,0 +1,2 @@ +minRunners: 0 +maxRunners: 0 diff --git a/.github/scripts/query-regression-comment.cjs b/.github/scripts/query-regression-comment.cjs index 5a98e0b8161..6f4386fc3be 100644 --- a/.github/scripts/query-regression-comment.cjs +++ b/.github/scripts/query-regression-comment.cjs @@ -241,8 +241,10 @@ module.exports = async function validateQueryRegressionComment({ github, context '> Rendered by a trusted workflow from JSON artifacts produced by the query-regression run. Results from untrusted PR code are advisory until reviewed.', '', `- **Workflow run:** ${serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${expectedRunId}`, - `- **Base SHA:** \`${text(metadata.base_sha)}\``, + `- **Built base SHA:** \`${text(metadata.built_base_sha)}\``, + `- **Event base SHA:** \`${text(metadata.event_base_sha)}\``, `- **Head SHA:** \`${text(metadata.head_sha)}\``, + `- **Candidate merge SHA:** \`${text(metadata.candidate_sha)}\``, '', ].join('\n'); diff --git a/.github/scripts/query-regression-pr-metadata.py b/.github/scripts/query-regression-pr-metadata.py index 2d8ed3f2e70..ba7ab113b6c 100644 --- a/.github/scripts/query-regression-pr-metadata.py +++ b/.github/scripts/query-regression-pr-metadata.py @@ -31,7 +31,10 @@ def main() -> int: metadata = { "pr_number": int(os.environ["PR_NUMBER"]), "head_sha": os.environ["HEAD_SHA"], - "base_sha": os.environ["BASE_SHA"], + "base_sha": os.environ["EVENT_BASE_SHA"], + "event_base_sha": os.environ["EVENT_BASE_SHA"], + "built_base_sha": os.environ["VERIFIED_BASE_SHA"], + "candidate_sha": os.environ["VERIFIED_CANDIDATE_SHA"], "head_repo": os.environ["HEAD_REPO"], "base_repo": os.environ["BASE_REPO"], "run_id": int(os.environ["RUN_ID"]), diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index 619442e2f95..aa6a7724b24 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -45,10 +45,6 @@ on: permissions: contents: read -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.run_id }} - cancel-in-progress: true - jobs: query-regression: if: >- @@ -59,12 +55,29 @@ jobs: (github.event.action != 'labeled' || github.event.label.name == 'query-regression')) }} runs-on: ${{ github.event_name == 'workflow_dispatch' && inputs.runner || 'perf-regression-8-cores' }} timeout-minutes: 180 + concurrency: + group: query-regression-persistent-cache-v1 + queue: max + cancel-in-progress: false env: CARGO_PROFILE: ${{ github.event_name == 'pull_request' && 'nightly' || inputs.cargo_profile }} - CARGO_TARGET_DIR: ${{ github.workspace }}/query-regression-target - CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" + CARGO_HOME: /home/runner/.cargo + RUSTUP_HOME: /opt/rustup + RUSTUP_TOOLCHAIN: nightly-2026-03-21 + RUSTUP_AUTO_INSTALL: "0" + CARGO_TARGET_DIR: /home/runner/query-regression-target + QUERY_REGRESSION_CACHE_META: /home/runner/query-regression-cache-meta + RUSTC_WRAPPER: /usr/local/bin/sccache + SCCACHE_DIR: /home/runner/.cache/sccache + SCCACHE_CACHE_SIZE: 40G + CARGO_INCREMENTAL: "0" + RUSTFLAGS: -D warnings -C link-arg=-fuse-ld=mold + QUERY_REGRESSION_CACHE_EPOCH: "1" BASE_REF: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || inputs.base_ref }} - CANDIDATE_REF: ${{ github.event_name == 'pull_request' && format('refs/pull/{0}/merge', github.event.pull_request.number) || inputs.candidate_ref || github.ref }} + CANDIDATE_REF: ${{ github.event_name == 'workflow_dispatch' && (inputs.candidate_ref || github.ref) || '' }} + EVENT_MERGE_SHA: ${{ github.event_name == 'pull_request' && github.sha || '' }} + EVENT_HEAD_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || '' }} + EVENT_BASE_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || '' }} steps: - name: Checkout base source uses: actions/checkout@v4 @@ -82,28 +95,338 @@ jobs: cp src/.github/scripts/query-regression-summary.py \ query-regression-trusted-scripts/query-regression-summary.py - - uses: arduino/setup-protoc@v3 - with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + - name: Resolve immutable candidate and build base + working-directory: src + shell: bash + run: | + set -euo pipefail - - uses: actions-rust-lang/setup-rust-toolchain@v1 - with: - cache: false + fail_closed() { + printf '%s\n' \ + "PR changed or merge result unavailable; maintainer must remove and re-add query-regression after reviewing current revision. $1" >&2 + exit 1 + } - - uses: astral-sh/setup-uv@v5 + is_full_sha() { + [[ "$1" =~ ^[[:xdigit:]]{40}$ ]] + } - - name: Rust cache - uses: Swatinem/rust-cache@v2 - with: - workspaces: src - shared-key: query-regression - cache-all-crates: "true" - save-if: false + initially_checked_out_base_sha="$(git rev-parse --verify HEAD)" + is_full_sha "${initially_checked_out_base_sha}" || { + printf 'Could not resolve initially checked-out BASE_REF commit.\n' >&2 + exit 1 + } + + if [[ "${GITHUB_EVENT_NAME}" == "pull_request" ]]; then + for identity in EVENT_MERGE_SHA EVENT_HEAD_SHA EVENT_BASE_SHA; do + value="${!identity-}" + is_full_sha "${value}" || fail_closed "Missing or invalid ${identity}." + done + + if ! git fetch --no-tags --prune --depth=2 origin "${EVENT_MERGE_SHA}"; then + fail_closed "Could not fetch event merge SHA ${EVENT_MERGE_SHA}." + fi + fetched_sha="$(git rev-parse --verify FETCH_HEAD 2>/dev/null)" \ + || fail_closed "Fetched event merge SHA could not be resolved." + [[ "${fetched_sha,,}" == "${EVENT_MERGE_SHA,,}" ]] \ + || fail_closed "FETCH_HEAD ${fetched_sha} does not match event merge SHA ${EVENT_MERGE_SHA}." + + read -r -a parents <<< "$(git show -s --format=%P FETCH_HEAD)" + (( ${#parents[@]} == 2 )) \ + || fail_closed "Event merge SHA ${EVENT_MERGE_SHA} is not a two-parent merge commit." + head_parent_count=0 + for parent in "${parents[@]}"; do + if [[ "${parent,,}" == "${EVENT_HEAD_SHA,,}" ]]; then + ((head_parent_count += 1)) + else + VERIFIED_BASE_SHA="${parent,,}" + fi + done + (( head_parent_count == 1 )) \ + || fail_closed "Event merge SHA must have exactly one EVENT_HEAD_SHA parent." + is_full_sha "${VERIFIED_BASE_SHA-}" \ + || fail_closed "Event merge SHA has no non-head parent for the base build." + VERIFIED_CANDIDATE_SHA="${fetched_sha,,}" + else + git fetch --no-tags --prune --depth=1 origin "${CANDIDATE_REF}" + VERIFIED_CANDIDATE_SHA="$(git rev-parse --verify FETCH_HEAD)" + VERIFIED_BASE_SHA="${initially_checked_out_base_sha}" + fi + + is_full_sha "${VERIFIED_CANDIDATE_SHA}" || { + printf 'Could not resolve candidate SHA.\n' >&2 + exit 1 + } + is_full_sha "${VERIFIED_BASE_SHA}" || { + printf 'Could not resolve build base SHA.\n' >&2 + exit 1 + } + printf 'VERIFIED_CANDIDATE_SHA=%s\n' "${VERIFIED_CANDIDATE_SHA,,}" >> "${GITHUB_ENV}" + printf 'VERIFIED_BASE_SHA=%s\n' "${VERIFIED_BASE_SHA,,}" >> "${GITHUB_ENV}" + printf 'Verified candidate SHA: %s\n' "${VERIFIED_CANDIDATE_SHA,,}" + printf 'Verified build base SHA: %s\n' "${VERIFIED_BASE_SHA,,}" + printf -- '- Verified candidate SHA: `%s`; built base SHA: `%s`\n' \ + "${VERIFIED_CANDIDATE_SHA,,}" "${VERIFIED_BASE_SHA,,}" >> "${GITHUB_STEP_SUMMARY}" + git reset --hard "${VERIFIED_BASE_SHA}" + git clean -ffdx + + - name: Verify runner image tools + shell: bash + run: | + set -euo pipefail + [[ "$(id -u)" == "1001" ]] + [[ "$(id -g)" == "1001" ]] + [[ "$(protoc --version)" == "libprotoc 3.21.12" ]] + [[ "$(uv --version)" =~ ^uv[[:space:]]0\.11\.26([[:space:]]|$) ]] + mold_version="$(mold --version)" + [[ "${mold_version}" =~ ^mold[[:space:]]2\.30\.0([[:space:]]|$) ]] + [[ "$(python3 --version)" == "Python 3.12.3" ]] + sccache_version="$(sccache --version)" + [[ "${sccache_version}" =~ ^sccache[[:space:]]0\.16\.0([[:space:]]|$) ]] + [[ "$(command -v rustup)" == "/opt/cargo/bin/rustup" ]] + [[ "$(command -v cargo)" == "/opt/cargo/bin/cargo" ]] + [[ "$(command -v rustc)" == "/opt/cargo/bin/rustc" ]] + [[ "$(rustup --version)" =~ ^rustup[[:space:]]1\.29\.0([[:space:]]|$) ]] + cargo_version="$(cargo --version)" + [[ "${cargo_version}" =~ ^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] + rustc_version="$(rustc --version)" + [[ "${rustc_version}" =~ ^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] + active_toolchain="$(rustup show active-toolchain)" + [[ "${active_toolchain}" =~ ^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$) ]] + [[ "${RUSTUP_HOME}" == "/opt/rustup" ]] + [[ "${RUSTUP_TOOLCHAIN}" == "nightly-2026-03-21" ]] + [[ "${RUSTUP_AUTO_INSTALL}" == "0" ]] + test -r /opt/rustup && test -x /opt/rustup + test ! -w /opt/rustup + test ! -w /opt/cargo/bin + for entry in config config.toml credentials credentials.toml bin .crates.toml .crates2.json .global-cache .package-cache; do + test ! -e "${CARGO_HOME}/${entry}" + done + mkdir -p "${CARGO_HOME}/registry" "${CARGO_HOME}/git" + + - name: Prepare persistent query regression cache + shell: bash + working-directory: src + run: | + set -euo pipefail + + readonly EXPECTED_CARGO_HOME="/home/runner/.cargo" + readonly EXPECTED_CARGO_REGISTRY="/home/runner/.cargo/registry" + readonly EXPECTED_CARGO_GIT="/home/runner/.cargo/git" + readonly EXPECTED_RUSTUP_HOME="/opt/rustup" + readonly EXPECTED_TARGET_DIR="/home/runner/query-regression-target" + readonly EXPECTED_CACHE_META="/home/runner/query-regression-cache-meta" + readonly EXPECTED_SCCACHE_DIR="/home/runner/.cache/sccache" + readonly EXPECTED_RUSTC_WRAPPER="/usr/local/bin/sccache" + readonly RUNNER_IMAGE_DIGEST="sha256:7022b551771dc82ec4c493f0c3db6ee3050ececd336598690a2165240ebe953e" + readonly RUNNER_IMAGE_EPOCH="3" + + require_expected_root() { + local name="$1" + local actual="$2" + local expected="$3" + [[ "${actual}" == /* && "${actual}" == "${expected}" ]] || { + printf 'Refusing unexpected %s root: %s\n' "${name}" "${actual}" >&2 + exit 1 + } + } + + clear_directory() { + local root="$1" + [[ "${root}" == /* && "${root}" != "/" ]] || { + printf 'Refusing to clear non-absolute or root path: %s\n' "${root}" >&2 + exit 1 + } + case "${root}" in + "${EXPECTED_TARGET_DIR}"|"${EXPECTED_TARGET_DIR}/debug"|"${EXPECTED_TARGET_DIR}/release"|"${EXPECTED_TARGET_DIR}/nightly") ;; + *) + printf 'Refusing to clear unexpected cache path: %s\n' "${root}" >&2 + exit 1 + ;; + esac + find "${root}" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + + } + + clear_cargo_extracted_trees() { + require_expected_root CARGO_REGISTRY "${CARGO_HOME}/registry" "${EXPECTED_CARGO_REGISTRY}" + require_expected_root CARGO_GIT "${CARGO_HOME}/git" "${EXPECTED_CARGO_GIT}" + rm -rf -- "${CARGO_HOME}/registry/src" "${CARGO_HOME}/git/checkouts" + } + + cargo_size_kib() { + du -sk -- "${CARGO_HOME}/registry" "${CARGO_HOME}/git" | awk '{ total += $1 } END { print total }' + } + + target_size_kib() { + du -sk -- "${CARGO_TARGET_DIR}" | cut -f1 + } + + free_kib() { + df -Pk "${CARGO_TARGET_DIR}" | awk 'NR == 2 { print $4 }' + } + + report_cache_usage() { + du -sh -- "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ + "${RUSTUP_HOME}" "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}" + df -P "${CARGO_TARGET_DIR}" + df -hP "${CARGO_TARGET_DIR}" + df -Pi "${CARGO_TARGET_DIR}" + } + + require_expected_root CARGO_HOME "${CARGO_HOME}" "${EXPECTED_CARGO_HOME}" + require_expected_root CARGO_REGISTRY "${CARGO_HOME}/registry" "${EXPECTED_CARGO_REGISTRY}" + require_expected_root CARGO_GIT "${CARGO_HOME}/git" "${EXPECTED_CARGO_GIT}" + require_expected_root RUSTUP_HOME "${RUSTUP_HOME}" "${EXPECTED_RUSTUP_HOME}" + require_expected_root CARGO_TARGET_DIR "${CARGO_TARGET_DIR}" "${EXPECTED_TARGET_DIR}" + require_expected_root QUERY_REGRESSION_CACHE_META "${QUERY_REGRESSION_CACHE_META}" "${EXPECTED_CACHE_META}" + require_expected_root SCCACHE_DIR "${SCCACHE_DIR}" "${EXPECTED_SCCACHE_DIR}" + [[ "${RUSTC_WRAPPER}" == "${EXPECTED_RUSTC_WRAPPER}" ]] || { + printf 'Refusing unexpected RUSTC_WRAPPER: %s\n' "${RUSTC_WRAPPER}" >&2 + exit 1 + } + [[ "${SCCACHE_CACHE_SIZE}" == "40G" ]] || { + printf 'Refusing unexpected SCCACHE_CACHE_SIZE: %s\n' "${SCCACHE_CACHE_SIZE}" >&2 + exit 1 + } + [[ "${CARGO_INCREMENTAL}" == "0" ]] || { + printf 'Refusing unexpected CARGO_INCREMENTAL: %s\n' "${CARGO_INCREMENTAL}" >&2 + exit 1 + } + test -r "${RUSTUP_HOME}" + test -x "${RUSTUP_HOME}" + test ! -w "${RUSTUP_HOME}" + mkdir -p "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ + "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}" + for root in "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ + "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}"; do + test -w "${root}" + touch "${root}/.query-regression-write-test" + rm -f "${root}/.query-regression-write-test" + done + + case "${CARGO_PROFILE}" in + dev) target_profile="debug" ;; + release|nightly) target_profile="${CARGO_PROFILE}" ;; + *) printf 'Unsupported CARGO_PROFILE: %s\n' "${CARGO_PROFILE}" >&2; exit 1 ;; + esac + + rustc_vv="$(rustc -vV)" + if cargo -Vv >/dev/null 2>&1; then + cargo_vv="$(cargo -Vv)" + else + cargo_vv="$(cargo -V)" + fi + host_triple="$(printf '%s\n' "${rustc_vv}" | grep '^host: ' | cut -d' ' -f2-)" + cc_version="$(cc --version)" + mold_version="$(mold --version | grep -m1 .)" + protoc_version="$(protoc --version)" + sccache_version="$(sccache --version)" + rustup_version="$(rustup --version)" + active_toolchain="$(rustup show active-toolchain)" + rustup_path="$(command -v rustup)" + cargo_path="$(command -v cargo)" + rustc_path="$(command -v rustc)" + profile_target_dir="${CARGO_TARGET_DIR}/${target_profile}" + marker="${QUERY_REGRESSION_CACHE_META}/target-${target_profile}.abi" + new_marker="$(mktemp "${QUERY_REGRESSION_CACHE_META}/.${target_profile}.abi.XXXXXX")" + { + printf 'cache_epoch=%s\n' "${QUERY_REGRESSION_CACHE_EPOCH}" + printf 'runner_image_digest=%s\n' "${RUNNER_IMAGE_DIGEST}" + printf 'runner_image_epoch=%s\n' "${RUNNER_IMAGE_EPOCH}" + printf 'host_triple=%s\n' "${host_triple}" + printf 'architecture=%s\n' "$(uname -m)" + printf 'cargo_profile=%s\n' "${CARGO_PROFILE}" + printf 'target_profile=%s\n' "${target_profile}" + printf 'RUSTFLAGS=%s\n' "${RUSTFLAGS-}" + printf 'CARGO_ENCODED_RUSTFLAGS=%s\n' "${CARGO_ENCODED_RUSTFLAGS-}" + printf 'CARGO_INCREMENTAL=%s\n' "${CARGO_INCREMENTAL-}" + printf 'RUSTC_WRAPPER=%s\n' "${RUSTC_WRAPPER-}" + printf 'SCCACHE_DIR=%s\n' "${SCCACHE_DIR-}" + printf 'SCCACHE_CACHE_SIZE=%s\n' "${SCCACHE_CACHE_SIZE-}" + printf 'RUSTUP_HOME=%s\n' "${RUSTUP_HOME-}" + printf 'RUSTUP_TOOLCHAIN=%s\n' "${RUSTUP_TOOLCHAIN-}" + printf 'RUSTUP_AUTO_INSTALL=%s\n' "${RUSTUP_AUTO_INSTALL-}" + printf 'PATH=%s\n' "${PATH-}" + printf '%s\n' "rustc_vV=${rustc_vv}" + printf '%s\n' "cargo_vV=${cargo_vv}" + printf '%s\n' "cc_version=${cc_version}" + printf '%s\n' "mold_version=${mold_version}" + printf '%s\n' "protoc_version=${protoc_version}" + printf '%s\n' "sccache_version=${sccache_version}" + printf '%s\n' "rustup_version=${rustup_version}" + printf '%s\n' "active_toolchain=${active_toolchain}" + printf '%s\n' "rustup_path=${rustup_path}" + printf '%s\n' "cargo_path=${cargo_path}" + printf '%s\n' "rustc_path=${rustc_path}" + } > "${new_marker}" + + if ! cmp -s "${marker}" "${new_marker}"; then + printf 'Cache ABI marker mismatch; invalidating target profile: %s\n' "${target_profile}" + mkdir -p "${profile_target_dir}" + clear_directory "${profile_target_dir}" + else + printf 'Cache ABI marker match; reusing target profile: %s\n' "${target_profile}" + fi + mv -f "${new_marker}" "${marker}" + + report_cache_usage + + target_size="$(target_size_kib)" + if (( target_size >= 400 * 1024 * 1024 )); then + printf 'Warning: target cache is at least 400 GiB (%s KiB)\n' "${target_size}" >&2 + fi + if (( target_size >= 450 * 1024 * 1024 )); then + printf 'Target cache is at least 450 GiB; clearing complete target root\n' >&2 + clear_directory "${CARGO_TARGET_DIR}" + fi + + cargo_size="$(cargo_size_kib)" + if (( cargo_size >= 60 * 1024 * 1024 )); then + printf 'Warning: Cargo cache is at least 60 GiB (%s KiB)\n' "${cargo_size}" >&2 + fi + if (( cargo_size >= 80 * 1024 * 1024 )); then + printf 'Cargo cache is at least 80 GiB; removing extracted sources and checkouts\n' >&2 + clear_cargo_extracted_trees + cargo_size="$(cargo_size_kib)" + if (( cargo_size >= 80 * 1024 * 1024 )); then + printf 'Cargo cache remains at least 80 GiB after cleanup (%s KiB)\n' "${cargo_size}" >&2 + exit 1 + fi + fi + + free_space="$(free_kib)" + if (( free_space < 300 * 1024 * 1024 )); then + printf 'Free space is below 300 GiB; clearing complete target root\n' >&2 + clear_directory "${CARGO_TARGET_DIR}" + free_space="$(free_kib)" + if (( free_space < 300 * 1024 * 1024 )); then + printf 'Free space remains below 300 GiB; removing Cargo extracted sources and checkouts\n' >&2 + clear_cargo_extracted_trees + free_space="$(free_kib)" + if (( free_space < 300 * 1024 * 1024 )); then + printf 'Free space remains below 300 GiB after cleanup (%s KiB)\n' "${free_space}" >&2 + exit 1 + fi + fi + fi + + report_cache_usage + + sccache --start-server + sccache --zero-stats + sccache --show-stats - name: Build base greptime working-directory: src run: | + set -euo pipefail + SECONDS=0 cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime + base_build_elapsed="${SECONDS}" + printf 'Base greptime cargo build elapsed: %s seconds\n' "${base_build_elapsed}" + printf -- '- Base greptime cargo build: %s seconds\n' "${base_build_elapsed}" >> "${GITHUB_STEP_SUMMARY}" + sccache --show-stats + sccache --zero-stats target_dir="${CARGO_PROFILE}" if [[ "${CARGO_PROFILE}" == "dev" ]]; then target_dir="debug" @@ -114,9 +437,18 @@ jobs: - name: Switch source to candidate working-directory: src + shell: bash run: | - git fetch --no-tags --prune --depth=1 origin "${CANDIDATE_REF}" - git reset --hard FETCH_HEAD + set -euo pipefail + [[ "${VERIFIED_CANDIDATE_SHA}" =~ ^[[:xdigit:]]{40}$ ]] || { + printf 'Verified candidate SHA is missing or invalid.\n' >&2 + exit 1 + } + git cat-file -e "${VERIFIED_CANDIDATE_SHA}^{commit}" || { + printf 'Verified candidate SHA is unavailable in the local repository.\n' >&2 + exit 1 + } + git reset --hard "${VERIFIED_CANDIDATE_SHA}" git clean -ffdx - name: Test query regression runner lifecycle @@ -126,8 +458,14 @@ jobs: - name: Build candidate greptime and fixture generators working-directory: src run: | + set -euo pipefail + SECONDS=0 cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime cargo build --profile "${CARGO_PROFILE}" -p cmd --bin query_perf_fixture --features dev-tools + candidate_build_elapsed="${SECONDS}" + printf 'Candidate greptime and fixture generator cargo builds elapsed: %s seconds\n' "${candidate_build_elapsed}" + printf -- '- Candidate greptime and fixture generator cargo builds: %s seconds\n' "${candidate_build_elapsed}" >> "${GITHUB_STEP_SUMMARY}" + sccache --show-stats target_dir="${CARGO_PROFILE}" if [[ "${CARGO_PROFILE}" == "dev" ]]; then target_dir="debug" @@ -192,6 +530,46 @@ jobs: if-no-files-found: warn retention-days: 7 + - name: Report persistent cache usage + if: ${{ always() }} + shell: bash + run: | + set -euo pipefail + for root in \ + /home/runner/.cargo \ + /home/runner/.cargo/registry \ + /home/runner/.cargo/git \ + /opt/rustup \ + /home/runner/query-regression-target \ + /home/runner/query-regression-cache-meta \ + /home/runner/.cache/sccache; do + if [[ -e "${root}" ]]; then + du -sh -- "${root}" + else + printf 'Missing cache path (report only): %s\n' "${root}" + fi + done + if [[ -e /opt/rustup ]]; then + if [[ -w /opt/rustup ]]; then + printf 'Immutable Rustup path is unexpectedly writable: /opt/rustup\n' >&2 + else + printf 'Immutable Rustup path is read-only to runner: /opt/rustup\n' + fi + fi + report_root=/home/runner/query-regression-target + if [[ ! -e "${report_root}" ]]; then + report_root=/home/runner + printf 'Target cache path is missing; reporting its filesystem through %s\n' "${report_root}" + fi + df -P "${report_root}" + df -hP "${report_root}" + df -Pi "${report_root}" + if command -v sccache >/dev/null 2>&1; then + sccache --show-stats || printf 'Unable to show sccache statistics (report only)\n' >&2 + else + printf 'sccache is unavailable (report only)\n' >&2 + fi + - name: Fail on regression failure if: ${{ steps.run.outputs.status != '0' }} run: exit 1 From 1a476f76e235f0c5a311f670fe5544edc98ca77e Mon Sep 17 00:00:00 2001 From: raphaelroshan <49832307+raphaelroshan@users.noreply.github.com> Date: Wed, 15 Jul 2026 09:49:34 +0800 Subject: [PATCH 30/48] fix: avoid panic when negating MIN-valued literals (#8484) Value::try_negative and the temporal negative() helpers negated with raw unary minus, which panics (debug) or wraps (release) on MIN values such as -i64::MIN. try_negative already returns None for the unsigned arms; make the signed and temporal arms honor that contract via checked_neg / checked_negative so a MIN literal produces a clean error instead. Signed-off-by: raphaelroshan --- src/common/time/src/date.rs | 4 ++ src/common/time/src/duration.rs | 5 +++ src/common/time/src/interval.rs | 19 ++++++++++ src/common/time/src/time.rs | 5 +++ src/common/time/src/timestamp.rs | 5 +++ src/datatypes/src/value.rs | 65 ++++++++++++++++++++++++++------ 6 files changed, 92 insertions(+), 11 deletions(-) diff --git a/src/common/time/src/date.rs b/src/common/time/src/date.rs index bda9105f559..38667abbcac 100644 --- a/src/common/time/src/date.rs +++ b/src/common/time/src/date.rs @@ -198,6 +198,10 @@ impl Date { pub fn negative(&self) -> Self { Self(-self.0) } + + pub fn checked_negative(&self) -> Option { + self.0.checked_neg().map(Self) + } } #[cfg(test)] diff --git a/src/common/time/src/duration.rs b/src/common/time/src/duration.rs index a1e582ca82c..b34ced336bd 100644 --- a/src/common/time/src/duration.rs +++ b/src/common/time/src/duration.rs @@ -97,6 +97,11 @@ impl Duration { self.value = -self.value; self } + + pub fn checked_negative(mut self) -> Option { + self.value = self.value.checked_neg()?; + Some(self) + } } /// Convert i64 to Duration Type. diff --git a/src/common/time/src/interval.rs b/src/common/time/src/interval.rs index 4cb8c45c55f..21bdc656d88 100644 --- a/src/common/time/src/interval.rs +++ b/src/common/time/src/interval.rs @@ -80,6 +80,10 @@ impl IntervalYearMonth { Self::new(-self.months) } + pub fn checked_negative(&self) -> Option { + self.months.checked_neg().map(Self::new) + } + pub fn to_iso8601_string(&self) -> String { IntervalFormat::from(*self).to_iso8601_string() } @@ -157,6 +161,13 @@ impl IntervalDayTime { Self::new(-self.days, -self.milliseconds) } + pub fn checked_negative(&self) -> Option { + Some(Self::new( + self.days.checked_neg()?, + self.milliseconds.checked_neg()?, + )) + } + pub fn to_iso8601_string(&self) -> String { IntervalFormat::from(*self).to_iso8601_string() } @@ -278,6 +289,14 @@ impl IntervalMonthDayNano { Self::new(-self.months, -self.days, -self.nanoseconds) } + pub fn checked_negative(&self) -> Option { + Some(Self::new( + self.months.checked_neg()?, + self.days.checked_neg()?, + self.nanoseconds.checked_neg()?, + )) + } + pub fn to_iso8601_string(&self) -> String { IntervalFormat::from(*self).to_iso8601_string() } diff --git a/src/common/time/src/time.rs b/src/common/time/src/time.rs index 00b123f6607..a715e2efec0 100644 --- a/src/common/time/src/time.rs +++ b/src/common/time/src/time.rs @@ -145,6 +145,11 @@ impl Time { self.value = -self.value; self } + + pub fn checked_negative(mut self) -> Option { + self.value = self.value.checked_neg()?; + Some(self) + } } impl From for Time { diff --git a/src/common/time/src/timestamp.rs b/src/common/time/src/timestamp.rs index 120dc8e1e97..8ad38e280fe 100644 --- a/src/common/time/src/timestamp.rs +++ b/src/common/time/src/timestamp.rs @@ -484,6 +484,11 @@ impl Timestamp { self.value = -self.value; self } + + pub fn checked_negative(mut self) -> Option { + self.value = self.value.checked_neg()?; + Some(self) + } } impl Timestamp { diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 5d703480d3f..9d50ac78af4 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -543,20 +543,20 @@ impl Value { None } } - Value::Int8(x) => Some(Value::Int8(-*x)), - Value::Int16(x) => Some(Value::Int16(-*x)), - Value::Int32(x) => Some(Value::Int32(-*x)), - Value::Int64(x) => Some(Value::Int64(-*x)), + Value::Int8(x) => x.checked_neg().map(Value::Int8), + Value::Int16(x) => x.checked_neg().map(Value::Int16), + Value::Int32(x) => x.checked_neg().map(Value::Int32), + Value::Int64(x) => x.checked_neg().map(Value::Int64), Value::Float32(x) => Some(Value::Float32(-*x)), Value::Float64(x) => Some(Value::Float64(-*x)), Value::Decimal128(x) => Some(Value::Decimal128(x.negative())), - Value::Date(x) => Some(Value::Date(x.negative())), - Value::Timestamp(x) => Some(Value::Timestamp(x.negative())), - Value::Time(x) => Some(Value::Time(x.negative())), - Value::Duration(x) => Some(Value::Duration(x.negative())), - Value::IntervalYearMonth(x) => Some(Value::IntervalYearMonth(x.negative())), - Value::IntervalDayTime(x) => Some(Value::IntervalDayTime(x.negative())), - Value::IntervalMonthDayNano(x) => Some(Value::IntervalMonthDayNano(x.negative())), + Value::Date(x) => x.checked_negative().map(Value::Date), + Value::Timestamp(x) => x.checked_negative().map(Value::Timestamp), + Value::Time(x) => x.checked_negative().map(Value::Time), + Value::Duration(x) => x.checked_negative().map(Value::Duration), + Value::IntervalYearMonth(x) => x.checked_negative().map(Value::IntervalYearMonth), + Value::IntervalDayTime(x) => x.checked_negative().map(Value::IntervalDayTime), + Value::IntervalMonthDayNano(x) => x.checked_negative().map(Value::IntervalMonthDayNano), Value::Binary(_) | Value::String(_) @@ -1744,6 +1744,49 @@ pub(crate) mod tests { use crate::types::json_type::{JsonNativeType, JsonObjectType}; use crate::vectors::ListVectorBuilder; + #[test] + fn test_try_negative_overflow() { + // Negating a MIN value overflows, so try_negative returns None instead + // of panicking (consistent with the unsigned arms). + assert_eq!(Value::Int8(i8::MIN).try_negative(), None); + assert_eq!(Value::Int16(i16::MIN).try_negative(), None); + assert_eq!(Value::Int32(i32::MIN).try_negative(), None); + assert_eq!(Value::Int64(i64::MIN).try_negative(), None); + assert_eq!( + Value::Timestamp(Timestamp::new_nanosecond(i64::MIN)).try_negative(), + None + ); + assert_eq!(Value::Date(Date::new(i32::MIN)).try_negative(), None); + assert_eq!( + Value::Time(Time::new_nanosecond(i64::MIN)).try_negative(), + None + ); + assert_eq!( + Value::Duration(Duration::new_nanosecond(i64::MIN)).try_negative(), + None + ); + assert_eq!( + Value::IntervalYearMonth(IntervalYearMonth::new(i32::MIN)).try_negative(), + None + ); + assert_eq!( + Value::IntervalDayTime(IntervalDayTime::new(i32::MIN, i32::MIN)).try_negative(), + None + ); + assert_eq!( + Value::IntervalMonthDayNano(IntervalMonthDayNano::new(i32::MIN, i32::MIN, i64::MIN)) + .try_negative(), + None + ); + + // Non-MIN values still negate. + assert_eq!(Value::Int64(5).try_negative(), Some(Value::Int64(-5))); + assert_eq!( + Value::Timestamp(Timestamp::new_nanosecond(5)).try_negative(), + Some(Value::Timestamp(Timestamp::new_nanosecond(-5))) + ); + } + pub(crate) fn build_struct_type() -> StructType { StructType::new(Arc::new(vec![ StructField::new("id".to_string(), ConcreteDataType::int32_datatype(), false), From 2123108db08711fac45aa8030658353aa2d7a67c Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 15 Jul 2026 09:54:49 +0800 Subject: [PATCH 31/48] fix(promql): preserve ordinary NaN samples (#8494) * fix(promql): distinguish stale markers from NaN Signed-off-by: discord9 * fix(promql): preserve ordinary NaN samples Signed-off-by: discord9 --------- Signed-off-by: discord9 --- Cargo.lock | 1 + src/common/query/src/lib.rs | 1 + src/common/query/src/prometheus.rs | 37 +++ src/promql/Cargo.toml | 1 + .../src/extension_plan/instant_manipulate.rs | 215 +++++++++++++++++- src/promql/src/extension_plan/normalize.rs | 127 ++++++++--- src/promql/src/extension_plan/test_util.rs | 11 +- src/promql/src/functions/aggr_over_time.rs | 72 +++++- .../src/http/result/prometheus_resp.rs | 61 ++++- .../standalone/common/promql/scalar.result | 17 +- .../cases/standalone/common/promql/scalar.sql | 7 +- .../common/tql/aggr_over_time.result | 5 +- .../standalone/common/tql/operator.result | 9 +- .../cases/standalone/common/tql/operator.sql | 3 +- 14 files changed, 507 insertions(+), 60 deletions(-) create mode 100644 src/common/query/src/prometheus.rs diff --git a/Cargo.lock b/Cargo.lock index cb6cb42c2e2..a813b9468a8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11020,6 +11020,7 @@ dependencies = [ "bytemuck", "common-error", "common-macro", + "common-query", "common-recordbatch", "common-telemetry", "criterion 0.7.0", diff --git a/src/common/query/src/lib.rs b/src/common/query/src/lib.rs index bdba49eaf70..e70c4d97e7b 100644 --- a/src/common/query/src/lib.rs +++ b/src/common/query/src/lib.rs @@ -17,6 +17,7 @@ pub mod error; pub mod logical_plan; pub mod native_histogram; pub mod prelude; +pub mod prometheus; pub mod request; pub mod stream; #[cfg(any(test, feature = "testing"))] diff --git a/src/common/query/src/prometheus.rs b/src/common/query/src/prometheus.rs new file mode 100644 index 00000000000..709aa5d7693 --- /dev/null +++ b/src/common/query/src/prometheus.rs @@ -0,0 +1,37 @@ +// 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. + +/// Canonical Prometheus stale-marker NaN bit pattern. +pub const PROMETHEUS_STALE_NAN_BITS: u64 = 0x7ff0_0000_0000_0002; + +/// Returns whether `value` is the canonical Prometheus stale-marker NaN. +#[inline] +pub fn is_prometheus_stale_nan(value: f64) -> bool { + value.to_bits() == PROMETHEUS_STALE_NAN_BITS +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn recognizes_only_the_canonical_stale_marker() { + assert!(is_prometheus_stale_nan(f64::from_bits( + 0x7ff0_0000_0000_0002 + ))); + assert!(!is_prometheus_stale_nan(f64::from_bits( + 0x7ff8_0000_0000_0000 + ))); + } +} diff --git a/src/promql/Cargo.toml b/src/promql/Cargo.toml index 460be8ddd95..36ac213fef4 100644 --- a/src/promql/Cargo.toml +++ b/src/promql/Cargo.toml @@ -13,6 +13,7 @@ async-trait.workspace = true bytemuck.workspace = true common-error.workspace = true common-macro.workspace = true +common-query.workspace = true common-telemetry.workspace = true datafusion.workspace = true datafusion-common.workspace = true diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 7318b9cba5d..1897b453f38 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -18,6 +18,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use common_query::prometheus::is_prometheus_stale_nan; use datafusion::arrow::array::{Array, Float64Array, TimestampMillisecondArray, UInt64Array}; use datafusion::arrow::datatypes::{DataType, SchemaRef}; use datafusion::arrow::record_batch::RecordBatch; @@ -579,9 +580,10 @@ impl InstantManipulateStream { match curr.cmp(&expected_ts) { Ordering::Equal => { if let Some(field_column) = &field_column - && field_column.value(cursor).is_nan() + && field_column.is_valid(cursor) + && is_prometheus_stale_nan(field_column.value(cursor)) { - // ignore the NaN value + // Ignore the stale marker. } else { take_indices.push(cursor as u64); aligned_ts.push(expected_ts); @@ -613,9 +615,10 @@ impl InstantManipulateStream { if prev_ts + self.lookback_delta > expected_ts { // only use the point in the time range if let Some(field_column) = &field_column - && field_column.value(prev_cursor).is_nan() + && field_column.is_valid(prev_cursor) + && is_prometheus_stale_nan(field_column.value(prev_cursor)) { - // if the newest value is NaN, it means the value is stale, so we should not use it + // Do not use a stale marker as the newest value. continue; } // use this point @@ -624,9 +627,10 @@ impl InstantManipulateStream { } } } else if let Some(field_column) = &field_column - && field_column.value(cursor).is_nan() + && field_column.is_valid(cursor) + && is_prometheus_stale_nan(field_column.value(cursor)) { - // if the newest value is NaN, it means the value is stale, so we should not use it + // Do not use a stale marker as the newest value. } else { // use this point take_indices.push(cursor as u64); @@ -688,6 +692,7 @@ fn reuse_constant_column(array: &Arc, len: usize) -> DataFusionResult #[cfg(test)] mod test { + use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::ToDFSchema; use datafusion::datasource::memory::MemorySourceConfig; @@ -697,7 +702,7 @@ mod test { use super::*; use crate::extension_plan::test_util::{ - TIME_INDEX_COLUMN, prepare_test_data, prepare_test_data_with_nan, + TIME_INDEX_COLUMN, prepare_test_data, prepare_test_data_with_stale_marker, }; async fn do_normalize_test( @@ -706,10 +711,10 @@ mod test { lookback_delta: Millisecond, interval: Millisecond, expected: String, - contains_nan: bool, + contains_stale_marker: bool, ) { - let memory_exec = if contains_nan { - Arc::new(prepare_test_data_with_nan()) + let memory_exec = if contains_stale_marker { + Arc::new(prepare_test_data_with_stale_marker()) } else { Arc::new(prepare_test_data()) }; @@ -1255,7 +1260,7 @@ mod test { } #[tokio::test] - async fn lookback_10s_interval_10s_with_nan() { + async fn lookback_10s_interval_10s_with_stale_marker() { let expected = String::from( "+---------------------+-------+\ \n| timestamp | value |\ @@ -1269,7 +1274,7 @@ mod test { } #[tokio::test] - async fn lookback_10s_interval_10s_with_nan_unaligned() { + async fn lookback_10s_interval_10s_with_stale_marker_unaligned() { let expected = String::from( "+-------------------------+-------+\ \n| timestamp | value |\ @@ -1303,4 +1308,190 @@ mod test { ) .await; } + + #[tokio::test] + async fn ordinary_nan_is_selected_for_exact_and_lookback() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000])), + Arc::new(Float64Array::from(vec![f64::NAN])), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let exec = Arc::new(InstantManipulateExec { + start: 1_000, + end: 1_500, + lookback_delta: 1_000, + interval: 500, + time_index_column: TIME_INDEX_COLUMN.to_string(), + field_column: Some("value".to_string()), + reuse_tsid_column: false, + input, + metric: ExecutionPlanMetricsSet::new(), + }); + + let context = SessionContext::default(); + let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) + .await + .unwrap(); + let values = batches + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + let timestamps = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + + assert_eq!(values.len(), 2); + assert_eq!(timestamps, vec![1_000, 1_500]); + assert!(values.iter().all(|value| value.is_nan())); + assert_eq!( + values + .iter() + .map(|value| value.to_bits()) + .collect::>(), + vec![f64::NAN.to_bits(); 2] + ); + } + + #[tokio::test] + async fn prometheus_stale_nan_suppresses_exact_and_lookback() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![500, 1_000])), + Arc::new(Float64Array::from(vec![ + 42.0, + f64::from_bits(0x7ff0_0000_0000_0002), + ])), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let exec = Arc::new(InstantManipulateExec { + start: 1_000, + end: 1_500, + lookback_delta: 1_001, + interval: 500, + time_index_column: TIME_INDEX_COLUMN.to_string(), + field_column: Some("value".to_string()), + reuse_tsid_column: false, + input, + metric: ExecutionPlanMetricsSet::new(), + }); + + let context = SessionContext::default(); + let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) + .await + .unwrap(); + + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "the stale marker must suppress both the exact and lookback selections rather than \ + falling back to 42.0" + ); + } + + #[tokio::test] + async fn null_value_backed_by_stale_bits_is_selected_for_exact_and_lookback() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let field_column = Float64Array::new( + vec![f64::from_bits(0x7ff0_0000_0000_0002)].into(), + Some(NullBuffer::from(vec![false])), + ); + assert!(!field_column.is_valid(0)); + assert_eq!(field_column.value(0).to_bits(), 0x7ff0_0000_0000_0002); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000])), + Arc::new(field_column), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let exec = Arc::new(InstantManipulateExec { + start: 1_000, + end: 1_500, + lookback_delta: 1_000, + interval: 500, + time_index_column: TIME_INDEX_COLUMN.to_string(), + field_column: Some("value".to_string()), + reuse_tsid_column: false, + input, + metric: ExecutionPlanMetricsSet::new(), + }); + + let context = SessionContext::default(); + let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let batch = batches.iter().find(|batch| batch.num_rows() == 2).unwrap(); + let timestamps = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(timestamps.values(), &[1_000, 1_500]); + assert!(!values.is_valid(0)); + assert!(!values.is_valid(1)); + } } diff --git a/src/promql/src/extension_plan/normalize.rs b/src/promql/src/extension_plan/normalize.rs index 6e24f3869b0..a84d5fc5c01 100644 --- a/src/promql/src/extension_plan/normalize.rs +++ b/src/promql/src/extension_plan/normalize.rs @@ -17,7 +17,8 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use datafusion::arrow::array::{BooleanArray, Float64Array}; +use common_query::prometheus::is_prometheus_stale_nan; +use datafusion::arrow::array::{Array, BooleanArray, Float64Array}; use datafusion::arrow::compute; use datafusion::common::{DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics}; use datafusion::error::DataFusionError; @@ -52,12 +53,12 @@ use crate::metrics::PROMQL_SERIES_COUNT; /// Roughly speaking, this method does these things: /// - bias sample's timestamp by offset /// - sort the record batch based on timestamp column -/// - remove NaN values (optional) +/// - remove Prometheus stale markers (optional) #[derive(Debug, PartialEq, Eq, Hash, PartialOrd)] pub struct SeriesNormalize { offset: Millisecond, time_index_column_name: String, - need_filter_out_nan: bool, + filter_stale_markers: bool, tag_columns: Vec, input: LogicalPlan, @@ -122,7 +123,7 @@ impl UserDefinedLogicalNodeCore for SeriesNormalize { write!( f, "PromSeriesNormalize: offset=[{}], time index=[{}], filter NaN: [{}]", - self.offset, self.time_index_column_name, self.need_filter_out_nan + self.offset, self.time_index_column_name, self.filter_stale_markers ) } @@ -158,7 +159,7 @@ impl UserDefinedLogicalNodeCore for SeriesNormalize { Ok(Self { offset: self.offset, time_index_column_name, - need_filter_out_nan: self.need_filter_out_nan, + filter_stale_markers: self.filter_stale_markers, tag_columns, input, unfix: None, @@ -167,7 +168,7 @@ impl UserDefinedLogicalNodeCore for SeriesNormalize { Ok(Self { offset: self.offset, time_index_column_name: self.time_index_column_name.clone(), - need_filter_out_nan: self.need_filter_out_nan, + filter_stale_markers: self.filter_stale_markers, tag_columns: self.tag_columns.clone(), input, unfix: None, @@ -180,14 +181,14 @@ impl SeriesNormalize { pub fn new>( offset: Millisecond, time_index_column_name: N, - need_filter_out_nan: bool, + filter_stale_markers: bool, tag_columns: Vec, input: LogicalPlan, ) -> Self { Self { offset, time_index_column_name: time_index_column_name.as_ref().to_string(), - need_filter_out_nan, + filter_stale_markers, tag_columns, input, unfix: None, @@ -202,7 +203,7 @@ impl SeriesNormalize { Arc::new(SeriesNormalizeExec { offset: self.offset, time_index_column_name: self.time_index_column_name.clone(), - need_filter_out_nan: self.need_filter_out_nan, + filter_stale_markers: self.filter_stale_markers, input: exec_input, tag_columns: self.tag_columns.clone(), metric: ExecutionPlanMetricsSet::new(), @@ -222,7 +223,7 @@ impl SeriesNormalize { pb::SeriesNormalize { offset: self.offset, time_index_idx, - filter_nan: self.need_filter_out_nan, + filter_nan: self.filter_stale_markers, tag_column_indices, ..Default::default() } @@ -244,7 +245,7 @@ impl SeriesNormalize { Ok(Self { offset: pb_normalize.offset, time_index_column_name: String::new(), - need_filter_out_nan: pb_normalize.filter_nan, + filter_stale_markers: pb_normalize.filter_nan, tag_columns: Vec::new(), input: placeholder_plan, unfix: Some(unfix), @@ -256,7 +257,7 @@ impl SeriesNormalize { pub struct SeriesNormalizeExec { offset: Millisecond, time_index_column_name: String, - need_filter_out_nan: bool, + filter_stale_markers: bool, tag_columns: Vec, input: Arc, @@ -303,7 +304,7 @@ impl ExecutionPlan for SeriesNormalizeExec { Ok(Arc::new(Self { offset: self.offset, time_index_column_name: self.time_index_column_name.clone(), - need_filter_out_nan: self.need_filter_out_nan, + filter_stale_markers: self.filter_stale_markers, input: children[0].clone(), tag_columns: self.tag_columns.clone(), metric: self.metric.clone(), @@ -334,7 +335,7 @@ impl ExecutionPlan for SeriesNormalizeExec { Ok(Box::pin(SeriesNormalizeStream { offset: self.offset, time_index, - need_filter_out_nan: self.need_filter_out_nan, + filter_stale_markers: self.filter_stale_markers, schema, input, metric: baseline_metric, @@ -364,7 +365,7 @@ impl DisplayAs for SeriesNormalizeExec { write!( f, "PromSeriesNormalizeExec: offset=[{}], time index=[{}], filter NaN: [{}]", - self.offset, self.time_index_column_name, self.need_filter_out_nan + self.offset, self.time_index_column_name, self.filter_stale_markers ) } } @@ -375,7 +376,7 @@ pub struct SeriesNormalizeStream { offset: Millisecond, // Column index of TIME INDEX column's position in schema time_index: usize, - need_filter_out_nan: bool, + filter_stale_markers: bool, schema: SchemaRef, input: SendableRecordBatchStream, @@ -408,25 +409,25 @@ impl SeriesNormalizeStream { columns[self.time_index] = ts_column_biased; let result_batch = RecordBatch::try_new(input.schema(), columns)?; - if !self.need_filter_out_nan { + if !self.filter_stale_markers { return Ok(result_batch); } - // TODO(ruihang): consider the "special NaN" - // filter out NaN - let mut filter = vec![true; input.num_rows()]; + // Filter out Prometheus stale markers. + let mut stale_marker_filter = vec![true; input.num_rows()]; for column in result_batch.columns() { if let Some(float_column) = column.as_any().downcast_ref::() { - for (i, flag) in filter.iter_mut().enumerate() { - if float_column.value(i).is_nan() { + for (i, flag) in stale_marker_filter.iter_mut().enumerate() { + if float_column.is_valid(i) && is_prometheus_stale_nan(float_column.value(i)) { *flag = false; } } } } - let result = compute::filter_record_batch(&result_batch, &BooleanArray::from(filter)) - .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + let result = + compute::filter_record_batch(&result_batch, &BooleanArray::from(stale_marker_filter)) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; Ok(result) } } @@ -462,6 +463,7 @@ impl Stream for SeriesNormalizeStream { #[cfg(test)] mod test { use datafusion::arrow::array::Float64Array; + use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{ ArrowPrimitiveType, DataType, Field, Schema, TimestampMillisecondType, }; @@ -522,7 +524,7 @@ mod test { let normalize_exec = Arc::new(SeriesNormalizeExec { offset: 0, time_index_column_name: TIME_INDEX_COLUMN.to_string(), - need_filter_out_nan: true, + filter_stale_markers: true, input: memory_exec, tag_columns: vec!["path".to_string()], metric: ExecutionPlanMetricsSet::new(), @@ -556,7 +558,7 @@ mod test { let normalize_exec = Arc::new(SeriesNormalizeExec { offset: 1_000, time_index_column_name: TIME_INDEX_COLUMN.to_string(), - need_filter_out_nan: true, + filter_stale_markers: true, input: memory_exec, metric: ExecutionPlanMetricsSet::new(), tag_columns: vec!["path".to_string()], @@ -583,4 +585,77 @@ mod test { assert_eq!(result_literal, expected); } + + #[tokio::test] + async fn filters_stale_markers_and_preserves_ordinary_nan() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + TimestampMillisecondType::DATA_TYPE, + false, + ), + Field::new("value", DataType::Float64, true), + Field::new("auxiliary", DataType::Float64, true), + ])); + let value_column = Float64Array::new( + vec![ + 42.0, + f64::from_bits(0x7ff0_0000_0000_0002), + 24.0, + f64::from_bits(0x7ff0_0000_0000_0002), + ] + .into(), + Some(NullBuffer::from(vec![true, true, true, false])), + ); + assert!(!value_column.is_valid(3)); + assert_eq!(value_column.value(3).to_bits(), 0x7ff0_0000_0000_0002); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 1_000, 2_000, 3_000, 4_000, + ])), + Arc::new(value_column), + Arc::new(Float64Array::from(vec![ + f64::from_bits(0x7ff8_0000_0000_0000), + 1.0, + f64::from_bits(0x7ff0_0000_0000_0002), + 2.0, + ])), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let exec = Arc::new(SeriesNormalizeExec { + offset: 0, + time_index_column_name: TIME_INDEX_COLUMN.to_string(), + filter_stale_markers: true, + tag_columns: Vec::new(), + input, + metric: ExecutionPlanMetricsSet::new(), + }); + + let context = SessionContext::default(); + let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let batch = batches.iter().find(|batch| batch.num_rows() == 2).unwrap(); + let value = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let auxiliary = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(value.value(0), 42.0); + assert_eq!(auxiliary.value(0).to_bits(), 0x7ff8_0000_0000_0000); + assert!(!value.is_valid(1)); + } } diff --git a/src/promql/src/extension_plan/test_util.rs b/src/promql/src/extension_plan/test_util.rs index 5521d151c21..43d0d686038 100644 --- a/src/promql/src/extension_plan/test_util.rs +++ b/src/promql/src/extension_plan/test_util.rs @@ -16,6 +16,7 @@ use std::sync::Arc; +use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use common_recordbatch::DfRecordBatch as RecordBatch; use datafusion::arrow::array::Float64Array; use datafusion::arrow::datatypes::{ @@ -52,7 +53,7 @@ pub(crate) fn prepare_test_data() -> DataSourceExec { )) } -pub(crate) fn prepare_test_data_with_nan() -> DataSourceExec { +pub(crate) fn prepare_test_data_with_stale_marker() -> DataSourceExec { let schema = Arc::new(Schema::new(vec![ Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true), Field::new("value", DataType::Float64, true), @@ -60,7 +61,13 @@ pub(crate) fn prepare_test_data_with_nan() -> DataSourceExec { let timestamp_column = Arc::new(TimestampMillisecondArray::from(vec![ 0, 30_000, 60_000, 90_000, 120_000, // every 30s ])) as _; - let field_column = Arc::new(Float64Array::from(vec![0.0, f64::NAN, 6.0, f64::NAN, 12.0])) as _; + let field_column = Arc::new(Float64Array::from(vec![ + 0.0, + f64::from_bits(PROMETHEUS_STALE_NAN_BITS), + 6.0, + f64::from_bits(PROMETHEUS_STALE_NAN_BITS), + 12.0, + ])) as _; let data = RecordBatch::try_new(schema.clone(), vec![timestamp_column, field_column]).unwrap(); DataSourceExec::new(Arc::new( diff --git a/src/promql/src/functions/aggr_over_time.rs b/src/promql/src/functions/aggr_over_time.rs index 518c8689eff..05b2b1b2e2d 100644 --- a/src/promql/src/functions/aggr_over_time.rs +++ b/src/promql/src/functions/aggr_over_time.rs @@ -43,7 +43,14 @@ pub fn avg_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Op display_name = prom_min_over_time )] pub fn min_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option { - compute::min(values) + let mut valid_values = values.iter().flatten(); + let mut min = valid_values.next()?; + for value in valid_values { + if value < min || min.is_nan() { + min = value; + } + } + Some(min) } /// The maximum value of all points in the specified interval. @@ -53,7 +60,14 @@ pub fn min_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Op display_name = prom_max_over_time )] pub fn max_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option { - compute::max(values) + let mut valid_values = values.iter().flatten(); + let mut max = valid_values.next()?; + for value in valid_values { + if value > max || max.is_nan() { + max = value; + } + } + Some(max) } /// The sum of all values in the specified interval. @@ -183,6 +197,60 @@ mod test { use super::*; use crate::functions::test_util::simple_range_udf_runner; + fn assert_over_time_value(actual: Option, expected: Option) { + match (actual, expected) { + (Some(actual), Some(expected)) if expected.is_nan() => assert!(actual.is_nan()), + (Some(actual), Some(expected)) => assert_eq!(actual, expected), + (None, None) => {} + (actual, expected) => panic!("expected {expected:?}, got {actual:?}"), + } + } + + fn assert_min_max( + values: Vec>, + expected_min: Option, + expected_max: Option, + ) { + let timestamps = TimestampMillisecondArray::from(vec![0; values.len()]); + let values = Float64Array::from(values); + + assert_over_time_value(min_over_time(×tamps, &values), expected_min); + assert_over_time_value(max_over_time(×tamps, &values), expected_max); + } + + #[test] + fn min_max_over_time_ignore_ordinary_nan_when_finite_values_exist() { + let ordinary_nan = f64::from_bits(0x7ff8_0000_0000_0000); + + assert_min_max( + vec![Some(ordinary_nan), Some(3.0), Some(-2.0)], + Some(-2.0), + Some(3.0), + ); + assert_min_max( + vec![Some(3.0), Some(ordinary_nan), Some(-2.0)], + Some(-2.0), + Some(3.0), + ); + assert_min_max( + vec![Some(-2.0), Some(3.0), Some(ordinary_nan)], + Some(-2.0), + Some(3.0), + ); + assert_min_max( + vec![Some(ordinary_nan), Some(ordinary_nan)], + Some(ordinary_nan), + Some(ordinary_nan), + ); + assert_min_max( + vec![Some(3.0), Some(-2.0), Some(1.0)], + Some(-2.0), + Some(3.0), + ); + assert_min_max(vec![], None, None); + assert_min_max(vec![None, None], None, None); + } + // build timestamp range and value range arrays for test fn build_test_range_arrays() -> (RangeArray, RangeArray) { let ts_array = Arc::new(TimestampMillisecondArray::from_iter( diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index 6f2b1156862..f0d16894363 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -24,6 +24,7 @@ use axum::http::HeaderValue; use axum::response::{IntoResponse, Response}; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; +use common_query::prometheus::is_prometheus_stale_nan; use common_query::{Output, OutputData}; use common_recordbatch::RecordBatches; use datatypes::prelude::ConcreteDataType; @@ -268,8 +269,8 @@ impl PrometheusJsonResponse { // retrieve value if field_column.is_valid(row_index) { let v = field_column.value(row_index); - // ignore all NaN values to reduce the amount of data to be sent. - if v.is_nan() { + // Ignore Prometheus stale markers. + if is_prometheus_stale_nan(v) { continue; } @@ -350,3 +351,59 @@ impl PrometheusJsonResponse { Ok(data) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_recordbatch::{RecordBatch, RecordBatches}; + use datatypes::data_type::ConcreteDataType; + use datatypes::schema::{ColumnSchema, Schema}; + use datatypes::vectors::{Float64Vector, TimestampMillisecondVector}; + + use super::*; + + #[test] + fn matrix_response_preserves_ordinary_nan_and_filters_stale_markers() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ])); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![ + 1_000, 2_000, 3_000, 4_000, + ])) as _, + Arc::new(Float64Vector::from(vec![ + Some(1.0), + Some(f64::from_bits(0x7ff8_0000_0000_0000)), + Some(f64::from_bits(0x7ff0_0000_0000_0002)), + None, + ])) as _, + ], + ) + .unwrap(); + let batches = RecordBatches::try_new(schema, vec![batch]).unwrap(); + + let response = + PrometheusJsonResponse::record_batches_to_data(batches, None, ValueType::Matrix) + .unwrap(); + let PrometheusResponse::PromData(data) = response else { + panic!("expected Prometheus data response"); + }; + let PromQueryResult::Matrix(series) = data.result else { + panic!("expected matrix result"); + }; + + assert_eq!(series.len(), 1); + assert_eq!( + series[0].values, + vec![(1.0, "1".to_string()), (2.0, "NaN".to_string())] + ); + } +} diff --git a/tests/cases/standalone/common/promql/scalar.result b/tests/cases/standalone/common/promql/scalar.result index c3292b4f5c5..ec1c768f693 100644 --- a/tests/cases/standalone/common/promql/scalar.result +++ b/tests/cases/standalone/common/promql/scalar.result @@ -586,8 +586,8 @@ INSERT INTO TABLE presence_metric VALUES Affected Rows: 15 --- NaN drops `cpu0` from the grouped count, while the NULL sample on `cpu2` --- still leaves a zero-valued row in `count(...) by (cpu)`. +-- Ordinary NaN keeps `cpu0` present in the grouped count, while the NULL sample +-- on `cpu2` still leaves a zero-valued row in `count(...) by (cpu)`. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 600, '200s') count(presence_metric{instance="i1"}) by (cpu); @@ -595,6 +595,8 @@ TQL EVAL (0, 600, '200s') count(presence_metric{instance="i1"}) by (cpu); | cpu | ts | count(presence_metric.val) | +------+---------------------+----------------------------+ | cpu0 | 1970-01-01T00:00:00 | 2 | +| cpu0 | 1970-01-01T00:03:20 | 2 | +| cpu0 | 1970-01-01T00:06:40 | 2 | | cpu0 | 1970-01-01T00:10:00 | 2 | | cpu1 | 1970-01-01T00:00:00 | 1 | | cpu1 | 1970-01-01T00:03:20 | 1 | @@ -605,7 +607,8 @@ TQL EVAL (0, 600, '200s') count(presence_metric{instance="i1"}) by (cpu); | cpu2 | 1970-01-01T00:06:40 | 0 | +------+---------------------+----------------------------+ --- Nested-count rewrite should preserve grouped presence after stale-NaN filtering and null-value pruning. +-- Nested-count rewrite should preserve grouped presence after stale-marker filtering +-- and null-value pruning. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 600, '200s') scalar(count(count(presence_metric{instance="i1"}) by (cpu))); @@ -613,8 +616,8 @@ TQL EVAL (0, 600, '200s') scalar(count(count(presence_metric{instance="i1"}) by | ts | scalar(count(count(presence_metric.val))) | +---------------------+-------------------------------------------+ | 1970-01-01T00:00:00 | 3.0 | -| 1970-01-01T00:03:20 | 2.0 | -| 1970-01-01T00:06:40 | 2.0 | +| 1970-01-01T00:03:20 | 3.0 | +| 1970-01-01T00:06:40 | 3.0 | | 1970-01-01T00:10:00 | 2.0 | +---------------------+-------------------------------------------+ @@ -626,8 +629,8 @@ TQL EVAL (0, 600, '200s') scalar(count(sum(presence_metric{instance="i1"}) by (c | ts | scalar(count(sum(presence_metric.val))) | +---------------------+-----------------------------------------+ | 1970-01-01T00:00:00 | 3.0 | -| 1970-01-01T00:03:20 | 1.0 | -| 1970-01-01T00:06:40 | 1.0 | +| 1970-01-01T00:03:20 | 2.0 | +| 1970-01-01T00:06:40 | 2.0 | | 1970-01-01T00:10:00 | 2.0 | +---------------------+-----------------------------------------+ diff --git a/tests/cases/standalone/common/promql/scalar.sql b/tests/cases/standalone/common/promql/scalar.sql index 662f9665fee..b0c85d89e35 100644 --- a/tests/cases/standalone/common/promql/scalar.sql +++ b/tests/cases/standalone/common/promql/scalar.sql @@ -186,12 +186,13 @@ INSERT INTO TABLE presence_metric VALUES (600000, 'i1', 'cpu0', 'b', 8.0), (600000, 'i2', 'cpu9', 'a', 103.0); --- NaN drops `cpu0` from the grouped count, while the NULL sample on `cpu2` --- still leaves a zero-valued row in `count(...) by (cpu)`. +-- Ordinary NaN keeps `cpu0` present in the grouped count, while the NULL sample +-- on `cpu2` still leaves a zero-valued row in `count(...) by (cpu)`. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 600, '200s') count(presence_metric{instance="i1"}) by (cpu); --- Nested-count rewrite should preserve grouped presence after stale-NaN filtering and null-value pruning. +-- Nested-count rewrite should preserve grouped presence after stale-marker filtering +-- and null-value pruning. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 600, '200s') scalar(count(count(presence_metric{instance="i1"}) by (cpu))); diff --git a/tests/cases/standalone/common/tql/aggr_over_time.result b/tests/cases/standalone/common/tql/aggr_over_time.result index c846d73c051..e39c9d52384 100644 --- a/tests/cases/standalone/common/tql/aggr_over_time.result +++ b/tests/cases/standalone/common/tql/aggr_over_time.result @@ -295,6 +295,7 @@ tql eval (60, 60, '1s') min_over_time(data[2m]); | 1970-01-01T00:01:00 | 0.0 | some_nan | | 1970-01-01T00:01:00 | 0.0 | some_nan3 | | 1970-01-01T00:01:00 | 1.0 | some_nan2 | +| 1970-01-01T00:01:00 | NaN | only_nan | +---------------------+----------------------------------+-----------+ -- eval instant at 1m max_over_time(data[2m]) @@ -313,6 +314,7 @@ tql eval (60, 60, '1s') max_over_time(data[2m]); | 1970-01-01T00:01:00 | 2.0 | some_nan | | 1970-01-01T00:01:00 | 2.0 | some_nan2 | | 1970-01-01T00:01:00 | 3.0 | numbers | +| 1970-01-01T00:01:00 | NaN | only_nan | +---------------------+----------------------------------+-----------+ -- eval instant at 1m last_over_time(data[2m]) @@ -327,10 +329,11 @@ tql eval (60, 60, '1s') last_over_time(data[2m]); +---------------------+-----------------------------------+-----------+ | ts | prom_last_over_time(ts_range,val) | ty | +---------------------+-----------------------------------+-----------+ -| 1970-01-01T00:01:00 | 0.0 | some_nan | | 1970-01-01T00:01:00 | 1.0 | some_nan2 | | 1970-01-01T00:01:00 | 1.0 | some_nan3 | | 1970-01-01T00:01:00 | 3.0 | numbers | +| 1970-01-01T00:01:00 | NaN | only_nan | +| 1970-01-01T00:01:00 | NaN | some_nan | +---------------------+-----------------------------------+-----------+ drop table data; diff --git a/tests/cases/standalone/common/tql/operator.result b/tests/cases/standalone/common/tql/operator.result index bc105b14549..bc0971c3f3b 100644 --- a/tests/cases/standalone/common/tql/operator.result +++ b/tests/cases/standalone/common/tql/operator.result @@ -40,11 +40,14 @@ tql eval (60, 60, '1s') trigy atan2 trigx; -- eval instant at 1m trigy atan2 trigNaN -- trigy{} NaN --- This query doesn't have result because `trignan` is NaN and will be filtered out. +-- Ordinary NaN is a valid sample and propagates through `atan2`. tql eval (60, 60, '1s') trigy atan2 trignan; -++ -++ ++---------------------+------------------------------+ +| ts | atan2(trigy.val,trignan.val) | ++---------------------+------------------------------+ +| 1970-01-01T00:01:00 | NaN | ++---------------------+------------------------------+ -- eval instant at 1m 10 atan2 20 -- 0.4636476090008061 diff --git a/tests/cases/standalone/common/tql/operator.sql b/tests/cases/standalone/common/tql/operator.sql index 2ab77dceb9d..8d10f62bcc7 100644 --- a/tests/cases/standalone/common/tql/operator.sql +++ b/tests/cases/standalone/common/tql/operator.sql @@ -23,7 +23,7 @@ tql eval (60, 60, '1s') trigy atan2 trigx; -- eval instant at 1m trigy atan2 trigNaN -- trigy{} NaN --- This query doesn't have result because `trignan` is NaN and will be filtered out. +-- Ordinary NaN is a valid sample and propagates through `atan2`. tql eval (60, 60, '1s') trigy atan2 trignan; -- eval instant at 1m 10 atan2 20 @@ -73,4 +73,3 @@ VALUES tql eval (1743465600.5, 1743465610, '1s') irate(t[2s]); drop table t; - From 623145e63593bf9204a4f37719c8fb4715080aff Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 15 Jul 2026 09:59:24 +0800 Subject: [PATCH 32/48] fix(query): validate DistAnalyzeExec child count (#8510) * fix(query): validate DistAnalyzeExec child count Signed-off-by: discord9 * test(query): avoid implicit clone lint Signed-off-by: discord9 --------- Signed-off-by: discord9 --- src/query/src/analyze.rs | 77 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 75 insertions(+), 2 deletions(-) diff --git a/src/query/src/analyze.rs b/src/query/src/analyze.rs index 6f6864ef2b9..e71a10e2f20 100644 --- a/src/query/src/analyze.rs +++ b/src/query/src/analyze.rs @@ -33,7 +33,7 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, accept, }; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; -use datafusion_common::{DataFusionError, internal_err}; +use datafusion_common::{DataFusionError, assert_eq_or_internal_err, internal_err}; use datafusion_physical_expr::{Distribution, EquivalenceProperties, Partitioning}; use futures::StreamExt; use serde::Serialize; @@ -172,8 +172,13 @@ impl ExecutionPlan for DistAnalyzeExec { self: Arc, mut children: Vec>, ) -> DfResult> { + assert_eq_or_internal_err!( + children.len(), + 1, + "DistAnalyzeExec requires exactly one child" + ); Ok(Arc::new(Self::new( - children.pop().unwrap(), + children.swap_remove(0), self.verbose, self.format, ))) @@ -381,3 +386,71 @@ impl Display for JsonMetrics { write!(f, "{}", serde_json::to_string(self).unwrap()) } } + +#[cfg(test)] +mod tests { + use datafusion::physical_plan::empty::EmptyExec; + + use super::*; + + fn empty_plan(name: &str) -> Arc { + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( + name, + DataType::Utf8, + true, + )])))) + } + + #[test] + fn qbs_dist_analyze_rejects_zero_children() { + let analyze = Arc::new(DistAnalyzeExec::new( + empty_plan("original"), + false, + AnalyzeFormat::TEXT, + )); + + assert!(ExecutionPlan::with_new_children(analyze, vec![]).is_err()); + } + + #[test] + fn qbs_dist_analyze_rejects_multiple_children() { + let analyze = Arc::new(DistAnalyzeExec::new( + empty_plan("original"), + false, + AnalyzeFormat::TEXT, + )); + + let result = ExecutionPlan::with_new_children( + analyze, + vec![empty_plan("first"), empty_plan("second")], + ); + + if let Ok(plan) = result { + let retained = plan + .as_any() + .downcast_ref::() + .unwrap() + .input() + .schema() + .field(0) + .name() + .clone(); + panic!("expected an arity error for multiple children, but retained `{retained}`"); + } + } + + #[test] + fn qbs_dist_analyze_accepts_exactly_one_child() { + let analyze = Arc::new(DistAnalyzeExec::new( + empty_plan("original"), + false, + AnalyzeFormat::TEXT, + )); + let replacement = empty_plan("replacement"); + + let rebuilt = ExecutionPlan::with_new_children(analyze, vec![replacement]).unwrap(); + let rebuilt = rebuilt.as_any().downcast_ref::().unwrap(); + + assert_eq!(rebuilt.input().schema().field(0).name(), "replacement"); + } +} From d48f2961444e56f7c3d22761f61a6315954eaa90 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 15 Jul 2026 10:05:15 +0800 Subject: [PATCH 33/48] fix: reject out-of-range PostgreSQL numeric UInt64 (#8517) Signed-off-by: discord9 --- src/servers/src/postgres/types.rs | 138 +++++++++++++++++++++++++++--- tests-integration/tests/sql.rs | 69 +++++++++++++++ 2 files changed, 197 insertions(+), 10 deletions(-) diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index 74e230f83af..b282bdea68a 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -582,9 +582,10 @@ pub(super) fn parameters_to_scalar_values( Some(st @ ConcreteDataType::Timestamp(unit)) => { to_timestamp_scalar_value(data.and_then(|n| n.to_i64()), unit, st)? } - Some(ConcreteDataType::UInt64(_)) | None => { - ScalarValue::UInt64(data.and_then(|n| n.to_u64())) - } + Some(ConcreteDataType::UInt64(_)) | None => ScalarValue::UInt64( + data.map(|n| n.to_u64().ok_or_else(|| numeric_out_of_range_error(n))) + .transpose()?, + ), Some(st) => { return Err(invalid_parameter_error( "invalid_parameter_type", @@ -924,21 +925,28 @@ pub(super) fn parameters_to_scalar_values( &Type::NUMERIC_ARRAY => { let data = portal.parameter::>>(idx, &client_type)?; if let Some(data) = data { - let build_u64_list = |data: Vec>| { + let build_u64_list = |data: Vec>| -> PgWireResult { let values = data .into_iter() - .map(|n| ScalarValue::UInt64(n.and_then(|n| n.to_u64()))) - .collect::>(); - ScalarValue::List(ScalarValue::new_list( + .map(|n| { + Ok(ScalarValue::UInt64( + n.map(|n| { + n.to_u64().ok_or_else(|| numeric_out_of_range_error(n)) + }) + .transpose()?, + )) + }) + .collect::>>()?; + Ok(ScalarValue::List(ScalarValue::new_list( &values, &ArrowDataType::UInt64, true, - )) + ))) }; if let Some(server_type) = &server_type { match server_type { ConcreteDataType::List(list_type) => match list_type.item_type() { - ConcreteDataType::UInt64(_) => build_u64_list(data), + ConcreteDataType::UInt64(_) => build_u64_list(data)?, ConcreteDataType::Decimal128(dt) => { let values = data .into_iter() @@ -975,7 +983,7 @@ pub(super) fn parameters_to_scalar_values( } } else { // server type not provided - build_u64_list(data) + build_u64_list(data)? } } else { ScalarValue::Null @@ -2025,4 +2033,114 @@ mod test { let values = parameters_to_scalar_values(&plan, &portal).unwrap(); assert_eq!(values[0], ScalarValue::Int8(None)); } + + fn numeric_uint64_array_plan() -> LogicalPlan { + build_plan_with_params(vec![( + "$1", + DataType::List(Arc::new(Field::new("item", DataType::UInt64, true))), + )]) + } + + fn assert_numeric_out_of_range(result: PgWireResult>) { + match result.unwrap_err() { + PgWireError::UserError(error) => { + assert_eq!("22023", error.code); + assert_eq!("numeric_value_out_of_range", error.message); + } + error => panic!("expected numeric out-of-range error, got {error:?}"), + } + } + + #[test] + fn test_numeric_uint64_scalar_negative_rejected() { + let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]); + let portal = make_portal(vec![Some(Type::NUMERIC)], vec![s("-1")]); + + assert_numeric_out_of_range(parameters_to_scalar_values(&plan, &portal)); + } + + #[test] + fn test_numeric_uint64_scalar_above_u64_max_rejected() { + let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]); + let portal = make_portal(vec![Some(Type::NUMERIC)], vec![s("18446744073709551616")]); + + assert_numeric_out_of_range(parameters_to_scalar_values(&plan, &portal)); + } + + #[test] + fn test_numeric_uint64_scalar_null_preserved() { + let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]); + let portal = make_portal(vec![Some(Type::NUMERIC)], vec![None]); + + let values = parameters_to_scalar_values(&plan, &portal).unwrap(); + assert_eq!(ScalarValue::UInt64(None), values[0]); + } + + #[test] + fn test_numeric_uint64_scalar_u64_max_preserved() { + let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]); + let portal = make_portal(vec![Some(Type::NUMERIC)], vec![s(&u64::MAX.to_string())]); + + let values = parameters_to_scalar_values(&plan, &portal).unwrap(); + assert_eq!(ScalarValue::UInt64(Some(u64::MAX)), values[0]); + } + + #[test] + fn test_numeric_uint64_array_outer_null_preserved() { + let portal = make_portal(vec![Some(Type::NUMERIC_ARRAY)], vec![None]); + + let values = parameters_to_scalar_values(&numeric_uint64_array_plan(), &portal).unwrap(); + assert_eq!(ScalarValue::Null, values[0]); + } + + #[test] + fn test_numeric_uint64_array_preserves_values_and_null_slots() { + let portal = make_portal( + vec![Some(Type::NUMERIC_ARRAY)], + vec![s("{42,NULL,18446744073709551615}")], + ); + + let values = parameters_to_scalar_values(&numeric_uint64_array_plan(), &portal).unwrap(); + let expected = ScalarValue::List(ScalarValue::new_list( + &[ + ScalarValue::UInt64(Some(42)), + ScalarValue::UInt64(None), + ScalarValue::UInt64(Some(u64::MAX)), + ], + &ArrowDataType::UInt64, + true, + )); + assert_eq!(expected, values[0]); + } + + #[test] + fn test_numeric_uint64_array_invalid_after_valid_and_null_prefix_rejected() { + let portal = make_portal(vec![Some(Type::NUMERIC_ARRAY)], vec![s("{42,NULL,-1}")]); + + assert_numeric_out_of_range(parameters_to_scalar_values( + &numeric_uint64_array_plan(), + &portal, + )); + } + + #[test] + fn test_numeric_uint64_array_above_u64_max_rejected() { + let portal = make_portal( + vec![Some(Type::NUMERIC_ARRAY)], + vec![s("{18446744073709551616}")], + ); + + assert_numeric_out_of_range(parameters_to_scalar_values( + &numeric_uint64_array_plan(), + &portal, + )); + } + + #[test] + fn test_numeric_uint64_uninferred_array_invalid_value_rejected() { + let plan = LogicalPlanBuilder::empty(true).build().unwrap(); + let portal = make_portal(vec![Some(Type::NUMERIC_ARRAY)], vec![s("{-1}")]); + + assert_numeric_out_of_range(parameters_to_scalar_values(&plan, &portal)); + } } diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 156de7c415c..0e477061bb7 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashMap; +use std::str::FromStr; use std::time::Duration; use auth::user_provider_from_option; @@ -1353,6 +1354,68 @@ pub async fn test_postgres_uint64_parameter(store_type: StorageType) { let count: i64 = rows[0].get(0); assert_eq!(count, 1); + let scalar_statement = client + .prepare("SELECT arrow_cast($1, 'UInt64')") + .await + .unwrap(); + let scalar_null_statement = client + .prepare("SELECT arrow_cast($1, 'UInt64') IS NULL") + .await + .unwrap(); + let array_statement = client + .prepare("SELECT arrow_cast($1, 'List(UInt64)')") + .await + .unwrap(); + let array_null_statement = client + .prepare("SELECT arrow_cast($1, 'List(UInt64)') IS NULL") + .await + .unwrap(); + + let ordinary = Decimal::from(42); + let max = Decimal::from_str("18446744073709551615").unwrap(); + for value in [&ordinary, &max] { + let row = client.query_one(&scalar_statement, &[value]).await.unwrap(); + assert_eq!(*value, row.get::<_, Decimal>(0)); + } + + let scalar_null: Option = None; + let row = client + .query_one(&scalar_null_statement, &[&scalar_null]) + .await + .unwrap(); + assert!(row.get::<_, bool>(0)); + + let array_values = vec![Some(ordinary), None, Some(max)]; + let row = client + .query_one(&array_statement, &[&array_values]) + .await + .unwrap(); + assert_eq!(array_values, row.get::<_, Vec>>(0)); + + let array_null: Option>> = None; + let row = client + .query_one(&array_null_statement, &[&array_null]) + .await + .unwrap(); + assert!(row.get::<_, bool>(0)); + + for value in [ + Decimal::from(-1), + Decimal::from_str("18446744073709551616").unwrap(), + ] { + let error = client + .query(&scalar_statement, &[&value]) + .await + .unwrap_err(); + assert_pg_numeric_range_error(error); + + let error = client + .query(&array_statement, &[&vec![Some(value)]]) + .await + .unwrap_err(); + assert_pg_numeric_range_error(error); + } + drop(client); rx.await.unwrap(); @@ -1360,6 +1423,12 @@ pub async fn test_postgres_uint64_parameter(store_type: StorageType) { guard.remove_all().await; } +fn assert_pg_numeric_range_error(error: tokio_postgres::Error) { + let error = error.as_db_error().expect("expected PostgreSQL user error"); + assert_eq!("22023", error.code().code()); + assert_eq!("numeric_value_out_of_range", error.message()); +} + pub async fn test_postgres_explain_bind_parameter(store_type: StorageType) { // Regression test for #8029: EXPLAIN / EXPLAIN ANALYZE must accept bind // parameters over the Postgres extended query protocol. From 6e8bd8a5f989a6e45650362b622a7e027977149f Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 15 Jul 2026 11:01:35 +0800 Subject: [PATCH 34/48] test: add legacy JSON2 non-append compat case (#8513) * test: add legacy json2 non-append compat case Signed-off-by: discord9 * Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --------- Signed-off-by: discord9 Co-authored-by: dennis zhuang Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../legacy_json2_non_append_table/case.toml | 8 +++ .../legacy_json2_non_append_table/setup.sql | 10 +++ .../verify.result | 63 +++++++++++++++++++ .../legacy_json2_non_append_table/verify.sql | 16 +++++ 4 files changed, 97 insertions(+) create mode 100644 tests/compatibility/cases/legacy_json2_non_append_table/case.toml create mode 100644 tests/compatibility/cases/legacy_json2_non_append_table/setup.sql create mode 100644 tests/compatibility/cases/legacy_json2_non_append_table/verify.result create mode 100644 tests/compatibility/cases/legacy_json2_non_append_table/verify.sql diff --git a/tests/compatibility/cases/legacy_json2_non_append_table/case.toml b/tests/compatibility/cases/legacy_json2_non_append_table/case.toml new file mode 100644 index 00000000000..8bb090ea788 --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_non_append_table/case.toml @@ -0,0 +1,8 @@ +name = "legacy_json2_non_append_table" +reason = "Verify current binaries can read a non-append JSON2 table created by old binaries and migrate it to append_mode=true." +introduced_by = "PR #8434" +topologies = ["distributed"] +from_range = ["=v1.1.0", "=v1.1.1", "=v1.1.2"] +to_range = [">v1.1.2"] +features = ["json", "json2", "table", "append_mode", "read"] +owner = "query" diff --git a/tests/compatibility/cases/legacy_json2_non_append_table/setup.sql b/tests/compatibility/cases/legacy_json2_non_append_table/setup.sql new file mode 100644 index 00000000000..6eeeed1f908 --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_non_append_table/setup.sql @@ -0,0 +1,10 @@ +CREATE TABLE t_legacy_json2_non_append_table ( + ts TIMESTAMP TIME INDEX, + j JSON2 +); + +INSERT INTO t_legacy_json2_non_append_table (ts, j) VALUES + ('2026-07-08 00:00:00+0000', '{"a": 1, "nested": {"s": "old-1"}}'), + ('2026-07-08 00:01:00+0000', '{"a": 2, "nested": {"s": "old-2"}}'); + +ADMIN FLUSH_TABLE('t_legacy_json2_non_append_table'); diff --git a/tests/compatibility/cases/legacy_json2_non_append_table/verify.result b/tests/compatibility/cases/legacy_json2_non_append_table/verify.result new file mode 100644 index 00000000000..2cc3c230e2b --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_non_append_table/verify.result @@ -0,0 +1,63 @@ +SELECT ts, j.a AS a, j.nested.s AS nested_s +FROM t_legacy_json2_non_append_table +ORDER BY ts; + ++---------------------+---+----------+ +| ts | a | nested_s | ++---------------------+---+----------+ +| 2026-07-08T00:00:00 | 1 | old-1 | +| 2026-07-08T00:01:00 | 2 | old-2 | ++---------------------+---+----------+ + +SHOW CREATE TABLE t_legacy_json2_non_append_table; + ++---------------------------------+----------------------------------------------------------------+ +| Table | Create Table | ++---------------------------------+----------------------------------------------------------------+ +| t_legacy_json2_non_append_table | CREATE TABLE IF NOT EXISTS "t_legacy_json2_non_append_table" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2 NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++---------------------------------+----------------------------------------------------------------+ + +ALTER TABLE t_legacy_json2_non_append_table SET 'append_mode' = 'true'; + +Affected Rows: 0 + +SHOW CREATE TABLE t_legacy_json2_non_append_table; + ++---------------------------------+----------------------------------------------------------------+ +| Table | Create Table | ++---------------------------------+----------------------------------------------------------------+ +| t_legacy_json2_non_append_table | CREATE TABLE IF NOT EXISTS "t_legacy_json2_non_append_table" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2 NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++---------------------------------+----------------------------------------------------------------+ + +INSERT INTO t_legacy_json2_non_append_table (ts, j) VALUES + ('2026-07-08 00:02:00+0000', '{"a": 3, "nested": {"s": "current-3"}}'); + +Affected Rows: 1 + +SELECT ts, j.a AS a, j.nested.s AS nested_s +FROM t_legacy_json2_non_append_table +ORDER BY ts; + ++---------------------+---+-----------+ +| ts | a | nested_s | ++---------------------+---+-----------+ +| 2026-07-08T00:00:00 | 1 | old-1 | +| 2026-07-08T00:01:00 | 2 | old-2 | +| 2026-07-08T00:02:00 | 3 | current-3 | ++---------------------+---+-----------+ diff --git a/tests/compatibility/cases/legacy_json2_non_append_table/verify.sql b/tests/compatibility/cases/legacy_json2_non_append_table/verify.sql new file mode 100644 index 00000000000..2dbc7fc729e --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_non_append_table/verify.sql @@ -0,0 +1,16 @@ +SELECT ts, j.a AS a, j.nested.s AS nested_s +FROM t_legacy_json2_non_append_table +ORDER BY ts; + +SHOW CREATE TABLE t_legacy_json2_non_append_table; + +ALTER TABLE t_legacy_json2_non_append_table SET 'append_mode' = 'true'; + +SHOW CREATE TABLE t_legacy_json2_non_append_table; + +INSERT INTO t_legacy_json2_non_append_table (ts, j) VALUES + ('2026-07-08 00:02:00+0000', '{"a": 3, "nested": {"s": "current-3"}}'); + +SELECT ts, j.a AS a, j.nested.s AS nested_s +FROM t_legacy_json2_non_append_table +ORDER BY ts; From 912db22417f2d0d8b5fdc369a67a6026f682bb4a Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 15 Jul 2026 11:04:11 +0800 Subject: [PATCH 35/48] fix(log-store): deduplicate Kafka WAL multipart records (#8514) * fix(log-store): deduplicate Kafka WAL multipart records Signed-off-by: WenyXu * fix(log-store): preserve delayed Kafka WAL entry offsets Signed-off-by: WenyXu * fix(log-store): reject conflicting Kafka WAL last records Signed-off-by: WenyXu * fix(log-store): emit Kafka WAL entries on last record Signed-off-by: WenyXu * fix(log-store): discard duplicate Kafka WAL first records Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/log-store/src/kafka/util/record.rs | 166 ++++++++++++++++++++----- 1 file changed, 137 insertions(+), 29 deletions(-) diff --git a/src/log-store/src/kafka/util/record.rs b/src/log-store/src/kafka/util/record.rs index 720f989139d..aefb098e0db 100644 --- a/src/log-store/src/kafka/util/record.rs +++ b/src/log-store/src/kafka/util/record.rs @@ -227,13 +227,17 @@ pub fn remaining_entries( /// - Emits a [RecordType::Full] type record immediately. /// /// For type of [Entry::MultiplePart] Entry: -/// - Emits a complete or incomplete [Entry] while the next same [RegionId] record arrives. +/// - Emits a complete [Entry] immediately when a [RecordType::Last] record arrives with +/// buffered parts. +/// - Emits an incomplete [Entry] when a new [RecordType::First] record replaces +/// buffered incomplete records for the same [RegionId]. /// -/// **Incomplete Entry:** -/// If the records arrive in the following order, it emits **the incomplete [Entry]** when the next record arrives. -/// - **[RecordType::First], [RecordType::Middle]**, [RecordType::First] -/// - **[RecordType::Middle]**, [RecordType::First] -/// - **[RecordType::Last]** +/// **Discarded records:** +/// - A standalone [RecordType::First] record is discarded when another [RecordType::First] +/// record with the same payload for the same [RegionId] arrives. +/// - A [RecordType::Last] record without buffered parts is discarded. +/// +/// A trailing incomplete entry is emitted by [`remaining_entries`]. pub(crate) fn maybe_emit_entry( provider: &Arc, record: Record, @@ -244,7 +248,16 @@ pub(crate) fn maybe_emit_entry( RecordType::Full => entry = Some(convert_to_naive_entry(provider.clone(), record)), RecordType::First => { let region_id = record.meta.ns.region_id.into(); + let duplicate = buffered_records.get(®ion_id).is_some_and(|records| { + records.len() == 1 + && records[0].meta.tp == RecordType::First + && records[0].data == record.data + }); if let Some(records) = buffered_records.insert(region_id, vec![record]) { + if duplicate { + // A duplicate standalone First cannot form a complete entry. + return Ok(None); + } // Incomplete entry entry = Some(convert_to_multiple_entry( provider.clone(), @@ -261,6 +274,11 @@ pub(crate) fn maybe_emit_entry( if !records.is_empty() { // Safety: the records are guaranteed not empty if the key exists. let last_record = records.last().unwrap(); + if matches!(last_record.meta.tp, RecordType::Middle(last_seq) if last_seq == seq) + && last_record.data == record.data + { + return Ok(None); + } let legal = match last_record.meta.tp { // Legal if this record follows a First record. RecordType::First => seq == 1, @@ -290,14 +308,10 @@ pub(crate) fn maybe_emit_entry( provider.clone(), region_id, records, - )) + )); } else { - // Incomplete entry - entry = Some(convert_to_multiple_entry( - provider.clone(), - region_id, - vec![record], - )) + // Intentionally discard a Last record without buffered parts. It is either a + // duplicate Last record or the tail of an incomplete multipart entry. } } } @@ -359,7 +373,16 @@ mod tests { .unwrap() .is_none() ); - let record = new_test_record(RecordType::First, 2, region_id.as_u64(), vec![2; 100]); + let record = new_test_record(RecordType::First, 2, region_id.as_u64(), vec![1; 100]); + // An identical First is a duplicate and does not emit an incomplete entry. + assert!( + maybe_emit_entry(&provider, record, &mut buffer) + .unwrap() + .is_none() + ); + assert_eq!(buffer[®ion_id][0].meta.entry_id, 2); + + let record = new_test_record(RecordType::First, 3, region_id.as_u64(), vec![2; 100]); let incomplete_entry = maybe_emit_entry(&provider, record, &mut buffer) .unwrap() .unwrap(); @@ -369,29 +392,21 @@ mod tests { Entry::MultiplePart(MultiplePartEntry { provider: Provider::Kafka(provider.clone()), region_id, - entry_id: 1, + entry_id: 2, headers: vec![MultiplePartHeader::First], parts: vec![vec![1; 100]], }) ); - // `Last` overwrite `None` + // `Last` after `None` let mut buffer = HashMap::new(); let record = new_test_record(RecordType::Last, 1, region_id.as_u64(), vec![1; 100]); - let incomplete_entry = maybe_emit_entry(&provider, record, &mut buffer) - .unwrap() - .unwrap(); - - assert_eq!( - incomplete_entry, - Entry::MultiplePart(MultiplePartEntry { - provider: Provider::Kafka(provider.clone()), - region_id, - entry_id: 1, - headers: vec![MultiplePartHeader::Last], - parts: vec![vec![1; 100]], - }) + assert!( + maybe_emit_entry(&provider, record, &mut buffer) + .unwrap() + .is_none() ); + assert!(buffer.is_empty()); // `First` overwrite `Middle(0)` let mut buffer = HashMap::new(); @@ -451,6 +466,99 @@ mod tests { assert_matches!(err, error::Error::IllegalSequence { .. }); } + #[test] + fn test_maybe_emit_entry_deduplicates_multipart_records() { + let provider = Arc::new(KafkaProvider::new("my_topic".to_string())); + let region_id = RegionId::new(1, 1); + let mut buffer = HashMap::new(); + + for record in [ + new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]), + new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![2; 100]), + new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![2; 100]), + ] { + assert!( + maybe_emit_entry(&provider, record, &mut buffer) + .unwrap() + .is_none() + ); + } + + let last = new_test_record(RecordType::Last, 1, region_id.as_u64(), vec![3; 100]); + let entry = maybe_emit_entry(&provider, last, &mut buffer) + .unwrap() + .unwrap(); + + assert_eq!( + entry, + Entry::MultiplePart(MultiplePartEntry { + provider: Provider::Kafka(provider.clone()), + region_id, + entry_id: 1, + headers: vec![ + MultiplePartHeader::First, + MultiplePartHeader::Middle(1), + MultiplePartHeader::Last, + ], + parts: vec![vec![1; 100], vec![2; 100], vec![3; 100]], + }) + ); + let duplicate = new_test_record(RecordType::Last, 1, region_id.as_u64(), vec![3; 100]); + assert!( + maybe_emit_entry(&provider, duplicate, &mut buffer) + .unwrap() + .is_none() + ); + } + + #[test] + fn test_maybe_emit_entry_rejects_conflicting_duplicate_middle_record() { + let provider = Arc::new(KafkaProvider::new("my_topic".to_string())); + let region_id = RegionId::new(1, 1); + let mut buffer = HashMap::new(); + + for record in [ + new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]), + new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![2; 100]), + ] { + assert!( + maybe_emit_entry(&provider, record, &mut buffer) + .unwrap() + .is_none() + ); + } + + let duplicate = new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![3; 100]); + let err = maybe_emit_entry(&provider, duplicate, &mut buffer).unwrap_err(); + + assert_matches!(err, error::Error::IllegalSequence { .. }); + } + + #[test] + fn test_maybe_emit_entry_preserves_multipart_entry_order_before_full_record() { + let provider = Arc::new(KafkaProvider::new("my_topic".to_string())); + let region_id = RegionId::new(1, 1); + let mut buffer = HashMap::new(); + + let first = new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]); + assert!( + maybe_emit_entry(&provider, first, &mut buffer) + .unwrap() + .is_none() + ); + let last = new_test_record(RecordType::Last, 2, region_id.as_u64(), vec![2; 100]); + let multipart_entry = maybe_emit_entry(&provider, last, &mut buffer) + .unwrap() + .unwrap(); + let full = new_test_record(RecordType::Full, 3, region_id.as_u64(), vec![3; 100]); + let full_entry = maybe_emit_entry(&provider, full, &mut buffer) + .unwrap() + .unwrap(); + + assert_eq!(multipart_entry.entry_id(), 2); + assert_eq!(full_entry.entry_id(), 3); + } + #[test] fn test_meta_size() { let meta = RecordMeta { From 1122d1e7d7e13380b443d2463cf4830b69e6baf3 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 15 Jul 2026 13:02:03 +0800 Subject: [PATCH 36/48] fix: apply log query limit after expressions (#8496) Signed-off-by: discord9 --- src/query/src/log_query/planner.rs | 38 ++++++------ src/servers/src/http/logs.rs | 42 +++++++++++++- tests-integration/tests/http.rs | 92 +++++++++++++++++++++++++++++- 3 files changed, 148 insertions(+), 24 deletions(-) diff --git a/src/query/src/log_query/planner.rs b/src/query/src/log_query/planner.rs index 53d49f98db0..05aeb0795d2 100644 --- a/src/query/src/log_query/planner.rs +++ b/src/query/src/log_query/planner.rs @@ -34,8 +34,6 @@ use crate::log_query::error::{ UnknownTableSnafu, }; -const DEFAULT_LIMIT: usize = 1000; - pub struct LogQueryPlanner { table_provider: DfTableSourceProvider, session_state: SessionState, @@ -99,19 +97,18 @@ impl LogQueryPlanner { .context(DataFusionPlanningSnafu)?; } - // Apply limit - plan_builder = plan_builder - .limit( - query.limit.skip.unwrap_or(0), - Some(query.limit.fetch.unwrap_or(DEFAULT_LIMIT)), - ) - .context(DataFusionPlanningSnafu)?; - // Apply log expressions for expr in &query.exprs { plan_builder = self.process_log_expr(plan_builder, expr)?; } + // Apply pagination to the final result after all log expressions. + if query.limit.skip.is_some() || query.limit.fetch.is_some() { + plan_builder = plan_builder + .limit(query.limit.skip.unwrap_or(0), query.limit.fetch) + .context(DataFusionPlanningSnafu)?; + } + // Build the final plan let plan = plan_builder.build().context(DataFusionPlanningSnafu)?; @@ -849,7 +846,7 @@ mod tests { }; let plan = planner.query_to_plan(log_query).await.unwrap(); - let expected = "Limit: skip=10, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ + let expected = "Limit: skip=10, fetch=None [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; @@ -884,9 +881,8 @@ mod tests { }; let plan = planner.query_to_plan(log_query).await.unwrap(); - let expected = "Limit: skip=0, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; + let expected = "Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ +\n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); } @@ -931,8 +927,8 @@ mod tests { }; let plan = planner.query_to_plan(log_query).await.unwrap(); - let expected = "Aggregate: groupBy=[[greptime.public.test_table.host]], aggr=[[count(greptime.public.test_table.message) AS count_result]] [host:Utf8;N, count_result:Int64]\ -\n Limit: skip=0, fetch=100 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ + let expected = "Limit: skip=0, fetch=100 [host:Utf8;N, count_result:Int64]\ +\n Aggregate: groupBy=[[greptime.public.test_table.host]], aggr=[[count(greptime.public.test_table.message) AS count_result]] [host:Utf8;N, count_result:Int64]\ \n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; @@ -971,8 +967,8 @@ mod tests { }; let plan = planner.query_to_plan(log_query).await.unwrap(); - let expected = "Projection: date_trunc(Utf8(\"day\"), greptime.public.test_table.timestamp) AS time_bucket [time_bucket:Timestamp(ms)]\ - \n Limit: skip=0, fetch=100 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ + let expected = "Limit: skip=0, fetch=100 [time_bucket:Timestamp(ms)]\ + \n Projection: date_trunc(Utf8(\"day\"), greptime.public.test_table.timestamp) AS time_bucket [time_bucket:Timestamp(ms)]\ \n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; @@ -1054,9 +1050,9 @@ mod tests { }; let plan = planner.query_to_plan(log_query).await.unwrap(); - let expected = "Aggregate: groupBy=[[2__date_histogram__time_bucket]], aggr=[[count(2__date_histogram__time_bucket) AS count_result]] [2__date_histogram__time_bucket:Timestamp(ns);N, count_result:Int64]\ -\n Projection: date_bin(Utf8(\"30 seconds\"), greptime.public.test_table.timestamp) AS 2__date_histogram__time_bucket [2__date_histogram__time_bucket:Timestamp(ns);N]\ -\n Limit: skip=0, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ + let expected = "Limit: skip=0, fetch=None [2__date_histogram__time_bucket:Timestamp(ns);N, count_result:Int64]\ +\n Aggregate: groupBy=[[2__date_histogram__time_bucket]], aggr=[[count(2__date_histogram__time_bucket) AS count_result]] [2__date_histogram__time_bucket:Timestamp(ns);N, count_result:Int64]\ +\n Projection: date_bin(Utf8(\"30 seconds\"), greptime.public.test_table.timestamp) AS 2__date_histogram__time_bucket [2__date_histogram__time_bucket:Timestamp(ns);N]\ \n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; diff --git a/src/servers/src/http/logs.rs b/src/servers/src/http/logs.rs index b9f23b68089..ca7c061da05 100644 --- a/src/servers/src/http/logs.rs +++ b/src/servers/src/http/logs.rs @@ -19,24 +19,31 @@ use axum::extract::State; use axum::response::{IntoResponse, Response}; use axum::{Extension, Json}; use common_telemetry::tracing; -use log_query::LogQuery; +use log_query::{Limit, LogQuery}; use session::context::{Channel, QueryContext}; use crate::http::result::greptime_result_v1::GreptimedbV1Response; use crate::query_handler::LogQueryHandlerRef; +const DEFAULT_FETCH: usize = 1000; + +fn apply_default_fetch(limit: &mut Limit) { + limit.fetch.get_or_insert(DEFAULT_FETCH); +} + #[axum_macros::debug_handler] #[tracing::instrument(skip_all, fields(protocol = "http", request_type = "logs"))] pub async fn logs( State(handler): State, Extension(mut query_ctx): Extension, - Json(params): Json, + Json(mut params): Json, ) -> Response { let exec_start = Instant::now(); let db = query_ctx.get_db_string(); query_ctx.set_channel(Channel::Log); let query_ctx = Arc::new(query_ctx); + apply_default_fetch(&mut params.limit); let _timer = crate::metrics::METRIC_HTTP_LOGS_ELAPSED .with_label_values(&[db.as_str()]) @@ -48,3 +55,34 @@ pub async fn logs( resp.with_execution_time(exec_start.elapsed().as_millis() as u64) .into_response() } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_apply_default_fetch() { + let mut limit = Limit { + skip: None, + fetch: None, + }; + apply_default_fetch(&mut limit); + assert_eq!(limit.fetch, Some(DEFAULT_FETCH)); + + let mut limit = Limit { + skip: Some(10), + fetch: None, + }; + apply_default_fetch(&mut limit); + assert_eq!(limit.skip, Some(10)); + assert_eq!(limit.fetch, Some(DEFAULT_FETCH)); + + let mut limit = Limit { + skip: Some(10), + fetch: Some(42), + }; + apply_default_fetch(&mut limit); + assert_eq!(limit.skip, Some(10)); + assert_eq!(limit.fetch, Some(42)); + } +} diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index a541b82b09e..133afd7822d 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -44,7 +44,7 @@ use common_memory_manager::OnExhaustedPolicy; use common_options::plugin_options::StandaloneFlag; use flate2::Compression; use flate2::write::GzEncoder; -use log_query::{Context, Limit, LogQuery, TimeFilter}; +use log_query::{AggFunc, Context, Limit, LogExpr, LogQuery, TimeFilter}; use loki_proto::logproto::{EntryAdapter, LabelPairAdapter, PushRequest, StreamAdapter}; use loki_proto::prost_types::Timestamp; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; @@ -7880,6 +7880,84 @@ pub async fn test_log_query(store_type: StorageType) { ] ); + let res = client + .get("/v1/sql?sql=create table logs_limit (`ts` timestamp time index, `message` string);") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + let res = client + .get("/v1/sql?sql=insert into logs_limit select number, 'hello' from numbers limit 1001;") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + + let aggregate_query = LogQuery { + table: TableName::new("greptime", "public", "logs_limit"), + time_filter: TimeFilter { + start: Some("1970-01-01T00:00:00Z".to_string()), + end: Some("1970-01-01T00:00:02Z".to_string()), + span: None, + }, + limit: Limit { + skip: None, + fetch: None, + }, + columns: vec![], + filters: Default::default(), + context: Context::None, + exprs: vec![LogExpr::AggrFunc { + expr: vec![AggFunc::new( + "count".to_string(), + vec![LogExpr::NamedIdent("message".to_string())], + Some("count_result".to_string()), + )], + by: vec![], + }], + }; + let res = client + .post("/v1/logs") + .header("Content-Type", "application/json") + .body(serde_json::to_string(&aggregate_query).unwrap()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + assert_eq!(get_rows_from_output(&res.text().await), "[[1001]]"); + + let mut output_query = LogQuery { + table: TableName::new("greptime", "public", "logs_limit"), + time_filter: TimeFilter { + start: Some("1970-01-01T00:00:00Z".to_string()), + end: Some("1970-01-01T00:00:02Z".to_string()), + span: None, + }, + limit: Limit { + skip: None, + fetch: None, + }, + columns: vec!["ts".to_string(), "message".to_string()], + filters: Default::default(), + context: Context::None, + exprs: vec![], + }; + let res = client + .post("/v1/logs") + .header("Content-Type", "application/json") + .body(serde_json::to_string(&output_query).unwrap()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + assert_eq!(get_output_row_count(&res.text().await), 1000); + + output_query.limit.fetch = Some(17); + let res = client + .post("/v1/logs") + .header("Content-Type", "application/json") + .body(serde_json::to_string(&output_query).unwrap()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + assert_eq!(get_output_row_count(&res.text().await), 17); + guard.remove_all().await; } @@ -9662,6 +9740,18 @@ fn get_rows_from_output(output: &str) -> String { .to_string() } +fn get_output_row_count(output: &str) -> usize { + let resp: Value = serde_json::from_str(output).unwrap(); + resp.get("output") + .and_then(Value::as_array) + .and_then(|v| v.first()) + .and_then(|v| v.get("records")) + .and_then(|v| v.get("rows")) + .and_then(Value::as_array) + .unwrap() + .len() +} + fn compress_vec_with_gzip(data: Vec) -> Vec { let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); encoder.write_all(&data).unwrap(); From 3c1859b626a0939e2528c4a49f06542bf1c794e2 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Wed, 15 Jul 2026 15:17:46 +0800 Subject: [PATCH 37/48] ci: nightly jsonbench test 100m dataset (#8529) * ci: attach the jsonbench data volume Signed-off-by: luofucong * Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --------- Signed-off-by: luofucong Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .github/workflows/nightly-jsonbench.yaml | 94 +++++++++++++++++++++--- 1 file changed, 82 insertions(+), 12 deletions(-) diff --git a/.github/workflows/nightly-jsonbench.yaml b/.github/workflows/nightly-jsonbench.yaml index da13a428e38..d5db6f9ba4a 100644 --- a/.github/workflows/nightly-jsonbench.yaml +++ b/.github/workflows/nightly-jsonbench.yaml @@ -10,14 +10,19 @@ on: description: The nightly build workflow run id to download GreptimeDB artifacts from required: true type: string + download_dataset: + description: Download the JSONBench dataset before running the benchmark + required: false + default: false + type: boolean permissions: actions: read contents: read concurrency: - group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} - cancel-in-progress: true + group: ${{ github.workflow }} + cancel-in-progress: false jobs: resolve-artifact: @@ -115,6 +120,34 @@ jobs: security-group-id: ${{ vars.EC2_RUNNER_SECURITY_GROUP_ID }} subnet-id: ${{ vars.EC2_RUNNER_SUBNET_ID }} + - name: Attach JSONBench data volume + if: ${{ steps.start-linux-runner.outputs.ec2-instance-id != '' && inputs.download_dataset != true }} + shell: bash + env: + INSTANCE_ID: ${{ steps.start-linux-runner.outputs.ec2-instance-id }} + VOLUME_ID: ${{ secrets.JSONBENCH_DATA_VOLUME_ID }} + run: | + set -euo pipefail + + if [[ -z "${VOLUME_ID}" ]]; then + echo "JSONBENCH_DATA_VOLUME_ID is not configured" + exit 1 + fi + + echo "Waiting for JSONBench data volume to become available" + aws ec2 wait volume-available --volume-ids "${VOLUME_ID}" + + echo "Attaching JSONBench data volume to EC2 instance ${INSTANCE_ID} as /dev/sdf" + aws ec2 attach-volume \ + --instance-id "${INSTANCE_ID}" \ + --volume-id "${VOLUME_ID}" \ + --device /dev/sdf \ + > /dev/null + + echo "Waiting for JSONBench data volume attachment to complete" + aws ec2 wait volume-in-use --volume-ids "${VOLUME_ID}" + echo "JSONBench data volume is attached to EC2 instance ${INSTANCE_ID}" + jsonbench: name: Run JSONBench if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success') }} @@ -123,7 +156,37 @@ jobs: timeout-minutes: 120 env: JSONBENCH_OUTPUT_PREFIX: _linux-${{ needs.resolve-artifact.outputs.artifact-arch }} + JSONBENCH_DATASET: 3 + JSONBENCH_DATA_DIR: /root/data/bluesky steps: + - name: Mount JSONBench data volume + if: ${{ needs.allocate-runner.outputs.linux-ec2-runner-instance-id != '' && inputs.download_dataset != true }} + shell: bash + env: + VOLUME_ID: ${{ secrets.JSONBENCH_DATA_VOLUME_ID }} + run: | + set -euo pipefail + + volume_serial="${VOLUME_ID//-/}" + device="" + echo "Waiting for the JSONBench data volume block device" + for _ in {1..30}; do + device=$(lsblk -dpno NAME,SERIAL | awk -v serial="${volume_serial}" '$2 == serial { print $1; exit }') + [[ -n "${device}" ]] && break + sleep 2 + done + + if [[ -z "${device}" ]]; then + echo "Cannot find the JSONBench data volume block device" + exit 1 + fi + + echo "Resolved JSONBench data volume to ${device}" + mkdir -p "${JSONBENCH_DATA_DIR}" + echo "Mounting ${device} read-only at ${JSONBENCH_DATA_DIR}" + mount -o ro "${device}" "${JSONBENCH_DATA_DIR}" + echo "JSONBench data volume is mounted at ${JSONBENCH_DATA_DIR}" + - name: Download GreptimeDB nightly artifact uses: actions/download-artifact@v4 with: @@ -142,26 +205,33 @@ jobs: chmod +x ./greptime rm -rf greptimedb-artifact "${{ needs.resolve-artifact.outputs.artifact-name }}" - - name: Run JSONBench - env: - # TODO(LFC): Change to "3" (100m) when JSON2 ingestion performance is optimized. - JSONBENCH_DATASET: 2 + - name: Clone JSONBench shell: bash run: | set -euo pipefail - export JSONBENCH_DATA_DIR="/root/data/bluesky" - echo "Use JSONBench data directory ${JSONBENCH_DATA_DIR}" - echo "${JSONBENCH_DATASET}" > jsonbench.dataset - - echo "Cloning JSONBench" git clone --branch greptimedb-new-json --depth 1 https://github.com/GreptimeTeam/JSONBench.git JSONBench + - name: Download JSONBench dataset + if: ${{ inputs.download_dataset == true }} + shell: bash + run: | + set -euo pipefail + echo "Downloading JSONBench dataset choice ${JSONBENCH_DATASET} to ${JSONBENCH_DATA_DIR}" mkdir -p "${JSONBENCH_DATA_DIR}" printf "${JSONBENCH_DATASET}\n" | ./JSONBench/download_data.sh + + - name: Run JSONBench + shell: bash + run: | + set -euo pipefail + + echo "Use JSONBench data directory ${JSONBENCH_DATA_DIR}" + echo "${JSONBENCH_DATASET}" > jsonbench.dataset + downloaded_files=$(find "${JSONBENCH_DATA_DIR}" -type f | wc -l) - echo "Downloaded JSONBench dataset files: ${downloaded_files}" + echo "Available JSONBench dataset files: ${downloaded_files}" export GREPTIMEDB_STANDALONE__WAL__DIR=greptimedb_data/wal export GREPTIMEDB_STANDALONE__STORAGE__DATA_HOME=greptimedb_data From 0a002f4d65f83aec561567d8e90497fd0a99c0b1 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Wed, 15 Jul 2026 15:31:25 +0800 Subject: [PATCH 38/48] feat: support per-region write buffer limits (#8473) * feat(mito): add per-region write buffer limit Signed-off-by: evenyag * feat(mito): add default region write buffer size Signed-off-by: evenyag * docs: expand configuration change checklist Signed-off-by: evenyag * test: cover table write buffer size option Signed-off-by: evenyag * fix(mito): refine region write buffer checks Signed-off-by: evenyag * fix(mito): keep region-stalled writes queued Signed-off-by: evenyag * fix(mito): allow zero region write buffer size Signed-off-by: evenyag * fix(mito): validate region state before write stall Signed-off-by: evenyag * fix(mito): reject writes beyond region hard limit Signed-off-by: evenyag * docs: update example Signed-off-by: evenyag --------- Signed-off-by: evenyag --- AGENTS.md | 6 +- config/config.md | 2 + config/datanode.example.toml | 3 + config/standalone.example.toml | 3 + src/cmd/tests/load_config_test.rs | 3 + src/mito2/src/compaction/window.rs | 1 + src/mito2/src/config.rs | 5 + src/mito2/src/engine/flush_test.rs | 577 +++++++++++++++++- src/mito2/src/flush.rs | 2 + src/mito2/src/region/options.rs | 41 ++ src/mito2/src/worker.rs | 8 + src/mito2/src/worker/handle_alter.rs | 17 +- src/mito2/src/worker/handle_flush.rs | 165 ++++- src/mito2/src/worker/handle_write.rs | 119 +++- src/store-api/src/mito_engine_options.rs | 4 + src/store-api/src/region_request.rs | 49 +- src/table/src/metadata.rs | 3 + src/table/src/requests.rs | 4 +- tests-integration/tests/http.rs | 1 + .../common/alter/alter_table_options.result | 73 +++ .../common/alter/alter_table_options.sql | 19 + .../information_schema/region_info.result | 2 +- .../common/information_schema/region_info.sql | 2 +- .../standalone/common/show/show_create.result | 2 +- .../standalone/common/show/show_create.sql | 2 +- 25 files changed, 1087 insertions(+), 26 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index ce25569aff1..fd2ec44aba1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -77,8 +77,10 @@ change-coupling points, and gotchas: 2. `make clippy` 3. `make test` (or `cargo nextest run`) 4. `make check-udeps` (run `make fix-udeps` if it reports unused dependencies). -5. If you changed sample config under `config/`, run `make config-docs` (needs - Docker) and commit the regenerated `config/config.md`. +5. If you added or changed a public configuration option, update the applicable + example TOMLs, configuration-loading and serialized-config snapshot tests, + and related user-facing documentation. Run `make config-docs` (needs Docker) + and commit the regenerated `config/config.md`. 6. If you changed a persisted or wire format, add a compatibility test case (see `.agents/architecture-invariants.md`). 7. Use a conventional-commit title, sign off commits (`git commit -s`), and sign diff --git a/config/config.md b/config/config.md index b0da5722f98..be287826fbc 100644 --- a/config/config.md +++ b/config/config.md @@ -157,6 +157,7 @@ | `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size`. | +| `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. | | `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. | | `region_engine.mito.vector_cache_size` | String | Auto | Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. | | `region_engine.mito.page_cache_size` | String | Auto | Cache size for pages of SST row groups. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory. | @@ -560,6 +561,7 @@ | `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size` | +| `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. | | `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. | | `region_engine.mito.vector_cache_size` | String | Auto | Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. | | `region_engine.mito.page_cache_size` | String | Auto | Cache size for pages of SST row groups. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index f823d3114be..cfa9f695326 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -468,6 +468,9 @@ auto_flush_interval = "1h" ## @toml2docs:none-default="Auto" #+ global_write_buffer_reject_size = "2GB" +## Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. +#+ default_region_write_buffer_size = "0" + ## Cache size for SST metadata. Setting it to 0 to disable the cache. ## If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. ## @toml2docs:none-default="Auto" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 822f9ba568d..439c947c15c 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -597,6 +597,9 @@ auto_flush_interval = "1h" ## @toml2docs:none-default="Auto" #+ global_write_buffer_reject_size = "2GB" +## Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. +#+ default_region_write_buffer_size = "0" + ## Cache size for SST metadata. Setting it to 0 to disable the cache. ## If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. ## @toml2docs:none-default="Auto" diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 3570e88cc1e..40188dd0489 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -16,6 +16,7 @@ use std::time::Duration; use cmd::options::GreptimeOptions; use common_base::memory_limit::MemoryLimit; +use common_base::readable_size::ReadableSize; use common_config::{Configurable, DEFAULT_DATA_HOME, ENV_VAR_SEP}; use common_options::datanode::{ClientOptions, DatanodeClientOptions}; use common_telemetry::logging::{DEFAULT_LOGGING_DIR, DEFAULT_OTLP_HTTP_ENDPOINT, LoggingOptions}; @@ -90,6 +91,7 @@ fn test_load_datanode_example_config() { region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { auto_flush_interval: Duration::from_secs(3600), + default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), scan_memory_limit: MemoryLimit::Unlimited, ..Default::default() @@ -299,6 +301,7 @@ fn test_load_standalone_example_config() { region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { auto_flush_interval: Duration::from_secs(3600), + default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), scan_memory_limit: MemoryLimit::Unlimited, ..Default::default() diff --git a/src/mito2/src/compaction/window.rs b/src/mito2/src/compaction/window.rs index 49e43f41fd4..f802601df30 100644 --- a/src/mito2/src/compaction/window.rs +++ b/src/mito2/src/compaction/window.rs @@ -254,6 +254,7 @@ mod tests { merge_mode: None, sst_format: None, primary_key_encoding: None, + write_buffer_size: None, }, compaction_time_window: None, } diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 407fedc955d..f7489709ee3 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -109,6 +109,10 @@ pub struct MitoConfig { pub global_write_buffer_size: ReadableSize, /// Global write buffer size threshold to reject write requests. pub global_write_buffer_reject_size: ReadableSize, + /// Default write buffer size for each region. Regions stall at this size and + /// reject writes at twice this size. Setting it to 0 disables both limits + /// unless the table specifies `write_buffer_size`. + pub default_region_write_buffer_size: ReadableSize, // Cache configs: /// Cache size for SST metadata. Setting it to 0 to disable the cache. @@ -206,6 +210,7 @@ impl Default for MitoConfig { auto_flush_interval: Duration::from_secs(30 * 60), global_write_buffer_size: ReadableSize::gb(1), global_write_buffer_reject_size: ReadableSize::gb(2), + default_region_write_buffer_size: ReadableSize::mb(0), sst_meta_cache_size: ReadableSize::mb(128), vector_cache_size: ReadableSize::mb(512), page_cache_size: ReadableSize::mb(512), diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index 6d68a33abda..ccac8982970 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -14,6 +14,7 @@ //! Flush tests for mito engine. +use std::assert_matches; use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering}; @@ -22,6 +23,7 @@ use std::time::Duration; use api::v1::Rows; use async_trait::async_trait; use common_base::Plugins; +use common_base::readable_size::ReadableSize; use common_recordbatch::RecordBatches; use common_time::util::current_time_millis; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; @@ -29,17 +31,20 @@ use rstest::rstest; use rstest_reuse::{self, apply}; use store_api::ManifestVersion; use store_api::metadata::RegionMetadataRef; -use store_api::region_engine::RegionEngine; +use store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY; +use store_api::region_engine::{RegionEngine, RegionRole}; use store_api::region_request::{ - PathType, RegionCloseRequest, RegionFlushRequest, RegionOpenRequest, RegionRequest, - ReplayCheckpoint, + AlterKind, PathType, RegionAlterRequest, RegionCloseRequest, RegionFlushRequest, + RegionOpenRequest, RegionPutRequest, RegionRequest, ReplayCheckpoint, SetRegionOption, }; use store_api::storage::{RegionId, ScanRequest}; use tokio::sync::Notify; use crate::config::MitoConfig; -use crate::engine::listener::{FlushListener, StallListener}; +use crate::engine::MitoEngine; +use crate::engine::listener::{EventListener, FlushListener, StallListener}; use crate::engine::region_hook::{RegionHook, RegionHookRef, SstFileInfo}; +use crate::error::Error; use crate::manifest::action::RegionMetaActionList; use crate::test_util::{ CreateRequestBuilder, LogStoreFactory, MockWriteBufferManager, TestEnv, build_rows, @@ -50,6 +55,73 @@ use crate::test_util::{ use crate::time_provider::TimeProvider; use crate::worker::MAX_INITIAL_CHECK_DELAY_SECS; +async fn set_write_buffer_size_to_current_usage(engine: &MitoEngine, region_id: RegionId) { + let version = engine.get_region(region_id).unwrap().version(); + let memory_usage = version.memtables.mutable_usage() + version.memtables.immutables_usage(); + assert!(memory_usage > 0); + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::SetRegionOptions { + options: vec![SetRegionOption::WriteBufferSize(Some(ReadableSize( + memory_usage as u64, + )))], + }, + }), + ) + .await + .unwrap(); +} + +struct RegionFlushGate { + blocked_region_id: RegionId, + observed_region_id: RegionId, + blocked_flush_started: Notify, + resume_blocked_flush: Notify, + observed_flush_finished: Notify, +} + +impl RegionFlushGate { + fn new(blocked_region_id: RegionId, observed_region_id: RegionId) -> Self { + Self { + blocked_region_id, + observed_region_id, + blocked_flush_started: Notify::new(), + resume_blocked_flush: Notify::new(), + observed_flush_finished: Notify::new(), + } + } + + async fn wait_blocked_flush_started(&self) { + self.blocked_flush_started.notified().await; + } + + fn resume_blocked_flush(&self) { + self.resume_blocked_flush.notify_one(); + } + + async fn wait_observed_flush_finished(&self) { + self.observed_flush_finished.notified().await; + } +} + +#[async_trait] +impl EventListener for RegionFlushGate { + fn on_flush_success(&self, region_id: RegionId) { + if region_id == self.observed_region_id { + self.observed_flush_finished.notify_one(); + } + } + + async fn on_flush_begin(&self, region_id: RegionId) { + if region_id == self.blocked_region_id { + self.blocked_flush_started.notify_one(); + self.resume_blocked_flush.notified().await; + } + } +} + #[tokio::test] async fn test_manual_flush() { test_manual_flush_with_format(false).await; @@ -273,6 +345,503 @@ async fn test_write_stall_with_format(flat_format: bool) { assert_eq!(expected, batches.pretty_print().unwrap()); } +#[tokio::test] +async fn test_region_write_buffer_limit_flushes_hot_region() { + let mut env = TestEnv::new().await; + let listener = Arc::new(FlushListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + + let hot_region_id = RegionId::new(1, 1); + let normal_region_id = RegionId::new(2, 1); + let idle_region_id = RegionId::new(3, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + hot_region_id.table_id(), + "hot_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + idle_region_id.table_id(), + "idle_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + normal_region_id.table_id(), + "normal_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + // The hot region sets its limit to its memory usage after the initial write. + let hot_request = CreateRequestBuilder::new().build(); + let hot_schema = rows_schema(&hot_request); + engine + .handle_request(hot_region_id, RegionRequest::Create(hot_request)) + .await + .unwrap(); + + // An explicit table option takes precedence over the engine-level default. + let normal_request = CreateRequestBuilder::new() + .table_dir("normal") + .insert_option(WRITE_BUFFER_SIZE_KEY, "1GiB") + .build(); + let normal_schema = rows_schema(&normal_request); + engine + .handle_request(normal_region_id, RegionRequest::Create(normal_request)) + .await + .unwrap(); + + let idle_request = CreateRequestBuilder::new() + .table_dir("idle") + .insert_option(WRITE_BUFFER_SIZE_KEY, "1B") + .build(); + let idle_schema = rows_schema(&idle_request); + engine + .handle_request(idle_region_id, RegionRequest::Create(idle_request)) + .await + .unwrap(); + + put_rows( + &engine, + hot_region_id, + Rows { + schema: hot_schema.clone(), + rows: build_rows_for_key("hot", 0, 2, 0), + }, + ) + .await; + set_write_buffer_size_to_current_usage(&engine, hot_region_id).await; + put_rows( + &engine, + normal_region_id, + Rows { + schema: normal_schema, + rows: build_rows_for_key("normal", 0, 2, 0), + }, + ) + .await; + put_rows( + &engine, + idle_region_id, + Rows { + schema: idle_schema, + rows: build_rows_for_key("idle", 0, 2, 0), + }, + ) + .await; + + // The next write triggers a flush after the first write crosses the threshold. + put_rows( + &engine, + hot_region_id, + Rows { + schema: hot_schema, + rows: build_rows_for_key("hot", 2, 2, 0), + }, + ) + .await; + listener.wait().await; + + let scanner = engine + .scanner(hot_region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(1, scanner.num_files()); + + let scanner = engine + .scanner(normal_region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(0, scanner.num_files()); + assert_eq!(1, scanner.num_memtables()); + + let scanner = engine + .scanner(idle_region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(0, scanner.num_files()); + assert_eq!(1, scanner.num_memtables()); +} + +#[tokio::test] +async fn test_region_write_buffer_stall_does_not_block_other_region() { + let mut env = TestEnv::new().await; + let listener = Arc::new(StallListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + + let stalled_region_id = RegionId::new(1, 1); + let normal_region_id = RegionId::new(2, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + stalled_region_id.table_id(), + "stalled_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + normal_region_id.table_id(), + "normal_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let stalled_request = CreateRequestBuilder::new().build(); + let stalled_schema = rows_schema(&stalled_request); + engine + .handle_request(stalled_region_id, RegionRequest::Create(stalled_request)) + .await + .unwrap(); + + let normal_request = CreateRequestBuilder::new() + .table_dir("normal") + .insert_option(WRITE_BUFFER_SIZE_KEY, "1GiB") + .build(); + let normal_schema = rows_schema(&normal_request); + engine + .handle_request(normal_region_id, RegionRequest::Create(normal_request)) + .await + .unwrap(); + + put_rows( + &engine, + stalled_region_id, + Rows { + schema: stalled_schema.clone(), + rows: build_rows_for_key("stalled", 0, 2, 0), + }, + ) + .await; + set_write_buffer_size_to_current_usage(&engine, stalled_region_id).await; + + let engine_cloned = engine.clone(); + let stalled_write = tokio::spawn(async move { + put_rows( + &engine_cloned, + stalled_region_id, + Rows { + schema: stalled_schema, + rows: build_rows_for_key("stalled", 2, 2, 0), + }, + ) + .await; + }); + + listener.wait().await; + + put_rows( + &engine, + normal_region_id, + Rows { + schema: normal_schema, + rows: build_rows_for_key("normal", 0, 2, 0), + }, + ) + .await; + + flush_region(&engine, stalled_region_id, None).await; + stalled_write.await.unwrap(); + + let scanner = engine + .scanner(normal_region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(0, scanner.num_files()); + assert_eq!(1, scanner.num_memtables()); +} + +#[tokio::test] +async fn test_region_write_buffer_does_not_stall_follower_write() { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_region_write_buffer_size: ReadableSize(1), + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("follower", 0, 2, 0), + }, + ) + .await; + + engine + .set_region_role(region_id, RegionRole::Follower) + .unwrap(); + + let err = tokio::time::timeout( + Duration::from_secs(1), + engine.handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema, + rows: build_rows_for_key("follower", 2, 4, 0), + }, + hint: None, + partition_expr_version: None, + }), + ), + ) + .await + .expect("write to follower should not remain stalled") + .unwrap_err(); + assert_matches!( + err.into_inner().as_any().downcast_ref::().unwrap(), + Error::RegionState { .. } + ); +} + +#[tokio::test] +async fn test_unrelated_flush_does_not_resume_region_stalled_write() { + let mut env = TestEnv::new().await; + let stalled_region_id = RegionId::new(1, 1); + let normal_region_id = RegionId::new(2, 1); + let listener = Arc::new(RegionFlushGate::new(stalled_region_id, normal_region_id)); + let engine = env + .create_engine_with( + MitoConfig { + max_background_flushes: 2, + ..Default::default() + }, + None, + Some(listener.clone()), + None, + ) + .await; + + env.get_schema_metadata_manager() + .register_region_table_info( + stalled_region_id.table_id(), + "stalled_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + normal_region_id.table_id(), + "normal_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let stalled_request = CreateRequestBuilder::new().build(); + let stalled_schema = rows_schema(&stalled_request); + engine + .handle_request(stalled_region_id, RegionRequest::Create(stalled_request)) + .await + .unwrap(); + + let normal_request = CreateRequestBuilder::new() + .table_dir("normal") + .insert_option(WRITE_BUFFER_SIZE_KEY, "1GiB") + .build(); + let normal_schema = rows_schema(&normal_request); + engine + .handle_request(normal_region_id, RegionRequest::Create(normal_request)) + .await + .unwrap(); + + put_rows( + &engine, + stalled_region_id, + Rows { + schema: stalled_schema.clone(), + rows: build_rows_for_key("stalled", 0, 2, 0), + }, + ) + .await; + set_write_buffer_size_to_current_usage(&engine, stalled_region_id).await; + put_rows( + &engine, + normal_region_id, + Rows { + schema: normal_schema, + rows: build_rows_for_key("normal", 0, 2, 0), + }, + ) + .await; + + let engine_cloned = engine.clone(); + let mut stalled_write = tokio::spawn(async move { + put_rows( + &engine_cloned, + stalled_region_id, + Rows { + schema: stalled_schema, + rows: build_rows_for_key("stalled", 2, 2, 0), + }, + ) + .await; + }); + + listener.wait_blocked_flush_started().await; + flush_region(&engine, normal_region_id, None).await; + listener.wait_observed_flush_finished().await; + + assert!( + tokio::time::timeout(Duration::from_secs(1), &mut stalled_write) + .await + .is_err(), + "an unrelated region flush resumed the region-stalled write" + ); + + listener.resume_blocked_flush(); + tokio::time::timeout(Duration::from_secs(10), &mut stalled_write) + .await + .expect("the stalled write should resume after its region flushes") + .unwrap(); +} + +#[tokio::test] +async fn test_region_write_buffer_rejects_only_full_region_queue() { + let mut env = TestEnv::new().await; + let hot_region_id = RegionId::new(1, 1); + let normal_region_id = RegionId::new(2, 1); + let listener = Arc::new(RegionFlushGate::new(hot_region_id, hot_region_id)); + let engine = env + .create_engine_with( + MitoConfig { + max_background_flushes: 2, + ..Default::default() + }, + None, + Some(listener.clone()), + None, + ) + .await; + + env.get_schema_metadata_manager() + .register_region_table_info( + hot_region_id.table_id(), + "hot_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + normal_region_id.table_id(), + "normal_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let hot_request = CreateRequestBuilder::new().build(); + let hot_schema = rows_schema(&hot_request); + engine + .handle_request(hot_region_id, RegionRequest::Create(hot_request)) + .await + .unwrap(); + let normal_request = CreateRequestBuilder::new().table_dir("normal").build(); + let normal_schema = rows_schema(&normal_request); + engine + .handle_request(normal_region_id, RegionRequest::Create(normal_request)) + .await + .unwrap(); + + put_rows( + &engine, + hot_region_id, + Rows { + schema: hot_schema.clone(), + rows: build_rows_for_key("hot", 0, 2, 0), + }, + ) + .await; + set_write_buffer_size_to_current_usage(&engine, hot_region_id).await; + + let engine_cloned = engine.clone(); + let stalled_schema = hot_schema.clone(); + let stalled_write = tokio::spawn(async move { + engine_cloned + .handle_request( + hot_region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: stalled_schema, + rows: build_rows_for_key("hot", 2, 1026, 2), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + }); + + listener.wait_blocked_flush_started().await; + + put_rows( + &engine, + normal_region_id, + Rows { + schema: normal_schema, + rows: build_rows_for_key("normal", 0, 2, 0), + }, + ) + .await; + + listener.resume_blocked_flush(); + listener.wait_observed_flush_finished().await; + + let err = tokio::time::timeout(Duration::from_secs(1), stalled_write) + .await + .expect("the stalled queue should be rejected at the region hard limit") + .unwrap() + .unwrap_err(); + assert_matches!( + err.into_inner().as_any().downcast_ref::().unwrap(), + Error::RejectWrite { .. } + ); +} + #[tokio::test] async fn test_flush_empty() { test_flush_empty_with_format(false).await; diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index e6aa15f9d48..71abc0b0a13 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -211,6 +211,8 @@ impl WriteBufferManager for WriteBufferManagerImpl { pub enum FlushReason { /// Engine reaches flush threshold. EngineFull, + /// Region reaches its write buffer threshold. + RegionFull, /// Manual flush. Manual, /// Flush to alter table. diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 92dfd460fe6..2ac881c4072 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -107,6 +107,11 @@ pub struct RegionOptions { /// Internal primary key encoding override used by metric-engine. #[serde(skip_serializing_if = "Option::is_none")] pub primary_key_encoding: Option, + /// Per-region write buffer size. A positive size flushes/stalls this region + /// independently of the global write buffer limit and rejects writes at twice + /// the configured size; zero disables both limits. + #[serde(skip_serializing_if = "Option::is_none")] + pub write_buffer_size: Option, } impl RegionOptions { @@ -256,6 +261,7 @@ impl RegionOptions { merge_mode: options.merge_mode, sst_format, primary_key_encoding, + write_buffer_size: options.write_buffer_size, }; opts.validate()?; @@ -354,6 +360,7 @@ impl Default for TwcsOptions { #[derive(Debug, Deserialize)] #[serde(default)] struct RegionOptionsWithoutEnum { + write_buffer_size: Option, /// Region SST files TTL. ttl: Option, #[serde(with = "humantime_serde")] @@ -371,6 +378,7 @@ impl Default for RegionOptionsWithoutEnum { fn default() -> Self { let options = RegionOptions::default(); RegionOptionsWithoutEnum { + write_buffer_size: options.write_buffer_size, ttl: options.ttl, auto_flush_interval: options.auto_flush_interval, storage: options.storage, @@ -520,6 +528,7 @@ mod tests { use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_wal::options::KafkaWalOptions; + use store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY; use super::*; @@ -569,6 +578,28 @@ mod tests { ); } + #[test] + fn test_with_write_buffer_size() { + let map = make_map(&[(WRITE_BUFFER_SIZE_KEY, "128MiB")]); + let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap(); + let expect = RegionOptions { + write_buffer_size: Some(ReadableSize::mb(128)), + ..Default::default() + }; + assert_eq!(expect, options); + } + + #[test] + fn test_with_zero_write_buffer_size() { + let map = make_map(&[(WRITE_BUFFER_SIZE_KEY, "0")]); + let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap(); + let expect = RegionOptions { + write_buffer_size: Some(ReadableSize::mb(0)), + ..Default::default() + }; + assert_eq!(expect, options); + } + #[test] fn test_with_storage() { let map = make_map(&[("storage", "S3")]); @@ -898,6 +929,7 @@ mod tests { merge_mode: Some(MergeMode::LastNonNull), sst_format: Some(FormatType::Flat), primary_key_encoding: None, + write_buffer_size: None, }; assert_eq!(expect, options); } @@ -928,10 +960,18 @@ mod tests { merge_mode: Some(MergeMode::LastNonNull), sst_format: None, primary_key_encoding: None, + write_buffer_size: Some(ReadableSize::mb(128)), }; let region_options_json_str = serde_json::to_string(&options).unwrap(); let got: RegionOptions = serde_json::from_str(®ion_options_json_str).unwrap(); assert_eq!(options, got); + + let old_region_options_json_str = r#"{"ttl":null}"#; + let got: RegionOptions = serde_json::from_str(old_region_options_json_str).unwrap(); + assert_eq!(None, got.write_buffer_size); + + let default_json = serde_json::to_value(RegionOptions::default()).unwrap(); + assert!(default_json.get(WRITE_BUFFER_SIZE_KEY).is_none()); } #[test] @@ -985,6 +1025,7 @@ mod tests { merge_mode: Some(MergeMode::LastNonNull), sst_format: None, primary_key_encoding: None, + write_buffer_size: None, }; assert_eq!(options, got); } diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 23b6dd4d7b8..a3602d29ff6 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -771,6 +771,14 @@ pub(crate) struct StalledRequests { } impl StalledRequests { + /// Returns the estimated size of stalled requests for a region. + pub(crate) fn estimated_size(&self, region_id: &RegionId) -> usize { + self.requests + .get(region_id) + .map(|(size, _, _)| *size) + .unwrap_or_default() + } + /// Appends stalled requests. pub(crate) fn append( &mut self, diff --git a/src/mito2/src/worker/handle_alter.rs b/src/mito2/src/worker/handle_alter.rs index e022d166315..87044d768de 100644 --- a/src/mito2/src/worker/handle_alter.rs +++ b/src/mito2/src/worker/handle_alter.rs @@ -188,6 +188,20 @@ impl RegionWorkerLoop { let mut current_options = version.options.clone(); for option in options { match option { + SetRegionOption::WriteBufferSize(new_write_buffer_size) => { + info!( + "Update region write_buffer_size: {}, previous: {:?} new: {:?}", + region.region_id, current_options.write_buffer_size, new_write_buffer_size + ); + current_options.write_buffer_size = new_write_buffer_size; + current_options.validate().map_err(|e| { + store_api::metadata::InvalidRegionRequestSnafu { + region_id: region.region_id, + err: e.to_string(), + } + .build() + })?; + } SetRegionOption::Ttl(new_ttl) => { info!( "Update region ttl: {}, previous: {:?} new: {:?}", @@ -273,7 +287,8 @@ fn new_region_options_on_empty_memtable( let mut current_options = current_options.clone(); for option in options { match option { - SetRegionOption::Ttl(_) + SetRegionOption::WriteBufferSize(_) + | SetRegionOption::Ttl(_) | SetRegionOption::Twsc(_, _) | SetRegionOption::AutoFlushInterval(_) => (), SetRegionOption::Format(format_str) => { diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs index c8e86699ea1..ccbe31c3a32 100644 --- a/src/mito2/src/worker/handle_flush.rs +++ b/src/mito2/src/worker/handle_flush.rs @@ -14,9 +14,11 @@ //! Handling flush related requests. +use std::collections::HashSet; use std::sync::Arc; use std::sync::atomic::Ordering; +use common_base::readable_size::ReadableSize; use common_telemetry::{debug, error, info}; use store_api::logstore::LogStore; use store_api::region_request::{RegionFlushReason, RegionFlushRequest}; @@ -26,10 +28,24 @@ use crate::config::{IndexBuildMode, MitoConfig}; use crate::error::{RegionNotFoundSnafu, Result}; use crate::flush::{FlushReason, RegionFlushTask}; use crate::region::MitoRegionRef; +use crate::region::version::VersionRef; use crate::request::{BuildIndexRequest, FlushFailed, FlushFinished, OnFailure, OptionOutputTx}; use crate::sst::index::IndexBuildType; use crate::worker::RegionWorkerLoop; +#[derive(Debug, Default)] +pub(crate) struct RegionWriteBufferPressure { + pub(crate) stalled_region_ids: HashSet, + pub(crate) rejected_region_ids: HashSet, +} + +#[derive(Debug, Default, PartialEq, Eq)] +struct RegionWriteBufferStatus { + should_flush: bool, + should_stall: bool, + should_reject: bool, +} + fn resolve_flush_reason( request_reason: Option, is_downgrading: bool, @@ -85,8 +101,7 @@ impl RegionWorkerLoop { } let version = region.version(); - let region_memtable_size = - version.memtables.mutable_usage() + version.memtables.immutables_usage(); + let region_memtable_size = region_memtable_usage(&version); let auto_flush_interval = version .options @@ -152,6 +167,53 @@ impl RegionWorkerLoop { Ok(()) } + /// Flushes write regions that exceed their flush threshold and returns region pressure. + pub(crate) fn maybe_flush_write_regions( + &mut self, + region_ids: HashSet, + ) -> RegionWriteBufferPressure { + let mut pressure = RegionWriteBufferPressure::default(); + + for region_id in region_ids { + let Some(region) = self.regions.get_region(region_id) else { + continue; + }; + if !region.is_writable() { + continue; + } + + let status = region_write_buffer_status( + ®ion.version(), + self.config.default_region_write_buffer_size, + self.stalled_requests.estimated_size(®ion_id), + ); + + if status.should_flush && !self.flush_scheduler.is_flush_requested(region.region_id) { + let task = self.new_flush_task( + ®ion, + FlushReason::RegionFull, + None, + self.config.clone(), + ); + if let Err(e) = self.flush_scheduler.schedule_flush( + region.region_id, + ®ion.version_control, + task, + ) { + error!(e; "Failed to schedule flush task for region {}", region.region_id); + } + } + + if status.should_reject { + pressure.rejected_region_ids.insert(region_id); + } else if status.should_stall { + pressure.stalled_region_ids.insert(region_id); + } + } + + pressure + } + /// Creates a flush task with specific `reason` for the `region`. pub(crate) fn new_flush_task( &self, @@ -179,6 +241,61 @@ impl RegionWorkerLoop { } } +pub(crate) fn region_memtable_usage(version: &VersionRef) -> usize { + version.memtables.mutable_usage() + version.memtables.immutables_usage() +} + +pub(crate) fn region_write_buffer_size( + version: &VersionRef, + default_region_write_buffer_size: ReadableSize, +) -> Option { + let size = version + .options + .write_buffer_size + .unwrap_or(default_region_write_buffer_size); + (size.as_bytes() > 0).then_some(size.as_bytes() as usize) +} + +/// Returns the pressure state of a region's write buffer. +fn region_write_buffer_status( + version: &VersionRef, + default_region_write_buffer_size: ReadableSize, + stalled_request_size: usize, +) -> RegionWriteBufferStatus { + let Some(write_buffer_size) = + region_write_buffer_size(version, default_region_write_buffer_size) + else { + return RegionWriteBufferStatus::default(); + }; + + let mutable_usage = version.memtables.mutable_usage(); + let memory_usage = region_memtable_usage(version); + region_write_buffer_status_from_usage( + write_buffer_size, + mutable_usage, + memory_usage, + stalled_request_size, + ) +} + +fn region_write_buffer_status_from_usage( + write_buffer_size: usize, + mutable_usage: usize, + memory_usage: usize, + stalled_request_size: usize, +) -> RegionWriteBufferStatus { + let mutable_limit = std::cmp::max(1, write_buffer_size / 2); + let should_stall = memory_usage >= write_buffer_size; + let reject_limit = write_buffer_size.saturating_mul(2); + let should_reject = memory_usage.saturating_add(stalled_request_size) >= reject_limit; + + RegionWriteBufferStatus { + should_flush: mutable_usage >= mutable_limit || should_stall, + should_stall, + should_reject, + } +} + impl RegionWorkerLoop { /// Handles manual flush request. pub(crate) fn handle_flush_request( @@ -438,4 +555,48 @@ mod tests { assert_eq!(resolve_flush_reason(None, true), FlushReason::Downgrading); assert_eq!(resolve_flush_reason(None, false), FlushReason::Manual); } + + #[test] + fn test_region_write_buffer_status_boundaries() { + assert_eq!( + RegionWriteBufferStatus::default(), + region_write_buffer_status_from_usage(100, 49, 99, 0) + ); + assert_eq!( + RegionWriteBufferStatus { + should_flush: true, + should_stall: false, + should_reject: false, + }, + region_write_buffer_status_from_usage(100, 50, 99, 0) + ); + assert_eq!( + RegionWriteBufferStatus { + should_flush: true, + should_stall: true, + should_reject: false, + }, + region_write_buffer_status_from_usage(100, 50, 100, 99) + ); + assert_eq!( + RegionWriteBufferStatus { + should_flush: true, + should_stall: true, + should_reject: true, + }, + region_write_buffer_status_from_usage(100, 50, 100, 100) + ); + } + + #[test] + fn test_region_write_buffer_status_saturates_reject_limit() { + assert_eq!( + RegionWriteBufferStatus { + should_flush: true, + should_stall: true, + should_reject: true, + }, + region_write_buffer_status_from_usage(usize::MAX, usize::MAX, usize::MAX, usize::MAX,) + ); + } } diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index dd0acd63000..0eb4c7bc39e 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -14,7 +14,7 @@ //! Handling write requests. -use std::collections::{HashMap, hash_map}; +use std::collections::{HashMap, HashSet, hash_map}; use std::sync::Arc; use api::v1::OpType; @@ -50,8 +50,11 @@ impl RegionWorkerLoop { return; } - // Flush this worker if the engine needs to flush. + let write_region_ids = write_region_ids(write_requests, bulk_requests); + + // Check region pressure before writes to match the global write buffer behavior. self.maybe_flush_worker(); + let pressure = self.maybe_flush_write_regions(write_region_ids); if self.should_reject_write() { // The memory pressure is still too high, reject write requests. @@ -61,6 +64,20 @@ impl RegionWorkerLoop { return; } + if !pressure.rejected_region_ids.is_empty() { + reject_region_write_requests( + &pressure.rejected_region_ids, + write_requests, + bulk_requests, + ); + for region_id in &pressure.rejected_region_ids { + self.reject_region_stalled_requests(region_id); + } + if write_requests.is_empty() && bulk_requests.is_empty() { + return; + } + } + if self.write_buffer_manager.should_stall() && allow_stall { let stalled_count = (write_requests.len() + bulk_requests.len()) as i64; self.stalling_count.add(stalled_count); @@ -70,6 +87,17 @@ impl RegionWorkerLoop { return; } + if allow_stall { + self.stall_region_write_requests( + &pressure.stalled_region_ids, + write_requests, + bulk_requests, + ); + if write_requests.is_empty() && bulk_requests.is_empty() { + return; + } + } + // Prepare write context. let mut region_ctxs = { let _timer = WRITE_STAGE_ELAPSED @@ -163,15 +191,31 @@ impl RegionWorkerLoop { .inc_by(delete_rows as u64); } - /// Handles all stalled write requests. + /// Handles stalled write requests whose regions no longer need to stall. pub(crate) async fn handle_stalled_requests(&mut self) { - // Handle stalled requests. - let stalled = std::mem::take(&mut self.stalled_requests); - self.stalling_count.sub(stalled.stalled_count() as i64); - // We already stalled these requests, don't stall them again. - for (_, (_, mut requests, mut bulk)) in stalled.requests { - self.handle_write_requests(&mut requests, &mut bulk, false) - .await; + let region_ids = self + .stalled_requests + .requests + .keys() + .copied() + .collect::>(); + let pressure = self.maybe_flush_write_regions(region_ids); + for region_id in &pressure.rejected_region_ids { + self.reject_region_stalled_requests(region_id); + } + let ready_region_ids = self + .stalled_requests + .requests + .keys() + .filter(|region_id| !pressure.stalled_region_ids.contains(region_id)) + .copied() + .collect::>(); + + // These requests have already been stalled. Retry ready regions without stalling the + // same requests again. Regions that still exceed their limit remain in the queue until + // their own flush releases the pressure. + for region_id in ready_region_ids { + self.handle_region_stalled_requests(®ion_id, false).await; } } @@ -553,6 +597,34 @@ impl RegionWorkerLoop { self.write_buffer_manager.memory_usage() + self.stalled_requests.estimated_size >= self.config.global_write_buffer_reject_size.as_bytes() as usize } + + fn stall_region_write_requests( + &mut self, + stalled_region_ids: &HashSet, + write_requests: &mut Vec, + bulk_requests: &mut Vec, + ) { + let mut stalled_count = 0; + let mut stalled_write_requests = write_requests + .extract_if(.., |req| { + stalled_region_ids.contains(&req.request.region_id) + }) + .collect::>(); + let mut stalled_bulk_requests = bulk_requests + .extract_if(.., |req| stalled_region_ids.contains(&req.region_id)) + .collect::>(); + + stalled_count += stalled_write_requests.len() + stalled_bulk_requests.len(); + self.stalled_requests + .append(&mut stalled_write_requests, &mut stalled_bulk_requests); + + if stalled_count > 0 { + let stalled_count = stalled_count as i64; + self.stalling_count.add(stalled_count); + WRITE_STALL_TOTAL.inc_by(stalled_count as u64); + self.listener.on_write_stall(); + } + } } /// Send rejected error to all `write_requests`. @@ -576,6 +648,33 @@ fn reject_write_requests( } } +fn reject_region_write_requests( + rejected_region_ids: &HashSet, + write_requests: &mut Vec, + bulk_requests: &mut Vec, +) { + let mut rejected_write_requests = write_requests + .extract_if(.., |req| { + rejected_region_ids.contains(&req.request.region_id) + }) + .collect::>(); + let mut rejected_bulk_requests = bulk_requests + .extract_if(.., |req| rejected_region_ids.contains(&req.region_id)) + .collect::>(); + reject_write_requests(&mut rejected_write_requests, &mut rejected_bulk_requests); +} + +fn write_region_ids( + write_requests: &[SenderWriteRequest], + bulk_requests: &[SenderBulkRequest], +) -> HashSet { + write_requests + .iter() + .map(|req| req.request.region_id) + .chain(bulk_requests.iter().map(|req| req.region_id)) + .collect() +} + /// Rejects delete request under append mode. fn check_op_type(append_mode: bool, request: &WriteRequest) -> Result<()> { if append_mode { diff --git a/src/store-api/src/mito_engine_options.rs b/src/store-api/src/mito_engine_options.rs index 8762ed69779..4320a40d8de 100644 --- a/src/store-api/src/mito_engine_options.rs +++ b/src/store-api/src/mito_engine_options.rs @@ -19,6 +19,8 @@ pub use common_wal::options::WAL_OPTIONS_KEY; /// Option key for append mode. pub const APPEND_MODE_KEY: &str = "append_mode"; +/// Option key for the per-region write buffer size. +pub const WRITE_BUFFER_SIZE_KEY: &str = "write_buffer_size"; /// Option key for merge mode. pub const MERGE_MODE_KEY: &str = "merge_mode"; /// Option key for TTL(time-to-live) @@ -72,6 +74,7 @@ pub const SST_FORMAT_KEY: &str = "sst_format"; pub fn is_mito_engine_option_key(key: &str) -> bool { [ "ttl", + WRITE_BUFFER_SIZE_KEY, AUTO_FLUSH_INTERVAL_KEY, COMPACTION_TYPE, COMPACTION_OVERRIDE, @@ -107,6 +110,7 @@ mod tests { #[test] fn test_is_mito_engine_option_key() { assert!(is_mito_engine_option_key("ttl")); + assert!(is_mito_engine_option_key("write_buffer_size")); assert!(is_mito_engine_option_key("auto_flush_interval")); assert!(is_mito_engine_option_key("compaction.type")); assert!(is_mito_engine_option_key("compaction.override")); diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 6aa3e3f49d0..d63c9058a33 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -33,6 +33,7 @@ use api::v1::{ SemanticType, SkippingIndexType as PbSkippingIndexType, WriteHint, }; pub use common_base::AffectedRows; +use common_base::readable_size::ReadableSize; use common_grpc::flight::FlightDecoder; use common_recordbatch::DfRecordBatch; use common_time::{TimeToLive, Timestamp}; @@ -54,7 +55,7 @@ use crate::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY; use crate::metrics; use crate::mito_engine_options::{ APPEND_MODE_KEY, AUTO_FLUSH_INTERVAL_KEY, SST_FORMAT_KEY, TTL_KEY, TWCS_MAX_OUTPUT_FILE_SIZE, - TWCS_TIME_WINDOW, TWCS_TRIGGER_FILE_NUM, + TWCS_TIME_WINDOW, TWCS_TRIGGER_FILE_NUM, WRITE_BUFFER_SIZE_KEY, }; use crate::path_utils::table_dir; use crate::storage::{ColumnId, RegionId, ScanRequest}; @@ -1405,6 +1406,7 @@ impl From for ModifyColumnType { #[derive(Debug, Eq, PartialEq, Clone, Serialize, Deserialize)] pub enum SetRegionOption { + WriteBufferSize(Option), Ttl(Option), // Modifying TwscOptions with values as (option name, new value). Twsc(String, String), @@ -1422,6 +1424,12 @@ impl TryFrom<&PbOption> for SetRegionOption { fn try_from(value: &PbOption) -> std::result::Result { let PbOption { key, value } = value; match key.as_str() { + WRITE_BUFFER_SIZE_KEY => { + let size = value + .parse::() + .map_err(|_| InvalidSetRegionOptionRequestSnafu { key, value }.build())?; + Ok(Self::WriteBufferSize(Some(size))) + } TTL_KEY => { let ttl = TimeToLive::from_humantime_or_str(value) .map_err(|_| InvalidSetRegionOptionRequestSnafu { key, value }.build())?; @@ -1470,6 +1478,7 @@ impl From<&UnsetRegionOption> for SetRegionOption { SetRegionOption::Twsc(unset_option.to_string(), String::new()) } UnsetRegionOption::Ttl => SetRegionOption::Ttl(Default::default()), + UnsetRegionOption::WriteBufferSize => SetRegionOption::WriteBufferSize(None), } } } @@ -1480,6 +1489,7 @@ impl TryFrom<&str> for UnsetRegionOption { fn try_from(key: &str) -> Result { match key.to_ascii_lowercase().as_str() { TTL_KEY => Ok(Self::Ttl), + WRITE_BUFFER_SIZE_KEY => Ok(Self::WriteBufferSize), TWCS_TRIGGER_FILE_NUM => Ok(Self::TwcsTriggerFileNum), TWCS_MAX_OUTPUT_FILE_SIZE => Ok(Self::TwcsMaxOutputFileSize), TWCS_TIME_WINDOW => Ok(Self::TwcsTimeWindow), @@ -1494,12 +1504,14 @@ pub enum UnsetRegionOption { TwcsMaxOutputFileSize, TwcsTimeWindow, Ttl, + WriteBufferSize, } impl UnsetRegionOption { pub fn as_str(&self) -> &str { match self { Self::Ttl => TTL_KEY, + Self::WriteBufferSize => WRITE_BUFFER_SIZE_KEY, Self::TwcsTriggerFileNum => TWCS_TRIGGER_FILE_NUM, Self::TwcsMaxOutputFileSize => TWCS_MAX_OUTPUT_FILE_SIZE, Self::TwcsTimeWindow => TWCS_TIME_WINDOW, @@ -1831,6 +1843,41 @@ mod tests { ); } + #[test] + fn test_write_buffer_size_region_options() { + let option = PbOption { + key: WRITE_BUFFER_SIZE_KEY.to_string(), + value: "128MiB".to_string(), + }; + assert_eq!( + SetRegionOption::WriteBufferSize(Some(ReadableSize::mb(128))), + SetRegionOption::try_from(&option).unwrap() + ); + + let option = PbOption { + key: WRITE_BUFFER_SIZE_KEY.to_string(), + value: "invalid".to_string(), + }; + SetRegionOption::try_from(&option).unwrap_err(); + + // Clearing the option uses UNSET. Empty SET values, including SQL NULL, + // remain invalid because the SQL layer does not distinguish NULL from ''. + let option = PbOption { + key: WRITE_BUFFER_SIZE_KEY.to_string(), + value: String::new(), + }; + SetRegionOption::try_from(&option).unwrap_err(); + + assert_eq!( + UnsetRegionOption::WriteBufferSize, + UnsetRegionOption::try_from(WRITE_BUFFER_SIZE_KEY).unwrap() + ); + assert_eq!( + SetRegionOption::WriteBufferSize(None), + SetRegionOption::from(&UnsetRegionOption::WriteBufferSize) + ); + } + /// Returns a new region metadata for testing. Metadata: /// `[(ts, ms, 1), (tag_0, string, 2), (field_0, string, 3), (field_1, bool, 4)]` fn new_metadata() -> RegionMetadata { diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index 006cfb34962..b2801dd8654 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -350,6 +350,9 @@ impl TableMeta { for request in requests { match request { + SetRegionOption::WriteBufferSize(new_write_buffer_size) => { + new_options.write_buffer_size = *new_write_buffer_size; + } SetRegionOption::Ttl(new_ttl) => { new_options.ttl = *new_ttl; } diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 12214a92f9d..12331bf8bb9 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -145,7 +145,7 @@ pub fn validate_table_option(key: &str) -> bool { #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] pub struct TableOptions { - /// Memtable size of memtable. + /// Per-region write buffer stall threshold. Writes are rejected at twice this size. pub write_buffer_size: Option, /// Time-to-live of table. Expired data will be automatically purged. pub ttl: Option, @@ -155,7 +155,7 @@ pub struct TableOptions { pub extra_options: HashMap, } -pub const WRITE_BUFFER_SIZE_KEY: &str = "write_buffer_size"; +pub const WRITE_BUFFER_SIZE_KEY: &str = store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY; pub const TTL_KEY: &str = store_api::mito_engine_options::TTL_KEY; pub const STORAGE_KEY: &str = "storage"; pub const COMMENT_KEY: &str = "comment"; diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 133afd7822d..ba6710da34b 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -1939,6 +1939,7 @@ compress_manifest = false experimental_compaction_memory_limit = "unlimited" experimental_compaction_on_exhausted = "wait" auto_flush_interval = "30m" +default_region_write_buffer_size = "0KiB" enable_write_cache = false write_cache_path = "" write_cache_size = "5GiB" diff --git a/tests/cases/standalone/common/alter/alter_table_options.result b/tests/cases/standalone/common/alter/alter_table_options.result index f91fd9da88e..b169562c894 100644 --- a/tests/cases/standalone/common/alter/alter_table_options.result +++ b/tests/cases/standalone/common/alter/alter_table_options.result @@ -259,6 +259,79 @@ DROP TABLE ato; Affected Rows: 0 +CREATE TABLE write_buffer_size_options( + i INTEGER, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY(i) +); + +Affected Rows: 0 + +ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; + +Affected Rows: 0 + +SHOW CREATE TABLE write_buffer_size_options; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | +| | "i" INT NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("i") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | write_buffer_size = '128.0MiB' | +| | ) | ++---------------------------+----------------------------------------------------------+ + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE write_buffer_size_options; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | +| | "i" INT NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("i") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | write_buffer_size = '128.0MiB' | +| | ) | ++---------------------------+----------------------------------------------------------+ + +ALTER TABLE write_buffer_size_options UNSET 'write_buffer_size'; + +Affected Rows: 0 + +SHOW CREATE TABLE write_buffer_size_options; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | +| | "i" INT NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("i") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++---------------------------+----------------------------------------------------------+ + +DROP TABLE write_buffer_size_options; + +Affected Rows: 0 + CREATE TABLE phy (ts timestamp time index, val double) engine=metric with ("physical_metric_table" = ""); Affected Rows: 0 diff --git a/tests/cases/standalone/common/alter/alter_table_options.sql b/tests/cases/standalone/common/alter/alter_table_options.sql index badf1275cdb..c3f67faeee5 100644 --- a/tests/cases/standalone/common/alter/alter_table_options.sql +++ b/tests/cases/standalone/common/alter/alter_table_options.sql @@ -55,6 +55,25 @@ SHOW CREATE TABLE ato; DROP TABLE ato; +CREATE TABLE write_buffer_size_options( + i INTEGER, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY(i) +); + +ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; + +SHOW CREATE TABLE write_buffer_size_options; + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE write_buffer_size_options; + +ALTER TABLE write_buffer_size_options UNSET 'write_buffer_size'; + +SHOW CREATE TABLE write_buffer_size_options; + +DROP TABLE write_buffer_size_options; + CREATE TABLE phy (ts timestamp time index, val double) engine=metric with ("physical_metric_table" = ""); ALTER TABLE phy set ttl='2years'; diff --git a/tests/cases/standalone/common/information_schema/region_info.result b/tests/cases/standalone/common/information_schema/region_info.result index ed9a07add60..91cab62336d 100644 --- a/tests/cases/standalone/common/information_schema/region_info.result +++ b/tests/cases/standalone/common/information_schema/region_info.result @@ -41,7 +41,7 @@ ADMIN FLUSH_TABLE('region_info_case'); +---------------------------------------+ -- SQLNESS REPLACE (\s+\d+\s+) --- SQLNESS REPLACE (\{".*"\}) +-- SQLNESS REPLACE (\{.*\}) -- SQLNESS REPLACE (-{40,}) ---------------- -- SQLNESS REPLACE (region_options\s+\|) region_options | SELECT region_id, state, role, writable, committed_sequence, flushed_sequence, manifest_version, compaction_time_window, region_options, sst_format diff --git a/tests/cases/standalone/common/information_schema/region_info.sql b/tests/cases/standalone/common/information_schema/region_info.sql index 1aa682393f2..6c901387ce9 100644 --- a/tests/cases/standalone/common/information_schema/region_info.sql +++ b/tests/cases/standalone/common/information_schema/region_info.sql @@ -11,7 +11,7 @@ INSERT INTO region_info_case VALUES (1, 1), (2, 2); ADMIN FLUSH_TABLE('region_info_case'); -- SQLNESS REPLACE (\s+\d+\s+) --- SQLNESS REPLACE (\{".*"\}) +-- SQLNESS REPLACE (\{.*\}) -- SQLNESS REPLACE (-{40,}) ---------------- -- SQLNESS REPLACE (region_options\s+\|) region_options | SELECT region_id, state, role, writable, committed_sequence, flushed_sequence, manifest_version, compaction_time_window, region_options, sst_format diff --git a/tests/cases/standalone/common/show/show_create.result b/tests/cases/standalone/common/show/show_create.result index 72392a0997d..18879ed95e8 100644 --- a/tests/cases/standalone/common/show/show_create.result +++ b/tests/cases/standalone/common/show/show_create.result @@ -15,7 +15,7 @@ PARTITION ON COLUMNS (`id`) ( ENGINE=mito WITH( ttl = '7d', - write_buffer_size = 1024 + write_buffer_size = '1KiB' ); Affected Rows: 0 diff --git a/tests/cases/standalone/common/show/show_create.sql b/tests/cases/standalone/common/show/show_create.sql index eb0f90c6746..fd8b1c16ea2 100644 --- a/tests/cases/standalone/common/show/show_create.sql +++ b/tests/cases/standalone/common/show/show_create.sql @@ -15,7 +15,7 @@ PARTITION ON COLUMNS (`id`) ( ENGINE=mito WITH( ttl = '7d', - write_buffer_size = 1024 + write_buffer_size = '1KiB' ); SHOW CREATE TABLE system_metrics; From f8aff344e9fba6833f1c96c3850031f3e7be8c69 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 15 Jul 2026 17:28:24 +0800 Subject: [PATCH 39/48] feat!: remove configuration of sparse_primary_key_encoding (#8470) * feat: remove configuration of sparse_primary_key_encoding, always use true * Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * fix: merge --------- Co-authored-by: dennis zhuang Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- config/config.md | 2 -- config/datanode.example.toml | 2 -- config/standalone.example.toml | 2 -- src/cmd/tests/load_config_test.rs | 43 +++++++++++++++++++++++-- src/metric-engine/src/config.rs | 8 ----- src/metric-engine/src/engine.rs | 7 ---- src/metric-engine/src/engine/create.rs | 5 +-- src/metric-engine/src/engine/open.rs | 10 ++---- src/metric-engine/src/engine/options.rs | 21 ++++++------ src/metric-engine/src/repeated_task.rs | 20 ++++-------- 10 files changed, 59 insertions(+), 61 deletions(-) diff --git a/config/config.md b/config/config.md index be287826fbc..c999e32a9c4 100644 --- a/config/config.md +++ b/config/config.md @@ -206,7 +206,6 @@ | `region_engine.mito.bloom_filter_index.mem_threshold_on_create` | String | `auto` | Memory threshold for bloom filter creation.
- `auto`: automatically determine the threshold based on the system memory size (default)
- `unlimited`: no memory limit
- `[size]` e.g. `64MB`: fixed memory threshold | | `region_engine.file` | -- | -- | Enable the file engine. | | `region_engine.metric` | -- | -- | Metric engine options. | -| `region_engine.metric.sparse_primary_key_encoding` | Bool | `true` | Whether to use sparse primary key encoding. | | `logging` | -- | -- | The logging options. | | `logging.dir` | String | `./greptimedb_data/logs` | The directory to store the log files. If set to empty, logs will not be written to files. | | `logging.level` | String | Unset | The log level. Can be `info`/`debug`/`warn`/`error`. | @@ -614,7 +613,6 @@ | `region_engine.mito.gc.unknown_file_lingering_time` | String | `1d` | Lingering time before deleting unknown files (files with undetermined expel time).
Only applies during full file listing GC.
This uses the object's last-modified timestamp as a heuristic (strict less-than comparison);
do not configure this value too small in production to avoid deleting pre-manifest files
from in-progress compaction or flush.
For active/open regions, an unknown file is deleted only if its object last-modified time exceeds this TTL.
If the object store does not provide a last-modified timestamp, the file is conservatively kept.
For dropped regions, unknown files are deleted immediately. | | `region_engine.file` | -- | -- | Enable the file engine. | | `region_engine.metric` | -- | -- | Metric engine options. | -| `region_engine.metric.sparse_primary_key_encoding` | Bool | `true` | Whether to use sparse primary key encoding. | | `logging` | -- | -- | The logging options. | | `logging.dir` | String | `./greptimedb_data/logs` | The directory to store the log files. If set to empty, logs will not be written to files. | | `logging.level` | String | Unset | The log level. Can be `info`/`debug`/`warn`/`error`. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index cfa9f695326..9646936f00e 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -700,8 +700,6 @@ unknown_file_lingering_time = "1d" [[region_engine]] ## Metric engine options. [region_engine.metric] -## Whether to use sparse primary key encoding. -sparse_primary_key_encoding = true ## The logging options. [logging] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 439c947c15c..800543153e3 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -810,8 +810,6 @@ mem_threshold_on_create = "auto" [[region_engine]] ## Metric engine options. [region_engine.metric] -## Whether to use sparse primary key encoding. -sparse_primary_key_encoding = true ## The logging options. [logging] diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 40188dd0489..434c6507b7f 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::io::Write; use std::time::Duration; use cmd::options::GreptimeOptions; @@ -20,6 +21,7 @@ use common_base::readable_size::ReadableSize; use common_config::{Configurable, DEFAULT_DATA_HOME, ENV_VAR_SEP}; use common_options::datanode::{ClientOptions, DatanodeClientOptions}; use common_telemetry::logging::{DEFAULT_LOGGING_DIR, DEFAULT_OTLP_HTTP_ENDPOINT, LoggingOptions}; +use common_test_util::temp_dir::create_named_temp_file; use common_wal::config::DatanodeWalConfig; use common_wal::config::raft_engine::RaftEngineConfig; use datanode::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; @@ -98,7 +100,6 @@ fn test_load_datanode_example_config() { }), RegionEngineConfig::File(FileEngineConfig {}), RegionEngineConfig::Metric(MetricEngineConfig { - sparse_primary_key_encoding: true, flush_metadata_region_interval: Duration::from_secs(30), }), ], @@ -308,7 +309,6 @@ fn test_load_standalone_example_config() { }), RegionEngineConfig::File(FileEngineConfig {}), RegionEngineConfig::Metric(MetricEngineConfig { - sparse_primary_key_encoding: true, flush_metadata_region_interval: Duration::from_secs(30), }), ], @@ -386,3 +386,42 @@ fn test_load_heartbeat_env_vars_from_env() { similar_asserts::assert_eq!(standalone.component.heartbeat_env_vars, expected); }); } + +#[test] +fn test_load_metric_config_with_removed_sparse_primary_key_encoding() { + // The `sparse_primary_key_encoding` option was removed from the metric + // engine config and is now always-on. Existing user configs that still set + // it must be tolerated (ignored), not rejected. + let mut file = create_named_temp_file(); + let toml_str = r#" + node_id = 42 + + [meta_client] + metasrv_addrs = ["127.0.0.1:3002"] + + [[region_engine]] + [region_engine.metric] + sparse_primary_key_encoding = false + flush_metadata_region_interval = "30s" + "#; + write!(file, "{}", toml_str).unwrap(); + + let options = + GreptimeOptions::::load_layered_options(file.path().to_str(), "") + .expect("config with removed 'sparse_primary_key_encoding' should load without error"); + + // The unknown option is ignored; known metric engine options are still parsed/applied. + let metric = options + .component + .region_engine + .iter() + .find_map(|c| match c { + RegionEngineConfig::Metric(c) => Some(c), + _ => None, + }) + .expect("metric engine config should be present"); + assert_eq!( + metric.flush_metadata_region_interval, + Duration::from_secs(30) + ); +} diff --git a/src/metric-engine/src/config.rs b/src/metric-engine/src/config.rs index e342cd9d732..b4e79c34b55 100644 --- a/src/metric-engine/src/config.rs +++ b/src/metric-engine/src/config.rs @@ -23,9 +23,6 @@ pub(crate) const DEFAULT_FLUSH_METADATA_REGION_INTERVAL: Duration = Duration::fr /// Configuration for the metric engine. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct EngineConfig { - /// Whether to use sparse primary key encoding. - #[serde(default = "EngineConfig::default_sparse_primary_key_encoding")] - pub sparse_primary_key_encoding: bool, /// The flush interval of the metadata region. #[serde( with = "humantime_serde", @@ -38,7 +35,6 @@ impl Default for EngineConfig { fn default() -> Self { Self { flush_metadata_region_interval: DEFAULT_FLUSH_METADATA_REGION_INTERVAL, - sparse_primary_key_encoding: Self::default_sparse_primary_key_encoding(), } } } @@ -48,10 +44,6 @@ impl EngineConfig { DEFAULT_FLUSH_METADATA_REGION_INTERVAL } - fn default_sparse_primary_key_encoding() -> bool { - true - } - /// Sanitizes the configuration. pub fn sanitize(&mut self) { if self.flush_metadata_region_interval.is_zero() { diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index f5c7863e825..7d5d74d0fa1 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -479,7 +479,6 @@ impl MetricEngine { metadata_region, data_region, state: state.clone(), - config, row_modifier: RowModifier::default(), flush_task: RepeatedTask::new( flush_interval, @@ -559,11 +558,6 @@ impl MetricEngine { ) -> Result { self.inner.scan_to_stream(region_id, request).await } - - /// Returns the configuration of the engine. - pub fn config(&self) -> &EngineConfig { - &self.inner.config - } } struct MetricEngineInner { @@ -571,7 +565,6 @@ struct MetricEngineInner { metadata_region: MetadataRegion, data_region: DataRegion, state: Arc>, - config: EngineConfig, row_modifier: RowModifier, flush_task: RepeatedTask, } diff --git a/src/metric-engine/src/engine/create.rs b/src/metric-engine/src/engine/create.rs index a5cbd566723..a9b5e16bab3 100644 --- a/src/metric-engine/src/engine/create.rs +++ b/src/metric-engine/src/engine/create.rs @@ -554,10 +554,7 @@ impl MetricEngineInner { data_region_request.primary_key = primary_key; // set data region options - set_data_region_options( - &mut data_region_request.options, - self.config.sparse_primary_key_encoding, - ); + set_data_region_options(&mut data_region_request.options); data_region_request } diff --git a/src/metric-engine/src/engine/open.rs b/src/metric-engine/src/engine/open.rs index 54bbdfd7038..85b538d05b0 100644 --- a/src/metric-engine/src/engine/open.rs +++ b/src/metric-engine/src/engine/open.rs @@ -274,10 +274,7 @@ impl MetricEngineInner { }; let mut data_region_options = request.options; - set_data_region_options( - &mut data_region_options, - self.config.sparse_primary_key_encoding, - ); + set_data_region_options(&mut data_region_options); let clean_up_data_region_request = RegionCleanUpRequest { table_dir: request.table_dir, path_type: PathType::Data, @@ -317,10 +314,7 @@ impl MetricEngineInner { }; let mut data_region_options = request.options; - set_data_region_options( - &mut data_region_options, - self.config.sparse_primary_key_encoding, - ); + set_data_region_options(&mut data_region_options); let open_data_region_request = RegionOpenRequest { table_dir: request.table_dir.clone(), path_type: PathType::Data, diff --git a/src/metric-engine/src/engine/options.rs b/src/metric-engine/src/engine/options.rs index 7008325bc54..a952ffab82d 100644 --- a/src/metric-engine/src/engine/options.rs +++ b/src/metric-engine/src/engine/options.rs @@ -60,10 +60,7 @@ pub enum IndexOptions { } /// Sets data region specific options. -pub fn set_data_region_options( - options: &mut HashMap, - sparse_primary_key_encoding_if_absent: bool, -) { +pub fn set_data_region_options(options: &mut HashMap) { options.remove(METRIC_ENGINE_INDEX_TYPE_OPTION); options.remove(METRIC_ENGINE_INDEX_SKIPPING_INDEX_GRANULARITY_OPTION); options.remove(METRIC_ENGINE_INDEX_SKIPPING_INDEX_FALSE_POSITIVE_RATE_OPTION); @@ -83,11 +80,11 @@ pub fn set_data_region_options( options.insert("sst_format".to_string(), "flat".to_string()); // Decide the top-level primary key encoding: caller-supplied top-level key wins, - // then extracted legacy value, then the `sparse` default if requested. + // then extracted legacy value, then the `sparse` default. if !options.contains_key(PRIMARY_KEY_ENCODING) { if let Some(encoding) = legacy_encoding { options.insert(PRIMARY_KEY_ENCODING.to_string(), encoding); - } else if sparse_primary_key_encoding_if_absent { + } else { options.insert(PRIMARY_KEY_ENCODING.to_string(), "sparse".to_string()); } } @@ -175,7 +172,7 @@ mod tests { METRIC_ENGINE_INDEX_SKIPPING_INDEX_FALSE_POSITIVE_RATE_OPTION.to_string(), "0.01".to_string(), ); - set_data_region_options(&mut options, false); + set_data_region_options(&mut options); for key in [ METRIC_ENGINE_INDEX_TYPE_OPTION, @@ -227,7 +224,7 @@ mod tests { fn test_set_data_region_options_default_compaction_time_window() { // Test that default time window is set when not specified let mut options = HashMap::new(); - set_data_region_options(&mut options, false); + set_data_region_options(&mut options); assert_eq!(options.get("memtable.type"), Some(&"bulk".to_string())); assert_eq!(options.get("sst_format"), Some(&"flat".to_string())); @@ -241,7 +238,7 @@ mod tests { #[test] fn test_set_data_region_options_sparse_primary_key_encoding() { let mut options = HashMap::new(); - set_data_region_options(&mut options, true); + set_data_region_options(&mut options); assert_eq!(options.get("memtable.type"), Some(&"bulk".to_string())); assert_eq!(options.get("sst_format"), Some(&"flat".to_string())); @@ -264,7 +261,7 @@ mod tests { "memtable.partition_tree.index_max_keys_per_shard".to_string(), "2048".to_string(), ); - set_data_region_options(&mut options, false); + set_data_region_options(&mut options); assert_eq!(options.get("memtable.type"), Some(&"bulk".to_string())); assert_eq!(options.get("sst_format"), Some(&"flat".to_string())); @@ -282,7 +279,7 @@ mod tests { let mut options = HashMap::new(); options.insert(PRIMARY_KEY_ENCODING.to_string(), "dense".to_string()); // Sparse flag is on but caller already specified dense. - set_data_region_options(&mut options, true); + set_data_region_options(&mut options); assert_eq!( options.get(PRIMARY_KEY_ENCODING), @@ -296,7 +293,7 @@ mod tests { let mut options = HashMap::new(); options.insert(TWCS_TIME_WINDOW.to_string(), "2h".to_string()); options.insert(COMPACTION_TYPE.to_string(), "twcs".to_string()); - set_data_region_options(&mut options, false); + set_data_region_options(&mut options); // User's time window should be preserved assert_eq!(options.get(TWCS_TIME_WINDOW), Some(&"2h".to_string())); diff --git a/src/metric-engine/src/repeated_task.rs b/src/metric-engine/src/repeated_task.rs index d0299035d95..e3e483ec111 100644 --- a/src/metric-engine/src/repeated_task.rs +++ b/src/metric-engine/src/repeated_task.rs @@ -97,7 +97,6 @@ mod tests { "test_flush_metadata_region_task", EngineConfig { flush_metadata_region_interval: Duration::from_millis(10), - ..Default::default() }, ) .await; @@ -124,7 +123,6 @@ mod tests { "test_flush_metadata_region_task_with_long_interval", EngineConfig { flush_metadata_region_interval: Duration::from_secs(60), - ..Default::default() }, ) .await; @@ -145,18 +143,12 @@ mod tests { ) } - #[tokio::test] - async fn test_flush_metadata_region_sanitize() { - let env = TestEnv::with_prefix_and_config( - "test_flush_metadata_region_sanitize", - EngineConfig { - flush_metadata_region_interval: Duration::from_secs(0), - ..Default::default() - }, - ) - .await; - let metric = env.metric(); - let config = metric.config(); + #[test] + fn test_flush_metadata_region_sanitize() { + let mut config = EngineConfig { + flush_metadata_region_interval: Duration::from_secs(0), + }; + config.sanitize(); assert_eq!( config.flush_metadata_region_interval, DEFAULT_FLUSH_METADATA_REGION_INTERVAL From 67683cef2e445f555edf0506c1c302ac61b5548f Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 15 Jul 2026 19:49:43 +0800 Subject: [PATCH 40/48] feat: support SCRAM auth for Postgres (#8304) * feat: support SCRAM auth for Postgres Signed-off-by: Dennis Zhuang * feat: add pg_scram_sha256 format to hash-password command Signed-off-by: Dennis Zhuang * fix: harden Postgres SCRAM auth - Verify the client-final nonce matches the server-issued nonce, per RFC 5802 transcript validation, instead of only checking the channel-binding field. - Replace the per-connection PBKDF2 over a random password for unknown users with a deterministic mock verifier keyed by the username and a process-wide secret. This avoids a CPU-exhaustion DoS on unknown usernames and removes a username-enumeration oracle: the SCRAM server-first salt and iteration count are now stable per username and indistinguishable from a real user, with no PBKDF2 cost and random keys that never accept a proof. Signed-off-by: Dennis Zhuang * style: format PG_SCRAM_MOCK_SECRET declaration Signed-off-by: Dennis Zhuang * fix: precompute stable SCRAM verifier for plaintext users Plaintext-backed credentials derived a Postgres SCRAM verifier on the fly on every connection, using a fresh random salt and running PBKDF2 each time. That made a known plaintext user distinguishable from stored-hash and unknown (mock) users through both the unstable server-first salt and the per-connection timing, enabling username enumeration. Precompute the SCRAM verifier once at load time (stable salt, default iteration count) and reuse it, matching the mock verifier handed to unknown users. Document that non-default iteration counts remain observable in the SCRAM handshake and weaken enumeration resistance. Signed-off-by: Dennis Zhuang * fix: normalize passwords for Postgres SCRAM Signed-off-by: Dennis Zhuang * chore: docs Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- Cargo.lock | 3 + Cargo.toml | 1 + config/config.md | 4 +- config/frontend.example.toml | 9 +- config/standalone.example.toml | 9 +- src/auth/Cargo.toml | 3 + src/auth/src/common.rs | 276 +++++++++++- src/auth/src/lib.rs | 11 +- src/auth/src/user_provider.rs | 351 ++++++++++++++- .../src/user_provider/static_user_provider.rs | 7 +- .../user_provider/watch_file_user_provider.rs | 10 +- src/cmd/src/user.rs | 45 +- src/servers/src/postgres/auth_handler.rs | 399 ++++++++++++++++-- src/servers/tests/postgres/mod.rs | 105 ++++- 14 files changed, 1164 insertions(+), 69 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a813b9468a8..29640722ba1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -948,11 +948,14 @@ dependencies = [ "common-test-util", "digest 0.10.7", "hex", + "hmac 0.12.1", "pbkdf2", + "rand 0.9.4", "sha1 0.10.6", "sha2 0.10.9", "snafu 0.8.6", "sql", + "stringprep", "subtle", "tokio", ] diff --git a/Cargo.toml b/Cargo.toml index 31093a556c6..12a1d2d346c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -242,6 +242,7 @@ sqlx = { version = "0.8", default-features = false, features = [ "tls-rustls-aws-lc-rs", "rust_decimal", ] } +stringprep = "0.1" strum = { version = "0.27", features = ["derive"] } sysinfo = "0.33" tempfile = "3" diff --git a/config/config.md b/config/config.md index c999e32a9c4..3d406494cfe 100644 --- a/config/config.md +++ b/config/config.md @@ -15,7 +15,7 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index and value columns. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | -| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:"
"pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password. | +| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest.
"mysql_native_password" is MySQL-specific and falls back to cleartext password auth for PostgreSQL.
For PostgreSQL SCRAM, keep all users on the default iteration count (4096) and salt length (16);
a non-default "pbkdf2_sha256"/"pg_scram_sha256" iteration count or salt length is observable in the
SCRAM handshake and weakens resistance to username enumeration. | | `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | | `init_regions_in_background` | Bool | `false` | Initialize all regions in the background during the startup.
By default, it provides services after all regions have been initialized. | @@ -239,7 +239,7 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index and value columns. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | -| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:"
"pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password. | +| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest.
"mysql_native_password" is MySQL-specific and falls back to cleartext password auth for PostgreSQL.
For PostgreSQL SCRAM, keep all users on the default iteration count (4096) and salt length (16);
a non-default "pbkdf2_sha256"/"pg_scram_sha256" iteration count or salt length is observable in the
SCRAM handshake and weakens resistance to username enumeration. | | `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | | `runtime` | -- | -- | The runtime options. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 9fc65ece8af..c856c24b0f2 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -13,8 +13,13 @@ default_column_prefix = "greptime" ## The user provider for authentication. ## Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", -## "mysql_native_password:" -## "pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password. +## "mysql_native_password:", +## "pg_scram_sha256::::" +## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest. +## "mysql_native_password" is MySQL-specific and falls back to cleartext password auth for PostgreSQL. +## For PostgreSQL SCRAM, keep all users on the default iteration count (4096) and salt length (16); +## a non-default "pbkdf2_sha256"/"pg_scram_sha256" iteration count or salt length is observable in the +## SCRAM handshake and weakens resistance to username enumeration. ## @toml2docs:none-default #+ user_provider = "static_user_provider:file:/path/to/users" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 800543153e3..17e830a2c1d 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -13,8 +13,13 @@ default_column_prefix = "greptime" ## The user provider for authentication. ## Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", -## "mysql_native_password:" -## "pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password. +## "mysql_native_password:", +## "pg_scram_sha256::::" +## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest. +## "mysql_native_password" is MySQL-specific and falls back to cleartext password auth for PostgreSQL. +## For PostgreSQL SCRAM, keep all users on the default iteration count (4096) and salt length (16); +## a non-default "pbkdf2_sha256"/"pg_scram_sha256" iteration count or salt length is observable in the +## SCRAM handshake and weakens resistance to username enumeration. ## @toml2docs:none-default #+ user_provider = "static_user_provider:file:/path/to/users" diff --git a/src/auth/Cargo.toml b/src/auth/Cargo.toml index 77a676ea25b..959eb25078e 100644 --- a/src/auth/Cargo.toml +++ b/src/auth/Cargo.toml @@ -21,11 +21,14 @@ common-macro.workspace = true common-telemetry.workspace = true digest = "0.10" hex.workspace = true +hmac = "0.12" pbkdf2 = "0.12" +rand.workspace = true sha1 = "0.10" sha2 = "0.10" snafu.workspace = true sql.workspace = true +stringprep.workspace = true subtle = "2.6" tokio.workspace = true diff --git a/src/auth/src/common.rs b/src/auth/src/common.rs index 3122822a520..630995b2afc 100644 --- a/src/auth/src/common.rs +++ b/src/auth/src/common.rs @@ -12,14 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use common_base::secrets::SecretString; use digest::Digest; +use hmac::{Hmac, Mac}; use pbkdf2::pbkdf2_hmac; use sha1::Sha1; use sha2::Sha256; use snafu::{OptionExt, ensure}; +use subtle::ConstantTimeEq; use crate::error::{IllegalParamSnafu, InvalidConfigSnafu, Result, UserPasswordMismatchSnafu}; use crate::user_info::DefaultUserInfo; @@ -31,9 +33,17 @@ use crate::{UserInfoRef, UserProviderRef}; pub(crate) const DEFAULT_USERNAME: &str = "greptime"; pub const DEFAULT_PBKDF2_SHA256_ITERATIONS: u32 = 4096; +pub const DEFAULT_PBKDF2_SHA256_SALT_LEN: usize = 16; pub const PBKDF2_SHA256_HASH_LEN: usize = 32; pub const MAX_PBKDF2_SHA256_ITERATIONS: u32 = 1_000_000; pub const MAX_PBKDF2_SHA256_SALT_LEN: usize = 1024; +pub const PG_SCRAM_SHA256_KEY_LEN: usize = 32; + +type HmacSha256 = Hmac; + +/// Process-wide secret used to derive a stable mock salt for unknown users, so +/// the SCRAM `server-first-message` can't be used to enumerate usernames. +static PG_SCRAM_MOCK_SECRET: LazyLock<[u8; PG_SCRAM_SHA256_KEY_LEN]> = LazyLock::new(rand::random); /// construct a [`UserInfo`](crate::user_info::UserInfo) impl with name /// use default username `greptime` if None is provided @@ -200,6 +210,212 @@ pub fn format_pbkdf2_sha256_password_verifier( )) } +#[derive(Clone, PartialEq, Eq)] +pub struct PgScramSha256Verifier { + iterations: u32, + salt: Vec, + stored_key: Vec, + server_key: Vec, +} + +impl std::fmt::Debug for PgScramSha256Verifier { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PgScramSha256Verifier") + .field("iterations", &self.iterations) + .field("salt", &"") + .field("stored_key", &"") + .field("server_key", &"") + .finish() + } +} + +impl PgScramSha256Verifier { + pub fn new( + iterations: u32, + salt: Vec, + stored_key: Vec, + server_key: Vec, + ) -> Result { + ensure!( + iterations > 0 && iterations <= MAX_PBKDF2_SHA256_ITERATIONS, + IllegalParamSnafu { + msg: format!( + "pg_scram_sha256 iterations must be in 1..={}", + MAX_PBKDF2_SHA256_ITERATIONS + ) + } + ); + ensure!( + !salt.is_empty() && salt.len() <= MAX_PBKDF2_SHA256_SALT_LEN, + IllegalParamSnafu { + msg: format!( + "pg_scram_sha256 salt length must be in 1..={}", + MAX_PBKDF2_SHA256_SALT_LEN + ) + } + ); + ensure!( + stored_key.len() == PG_SCRAM_SHA256_KEY_LEN, + IllegalParamSnafu { + msg: "pg_scram_sha256 stored key must be 32 bytes" + } + ); + ensure!( + server_key.len() == PG_SCRAM_SHA256_KEY_LEN, + IllegalParamSnafu { + msg: "pg_scram_sha256 server key must be 32 bytes" + } + ); + + Ok(Self { + iterations, + salt, + stored_key, + server_key, + }) + } + + pub fn from_password(password: &[u8], salt: &[u8], iterations: u32) -> Result { + let salted_password = pg_scram_sha256_salted_password(password, salt, iterations)?; + Self::from_salted_password(salted_password, salt.to_vec(), iterations) + } + + pub fn from_salted_password( + salted_password: Vec, + salt: Vec, + iterations: u32, + ) -> Result { + ensure!( + salted_password.len() == PG_SCRAM_SHA256_KEY_LEN, + IllegalParamSnafu { + msg: "pg_scram_sha256 salted password must be 32 bytes" + } + ); + let client_key = hmac_sha256(&salted_password, b"Client Key"); + let stored_key = sha256(&client_key); + let server_key = hmac_sha256(&salted_password, b"Server Key"); + Self::new(iterations, salt, stored_key, server_key) + } + + /// Builds a throwaway verifier for an unknown user. The salt is derived + /// deterministically from the username and a process-wide secret, so the + /// SCRAM `server-first-message` (salt + iterations) stays stable per + /// username and indistinguishable from a real user across reconnects. No + /// PBKDF2 is run (avoids a CPU-exhaustion DoS keyed on unknown usernames), + /// and the random keys guarantee the client proof never matches. + pub fn mock_for_unknown_user(username: &[u8]) -> Self { + let mock_salt = hmac_sha256(PG_SCRAM_MOCK_SECRET.as_slice(), username); + Self { + iterations: DEFAULT_PBKDF2_SHA256_ITERATIONS, + salt: mock_salt[..DEFAULT_PBKDF2_SHA256_SALT_LEN].to_vec(), + stored_key: rand::random::<[u8; PG_SCRAM_SHA256_KEY_LEN]>().to_vec(), + server_key: rand::random::<[u8; PG_SCRAM_SHA256_KEY_LEN]>().to_vec(), + } + } + + pub fn iterations(&self) -> u32 { + self.iterations + } + + pub fn salt(&self) -> &[u8] { + &self.salt + } + + pub fn verify_plain_password(&self, password: &[u8]) -> Result { + let salted_password = + pg_scram_sha256_salted_password(password, &self.salt, self.iterations)?; + let client_key = hmac_sha256(&salted_password, b"Client Key"); + let stored_key = sha256(&client_key); + Ok(self.stored_key.ct_eq(&stored_key).into()) + } + + pub fn verify_client_proof( + &self, + auth_message: &[u8], + client_proof: &[u8], + ) -> Result>> { + if client_proof.len() != PG_SCRAM_SHA256_KEY_LEN { + return Ok(None); + } + + let client_signature = hmac_sha256(&self.stored_key, auth_message); + let client_key = xor(client_proof, &client_signature); + let stored_key = sha256(&client_key); + if self.stored_key.ct_eq(&stored_key).into() { + Ok(Some(hmac_sha256(&self.server_key, auth_message))) + } else { + Ok(None) + } + } +} + +pub fn format_pg_scram_sha256_password_verifier( + password: &[u8], + salt: &[u8], + iterations: u32, +) -> Result { + let verifier = PgScramSha256Verifier::from_password(password, salt, iterations)?; + Ok(format!( + "pg_scram_sha256:{iterations}:{}:{}:{}", + hex::encode(verifier.salt), + hex::encode(verifier.stored_key), + hex::encode(verifier.server_key) + )) +} + +fn pg_scram_sha256_salted_password( + password: &[u8], + salt: &[u8], + iterations: u32, +) -> Result> { + ensure!( + iterations > 0 && iterations <= MAX_PBKDF2_SHA256_ITERATIONS, + IllegalParamSnafu { + msg: format!( + "pg_scram_sha256 iterations must be in 1..={}", + MAX_PBKDF2_SHA256_ITERATIONS + ) + } + ); + ensure!( + !salt.is_empty() && salt.len() <= MAX_PBKDF2_SHA256_SALT_LEN, + IllegalParamSnafu { + msg: format!( + "pg_scram_sha256 salt length must be in 1..={}", + MAX_PBKDF2_SHA256_SALT_LEN + ) + } + ); + + let prepared_password = std::str::from_utf8(password) + .ok() + .and_then(|password| stringprep::saslprep(password).ok()); + let password = prepared_password + .as_deref() + .map(str::as_bytes) + .unwrap_or(password); + + let mut salted_password = [0u8; PG_SCRAM_SHA256_KEY_LEN]; + pbkdf2_hmac::(password, salt, iterations, &mut salted_password); + Ok(salted_password.to_vec()) +} + +fn hmac_sha256(key: &[u8], msg: &[u8]) -> Vec { + let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length"); + mac.update(msg); + mac.finalize().into_bytes().to_vec() +} + +fn sha256(data: &[u8]) -> Vec { + let mut hasher = Sha256::new(); + hasher.update(data); + hasher.finalize().to_vec() +} + +fn xor(lhs: &[u8], rhs: &[u8]) -> Vec { + lhs.iter().zip(rhs).map(|(l, r)| l ^ r).collect() +} + fn sha1_two(input_1: &[u8], input_2: &[u8]) -> Vec { let mut hasher = Sha1::new(); hasher.update(input_1); @@ -274,4 +490,62 @@ mod tests { .is_err() ); } + + #[test] + fn test_format_pg_scram_sha256_password_verifier() { + let verifier = + format_pg_scram_sha256_password_verifier("password".as_bytes(), b"salt", 4096).unwrap(); + assert_eq!( + "pg_scram_sha256:4096:73616c74:945e1c466fc9932efadc23781edc5d1e78d5e10f005933652af1a6105154f084:b9bf0e811b1fb6793671c0cc3adedf7c75cd72291191092ad65878c5a02aad2c", + verifier + ); + } + + #[test] + fn test_pg_scram_sha256_applies_saslprep() { + let normalized = PgScramSha256Verifier::from_password(b"pass word", b"salt", 4096).unwrap(); + let non_breaking_space = + PgScramSha256Verifier::from_password("pass\u{00a0}word".as_bytes(), b"salt", 4096) + .unwrap(); + + assert_eq!(normalized, non_breaking_space); + assert!( + non_breaking_space + .verify_plain_password(b"pass word") + .unwrap() + ); + } + + #[test] + fn test_pg_scram_sha256_uses_original_bytes_when_saslprep_is_not_possible() { + let invalid_utf8 = b"password\xff"; + let prohibited = b"password\x07"; + + for password in [invalid_utf8.as_slice(), prohibited.as_slice()] { + let verifier = PgScramSha256Verifier::from_password(password, b"salt", 4096).unwrap(); + assert!(verifier.verify_plain_password(password).unwrap()); + } + } + + #[test] + fn test_mock_verifier_is_stable_per_username() { + let alice = PgScramSha256Verifier::mock_for_unknown_user(b"alice"); + let alice_again = PgScramSha256Verifier::mock_for_unknown_user(b"alice"); + let bob = PgScramSha256Verifier::mock_for_unknown_user(b"bob"); + + // Same username must yield the same salt and iterations across reconnects, + // otherwise the SCRAM server-first message leaks user (non-)existence. + assert_eq!(alice.salt(), alice_again.salt()); + assert_ne!(alice.salt(), bob.salt()); + assert_eq!(alice.iterations(), DEFAULT_PBKDF2_SHA256_ITERATIONS); + assert_eq!(alice.salt().len(), DEFAULT_PBKDF2_SHA256_SALT_LEN); + + // The mock verifier must never accept a client proof. + assert!( + alice + .verify_client_proof(b"auth-message", &[0u8; PG_SCRAM_SHA256_KEY_LEN]) + .unwrap() + .is_none() + ); + } } diff --git a/src/auth/src/lib.rs b/src/auth/src/lib.rs index 180042eefd8..f1dab981452 100644 --- a/src/auth/src/lib.rs +++ b/src/auth/src/lib.rs @@ -22,16 +22,17 @@ mod user_provider; pub mod tests; pub use common::{ - DEFAULT_PBKDF2_SHA256_ITERATIONS, HashedPassword, Identity, MAX_PBKDF2_SHA256_ITERATIONS, - MAX_PBKDF2_SHA256_SALT_LEN, PBKDF2_SHA256_HASH_LEN, Password, auth_mysql, + DEFAULT_PBKDF2_SHA256_ITERATIONS, DEFAULT_PBKDF2_SHA256_SALT_LEN, HashedPassword, Identity, + MAX_PBKDF2_SHA256_ITERATIONS, MAX_PBKDF2_SHA256_SALT_LEN, PBKDF2_SHA256_HASH_LEN, + PG_SCRAM_SHA256_KEY_LEN, Password, PgScramSha256Verifier, auth_mysql, format_mysql_native_password_verifier, format_pbkdf2_sha256_password_verifier, - mysql_native_password_hash, static_user_provider_from_option, user_provider_from_option, - userinfo_by_name, + format_pg_scram_sha256_password_verifier, mysql_native_password_hash, + static_user_provider_from_option, user_provider_from_option, userinfo_by_name, }; pub use permission::{DefaultPermissionChecker, PermissionChecker, PermissionReq, PermissionResp}; pub use user_info::UserInfo; -pub use user_provider::UserProvider; pub use user_provider::static_user_provider::StaticUserProvider; +pub use user_provider::{PgAuthInfo, UserProvider}; /// pub type alias pub type UserInfoRef = std::sync::Arc; diff --git a/src/auth/src/user_provider.rs b/src/auth/src/user_provider.rs index 2ff61bbe9bc..344e6669015 100644 --- a/src/auth/src/user_provider.rs +++ b/src/auth/src/user_provider.rs @@ -28,8 +28,9 @@ use snafu::{OptionExt, ResultExt, ensure}; use subtle::ConstantTimeEq; use crate::common::{ - Identity, MAX_PBKDF2_SHA256_ITERATIONS, MAX_PBKDF2_SHA256_SALT_LEN, PBKDF2_SHA256_HASH_LEN, - Password, auth_mysql_with_hash_stage_2, + DEFAULT_PBKDF2_SHA256_SALT_LEN, Identity, MAX_PBKDF2_SHA256_ITERATIONS, + MAX_PBKDF2_SHA256_SALT_LEN, PBKDF2_SHA256_HASH_LEN, PG_SCRAM_SHA256_KEY_LEN, Password, + PgScramSha256Verifier, auth_mysql_with_hash_stage_2, }; use crate::error::{ IllegalParamSnafu, InvalidConfigSnafu, IoSnafu, Result, UnsupportedPasswordTypeSnafu, @@ -64,15 +65,35 @@ pub trait UserProvider: Send + Sync { Ok(user_info) } + async fn postgres_auth_info(&self, _id: Identity<'_>) -> Result { + Ok(PgAuthInfo::Cleartext) + } + /// Returns whether this user provider implementation is backed by an external system. fn external(&self) -> bool { false } } -#[derive(Clone, PartialEq, Eq)] +pub enum PgAuthInfo { + ScramSha256 { + verifier: PgScramSha256Verifier, + user_info: Option, + }, + Cleartext, +} + +#[derive(Clone)] pub(crate) enum PasswordVerifier { - PlainText(String), + PlainText { + password: String, + /// SCRAM verifier derived once at load time. Precomputing it keeps the + /// Postgres SCRAM `server-first-message` (stable salt, fixed iterations) + /// and per-connection cost indistinguishable from stored-hash and + /// unknown users, instead of running PBKDF2 with a fresh salt on every + /// connection. + scram: PgScramSha256Verifier, + }, Pbkdf2Sha256 { iterations: u32, salt: Vec, @@ -81,12 +102,46 @@ pub(crate) enum PasswordVerifier { MysqlNativePassword { hash_stage_2: Vec, }, + PgScramSha256(PgScramSha256Verifier), } +impl PartialEq for PasswordVerifier { + fn eq(&self, other: &Self) -> bool { + // The precomputed SCRAM verifier is a cache derived from the password, so + // two plaintext verifiers are equal iff their passwords match. + match (self, other) { + ( + PasswordVerifier::PlainText { password: a, .. }, + PasswordVerifier::PlainText { password: b, .. }, + ) => a == b, + ( + PasswordVerifier::Pbkdf2Sha256 { + iterations: i1, + salt: s1, + hash: h1, + }, + PasswordVerifier::Pbkdf2Sha256 { + iterations: i2, + salt: s2, + hash: h2, + }, + ) => i1 == i2 && s1 == s2 && h1 == h2, + ( + PasswordVerifier::MysqlNativePassword { hash_stage_2: a }, + PasswordVerifier::MysqlNativePassword { hash_stage_2: b }, + ) => a == b, + (PasswordVerifier::PgScramSha256(a), PasswordVerifier::PgScramSha256(b)) => a == b, + _ => false, + } + } +} + +impl Eq for PasswordVerifier {} + impl fmt::Debug for PasswordVerifier { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { - PasswordVerifier::PlainText(_) => { + PasswordVerifier::PlainText { .. } => { f.debug_tuple("PlainText").field(&"").finish() } PasswordVerifier::Pbkdf2Sha256 { iterations, .. } => f @@ -99,14 +154,33 @@ impl fmt::Debug for PasswordVerifier { .debug_struct("MysqlNativePassword") .field("hash_stage_2", &"") .finish(), + PasswordVerifier::PgScramSha256(_) => f + .debug_struct("PgScramSha256") + .field("verifier", &"") + .finish(), } } } impl PasswordVerifier { + /// Builds a plaintext verifier, precomputing a stable SCRAM verifier so the + /// Postgres SCRAM handshake for this user carries a stable salt and runs no + /// per-connection PBKDF2. Uses [`DEFAULT_PBKDF2_SHA256_ITERATIONS`] to match + /// the mock verifier handed to unknown users. + fn plain_text(password: String) -> Option { + let salt = rand::random::<[u8; DEFAULT_PBKDF2_SHA256_SALT_LEN]>(); + let scram = PgScramSha256Verifier::from_password( + password.as_bytes(), + &salt, + crate::DEFAULT_PBKDF2_SHA256_ITERATIONS, + ) + .ok()?; + Some(Self::PlainText { password, scram }) + } + fn parse(input: &str) -> Option { if let Some(password) = input.strip_prefix("plain:") { - return Some(Self::PlainText(password.to_string())); + return Self::plain_text(password.to_string()); } if let Some(verifier) = input.strip_prefix("pbkdf2_sha256:") { @@ -140,14 +214,50 @@ impl PasswordVerifier { return Some(Self::MysqlNativePassword { hash_stage_2 }); } - Some(Self::PlainText(input.to_string())) + if let Some(verifier) = input.strip_prefix("pg_scram_sha256:") { + let mut parts = verifier.split(':'); + let iterations = parts.next()?.parse::().ok()?; + let salt = hex::decode(parts.next()?).ok()?; + let stored_key = hex::decode(parts.next()?).ok()?; + let server_key = hex::decode(parts.next()?).ok()?; + if parts.next().is_some() + || stored_key.len() != PG_SCRAM_SHA256_KEY_LEN + || server_key.len() != PG_SCRAM_SHA256_KEY_LEN + { + return None; + } + + return PgScramSha256Verifier::new(iterations, salt, stored_key, server_key) + .ok() + .map(Self::PgScramSha256); + } + + Self::plain_text(input.to_string()) + } + + fn supports_pg_scram_sha256(&self) -> bool { + matches!( + self, + PasswordVerifier::PlainText { .. } | PasswordVerifier::PgScramSha256(_) + ) + } + + fn to_pg_scram_sha256_verifier(&self) -> Option { + match self { + PasswordVerifier::PlainText { scram, .. } => Some(scram.clone()), + // Legacy PBKDF2 verifiers were derived without SASLprep, and the + // original password is unavailable to normalize them safely. + PasswordVerifier::Pbkdf2Sha256 { .. } => None, + PasswordVerifier::PgScramSha256(verifier) => Some(verifier.clone()), + PasswordVerifier::MysqlNativePassword { .. } => None, + } } fn verify_plain_text(&self, password: &str) -> bool { match self { - PasswordVerifier::PlainText(expected) => { - expected.as_bytes().ct_eq(password.as_bytes()).into() - } + PasswordVerifier::PlainText { + password: expected, .. + } => expected.as_bytes().ct_eq(password.as_bytes()).into(), PasswordVerifier::Pbkdf2Sha256 { iterations, salt, @@ -161,6 +271,9 @@ impl PasswordVerifier { hash.as_slice().ct_eq(&actual[..]).into() } PasswordVerifier::MysqlNativePassword { .. } => false, + PasswordVerifier::PgScramSha256(verifier) => verifier + .verify_plain_password(password.as_bytes()) + .unwrap_or(false), } } @@ -171,7 +284,7 @@ impl PasswordVerifier { username: &str, ) -> Result<()> { match self { - PasswordVerifier::PlainText(password) => { + PasswordVerifier::PlainText { password, .. } => { auth_mysql(auth_data, salt, username, password.as_bytes()) } PasswordVerifier::MysqlNativePassword { hash_stage_2 } => { @@ -181,6 +294,10 @@ impl PasswordVerifier { password_type: "mysql_native_password_with_pbkdf2_sha256_verifier", } .fail(), + PasswordVerifier::PgScramSha256(_) => UnsupportedPasswordTypeSnafu { + password_type: "mysql_native_password_with_pg_scram_sha256_verifier", + } + .fail(), } } } @@ -243,6 +360,7 @@ fn load_credential_from_file(filepath: &str) -> Result { /// - `plain:` /// - `pbkdf2_sha256:::` /// - `mysql_native_password:` +/// - `pg_scram_sha256::::` pub(crate) fn parse_credential_line( line: &str, ) -> Option<(String, (PasswordVerifier, PermissionMode))> { @@ -263,6 +381,53 @@ pub(crate) fn parse_credential_line( Some((username.to_string(), (verifier, permission_mode))) } +pub(crate) fn postgres_auth_info_with_credential( + users: &UserInfoMap, + input_id: Identity<'_>, +) -> Result { + match input_id { + Identity::UserId(username, _) => { + ensure!( + !username.is_empty(), + IllegalParamSnafu { + msg: "blank username" + } + ); + + if !users + .values() + .all(|(verifier, _)| verifier.supports_pg_scram_sha256()) + { + // PostgreSQL chooses one auth method during startup. Selecting it + // per username would expose user or verifier existence. + return Ok(PgAuthInfo::Cleartext); + } + + if let Some((verifier, permission_mode)) = users.get(username) { + if let Some(verifier) = verifier.to_pg_scram_sha256_verifier() { + return Ok(PgAuthInfo::ScramSha256 { + verifier, + user_info: Some(DefaultUserInfo::with_name_and_permission( + username, + *permission_mode, + )), + }); + } + + return Ok(PgAuthInfo::Cleartext); + } + + // Unknown user: hand back a deterministic mock verifier so the SCRAM + // handshake is indistinguishable from a real user, without running + // PBKDF2 or leaking existence through an unstable salt. + Ok(PgAuthInfo::ScramSha256 { + verifier: PgScramSha256Verifier::mock_for_unknown_user(username.as_bytes()), + user_info: None, + }) + } + } +} + fn authenticate_with_credential( users: &UserInfoMap, input_id: Identity<'_>, @@ -317,10 +482,10 @@ mod tests { use sha1::Sha1; use super::*; - use crate::common::mysql_native_password_hash; + use crate::common::{format_pg_scram_sha256_password_verifier, mysql_native_password_hash}; fn plain(password: &str) -> PasswordVerifier { - PasswordVerifier::PlainText(password.to_string()) + PasswordVerifier::plain_text(password.to_string()).unwrap() } fn sha1_one(data: &[u8]) -> Vec { @@ -484,6 +649,23 @@ mod tests { let result = parse_credential_line("user=mysql_native_password:abcd"); assert_eq!(result, None); + let verifier = + format_pg_scram_sha256_password_verifier(b"password", b"salt", 4096).unwrap(); + let result = parse_credential_line(&format!("user={verifier}")); + assert!(matches!( + result, + Some(( + _, + ( + PasswordVerifier::PgScramSha256(PgScramSha256Verifier { .. }), + PermissionMode::ReadWrite + ) + )) + )); + + let result = parse_credential_line("user=pg_scram_sha256:4096:73616c74:abcd:abcd"); + assert_eq!(result, None); + // Invalid format - no equals sign let result = parse_credential_line("invalid_line"); assert_eq!(result, None); @@ -541,7 +723,7 @@ mod tests { let users = HashMap::from([( "user".to_string(), ( - PasswordVerifier::PlainText(password.to_string()), + PasswordVerifier::plain_text(password.to_string()).unwrap(), PermissionMode::default(), ), )]); @@ -585,9 +767,148 @@ mod tests { assert!(result.is_err()); } + #[test] + fn test_authenticate_with_pg_scram_sha256_verifier_plain_text() { + let verifier = + format_pg_scram_sha256_password_verifier(b"password", b"salt", 4096).unwrap(); + let (_, user) = parse_credential_line(&format!("user={verifier}")).unwrap(); + let users = HashMap::from([("user".to_string(), user)]); + + let result = authenticate_with_credential( + &users, + Identity::UserId("user", None), + Password::PlainText("password".to_string().into()), + ); + assert!(result.is_ok()); + + let result = authenticate_with_credential( + &users, + Identity::UserId("user", None), + Password::PlainText("wrong".to_string().into()), + ); + assert!(result.is_err()); + } + + #[test] + fn test_plain_text_scram_verifier_is_stable() { + let verifier = PasswordVerifier::plain_text("password".to_string()).unwrap(); + let first = verifier.to_pg_scram_sha256_verifier().unwrap(); + let second = verifier.to_pg_scram_sha256_verifier().unwrap(); + + // A plaintext-backed user must present a stable salt and fixed iterations + // across connections without re-running PBKDF2, otherwise the SCRAM + // server-first message (and its timing) leaks that the user exists. + assert_eq!(first.salt(), second.salt()); + assert_eq!(first.salt().len(), DEFAULT_PBKDF2_SHA256_SALT_LEN); + assert_eq!(first.iterations(), crate::DEFAULT_PBKDF2_SHA256_ITERATIONS); + + // The derived verifier must still accept the real password. + assert!(first.verify_plain_password(b"password").unwrap()); + assert!(!first.verify_plain_password(b"wrong").unwrap()); + } + + #[test] + fn test_postgres_auth_info_uses_scram_for_unknown_user() { + let verifier = + format_pg_scram_sha256_password_verifier(b"password", b"salt", 4096).unwrap(); + let (_, user) = parse_credential_line(&format!("user={verifier}")).unwrap(); + let users = HashMap::from([("user".to_string(), user)]); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("unknown", None)).unwrap(); + assert!(matches!( + auth_info, + PgAuthInfo::ScramSha256 { + user_info: None, + .. + } + )); + } + + #[test] + fn test_postgres_auth_info_falls_back_to_cleartext() { + let hash_stage_2 = mysql_native_password_hash("password".as_bytes()); + let users = HashMap::from([( + "user".to_string(), + ( + PasswordVerifier::MysqlNativePassword { hash_stage_2 }, + PermissionMode::default(), + ), + )]); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("user", None)).unwrap(); + assert!(matches!(auth_info, PgAuthInfo::Cleartext)); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("unknown", None)).unwrap(); + assert!(matches!(auth_info, PgAuthInfo::Cleartext)); + } + + #[test] + fn test_postgres_auth_info_with_pbkdf2_falls_back_to_cleartext() { + let iterations = 4096; + let salt = b"salt"; + let password = "pass\u{00a0}word"; + let mut hash = [0u8; PBKDF2_SHA256_HASH_LEN]; + pbkdf2_hmac::(password.as_bytes(), salt, iterations, &mut hash); + let users = HashMap::from([( + "user".to_string(), + ( + PasswordVerifier::Pbkdf2Sha256 { + iterations, + salt: salt.to_vec(), + hash: hash.to_vec(), + }, + PermissionMode::default(), + ), + )]); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("user", None)).unwrap(); + assert!(matches!(auth_info, PgAuthInfo::Cleartext)); + + assert!( + authenticate_with_credential( + &users, + Identity::UserId("user", None), + Password::PlainText(password.to_string().into()), + ) + .is_ok() + ); + } + + #[test] + fn test_postgres_auth_info_mixed_verifiers_fall_back_to_cleartext() { + let verifier = + format_pg_scram_sha256_password_verifier(b"password", b"salt", 4096).unwrap(); + let (_, scram_user) = parse_credential_line(&format!("scram={verifier}")).unwrap(); + let mysql_user = ( + PasswordVerifier::MysqlNativePassword { + hash_stage_2: mysql_native_password_hash("password".as_bytes()), + }, + PermissionMode::default(), + ); + let users = HashMap::from([ + ("scram".to_string(), scram_user), + ("mysql".to_string(), mysql_user), + ]); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("scram", None)).unwrap(); + assert!(matches!(auth_info, PgAuthInfo::Cleartext)); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("unknown", None)).unwrap(); + assert!(matches!(auth_info, PgAuthInfo::Cleartext)); + } + #[test] fn test_password_verifier_debug_redacts_secrets() { - let debug = format!("{:?}", PasswordVerifier::PlainText("secret".to_string())); + let debug = format!( + "{:?}", + PasswordVerifier::plain_text("secret".to_string()).unwrap() + ); assert!(debug.contains("")); assert!(!debug.contains("secret")); diff --git a/src/auth/src/user_provider/static_user_provider.rs b/src/auth/src/user_provider/static_user_provider.rs index 689b9627261..4ea7efb4345 100644 --- a/src/auth/src/user_provider/static_user_provider.rs +++ b/src/auth/src/user_provider/static_user_provider.rs @@ -17,7 +17,8 @@ use snafu::OptionExt; use crate::error::{InvalidConfigSnafu, Result}; use crate::user_provider::{ - UserInfoMap, authenticate_with_credential, load_credential_from_file, parse_credential_line, + PgAuthInfo, UserInfoMap, authenticate_with_credential, load_credential_from_file, + parse_credential_line, postgres_auth_info_with_credential, }; use crate::{Identity, Password, UserInfoRef, UserProvider}; @@ -67,6 +68,10 @@ impl UserProvider for StaticUserProvider { authenticate_with_credential(&self.users, id, pwd) } + async fn postgres_auth_info(&self, id: Identity<'_>) -> Result { + postgres_auth_info_with_credential(&self.users, id) + } + async fn authorize( &self, _catalog: &str, diff --git a/src/auth/src/user_provider/watch_file_user_provider.rs b/src/auth/src/user_provider/watch_file_user_provider.rs index 451efd5cc4d..4d39ae3e4bc 100644 --- a/src/auth/src/user_provider/watch_file_user_provider.rs +++ b/src/auth/src/user_provider/watch_file_user_provider.rs @@ -20,7 +20,10 @@ use common_telemetry::{info, warn}; use snafu::ResultExt; use crate::error::{FileWatchSnafu, Result}; -use crate::user_provider::{UserInfoMap, authenticate_with_credential, load_credential_from_file}; +use crate::user_provider::{ + PgAuthInfo, UserInfoMap, authenticate_with_credential, load_credential_from_file, + postgres_auth_info_with_credential, +}; use crate::{Identity, Password, UserInfoRef, UserProvider}; pub(crate) const WATCH_FILE_USER_PROVIDER: &str = "watch_file_user_provider"; @@ -83,6 +86,11 @@ impl UserProvider for WatchFileUserProvider { authenticate_with_credential(&users, id, password) } + async fn postgres_auth_info(&self, id: Identity<'_>) -> Result { + let users = self.users.lock().expect("users credential must be valid"); + postgres_auth_info_with_credential(&users, id) + } + async fn authorize(&self, _: &str, _: &str, _: &UserInfoRef) -> Result<()> { // default allow all Ok(()) diff --git a/src/cmd/src/user.rs b/src/cmd/src/user.rs index 6f7e28f9687..cf7a85b11d2 100644 --- a/src/cmd/src/user.rs +++ b/src/cmd/src/user.rs @@ -17,8 +17,9 @@ use std::io; use auth::{ - DEFAULT_PBKDF2_SHA256_ITERATIONS, MAX_PBKDF2_SHA256_SALT_LEN, + DEFAULT_PBKDF2_SHA256_ITERATIONS, DEFAULT_PBKDF2_SHA256_SALT_LEN, MAX_PBKDF2_SHA256_SALT_LEN, format_mysql_native_password_verifier, format_pbkdf2_sha256_password_verifier, + format_pg_scram_sha256_password_verifier, }; use clap::{ArgGroup, Parser, Subcommand, ValueEnum}; use rand::RngCore; @@ -65,15 +66,15 @@ pub struct HashPasswordCommand { #[clap(long)] password_stdin: bool, - /// PBKDF2-SHA256 iteration count. + /// PBKDF2-SHA256 / SCRAM-SHA-256 iteration count. #[clap(long, default_value_t = DEFAULT_PBKDF2_SHA256_ITERATIONS)] iterations: u32, - /// PBKDF2-SHA256 random salt length in bytes. - #[clap(long, default_value_t = 16)] + /// PBKDF2-SHA256 / SCRAM-SHA-256 random salt length in bytes. + #[clap(long, default_value_t = DEFAULT_PBKDF2_SHA256_SALT_LEN)] salt_len: usize, - /// PBKDF2-SHA256 salt as hex. Mainly useful for deterministic automation. + /// PBKDF2-SHA256 / SCRAM-SHA-256 salt as hex. Mainly useful for deterministic automation. #[clap(long)] salt_hex: Option, } @@ -83,6 +84,7 @@ pub struct HashPasswordCommand { enum PasswordFormat { Pbkdf2Sha256, MysqlNativePassword, + PgScramSha256, } impl HashPasswordCommand { @@ -98,6 +100,16 @@ impl HashPasswordCommand { PasswordFormat::MysqlNativePassword => { format_mysql_native_password_verifier(password.as_bytes()) } + PasswordFormat::PgScramSha256 => { + let salt = self.pbkdf2_salt()?; + format_pg_scram_sha256_password_verifier( + password.as_bytes(), + &salt, + self.iterations, + ) + .map_err(common_error::ext::BoxedError::new) + .context(error::OtherSnafu)? + } }; println!("{verifier}"); @@ -186,6 +198,29 @@ mod tests { ); } + #[test] + fn test_hash_password_command_with_pg_scram_sha256() { + let cmd = HashPasswordCommand { + format: PasswordFormat::PgScramSha256, + password: Some("password".to_string()), + password_stdin: false, + iterations: 4096, + salt_len: 16, + salt_hex: Some("73616c74".to_string()), + }; + + let password = cmd.read_password().unwrap(); + let salt = cmd.pbkdf2_salt().unwrap(); + let verifier = + format_pg_scram_sha256_password_verifier(password.as_bytes(), &salt, cmd.iterations) + .unwrap(); + + assert_eq!( + "pg_scram_sha256:4096:73616c74:945e1c466fc9932efadc23781edc5d1e78d5e10f005933652af1a6105154f084:b9bf0e811b1fb6793671c0cc3adedf7c75cd72291191092ad65878c5a02aad2c", + verifier + ); + } + #[test] fn test_hash_password_command_with_mysql_native_password() { let verifier = format_mysql_native_password_verifier("123456".as_bytes()); diff --git a/src/servers/src/postgres/auth_handler.rs b/src/servers/src/postgres/auth_handler.rs index 286c39da64b..e221f451732 100644 --- a/src/servers/src/postgres/auth_handler.rs +++ b/src/servers/src/postgres/auth_handler.rs @@ -15,20 +15,27 @@ use std::fmt::Debug; use std::sync::Exclusive; -use ::auth::{Identity, Password, UserInfoRef, UserProviderRef, userinfo_by_name}; +use ::auth::{ + Identity, Password, PgAuthInfo, PgScramSha256Verifier, UserInfoRef, UserProviderRef, + userinfo_by_name, +}; use async_trait::async_trait; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64; use common_catalog::parse_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; use common_time::Timezone; use futures::{Sink, SinkExt}; use pgwire::api::auth::StartupHandler; +use pgwire::api::auth::sasl::SCRAM_SHA_256_METHOD; use pgwire::api::{ClientInfo, PgWireConnectionState, auth}; use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::response::ErrorResponse; -use pgwire::messages::startup::{Authentication, SecretKey}; +use pgwire::messages::startup::{Authentication, PasswordMessageFamily, SecretKey}; use pgwire::messages::{PgWireBackendMessage, PgWireFrontendMessage}; use session::Session; use snafu::IntoError; +use tokio::sync::Mutex; use crate::error::{AuthSnafu, Result}; use crate::metrics::METRIC_AUTH_FAILURE; @@ -39,14 +46,34 @@ use crate::query_handler::sql::ServerSqlQueryHandlerRef; pub(crate) struct PgLoginVerifier { user_provider: Option, + state: Mutex, } impl PgLoginVerifier { pub(crate) fn new(user_provider: Option) -> Self { - Self { user_provider } + Self { + user_provider, + state: Mutex::new(PgAuthenticationState::Initial), + } } } +enum PgAuthenticationState { + Initial, + Cleartext, + SaslInitial { + auth_info: PgAuthInfo, + }, + SaslFinal { + verifier: PgScramSha256Verifier, + user_info: Option, + channel_binding: String, + nonce: String, + client_first_bare: String, + server_first: String, + }, +} + #[allow(dead_code)] struct LoginInfo { user: Option, @@ -113,6 +140,57 @@ impl PgLoginVerifier { Ok(user_info) => Ok(Some(user_info)), } } + + async fn postgres_auth_info(&self, login: &LoginInfo) -> Result { + let user_provider = match &self.user_provider { + Some(provider) => provider, + None => return Ok(PgAuthInfo::Cleartext), + }; + + let user_name = match &login.user { + Some(name) => name, + None => return Ok(PgAuthInfo::Cleartext), + }; + + match user_provider + .postgres_auth_info(Identity::UserId(user_name, None)) + .await + { + Err(e) => { + METRIC_AUTH_FAILURE + .with_label_values(&[e.status_code().as_ref()]) + .inc(); + Err(AuthSnafu.into_error(e)) + } + Ok(auth_info) => Ok(auth_info), + } + } + + async fn authorize(&self, login: &LoginInfo, user_info: &UserInfoRef) -> Result<()> { + let user_provider = match &self.user_provider { + Some(provider) => provider, + None => return Ok(()), + }; + + let catalog = match &login.catalog { + Some(name) => name, + None => return Ok(()), + }; + let schema = match &login.schema { + Some(name) => name, + None => return Ok(()), + }; + + match user_provider.authorize(catalog, schema, user_info).await { + Err(e) => { + METRIC_AUTH_FAILURE + .with_label_values(&[e.status_code().as_ref()]) + .inc(); + Err(AuthSnafu.into_error(e)) + } + Ok(()) => Ok(()), + } + } } fn set_client_info(client: &mut C, session: &Session) @@ -190,12 +268,36 @@ impl StartupHandler for PostgresServerHandlerInner { } if self.login_verifier.user_provider.is_some() { + let login_info = LoginInfo::from_client_info(client); + let auth_info = match self.login_verifier.postgres_auth_info(&login_info).await + { + Ok(auth_info) => auth_info, + Err(_) => { + return send_password_authentication_failed(client).await; + } + }; + client.set_state(PgWireConnectionState::AuthenticationInProgress); - client - .send(PgWireBackendMessage::Authentication( - Authentication::CleartextPassword, - )) - .await?; + match auth_info { + PgAuthInfo::ScramSha256 { .. } => { + *self.login_verifier.state.lock().await = + PgAuthenticationState::SaslInitial { auth_info }; + client + .send(PgWireBackendMessage::Authentication(Authentication::SASL( + vec![SCRAM_SHA_256_METHOD.to_string()], + ))) + .await?; + } + PgAuthInfo::Cleartext => { + *self.login_verifier.state.lock().await = + PgAuthenticationState::Cleartext; + client + .send(PgWireBackendMessage::Authentication( + Authentication::CleartextPassword, + )) + .await?; + } + } } else { self.session.set_user_info(userinfo_by_name( client.metadata().get(super::METADATA_USER).cloned(), @@ -205,27 +307,20 @@ impl StartupHandler for PostgresServerHandlerInner { } } PgWireFrontendMessage::PasswordMessageFamily(pwd) => { - // the newer version of pgwire has a few variant password - // message like cleartext/md5 password, saslresponse, etc. Here - // we must manually coerce it into password - let pwd = pwd.into_password()?; - let login_info = LoginInfo::from_client_info(client); - - // do authenticate - let auth_result = self.login_verifier.auth(&login_info, &pwd.password).await; - - if let Ok(Some(user_info)) = auth_result { - self.session.set_user_info(user_info); - set_client_info(client, &self.session); - auth::finish_authentication(client, self.param_provider.as_ref()).await?; - } else { - return send_error( - client, - PgErrorCode::Ec28P01 - .to_err_info("password authentication failed".to_string()), - ) - .await; + match self + .authenticate_password_message(client, &login_info, pwd) + .await? + { + PgAuthenticationResult::Continue => {} + PgAuthenticationResult::Success(user_info) => { + self.session.set_user_info(user_info); + set_client_info(client, &self.session); + auth::finish_authentication(client, self.param_provider.as_ref()).await?; + } + PgAuthenticationResult::Failed => { + return send_password_authentication_failed(client).await; + } } } _ => {} @@ -234,6 +329,201 @@ impl StartupHandler for PostgresServerHandlerInner { } } +enum PgAuthenticationResult { + Continue, + Success(UserInfoRef), + Failed, +} + +impl PostgresServerHandlerInner { + async fn authenticate_password_message( + &self, + client: &mut C, + login_info: &LoginInfo, + pwd: PasswordMessageFamily, + ) -> PgWireResult + where + C: ClientInfo + Sink + Unpin + Send, + C::Error: Debug, + PgWireError: From<>::Error>, + { + let mut state = self.login_verifier.state.lock().await; + let current_state = std::mem::replace(&mut *state, PgAuthenticationState::Initial); + + match current_state { + PgAuthenticationState::Cleartext => { + let pwd = pwd.into_password()?; + drop(state); + match self.login_verifier.auth(login_info, &pwd.password).await { + Ok(Some(user_info)) => Ok(PgAuthenticationResult::Success(user_info)), + Ok(None) | Err(_) => Ok(PgAuthenticationResult::Failed), + } + } + PgAuthenticationState::SaslInitial { auth_info } => { + let sasl_initial = pwd.into_sasl_initial_response()?; + if sasl_initial.auth_method != SCRAM_SHA_256_METHOD { + return Ok(PgAuthenticationResult::Failed); + } + let Some(data) = sasl_initial.data else { + return Ok(PgAuthenticationResult::Failed); + }; + let Some(client_first) = ScramClientFirst::parse(&data) else { + return Ok(PgAuthenticationResult::Failed); + }; + let PgAuthInfo::ScramSha256 { + verifier, + user_info, + } = auth_info + else { + return Ok(PgAuthenticationResult::Failed); + }; + + let server_nonce = BASE64.encode(rand::random::<[u8; 18]>()); + let nonce = format!("{}{}", client_first.nonce, server_nonce); + let server_first = format!( + "r={},s={},i={}", + nonce, + BASE64.encode(verifier.salt()), + verifier.iterations() + ); + client + .send(PgWireBackendMessage::Authentication( + Authentication::SASLContinue(server_first.clone().into()), + )) + .await?; + *state = PgAuthenticationState::SaslFinal { + verifier, + user_info, + channel_binding: client_first.channel_binding, + nonce, + client_first_bare: client_first.bare, + server_first, + }; + Ok(PgAuthenticationResult::Continue) + } + PgAuthenticationState::SaslFinal { + verifier, + user_info, + channel_binding, + nonce, + client_first_bare, + server_first, + } => { + let sasl_response = pwd.into_sasl_response()?; + let Some(client_final) = ScramClientFinal::parse(&sasl_response.data) else { + return Ok(PgAuthenticationResult::Failed); + }; + if client_final.channel_binding != channel_binding { + return Ok(PgAuthenticationResult::Failed); + } + if client_final.nonce != nonce { + return Ok(PgAuthenticationResult::Failed); + } + + let auth_message = format!( + "{},{},{}", + client_first_bare, server_first, client_final.without_proof + ); + let Ok(client_proof) = BASE64.decode(client_final.proof.as_bytes()) else { + return Ok(PgAuthenticationResult::Failed); + }; + let Ok(Some(server_signature)) = + verifier.verify_client_proof(auth_message.as_bytes(), &client_proof) + else { + return Ok(PgAuthenticationResult::Failed); + }; + let Some(user_info) = user_info else { + return Ok(PgAuthenticationResult::Failed); + }; + + drop(state); + if self + .login_verifier + .authorize(login_info, &user_info) + .await + .is_err() + { + return Ok(PgAuthenticationResult::Failed); + } + + client + .send(PgWireBackendMessage::Authentication( + Authentication::SASLFinal( + format!("v={}", BASE64.encode(server_signature)).into(), + ), + )) + .await?; + Ok(PgAuthenticationResult::Success(user_info)) + } + PgAuthenticationState::Initial => Ok(PgAuthenticationResult::Failed), + } + } +} + +struct ScramClientFirst { + channel_binding: String, + bare: String, + nonce: String, +} + +impl ScramClientFirst { + fn parse(data: &[u8]) -> Option { + let message = std::str::from_utf8(data).ok()?; + let mut parts = message.splitn(3, ','); + let cbind = parts.next()?; + if !matches!(cbind, "n" | "y") { + return None; + } + let authzid = parts.next()?; + if !authzid.is_empty() { + return None; + } + let channel_binding = BASE64.encode(format!("{cbind},,")); + let bare = parts.next()?.to_string(); + let nonce = bare + .split(',') + .find_map(|chunk| chunk.strip_prefix("r="))? + .to_string(); + Some(Self { + channel_binding, + bare, + nonce, + }) + } +} + +struct ScramClientFinal { + channel_binding: String, + nonce: String, + without_proof: String, + proof: String, +} + +impl ScramClientFinal { + fn parse(data: &[u8]) -> Option { + let message = std::str::from_utf8(data).ok()?; + let proof_pos = message.rfind(",p=")?; + let without_proof = message[..proof_pos].to_string(); + let proof = message[proof_pos + 3..].to_string(); + let channel_binding = without_proof + .split(',') + .find_map(|chunk| chunk.strip_prefix("c="))?; + let nonce = without_proof + .split(',') + .find_map(|chunk| chunk.strip_prefix("r="))?; + if nonce.is_empty() || proof.is_empty() { + return None; + } + + Some(Self { + channel_binding: channel_binding.to_string(), + nonce: nonce.to_string(), + without_proof, + proof, + }) + } +} + async fn send_error(client: &mut C, err_info: ErrorInfo) -> PgWireResult<()> where C: ClientInfo + Sink + Unpin + Send, @@ -248,6 +538,19 @@ where Ok(()) } +async fn send_password_authentication_failed(client: &mut C) -> PgWireResult<()> +where + C: ClientInfo + Sink + Unpin + Send, + C::Error: Debug, + PgWireError: From<>::Error>, +{ + send_error( + client, + PgErrorCode::Ec28P01.to_err_info("password authentication failed".to_string()), + ) + .await +} + enum DbResolution { Resolved(String, String), NotFound(String), @@ -277,3 +580,43 @@ where Ok(DbResolution::NotFound("Database not specified".to_owned())) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_scram_client_first_parse() { + let message = b"n,,n=greptime,r=clientnonce"; + let client_first = ScramClientFirst::parse(message).unwrap(); + assert_eq!("biws", client_first.channel_binding); + assert_eq!("n=greptime,r=clientnonce", client_first.bare); + assert_eq!("clientnonce", client_first.nonce); + + let message = b"y,,n=greptime,r=clientnonce"; + let client_first = ScramClientFirst::parse(message).unwrap(); + assert_eq!("eSws", client_first.channel_binding); + assert_eq!("n=greptime,r=clientnonce", client_first.bare); + assert_eq!("clientnonce", client_first.nonce); + + assert!(ScramClientFirst::parse(b"p=tls-server-end-point,,n=greptime,r=nonce").is_none()); + assert!(ScramClientFirst::parse(b"n,a=authzid,n=greptime,r=nonce").is_none()); + } + + #[test] + fn test_scram_client_final_parse() { + let message = b"c=biws,r=clientnonceservernonce,p=dGVzdA=="; + let client_final = ScramClientFinal::parse(message).unwrap(); + assert_eq!("biws", client_final.channel_binding); + assert_eq!("clientnonceservernonce", client_final.nonce); + assert_eq!( + "c=biws,r=clientnonceservernonce", + client_final.without_proof + ); + assert_eq!("dGVzdA==", client_final.proof); + + assert!(ScramClientFinal::parse(b"r=nonce,p=dGVzdA==").is_none()); + assert!(ScramClientFinal::parse(b"c=biws,r=nonce").is_none()); + assert!(ScramClientFinal::parse(b"c=biws,r=,p=dGVzdA==").is_none()); + } +} diff --git a/src/servers/tests/postgres/mod.rs b/src/servers/tests/postgres/mod.rs index 5a3e97e3e71..35489687d97 100644 --- a/src/servers/tests/postgres/mod.rs +++ b/src/servers/tests/postgres/mod.rs @@ -16,8 +16,8 @@ use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; -use auth::UserProviderRef; use auth::tests::{DatabaseAuthInfo, MockUserProvider}; +use auth::{UserProviderRef, format_pg_scram_sha256_password_verifier, user_provider_from_option}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_runtime::Builder as RuntimeBuilder; use common_runtime::runtime::BuilderBuild; @@ -43,12 +43,6 @@ fn create_postgres_server( tls: TlsOption, auth_info: Option, ) -> Result> { - let instance = Arc::new(create_testing_instance(table)); - let io_runtime = RuntimeBuilder::default() - .worker_threads(4) - .thread_name("postgres-io-handlers") - .build() - .unwrap(); let user_provider: Option = if check_pwd { let mut provider = MockUserProvider::default(); if let Some(info) = auth_info { @@ -58,6 +52,20 @@ fn create_postgres_server( } else { None }; + create_postgres_server_with_user_provider(table, tls, user_provider) +} + +fn create_postgres_server_with_user_provider( + table: TableRef, + tls: TlsOption, + user_provider: Option, +) -> Result> { + let instance = Arc::new(create_testing_instance(table)); + let io_runtime = RuntimeBuilder::default() + .worker_threads(4) + .thread_name("postgres-io-handlers") + .build() + .unwrap(); let tls_server_config = Arc::new( ReloadableTlsServerConfig::try_new(tls.clone()) @@ -75,6 +83,21 @@ fn create_postgres_server( ))) } +async fn start_test_server_with_user_provider( + user_provider: UserProviderRef, +) -> Result<(Box, u16)> { + common_telemetry::init_default_ut_logging(); + let _ = install_default_crypto_provider(); + + let table = MemTable::default_numbers_table(); + let mut postgres_server = + create_postgres_server_with_user_provider(table, Default::default(), Some(user_provider))?; + let listening = "127.0.0.1:0".parse::().unwrap(); + postgres_server.start(listening).await.unwrap(); + let server_addr = postgres_server.bind_addr().unwrap(); + Ok((postgres_server, server_addr.port())) +} + #[tokio::test] pub async fn test_start_postgres_server() -> Result<()> { let table = MemTable::default_numbers_table(); @@ -138,6 +161,61 @@ async fn test_schema_validating() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_pg_scram_sha256_auth() -> Result<()> { + let verifier = + format_pg_scram_sha256_password_verifier(b"greptime", b"pg-scram-salt", 4096).unwrap(); + let user_provider = + user_provider_from_option(&format!("static_user_provider:cmd:greptime={verifier}")) + .unwrap(); + let (postgres_server, server_port) = + start_test_server_with_user_provider(user_provider).await?; + + let client = create_plain_connection_with_credentials(server_port, "greptime", "greptime") + .await + .unwrap(); + let rows = client + .simple_query("SELECT uint32s FROM numbers LIMIT 1") + .await; + assert_eq!(unwrap_results(rows.unwrap().as_ref())[0], "0"); + + let wrong_password = + create_plain_connection_with_credentials(server_port, "greptime", "wrong").await; + assert!(wrong_password.is_err()); + + let unknown_user = + create_plain_connection_with_credentials(server_port, "not_found", "greptime").await; + assert!(unknown_user.is_err()); + + postgres_server.shutdown().await.unwrap(); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_pg_cleartext_auth_fallback() -> Result<()> { + common_telemetry::init_default_ut_logging(); + let table = MemTable::default_numbers_table(); + let mut postgres_server = create_postgres_server(table, true, Default::default(), None)?; + let listening = "127.0.0.1:0".parse::().unwrap(); + postgres_server.start(listening).await.unwrap(); + let server_port = postgres_server.bind_addr().unwrap().port(); + + let client = create_plain_connection_with_credentials(server_port, "greptime", "greptime") + .await + .unwrap(); + let rows = client + .simple_query("SELECT uint32s FROM numbers LIMIT 1") + .await; + assert_eq!(unwrap_results(rows.unwrap().as_ref())[0], "0"); + + let wrong_password = + create_plain_connection_with_credentials(server_port, "greptime", "wrong").await; + assert!(wrong_password.is_err()); + + postgres_server.shutdown().await.unwrap(); + Ok(()) +} + // #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_shutdown_pg_server(with_pwd: bool) -> Result<()> { common_telemetry::init_default_ut_logging(); @@ -433,6 +511,19 @@ async fn create_plain_connection( Ok(client) } +async fn create_plain_connection_with_credentials( + port: u16, + username: &str, + password: &str, +) -> std::result::Result { + let url = format!( + "host=127.0.0.1 port={port} user={username} password={password} connect_timeout=2 dbname={DEFAULT_SCHEMA_NAME}", + ); + let (client, conn) = tokio_postgres::connect(&url, NoTls).await?; + let _handle = tokio::spawn(conn); + Ok(client) +} + async fn create_connection_with_given_db( port: u16, db: &str, From 887fdf751308ce215be0b883eebd4515f71171ea Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 16 Jul 2026 10:27:09 +0800 Subject: [PATCH 41/48] fix(query): delegate merge sort cardinality effect (#8532) Signed-off-by: discord9 --- src/query/src/dist_plan/merge_sort.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/query/src/dist_plan/merge_sort.rs b/src/query/src/dist_plan/merge_sort.rs index fca81dc4a52..2c8f4d9fd8a 100644 --- a/src/query/src/dist_plan/merge_sort.rs +++ b/src/query/src/dist_plan/merge_sort.rs @@ -21,6 +21,7 @@ use std::fmt; use std::sync::Arc; use datafusion::execution::TaskContext; +use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::metrics::MetricsSet; use datafusion::physical_plan::projection::{ProjectionExec, make_with_child, update_ordering}; use datafusion::physical_plan::sorts::sort::SortExec; @@ -253,6 +254,10 @@ impl ExecutionPlan for MergeSortExec { self.inner.partition_statistics(partition) } + fn cardinality_effect(&self) -> CardinalityEffect { + self.inner.cardinality_effect() + } + /// Intentionally keeps DataFusion's generic limit pushdown disabled. /// /// `MergeSortExec` still supports its own global fetch through From a598690bf7805714917bd824dd5bcdf64ebffd70 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 16 Jul 2026 10:29:43 +0800 Subject: [PATCH 42/48] fix(promql): handle missing labels in or matching (#8504) * fix(promql): handle missing labels in or matching Signed-off-by: discord9 * test(promql): streamline or matching coverage Signed-off-by: discord9 * fix(promql): preserve unmatched rhs series in or Signed-off-by: discord9 * perf(promql): stream union distinct inputs Signed-off-by: discord9 --------- Signed-off-by: discord9 --- .../src/extension_plan/union_distinct_on.rs | 1649 +++++++++++++---- src/promql/src/lib.rs | 2 - src/query/src/promql/planner.rs | 846 ++++++++- .../common/promql/set_operation.result | 24 +- .../promql/tsid_binary_join_regression.result | 24 +- 5 files changed, 2048 insertions(+), 497 deletions(-) diff --git a/src/promql/src/extension_plan/union_distinct_on.rs b/src/promql/src/extension_plan/union_distinct_on.rs index 149475bfe49..a7d8fd48594 100644 --- a/src/promql/src/extension_plan/union_distinct_on.rs +++ b/src/promql/src/extension_plan/union_distinct_on.rs @@ -17,7 +17,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use ahash::{HashMap, RandomState}; +use ahash::{HashSet, RandomState}; use datafusion::arrow::array::UInt64Array; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; @@ -34,14 +34,12 @@ use datafusion::physical_plan::{ }; use datafusion_expr::col; use datatypes::arrow::compute; -use futures::future::BoxFuture; -use futures::{Stream, StreamExt, TryStreamExt, ready}; +use futures::{Stream, StreamExt, ready}; use greptime_proto::substrait_extension as pb; use prost::Message; use snafu::ResultExt; -use crate::error::{DeserializeSnafu, Result}; -use crate::extension_plan::{resolve_column_name, serialize_column_index}; +use crate::error::{DataFusionPlanningSnafu, DeserializeSnafu, Result}; /// A special kind of `UNION`(`OR` in PromQL) operator, for PromQL specific use case. /// @@ -49,35 +47,23 @@ use crate::extension_plan::{resolve_column_name, serialize_column_index}; /// most different part is that it treat left child and right child differently: /// - All columns from left child will be outputted. /// - Only check collisions (when not distinct) on the columns specified by `compare_keys`. -/// - When there is a collision: -/// - If the collision is from right child itself, only the first observed row will be -/// preserved. All others are discarded. -/// - If the collision is from left child, the row in right child will be discarded. -/// - The output order is not maintained. This plan will output left child first, then right child. -/// - The output schema contains all columns from left or right child plans. +/// - Rows from the right child with the same comparison signature are all preserved. +/// - If a signature occurs in the left child, all matching rows from the right child are discarded. +/// - Output preserves each input's batch and row order, with all left rows before right rows. +/// - The output schema is based on the left child schema, with nullability widened from both inputs. /// -/// From the implementation perspective, this operator is similar to `HashJoin`, but the -/// probe side is the right child, and the build side is the left child. Another difference -/// is that the probe is opting-out. -/// -/// This plan will exhaust the right child first to build probe hash table, then streaming -/// on left side, and use the left side to "mask" the hash table. +/// Execution streams the left child first while retaining its comparison signatures. The +/// right child is polled only after the left child completes, then rows whose signatures were +/// observed on the left are omitted. #[derive(Debug, PartialEq, Eq, Hash)] pub struct UnionDistinctOn { left: LogicalPlan, right: LogicalPlan, /// The columns to compare for equality. /// TIME INDEX is included. - compare_keys: Vec, - ts_col: String, + compare_key_indices: Vec, + ts_col_idx: usize, output_schema: DFSchemaRef, - unfix: Option, -} - -#[derive(Debug, PartialEq, Eq, Hash, PartialOrd)] -struct UnfixIndices { - pub compare_key_indices: Vec, - pub ts_col_idx: u64, } impl UnionDistinctOn { @@ -85,21 +71,92 @@ impl UnionDistinctOn { "UnionDistinctOn" } - pub fn new( + pub fn try_new( left: LogicalPlan, right: LogicalPlan, - compare_keys: Vec, - ts_col: String, - output_schema: DFSchemaRef, - ) -> Self { - Self { + compare_key_indices: Vec, + ts_col_idx: usize, + ) -> DataFusionResult { + let output_schema = + Self::validate_children(&left, &right, &compare_key_indices, ts_col_idx)?; + Ok(Self { left, right, - compare_keys, - ts_col, + compare_key_indices, + ts_col_idx, output_schema, - unfix: None, + }) + } + + fn validate_children( + left: &LogicalPlan, + right: &LogicalPlan, + compare_key_indices: &[usize], + ts_col_idx: usize, + ) -> DataFusionResult { + let left_schema = left.schema(); + let right_schema = right.schema(); + let left_fields = left_schema.fields(); + let right_fields = right_schema.fields(); + + if left_fields.len() != right_fields.len() { + return Err(DataFusionError::Plan(format!( + "UnionDistinctOn inputs have different field counts: left={}, right={}", + left_fields.len(), + right_fields.len() + ))); } + + for (column_type, index) in compare_key_indices + .iter() + .map(|index| ("compare key", *index)) + .chain(std::iter::once(("timestamp", ts_col_idx))) + { + if index >= left_fields.len() || index >= right_fields.len() { + return Err(DataFusionError::Plan(format!( + "UnionDistinctOn {column_type} index {index} is out of bounds for inputs with {} fields", + left_fields.len() + ))); + } + } + + for (index, (left_field, right_field)) in left_fields.iter().zip(right_fields).enumerate() { + if left_field.data_type() != right_field.data_type() { + return Err(DataFusionError::Plan(format!( + "UnionDistinctOn input field at index {index} has incompatible data types: left={:?}, right={:?}", + left_field.data_type(), + right_field.data_type() + ))); + } + } + + let output_fields = left_fields + .iter() + .zip(right_fields) + .enumerate() + .map(|(index, (left_field, right_field))| { + let (qualifier, _) = left_schema.qualified_field(index); + ( + qualifier.cloned(), + Arc::new( + left_field + .as_ref() + .clone() + .with_nullable(left_field.is_nullable() || right_field.is_nullable()), + ), + ) + }) + .collect(); + let output_schema = + DFSchema::new_with_metadata(output_fields, left_schema.metadata().clone()).map_err( + |error| { + DataFusionError::Plan(format!( + "Failed to construct UnionDistinctOn output schema: {error}" + )) + }, + )?; + + Ok(Arc::new(output_schema)) } pub fn to_execution_plan( @@ -117,8 +174,8 @@ impl UnionDistinctOn { Arc::new(UnionDistinctOnExec { left: left_exec, right: right_exec, - compare_keys: self.compare_keys.clone(), - ts_col: self.ts_col.clone(), + compare_key_indices: self.compare_key_indices.clone(), + ts_col_idx: self.ts_col_idx, output_schema, metric: ExecutionPlanMetricsSet::new(), properties, @@ -128,12 +185,11 @@ impl UnionDistinctOn { pub fn serialize(&self) -> Vec { let compare_key_indices = self - .compare_keys + .compare_key_indices .iter() - .map(|name| serialize_column_index(&self.output_schema, name)) - .collect::>(); - - let ts_col_idx = serialize_column_index(&self.output_schema, &self.ts_col); + .map(|index| u64::try_from(*index).expect("usize always fits in u64")) + .collect(); + let ts_col_idx = u64::try_from(self.ts_col_idx).expect("usize always fits in u64"); pb::UnionDistinctOn { compare_key_indices, @@ -149,18 +205,33 @@ impl UnionDistinctOn { schema: Arc::new(DFSchema::empty()), }); - let unfix = UnfixIndices { - compare_key_indices: pb_union.compare_key_indices.clone(), - ts_col_idx: pb_union.ts_col_idx, - }; + let compare_key_indices = pb_union + .compare_key_indices + .into_iter() + .map(|index| { + usize::try_from(index).map_err(|_| { + DataFusionError::Plan(format!( + "UnionDistinctOn compare key index {index} does not fit in usize" + )) + }) + }) + .collect::>>() + .context(DataFusionPlanningSnafu)?; + let ts_col_idx = usize::try_from(pb_union.ts_col_idx) + .map_err(|_| { + DataFusionError::Plan(format!( + "UnionDistinctOn timestamp index {} does not fit in usize", + pb_union.ts_col_idx + )) + }) + .context(DataFusionPlanningSnafu)?; Ok(Self { left: placeholder_plan.clone(), right: placeholder_plan, - compare_keys: Vec::new(), - ts_col: String::new(), + compare_key_indices, + ts_col_idx, output_schema: Arc::new(DFSchema::empty()), - unfix: Some(unfix), }) } } @@ -176,11 +247,14 @@ impl PartialOrd for UnionDistinctOn { Some(core::cmp::Ordering::Equal) => {} ord => return ord, } - match self.compare_keys.partial_cmp(&other.compare_keys) { + match self + .compare_key_indices + .partial_cmp(&other.compare_key_indices) + { Some(core::cmp::Ordering::Equal) => {} ord => return ord, } - self.ts_col.partial_cmp(&other.ts_col) + self.ts_col_idx.partial_cmp(&other.ts_col_idx) } } @@ -198,22 +272,21 @@ impl UserDefinedLogicalNodeCore for UnionDistinctOn { } fn expressions(&self) -> Vec { - if self.unfix.is_some() { - return vec![]; - } - - let mut exprs: Vec = self.compare_keys.iter().map(col).collect(); - if !self.compare_keys.iter().any(|key| key == &self.ts_col) { - exprs.push(col(&self.ts_col)); + let fields = self.left.schema().fields(); + let mut exprs = self + .compare_key_indices + .iter() + .filter_map(|index| fields.get(*index).map(|field| col(field.name()))) + .collect::>(); + if !self.compare_key_indices.contains(&self.ts_col_idx) + && let Some(field) = fields.get(self.ts_col_idx) + { + exprs.push(col(field.name())); } exprs } fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option>> { - if self.unfix.is_some() { - return None; - } - let left_len = self.left.schema().fields().len(); let right_len = self.right.schema().fields().len(); Some(vec![ @@ -223,10 +296,20 @@ impl UserDefinedLogicalNodeCore for UnionDistinctOn { } fn fmt_for_explain(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let fields = self.left.schema().fields(); + let display_column = |index: usize| match fields.get(index) { + Some(field) => format!("{}@{index}", field.name()), + None => format!("@{index}"), + }; + let compare_keys = self + .compare_key_indices + .iter() + .map(|index| display_column(*index)) + .collect::>(); write!( f, - "UnionDistinctOn: on col=[{:?}], ts_col=[{}]", - self.compare_keys, self.ts_col + "UnionDistinctOn: on col={compare_keys:?}, ts_col={}", + display_column(self.ts_col_idx) ) } @@ -245,38 +328,12 @@ impl UserDefinedLogicalNodeCore for UnionDistinctOn { let left = inputs.next().unwrap(); let right = inputs.next().unwrap(); - if let Some(unfix) = &self.unfix { - let output_schema = left.schema().clone(); - - let compare_keys = unfix - .compare_key_indices - .iter() - .map(|idx| { - resolve_column_name(*idx, &output_schema, "UnionDistinctOn", "compare key") - }) - .collect::>>()?; - - let ts_col = - resolve_column_name(unfix.ts_col_idx, &output_schema, "UnionDistinctOn", "ts")?; - - Ok(Self { - left, - right, - compare_keys, - ts_col, - output_schema, - unfix: None, - }) - } else { - Ok(Self { - left, - right, - compare_keys: self.compare_keys.clone(), - ts_col: self.ts_col.clone(), - output_schema: self.output_schema.clone(), - unfix: None, - }) - } + Self::try_new( + left, + right, + self.compare_key_indices.clone(), + self.ts_col_idx, + ) } } @@ -284,8 +341,8 @@ impl UserDefinedLogicalNodeCore for UnionDistinctOn { pub struct UnionDistinctOnExec { left: Arc, right: Arc, - compare_keys: Vec, - ts_col: String, + compare_key_indices: Vec, + ts_col_idx: usize, output_schema: SchemaRef, metric: ExecutionPlanMetricsSet, properties: Arc, @@ -326,8 +383,8 @@ impl ExecutionPlan for UnionDistinctOnExec { Ok(Arc::new(UnionDistinctOnExec { left, right, - compare_keys: self.compare_keys.clone(), - ts_col: self.ts_col.clone(), + compare_key_indices: self.compare_key_indices.clone(), + ts_col_idx: self.ts_col_idx, output_schema: self.output_schema.clone(), metric: self.metric.clone(), properties: self.properties.clone(), @@ -341,41 +398,23 @@ impl ExecutionPlan for UnionDistinctOnExec { context: Arc, ) -> DataFusionResult { let left_stream = self.left.execute(partition, context.clone())?; - let right_stream = self.right.execute(partition, context.clone())?; - // Convert column name to column index. Add one for the time column. - let mut key_indices = Vec::with_capacity(self.compare_keys.len() + 1); - for key in &self.compare_keys { - let index = self - .output_schema - .column_with_name(key) - .map(|(i, _)| i) - .ok_or_else(|| DataFusionError::Internal(format!("Column {} not found", key)))?; - key_indices.push(index); - } - let ts_index = self - .output_schema - .column_with_name(&self.ts_col) - .map(|(i, _)| i) - .ok_or_else(|| { - DataFusionError::Internal(format!("Column {} not found", self.ts_col)) - })?; - key_indices.push(ts_index); + let mut key_indices = self.compare_key_indices.clone(); + key_indices.push(self.ts_col_idx); - // Build right hash table future. - let hashed_data_future = HashedDataFut::Pending(Box::pin(HashedData::new( - right_stream, - self.random_state.clone(), - key_indices.clone(), - ))); - - let baseline_metric = BaselineMetrics::new(&self.metric, partition); Ok(Box::pin(UnionDistinctOnStream { left: left_stream, - right: hashed_data_future, + right_plan: self.right.clone(), + right_partition: partition, + right_context: context, + right: None, compare_keys: key_indices, output_schema: self.output_schema.clone(), - metric: baseline_metric, + random_state: self.random_state.clone(), + lhs_signatures: HashSet::default(), + hashes: Vec::new(), + phase: StreamPhase::Left, + metric: BaselineMetrics::new(&self.metric, partition), })) } @@ -396,60 +435,139 @@ impl DisplayAs for UnionDistinctOnExec { | DisplayFormatType::TreeRender => { write!( f, - "UnionDistinctOnExec: on col=[{:?}], ts_col=[{}]", - self.compare_keys, self.ts_col + "UnionDistinctOnExec: on col={:?}, ts_col={}", + self.compare_key_indices, self.ts_col_idx ) } } } } -// TODO(ruihang): some unused fields are for metrics, which will be implemented later. -#[allow(dead_code)] pub struct UnionDistinctOnStream { left: SendableRecordBatchStream, - right: HashedDataFut, + right_plan: Arc, + right_partition: usize, + right_context: Arc, + right: Option, /// Include time index compare_keys: Vec, output_schema: SchemaRef, + random_state: RandomState, + lhs_signatures: HashSet, + hashes: Vec, + phase: StreamPhase, metric: BaselineMetrics, } -impl UnionDistinctOnStream { - fn poll_impl(&mut self, cx: &mut Context<'_>) -> Poll::Item>> { - // resolve the right stream - let right = match self.right { - HashedDataFut::Pending(ref mut fut) => { - let right = ready!(fut.as_mut().poll(cx))?; - self.right = HashedDataFut::Ready(right); - let HashedDataFut::Ready(right_ref) = &mut self.right else { - unreachable!() - }; - right_ref - } - HashedDataFut::Ready(ref mut right) => right, - HashedDataFut::Empty => return Poll::Ready(None), - }; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StreamPhase { + Left, + Right, + Done, +} - // poll left and probe with right - let next_left = ready!(self.left.poll_next_unpin(cx)); - match next_left { - Some(Ok(left)) => { - // observe left batch and return it - right.update_map(&left)?; - Poll::Ready(Some(Ok(left))) +impl UnionDistinctOnStream { + fn hash_batch(&mut self, batch: &RecordBatch) -> DataFusionResult<()> { + let arrays = self + .compare_keys + .iter() + .map(|index| batch.column(*index).clone()) + .collect::>(); + self.hashes.clear(); + self.hashes.resize(batch.num_rows(), 0); + hash_utils::create_hashes(&arrays, &self.random_state, &mut self.hashes)?; + Ok(()) + } + + fn filter_rhs_batch(&mut self, batch: RecordBatch) -> DataFusionResult> { + self.hash_batch(&batch)?; + let mut survivor_indices: Option> = None; + for (index, hash) in self.hashes.iter().enumerate() { + if self.lhs_signatures.contains(hash) { + survivor_indices.get_or_insert_with(|| (0..index).collect()); + } else if let Some(indices) = &mut survivor_indices { + indices.push(index); } - Some(Err(e)) => Poll::Ready(Some(Err(e))), + } + + match survivor_indices { + None => Ok(Some(with_schema(batch, self.output_schema.clone())?)), + Some(indices) if indices.is_empty() => Ok(None), + Some(indices) => Ok(Some(with_schema( + take_batch(&batch, &indices)?, + self.output_schema.clone(), + )?)), + } + } + + fn terminal_error(&mut self, error: DataFusionError) -> Poll::Item>> { + self.phase = StreamPhase::Done; + Poll::Ready(Some(Err(error))) + } + + fn poll_right(&mut self, cx: &mut Context<'_>) -> Poll::Item>> { + if self.right.is_none() { + let right = match self + .right_plan + .execute(self.right_partition, self.right_context.clone()) + { + Ok(right) => right, + Err(error) => return self.terminal_error(error), + }; + self.right = Some(right); + } + + let right = self.right.as_mut().expect("right stream is initialized"); + match ready!(right.poll_next_unpin(cx)) { + Some(Ok(batch)) if self.lhs_signatures.is_empty() => { + match with_schema(batch, self.output_schema.clone()) { + Ok(batch) => Poll::Ready(Some(Ok(batch))), + Err(error) => self.terminal_error(error), + } + } + Some(Ok(batch)) => match self.filter_rhs_batch(batch) { + Ok(Some(batch)) => Poll::Ready(Some(Ok(batch))), + Ok(None) => { + // One fully filtered input batch has been consumed. Yield so a sequence + // of filtered batches cannot monopolize a single downstream poll. + cx.waker().wake_by_ref(); + Poll::Pending + } + Err(error) => self.terminal_error(error), + }, + Some(Err(error)) => self.terminal_error(error), None => { - // left stream is exhausted, so we can send the right part - let right = std::mem::replace(&mut self.right, HashedDataFut::Empty); - let HashedDataFut::Ready(data) = right else { - unreachable!() - }; - Poll::Ready(Some(data.finish())) + self.phase = StreamPhase::Done; + Poll::Ready(None) } } } + + fn poll_impl(&mut self, cx: &mut Context<'_>) -> Poll::Item>> { + match self.phase { + StreamPhase::Left => match ready!(self.left.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + if batch.num_rows() > 0 { + if let Err(error) = self.hash_batch(&batch) { + return self.terminal_error(error); + } + self.lhs_signatures.extend(self.hashes.iter().copied()); + } + match with_schema(batch, self.output_schema.clone()) { + Ok(batch) => Poll::Ready(Some(Ok(batch))), + Err(error) => self.terminal_error(error), + } + } + Some(Err(error)) => self.terminal_error(error), + None => { + self.phase = StreamPhase::Right; + self.poll_right(cx) + } + }, + StreamPhase::Right => self.poll_right(cx), + StreamPhase::Done => Poll::Ready(None), + } + } } impl RecordBatchStream for UnionDistinctOnStream { @@ -462,185 +580,49 @@ impl Stream for UnionDistinctOnStream { type Item = DataFusionResult; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.poll_impl(cx) + let poll = self.poll_impl(cx); + self.metric.record_poll(poll) } } -/// Simple future state for [HashedData] -enum HashedDataFut { - /// The result is not ready - Pending(BoxFuture<'static, DataFusionResult>), - /// The result is ready - Ready(HashedData), - /// The result is taken - Empty, -} - -/// ALL input batches and its hash table -struct HashedData { - // TODO(ruihang): use `JoinHashMap` instead after upgrading to DF 34.0 - /// Hash table for all input batches. The key is hash value, and the value - /// is the index of `bathc`. - hash_map: HashMap, - /// Output batch. - batch: RecordBatch, - /// The indices of the columns to be hashed. - hash_key_indices: Vec, - random_state: RandomState, -} - -impl HashedData { - pub async fn new( - input: SendableRecordBatchStream, - random_state: RandomState, - hash_key_indices: Vec, - ) -> DataFusionResult { - // Collect all batches from the input stream - let initial = (Vec::new(), 0); - let schema = input.schema(); - let (batches, _num_rows) = input - .try_fold(initial, |mut acc, batch| async { - // Update rowcount - acc.1 += batch.num_rows(); - // Push batch to output - acc.0.push(batch); - Ok(acc) - }) - .await?; - - // Create hash for each batch - let mut hash_map = HashMap::default(); - let mut hashes_buffer = Vec::new(); - let mut interleave_indices = Vec::new(); - for (batch_number, batch) in batches.iter().enumerate() { - hashes_buffer.resize(batch.num_rows(), 0); - // get columns for hashing - let arrays = hash_key_indices - .iter() - .map(|i| batch.column(*i).clone()) - .collect::>(); - - // compute hash - let hash_values = - hash_utils::create_hashes(&arrays, &random_state, &mut hashes_buffer)?; - for (row_number, hash_value) in hash_values.iter().enumerate() { - // Only keeps the first observed row for each hash value - if hash_map - .try_insert(*hash_value, interleave_indices.len()) - .is_ok() - { - interleave_indices.push((batch_number, row_number)); - } - } - } - - // Finalize the hash map - let batch = interleave_batches(schema, batches, interleave_indices)?; - - Ok(Self { - hash_map, - batch, - hash_key_indices, - random_state, - }) - } - - /// Remove rows that hash value present in the input - /// record batch from the hash map. - pub fn update_map(&mut self, input: &RecordBatch) -> DataFusionResult<()> { - // get columns for hashing - let mut hashes_buffer = Vec::new(); - let arrays = self - .hash_key_indices - .iter() - .map(|i| input.column(*i).clone()) - .collect::>(); - - // compute hash - hashes_buffer.resize(input.num_rows(), 0); - let hash_values = - hash_utils::create_hashes(&arrays, &self.random_state, &mut hashes_buffer)?; - - // remove those hashes - for hash in hash_values { - self.hash_map.remove(hash); - } - - Ok(()) - } - - pub fn finish(self) -> DataFusionResult { - let valid_indices = self.hash_map.values().copied().collect::>(); - let result = take_batch(&self.batch, &valid_indices)?; - Ok(result) - } -} - -/// Utility function to interleave batches. Based on [interleave](datafusion::arrow::compute::interleave) -fn interleave_batches( - schema: SchemaRef, - batches: Vec, - indices: Vec<(usize, usize)>, -) -> DataFusionResult { - if batches.is_empty() { - if indices.is_empty() { - return Ok(RecordBatch::new_empty(schema)); - } else { - return Err(DataFusionError::Internal( - "Cannot interleave empty batches with non-empty indices".to_string(), - )); - } - } - - // transform batches into arrays - let mut arrays = vec![vec![]; schema.fields().len()]; - for batch in &batches { - for (i, array) in batch.columns().iter().enumerate() { - arrays[i].push(array.as_ref()); - } - } - - // interleave arrays - let interleaved_arrays: Vec<_> = arrays - .into_iter() - .map(|array| compute::interleave(&array, &indices)) - .collect::>()?; - - // assemble new record batch - RecordBatch::try_new(schema, interleaved_arrays) +fn with_schema(batch: RecordBatch, schema: SchemaRef) -> DataFusionResult { + RecordBatch::try_new(schema, batch.columns().to_vec()) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) } -/// Utility function to take rows from a record batch. Based on [take](datafusion::arrow::compute::take) +/// Utility function to take ordered rows from a record batch. fn take_batch(batch: &RecordBatch, indices: &[usize]) -> DataFusionResult { - // fast path if batch.num_rows() == indices.len() { return Ok(batch.clone()); } - let schema = batch.schema(); - - let indices_array = UInt64Array::from_iter(indices.iter().map(|i| *i as u64)); + let indices_array = UInt64Array::from_iter(indices.iter().map(|index| *index as u64)); let arrays = batch .columns() .iter() .map(|array| compute::take(array, &indices_array, None)) .collect::, _>>() - .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + .map_err(|error| DataFusionError::ArrowError(Box::new(error), None))?; - let result = RecordBatch::try_new(schema, arrays) - .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; - Ok(result) + RecordBatch::try_new(batch.schema(), arrays) + .map_err(|error| DataFusionError::ArrowError(Box::new(error), None)) } #[cfg(test)] mod test { + use std::collections::VecDeque; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::task::{Wake, Waker}; - use datafusion::arrow::array::Int32Array; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::array::{Array, Float64Array, Int32Array, Int64Array, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::common::ToDFSchema; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; + use datafusion::prelude::SessionContext; + use futures::StreamExt; use super::*; @@ -660,13 +642,7 @@ mod test { produce_one_row: false, schema: df_schema.clone(), }); - let plan = UnionDistinctOn::new( - left, - right, - vec!["k".to_string()], - "ts".to_string(), - df_schema, - ); + let plan = UnionDistinctOn::try_new(left, right, vec![1], 0).unwrap(); // Simulate a parent projection requesting only one output column. let output_columns = [2usize]; @@ -676,56 +652,6 @@ mod test { assert_eq!(required[1].as_slice(), &[0, 1, 2]); } - #[test] - fn test_interleave_batches() { - let schema = Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - ]); - - let batch1 = RecordBatch::try_new( - Arc::new(schema.clone()), - vec![ - Arc::new(Int32Array::from(vec![1, 2, 3])), - Arc::new(Int32Array::from(vec![4, 5, 6])), - ], - ) - .unwrap(); - - let batch2 = RecordBatch::try_new( - Arc::new(schema.clone()), - vec![ - Arc::new(Int32Array::from(vec![7, 8, 9])), - Arc::new(Int32Array::from(vec![10, 11, 12])), - ], - ) - .unwrap(); - - let batch3 = RecordBatch::try_new( - Arc::new(schema.clone()), - vec![ - Arc::new(Int32Array::from(vec![13, 14, 15])), - Arc::new(Int32Array::from(vec![16, 17, 18])), - ], - ) - .unwrap(); - - let batches = vec![batch1, batch2, batch3]; - let indices = vec![(0, 0), (1, 0), (2, 0), (0, 1), (1, 1), (2, 1)]; - let result = interleave_batches(Arc::new(schema.clone()), batches, indices).unwrap(); - - let expected = RecordBatch::try_new( - Arc::new(schema), - vec![ - Arc::new(Int32Array::from(vec![1, 7, 13, 2, 8, 14])), - Arc::new(Int32Array::from(vec![4, 10, 16, 5, 11, 17])), - ], - ) - .unwrap(); - - assert_eq!(result, expected); - } - #[test] fn test_take_batch() { let schema = Schema::new(vec![ @@ -757,39 +683,958 @@ mod test { assert_eq!(result, expected); } + fn empty_plan(schema: datafusion::common::DFSchemaRef) -> LogicalPlan { + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema, + }) + } + + fn schemas(left: SchemaRef, right: SchemaRef) -> (LogicalPlan, LogicalPlan) { + ( + empty_plan(left.to_dfschema_ref().unwrap()), + empty_plan(right.to_dfschema_ref().unwrap()), + ) + } + + fn schema(prefix: &str, key: &str, nullable: bool) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(format!("{prefix}_ts"), DataType::Int64, false), + Field::new(format!("{prefix}_{key}"), DataType::Utf8, nullable), + Field::new(format!("{prefix}_value"), DataType::Float64, nullable), + ])) + } + + fn batch(schema: SchemaRef, ts: i64, key: Option<&str>, value: Option) -> RecordBatch { + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![ts])), + Arc::new(StringArray::from(vec![key])), + Arc::new(Float64Array::from(vec![value])), + ], + ) + .unwrap() + } + + fn source_exec(batch: RecordBatch) -> Arc { + source_exec_batches(batch.schema(), vec![batch]) + } + + async fn execute( + plan: &UnionDistinctOn, + left: RecordBatch, + right: RecordBatch, + ) -> Vec { + datafusion::physical_plan::collect( + plan.to_execution_plan(source_exec(left), source_exec(right)), + SessionContext::default().task_ctx(), + ) + .await + .unwrap() + } + + fn source_exec_batches(schema: SchemaRef, batches: Vec) -> Arc { + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[batches], schema, None).unwrap(), + ))) + } + + async fn execute_batches( + plan: &UnionDistinctOn, + left_schema: SchemaRef, + left: Vec, + right_schema: SchemaRef, + right: Vec, + ) -> Vec { + datafusion::physical_plan::collect( + plan.to_execution_plan( + source_exec_batches(left_schema, left), + source_exec_batches(right_schema, right), + ), + SessionContext::default().task_ctx(), + ) + .await + .unwrap() + } + + fn simple_schema(prefix: &str, nullable: bool) -> SchemaRef { + schema(prefix, "label", nullable) + } + + fn simple_batch(schema: SchemaRef, rows: &[(i64, &str, f64)]) -> RecordBatch { + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from_iter_values(rows.iter().map(|row| row.0))), + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.1))), + Arc::new(Float64Array::from_iter_values(rows.iter().map(|row| row.2))), + ], + ) + .unwrap() + } + + fn simple_rows(batches: &[RecordBatch]) -> Vec<(i64, String, f64)> { + batches + .iter() + .flat_map(|batch| { + let timestamps = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let labels = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(move |row| { + ( + timestamps.value(row), + labels.value(row).to_string(), + values.value(row), + ) + }) + }) + .collect() + } + + fn simple_plan(left_schema: SchemaRef, right_schema: SchemaRef) -> UnionDistinctOn { + let (left, right) = schemas(left_schema, right_schema); + UnionDistinctOn::try_new(left, right, vec![1], 0).unwrap() + } + + #[derive(Clone, Debug)] + enum TestEvent { + Batch(RecordBatch), + Error(&'static str), + } + + struct TestStream { + schema: SchemaRef, + events: VecDeque, + polls: Arc, + } + + impl Stream for TestStream { + type Item = DataFusionResult; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.polls.fetch_add(1, Ordering::SeqCst); + match self.events.pop_front() { + Some(TestEvent::Batch(batch)) => Poll::Ready(Some(Ok(batch))), + Some(TestEvent::Error(message)) => { + Poll::Ready(Some(Err(DataFusionError::Internal(message.to_string())))) + } + None => Poll::Ready(None), + } + } + } + + impl RecordBatchStream for TestStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + } + + struct CountingWake(AtomicUsize); + + impl Wake for CountingWake { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } + + fn wake_by_ref(self: &Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + fn test_stream( + schema: SchemaRef, + events: Vec, + polls: Arc, + ) -> SendableRecordBatchStream { + Box::pin(TestStream { + schema, + events: events.into(), + polls, + }) + } + + #[derive(Debug)] + struct TestExec { + schema: SchemaRef, + events: Vec, + polls: Arc, + executions: Arc, + execute_error: Option<&'static str>, + properties: Arc, + } + + impl TestExec { + fn new( + schema: SchemaRef, + events: Vec, + polls: Arc, + executions: Arc, + ) -> Self { + Self { + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )), + schema, + events, + polls, + executions, + execute_error: None, + } + } + + fn with_execute_error(mut self, error: &'static str) -> Self { + self.execute_error = Some(error); + self + } + } + + impl DisplayAs for TestExec { + fn fmt_as(&self, _t: DisplayFormatType, _f: &mut std::fmt::Formatter) -> std::fmt::Result { + Ok(()) + } + } + + impl ExecutionPlan for TestExec { + fn name(&self) -> &str { + "TestExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DataFusionResult> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> DataFusionResult { + self.executions.fetch_add(1, Ordering::SeqCst); + if let Some(error) = self.execute_error { + return Err(DataFusionError::Internal(error.to_string())); + } + Ok(test_stream( + self.schema.clone(), + self.events.clone(), + self.polls.clone(), + )) + } + } + + fn test_union_stream( + left: SendableRecordBatchStream, + right: Arc, + output_schema: SchemaRef, + ) -> UnionDistinctOnStream { + let metrics = ExecutionPlanMetricsSet::new(); + UnionDistinctOnStream { + left, + right_plan: right, + right_partition: 0, + right_context: SessionContext::default().task_ctx(), + right: None, + compare_keys: vec![1, 0], + output_schema, + random_state: RandomState::new(), + lhs_signatures: HashSet::default(), + hashes: Vec::new(), + phase: StreamPhase::Left, + metric: BaselineMetrics::new(&metrics, 0), + } + } + + fn series_rows(batches: &[RecordBatch]) -> Vec<(i64, &str, f64, &str)> { + batches + .iter() + .flat_map(|batch| { + let timestamps = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let series = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let keys = batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(move |row| { + ( + timestamps.value(row), + series.value(row), + values.value(row), + keys.value(row), + ) + }) + }) + .collect() + } + #[tokio::test] - async fn encode_decode_union_distinct_on() { + async fn serialize_deserialize_and_execute_with_different_input_names() { + let (left_schema, right_schema) = + (schema("left", "job", false), schema("right", "job", false)); + let (left_plan, right_plan) = schemas(left_schema.clone(), right_schema.clone()); + let decoded = UnionDistinctOn::deserialize( + &UnionDistinctOn::try_new(left_plan.clone(), right_plan.clone(), vec![1], 0) + .unwrap() + .serialize(), + ) + .unwrap() + .with_exprs_and_inputs(vec![], vec![left_plan, right_plan]) + .unwrap(); + assert_eq!( + (decoded.compare_key_indices.as_slice(), decoded.ts_col_idx), + (&[1usize][..], 0) + ); + assert_eq!( + decoded.output_schema, + left_schema.clone().to_dfschema_ref().unwrap() + ); + let result = execute( + &decoded, + batch(left_schema.clone(), 1, Some("left"), Some(10.0)), + batch(right_schema, 2, Some("right"), Some(20.0)), + ) + .await; + assert!(result.iter().all(|batch| batch.schema() == left_schema)); + } + + #[tokio::test] + async fn execute_widens_nullable_fields_and_emits_rhs_nulls() { + let (left_schema, right_schema) = ( + schema("left", "label", false), + schema("right", "label", true), + ); + let (left_plan, right_plan) = schemas(left_schema.clone(), right_schema.clone()); + let plan = UnionDistinctOn::try_new(left_plan, right_plan, vec![1, 2], 0).unwrap(); + let declared_schema = plan.output_schema.inner().clone(); + assert!( + declared_schema.fields()[1..] + .iter() + .all(|field| field.is_nullable()) + ); + let result = execute( + &plan, + batch(left_schema, 1, Some("present"), Some(10.0)), + batch(right_schema, 2, None, None), + ) + .await; + assert_eq!(result.len(), 2); + assert!(result.iter().all(|batch| batch.schema() == declared_schema)); + assert!(result[1].column(1).is_null(0)); + assert!(result[1].column(2).is_null(0)); + } + + #[tokio::test] + async fn empty_lhs_preserves_distinct_rhs_series_with_same_normalized_key() { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Int64, false), + Field::new("series", DataType::Utf8, false), + Field::new("value", DataType::Float64, false), + Field::new("__normalized_absent_label", DataType::Utf8, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("rhs_ts", DataType::Int64, false), + Field::new("rhs_series", DataType::Utf8, false), + Field::new("rhs_value", DataType::Float64, false), + Field::new("rhs_normalized_absent_label", DataType::Utf8, false), + ])); + let (left, right) = schemas(left_schema.clone(), right_schema.clone()); + let plan = UnionDistinctOn::try_new(left, right, vec![3], 0).unwrap(); + let declared_schema = plan.output_schema.inner().clone(); + let rhs = RecordBatch::try_new( + right_schema, + vec![ + Arc::new(Int64Array::from(vec![1_000, 1_000])), + Arc::new(StringArray::from(vec!["series_a", "series_b"])), + Arc::new(Float64Array::from(vec![10.0, 20.0])), + Arc::new(StringArray::from(vec!["", ""])), + ], + ) + .unwrap(); + + let result = execute(&plan, RecordBatch::new_empty(left_schema), rhs).await; + assert!(result.iter().all(|batch| batch.schema() == declared_schema)); + assert_eq!(result.iter().map(RecordBatch::num_rows).sum::(), 2); + + let mut rows = series_rows(&result); + rows.sort_by_key(|row| row.1); + assert_eq!( + rows, + vec![(1_000, "series_a", 10.0, ""), (1_000, "series_b", 20.0, "")] + ); + } + + #[tokio::test] + async fn lhs_signature_suppresses_all_matching_rhs_rows_and_preserves_lhs() { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Int64, false), + Field::new("series", DataType::Utf8, false), + Field::new("value", DataType::Float64, false), + Field::new("__normalized_absent_label", DataType::Utf8, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("rhs_ts", DataType::Int64, false), + Field::new("rhs_series", DataType::Utf8, false), + Field::new("rhs_value", DataType::Float64, false), + Field::new("rhs_normalized_absent_label", DataType::Utf8, false), + ])); + let (left, right) = schemas(left_schema.clone(), right_schema.clone()); + let plan = UnionDistinctOn::try_new(left, right, vec![3], 0).unwrap(); + let declared_schema = plan.output_schema.inner().clone(); + let lhs = RecordBatch::try_new( + left_schema, + vec![ + Arc::new(Int64Array::from(vec![1_000, 2_000])), + Arc::new(StringArray::from(vec!["lhs_match", "lhs_only"])), + Arc::new(Float64Array::from(vec![1.0, 2.0])), + Arc::new(StringArray::from(vec!["", "left-only"])), + ], + ) + .unwrap(); + let rhs = RecordBatch::try_new( + right_schema, + vec![ + Arc::new(Int64Array::from(vec![1_000, 1_000, 3_000])), + Arc::new(StringArray::from(vec![ + "rhs_match_a", + "rhs_match_b", + "rhs_unmatched", + ])), + Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0])), + Arc::new(StringArray::from(vec!["", "", ""])), + ], + ) + .unwrap(); + + let result = execute(&plan, lhs, rhs).await; + assert_eq!(result.iter().map(RecordBatch::num_rows).sum::(), 3); + assert!(result.iter().all(|batch| batch.schema() == declared_schema)); + let mut rows = series_rows(&result); + rows.sort_by_key(|row| row.1); + assert_eq!( + rows, + vec![ + (1_000, "lhs_match", 1.0, ""), + (2_000, "lhs_only", 2.0, "left-only"), + (3_000, "rhs_unmatched", 30.0, ""), + ] + ); + } + + #[tokio::test] + async fn empty_lhs_preserves_duplicate_rhs_rows_across_batches() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let result = execute_batches( + &plan, + left_schema, + vec![], + right_schema.clone(), + vec![ + simple_batch( + right_schema.clone(), + &[(1, "same", 10.0), (1, "same", 11.0)], + ), + simple_batch(right_schema, &[(1, "same", 12.0)]), + ], + ) + .await; + assert_eq!( + simple_rows(&result), + vec![ + (1, "same".to_string(), 10.0), + (1, "same".to_string(), 11.0), + (1, "same".to_string(), 12.0), + ] + ); + } + + #[tokio::test] + async fn lhs_signature_suppresses_rhs_duplicates_across_batches() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let result = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema, &[(1, "match", 1.0)])], + right_schema.clone(), + vec![ + simple_batch(right_schema.clone(), &[(1, "match", 10.0)]), + simple_batch(right_schema, &[(1, "match", 11.0)]), + ], + ) + .await; + assert_eq!(simple_rows(&result), vec![(1, "match".to_string(), 1.0)]); + } + + #[tokio::test] + async fn hash_scratch_reset_suppresses_null_label_after_same_sized_lhs_batches() { + let left_schema = simple_schema("left", true); + let right_schema = simple_schema("right", true); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let lhs_poison = RecordBatch::try_new( + left_schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1])), + Arc::new(StringArray::from(vec![Some("poison")])), + Arc::new(Float64Array::from(vec![1.0])), + ], + ) + .unwrap(); + let lhs_null = RecordBatch::try_new( + left_schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![42])), + Arc::new(StringArray::from(vec![None::<&str>])), + Arc::new(Float64Array::from(vec![2.0])), + ], + ) + .unwrap(); + let rhs_null = RecordBatch::try_new( + right_schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![42])), + Arc::new(StringArray::from(vec![None::<&str>])), + Arc::new(Float64Array::from(vec![999.0])), + ], + ) + .unwrap(); + + let result = execute_batches( + &plan, + left_schema, + vec![lhs_poison, lhs_null], + right_schema, + vec![rhs_null], + ) + .await; + assert_eq!(result.iter().map(RecordBatch::num_rows).sum::(), 2); + let values = result + .iter() + .flat_map(|batch| { + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!(values, vec![1.0, 2.0]); + } + + #[tokio::test] + async fn mixed_rhs_duplicates_retain_unmatched_rows_in_order() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let result = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema, &[(1, "match", 1.0)])], + right_schema.clone(), + vec![simple_batch( + right_schema, + &[ + (1, "match", 10.0), + (1, "keep", 11.0), + (1, "keep", 12.0), + (1, "match", 13.0), + (1, "later", 14.0), + ], + )], + ) + .await; + assert_eq!( + simple_rows(&result), + vec![ + (1, "match".to_string(), 1.0), + (1, "keep".to_string(), 11.0), + (1, "keep".to_string(), 12.0), + (1, "later".to_string(), 14.0), + ] + ); + } + + #[tokio::test] + async fn same_label_at_different_timestamps_is_unmatched() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let result = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema, &[(1, "label", 1.0)])], + right_schema.clone(), + vec![simple_batch(right_schema, &[(2, "label", 2.0)])], + ) + .await; + assert_eq!( + simple_rows(&result), + vec![(1, "label".to_string(), 1.0), (2, "label".to_string(), 2.0)] + ); + } + + #[tokio::test] + async fn empty_input_combinations_complete_cleanly() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + + let empty_left = execute_batches( + &plan, + left_schema.clone(), + vec![], + right_schema.clone(), + vec![simple_batch(right_schema.clone(), &[(1, "rhs", 1.0)])], + ) + .await; + assert_eq!(simple_rows(&empty_left), vec![(1, "rhs".to_string(), 1.0)]); + + let empty_right = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema.clone(), &[(1, "lhs", 1.0)])], + right_schema.clone(), + vec![], + ) + .await; + assert_eq!(simple_rows(&empty_right), vec![(1, "lhs".to_string(), 1.0)]); + + let both_empty = execute_batches(&plan, left_schema, vec![], right_schema, vec![]).await; + assert!(both_empty.is_empty()); + } + + #[tokio::test] + async fn zero_row_batches_on_either_side_are_handled() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + + let zero_row_left = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema.clone(), &[])], + right_schema.clone(), + vec![simple_batch(right_schema.clone(), &[(1, "rhs", 1.0)])], + ) + .await; + assert_eq!( + simple_rows(&zero_row_left), + vec![(1, "rhs".to_string(), 1.0)] + ); + + let zero_row_right = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema, &[(1, "lhs", 1.0)])], + right_schema.clone(), + vec![simple_batch(right_schema, &[])], + ) + .await; + assert_eq!( + simple_rows(&zero_row_right), + vec![(1, "lhs".to_string(), 1.0)] + ); + } + + #[tokio::test] + async fn metrics_record_output_rows_and_completion() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let exec = plan.to_execution_plan( + source_exec(simple_batch(left_schema, &[(1, "lhs", 1.0)])), + source_exec(simple_batch(right_schema, &[(2, "rhs", 2.0)])), + ); + + let output = + datafusion::physical_plan::collect(exec.clone(), SessionContext::default().task_ctx()) + .await + .unwrap(); + assert_eq!(simple_rows(&output).len(), 2); + + let metrics = exec.metrics().unwrap(); + assert_eq!(metrics.output_rows(), Some(2)); + assert!(metrics.iter().any(|metric| { + metric.value().name() == "end_timestamp" && metric.value().as_usize() > 0 + })); + } + + #[tokio::test] + async fn output_keeps_lhs_before_rhs_and_input_order() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let result = execute_batches( + &plan, + left_schema.clone(), + vec![ + simple_batch(left_schema.clone(), &[(1, "lhs-a", 1.0)]), + simple_batch(left_schema, &[(2, "lhs-b", 2.0)]), + ], + right_schema.clone(), + vec![ + simple_batch(right_schema.clone(), &[(3, "rhs-a", 3.0)]), + simple_batch(right_schema, &[(4, "rhs-b", 4.0)]), + ], + ) + .await; + assert_eq!( + simple_rows(&result), + vec![ + (1, "lhs-a".to_string(), 1.0), + (2, "lhs-b".to_string(), 2.0), + (3, "rhs-a".to_string(), 3.0), + (4, "rhs-b".to_string(), 4.0), + ] + ); + } + + #[test] + fn fully_filtered_rhs_batch_wakes_before_later_unmatched_batch() { + let schema = simple_schema("stream", false); + let left_polls = Arc::new(AtomicUsize::new(0)); + let right_polls = Arc::new(AtomicUsize::new(0)); + let right_executions = Arc::new(AtomicUsize::new(0)); + let left = test_stream( + schema.clone(), + vec![TestEvent::Batch(simple_batch( + schema.clone(), + &[(1, "match", 1.0)], + ))], + left_polls, + ); + let right = Arc::new(TestExec::new( + schema.clone(), + vec![ + TestEvent::Batch(simple_batch(schema.clone(), &[(1, "match", 2.0)])), + TestEvent::Batch(simple_batch(schema.clone(), &[(2, "keep", 3.0)])), + ], + right_polls.clone(), + right_executions.clone(), + )); + let mut stream = Box::pin(test_union_stream(left, right, schema)); + let wake = Arc::new(CountingWake(AtomicUsize::new(0))); + let waker = Waker::from(wake.clone()); + let mut context = Context::from_waker(&waker); + + assert!(matches!( + stream.as_mut().poll_next(&mut context), + Poll::Ready(Some(Ok(_))) + )); + assert!(matches!( + stream.as_mut().poll_next(&mut context), + Poll::Pending + )); + assert_eq!(right_executions.load(Ordering::SeqCst), 1); + assert_eq!(right_polls.load(Ordering::SeqCst), 1); + assert_eq!(wake.0.load(Ordering::SeqCst), 1); + assert!(matches!( + stream.as_mut().poll_next(&mut context), + Poll::Ready(Some(Ok(batch))) + if simple_rows(std::slice::from_ref(&batch)) + == vec![(2, "keep".to_string(), 3.0)] + )); + } + + #[tokio::test] + async fn rhs_is_not_polled_before_lhs_eof_and_drop_preserves_backpressure() { + let schema = simple_schema("stream", false); + let left_polls = Arc::new(AtomicUsize::new(0)); + let right_polls = Arc::new(AtomicUsize::new(0)); + let left = test_stream( + schema.clone(), + vec![TestEvent::Batch(simple_batch( + schema.clone(), + &[(1, "lhs", 1.0)], + ))], + left_polls.clone(), + ); + let right_executions = Arc::new(AtomicUsize::new(0)); + let right = Arc::new(TestExec::new( + schema.clone(), + vec![TestEvent::Batch(simple_batch( + schema.clone(), + &[(2, "rhs", 2.0)], + ))], + right_polls.clone(), + right_executions.clone(), + )); + let mut stream = Box::pin(test_union_stream(left, right, schema)); + + let first = stream.next().await.unwrap().unwrap(); + assert_eq!(simple_rows(&[first]), vec![(1, "lhs".to_string(), 1.0)]); + assert_eq!(left_polls.load(Ordering::SeqCst), 1); + assert_eq!(right_polls.load(Ordering::SeqCst), 0); + assert_eq!(right_executions.load(Ordering::SeqCst), 0); + drop(stream); + assert_eq!(right_polls.load(Ordering::SeqCst), 0); + assert_eq!(right_executions.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn lhs_and_delayed_rhs_errors_propagate() { + let schema = simple_schema("stream", false); + let polls = Arc::new(AtomicUsize::new(0)); + let mut left_error = Box::pin(test_union_stream( + test_stream( + schema.clone(), + vec![TestEvent::Error("left failed")], + polls.clone(), + ), + Arc::new(TestExec::new( + schema.clone(), + vec![], + polls.clone(), + Arc::new(AtomicUsize::new(0)), + )), + schema.clone(), + )); + assert!(left_error.next().await.unwrap().is_err()); + assert!(left_error.next().await.is_none()); + + let mut right_error = Box::pin(test_union_stream( + test_stream( + schema.clone(), + vec![TestEvent::Batch(simple_batch( + schema.clone(), + &[(1, "lhs", 1.0)], + ))], + polls.clone(), + ), + Arc::new(TestExec::new( + schema.clone(), + vec![TestEvent::Error("right failed")], + polls, + Arc::new(AtomicUsize::new(0)), + )), + schema.clone(), + )); + assert!(right_error.next().await.unwrap().is_ok()); + assert!(right_error.next().await.unwrap().is_err()); + assert!(right_error.next().await.is_none()); + + let right_execute_error = Arc::new( + TestExec::new( + schema.clone(), + vec![], + Arc::new(AtomicUsize::new(0)), + Arc::new(AtomicUsize::new(0)), + ) + .with_execute_error("right execute failed"), + ); + let mut right_execute_error = Box::pin(test_union_stream( + test_stream(schema.clone(), vec![], Arc::new(AtomicUsize::new(0))), + right_execute_error, + schema, + )); + assert!(right_execute_error.next().await.unwrap().is_err()); + assert!(right_execute_error.next().await.is_none()); + } + + #[tokio::test] + async fn rhs_partial_and_full_batches_rebind_to_declared_left_schema() { + let left_schema = schema("left", "label", false); + let right_schema = schema("right", "label", true); + let (left, right) = schemas(left_schema.clone(), right_schema.clone()); + let plan = UnionDistinctOn::try_new(left, right, vec![1], 0).unwrap(); + let declared_schema = plan.output_schema.inner().clone(); + let partial_rhs = RecordBatch::try_new( + right_schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec![Some("match"), None])), + Arc::new(Float64Array::from(vec![Some(10.0), None])), + ], + ) + .unwrap(); + let full_rhs = batch(right_schema.clone(), 3, Some("full"), Some(30.0)); + let result = execute_batches( + &plan, + left_schema.clone(), + vec![batch(left_schema, 1, Some("match"), Some(1.0))], + right_schema, + vec![partial_rhs, full_rhs], + ) + .await; + assert!(result.iter().all(|batch| batch.schema() == declared_schema)); + assert_eq!(result.iter().map(RecordBatch::num_rows).sum::(), 3); + assert!(result[1].column(1).is_null(0)); + } + + #[test] + fn malformed_indices_and_incompatible_inputs_fail_before_execution() { let schema = Arc::new(Schema::new(vec![ Field::new("ts", DataType::Int64, false), Field::new("job", DataType::Utf8, false), - Field::new("value", DataType::Float64, false), ])); - let df_schema = schema.clone().to_dfschema_ref().unwrap(); - let left_plan = LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: false, - schema: df_schema.clone(), - }); - let right_plan = LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: false, - schema: df_schema.clone(), - }); - let plan_node = UnionDistinctOn::new( - left_plan.clone(), - right_plan.clone(), - vec!["job".to_string()], - "ts".to_string(), - df_schema.clone(), - ); - - let bytes = plan_node.serialize(); - - let union_distinct_on = UnionDistinctOn::deserialize(&bytes).unwrap(); - let union_distinct_on = union_distinct_on - .with_exprs_and_inputs(vec![], vec![left_plan, right_plan]) + let invalid = |compare_key_indices, ts_col_idx| { + let decoded = UnionDistinctOn::deserialize( + &pb::UnionDistinctOn { + compare_key_indices, + ts_col_idx, + } + .encode_to_vec(), + ) .unwrap(); - - assert_eq!(union_distinct_on.compare_keys, vec!["job".to_string()]); - assert_eq!(union_distinct_on.ts_col, "ts"); - assert_eq!(union_distinct_on.output_schema, df_schema); + let (left, right) = schemas(schema.clone(), schema.clone()); + decoded + .with_exprs_and_inputs(vec![], vec![left, right]) + .is_err() + }; + assert!(invalid(vec![2], 0)); + assert!(invalid(vec![1], 2)); + let incompatible_schema = Arc::new(Schema::new(vec![ + Field::new("other_ts", DataType::Int64, false), + Field::new("other_job", DataType::Int64, false), + ])); + let (left, right) = schemas(schema, incompatible_schema); + assert!(UnionDistinctOn::try_new(left, right, vec![1], 0).is_err()); } } diff --git a/src/promql/src/lib.rs b/src/promql/src/lib.rs index 2de1c6508e6..78bd012c4d1 100644 --- a/src/promql/src/lib.rs +++ b/src/promql/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(map_try_insert)] - pub mod error; pub mod extension_plan; pub mod functions; diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 61913156cf8..284701f3375 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -50,6 +50,7 @@ use datafusion_expr::utils::conjunction; use datafusion_expr::{ ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit, }; +use datafusion_functions::core::coalesce; use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit}; use datatypes::data_type::ConcreteDataType; use itertools::Itertools; @@ -4423,6 +4424,59 @@ impl PromPlanner { // Take the name of first field column. The length is checked above. let left_field_col = left_context.field_columns.first().unwrap(); let right_field_col = right_context.field_columns.first().unwrap(); + let left_field = left + .schema() + .iter() + .find(|(_, field)| field.name() == left_field_col) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { + col: left_field_col.clone(), + })?; + let right_field = right + .schema() + .iter() + .find(|(_, field)| field.name() == right_field_col) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { + col: right_field_col.clone(), + })?; + let target_field_type = if left_field.1 == right_field.1 { + left_field.1.clone() + } else if matches!( + left_field.1, + ArrowDataType::Int8 + | ArrowDataType::Int16 + | ArrowDataType::Int32 + | ArrowDataType::Int64 + | ArrowDataType::UInt8 + | ArrowDataType::UInt16 + | ArrowDataType::UInt32 + | ArrowDataType::UInt64 + | ArrowDataType::Float32 + | ArrowDataType::Float64 + ) && matches!( + right_field.1, + ArrowDataType::Int8 + | ArrowDataType::Int16 + | ArrowDataType::Int32 + | ArrowDataType::Int64 + | ArrowDataType::UInt8 + | ArrowDataType::UInt16 + | ArrowDataType::UInt32 + | ArrowDataType::UInt64 + | ArrowDataType::Float32 + | ArrowDataType::Float64 + ) { + ArrowDataType::Float64 + } else { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR value fields have incompatible types: {:?} and {:?}", + left_field.1, right_field.1 + ), + } + .fail(); + }; let left_has_tsid = left .schema() .fields() @@ -4459,10 +4513,26 @@ impl PromPlanner { all_columns.sort_unstable(); // use left time index column name as the result time index column name all_columns.insert(0, left_time_index_column.clone()); + let mut occupied_column_names = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); // step 1: align schema using project, fill non-exist columns with null let left_proj_exprs = all_columns.iter().map(|col| { - if tags_not_in_left.contains(col) { + if col == left_field_col && left_field.1 != target_field_type { + DfExpr::Cast(Cast { + expr: Box::new(DfExpr::Column(Column::new( + left_field.0.clone(), + left_field_col, + ))), + data_type: target_field_type.clone(), + }) + .alias(left_field_col.clone()) + } else if tags_not_in_left.contains(col) { DfExpr::Literal(ScalarValue::Utf8(None), None).alias(col.clone()) } else { DfExpr::Column(Column::new(None::, col)) @@ -4475,25 +4545,22 @@ impl PromPlanner { .alias(left_time_index_column.clone()); // The field column in right side may not have qualifier (it may be removed by join operation), // so we need to find it from the schema. - let right_qualifier_for_field = right - .schema() - .iter() - .find(|(_, f)| f.name() == right_field_col) - .map(|(q, _)| q) - .with_context(|| ColumnNotFoundSnafu { - col: right_field_col.clone(), - })? - .cloned(); - // `skip（1)` to skip the time index column let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| { // expr - if col == left_field_col && left_field_col != right_field_col { - // qualify field in right side if necessary to handle different field name - DfExpr::Column(Column::new( - right_qualifier_for_field.clone(), - right_field_col, - )) + if col == left_field_col { + let expr = DfExpr::Column(Column::new(right_field.0.clone(), right_field_col)); + if right_field.1 != target_field_type { + DfExpr::Cast(Cast { + expr: Box::new(expr), + data_type: target_field_type.clone(), + }) + .alias(left_field_col.clone()) + } else if left_field_col != right_field_col { + expr.alias(left_field_col.clone()) + } else { + expr + } } else if tags_not_in_right.contains(col) { DfExpr::Literal(ScalarValue::Utf8(None), None).alias(col.clone()) } else { @@ -4537,24 +4604,168 @@ impl PromPlanner { }; // sort to ensure the generated plan is not volatile match_columns.sort_unstable(); - // step 3: build `UnionDistinctOn` plan - let schema = left_projected.schema().clone(); - let union_distinct_on = UnionDistinctOn::new( - left_projected, - right_projected, - match_columns, - left_time_index_column.clone(), - schema, + match_columns.dedup(); + occupied_column_names.extend( + left_projected + .schema() + .fields() + .iter() + .chain(right_projected.schema().fields().iter()) + .map(|field| field.name().clone()), ); - let result = LogicalPlan::Extension(Extension { + + let visible_schema = left_projected.schema().clone(); + let visible_left_exprs = left_projected + .schema() + .iter() + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + let visible_right_exprs = right_projected + .schema() + .iter() + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + let mut left_match_exprs = Vec::with_capacity(match_columns.len()); + let mut right_match_exprs = Vec::with_capacity(match_columns.len()); + let mut next_internal_column = 0; + + for label in &match_columns { + let left_field = if left_tag_cols_set.contains(label) { + Some( + left_projected + .schema() + .iter() + .find(|(_, field)| field.name() == label) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, + ) + } else { + None + }; + let right_field = if right_tag_cols_set.contains(label) { + Some( + right_projected + .schema() + .iter() + .find(|(_, field)| field.name() == label) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, + ) + } else { + None + }; + let data_type = match (left_field.as_ref(), right_field.as_ref()) { + (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => { + left_type.clone() + } + (Some((_, left_type)), Some((_, right_type))) => { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}" + ), + } + .fail(); + } + (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(), + (None, None) => ArrowDataType::Utf8, + }; + let empty = match data_type { + ArrowDataType::Utf8 => ScalarValue::Utf8(Some(String::new())), + ArrowDataType::LargeUtf8 => ScalarValue::LargeUtf8(Some(String::new())), + _ => { + return UnexpectedPlanExprSnafu { + desc: format!("OR match label {label} must be a string"), + } + .fail(); + } + }; + let internal_name = loop { + let name = format!("__promql_or_match_{next_internal_column}"); + next_internal_column += 1; + if occupied_column_names.insert(name.clone()) { + break name; + } + }; + let normalize = |field: Option<(Option, ArrowDataType)>| { + let expr = if let Some((qualifier, _)) = field { + DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: vec![ + DfExpr::Column(Column::new(qualifier, label.clone())), + DfExpr::Literal(empty.clone(), None), + ], + }) + } else { + DfExpr::Literal(empty.clone(), None) + }; + expr.alias(internal_name.clone()) + }; + left_match_exprs.push(normalize(left_field)); + right_match_exprs.push(normalize(right_field)); + } + + let left_augmented = LogicalPlanBuilder::from(left_projected) + .project(visible_left_exprs.into_iter().chain(left_match_exprs)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + let right_augmented = LogicalPlanBuilder::from(right_projected) + .project(visible_right_exprs.into_iter().chain(right_match_exprs)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + // step 3: build `UnionDistinctOn` with normalized internal match keys. + let visible_field_count = visible_schema.fields().len(); + let compare_key_indices = + (visible_field_count..visible_field_count + match_columns.len()).collect::>(); + let (time_qualifier, _) = visible_schema + .iter() + .find(|(_, field)| field.name() == &left_time_index_column) + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let ts_col_idx = left_augmented + .schema() + .iter() + .position(|(qualifier, field)| { + qualifier == time_qualifier && field.name() == &left_time_index_column + }) + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let union_distinct_on = UnionDistinctOn::try_new( + left_augmented, + right_augmented, + compare_key_indices, + ts_col_idx, + ) + .context(DataFusionPlanningSnafu)?; + let augmented_result = LogicalPlan::Extension(Extension { node: Arc::new(union_distinct_on), }); + let result = LogicalPlanBuilder::from(augmented_result) + .project(visible_schema.iter().map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + })) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; // step 4: update context - self.ctx.time_index_column = Some(left_time_index_column); - self.ctx.tag_columns = all_tags.into_iter().collect(); - self.ctx.field_columns = vec![left_field_col.clone()]; - self.ctx.use_tsid = left_has_tsid && right_has_tsid; + let output_field_col = left_field_col.clone(); + let mut output_context = left_context; + let mut visible_tags = all_tags.into_iter().collect::>(); + visible_tags.sort_unstable(); + output_context.time_index_column = Some(left_time_index_column); + output_context.tag_columns = visible_tags; + output_context.field_columns = vec![output_field_col]; + output_context.use_tsid = left_has_tsid && right_has_tsid; + self.ctx = output_context; Ok(result) } @@ -4761,15 +4972,23 @@ mod test { use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_query::prelude::greptime_timestamp; use common_query::test_util::DummyDecoder; - use datafusion::arrow::datatypes::Schema as ArrowSchema; + use datafusion::arrow::array::{ + Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, + }; + use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; + use datafusion::datasource::{MemTable, provider_as_source}; + use datafusion::execution::context::SessionContext; use datafusion::logical_expr::Extension; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use promql_parser::label::Labels; use promql_parser::parser; use session::context::QueryContext; + use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use table::metadata::{TableInfoBuilder, TableMetaBuilder}; use table::test_util::EmptyTable; @@ -4777,6 +4996,7 @@ mod test { use crate::QueryEngineContext; use crate::options::QueryOptions; use crate::parser::QueryLanguageParser; + use crate::query_engine::DefaultSerializer; fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> { if let LogicalPlan::Extension(Extension { node }) = plan @@ -4872,6 +5092,262 @@ mod test { } } + enum DirectOrValue { + Float64(f64), + Int64(i64), + Utf8(&'static str), + } + + impl DirectOrValue { + fn data_type(&self) -> ArrowDataType { + match self { + Self::Float64(_) => ArrowDataType::Float64, + Self::Int64(_) => ArrowDataType::Int64, + Self::Utf8(_) => ArrowDataType::Utf8, + } + } + fn array(&self) -> Arc { + match self { + Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])), + Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])), + Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])), + } + } + } + + struct DirectOrSource { + name: &'static str, + empty: bool, + timestamp: i64, + tags: Vec<(&'static str, Option<&'static str>)>, + value: DirectOrValue, + } + + fn source( + name: &'static str, + empty: bool, + timestamp: i64, + tags: Vec<(&'static str, Option<&'static str>)>, + value: DirectOrValue, + ) -> DirectOrSource { + DirectOrSource { + name, + empty, + timestamp, + tags, + value, + } + } + + fn tagged_source( + name: &'static str, + empty: bool, + tag: (&'static str, Option<&'static str>), + value: DirectOrValue, + ) -> DirectOrSource { + source(name, empty, 1, vec![("job", Some("job")), tag], value) + } + + fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource { + source(name, true, 1, vec![("job", Some("job"))], value) + } + + fn table(source: &DirectOrSource) -> Arc { + let mut fields = vec![Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )]; + fields.extend( + source + .tags + .iter() + .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)), + ); + fields.push(Field::new("v", source.value.data_type(), true)); + let schema = Arc::new(ArrowSchema::new(fields)); + let partitions = if source.empty { + vec![vec![]] + } else { + let mut columns: Vec> = + vec![Arc::new(TimestampMillisecondArray::from(vec![ + source.timestamp, + ]))]; + columns.extend( + source + .tags + .iter() + .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc), + ); + columns.push(source.value.array()); + vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]] + }; + Arc::new(MemTable::try_new(schema, partitions).unwrap()) + } + + fn scan(source: &DirectOrSource) -> LogicalPlan { + LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None) + .unwrap() + .build() + .unwrap() + } + + fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext { + PromPlannerContext { + table_name: Some(qualifier.to_string()), + time_index_column: Some("ts".to_string()), + field_columns: vec![field.to_string()], + tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(), + ..Default::default() + } + } + + fn or_modifier(expr: &str) -> Option { + let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else { + unreachable!() + }; + expr.modifier + } + + async fn plan_direct_or( + left: LogicalPlan, + right: LogicalPlan, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, + ) -> LogicalPlan { + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + }; + planner + .or_operator( + left, + right, + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + modifier, + ) + .unwrap() + } + + async fn execute( + plan: LogicalPlan, + state: &QueryEngineState, + ) -> (LogicalPlan, Vec) { + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); + let physical = state + .session_state() + .create_physical_plan(&optimized) + .await + .unwrap(); + let batches = + datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) + .await + .unwrap(); + (optimized, batches) + } + + async fn run( + left: &DirectOrSource, + right: &DirectOrSource, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, + ) -> (LogicalPlan, Vec) { + let plan = plan_direct_or( + scan(left), + scan(right), + left_context, + right_context, + modifier, + ) + .await; + execute(plan, &build_query_engine_state()).await + } + + fn assert_no_internal_or_keys(schema: &DFSchema) { + assert!( + schema + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")), + "{schema:?}" + ); + } + + fn values(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(column) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .collect() + } + + fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option)> { + let mut rows = batches + .iter() + .flat_map(|batch| { + let values = batch + .column_by_name("v") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let labels = batch + .column_by_name("k") + .map(|column| column.as_any().downcast_ref::().unwrap()); + (0..batch.num_rows()).map(move |i| { + ( + values.value(i), + labels.and_then(|labels| { + (!labels.is_null(i)).then(|| labels.value(i).to_string()) + }), + ) + }) + }) + .collect::>(); + rows.sort_by(|left, right| left.0.total_cmp(&right.0)); + rows + } + + fn matrix_source( + name: &'static str, + k: Option>, + timestamp: i64, + value: f64, + ) -> DirectOrSource { + let mut tags = vec![("job", Some("job"))]; + if let Some(k) = k { + tags.push(("k", k)); + } + source(name, false, timestamp, tags, DirectOrValue::Float64(value)) + } + + fn matrix_context(name: &str, k: Option>) -> PromPlannerContext { + direct_or_context( + name, + if k.is_some() { &["job", "k"] } else { &["job"] }, + "v", + ) + } + async fn build_test_table_provider( table_name_tuples: &[(String, String)], num_tag: usize, @@ -8191,47 +8667,33 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom #[tokio::test] async fn test_or_not_exists_table_label() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"sum by (job, tag0, tag2) (metric_exists) or sum by (job, tag0, tag2) (metric_not_exists)"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "metric_exists".to_string())], + let state = build_query_engine_state(); + let provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())], &["job"], ) .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = r#"UnionDistinctOn: on col=[["job"]], ts_col=[greptime_timestamp] [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N] - SubqueryAlias: metric_exists [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N] - Projection: metric_exists.greptime_timestamp, metric_exists.job, sum(metric_exists.greptime_value) [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N] - Sort: metric_exists.job ASC NULLS LAST, metric_exists.greptime_timestamp ASC NULLS LAST [job:Utf8, greptime_timestamp:Timestamp(ms), sum(metric_exists.greptime_value):Float64;N] - Aggregate: groupBy=[[metric_exists.job, metric_exists.greptime_timestamp]], aggr=[[sum(metric_exists.greptime_value)]] [job:Utf8, greptime_timestamp:Timestamp(ms), sum(metric_exists.greptime_value):Float64;N] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - PromSeriesDivide: tags=["job"] [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Sort: metric_exists.job ASC NULLS FIRST, metric_exists.greptime_timestamp ASC NULLS FIRST [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Filter: metric_exists.greptime_timestamp >= TimestampMillisecond(-999, None) AND metric_exists.greptime_timestamp <= TimestampMillisecond(100000000, None) [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - TableScan: metric_exists [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - SubqueryAlias: [greptime_timestamp:Timestamp(ms), job:Utf8;N, sum(.value):Float64;N] - Projection: .time AS greptime_timestamp, Utf8(NULL) AS job, sum(.value) [greptime_timestamp:Timestamp(ms), job:Utf8;N, sum(.value):Float64;N] - Sort: .time ASC NULLS LAST [time:Timestamp(ms), sum(.value):Float64;N] - Aggregate: groupBy=[[.time]], aggr=[[sum(.value)]] [time:Timestamp(ms), sum(.value):Float64;N] - EmptyMetric: range=[0..-1], interval=[5000] [time:Timestamp(ms), value:Float64;N] - TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#; - - assert_eq!(plan.display_indent_schema().to_string(), expected); + let raw = PromPlanner::stmt_to_plan( + provider, + &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#), + &state, + ) + .await + .unwrap(); + assert!( + raw.display_indent_schema() + .to_string() + .contains("__promql_or_match_0@") + ); + let (optimized, batches) = execute(raw, &state).await; + assert_no_internal_or_keys(optimized.schema()); + assert!(batches.iter().all(|batch| { + batch + .schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")) + })); } #[tokio::test] @@ -8283,4 +8745,250 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom _ => panic!("Expected EmptyRelation, but got: {:?}", plan), } } + + #[tokio::test] + async fn test_direct_or_normalizes_missing_match_labels() { + type Case<'a> = ( + Option>, + Option>, + i64, + i64, + &'a [(f64, Option<&'a str>)], + ); + + let modifier = or_modifier("lhs or on(k) rhs"); + #[rustfmt::skip] + let cases: &[Case<'_>] = &[ + (None, None, 1, 1, &[(1.0, None)]), + (None, Some(Some("")), 1, 1, &[(1.0, None)]), + (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]), + (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), + (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]), + (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]), + (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), + (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]), + ]; + for &(left, right, left_ts, right_ts, expected) in cases { + let (optimized, batches) = run( + &matrix_source("lhs", left, left_ts, 1.0), + &matrix_source("rhs", right, right_ts, 2.0), + matrix_context("lhs", left), + matrix_context("rhs", right), + &modifier, + ) + .await; + assert_no_internal_or_keys(optimized.schema()); + assert_eq!( + rows(&batches), + expected + .iter() + .map(|(value, label)| (*value, label.map(str::to_string))) + .collect::>() + ); + } + } + + #[tokio::test] + async fn test_direct_or_match_modifiers() { + for (modifier, left, right, expected) in [ + (None, "left", "right", 2), + (or_modifier("lhs or on(k) rhs"), "same", "same", 1), + (or_modifier("lhs or on() rhs"), "left", "right", 1), + (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1), + ] { + let (_, batches) = run( + &matrix_source("lhs", Some(Some(left)), 1, 1.0), + &matrix_source("rhs", Some(Some(right)), 1, 2.0), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &modifier, + ) + .await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected + ); + } + } + + #[tokio::test] + async fn test_direct_or_nested_projection_uses_left_context() { + let left = matrix_source("lhs", Some(Some("k")), 1, 1.0); + let right = matrix_source("rhs", Some(Some("k")), 1, 2.0); + let raw = plan_direct_or( + scan(&left), + scan(&right), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + assert!(raw.schema().iter().any(|(qualifier, field)| { + qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v" + })); + let nested = LogicalPlanBuilder::from(raw) + .project(vec![ + DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(DfExpr::Column(Column::new( + Some(TableReference::bare("lhs")), + "v", + ))), + op: Operator::Plus, + right: Box::new(lit(1.0)), + }) + .alias("v_plus"), + ]) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(nested, &build_query_engine_state()).await; + assert_eq!(values(&batches, "v_plus"), vec![2.0]); + } + + #[tokio::test] + async fn test_direct_or_skips_user_internal_key_name() { + const USER_TAG: &str = "__promql_or_match_0"; + let left = tagged_source( + "lhs", + false, + (USER_TAG, Some("left")), + DirectOrValue::Float64(1.0), + ); + let right = tagged_source( + "rhs", + false, + (USER_TAG, Some("right")), + DirectOrValue::Float64(2.0), + ); + let raw = plan_direct_or( + scan(&left), + scan(&right), + direct_or_context("lhs", &["job", USER_TAG], "v"), + direct_or_context("rhs", &["job", USER_TAG], "v"), + &or_modifier("lhs or on(missing_label) rhs"), + ) + .await; + assert!( + raw.display_indent_schema() + .to_string() + .contains("__promql_or_match_1@") + ); + let (_, batches) = execute(raw, &build_query_engine_state()).await; + assert!( + batches + .iter() + .all(|batch| batch.column_by_name(USER_TAG).is_some()) + ); + } + + #[tokio::test] + async fn test_direct_or_substrait_round_trip_with_normalized_key() { + let state = build_query_engine_state(); + let ctx = SessionContext::new_with_state(state.session_state()); + let catalog = Arc::new(MemoryCatalogProvider::new()); + catalog + .register_schema("public", Arc::new(MemorySchemaProvider::new())) + .unwrap(); + ctx.register_catalog("datafusion", catalog); + let left = matrix_source("lhs", Some(Some("")), 1, 1.0); + let right = matrix_source("rhs", None, 1, 2.0); + ctx.register_table( + TableReference::full("datafusion", "public", "lhs"), + table(&left), + ) + .unwrap(); + ctx.register_table( + TableReference::full("datafusion", "public", "rhs"), + table(&right), + ) + .unwrap(); + let raw = plan_direct_or( + ctx.table("datafusion.public.lhs") + .await + .unwrap() + .into_unoptimized_plan(), + ctx.table("datafusion.public.rhs") + .await + .unwrap() + .into_unoptimized_plan(), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + let decoded = DFLogicalSubstraitConvertor + .decode( + DFLogicalSubstraitConvertor + .encode(&raw, DefaultSerializer) + .unwrap(), + ctx.state(), + ) + .await + .unwrap(); + let (optimized, batches) = execute(decoded, &state).await; + assert_no_internal_or_keys(optimized.schema()); + assert!(batches.iter().all(|batch| { + batch + .schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")) + })); + assert_eq!(values(&batches, "v"), vec![1.0]); + } + + #[tokio::test] + async fn test_direct_or_numeric_value_types() { + let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0)); + let right = tagged_source( + "rhs", + false, + ("k", Some("rhs")), + DirectOrValue::Float64(0.5), + ); + let (optimized, batches) = run( + &left, + &right, + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + assert_eq!( + optimized + .schema() + .field_with_name(None, "v") + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!(values(&batches, "v"), vec![0.5]); + let provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider: provider, + ctx: PromPlannerContext::default(), + }; + let left_context = direct_or_context("lhs", &["job"], "v"); + let right_context = direct_or_context("rhs", &["job"], "v"); + let error = planner + .or_operator( + scan(&job_source("lhs", DirectOrValue::Utf8("x"))), + scan(&job_source("rhs", DirectOrValue::Float64(1.0))), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on() rhs"), + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("OR value fields have incompatible types") + ); + } } diff --git a/tests/cases/standalone/common/promql/set_operation.result b/tests/cases/standalone/common/promql/set_operation.result index 83af6fc8755..50a0b6cf3b6 100644 --- a/tests/cases/standalone/common/promql/set_operation.result +++ b/tests/cases/standalone/common/promql/set_operation.result @@ -466,12 +466,12 @@ tql eval (0, 2000, '400') sum(t1{job="a"}); -- SQLNESS SORT_RESULT 3 1 tql eval (0, 2000, '400') sum(t1{job="a"}) - sum(t1{job="e"} or vector(1)); -+---------------------+-----------------------------------------------------+ -| ts | t1.sum(t1.greptime_value) - .sum(t1.greptime_value) | -+---------------------+-----------------------------------------------------+ -| 1970-01-01T00:00:00 | 0.0 | -| 1970-01-01T00:20:00 | 2.0 | -+---------------------+-----------------------------------------------------+ ++---------------------+---------------------------------------------------------+ +| ts | lhs.sum(t1.greptime_value) - rhs.sum(t1.greptime_value) | ++---------------------+---------------------------------------------------------+ +| 1970-01-01T00:00:00 | 0.0 | +| 1970-01-01T00:20:00 | 2.0 | ++---------------------+---------------------------------------------------------+ drop table t1; @@ -545,12 +545,12 @@ tql eval (0, 2000, '400') count(max by (namespace) (stats_used_bytes{namespace=~ +---------------------+---------------------------------------------+ | ts | count(max(stats_used_bytes.greptime_value)) | +---------------------+---------------------------------------------+ -| 1970-01-01T00:00:00 | 0 | -| 1970-01-01T00:06:40 | 0 | -| 1970-01-01T00:13:20 | 0 | -| 1970-01-01T00:20:00 | 2 | -| 1970-01-01T00:26:40 | 0 | -| 1970-01-01T00:33:20 | 0 | +| 1970-01-01T00:00:00 | 0.0 | +| 1970-01-01T00:06:40 | 0.0 | +| 1970-01-01T00:13:20 | 0.0 | +| 1970-01-01T00:20:00 | 2.0 | +| 1970-01-01T00:26:40 | 0.0 | +| 1970-01-01T00:33:20 | 0.0 | +---------------------+---------------------------------------------+ -- SQLNESS SORT_RESULT 3 1 diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result index d5313ba41c6..8c90d97a325 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result @@ -512,10 +512,10 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_ProjectionExec: expr=[host@2 as host, job@3 as job, ts@5 as ts, __tsid@4 as __tsid, greptime_value@0 / greptime_value@1 as tsid_binary_join_right.greptime_value / tsid_binary_join_left.greptime_value] REDACTED +| 0_| 0_|_ProjectionExec: expr=[host@2 as host, job@3 as job, ts@5 as ts, __tsid@4 as __tsid, greptime_value@0 / greptime_value@1 as lhs.greptime_value / rhs.greptime_value] REDACTED |_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@0, ts@4)], projection=[greptime_value@2, greptime_value@3, host@4, job@5, __tsid@6, ts@7], NullsEqual: true REDACTED |_|_|_ProjectionExec: expr=[ts@0 as ts, __tsid@1 as __tsid, greptime_value@2 as greptime_value] REDACTED -|_|_|_UnionDistinctOnExec: on col=[["host", "job"]], ts_col=[ts] REDACTED +|_|_|_UnionDistinctOnExec: on col=[5, 6], ts_col=0 REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_CoalescePartitionsExec REDACTED @@ -523,14 +523,14 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| -| 1_| 0_|_ProjectionExec: expr=[ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job] REDACTED +| 1_| 0_|_ProjectionExec: expr=[ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job, CASE WHEN host@1 IS NOT NULL THEN host@1 ELSE_END as __promql_or_match_0, CASE WHEN job@2 IS NOT NULL THEN job@2 ELSE_END as __promql_or_match_1] REDACTED |_|_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED |_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED |_|_|_| -| 1_| 0_|_ProjectionExec: expr=[ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job] REDACTED +| 1_| 0_|_ProjectionExec: expr=[ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job, CASE WHEN host@1 IS NOT NULL THEN host@1 ELSE_END as __promql_or_match_0, CASE WHEN job@2 IS NOT NULL THEN job@2 ELSE_END as __promql_or_match_1] REDACTED |_|_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED @@ -672,14 +672,14 @@ TQL EVAL (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) + t -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsid_binary_join_left; -+-------+------+---------------------+------------------------------------------------------------------------------+ -| host | job | ts | tsid_binary_join_right.greptime_value / tsid_binary_join_left.greptime_value | -+-------+------+---------------------+------------------------------------------------------------------------------+ -| host1 | job1 | 1970-01-01T00:00:00 | 1.0 | -| host1 | job1 | 1970-01-01T00:00:05 | 1.0 | -| host2 | job2 | 1970-01-01T00:00:00 | 1.0 | -| host2 | job2 | 1970-01-01T00:00:05 | 1.0 | -+-------+------+---------------------+------------------------------------------------------------------------------+ ++-------+------+---------------------+-----------------------------------------+ +| host | job | ts | lhs.greptime_value / rhs.greptime_value | ++-------+------+---------------------+-----------------------------------------+ +| host1 | job1 | 1970-01-01T00:00:00 | 1.0 | +| host1 | job1 | 1970-01-01T00:00:05 | 1.0 | +| host2 | job2 | 1970-01-01T00:00:00 | 1.0 | +| host2 | job2 | 1970-01-01T00:00:05 | 1.0 | ++-------+------+---------------------+-----------------------------------------+ -- Range functions are outside the first island version; the range selector subtree must -- remain a barrier. From a3816c889e0d6fd29628280083cc7293e54c85a8 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 16 Jul 2026 11:31:24 +0800 Subject: [PATCH 43/48] fix(ci): harden query regression runner (#8534) Signed-off-by: discord9 --- .../query-regression/Dockerfile | 21 ++++++- .../query-regression/README.md | 34 +++++++---- .../query-regression/values-8-cores.yaml | 26 ++++++++- .github/workflows/query-regression.yml | 57 +++++++++++++++++-- 4 files changed, 117 insertions(+), 21 deletions(-) diff --git a/.github/runner-scale-sets/query-regression/Dockerfile b/.github/runner-scale-sets/query-regression/Dockerfile index 95161f6093b..61392919587 100644 --- a/.github/runner-scale-sets/query-regression/Dockerfile +++ b/.github/runner-scale-sets/query-regression/Dockerfile @@ -26,6 +26,7 @@ RUN apt-get update \ git \ gzip \ jq \ + libprotobuf-dev \ libssl-dev \ mold \ openssh-client \ @@ -71,7 +72,9 @@ USER runner RUN set -eu \ && test "$(id -u)" = "1001" \ && temporary_cargo_home="$(mktemp --directory)" \ - && trap 'rm -rf "${temporary_cargo_home}"' 0 \ + && temporary_proto_dir="" \ + && cleanup() { rm -rf "${temporary_cargo_home}" "${temporary_proto_dir}"; } \ + && trap cleanup 0 \ && export CARGO_HOME="${temporary_cargo_home}" \ && test "${RUSTUP_AUTO_INSTALL}" = "0" \ && rustup --version \ @@ -81,4 +84,18 @@ RUN set -eu \ && printf 'Active toolchain: %s\n' "${active_toolchain}" \ && case "${active_toolchain}" in "${RUST_TOOLCHAIN}-x86_64-unknown-linux-gnu"|"${RUST_TOOLCHAIN}-x86_64-unknown-linux-gnu "*) ;; *) exit 1;; esac \ && test ! -w /opt/rustup \ - && test ! -w /opt/cargo/bin + && test ! -w /opt/cargo/bin \ + && test -r /usr/include/google/protobuf/any.proto \ + && test -r /usr/include/google/protobuf/empty.proto \ + && temporary_proto_dir="$(mktemp --directory)" \ + && printf '%s\n' \ + 'syntax = "proto3";' \ + 'package smoke;' \ + 'import "google/protobuf/any.proto";' \ + 'import "google/protobuf/empty.proto";' \ + 'message Smoke { google.protobuf.Any any = 1; google.protobuf.Empty empty = 2; }' \ + > "${temporary_proto_dir}/smoke.proto" \ + && protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ + --descriptor_set_out="${temporary_proto_dir}/smoke.pb" \ + "${temporary_proto_dir}/smoke.proto" \ + && test -s "${temporary_proto_dir}/smoke.pb" diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index fd68b4c8321..4fe8c0f3afb 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -56,13 +56,14 @@ egress to required GitHub Actions, artifact/cache, Rust/crate/toolchain, DNS, and image-registry endpoints; block unrelated cluster services, private ranges, and metadata endpoints unless a case requires them. -### Office routing prerequisite +### Network routing prerequisite -Direct split routing through `.2` is an **external office-network prerequisite**. -The gateway at `192.168.50.2` must route GitHub Actions, GitHub content, -artifact/cache, crates.io, Rust toolchain, and image-registry traffic directly -rather than through the VPN. Neither this repository nor Kubernetes configures -that route. Verify it with the responsible network operator before any canary. +Required split routing is an **external environment-specific prerequisite**. The +responsible network operator must route GitHub Actions, GitHub content, +artifact/cache, crates.io, Rust toolchain, and image-registry traffic through +the approved path rather than the VPN where applicable. Neither this repository +nor Kubernetes configures that route. Verify it with the responsible network +operator before any canary. ## Runner image and workflow tools @@ -85,11 +86,20 @@ Both digest-pinned init and runner containers use `IfNotPresent`: the immutable digest makes a cached image safe and avoids adding a registry dependency to every runner startup. +The runner optionally imports only the non-sensitive `HTTP_PROXY`, +`HTTPS_PROXY`, and `NO_PROXY` variables from the +`query-regression-runner-local-env` ConfigMap. Manage that ConfigMap locally in +the target namespace; private endpoint configuration must not be committed, and +credentials or secrets must never be placed in a ConfigMap. + Before builds, the workflow asserts UID/GID 1001 and exact image tool versions: `libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.30.0`, `Python 3.12.3`, `sccache 0.16.0`, root-owned `rustup 1.29.0`, and the image-baked `nightly-2026-03-21` Rust toolchain. Rustup, Cargo, and Rustc must resolve from `/opt/cargo/bin`; the runner cannot write `/opt/rustup` or `/opt/cargo/bin`. +Protobuf well-known includes, including `google/protobuf/any.proto` and +`google/protobuf/empty.proto`, are an image contract and must compile with +`protoc`. `actions-rust-lang/setup-rust-toolchain@v1` is intentionally removed. The workflow sets its warning-denying mold `RUSTFLAGS` directly, disables automatic Rustup installation, and performs no runtime toolchain downloads. @@ -107,6 +117,9 @@ corresponding cache or image contract changes. `queue: max`, and `cancel-in-progress: false`; admitted jobs queue rather than replacing older pending jobs. During maintenance, cancel admitted queued runs as well as pausing ARC. Runner Pods have `activeDeadlineSeconds=12600`. +The runner requests 6 CPU and limits at 8 CPU to preserve `minipc-3` +allocatable-capacity scheduling headroom; do not reset the request to 8 CPU +without revalidating scheduling capacity. The cache claim `query-regression-build-cache` is a nominal 600Gi `local-path` PVC in `arc-runners`. It is `ReadWriteOnce`; `local-path` uses @@ -184,7 +197,7 @@ exists; they are not mounted by the current configuration. ## Deploy and pause safely -First verify external `.2` routing and the disk preflight. Apply the PVC; while +First verify the configured external network route and the disk preflight. Apply the PVC; while the scale set is paused, it is expected to remain `Pending` because WaitForFirstConsumer has no scheduled runner: @@ -263,8 +276,8 @@ ABI-marker output, initial/base/candidate sccache statistics, and cache report. Confirm the image tool contract (root-owned Rustup/Cargo paths, baked nightly, and non-writable `/opt` roots) and that the ephemeral Cargo home contains only the mounted registry/Git data before Cargo creates per-Pod state. -Obtain dependency and tool network byte counters from the -identified `.2` counter source, filtered to `minipc-3` and the dependency/tool +Obtain dependency and tool network byte counters from the configured +environment counter source, filtered to `minipc-3` and the dependency/tool destinations. Accept the canary only when all of the following hold: @@ -278,7 +291,8 @@ Accept the canary only when all of the following hold: - warm dependency/tool network bytes are at most 10% of cold fill bytes; - cache sizes remain below soft watermarks, node free space remains at least 300GiB, and the benchmark is correct without TLS EOFs or timeouts; -- `.2` confirms this traffic is outside VPN accounting. +- the configured environment counter source confirms this traffic is outside VPN + accounting. Immediately return to 0/0 after either canary unless ongoing normal operation has been explicitly approved; return immediately on any traffic, cache, disk, diff --git a/.github/runner-scale-sets/query-regression/values-8-cores.yaml b/.github/runner-scale-sets/query-regression/values-8-cores.yaml index 77ec4055e5b..057ff55b507 100644 --- a/.github/runner-scale-sets/query-regression/values-8-cores.yaml +++ b/.github/runner-scale-sets/query-regression/values-8-cores.yaml @@ -27,7 +27,7 @@ template: claimName: query-regression-build-cache initContainers: - name: initialize-build-cache - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:7022b551771dc82ec4c493f0c3db6ee3050ececd336598690a2165240ebe953e + image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:2436278bf0756700c718927907cdbedac64d55c593f80b515acb44dec8b02cc1 imagePullPolicy: IfNotPresent command: - /bin/sh @@ -65,12 +65,32 @@ template: mountPath: /cache containers: - name: runner - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:7022b551771dc82ec4c493f0c3db6ee3050ececd336598690a2165240ebe953e + image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:2436278bf0756700c718927907cdbedac64d55c593f80b515acb44dec8b02cc1 imagePullPolicy: IfNotPresent command: ["/home/runner/run.sh"] env: - name: CARGO_HOME value: /home/runner/.cargo + - name: HTTP_PROXY + valueFrom: + configMapKeyRef: + name: query-regression-runner-local-env + key: HTTP_PROXY + optional: true + - name: HTTPS_PROXY + valueFrom: + configMapKeyRef: + name: query-regression-runner-local-env + key: HTTPS_PROXY + optional: true + - name: NO_PROXY + valueFrom: + configMapKeyRef: + name: query-regression-runner-local-env + key: NO_PROXY + optional: true + - name: UV_CACHE_DIR + value: /home/runner/.cargo/uv-cache - name: RUSTUP_HOME value: /opt/rustup - name: RUSTUP_TOOLCHAIN @@ -99,7 +119,7 @@ template: - ALL resources: requests: - cpu: "8" + cpu: "6" memory: 24Gi ephemeral-storage: 80Gi limits: diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index aa6a7724b24..3485bb32434 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -62,6 +62,7 @@ jobs: env: CARGO_PROFILE: ${{ github.event_name == 'pull_request' && 'nightly' || inputs.cargo_profile }} CARGO_HOME: /home/runner/.cargo + UV_CACHE_DIR: /home/runner/.cargo/uv-cache RUSTUP_HOME: /opt/rustup RUSTUP_TOOLCHAIN: nightly-2026-03-21 RUSTUP_AUTO_INSTALL: "0" @@ -111,6 +112,29 @@ jobs: [[ "$1" =~ ^[[:xdigit:]]{40}$ ]] } + fetch_with_retry() { + local ref="$1" + local depth="$2" + local attempt=1 + local max_attempts=5 + local delay + + while (( attempt <= max_attempts )); do + if git fetch --no-tags --prune --depth="${depth}" origin "${ref}"; then + return 0 + fi + printf 'Git fetch failed for %s (attempt %d/%d).\n' \ + "${ref}" "${attempt}" "${max_attempts}" >&2 + if (( attempt < max_attempts )); then + delay=$(( 2 ** (attempt - 1) )) + printf 'Retrying git fetch for %s in %d seconds.\n' "${ref}" "${delay}" >&2 + sleep "${delay}" + fi + ((attempt += 1)) + done + return 1 + } + initially_checked_out_base_sha="$(git rev-parse --verify HEAD)" is_full_sha "${initially_checked_out_base_sha}" || { printf 'Could not resolve initially checked-out BASE_REF commit.\n' >&2 @@ -123,7 +147,7 @@ jobs: is_full_sha "${value}" || fail_closed "Missing or invalid ${identity}." done - if ! git fetch --no-tags --prune --depth=2 origin "${EVENT_MERGE_SHA}"; then + if ! fetch_with_retry "${EVENT_MERGE_SHA}" 2; then fail_closed "Could not fetch event merge SHA ${EVENT_MERGE_SHA}." fi fetched_sha="$(git rev-parse --verify FETCH_HEAD 2>/dev/null)" \ @@ -148,7 +172,10 @@ jobs: || fail_closed "Event merge SHA has no non-head parent for the base build." VERIFIED_CANDIDATE_SHA="${fetched_sha,,}" else - git fetch --no-tags --prune --depth=1 origin "${CANDIDATE_REF}" + if ! fetch_with_retry "${CANDIDATE_REF}" 1; then + printf 'Could not fetch candidate ref %s.\n' "${CANDIDATE_REF}" >&2 + exit 1 + fi VERIFIED_CANDIDATE_SHA="$(git rev-parse --verify FETCH_HEAD)" VERIFIED_BASE_SHA="${initially_checked_out_base_sha}" fi @@ -176,7 +203,23 @@ jobs: set -euo pipefail [[ "$(id -u)" == "1001" ]] [[ "$(id -g)" == "1001" ]] + [[ "${UV_CACHE_DIR}" == "/home/runner/.cargo/uv-cache" ]] [[ "$(protoc --version)" == "libprotoc 3.21.12" ]] + temporary_proto_dir="$(mktemp --directory)" + trap 'rm -rf "${temporary_proto_dir}"' EXIT + test -r /usr/include/google/protobuf/any.proto + test -r /usr/include/google/protobuf/empty.proto + printf '%s\n' \ + 'syntax = "proto3";' \ + 'package smoke;' \ + 'import "google/protobuf/any.proto";' \ + 'import "google/protobuf/empty.proto";' \ + 'message Smoke { google.protobuf.Any any = 1; google.protobuf.Empty empty = 2; }' \ + > "${temporary_proto_dir}/smoke.proto" + protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ + --descriptor_set_out="${temporary_proto_dir}/smoke.pb" \ + "${temporary_proto_dir}/smoke.proto" + test -s "${temporary_proto_dir}/smoke.pb" [[ "$(uv --version)" =~ ^uv[[:space:]]0\.11\.26([[:space:]]|$) ]] mold_version="$(mold --version)" [[ "${mold_version}" =~ ^mold[[:space:]]2\.30\.0([[:space:]]|$) ]] @@ -213,13 +256,14 @@ jobs: readonly EXPECTED_CARGO_HOME="/home/runner/.cargo" readonly EXPECTED_CARGO_REGISTRY="/home/runner/.cargo/registry" readonly EXPECTED_CARGO_GIT="/home/runner/.cargo/git" + readonly EXPECTED_UV_CACHE_DIR="/home/runner/.cargo/uv-cache" readonly EXPECTED_RUSTUP_HOME="/opt/rustup" readonly EXPECTED_TARGET_DIR="/home/runner/query-regression-target" readonly EXPECTED_CACHE_META="/home/runner/query-regression-cache-meta" readonly EXPECTED_SCCACHE_DIR="/home/runner/.cache/sccache" readonly EXPECTED_RUSTC_WRAPPER="/usr/local/bin/sccache" - readonly RUNNER_IMAGE_DIGEST="sha256:7022b551771dc82ec4c493f0c3db6ee3050ececd336598690a2165240ebe953e" - readonly RUNNER_IMAGE_EPOCH="3" + readonly RUNNER_IMAGE_DIGEST="sha256:2436278bf0756700c718927907cdbedac64d55c593f80b515acb44dec8b02cc1" + readonly RUNNER_IMAGE_EPOCH="4" require_expected_root() { local name="$1" @@ -276,6 +320,7 @@ jobs: require_expected_root CARGO_HOME "${CARGO_HOME}" "${EXPECTED_CARGO_HOME}" require_expected_root CARGO_REGISTRY "${CARGO_HOME}/registry" "${EXPECTED_CARGO_REGISTRY}" require_expected_root CARGO_GIT "${CARGO_HOME}/git" "${EXPECTED_CARGO_GIT}" + require_expected_root UV_CACHE_DIR "${UV_CACHE_DIR}" "${EXPECTED_UV_CACHE_DIR}" require_expected_root RUSTUP_HOME "${RUSTUP_HOME}" "${EXPECTED_RUSTUP_HOME}" require_expected_root CARGO_TARGET_DIR "${CARGO_TARGET_DIR}" "${EXPECTED_TARGET_DIR}" require_expected_root QUERY_REGRESSION_CACHE_META "${QUERY_REGRESSION_CACHE_META}" "${EXPECTED_CACHE_META}" @@ -295,10 +340,10 @@ jobs: test -r "${RUSTUP_HOME}" test -x "${RUSTUP_HOME}" test ! -w "${RUSTUP_HOME}" - mkdir -p "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ + mkdir -p "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" "${UV_CACHE_DIR}" \ "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}" for root in "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ - "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}"; do + "${UV_CACHE_DIR}" "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}"; do test -w "${root}" touch "${root}/.query-regression-write-test" rm -f "${root}/.query-regression-write-test" From c7bcddacc5a1576875843db1b411e73acb1a622c Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 16 Jul 2026 14:40:42 +0800 Subject: [PATCH 44/48] fix: preserve nulls in timestamp arrays (#8508) Signed-off-by: discord9 --- src/servers/src/postgres/types.rs | 62 ++++++++++++++++++++++++++++--- 1 file changed, 57 insertions(+), 5 deletions(-) diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index b282bdea68a..465f468f1d5 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -1055,17 +1055,22 @@ pub(super) fn parameters_to_scalar_values( TimestampType::Nanosecond(_) => { let values = data .into_iter() - .filter_map(|ts| { - ts.and_then(|ts| { - ts.and_utc().timestamp_nanos_opt().map(|nanos| { + .map(|ts| match ts { + None => { + Ok(ScalarValue::TimestampNanosecond(None, None)) + } + Some(ts) => ts + .and_utc() + .timestamp_nanos_opt() + .map(|nanos| { ScalarValue::TimestampNanosecond( Some(nanos), None, ) }) - }) + .ok_or_else(|| numeric_out_of_range_error(ts)), }) - .collect::>(); + .collect::>>()?; ScalarValue::List(ScalarValue::new_list( &values, &ArrowDataType::Timestamp(TimeUnit::Nanosecond, None), @@ -2051,6 +2056,53 @@ mod test { } } + fn qbs_timestamp_nanosecond_array_plan() -> LogicalPlan { + build_plan_with_params(vec![( + "$1", + DataType::List(Arc::new(Field::new( + "item", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + ))), + )]) + } + + #[test] + fn test_qbs_pg_timestamp_nanosecond_array_preserves_null_slot() { + let portal = make_portal( + vec![Some(Type::TIMESTAMP_ARRAY)], + vec![s( + r#"{"2024-01-01 00:00:00.000001",NULL,"2024-01-01 00:00:00.000003"}"#, + )], + ); + + let values = + parameters_to_scalar_values(&qbs_timestamp_nanosecond_array_plan(), &portal).unwrap(); + let expected = ScalarValue::List(ScalarValue::new_list( + &[ + ScalarValue::TimestampNanosecond(Some(1_704_067_200_000_001_000), None), + ScalarValue::TimestampNanosecond(None, None), + ScalarValue::TimestampNanosecond(Some(1_704_067_200_000_003_000), None), + ], + &ArrowDataType::Timestamp(TimeUnit::Nanosecond, None), + true, + )); + assert_eq!(expected, values[0]); + } + + #[test] + fn test_qbs_pg_timestamp_nanosecond_array_out_of_range_rejected() { + let portal = make_portal( + vec![Some(Type::TIMESTAMP_ARRAY)], + vec![s(r#"{"3000-01-01 00:00:00"}"#)], + ); + + assert_numeric_out_of_range(parameters_to_scalar_values( + &qbs_timestamp_nanosecond_array_plan(), + &portal, + )); + } + #[test] fn test_numeric_uint64_scalar_negative_rejected() { let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]); From a585b83bfc03534ace08ed2ceb79931c23b83c72 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 16 Jul 2026 15:59:35 +0800 Subject: [PATCH 45/48] fix: count Postgres SCRAM auth failures and correct auth config docs (#8538) * fix: count Postgres SCRAM auth failures in the auth failure metric The SCRAM SASL paths returned `Failed` without touching `METRIC_AUTH_FAILURE`, so once SCRAM is enabled wrong-password and unknown-user attempts disappeared from `greptime_servers_auth_failure_count`. Funnel every SCRAM rejection through `record_scram_failure`, which records one failure with a uniform `UserPasswordMismatch` label so the counter stays useful without revealing whether the user exists. Signed-off-by: Dennis Zhuang * docs: correct per-protocol auth notes in config examples The user_provider note no longer matched the implementation: pbkdf2_sha256 is excluded from Postgres SCRAM (so its iteration count is never exposed in a SCRAM handshake), and the warning that hash-only verifiers cannot use MySQL's native password handshake had been dropped. State the actual per-protocol fallbacks and incompatibilities, scope the iteration/salt enumeration caveat to pg_scram_sha256, and regenerate config.md. Signed-off-by: Dennis Zhuang * fix: avoid double-counting Postgres SCRAM authorization failures authorize() already increments METRIC_AUTH_FAILURE with its own status code, so routing the authorization-rejection path through record_scram_failure counted it twice, mislabeling the second increment as UserPasswordMismatch. Return early to bypass the recorder. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- config/config.md | 4 +- config/frontend.example.toml | 13 +- config/standalone.example.toml | 13 +- src/servers/src/postgres/auth_handler.rs | 185 +++++++++++++---------- 4 files changed, 123 insertions(+), 92 deletions(-) diff --git a/config/config.md b/config/config.md index 3d406494cfe..a04178bdb93 100644 --- a/config/config.md +++ b/config/config.md @@ -15,7 +15,7 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index and value columns. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | -| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest.
"mysql_native_password" is MySQL-specific and falls back to cleartext password auth for PostgreSQL.
For PostgreSQL SCRAM, keep all users on the default iteration count (4096) and salt length (16);
a non-default "pbkdf2_sha256"/"pg_scram_sha256" iteration count or salt length is observable in the
SCRAM handshake and weakens resistance to username enumeration. | +| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | | `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | | `init_regions_in_background` | Bool | `false` | Initialize all regions in the background during the startup.
By default, it provides services after all regions have been initialized. | @@ -239,7 +239,7 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index and value columns. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | -| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest.
"mysql_native_password" is MySQL-specific and falls back to cleartext password auth for PostgreSQL.
For PostgreSQL SCRAM, keep all users on the default iteration count (4096) and salt length (16);
a non-default "pbkdf2_sha256"/"pg_scram_sha256" iteration count or salt length is observable in the
SCRAM handshake and weakens resistance to username enumeration. | +| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | | `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | | `runtime` | -- | -- | The runtime options. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index c856c24b0f2..73efe8c1376 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -15,11 +15,14 @@ default_column_prefix = "greptime" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", ## "mysql_native_password:", ## "pg_scram_sha256::::" -## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest. -## "mysql_native_password" is MySQL-specific and falls back to cleartext password auth for PostgreSQL. -## For PostgreSQL SCRAM, keep all users on the default iteration count (4096) and salt length (16); -## a non-default "pbkdf2_sha256"/"pg_scram_sha256" iteration count or salt length is observable in the -## SCRAM handshake and weakens resistance to username enumeration. +## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's +## native password handshake; a MySQL client must send the password in cleartext for such users. +## "mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all. +## PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or +## "mysql_native_password", PostgreSQL falls back to cleartext password auth for every user. +## For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are +## observable in the SCRAM server-first message, and non-default values weaken resistance to username +## enumeration. ## @toml2docs:none-default #+ user_provider = "static_user_provider:file:/path/to/users" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 17e830a2c1d..571353ec20b 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -15,11 +15,14 @@ default_column_prefix = "greptime" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", ## "mysql_native_password:", ## "pg_scram_sha256::::" -## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest. -## "mysql_native_password" is MySQL-specific and falls back to cleartext password auth for PostgreSQL. -## For PostgreSQL SCRAM, keep all users on the default iteration count (4096) and salt length (16); -## a non-default "pbkdf2_sha256"/"pg_scram_sha256" iteration count or salt length is observable in the -## SCRAM handshake and weakens resistance to username enumeration. +## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's +## native password handshake; a MySQL client must send the password in cleartext for such users. +## "mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all. +## PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or +## "mysql_native_password", PostgreSQL falls back to cleartext password auth for every user. +## For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are +## observable in the SCRAM server-first message, and non-default values weaken resistance to username +## enumeration. ## @toml2docs:none-default #+ user_provider = "static_user_provider:file:/path/to/users" diff --git a/src/servers/src/postgres/auth_handler.rs b/src/servers/src/postgres/auth_handler.rs index e221f451732..1788da89364 100644 --- a/src/servers/src/postgres/auth_handler.rs +++ b/src/servers/src/postgres/auth_handler.rs @@ -24,6 +24,7 @@ use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64; use common_catalog::parse_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; use common_time::Timezone; use futures::{Sink, SinkExt}; use pgwire::api::auth::StartupHandler; @@ -336,6 +337,18 @@ enum PgAuthenticationResult { } impl PostgresServerHandlerInner { + /// Records a rejected SCRAM attempt in [`METRIC_AUTH_FAILURE`]. The label is + /// intentionally uniform (never `UserNotFound`), so the counter cannot be + /// used to distinguish a wrong password from an unknown user. + fn record_scram_failure(result: PgAuthenticationResult) -> PgAuthenticationResult { + if matches!(result, PgAuthenticationResult::Failed) { + METRIC_AUTH_FAILURE + .with_label_values(&[StatusCode::UserPasswordMismatch.as_ref()]) + .inc(); + } + result + } + async fn authenticate_password_message( &self, client: &mut C, @@ -360,46 +373,53 @@ impl PostgresServerHandlerInner { } } PgAuthenticationState::SaslInitial { auth_info } => { - let sasl_initial = pwd.into_sasl_initial_response()?; - if sasl_initial.auth_method != SCRAM_SHA_256_METHOD { - return Ok(PgAuthenticationResult::Failed); - } - let Some(data) = sasl_initial.data else { - return Ok(PgAuthenticationResult::Failed); - }; - let Some(client_first) = ScramClientFirst::parse(&data) else { - return Ok(PgAuthenticationResult::Failed); - }; - let PgAuthInfo::ScramSha256 { - verifier, - user_info, - } = auth_info - else { - return Ok(PgAuthenticationResult::Failed); - }; + // A labeled block funnels every rejection through a single + // `record_scram_failure`, so failed SCRAM attempts still show up + // in `METRIC_AUTH_FAILURE` without sprinkling the metric call + // across each early return. + let result = 'sasl: { + let sasl_initial = pwd.into_sasl_initial_response()?; + if sasl_initial.auth_method != SCRAM_SHA_256_METHOD { + break 'sasl PgAuthenticationResult::Failed; + } + let Some(data) = sasl_initial.data else { + break 'sasl PgAuthenticationResult::Failed; + }; + let Some(client_first) = ScramClientFirst::parse(&data) else { + break 'sasl PgAuthenticationResult::Failed; + }; + let PgAuthInfo::ScramSha256 { + verifier, + user_info, + } = auth_info + else { + break 'sasl PgAuthenticationResult::Failed; + }; - let server_nonce = BASE64.encode(rand::random::<[u8; 18]>()); - let nonce = format!("{}{}", client_first.nonce, server_nonce); - let server_first = format!( - "r={},s={},i={}", - nonce, - BASE64.encode(verifier.salt()), - verifier.iterations() - ); - client - .send(PgWireBackendMessage::Authentication( - Authentication::SASLContinue(server_first.clone().into()), - )) - .await?; - *state = PgAuthenticationState::SaslFinal { - verifier, - user_info, - channel_binding: client_first.channel_binding, - nonce, - client_first_bare: client_first.bare, - server_first, + let server_nonce = BASE64.encode(rand::random::<[u8; 18]>()); + let nonce = format!("{}{}", client_first.nonce, server_nonce); + let server_first = format!( + "r={},s={},i={}", + nonce, + BASE64.encode(verifier.salt()), + verifier.iterations() + ); + client + .send(PgWireBackendMessage::Authentication( + Authentication::SASLContinue(server_first.clone().into()), + )) + .await?; + *state = PgAuthenticationState::SaslFinal { + verifier, + user_info, + channel_binding: client_first.channel_binding, + nonce, + client_first_bare: client_first.bare, + server_first, + }; + PgAuthenticationResult::Continue }; - Ok(PgAuthenticationResult::Continue) + Ok(Self::record_scram_failure(result)) } PgAuthenticationState::SaslFinal { verifier, @@ -409,51 +429,56 @@ impl PostgresServerHandlerInner { client_first_bare, server_first, } => { - let sasl_response = pwd.into_sasl_response()?; - let Some(client_final) = ScramClientFinal::parse(&sasl_response.data) else { - return Ok(PgAuthenticationResult::Failed); - }; - if client_final.channel_binding != channel_binding { - return Ok(PgAuthenticationResult::Failed); - } - if client_final.nonce != nonce { - return Ok(PgAuthenticationResult::Failed); - } + let result = 'sasl: { + let sasl_response = pwd.into_sasl_response()?; + let Some(client_final) = ScramClientFinal::parse(&sasl_response.data) else { + break 'sasl PgAuthenticationResult::Failed; + }; + if client_final.channel_binding != channel_binding { + break 'sasl PgAuthenticationResult::Failed; + } + if client_final.nonce != nonce { + break 'sasl PgAuthenticationResult::Failed; + } - let auth_message = format!( - "{},{},{}", - client_first_bare, server_first, client_final.without_proof - ); - let Ok(client_proof) = BASE64.decode(client_final.proof.as_bytes()) else { - return Ok(PgAuthenticationResult::Failed); - }; - let Ok(Some(server_signature)) = - verifier.verify_client_proof(auth_message.as_bytes(), &client_proof) - else { - return Ok(PgAuthenticationResult::Failed); - }; - let Some(user_info) = user_info else { - return Ok(PgAuthenticationResult::Failed); - }; + let auth_message = format!( + "{},{},{}", + client_first_bare, server_first, client_final.without_proof + ); + let Ok(client_proof) = BASE64.decode(client_final.proof.as_bytes()) else { + break 'sasl PgAuthenticationResult::Failed; + }; + let Ok(Some(server_signature)) = + verifier.verify_client_proof(auth_message.as_bytes(), &client_proof) + else { + break 'sasl PgAuthenticationResult::Failed; + }; + let Some(user_info) = user_info else { + break 'sasl PgAuthenticationResult::Failed; + }; - drop(state); - if self - .login_verifier - .authorize(login_info, &user_info) - .await - .is_err() - { - return Ok(PgAuthenticationResult::Failed); - } + drop(state); + if self + .login_verifier + .authorize(login_info, &user_info) + .await + .is_err() + { + // `authorize` already recorded this failure; return early + // to bypass `record_scram_failure` and avoid double-counting. + return Ok(PgAuthenticationResult::Failed); + } - client - .send(PgWireBackendMessage::Authentication( - Authentication::SASLFinal( - format!("v={}", BASE64.encode(server_signature)).into(), - ), - )) - .await?; - Ok(PgAuthenticationResult::Success(user_info)) + client + .send(PgWireBackendMessage::Authentication( + Authentication::SASLFinal( + format!("v={}", BASE64.encode(server_signature)).into(), + ), + )) + .await?; + PgAuthenticationResult::Success(user_info) + }; + Ok(Self::record_scram_failure(result)) } PgAuthenticationState::Initial => Ok(PgAuthenticationResult::Failed), } From 4873fc4f181a706315746d93d5ce917f7cb2024c Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 16 Jul 2026 16:47:03 +0800 Subject: [PATCH 46/48] feat(meta): add retention GC for soft-dropped tables (#8526) * feat(meta): persist soft-drop retention metadata Signed-off-by: Lei, HUANG * fix(meta): complete soft-drop retention coverage Signed-off-by: Lei, HUANG * fix(meta): finalize soft-drop metadata after prepare Signed-off-by: Lei, HUANG * feat(meta): purge expired soft-dropped tables Signed-off-by: Lei, HUANG * fix(meta): keep soft-drop GC responsive Signed-off-by: Lei, HUANG * fix(meta): deduplicate soft-drop purge tasks Signed-off-by: Lei, HUANG * fix(meta): release purge reservations safely Signed-off-by: Lei, HUANG * refactor(meta): simplify soft-drop GC wiring Signed-off-by: Lei, HUANG * feat(meta): persist soft-drop retention deadlines Signed-off-by: Lei, HUANG * fix(meta): retain soft-drop config for recovery Signed-off-by: Lei, HUANG * fix: complete standalone DDL test context Signed-off-by: Lei, HUANG * fix(meta): recheck retention before automatic purge Signed-off-by: Lei, HUANG * fix(meta): keep experimental soft drop disabled Signed-off-by: Lei, HUANG * fix(meta): rotate soft-drop purge candidates Signed-off-by: Lei, HUANG * fix(meta): address soft-drop retention review Signed-off-by: Lei, HUANG * fix(meta): protect retained table tombstones Signed-off-by: Lei, HUANG * fix(meta): preserve post-cleanup purge state Signed-off-by: Lei, HUANG * fix(meta): fence purge by drop generation Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- config/config.md | 3 + config/metasrv.example.toml | 6 + src/cli/src/bench/metadata.rs | 1 + src/cli/src/metadata/control/del/table.rs | 8 +- src/cmd/src/standalone.rs | 1 + src/cmd/tests/load_config_test.rs | 5 + src/common/meta/Cargo.toml | 2 +- src/common/meta/src/ddl.rs | 3 + src/common/meta/src/ddl/drop_table.rs | 73 +- .../meta/src/ddl/drop_table/executor.rs | 33 +- .../meta/src/ddl/drop_table/metadata.rs | 4 +- .../meta/src/ddl/purge_dropped_table.rs | 116 +- src/common/meta/src/ddl/tests/drop_table.rs | 1048 ++++++++++++++++- src/common/meta/src/ddl/undrop_table.rs | 165 ++- src/common/meta/src/ddl_manager.rs | 28 + src/common/meta/src/key.rs | 342 +++++- src/common/meta/src/key/tombstone.rs | 360 +++++- src/common/meta/src/test_util.rs | 1 + src/meta-srv/src/gc.rs | 2 + src/meta-srv/src/gc/ctx.rs | 165 ++- src/meta-srv/src/gc/mock.rs | 20 + src/meta-srv/src/gc/options.rs | 69 ++ src/meta-srv/src/gc/scheduler.rs | 465 +++++++- src/meta-srv/src/metasrv/builder.rs | 66 +- src/meta-srv/src/metrics.rs | 8 + src/meta-srv/src/procedure/utils.rs | 1 + tests-integration/src/standalone.rs | 1 + 27 files changed, 2911 insertions(+), 85 deletions(-) diff --git a/config/config.md b/config/config.md index a04178bdb93..93575fcce44 100644 --- a/config/config.md +++ b/config/config.md @@ -437,6 +437,9 @@ | `gc` | -- | -- | -- | | `gc.enable` | Bool | `false` | Whether GC is enabled. Default to false. Need to be the same with datanode's `mito.gc.enable`
If set to false, no GC will be performed | | `gc.gc_cooldown_period` | String | `5m` | Cooldown period between GC operations on the same region. | +| `gc.experimental_soft_drop` | -- | -- | -- | +| `gc.experimental_soft_drop.enable` | Bool | `false` | Reserved experimental option. Currently ignored. | +| `gc.experimental_soft_drop.retention` | String | `7d` | Reserved retention duration. Currently ignored. | | `logging` | -- | -- | The logging options. | | `logging.dir` | String | `./greptimedb_data/logs` | The directory to store the log files. If set to empty, logs will not be written to files. | | `logging.level` | String | Unset | The log level. Can be `info`/`debug`/`warn`/`error`. | diff --git a/config/metasrv.example.toml b/config/metasrv.example.toml index 28d831d4f2a..55e6497078b 100644 --- a/config/metasrv.example.toml +++ b/config/metasrv.example.toml @@ -326,6 +326,12 @@ enable = false ## Cooldown period between GC operations on the same region. gc_cooldown_period = "5m" +[gc.experimental_soft_drop] +## Reserved experimental option. Currently ignored. +enable = false +## Reserved retention duration. Currently ignored. +retention = "7d" + ## The logging options. [logging] ## The directory to store the log files. If set to empty, logs will not be written to files. diff --git a/src/cli/src/bench/metadata.rs b/src/cli/src/bench/metadata.rs index c21db072b1f..18c25c24ae9 100644 --- a/src/cli/src/bench/metadata.rs +++ b/src/cli/src/bench/metadata.rs @@ -119,6 +119,7 @@ impl TableMetadataBencher { &table_info.table_name(), &table_route, ®ion_wal_options, + None, ) .await; start.elapsed() diff --git a/src/cli/src/metadata/control/del/table.rs b/src/cli/src/metadata/control/del/table.rs index c92e5732248..59ebb2459ed 100644 --- a/src/cli/src/metadata/control/del/table.rs +++ b/src/cli/src/metadata/control/del/table.rs @@ -118,7 +118,13 @@ impl TableMetadataDeleter { .map_err(BoxedError::new)?; self.table_metadata_manager - .delete_table_metadata(table_id, &table_name, &table_route, ®ion_wal_options) + .delete_table_metadata( + table_id, + &table_name, + &table_route, + ®ion_wal_options, + None, + ) .await .map_err(BoxedError::new)?; Ok(()) diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 5096c9221cd..7a28366eb86 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -552,6 +552,7 @@ impl StartCommand { flow_metadata_allocator: flow_metadata_allocator.clone(), region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), soft_drop_enabled: false, + soft_drop_retention: None, }; let ddl_manager = DdlManager::try_new( diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 434c6507b7f..5814cc33b45 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -189,6 +189,11 @@ fn test_load_metasrv_example_config() { let options = GreptimeOptions::::load_layered_options(example_config.to_str(), "") .unwrap(); + assert!(!options.component.gc.experimental_soft_drop.enable); + assert_eq!( + Duration::from_secs(7 * 24 * 60 * 60), + options.component.gc.experimental_soft_drop.retention + ); let expected = GreptimeOptions:: { component: MetasrvOptions { selector: SelectorType::default(), diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 5134c40a804..90e631679a7 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -92,6 +92,7 @@ tokio-postgres-rustls = { version = "0.12", optional = true } tonic.workspace = true tracing.workspace = true typetag.workspace = true +uuid.workspace = true [dev-dependencies] chrono.workspace = true @@ -100,4 +101,3 @@ common-procedure-test.workspace = true common-test-util.workspace = true common-wal = { workspace = true, features = ["testing"] } datatypes.workspace = true -uuid.workspace = true diff --git a/src/common/meta/src/ddl.rs b/src/common/meta/src/ddl.rs index cc578b7552f..7f168e921e5 100644 --- a/src/common/meta/src/ddl.rs +++ b/src/common/meta/src/ddl.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::sync::Arc; +use std::time::Duration; use store_api::storage::{RegionId, TableId}; @@ -121,6 +122,8 @@ pub struct DdlContext { pub region_failure_detector_controller: RegionFailureDetectorControllerRef, /// Whether table drops should stop after tombstoning metadata. pub soft_drop_enabled: bool, + /// Fixed retention used to calculate new soft-drop deadlines. + pub soft_drop_retention: Option, } impl DdlContext { diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs index 8747bb111d2..f3f085d1f24 100644 --- a/src/common/meta/src/ddl/drop_table.rs +++ b/src/common/meta/src/ddl/drop_table.rs @@ -26,6 +26,7 @@ use common_procedure::{ }; use common_telemetry::info; use common_telemetry::tracing::warn; +use common_time::util::current_time_millis; use common_wal::options::WalOptions; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt}; @@ -33,13 +34,14 @@ use store_api::storage::RegionNumber; use strum::AsRefStr; use table::metadata::TableId; use table::table_reference::TableReference; +use uuid::Uuid; use self::executor::DropTableExecutor; use crate::ddl::DdlContext; use crate::ddl::utils::{convert_region_routes_to_detecting_regions, map_to_procedure_error}; use crate::error::{self, Result}; use crate::key::table_route::TableRouteValue; -use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; +use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; use crate::region_keeper::OperatingRegionGuard; use crate::rpc::ddl::DropTableTask; @@ -60,7 +62,7 @@ impl DropTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropTable"; pub fn new(task: DropTableTask, context: DdlContext) -> Self { - let data = DropTableData::new(task); + let data = DropTableData::new(task, context.soft_drop_enabled, context.soft_drop_retention); let executor = data.build_executor(); Self { context, @@ -71,7 +73,16 @@ impl DropTableProcedure { } pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult { - let data: DropTableData = serde_json::from_str(json).context(FromJsonSnafu)?; + let mut data: DropTableData = serde_json::from_str(json).context(FromJsonSnafu)?; + if data.state == DropTableState::Prepare + && data.soft_drop_enabled + && data.dropped_at.is_none() + && data.soft_drop_retention_millis.is_none() + { + data.soft_drop_retention_millis = context + .soft_drop_retention + .and_then(|retention| i64::try_from(retention.as_millis()).ok()); + } let executor = data.build_executor(); Ok(Self { @@ -83,10 +94,36 @@ impl DropTableProcedure { } pub(crate) async fn on_prepare(&mut self) -> Result { - if self.executor.on_prepare(&self.context).await?.stop() { + if self + .executor + .on_prepare(&self.context, self.data.soft_drop_enabled) + .await? + .stop() + { return Ok(Status::done()); } self.fill_table_metadata().await?; + if self.data.soft_drop_enabled && self.data.dropped_at.is_none() { + let dropped_at = current_time_millis(); + let retention_millis = + self.data + .soft_drop_retention_millis + .context(error::UnexpectedSnafu { + err_msg: "Soft-drop retention is missing from the procedure".to_string(), + })?; + let retention_expires_at = dropped_at.checked_add(retention_millis).context( + error::UnexpectedSnafu { + err_msg: format!( + "Soft-drop retention deadline overflows i64: dropped_at={dropped_at}, retention_millis={retention_millis}" + ), + }, + )?; + self.data.dropped_at = Some(dropped_at); + self.data.retention_expires_at = Some(retention_expires_at); + } + if self.data.soft_drop_enabled && self.data.drop_generation.is_none() { + self.data.drop_generation = Some(Uuid::new_v4().to_string()); + } self.data.state = DropTableState::DeleteMetadata; Ok(Status::executing(true)) @@ -139,6 +176,9 @@ impl DropTableProcedure { &self.context, table_route_value, &self.data.region_wal_options, + self.data.dropped_at, + self.data.retention_expires_at, + self.data.drop_generation.as_deref(), ) .await?; info!("Deleted table metadata for table {table_id}"); @@ -173,7 +213,7 @@ impl DropTableProcedure { .await?; } - if self.context.soft_drop_enabled { + if self.data.soft_drop_enabled { self.executor .on_close_regions( &self.context.node_manager, @@ -286,6 +326,7 @@ impl Procedure for DropTableProcedure { let lock_key = vec![ CatalogLock::Read(table_ref.catalog).into(), SchemaLock::read(table_ref.catalog, table_ref.schema).into(), + TableNameLock::new(table_ref.catalog, table_ref.schema, table_ref.table).into(), TableLock::Write(table_id).into(), ]; @@ -329,10 +370,24 @@ pub struct DropTableData { pub region_wal_options: HashMap, #[serde(default)] pub allow_rollback: bool, + #[serde(default)] + pub soft_drop_enabled: bool, + #[serde(default)] + pub dropped_at: Option, + #[serde(default)] + pub retention_expires_at: Option, + #[serde(default)] + pub soft_drop_retention_millis: Option, + #[serde(default)] + pub drop_generation: Option, } impl DropTableData { - pub fn new(task: DropTableTask) -> Self { + pub fn new( + task: DropTableTask, + soft_drop_enabled: bool, + soft_drop_retention: Option, + ) -> Self { Self { state: DropTableState::Prepare, task, @@ -340,6 +395,12 @@ impl DropTableData { physical_table_id: None, region_wal_options: HashMap::new(), allow_rollback: false, + soft_drop_enabled, + dropped_at: None, + retention_expires_at: None, + soft_drop_retention_millis: soft_drop_retention + .and_then(|retention| i64::try_from(retention.as_millis()).ok()), + drop_generation: None, } } diff --git a/src/common/meta/src/ddl/drop_table/executor.rs b/src/common/meta/src/ddl/drop_table/executor.rs index 1b0993170e6..61e74d1b5cf 100644 --- a/src/common/meta/src/ddl/drop_table/executor.rs +++ b/src/common/meta/src/ddl/drop_table/executor.rs @@ -36,6 +36,7 @@ use crate::ddl::utils::{ use crate::ddl::{CreateRequestBuilder, DdlContext, build_template_from_raw_table_info}; use crate::error::{self, Result}; use crate::instruction::CacheIdent; +use crate::key::DroppedTableLifecycle; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; use crate::node_manager::NodeManagerRef; @@ -86,7 +87,11 @@ impl DropTableExecutor { /// - Throws an error if table not exists and `drop_if_exists` is `false`. /// - Rejects dropping a recreated live table while an older tombstone still owns the same /// fully qualified name. - pub async fn on_prepare(&self, ctx: &DdlContext) -> Result> { + pub async fn on_prepare( + &self, + ctx: &DdlContext, + soft_drop_enabled: bool, + ) -> Result> { let table_ref = self.table.table_ref(); let exist = ctx @@ -110,12 +115,12 @@ impl DropTableExecutor { } ); - if ctx.soft_drop_enabled - && let Some(dropped_table) = ctx - .table_metadata_manager - .get_dropped_table(&self.table) - .await? + if let Some(dropped_table) = ctx + .table_metadata_manager + .get_dropped_table(&self.table) + .await? && dropped_table.table_id != self.table_id + && (soft_drop_enabled || dropped_table.dropped_at.is_some()) { return error::TableNameTombstoneConflictSnafu { table_name: table_ref.to_string(), @@ -134,13 +139,21 @@ impl DropTableExecutor { ctx: &DdlContext, table_route_value: &TableRouteValue, region_wal_options: &HashMap, + dropped_at: Option, + retention_expires_at: Option, + drop_generation: Option<&str>, ) -> Result<()> { ctx.table_metadata_manager - .delete_table_metadata( + .delete_table_metadata_with_retention_and_generation( self.table_id, &self.table, table_route_value, region_wal_options, + DroppedTableLifecycle { + dropped_at, + retention_expires_at, + drop_generation, + }, ) .await } @@ -560,7 +573,7 @@ mod tests { 1024, true, ); - let ctrl = executor.on_prepare(&ctx).await.unwrap(); + let ctrl = executor.on_prepare(&ctx, false).await.unwrap(); assert!(ctrl.stop()); // Drops a non-exists table @@ -569,7 +582,7 @@ mod tests { 1024, false, ); - let err = executor.on_prepare(&ctx).await.unwrap_err(); + let err = executor.on_prepare(&ctx, false).await.unwrap_err(); assert_matches!(err, error::Error::TableNotFound { .. }); // Drops a exists table @@ -587,7 +600,7 @@ mod tests { ) .await .unwrap(); - let ctrl = executor.on_prepare(&ctx).await.unwrap(); + let ctrl = executor.on_prepare(&ctx, false).await.unwrap(); assert!(!ctrl.stop()); } } diff --git a/src/common/meta/src/ddl/drop_table/metadata.rs b/src/common/meta/src/ddl/drop_table/metadata.rs index 982182ffbf1..213ebb461b3 100644 --- a/src/common/meta/src/ddl/drop_table/metadata.rs +++ b/src/common/meta/src/ddl/drop_table/metadata.rs @@ -50,7 +50,7 @@ impl DropTableProcedure { ); // TODO(hl): support soft-dropping logical tables. ensure!( - !(self.context.soft_drop_enabled + !(self.data.soft_drop_enabled && is_metric_engine_logical_table( &table_info_value.table_info, &table_route_value @@ -63,7 +63,7 @@ impl DropTableProcedure { if physical_table_id == self.data.table_id() { let engine = table_info_value.table_info.meta.engine; ensure!( - !(self.context.soft_drop_enabled && engine == FILE_ENGINE), + !(self.data.soft_drop_enabled && engine == FILE_ENGINE), error::UnsupportedSnafu { operation: "soft-dropping file-engine tables".to_string() } diff --git a/src/common/meta/src/ddl/purge_dropped_table.rs b/src/common/meta/src/ddl/purge_dropped_table.rs index 759e908c8b9..040a0d8e60d 100644 --- a/src/common/meta/src/ddl/purge_dropped_table.rs +++ b/src/common/meta/src/ddl/purge_dropped_table.rs @@ -19,6 +19,7 @@ use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ Context as ProcedureContext, LockKey, Procedure, Result as ProcedureResult, Status, }; +use common_time::util::current_time_millis; use common_wal::options::WalOptions; use serde::{Deserialize, Serialize}; use snafu::{ResultExt, ensure}; @@ -34,6 +35,7 @@ use crate::ddl::utils::{ map_to_procedure_error, }; use crate::error::{self, Result}; +use crate::key::DroppedTableMetadata; use crate::key::table_route::TableRouteValue; use crate::lock_key::TableLock; use crate::rpc::ddl::PurgeDroppedTableTask; @@ -46,6 +48,7 @@ pub struct PurgeDroppedTableProcedure { impl PurgeDroppedTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::PurgeDroppedTable"; + pub const EXPIRED_TYPE_NAME: &'static str = "metasrv-procedure::PurgeExpiredDroppedTable"; pub fn new(task: PurgeDroppedTableTask, context: DdlContext) -> Self { Self { @@ -54,12 +57,19 @@ impl PurgeDroppedTableProcedure { } } + pub fn new_if_expired(task: PurgeDroppedTableTask, context: DdlContext) -> Self { + Self { + context, + data: PurgeDroppedTableData::new_if_expired(task), + } + } + pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult { let data: PurgeDroppedTableData = serde_json::from_str(json).context(FromJsonSnafu)?; Ok(Self { context, data }) } - async fn on_prepare(&mut self) -> Result { + async fn eligible_dropped_table(&self) -> Result> { let dropped_table = self .context .table_metadata_manager @@ -67,8 +77,25 @@ impl PurgeDroppedTableProcedure { .await?; let Some(dropped_table) = dropped_table else { - return Ok(Status::done()); + return Ok(None); }; + if self.data.check_expired { + if !self.context.soft_drop_enabled { + return Ok(None); + } + let retention_millis = self + .context + .soft_drop_retention + .and_then(|retention| i64::try_from(retention.as_millis()).ok()); + let expires_at = dropped_table.retention_expires_at.or_else(|| { + dropped_table.dropped_at.and_then(|dropped_at| { + retention_millis.and_then(|retention| dropped_at.checked_add(retention)) + }) + }); + if expires_at.is_none_or(|expires_at| expires_at > current_time_millis()) { + return Ok(None); + } + } ensure!( !is_metric_engine_logical_table( &dropped_table.table_info_value.table_info, @@ -78,16 +105,58 @@ impl PurgeDroppedTableProcedure { operation: "purging metric logical tables".to_string() } ); + Ok(Some(dropped_table)) + } + + fn update_dropped_table(&mut self, dropped_table: DroppedTableMetadata) { self.data.table_id = Some(dropped_table.table_id); self.data.table_name = Some(dropped_table.table_name); self.data.table_info = Some(dropped_table.table_info_value.table_info); self.data.table_route_value = Some(dropped_table.table_route_value); self.data.region_wal_options = dropped_table.region_wal_options; + self.data.drop_generation = dropped_table.drop_generation; + self.data.drop_generation_loaded = true; + } + + async fn on_prepare(&mut self) -> Result { + let Some(dropped_table) = self.eligible_dropped_table().await? else { + return Ok(Status::done()); + }; + self.update_dropped_table(dropped_table); self.data.state = PurgeDroppedTableState::DropRegions; Ok(Status::executing(true)) } async fn on_drop_regions(&mut self) -> Result { + let expected_generation = self.data.drop_generation.as_deref().unwrap_or_default(); + let purge_claim = self + .context + .table_metadata_manager + .dropped_table_purge_claim(self.data.task.table_id) + .await?; + if let Some(purge_claim) = purge_claim { + if purge_claim != expected_generation { + return Ok(Status::done()); + } + self.data.purging_claimed = true; + } else { + // Revalidate after recovery before crossing the durable cleanup boundary. + let Some(dropped_table) = self.eligible_dropped_table().await? else { + return Ok(Status::done()); + }; + if dropped_table.drop_generation != self.data.drop_generation { + return Ok(Status::done()); + } + self.update_dropped_table(dropped_table); + self.context + .table_metadata_manager + .mark_dropped_table_purging( + self.data.task.table_id, + self.data.drop_generation.as_deref(), + ) + .await?; + self.data.purging_claimed = true; + } if let Some(region_routes) = self.data.physical_region_routes() { self.executor() .on_cleanup_regions_offline( @@ -109,6 +178,19 @@ impl PurgeDroppedTableProcedure { } async fn on_delete_tombstone(&mut self) -> Result { + if self.data.purging_claimed { + let expected_generation = self.data.drop_generation.as_deref().unwrap_or_default(); + if self + .context + .table_metadata_manager + .dropped_table_purge_claim(self.data.task.table_id) + .await? + .as_deref() + != Some(expected_generation) + { + return Ok(Status::done()); + } + } self.context .table_metadata_manager .delete_table_metadata_tombstone( @@ -129,10 +211,19 @@ impl PurgeDroppedTableProcedure { #[async_trait] impl Procedure for PurgeDroppedTableProcedure { fn type_name(&self) -> &str { - Self::TYPE_NAME + if self.data.check_expired { + Self::EXPIRED_TYPE_NAME + } else { + Self::TYPE_NAME + } } fn recover(&mut self) -> ProcedureResult<()> { + if self.data.state == PurgeDroppedTableState::DropRegions + && !self.data.drop_generation_loaded + { + self.data.state = PurgeDroppedTableState::Prepare; + } Ok(()) } @@ -164,6 +255,14 @@ pub struct PurgeDroppedTableData { table_route_value: Option, #[serde(default)] region_wal_options: HashMap, + #[serde(default)] + check_expired: bool, + #[serde(default)] + drop_generation: Option, + #[serde(default)] + drop_generation_loaded: bool, + #[serde(default)] + purging_claimed: bool, } impl PurgeDroppedTableData { @@ -176,6 +275,17 @@ impl PurgeDroppedTableData { table_info: None, table_route_value: None, region_wal_options: HashMap::new(), + check_expired: false, + drop_generation: None, + drop_generation_loaded: false, + purging_claimed: false, + } + } + + fn new_if_expired(task: PurgeDroppedTableTask) -> Self { + Self { + check_expired: true, + ..Self::new(task) } } diff --git a/src/common/meta/src/ddl/tests/drop_table.rs b/src/common/meta/src/ddl/tests/drop_table.rs index b14fe77cc5b..dd960df75c4 100644 --- a/src/common/meta/src/ddl/tests/drop_table.rs +++ b/src/common/meta/src/ddl/tests/drop_table.rs @@ -15,6 +15,7 @@ use std::assert_matches; use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use std::time::Duration; use api::v1::region::{RegionRequest, region_request}; use async_trait::async_trait; @@ -41,20 +42,297 @@ use crate::ddl::test_util::{ put_datanode_address, test_create_logical_table_task, test_create_physical_table_task, }; use crate::ddl::undrop_table::UndropTableProcedure; -use crate::ddl::{DetectingRegion, RegionFailureDetectorController, TableMetadata}; +use crate::ddl::{DdlContext, DetectingRegion, RegionFailureDetectorController, TableMetadata}; use crate::error::Error; +use crate::key::MetadataKey; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; +use crate::kv_backend::KvBackend; use crate::kv_backend::memory::MemoryKvBackend; use crate::peer::Peer; use crate::rpc::ddl::{DropTableTask, PurgeDroppedTableTask, UndropTableTask}; use crate::rpc::router::{Region, RegionRoute}; +use crate::rpc::store::{BatchDeleteRequest, PutRequest}; use crate::test_util::{MockDatanodeManager, new_ddl_context, new_ddl_context_with_kv_backend}; +fn dropped_at_marker_key(table_id: TableId) -> String { + format!("__tombstone/__dropped_at/{table_id}") +} + +fn retention_expires_at_marker_key(table_id: TableId) -> String { + format!("__tombstone/__retention_expires_at/{table_id}") +} + +fn purging_marker_key(table_id: TableId) -> String { + format!("__tombstone/__purging/{table_id}") +} + +fn drop_generation_marker_key(table_id: TableId) -> String { + format!("__tombstone/__drop_generation/{table_id}") +} + +async fn create_dropped_table( + context: &DdlContext, + table_id: TableId, + dropped_at: Option, + retention_expires_at: Option, +) { + let task = test_create_table_task("foo", table_id); + let table_name = task.table_name(); + let table_route = TableRouteValue::physical(vec![]); + context + .table_metadata_manager + .create_table_metadata(task.table_info, table_route.clone(), HashMap::new()) + .await + .unwrap(); + context + .table_metadata_manager + .delete_table_metadata_with_retention( + table_id, + &table_name, + &table_route, + &HashMap::new(), + dropped_at, + retention_expires_at, + ) + .await + .unwrap(); +} + +async fn undrop_at_restore_metadata( + context: &DdlContext, + table_id: TableId, +) -> UndropTableProcedure { + create_dropped_table(context, table_id, Some(1), Some(i64::MAX)).await; + let mut undrop = UndropTableProcedure::new(new_undrop_table_task(table_id), context.clone()); + let procedure_context = new_test_procedure_context(); + undrop.execute(&procedure_context).await.unwrap(); + undrop.execute(&procedure_context).await.unwrap(); + undrop +} + +fn legacy_undrop_snapshot(procedure: &UndropTableProcedure) -> String { + let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + let data = data.as_object_mut().unwrap(); + for field in [ + "dropped_at", + "retention_expires_at", + "drop_generation", + "tombstone_identity_loaded", + ] { + data.remove(field); + } + serde_json::to_string(data).unwrap() +} + +#[test] +fn test_old_drop_table_json_defaults_to_hard_drop() { + let mut runtime_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + runtime_context.soft_drop_enabled = true; + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", 1024, false), + runtime_context.clone(), + ); + let mut old_data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + old_data + .as_object_mut() + .unwrap() + .remove("soft_drop_enabled"); + old_data.as_object_mut().unwrap().remove("dropped_at"); + + let recovered = DropTableProcedure::from_json(&old_data.to_string(), runtime_context).unwrap(); + let recovered_data: serde_json::Value = + serde_json::from_str(&recovered.dump().unwrap()).unwrap(); + + assert_eq!(recovered_data["soft_drop_enabled"], false); + assert_eq!(recovered_data["dropped_at"], serde_json::Value::Null); +} + +#[test] +fn test_drop_table_lock_key_includes_table_name() { + let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); + + let keys = procedure + .lock_key() + .keys_to_lock() + .cloned() + .collect::>(); + assert!( + keys.iter().any(|key| matches!( + key, + StringKey::Exclusive(key) if key == "__table_name_lock/greptime.public.foo" + )), + "drop lock keys should include the table name: {keys:?}" + ); + assert!( + keys.iter().any(|key| matches!( + key, + StringKey::Exclusive(key) if key == "__table_lock/1024" + )), + "drop lock keys should include the table id: {keys:?}" + ); +} + +#[tokio::test] +async fn test_recovered_soft_drop_preserves_persisted_mode() { + let (tx, mut rx) = mpsc::channel(8); + let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut soft_context = + new_ddl_context_with_kv_backend(node_manager.clone(), kv_backend.clone()); + soft_context.soft_drop_enabled = true; + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + soft_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(RegionId::new(table_id, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + HashMap::new(), + ) + .await + .unwrap(); + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", table_id, false), + soft_context.clone(), + ); + + let hard_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); + let mut recovered = + DropTableProcedure::from_json(&procedure.dump().unwrap(), hard_context).unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + soft_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); + let (_, request) = rx.try_recv().unwrap(); + assert_matches!(request.body, Some(region_request::Body::Close(_))); +} + +#[tokio::test] +async fn test_legacy_soft_drop_prepare_preserves_persisted_mode() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let mut context = new_ddl_context(node_manager); + context.soft_drop_enabled = false; + context.soft_drop_retention = Some(Duration::from_millis(100)); + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + + let mut persisted: serde_json::Value = serde_json::from_str( + &DropTableProcedure::new(new_drop_table_task("foo", table_id, false), context.clone()) + .dump() + .unwrap(), + ) + .unwrap(); + persisted["soft_drop_enabled"] = true.into(); + persisted + .as_object_mut() + .unwrap() + .remove("soft_drop_retention_millis"); + persisted + .as_object_mut() + .unwrap() + .remove("retention_expires_at"); + + let mut recovered = + DropTableProcedure::from_json(&persisted.to_string(), context.clone()).unwrap(); + recovered.on_prepare().await.unwrap(); + + assert!(recovered.data.soft_drop_enabled); + assert_eq!(Some(100), recovered.data.soft_drop_retention_millis); + assert!(recovered.data.dropped_at.is_some()); + assert!(recovered.data.retention_expires_at.is_some()); +} + +#[tokio::test] +async fn test_disabled_soft_drop_creates_hard_drop_procedure() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let mut context = new_ddl_context(node_manager); + context.soft_drop_enabled = false; + context.soft_drop_retention = Some(Duration::from_millis(100)); + + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); + + assert!(!procedure.data.soft_drop_enabled); + assert_eq!(None, procedure.data.dropped_at); + assert_eq!(None, procedure.data.retention_expires_at); +} + +#[tokio::test] +async fn test_recovered_hard_drop_ignores_soft_runtime_context() { + let (tx, mut rx) = mpsc::channel(8); + let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let hard_context = new_ddl_context_with_kv_backend(node_manager.clone(), kv_backend.clone()); + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + hard_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(RegionId::new(table_id, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + HashMap::new(), + ) + .await + .unwrap(); + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", table_id, false), + hard_context.clone(), + ); + + let mut soft_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); + soft_context.soft_drop_enabled = true; + let mut recovered = + DropTableProcedure::from_json(&procedure.dump().unwrap(), soft_context).unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + hard_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_none() + ); + let (_, request) = rx.try_recv().unwrap(); + assert_matches!(request.body, Some(region_request::Body::Drop(_))); +} + #[tokio::test] async fn test_on_prepare_table_not_exists_err() { let node_manager = Arc::new(MockDatanodeManager::new(())); - let ddl_context = new_ddl_context(node_manager); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + ddl_context.soft_drop_retention = Some(Duration::from_millis(100)); let table_name = "foo"; let table_id = 1024; let task = test_create_table_task(table_name, table_id); @@ -71,8 +349,86 @@ async fn test_on_prepare_table_not_exists_err() { let task = new_drop_table_task("bar", table_id, false); let mut procedure = DropTableProcedure::new(task, ddl_context); + assert_eq!(procedure.data.dropped_at, None); + assert_eq!(procedure.data.retention_expires_at, None); let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::TableNotFound); + assert_eq!(procedure.data.dropped_at, None); + assert_eq!(procedure.data.retention_expires_at, None); +} + +#[tokio::test] +async fn test_soft_drop_prepare_assigns_stable_deadline_and_recovers_it() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + ddl_context.soft_drop_retention = Some(Duration::from_millis(100)); + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + let mut procedure = DropTableProcedure::new( + new_drop_table_task("foo", table_id, false), + ddl_context.clone(), + ); + + procedure.on_prepare().await.unwrap(); + let dropped_at = procedure.data.dropped_at.unwrap(); + let deadline = procedure.data.retention_expires_at.unwrap(); + assert_eq!(deadline, dropped_at + 100); + + let mut changed_context = ddl_context; + changed_context.soft_drop_retention = Some(Duration::from_secs(10)); + let recovered = + DropTableProcedure::from_json(&procedure.dump().unwrap(), changed_context).unwrap(); + assert_eq!(recovered.data.dropped_at, Some(dropped_at)); + assert_eq!(recovered.data.retention_expires_at, Some(deadline)); +} + +#[tokio::test] +async fn test_soft_drop_prepare_rejects_deadline_overflow_before_metadata_delete() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + ddl_context.soft_drop_retention = Some(Duration::from_millis(i64::MAX as u64)); + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + let mut procedure = DropTableProcedure::new( + new_drop_table_task("foo", table_id, false), + ddl_context.clone(), + ); + + let error = procedure.on_prepare().await.unwrap_err(); + assert!(error.to_string().contains("retention deadline"), "{error}"); + assert!( + ddl_context + .table_metadata_manager + .table_name_manager() + .get(TableNameKey::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + "foo", + )) + .await + .unwrap() + .is_some() + ); } #[tokio::test] @@ -246,6 +602,7 @@ async fn test_soft_drop_closes_regions_and_keeps_tombstone() { let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); let mut ddl_context = new_ddl_context(node_manager); ddl_context.soft_drop_enabled = true; + ddl_context.soft_drop_retention = Some(Duration::from_millis(100)); ddl_context.region_failure_detector_controller = detector_controller.clone(); let table_id = 1024; let table_name = "foo"; @@ -314,12 +671,35 @@ async fn test_soft_drop_closes_regions_and_keeps_tombstone() { .unwrap(); assert!(live_table.is_none()); + let dropped_at = procedure.data.dropped_at.unwrap(); let dropped_table = ddl_context .table_metadata_manager .get_dropped_table(&table_name) .await + .unwrap() .unwrap(); - assert!(dropped_table.is_some()); + assert_eq!(dropped_table.dropped_at, Some(dropped_at)); + assert_eq!(dropped_table.retention_expires_at, Some(dropped_at + 100)); + let dropped_tables = ddl_context + .table_metadata_manager + .list_dropped_tables() + .await + .unwrap(); + assert_eq!(dropped_tables[0].dropped_at, Some(dropped_at)); + assert_eq!( + dropped_tables[0].retention_expires_at, + Some(dropped_at + 100) + ); + + assert!( + ddl_context + .table_metadata_manager + .kv_backend() + .get(retention_expires_at_marker_key(table_id).as_bytes(),) + .await + .unwrap() + .is_some() + ); assert_eq!( detector_controller.deregistered().await, @@ -330,6 +710,69 @@ async fn test_soft_drop_closes_regions_and_keeps_tombstone() { ); } +#[tokio::test] +async fn test_soft_drop_timestamp_is_stable_across_retry_and_recovery() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", table_id, false), + ddl_context.clone(), + ); + assert_eq!(procedure.data.dropped_at, None); + let mut procedure = procedure; + procedure.on_prepare().await.unwrap(); + let dropped_at = procedure.data.dropped_at.unwrap(); + let persisted = procedure.dump().unwrap(); + let recovered = DropTableProcedure::from_json(&persisted, ddl_context.clone()).unwrap(); + assert_eq!(recovered.data.dropped_at, Some(dropped_at)); + let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + data["dropped_at"] = 1_234_567_890_i64.into(); + let mut recovered = + DropTableProcedure::from_json(&data.to_string(), ddl_context.clone()).unwrap(); + + recovered.on_delete_metadata().await.unwrap(); + recovered.on_delete_metadata().await.unwrap(); + let persisted = recovered.dump().unwrap(); + let mut recovered_again = + DropTableProcedure::from_json(&persisted, ddl_context.clone()).unwrap(); + recovered_again.on_delete_metadata().await.unwrap(); + + let marker = kv_backend + .get(dropped_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .unwrap(); + assert_eq!(marker.value, b"1234567890"); + let deadline_marker = kv_backend + .get(retention_expires_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .unwrap(); + assert_eq!( + deadline_marker.value, + recovered_again + .data + .retention_expires_at + .unwrap() + .to_string() + .as_bytes() + ); + assert_eq!(recovered_again.data.dropped_at, Some(1_234_567_890)); +} + #[tokio::test] async fn test_hard_drop_keeps_delete_tombstone_flow() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); @@ -430,7 +873,7 @@ async fn test_create_table_succeeds_while_tombstone_exists() { } #[tokio::test] -async fn test_drop_recreated_table_fails_when_previous_tombstone_exists() { +async fn test_hard_drop_recreated_table_fails_when_soft_tombstone_exists() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); let mut ddl_context = new_ddl_context(node_manager); ddl_context.soft_drop_enabled = true; @@ -451,6 +894,7 @@ async fn test_drop_recreated_table_fails_when_previous_tombstone_exists() { let drop_task = new_drop_table_task(table_name, original_table_id, false); let mut drop_procedure = DropTableProcedure::new(drop_task, ddl_context.clone()); execute_procedure_until_done(&mut drop_procedure).await; + ddl_context.soft_drop_enabled = false; ddl_context .table_metadata_manager @@ -520,7 +964,8 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { let datanode_handler = DatanodeWatcher::new(tx); let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); - let mut ddl_context = new_ddl_context(node_manager); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); ddl_context.soft_drop_enabled = true; ddl_context.region_failure_detector_controller = detector_controller.clone(); let table_id = 1024; @@ -548,6 +993,16 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { ddl_context.clone(), ); execute_procedure_until_done(&mut drop_procedure).await; + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .unwrap() + .dropped_at + .is_some() + ); detector_controller.clear().await; while rx.try_recv().is_ok() {} @@ -576,6 +1031,20 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { .unwrap() .is_none() ); + assert!( + kv_backend + .get(dropped_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); + assert!( + kv_backend + .get(retention_expires_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); let mut opened_regions = HashSet::new(); for _ in 0..2 { @@ -732,10 +1201,15 @@ async fn test_soft_drop_metric_logical_table_fails() { let mut procedure = DropTableProcedure::new( new_drop_table_task("foo", logical_table_id, false), - ddl_context, + ddl_context.clone(), ); let err = procedure.on_prepare().await.unwrap_err(); - + assert_eq!(err.status_code(), StatusCode::Unsupported); + let persisted_soft_drop = procedure.dump().unwrap(); + ddl_context.soft_drop_enabled = false; + let mut recovered = + DropTableProcedure::from_json(&persisted_soft_drop, ddl_context.clone()).unwrap(); + let err = recovered.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::Unsupported); } @@ -760,10 +1234,15 @@ async fn test_soft_drop_file_engine_table_fails() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), - ddl_context, + ddl_context.clone(), ); let err = procedure.on_prepare().await.unwrap_err(); - + assert_eq!(err.status_code(), StatusCode::Unsupported); + let persisted_soft_drop = procedure.dump().unwrap(); + ddl_context.soft_drop_enabled = false; + let mut recovered = + DropTableProcedure::from_json(&persisted_soft_drop, ddl_context.clone()).unwrap(); + let err = recovered.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::Unsupported); } @@ -1093,6 +1572,7 @@ async fn test_undrop_table_replayed_restore_metadata_is_idempotent() { UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); + procedure.execute(&ctx).await.unwrap(); let restore_metadata_data = procedure.dump().unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -1114,13 +1594,59 @@ async fn test_undrop_table_replayed_restore_metadata_is_idempotent() { assert_eq!(live_table.table_id(), table_id); } +#[tokio::test] +async fn test_legacy_undrop_replays_restore_metadata() { + for partial_restore in [false, true] { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + context.soft_drop_enabled = true; + let table_id = 1024; + let mut procedure = undrop_at_restore_metadata(&context, table_id).await; + let persisted = legacy_undrop_snapshot(&procedure); + + if partial_restore { + let live_key = + TableNameKey::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "foo").to_bytes(); + let tombstone_key = [b"__tombstone/".as_slice(), live_key.as_slice()].concat(); + let value = kv_backend.get(&tombstone_key).await.unwrap().unwrap().value; + kv_backend + .put(PutRequest::new().with_key(live_key).with_value(value)) + .await + .unwrap(); + kv_backend + .batch_delete(BatchDeleteRequest::new().with_keys(vec![tombstone_key])) + .await + .unwrap(); + } else { + procedure + .execute(&new_test_procedure_context()) + .await + .unwrap(); + } + + let mut replayed = UndropTableProcedure::from_json(&persisted, context.clone()).unwrap(); + replayed.recover().unwrap(); + execute_procedure_until_done(&mut replayed).await; + + let (table_info, table_route) = context + .table_metadata_manager + .get_full_table_info(table_id) + .await + .unwrap(); + assert!(table_info.is_some(), "partial_restore={partial_restore}"); + assert!(table_route.is_some(), "partial_restore={partial_restore}"); + } +} + #[tokio::test] async fn test_purge_dropped_table_cleans_regions_offline_and_deletes_tombstone() { let (tx, mut rx) = mpsc::channel(8); let datanode_handler = DatanodeWatcher::new(tx); let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); - let mut ddl_context = new_ddl_context(node_manager); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); ddl_context.soft_drop_enabled = true; ddl_context.region_failure_detector_controller = detector_controller.clone(); let table_id = 1024; @@ -1176,6 +1702,20 @@ async fn test_purge_dropped_table_cleans_regions_offline_and_deletes_tombstone() .unwrap() .is_none() ); + assert!( + kv_backend + .get(dropped_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); + assert!( + kv_backend + .get(retention_expires_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); assert_eq!( detector_controller.deregistered().await, @@ -1183,6 +1723,461 @@ async fn test_purge_dropped_table_cleans_regions_offline_and_deletes_tombstone() ); } +#[tokio::test] +async fn test_automatic_purge_rechecks_unexpired_tombstone() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let table_name = test_create_table_task("foo", table_id).table_name(); + let table_route_value = TableRouteValue::physical(vec![]); + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task("foo", table_id).table_info, + table_route_value.clone(), + HashMap::new(), + ) + .await + .unwrap(); + ddl_context + .table_metadata_manager + .delete_table_metadata_with_retention( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + Some(1), + Some(i64::MAX), + ) + .await + .unwrap(); + + let procedure = PurgeDroppedTableProcedure::new( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + data["check_expired"] = true.into(); + let mut procedure = + PurgeDroppedTableProcedure::from_json(&data.to_string(), ddl_context.clone()).unwrap(); + assert_eq!( + procedure.type_name(), + PurgeDroppedTableProcedure::EXPIRED_TYPE_NAME + ); + + execute_procedure_until_done(&mut procedure).await; + + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); +} + +#[tokio::test] +async fn test_recovered_automatic_purge_rechecks_unexpired_tombstone() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let table_name = test_create_table_task("foo", table_id).table_name(); + let table_route_value = TableRouteValue::physical(vec![]); + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task("foo", table_id).table_info, + table_route_value.clone(), + HashMap::new(), + ) + .await + .unwrap(); + ddl_context + .table_metadata_manager + .delete_table_metadata_with_retention( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + Some(0), + Some(0), + ) + .await + .unwrap(); + + let procedure = PurgeDroppedTableProcedure::new( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + data["check_expired"] = true.into(); + let mut procedure = + PurgeDroppedTableProcedure::from_json(&data.to_string(), ddl_context.clone()).unwrap(); + procedure + .execute(&new_test_procedure_context()) + .await + .unwrap(); + let persisted = procedure.dump().unwrap(); + + ddl_context + .table_metadata_manager + .restore_table_metadata(table_id, &table_name, &table_route_value, &HashMap::new()) + .await + .unwrap(); + ddl_context + .table_metadata_manager + .delete_table_metadata_with_retention( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + Some(1), + Some(i64::MAX), + ) + .await + .unwrap(); + + let mut recovered = + PurgeDroppedTableProcedure::from_json(&persisted, ddl_context.clone()).unwrap(); + recovered.recover().unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); +} + +#[tokio::test] +async fn test_recovered_automatic_purge_obeys_runtime_disable() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let ddl_context = new_ddl_context(node_manager); + let table_id = 1024; + let table_name = test_create_table_task("foo", table_id).table_name(); + let table_route_value = TableRouteValue::physical(vec![]); + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task("foo", table_id).table_info, + table_route_value.clone(), + HashMap::new(), + ) + .await + .unwrap(); + ddl_context + .table_metadata_manager + .delete_table_metadata_with_retention( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + Some(0), + Some(0), + ) + .await + .unwrap(); + + let procedure = PurgeDroppedTableProcedure::new_if_expired( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + let mut recovered = + PurgeDroppedTableProcedure::from_json(&procedure.dump().unwrap(), ddl_context.clone()) + .unwrap(); + recovered.recover().unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); +} + +#[tokio::test] +async fn test_recovered_automatic_purge_finishes_after_cleanup_is_claimed() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut enabled_context = + new_ddl_context_with_kv_backend(node_manager.clone(), kv_backend.clone()); + enabled_context.soft_drop_enabled = true; + let table_id = 1024; + create_dropped_table(&enabled_context, table_id, Some(0), Some(0)).await; + + let mut procedure = PurgeDroppedTableProcedure::new_if_expired( + new_purge_dropped_table_task(table_id), + enabled_context.clone(), + ); + procedure + .execute(&new_test_procedure_context()) + .await + .unwrap(); + let persisted_before_cleanup = procedure.dump().unwrap(); + procedure + .execute(&new_test_procedure_context()) + .await + .unwrap(); + assert!( + kv_backend + .get(purging_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_some() + ); + + let disabled_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + let mut recovered = + PurgeDroppedTableProcedure::from_json(&persisted_before_cleanup, disabled_context.clone()) + .unwrap(); + recovered.recover().unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + disabled_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_none() + ); + assert!( + kv_backend + .get(purging_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); +} + +#[tokio::test] +async fn test_recovered_automatic_purge_ignores_foreign_generation_claim() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut ddl_context = new_ddl_context_with_kv_backend(node_manager.clone(), kv_backend.clone()); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + create_dropped_table(&ddl_context, table_id, Some(0), Some(0)).await; + kv_backend + .put( + PutRequest::new() + .with_key(drop_generation_marker_key(table_id)) + .with_value("generation-a"), + ) + .await + .unwrap(); + + let mut procedure = PurgeDroppedTableProcedure::new_if_expired( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + procedure + .execute(&new_test_procedure_context()) + .await + .unwrap(); + let persisted = procedure.dump().unwrap(); + kv_backend + .put( + PutRequest::new() + .with_key(purging_marker_key(table_id)) + .with_value("generation-b"), + ) + .await + .unwrap(); + + let mut recovered = + PurgeDroppedTableProcedure::from_json(&persisted, ddl_context.clone()).unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); +} + +#[tokio::test] +async fn test_undrop_rejects_table_claimed_for_purge() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + let table_id = 1024; + create_dropped_table(&ddl_context, table_id, Some(0), Some(0)).await; + kv_backend + .put(PutRequest::new().with_key(purging_marker_key(table_id))) + .await + .unwrap(); + + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + let error = procedure.on_prepare().await.unwrap_err(); + + assert_eq!(StatusCode::TableNotFound, error.status_code()); + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); +} + +#[tokio::test] +async fn test_recovered_undrop_rejects_tombstone_deleted_by_purge() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let ddl_context = new_ddl_context(node_manager); + let table_id = 1024; + create_dropped_table(&ddl_context, table_id, Some(1), Some(i64::MAX)).await; + + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + undrop.on_prepare().await.unwrap(); + let open_regions_state = undrop.dump().unwrap(); + undrop.execute(&new_test_procedure_context()).await.unwrap(); + let restore_metadata_state = undrop.dump().unwrap(); + + let mut purge = PurgeDroppedTableProcedure::new( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut purge).await; + + for persisted in [open_regions_state, restore_metadata_state] { + let mut recovered = + UndropTableProcedure::from_json(&persisted, ddl_context.clone()).unwrap(); + let error = recovered + .execute(&new_test_procedure_context()) + .await + .unwrap_err(); + assert_eq!(StatusCode::TableNotFound, error.status_code()); + } +} + +#[tokio::test] +async fn test_recovered_undrop_rejects_replacement_generation() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + let table_id = 1024; + create_dropped_table(&ddl_context, table_id, Some(1), Some(i64::MAX)).await; + kv_backend + .put( + PutRequest::new() + .with_key(drop_generation_marker_key(table_id)) + .with_value("generation-a"), + ) + .await + .unwrap(); + + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + undrop.on_prepare().await.unwrap(); + let persisted = undrop.dump().unwrap(); + kv_backend + .put( + PutRequest::new() + .with_key(drop_generation_marker_key(table_id)) + .with_value("generation-b"), + ) + .await + .unwrap(); + + let mut recovered = UndropTableProcedure::from_json(&persisted, ddl_context).unwrap(); + let error = recovered + .execute(&new_test_procedure_context()) + .await + .unwrap_err(); + assert_eq!(StatusCode::TableNotFound, error.status_code()); +} + +#[tokio::test] +async fn test_legacy_unmarked_tombstone_can_be_undropped_and_purged() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let ddl_context = new_ddl_context(node_manager); + let table_id = 1024; + let table_name = test_create_table_task("foo", table_id).table_name(); + let table_route_value = TableRouteValue::physical(vec![]); + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task("foo", table_id).table_info, + table_route_value.clone(), + HashMap::new(), + ) + .await + .unwrap(); + ddl_context + .table_metadata_manager + .delete_table_metadata( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + None, + ) + .await + .unwrap(); + assert_eq!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .unwrap() + .dropped_at, + None + ); + let dropped_tables = ddl_context + .table_metadata_manager + .list_dropped_tables() + .await + .unwrap(); + assert_eq!(dropped_tables.len(), 1); + assert_eq!(dropped_tables[0].dropped_at, None); + + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + execute_procedure_until_done(&mut undrop).await; + + ddl_context + .table_metadata_manager + .delete_table_metadata( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + None, + ) + .await + .unwrap(); + let mut purge = PurgeDroppedTableProcedure::new( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut purge).await; + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_none() + ); +} + #[tokio::test] async fn test_purge_dropped_table_by_id_selects_tombstone_when_live_table_exists() { let (tx, mut rx) = mpsc::channel(8); @@ -1256,7 +2251,8 @@ async fn test_purge_dropped_table_by_id_selects_tombstone_when_live_table_exists async fn test_on_rollback() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); let kv_backend = Arc::new(MemoryKvBackend::new()); - let ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + let mut ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + ddl_context.soft_drop_enabled = true; // Prepares physical table metadata. let mut create_physical_table_task = test_create_physical_table_task("phy_table"); let TableMetadata { @@ -1297,8 +2293,36 @@ async fn test_on_rollback() { procedure.on_prepare().await.unwrap(); assert!(procedure.rollback_supported()); procedure.on_delete_metadata().await.unwrap(); + assert!( + kv_backend + .get(dropped_at_marker_key(physical_table_id).as_bytes()) + .await + .unwrap() + .is_some() + ); + assert!( + kv_backend + .get(retention_expires_at_marker_key(physical_table_id).as_bytes()) + .await + .unwrap() + .is_some() + ); assert!(procedure.rollback_supported()); procedure.rollback(&ctx).await.unwrap(); + assert!( + kv_backend + .get(dropped_at_marker_key(physical_table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); + assert!( + kv_backend + .get(retention_expires_at_marker_key(physical_table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); // Rollback again assert!(procedure.rollback_supported()); procedure.rollback(&ctx).await.unwrap(); @@ -1307,6 +2331,7 @@ async fn test_on_rollback() { } // Drops the logical table + ddl_context.soft_drop_enabled = false; let task = new_drop_table_task("foo", table_ids[0], false); let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); @@ -1347,6 +2372,7 @@ async fn create_metric_logical_table_tombstone( &task.table_name(), &TableRouteValue::logical(physical_table_id), &HashMap::new(), + None, ) .await .unwrap(); diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs index 9bd3b24822c..872a11ce65f 100644 --- a/src/common/meta/src/ddl/undrop_table.rs +++ b/src/common/meta/src/ddl/undrop_table.rs @@ -77,6 +77,7 @@ impl UndropTableProcedure { } pub(crate) async fn on_prepare(&mut self) -> Result { + self.ensure_not_purging().await?; let dropped_table = self .context .table_metadata_manager @@ -100,6 +101,10 @@ impl UndropTableProcedure { self.data.table_name = Some(dropped_table.table_name.clone()); self.data.table_route_value = Some(dropped_table.table_route_value.clone()); self.data.region_wal_options = dropped_table.region_wal_options; + self.data.dropped_at = dropped_table.dropped_at; + self.data.retention_expires_at = dropped_table.retention_expires_at; + self.data.drop_generation = dropped_table.drop_generation; + self.data.tombstone_identity_loaded = true; ensure!( !is_metric_engine_logical_table( &dropped_table.table_info_value.table_info, @@ -115,20 +120,36 @@ impl UndropTableProcedure { } async fn on_restore_metadata(&mut self) -> Result { + if !self.data.tombstone_identity_loaded + && self + .context + .table_metadata_manager + .get_dropped_table_by_id(self.data.task.table_id) + .await? + .is_some() + && !self.live_table_name_matches().await? + { + self.data.state = UndropTableState::Prepare; + return Ok(Status::executing(true)); + } self.ensure_table_name_loaded().await?; let table_name = self.data.table_name().clone(); let table_route_value = self.data.table_route_value(); - if let Err(err) = self - .context - .table_metadata_manager - .restore_table_metadata( - self.data.task.table_id, - &table_name, - table_route_value, - &self.data.region_wal_options, - ) - .await - { + let restore_result = async { + self.ensure_tombstone_available_for_restore().await?; + self.context + .table_metadata_manager + .restore_table_metadata( + self.data.task.table_id, + &table_name, + table_route_value, + &self.data.region_wal_options, + ) + .await?; + self.ensure_live_table_metadata_restored().await + } + .await; + if let Err(err) = restore_result { let should_cleanup_opened_regions = !err.is_retry_later(); let err = self.map_restore_metadata_error(err); if should_cleanup_opened_regions @@ -163,6 +184,90 @@ impl UndropTableProcedure { Ok(()) } + async fn ensure_not_purging(&self) -> Result<()> { + ensure!( + !self + .context + .table_metadata_manager + .is_dropped_table_purging(self.data.task.table_id) + .await?, + error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + } + ); + Ok(()) + } + + async fn ensure_tombstone_available(&self) -> Result<()> { + self.ensure_not_purging().await?; + let dropped_table = self + .context + .table_metadata_manager + .get_dropped_table_by_id(self.data.task.table_id) + .await? + .with_context(|| error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + })?; + ensure!( + dropped_table.table_name == *self.data.table_name() + && dropped_table.dropped_at == self.data.dropped_at + && dropped_table.retention_expires_at == self.data.retention_expires_at + && dropped_table.drop_generation == self.data.drop_generation, + error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + } + ); + Ok(()) + } + + async fn ensure_tombstone_available_for_restore(&self) -> Result<()> { + if self + .context + .table_metadata_manager + .get_dropped_table_by_id(self.data.task.table_id) + .await? + .is_some() + { + if !self.data.tombstone_identity_loaded && self.live_table_name_matches().await? { + return Ok(()); + } + return self.ensure_tombstone_available().await; + } + + ensure!( + self.live_table_name_matches().await?, + error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + } + ); + Ok(()) + } + + async fn live_table_name_matches(&self) -> Result { + let live_table = self + .context + .table_metadata_manager + .table_name_manager() + .get(TableNameKey::from(self.data.table_name())) + .await?; + Ok(live_table.is_some_and(|table| table.table_id() == self.data.task.table_id)) + } + + async fn ensure_live_table_metadata_restored(&self) -> Result<()> { + let (table_info, table_route) = self + .context + .table_metadata_manager + .get_full_table_info(self.data.task.table_id) + .await?; + ensure!( + table_info.is_some() && table_route.is_some(), + error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + } + ); + Ok(()) + } + fn map_restore_metadata_error(&self, err: error::Error) -> error::Error { match err { error::Error::TombstoneTargetAlreadyExists { .. } => { @@ -211,6 +316,7 @@ impl UndropTableProcedure { } async fn on_open_regions(&mut self) -> Result { + self.ensure_tombstone_available().await?; self.ensure_live_table_not_exists().await?; let TableRouteValue::Physical(route) = self.data.table_route_value() else { self.data.state = UndropTableState::RestoreMetadata; @@ -274,6 +380,10 @@ impl Procedure for UndropTableProcedure { } fn recover(&mut self) -> ProcedureResult<()> { + if self.data.state == UndropTableState::OpenRegions && !self.data.tombstone_identity_loaded + { + self.data.state = UndropTableState::Prepare; + } Ok(()) } @@ -402,6 +512,14 @@ pub struct UndropTableData { table_route_value: Option, #[serde(default)] region_wal_options: HashMap, + #[serde(default)] + dropped_at: Option, + #[serde(default)] + retention_expires_at: Option, + #[serde(default)] + drop_generation: Option, + #[serde(default)] + tombstone_identity_loaded: bool, } impl UndropTableData { @@ -413,6 +531,10 @@ impl UndropTableData { table_info: None, table_route_value: None, region_wal_options: HashMap::new(), + dropped_at: None, + retention_expires_at: None, + drop_generation: None, + tombstone_identity_loaded: false, } } @@ -489,6 +611,27 @@ mod tests { assert_eq!(StatusCode::TableAlreadyExists, err.status_code()); } + #[test] + fn test_recovered_legacy_snapshot_reloads_tombstone_identity() { + let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let mut procedure = + UndropTableProcedure::new(UndropTableTask { table_id: 42 }, context.clone()); + procedure.data.state = UndropTableState::OpenRegions; + let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + let data = data.as_object_mut().unwrap(); + data.remove("dropped_at"); + data.remove("retention_expires_at"); + data.remove("drop_generation"); + data.remove("tombstone_identity_loaded"); + + let mut recovered = + UndropTableProcedure::from_json(&serde_json::to_string(data).unwrap(), context) + .unwrap(); + recovered.recover().unwrap(); + + assert_eq!(UndropTableState::Prepare, recovered.data.state); + } + #[tokio::test] async fn test_cleanup_opened_regions_without_table_name_deregisters_detectors() { let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 3f421d3b028..455fc098b86 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -272,6 +272,18 @@ impl DdlManager { .context(RegisterProcedureLoaderSnafu { type_name })?; } + let type_name = PurgeDroppedTableProcedure::EXPIRED_TYPE_NAME; + let context = self.create_context(); + self.procedure_manager + .register_loader( + type_name, + Box::new(move |json: &str| { + PurgeDroppedTableProcedure::from_json(json, context.clone()) + .map(|procedure| Box::new(procedure) as _) + }), + ) + .context(RegisterProcedureLoaderSnafu { type_name })?; + self.repartition_procedure_factory .register_loaders(&self.ddl_context, &self.procedure_manager) .context(RegisterRepartitionProcedureLoaderSnafu)?; @@ -511,6 +523,20 @@ impl DdlManager { self.execute_procedure_and_wait(procedure_with_id).await } + /// Submits and executes a purge task that first rechecks the tombstone deadline. + #[tracing::instrument(skip_all)] + pub async fn submit_expired_purge_dropped_table_task( + &self, + purge_dropped_table_task: PurgeDroppedTableTask, + ) -> Result<(ProcedureId, Option)> { + let context = self.create_context(); + let procedure = + PurgeDroppedTableProcedure::new_if_expired(purge_dropped_table_task, context); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + + self.execute_procedure_and_wait(procedure_with_id).await + } + /// Submits and executes a create database task. #[tracing::instrument(skip_all)] pub async fn submit_create_database( @@ -1323,6 +1349,7 @@ mod tests { leader_region_registry: Arc::new(LeaderRegionRegistry::default()), region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), soft_drop_enabled: false, + soft_drop_retention: None, }, procedure_manager.clone(), Arc::new(DummyRepartitionProcedureFactory), @@ -1376,6 +1403,7 @@ mod tests { leader_region_registry: Arc::new(LeaderRegionRegistry::default()), region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), soft_drop_enabled: true, + soft_drop_retention: Some(std::time::Duration::from_secs(1)), }, procedure_manager, Arc::new(DummyRepartitionProcedureFactory), diff --git a/src/common/meta/src/key.rs b/src/common/meta/src/key.rs index 56fa35b9431..71e9ce4f4b0 100644 --- a/src/common/meta/src/key.rs +++ b/src/common/meta/src/key.rs @@ -164,8 +164,9 @@ use crate::key::topic_region::TopicRegionValue; use crate::key::txn_helper::TxnOpGetResponseSet; use crate::kv_backend::KvBackendRef; use crate::kv_backend::txn::{Txn, TxnOp}; +use crate::rpc::KeyValue; use crate::rpc::router::{LeaderState, RegionRoute, region_distribution}; -use crate::rpc::store::BatchDeleteRequest; +use crate::rpc::store::{BatchDeleteRequest, PutRequest}; use crate::state_store::PoisonValue; use crate::wal_provider::RegionWalOptions; @@ -440,6 +441,12 @@ pub struct DroppedTableName { pub table_id: TableId, /// Original fully qualified table name. pub table_name: TableName, + /// Unix timestamp in milliseconds when this table was soft-dropped. + pub dropped_at: Option, + /// Fixed automatic-GC deadline in Unix milliseconds. + pub retention_expires_at: Option, + /// Unique identity of this soft-drop generation. + pub drop_generation: Option, } #[derive(Debug, Clone)] @@ -454,6 +461,40 @@ pub struct DroppedTableMetadata { pub table_route_value: TableRouteValue, /// Per-region WAL options recovered from tombstoned datanode metadata. pub region_wal_options: HashMap, + /// Unix timestamp in milliseconds when this table was soft-dropped. + pub dropped_at: Option, + /// Fixed automatic-GC deadline in Unix milliseconds. + pub retention_expires_at: Option, + /// Unique identity of this soft-drop generation. + pub drop_generation: Option, +} + +/// Lifecycle markers stored with a logically deleted table. +pub struct DroppedTableLifecycle<'a> { + pub dropped_at: Option, + pub retention_expires_at: Option, + pub drop_generation: Option<&'a str>, +} + +const DROPPED_AT_KEY_PREFIX: &str = "__dropped_at"; +const RETENTION_EXPIRES_AT_KEY_PREFIX: &str = "__retention_expires_at"; +const DROP_GENERATION_KEY_PREFIX: &str = "__drop_generation"; +const PURGING_KEY_PREFIX: &str = "__purging"; + +fn dropped_at_key(table_id: TableId) -> Vec { + format!("{DROPPED_AT_KEY_PREFIX}/{table_id}").into_bytes() +} + +fn retention_expires_at_key(table_id: TableId) -> Vec { + format!("{RETENTION_EXPIRES_AT_KEY_PREFIX}/{table_id}").into_bytes() +} + +fn drop_generation_key(table_id: TableId) -> Vec { + format!("{DROP_GENERATION_KEY_PREFIX}/{table_id}").into_bytes() +} + +fn purging_key(table_id: TableId) -> Vec { + format!("{PURGING_KEY_PREFIX}/{table_id}").into_bytes() } impl Deref for DeserializedValueWithBytes { @@ -998,10 +1039,78 @@ impl TableMetadataManager { table_name: &TableName, table_route_value: &TableRouteValue, region_wal_options: &HashMap, + dropped_at: Option, + ) -> Result<()> { + self.delete_table_metadata_with_retention( + table_id, + table_name, + table_route_value, + region_wal_options, + dropped_at, + None, + ) + .await + } + + /// Deletes metadata logically and stores fixed soft-drop lifecycle markers. + pub async fn delete_table_metadata_with_retention( + &self, + table_id: TableId, + table_name: &TableName, + table_route_value: &TableRouteValue, + region_wal_options: &HashMap, + dropped_at: Option, + retention_expires_at: Option, + ) -> Result<()> { + self.delete_table_metadata_with_retention_and_generation( + table_id, + table_name, + table_route_value, + region_wal_options, + DroppedTableLifecycle { + dropped_at, + retention_expires_at, + drop_generation: None, + }, + ) + .await + } + + /// Deletes metadata logically and stores fixed soft-drop lifecycle markers. + pub async fn delete_table_metadata_with_retention_and_generation( + &self, + table_id: TableId, + table_name: &TableName, + table_route_value: &TableRouteValue, + region_wal_options: &HashMap, + lifecycle: DroppedTableLifecycle<'_>, ) -> Result<()> { let keys = self.table_metadata_keys(table_id, table_name, table_route_value, region_wal_options)?; - self.tombstone_manager.create(keys).await.map(|_| ()) + if let Some(dropped_at) = lifecycle.dropped_at { + let mut markers = vec![( + dropped_at_key(table_id), + dropped_at.to_string().into_bytes(), + )]; + if let Some(retention_expires_at) = lifecycle.retention_expires_at { + markers.push(( + retention_expires_at_key(table_id), + retention_expires_at.to_string().into_bytes(), + )); + } + if let Some(drop_generation) = lifecycle.drop_generation { + markers.push(( + drop_generation_key(table_id), + drop_generation.as_bytes().to_vec(), + )); + } + self.tombstone_manager + .create_with_markers(keys, markers) + .await + .map(|_| ()) + } else { + self.tombstone_manager.create(keys).await.map(|_| ()) + } } /// Lists dropped tables from tombstoned table-name entries. @@ -1016,8 +1125,39 @@ impl TableMetadataManager { dropped_tables.push(DroppedTableName { table_id, table_name, + dropped_at: None, + retention_expires_at: None, + drop_generation: None, }); } + if dropped_tables.is_empty() { + return Ok(dropped_tables); + } + + let marker_keys = dropped_tables + .iter() + .flat_map(|table| { + vec![ + dropped_at_key(table.table_id), + retention_expires_at_key(table.table_id), + drop_generation_key(table.table_id), + ] + }) + .collect::>(); + let marker_values = self.tombstone_manager.batch_get(&marker_keys).await?; + for table in &mut dropped_tables { + table.dropped_at = Self::parse_dropped_at( + table.table_id, + marker_values.get(&dropped_at_key(table.table_id)), + )?; + table.retention_expires_at = Self::parse_retention_expires_at( + table.table_id, + marker_values.get(&retention_expires_at_key(table.table_id)), + )?; + table.drop_generation = Self::parse_drop_generation( + marker_values.get(&drop_generation_key(table.table_id)), + ); + } Ok(dropped_tables) } @@ -1049,6 +1189,38 @@ impl TableMetadataManager { self.get_dropped_table_metadata(table_id, None).await } + /// Returns whether an automatic purge has durably claimed the dropped table. + pub async fn is_dropped_table_purging(&self, table_id: TableId) -> Result { + self.dropped_table_purge_claim(table_id) + .await + .map(|claim| claim.is_some()) + } + + /// Returns the soft-drop generation claimed by an automatic purge. + pub async fn dropped_table_purge_claim(&self, table_id: TableId) -> Result> { + self.tombstone_manager + .get(&purging_key(table_id)) + .await + .map(|marker| marker.map(|marker| String::from_utf8_lossy(&marker.value).into_owned())) + } + + /// Durably claims a dropped table before automatic purge starts cleaning its regions. + /// The caller MUST hold the table lock and revalidate the tombstone before calling this. + pub async fn mark_dropped_table_purging( + &self, + table_id: TableId, + drop_generation: Option<&str>, + ) -> Result<()> { + self.kv_backend + .put( + PutRequest::new() + .with_key(self.tombstone_manager.to_tombstone(&purging_key(table_id))) + .with_value(drop_generation.unwrap_or_default()), + ) + .await?; + Ok(()) + } + /// Deletes metadata tombstone for table **permanently**. /// The caller MUST ensure it has the exclusive access to `TableNameKey`. pub async fn delete_table_metadata_tombstone( @@ -1061,7 +1233,15 @@ impl TableMetadataManager { let table_metadata_keys = self.table_metadata_keys(table_id, table_name, table_route_value, region_wal_options)?; self.tombstone_manager - .delete(table_metadata_keys) + .delete_with_markers( + table_metadata_keys, + vec![ + dropped_at_key(table_id), + retention_expires_at_key(table_id), + drop_generation_key(table_id), + purging_key(table_id), + ], + ) .await .map(|_| ()) } @@ -1077,7 +1257,17 @@ impl TableMetadataManager { ) -> Result<()> { let keys = self.table_metadata_keys(table_id, table_name, table_route_value, region_wal_options)?; - self.tombstone_manager.restore(keys).await.map(|_| ()) + self.tombstone_manager + .restore_with_markers( + keys, + vec![ + dropped_at_key(table_id), + retention_expires_at_key(table_id), + drop_generation_key(table_id), + ], + ) + .await + .map(|_| ()) } /// Deletes metadata for table **permanently**. @@ -1138,6 +1328,23 @@ impl TableMetadataManager { let region_wal_options = self .dropped_region_wal_options(table_id, &table_route_value) .await?; + let dropped_at_key = dropped_at_key(table_id); + let retention_expires_at_key = retention_expires_at_key(table_id); + let drop_generation_key = drop_generation_key(table_id); + let marker_values = self + .tombstone_manager + .batch_get(&[ + dropped_at_key.clone(), + retention_expires_at_key.clone(), + drop_generation_key.clone(), + ]) + .await?; + let dropped_at = Self::parse_dropped_at(table_id, marker_values.get(&dropped_at_key))?; + let retention_expires_at = Self::parse_retention_expires_at( + table_id, + marker_values.get(&retention_expires_at_key), + )?; + let drop_generation = Self::parse_drop_generation(marker_values.get(&drop_generation_key)); Ok(Some(DroppedTableMetadata { table_id, @@ -1145,9 +1352,41 @@ impl TableMetadataManager { table_info_value, table_route_value, region_wal_options, + dropped_at, + retention_expires_at, + drop_generation, })) } + fn parse_dropped_at(table_id: TableId, kv: Option<&KeyValue>) -> Result> { + Self::parse_timestamp_marker(table_id, "dropped timestamp", kv) + } + + fn parse_retention_expires_at(table_id: TableId, kv: Option<&KeyValue>) -> Result> { + Self::parse_timestamp_marker(table_id, "retention deadline", kv) + } + + fn parse_drop_generation(kv: Option<&KeyValue>) -> Option { + kv.map(|kv| String::from_utf8_lossy(&kv.value).into_owned()) + } + + fn parse_timestamp_marker( + table_id: TableId, + description: &str, + kv: Option<&KeyValue>, + ) -> Result> { + let Some(kv) = kv else { + return Ok(None); + }; + let value = String::from_utf8_lossy(&kv.value); + value.parse().map(Some).map_err(|err| { + error::UnexpectedSnafu { + err_msg: format!("Invalid {description} '{value}' for table id {table_id}: {err}"), + } + .build() + }) + } + /// Rebuilds region WAL options from tombstoned datanode-table entries. async fn dropped_region_wal_options( &self, @@ -1633,6 +1872,7 @@ impl_optional_metadata_value! { mod tests { use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use bytes::Bytes; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; @@ -1662,9 +1902,10 @@ mod tests { use crate::kv_backend::KvBackend; use crate::kv_backend::memory::MemoryKvBackend; use crate::kv_backend::read_only::ReadOnlyKvBackend; + use crate::kv_backend::test_util::MockKvBackend; use crate::peer::Peer; use crate::rpc::router::{LeaderState, Region, RegionRoute, region_distribution}; - use crate::rpc::store::{PutRequest, RangeRequest}; + use crate::rpc::store::{BatchGetResponse, PutRequest, RangeRequest, RangeResponse}; use crate::wal_provider::{RegionWalOptions, WalProvider}; #[test] @@ -1837,6 +2078,7 @@ mod tests { &table_name, &table_route_value, ®ion_wal_options, + None, ) .await .unwrap(); @@ -2188,6 +2430,7 @@ mod tests { &table_name, table_route_value, ®ion_wal_options, + None, ) .await .unwrap(); @@ -2198,6 +2441,7 @@ mod tests { &table_name, table_route_value, ®ion_wal_options, + None, ) .await .unwrap(); @@ -2927,7 +3171,7 @@ mod tests { let table_name = TableName::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, table_name); let table_route_value = TableRouteValue::physical(region_routes.clone()); table_metadata_manager - .delete_table_metadata(table_id, &table_name, &table_route_value, &options) + .delete_table_metadata(table_id, &table_name, &table_route_value, &options, None) .await .unwrap(); table_metadata_manager @@ -2996,6 +3240,92 @@ mod tests { ); } + #[tokio::test] + async fn test_list_dropped_tables_batches_timestamp_lookup() { + let mem_kv = Arc::new(MemoryKvBackend::default()); + let table_metadata_manager = TableMetadataManager::new(mem_kv.clone()); + for (table_id, table_name, dropped_at) in + [(1025, "legacy", None), (1026, "marked", Some(1234))] + { + let task = test_create_table_task(table_name, table_id); + let table_name = task.table_name(); + let table_route = TableRouteValue::physical(vec![]); + table_metadata_manager + .create_table_metadata(task.table_info, table_route.clone(), HashMap::new()) + .await + .unwrap(); + table_metadata_manager + .delete_table_metadata( + table_id, + &table_name, + &table_route, + &HashMap::new(), + dropped_at, + ) + .await + .unwrap(); + } + let all_tombstones = mem_kv + .range(RangeRequest::new().with_prefix("__tombstone/")) + .await + .unwrap() + .kvs; + let range_calls = Arc::new(AtomicUsize::new(0)); + let batch_get_calls = Arc::new(AtomicUsize::new(0)); + let mock = MockKvBackend { + range_fn: Some({ + let all_tombstones = all_tombstones.clone(); + let range_calls = range_calls.clone(); + Arc::new(move |req| { + range_calls.fetch_add(1, Ordering::Relaxed); + let kvs = all_tombstones + .iter() + .filter(|kv| { + if req.range_end.is_empty() { + kv.key == req.key + } else { + kv.key >= req.key && kv.key < req.range_end + } + }) + .cloned() + .collect(); + Ok(RangeResponse { kvs, more: false }) + }) + }), + batch_get_fn: Some({ + let batch_get_calls = batch_get_calls.clone(); + Arc::new(move |req| { + batch_get_calls.fetch_add(1, Ordering::Relaxed); + let kvs = all_tombstones + .iter() + .filter(|kv| req.keys.contains(&kv.key)) + .cloned() + .collect(); + Ok(BatchGetResponse { kvs }) + }) + }), + put_fn: None, + batch_put_fn: None, + delete_range_fn: None, + batch_delete_fn: None, + txn: None, + max_txn_ops: None, + }; + let table_metadata_manager = TableMetadataManager::new(Arc::new(mock)); + + let dropped_tables = table_metadata_manager.list_dropped_tables().await.unwrap(); + + assert_eq!( + dropped_tables + .iter() + .map(|table| (table.table_id, table.dropped_at)) + .collect::>(), + vec![(1025, None), (1026, Some(1234))] + ); + assert_eq!(range_calls.load(Ordering::Relaxed), 1); + assert_eq!(batch_get_calls.load(Ordering::Relaxed), 1); + } + #[tokio::test] async fn test_dropped_table_lookup_survives_live_name_recreation() { let dropped_table_id = 1025; diff --git a/src/common/meta/src/key/tombstone.rs b/src/common/meta/src/key/tombstone.rs index 01fa4405776..284737dad0b 100644 --- a/src/common/meta/src/key/tombstone.rs +++ b/src/common/meta/src/key/tombstone.rs @@ -190,6 +190,7 @@ impl TombstoneManager { keys: &[Vec], dest_keys: &[Vec], require_dest_not_exists: bool, + extra_ops: Vec, ) -> Result { ensure!( keys.len() == dest_keys.len(), @@ -213,6 +214,9 @@ impl TombstoneManager { .into_iter() .map(|kv| (kv.key, kv.value)) .collect::>(); + if results.is_empty() && extra_ops.is_empty() { + return Ok(0); + } const MAX_RETRIES: usize = 8; for _ in 0..MAX_RETRIES { @@ -228,6 +232,10 @@ impl TombstoneManager { (txn, key.clone()) }) .unzip(); + let mut txns = txns; + if !extra_ops.is_empty() { + txns.push(Txn::new().and_then(extra_ops.clone())); + } let mut resp = self.kv_backend.txn(Txn::merge_all(txns)).await?; if resp.succeeded { return Ok(keys.len()); @@ -272,14 +280,30 @@ impl TombstoneManager { self.kv_backend.max_txn_ops() } + async fn execute_extra_ops(&self, extra_ops: Vec) -> Result<()> { + let max_txn_ops = self.max_txn_ops(); + ensure!( + max_txn_ops > 0, + error::UnexpectedSnafu { + err_msg: "max_txn_ops must be greater than 0".to_string(), + } + ); + for chunk in extra_ops.chunks(max_txn_ops) { + self.move_values_inner(&[], &[], false, chunk.to_vec()) + .await?; + } + Ok(()) + } + /// Moves values to `dest_key`. /// /// Returns the number of keys that were moved. - async fn move_values( + async fn move_values_with_extra( &self, keys: Vec>, dest_keys: Vec>, require_dest_not_exists: bool, + mut extra_ops: Vec, ) -> Result { ensure!( keys.len() == dest_keys.len(), @@ -292,6 +316,9 @@ impl TombstoneManager { } ); if keys.is_empty() { + if !extra_ops.is_empty() { + self.execute_extra_ops(extra_ops).await?; + } return Ok(0); } let txn_ops_per_key = if require_dest_not_exists { @@ -308,8 +335,15 @@ impl TombstoneManager { ) } ); - let chunk_size = max_txn_ops / txn_ops_per_key; - if keys.len() > chunk_size { + let merge_extra_ops = + !extra_ops.is_empty() && max_txn_ops.saturating_sub(txn_ops_per_key) >= extra_ops.len(); + let reserved_ops = if merge_extra_ops { extra_ops.len() } else { 0 }; + let chunk_size = (max_txn_ops - reserved_ops) / txn_ops_per_key; + let total_ops = keys + .len() + .saturating_mul(txn_ops_per_key) + .saturating_add(extra_ops.len()); + if keys.len() > chunk_size || total_ops > max_txn_ops { debug!( "Moving values with multiple chunks, keys len: {}, chunk_size: {}", keys.len(), @@ -318,18 +352,39 @@ impl TombstoneManager { let mut moved_keys = 0; let keys_chunks = keys.chunks(chunk_size).collect::>(); let dest_keys_chunks = dest_keys.chunks(chunk_size).collect::>(); - for (keys, dest_keys) in keys_chunks.into_iter().zip(dest_keys_chunks) { + let num_chunks = keys_chunks.len(); + for (index, (keys, dest_keys)) in + keys_chunks.into_iter().zip(dest_keys_chunks).enumerate() + { + let chunk_extra_ops = if merge_extra_ops && index + 1 == num_chunks { + std::mem::take(&mut extra_ops) + } else { + vec![] + }; moved_keys += self - .move_values_inner(keys, dest_keys, require_dest_not_exists) + .move_values_inner(keys, dest_keys, require_dest_not_exists, chunk_extra_ops) .await?; } + if !extra_ops.is_empty() { + self.execute_extra_ops(extra_ops).await?; + } Ok(moved_keys) } else { - self.move_values_inner(&keys, &dest_keys, require_dest_not_exists) + self.move_values_inner(&keys, &dest_keys, require_dest_not_exists, extra_ops) .await } } + async fn move_values( + &self, + keys: Vec>, + dest_keys: Vec>, + require_dest_not_exists: bool, + ) -> Result { + self.move_values_with_extra(keys, dest_keys, require_dest_not_exists, vec![]) + .await + } + /// Creates tombstones for keys. /// /// Preforms to: @@ -349,6 +404,39 @@ impl TombstoneManager { self.move_values(keys, dest_keys, false).await } + /// Creates tombstones and stores an additional marker in the tombstone namespace. + pub async fn create_with_marker( + &self, + keys: Vec>, + marker_key: Vec, + marker_value: Vec, + ) -> Result { + self.create_with_markers(keys, vec![(marker_key, marker_value)]) + .await + } + + /// Creates tombstones and stores additional markers in the tombstone namespace. + pub async fn create_with_markers( + &self, + keys: Vec>, + markers: Vec<(Vec, Vec)>, + ) -> Result { + let (keys, dest_keys): (Vec<_>, Vec<_>) = keys + .into_iter() + .map(|key| { + let tombstone_key = self.to_tombstone(&key); + (key, tombstone_key) + }) + .unzip(); + let marker_ops = markers + .into_iter() + .map(|(key, value)| TxnOp::Put(self.to_tombstone(&key), value)) + .collect(); + + self.move_values_with_extra(keys, dest_keys, false, marker_ops) + .await + } + /// Restores tombstones for keys. /// /// Preforms to: @@ -368,6 +456,46 @@ impl TombstoneManager { self.move_values(keys, dest_keys, true).await } + /// Restores tombstones and deletes an associated marker. + pub async fn restore_with_marker( + &self, + keys: Vec>, + marker_key: Vec, + ) -> Result { + let (keys, dest_keys): (Vec<_>, Vec<_>) = keys + .into_iter() + .map(|key| { + let tombstone_key = self.to_tombstone(&key); + (tombstone_key, key) + }) + .unzip(); + let marker_key = self.to_tombstone(&marker_key); + + self.move_values_with_extra(keys, dest_keys, true, vec![TxnOp::Delete(marker_key)]) + .await + } + + /// Restores tombstones and deletes associated markers. + pub async fn restore_with_markers( + &self, + keys: Vec>, + marker_keys: Vec>, + ) -> Result { + let (keys, dest_keys): (Vec<_>, Vec<_>) = keys + .into_iter() + .map(|key| { + let tombstone_key = self.to_tombstone(&key); + (tombstone_key, key) + }) + .unzip(); + let marker_ops = marker_keys + .into_iter() + .map(|key| TxnOp::Delete(self.to_tombstone(&key))) + .collect(); + self.move_values_with_extra(keys, dest_keys, true, marker_ops) + .await + } + /// Deletes tombstones values for the specified `keys`. /// /// Returns the number of keys that were deleted. @@ -385,6 +513,26 @@ impl TombstoneManager { Ok(num_keys) } + + /// Deletes tombstones and an associated marker. + pub async fn delete_with_marker( + &self, + mut keys: Vec>, + marker_key: Vec, + ) -> Result { + keys.push(marker_key); + self.delete(keys).await + } + + /// Deletes tombstones and associated markers. + pub async fn delete_with_markers( + &self, + mut keys: Vec>, + marker_keys: Vec>, + ) -> Result { + keys.extend(marker_keys); + self.delete(keys).await + } } #[cfg(test)] @@ -392,10 +540,12 @@ mod tests { use std::any::Any; use std::collections::HashMap; - use std::sync::Arc; + use std::sync::{Arc, Mutex}; use crate::error::{Error, Result}; - use crate::key::tombstone::TombstoneManager; + use crate::key::tombstone::{ + MOVE_VALUE_TXN_OPS_PER_KEY, RESTORE_VALUE_TXN_OPS_PER_KEY, TombstoneManager, + }; use crate::kv_backend::memory::MemoryKvBackend; use crate::kv_backend::txn::{Txn, TxnRequest, TxnResponse}; use crate::kv_backend::{KvBackend, TxnService}; @@ -408,6 +558,7 @@ mod tests { struct TxnOpLimitKvBackend { inner: Arc>, max_txn_ops: usize, + txn_op_counts: Mutex>, } #[async_trait::async_trait] @@ -426,6 +577,7 @@ mod tests { "txn ops {txn_ops} exceeds limit {}", self.max_txn_ops ); + self.txn_op_counts.lock().unwrap().push(txn_ops); self.inner.txn(txn).await } @@ -746,6 +898,7 @@ mod tests { let kv_backend = Arc::new(TxnOpLimitKvBackend { inner: inner.clone(), max_txn_ops: 6, + txn_op_counts: Mutex::new(Vec::new()), }); let tombstone_manager = TombstoneManager::new(kv_backend); let kvs = HashMap::from([ @@ -778,6 +931,7 @@ mod tests { let kv_backend = Arc::new(TxnOpLimitKvBackend { inner: inner.clone(), max_txn_ops: 5, + txn_op_counts: Mutex::new(Vec::new()), }); let tombstone_manager = TombstoneManager::new(kv_backend); let key = b"foo".to_vec(); @@ -805,6 +959,7 @@ mod tests { let kv_backend = Arc::new(TxnOpLimitKvBackend { inner: inner.clone(), max_txn_ops: 4, + txn_op_counts: Mutex::new(Vec::new()), }); let tombstone_manager = TombstoneManager::new(kv_backend); let kvs = HashMap::from([ @@ -831,6 +986,195 @@ mod tests { assert_eq!(kvs.len(), moved); } + #[tokio::test] + async fn test_create_with_marker_uses_final_transaction_at_exact_limit() { + let inner = Arc::new(MemoryKvBackend::default()); + let kv_backend = Arc::new(TxnOpLimitKvBackend { + inner: inner.clone(), + max_txn_ops: MOVE_VALUE_TXN_OPS_PER_KEY, + txn_op_counts: Mutex::new(Vec::new()), + }); + let tombstone_manager = TombstoneManager::new(kv_backend.clone()); + for (key, value) in [(b"bar".as_slice(), b"baz".as_slice()), (b"foo", b"hi")] { + inner + .put(PutRequest::new().with_key(key).with_value(value)) + .await + .unwrap(); + } + + tombstone_manager + .create_with_marker( + vec![b"bar".to_vec(), b"foo".to_vec()], + b"__dropped_at/42".to_vec(), + b"1234".to_vec(), + ) + .await + .unwrap(); + tombstone_manager + .create_with_marker( + vec![b"bar".to_vec(), b"foo".to_vec()], + b"__dropped_at/42".to_vec(), + b"1234".to_vec(), + ) + .await + .unwrap(); + + assert_eq!( + kv_backend.txn_op_counts.lock().unwrap().as_slice(), + &[4, 4, 1, 1] + ); + assert_eq!( + inner + .get(&tombstone_manager.to_tombstone(b"__dropped_at/42")) + .await + .unwrap() + .unwrap() + .value, + b"1234" + ); + } + + #[tokio::test] + async fn test_create_with_markers_chunks_extra_ops_after_key_move() { + let inner = Arc::new(MemoryKvBackend::default()); + let kv_backend = Arc::new(TxnOpLimitKvBackend { + inner: inner.clone(), + max_txn_ops: MOVE_VALUE_TXN_OPS_PER_KEY + 1, + txn_op_counts: Mutex::new(Vec::new()), + }); + let tombstone_manager = TombstoneManager::new(kv_backend.clone()); + inner + .put(PutRequest::new().with_key(b"foo").with_value(b"bar")) + .await + .unwrap(); + + tombstone_manager + .create_with_markers( + vec![b"foo".to_vec()], + (0..6) + .map(|index| (format!("marker-{index}").into_bytes(), vec![index])) + .collect(), + ) + .await + .unwrap(); + + assert_eq!( + kv_backend.txn_op_counts.lock().unwrap().as_slice(), + &[MOVE_VALUE_TXN_OPS_PER_KEY, 5, 1] + ); + } + + #[tokio::test] + async fn test_create_with_markers_chunks_marker_only_transactions() { + let inner = Arc::new(MemoryKvBackend::default()); + let kv_backend = Arc::new(TxnOpLimitKvBackend { + inner, + max_txn_ops: 2, + txn_op_counts: Mutex::new(Vec::new()), + }); + let tombstone_manager = TombstoneManager::new(kv_backend.clone()); + let markers = (0..5) + .map(|index| (format!("marker-{index}").into_bytes(), vec![index])) + .collect(); + + tombstone_manager + .create_with_markers(vec![], markers) + .await + .unwrap(); + + assert_eq!( + kv_backend.txn_op_counts.lock().unwrap().as_slice(), + &[2, 2, 1] + ); + } + + #[tokio::test] + async fn test_restore_with_marker_uses_final_transaction_at_exact_limit() { + let inner = Arc::new(MemoryKvBackend::default()); + let kv_backend = Arc::new(TxnOpLimitKvBackend { + inner: inner.clone(), + max_txn_ops: RESTORE_VALUE_TXN_OPS_PER_KEY, + txn_op_counts: Mutex::new(Vec::new()), + }); + let tombstone_manager = TombstoneManager::new(kv_backend.clone()); + for (key, value) in [(b"bar".as_slice(), b"baz".as_slice()), (b"foo", b"hi")] { + inner + .put( + PutRequest::new() + .with_key(tombstone_manager.to_tombstone(key)) + .with_value(value), + ) + .await + .unwrap(); + } + inner + .put( + PutRequest::new() + .with_key(tombstone_manager.to_tombstone(b"__dropped_at/42")) + .with_value(b"1234"), + ) + .await + .unwrap(); + + tombstone_manager + .restore_with_marker( + vec![b"bar".to_vec(), b"foo".to_vec()], + b"__dropped_at/42".to_vec(), + ) + .await + .unwrap(); + tombstone_manager + .restore_with_marker( + vec![b"bar".to_vec(), b"foo".to_vec()], + b"__dropped_at/42".to_vec(), + ) + .await + .unwrap(); + + assert_eq!( + kv_backend.txn_op_counts.lock().unwrap().as_slice(), + &[6, 6, 1, 1] + ); + assert!( + inner + .get(&tombstone_manager.to_tombstone(b"__dropped_at/42")) + .await + .unwrap() + .is_none() + ); + } + + #[tokio::test] + async fn test_delete_with_marker_removes_marker_with_limited_backend() { + let inner = Arc::new(MemoryKvBackend::default()); + let kv_backend = Arc::new(TxnOpLimitKvBackend { + inner: inner.clone(), + max_txn_ops: MOVE_VALUE_TXN_OPS_PER_KEY, + txn_op_counts: Mutex::new(Vec::new()), + }); + let tombstone_manager = TombstoneManager::new(kv_backend); + for key in [b"bar".as_slice(), b"foo", b"__dropped_at/42"] { + inner + .put( + PutRequest::new() + .with_key(tombstone_manager.to_tombstone(key)) + .with_value(b"value"), + ) + .await + .unwrap(); + } + + tombstone_manager + .delete_with_marker( + vec![b"bar".to_vec(), b"foo".to_vec()], + b"__dropped_at/42".to_vec(), + ) + .await + .unwrap(); + + assert!(inner.is_empty()); + } + #[tokio::test] async fn test_move_values_with_non_exists_values() { let kv_backend = Arc::new(MemoryKvBackend::default()); diff --git a/src/common/meta/src/test_util.rs b/src/common/meta/src/test_util.rs index 2300679b337..24c4b7b759c 100644 --- a/src/common/meta/src/test_util.rs +++ b/src/common/meta/src/test_util.rs @@ -207,6 +207,7 @@ pub fn new_ddl_context_with_kv_backend( flow_metadata_manager, region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), soft_drop_enabled: false, + soft_drop_retention: Some(std::time::Duration::from_secs(7 * 24 * 60 * 60)), } } diff --git a/src/meta-srv/src/gc.rs b/src/meta-srv/src/gc.rs index f9b3869bc63..f64f52eaeb5 100644 --- a/src/meta-srv/src/gc.rs +++ b/src/meta-srv/src/gc.rs @@ -29,6 +29,8 @@ mod scheduler; mod tracker; mod util; +pub(crate) use ctx::DefaultGcSchedulerCtx; +pub(crate) use options::EXPERIMENTAL_SOFT_DROP_ENABLED; pub use options::GcSchedulerOptions; pub use procedure::BatchGcProcedure; pub use scheduler::{Event, GcJobReport, GcScheduler, GcTickerRef}; diff --git a/src/meta-srv/src/gc/ctx.rs b/src/meta-srv/src/gc/ctx.rs index 6f00ff50a45..6b838e00478 100644 --- a/src/meta-srv/src/gc/ctx.rs +++ b/src/meta-srv/src/gc/ctx.rs @@ -12,13 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; use std::time::Duration; use common_meta::datanode::RegionStat; -use common_meta::key::TableMetadataManagerRef; +use common_meta::ddl_manager::DdlManagerRef; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; +use common_meta::key::{DroppedTableName, TableMetadataManagerRef}; +use common_meta::rpc::ddl::PurgeDroppedTableTask; use common_procedure::{ProcedureManagerRef, ProcedureWithId, watcher}; use common_telemetry::debug; use snafu::{OptionExt as _, ResultExt as _}; @@ -29,6 +33,7 @@ use crate::cluster::MetaPeerClientRef; use crate::error::{self, Result, TableMetadataManagerSnafu}; use crate::gc::Region2Peers; use crate::gc::procedure::BatchGcProcedure; +use crate::metrics::METRIC_META_GC_SOFT_DROP_PURGES_TOTAL; use crate::service::mailbox::MailboxRef; #[async_trait::async_trait] @@ -54,6 +59,77 @@ pub(crate) trait SchedulerCtx: Send + Sync { timeout: Duration, region_routes_override: Region2Peers, ) -> Result; + + async fn list_dropped_tables(&self) -> Result>; + + async fn purge_dropped_table(&self, table_id: TableId) -> Result<()>; + + fn try_reserve_purge( + &self, + table_id: TableId, + max_in_flight: usize, + ) -> Option; + + fn next_purge_scan_start(&self, _table_count: usize) -> usize { + 0 + } +} + +pub(crate) enum PurgeOutcome { + Succeeded, + Failed, +} + +pub(crate) struct PurgeReservation { + table_id: TableId, + in_flight: Arc>>, + outcome_recorded: bool, +} + +impl PurgeReservation { + pub(crate) fn try_new( + in_flight: Arc>>, + table_id: TableId, + max_in_flight: usize, + ) -> Option { + let mut tables = in_flight + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if tables.len() >= max_in_flight || !tables.insert(table_id) { + return None; + } + drop(tables); + Some(Self { + table_id, + in_flight, + outcome_recorded: false, + }) + } + + pub(crate) fn record_outcome(mut self, outcome: PurgeOutcome) { + let status = match outcome { + PurgeOutcome::Succeeded => "succeeded", + PurgeOutcome::Failed => "failed", + }; + METRIC_META_GC_SOFT_DROP_PURGES_TOTAL + .with_label_values(&[status]) + .inc(); + self.outcome_recorded = true; + } +} + +impl Drop for PurgeReservation { + fn drop(&mut self) { + if !self.outcome_recorded { + METRIC_META_GC_SOFT_DROP_PURGES_TOTAL + .with_label_values(&["cancelled"]) + .inc(); + } + self.in_flight + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .remove(&self.table_id); + } } pub(crate) struct DefaultGcSchedulerCtx { @@ -61,6 +137,12 @@ pub(crate) struct DefaultGcSchedulerCtx { pub(crate) table_metadata_manager: TableMetadataManagerRef, /// Procedure manager. pub(crate) procedure_manager: ProcedureManagerRef, + /// DDL manager used to submit the existing purge procedure. + pub(crate) ddl_manager: DdlManagerRef, + /// Process-local reservations for purge procedures submitted by this scheduler. + /// Procedure recovery after a metasrv restart may outlive this set. + in_flight_purges: Arc>>, + purge_scan_cursor: AtomicUsize, /// For getting `RegionStats`. pub(crate) meta_peer_client: MetaPeerClientRef, /// The mailbox to send messages. @@ -73,6 +155,7 @@ impl DefaultGcSchedulerCtx { pub fn try_new( table_metadata_manager: TableMetadataManagerRef, procedure_manager: ProcedureManagerRef, + ddl_manager: DdlManagerRef, meta_peer_client: MetaPeerClientRef, mailbox: MailboxRef, server_addr: String, @@ -80,6 +163,9 @@ impl DefaultGcSchedulerCtx { Ok(Self { table_metadata_manager, procedure_manager, + ddl_manager, + in_flight_purges: Arc::new(Mutex::new(HashSet::new())), + purge_scan_cursor: AtomicUsize::new(0), meta_peer_client, mailbox, server_addr, @@ -154,6 +240,36 @@ impl SchedulerCtx for DefaultGcSchedulerCtx { ) .await } + + async fn list_dropped_tables(&self) -> Result> { + self.table_metadata_manager + .list_dropped_tables() + .await + .context(TableMetadataManagerSnafu) + } + + async fn purge_dropped_table(&self, table_id: TableId) -> Result<()> { + self.ddl_manager + .submit_expired_purge_dropped_table_task(PurgeDroppedTableTask { table_id }) + .await + .context(error::SubmitDdlTaskSnafu)?; + Ok(()) + } + + fn try_reserve_purge( + &self, + table_id: TableId, + max_in_flight: usize, + ) -> Option { + PurgeReservation::try_new(self.in_flight_purges.clone(), table_id, max_in_flight) + } + + fn next_purge_scan_start(&self, table_count: usize) -> usize { + if table_count == 0 { + return 0; + } + self.purge_scan_cursor.fetch_add(1, Ordering::Relaxed) % table_count + } } impl DefaultGcSchedulerCtx { @@ -202,3 +318,48 @@ impl DefaultGcSchedulerCtx { Ok(gc_report) } } + +#[cfg(test)] +mod tests { + use std::panic::{AssertUnwindSafe, catch_unwind}; + + use super::*; + use crate::metrics::METRIC_META_GC_SOFT_DROP_PURGES_TOTAL; + + #[test] + fn test_purge_reservation_releases_slot_on_panic() { + let in_flight = Arc::new(std::sync::Mutex::new(HashSet::new())); + let cancelled = METRIC_META_GC_SOFT_DROP_PURGES_TOTAL.with_label_values(&["cancelled"]); + let before = cancelled.get(); + let reservation = PurgeReservation::try_new(in_flight.clone(), 1, 1).unwrap(); + + let result = catch_unwind(AssertUnwindSafe(|| { + let _reservation = reservation; + panic!("mock purge panic"); + })); + + assert!(result.is_err()); + assert!(PurgeReservation::try_new(in_flight, 2, 1).is_some()); + assert!(cancelled.get() > before); + } + + #[tokio::test] + async fn test_purge_reservation_releases_slot_on_task_abort() { + let in_flight = Arc::new(std::sync::Mutex::new(HashSet::new())); + let cancelled = METRIC_META_GC_SOFT_DROP_PURGES_TOTAL.with_label_values(&["cancelled"]); + let before = cancelled.get(); + let reservation = PurgeReservation::try_new(in_flight.clone(), 1, 1).unwrap(); + let handle = tokio::spawn(async move { + let _reservation = reservation; + std::future::pending::<()>().await; + }); + tokio::task::yield_now().await; + + handle.abort(); + let error = handle.await.unwrap_err(); + + assert!(error.is_cancelled()); + assert!(PurgeReservation::try_new(in_flight, 2, 1).is_some()); + assert!(cancelled.get() > before); + } +} diff --git a/src/meta-srv/src/gc/mock.rs b/src/meta-srv/src/gc/mock.rs index 0686d9a4653..f801df2707e 100644 --- a/src/meta-srv/src/gc/mock.rs +++ b/src/meta-srv/src/gc/mock.rs @@ -326,6 +326,26 @@ impl SchedulerCtx for MockSchedulerCtx { // Return the report with need_retry_regions populated - let the caller handle retry logic Ok(final_report) } + + async fn list_dropped_tables(&self) -> Result> { + Ok(vec![]) + } + + async fn purge_dropped_table(&self, _table_id: TableId) -> Result<()> { + Ok(()) + } + + fn try_reserve_purge( + &self, + table_id: TableId, + max_in_flight: usize, + ) -> Option { + crate::gc::ctx::PurgeReservation::try_new( + Arc::new(std::sync::Mutex::new(HashSet::new())), + table_id, + max_in_flight, + ) + } } pub struct TestEnv { diff --git a/src/meta-srv/src/gc/options.rs b/src/meta-srv/src/gc/options.rs index f6da6b850aa..4d1b93b9d95 100644 --- a/src/meta-srv/src/gc/options.rs +++ b/src/meta-srv/src/gc/options.rs @@ -23,6 +23,29 @@ use crate::error::{self, Result}; #[allow(unused)] pub(crate) const TICKER_INTERVAL: Duration = Duration::from_secs(60 * 5); +/// Soft drop remains disconnected from user configuration until it is ready. +pub(crate) const EXPERIMENTAL_SOFT_DROP_ENABLED: bool = false; + +/// Reserved configuration for garbage collecting soft-dropped tables. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(default)] +pub struct ExperimentalSoftDropGcOptions { + /// Reserved. This option is currently ignored. + pub enable: bool, + /// Reserved retention duration. This option is currently ignored. + #[serde(with = "humantime_serde")] + pub retention: Duration, +} + +impl Default for ExperimentalSoftDropGcOptions { + fn default() -> Self { + Self { + enable: false, + retention: Duration::from_days(7), + } + } +} + /// Configuration for GC operations. /// /// TODO(discord9): not expose most config to users for now, until GC scheduler is fully stable. @@ -33,6 +56,8 @@ pub struct GcSchedulerOptions { /// If set to false, no GC will be performed, and potentially some /// files from datanodes will never be deleted. pub enable: bool, + /// Reserved experimental soft-drop garbage collection options. + pub experimental_soft_drop: ExperimentalSoftDropGcOptions, /// Maximum number of tables to process concurrently. pub max_concurrent_tables: usize, /// Maximum number of retries per region when GC fails. @@ -73,6 +98,7 @@ impl Default for GcSchedulerOptions { fn default() -> Self { Self { enable: false, + experimental_soft_drop: ExperimentalSoftDropGcOptions::default(), max_concurrent_tables: 10, max_retries_per_region: 3, retry_backoff_duration: Duration::from_secs(5), @@ -174,3 +200,46 @@ impl GcSchedulerOptions { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_experimental_soft_drop_defaults() { + let options = GcSchedulerOptions::default(); + + assert!(!options.experimental_soft_drop.enable); + assert_eq!( + Duration::from_days(7), + options.experimental_soft_drop.retention + ); + } + + #[test] + fn test_experimental_soft_drop_valid_when_gc_is_enabled() { + let options = GcSchedulerOptions { + enable: true, + experimental_soft_drop: ExperimentalSoftDropGcOptions { + enable: true, + retention: Duration::from_days(1), + }, + ..Default::default() + }; + + assert!(options.validate().is_ok()); + } + + #[test] + fn test_experimental_soft_drop_options_are_ignored_by_validation() { + let options = GcSchedulerOptions { + experimental_soft_drop: ExperimentalSoftDropGcOptions { + enable: true, + retention: Duration::ZERO, + }, + ..Default::default() + }; + + assert!(options.validate().is_ok()); + } +} diff --git a/src/meta-srv/src/gc/scheduler.rs b/src/meta-srv/src/gc/scheduler.rs index 4f732f52bc0..e3e6a5f96ec 100644 --- a/src/meta-srv/src/gc/scheduler.rs +++ b/src/meta-srv/src/gc/scheduler.rs @@ -17,9 +17,7 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use common_meta::DatanodeId; -use common_meta::key::TableMetadataManagerRef; use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; -use common_procedure::ProcedureManagerRef; use common_telemetry::tracing::Instrument as _; use common_telemetry::{error, info}; use snafu::ResultExt; @@ -27,18 +25,19 @@ use store_api::storage::{GcReport, RegionId}; use tokio::sync::mpsc::{Receiver, Sender}; use tokio::sync::{Mutex, oneshot}; -use crate::cluster::MetaPeerClientRef; use crate::define_ticker; use crate::error::{self, Error, Result}; -use crate::gc::Region2Peers; -use crate::gc::ctx::{DefaultGcSchedulerCtx, SchedulerCtx}; +#[cfg(test)] +use crate::gc::ctx::PurgeReservation; +use crate::gc::ctx::{PurgeOutcome, SchedulerCtx}; use crate::gc::dropped::DroppedRegionCollector; use crate::gc::options::{GcSchedulerOptions, TICKER_INTERVAL}; use crate::gc::tracker::RegionGcTracker; +use crate::gc::{EXPERIMENTAL_SOFT_DROP_ENABLED, Region2Peers}; use crate::metrics::{ METRIC_META_GC_SCHEDULER_CYCLES_TOTAL, METRIC_META_GC_SCHEDULER_DURATION_SECONDS, + METRIC_META_GC_SOFT_DROP_PURGES_TOTAL, }; -use crate::service::mailbox::MailboxRef; /// Report for a GC job. #[derive(Debug)] @@ -132,11 +131,7 @@ pub struct GcScheduler { impl GcScheduler { /// Creates a new [`GcScheduler`] with custom configuration. pub(crate) fn new_with_config( - table_metadata_manager: TableMetadataManagerRef, - procedure_manager: ProcedureManagerRef, - meta_peer_client: MetaPeerClientRef, - mailbox: MailboxRef, - server_addr: String, + ctx: impl SchedulerCtx + 'static, runtime_switch_manager: RuntimeSwitchManagerRef, config: GcSchedulerOptions, ) -> Result<(Self, GcTicker)> { @@ -146,13 +141,7 @@ impl GcScheduler { let (tx, rx) = Self::channel(); let gc_ticker = GcTicker::new(TICKER_INTERVAL, tx); let gc_trigger = Self { - ctx: Arc::new(DefaultGcSchedulerCtx::try_new( - table_metadata_manager, - procedure_manager, - meta_peer_client, - mailbox, - server_addr, - )?), + ctx: Arc::new(ctx), runtime_switch_manager, receiver: rx, config, @@ -218,6 +207,10 @@ impl GcScheduler { info!("Skip gc trigger because maintenance mode is enabled"); return Ok(GcJobReport::default()); } + if EXPERIMENTAL_SOFT_DROP_ENABLED && self.config.experimental_soft_drop.enable { + self.purge_expired_soft_dropped_tables(common_time::util::current_time_millis()) + .await; + } let span = common_telemetry::tracing::info_span!("meta_gc_handle_tick"); let report = self.trigger_gc().instrument(span).await?; @@ -229,6 +222,60 @@ impl GcScheduler { Ok(report) } + async fn purge_expired_soft_dropped_tables(&self, now_millis: i64) { + // The scheduler is only constructed after GcSchedulerOptions::validate(). + let retention_millis = + i64::try_from(self.config.experimental_soft_drop.retention.as_millis()) + .unwrap_or(i64::MAX); + let dropped_tables = match self.ctx.list_dropped_tables().await { + Ok(dropped_tables) => dropped_tables, + Err(error) => { + error!(error; "Failed to list soft-dropped tables for GC"); + return; + } + }; + let table_count = dropped_tables.len(); + let scan_start = self.ctx.next_purge_scan_start(table_count); + + for table in dropped_tables + .iter() + .cycle() + .skip(scan_start) + .take(table_count) + .filter(|table| { + table + .retention_expires_at + .or_else(|| { + table + .dropped_at + .and_then(|dropped_at| dropped_at.checked_add(retention_millis)) + }) + .is_some_and(|expires_at| expires_at <= now_millis) + }) + { + let table_id = table.table_id; + let Some(reservation) = self + .ctx + .try_reserve_purge(table_id, self.config.max_concurrent_tables) + else { + continue; + }; + METRIC_META_GC_SOFT_DROP_PURGES_TOTAL + .with_label_values(&["submitted"]) + .inc(); + let ctx = self.ctx.clone(); + common_runtime::spawn_global(async move { + match ctx.purge_dropped_table(table_id).await { + Ok(()) => reservation.record_outcome(PurgeOutcome::Succeeded), + Err(error) => { + reservation.record_outcome(PurgeOutcome::Failed); + error!(error; "Failed to purge expired soft-dropped table {}", table_id); + } + } + }); + } + } + /// Handles a manual GC request with optional specific parameters. /// /// If `region_ids` is specified, GC will be performed only on those regions. @@ -332,14 +379,17 @@ pub(crate) fn new_test_runtime_switch_manager() -> RuntimeSwitchManagerRef { #[cfg(test)] mod tests { use std::collections::HashMap; + use std::sync::Mutex as StdMutex; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use common_meta::datanode::RegionStat; + use common_meta::key::DroppedTableName; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use store_api::storage::RegionId; use table::metadata::TableId; + use table::table_name::TableName; use super::*; @@ -348,6 +398,8 @@ mod tests { get_table_to_region_stats_calls: AtomicUsize, get_table_reparts_calls: AtomicUsize, gc_regions_calls: AtomicUsize, + list_dropped_tables_calls: AtomicUsize, + purge_dropped_table_calls: AtomicUsize, } impl CountingSchedulerCtx { @@ -367,6 +419,16 @@ mod tests { self.gc_regions_calls.load(Ordering::Relaxed), "gc_regions should not be called" ); + assert_eq!( + 0, + self.list_dropped_tables_calls.load(Ordering::Relaxed), + "list_dropped_tables should not be called" + ); + assert_eq!( + 0, + self.purge_dropped_table_calls.load(Ordering::Relaxed), + "purge_dropped_table should not be called" + ); } } @@ -407,6 +469,353 @@ mod tests { self.gc_regions_calls.fetch_add(1, Ordering::Relaxed); panic!("gc_regions should not be called in maintenance mode") } + + async fn list_dropped_tables(&self) -> Result> { + self.list_dropped_tables_calls + .fetch_add(1, Ordering::Relaxed); + panic!("list_dropped_tables should not be called in maintenance mode") + } + + async fn purge_dropped_table(&self, _table_id: TableId) -> Result<()> { + self.purge_dropped_table_calls + .fetch_add(1, Ordering::Relaxed); + panic!("purge_dropped_table should not be called in maintenance mode") + } + + fn try_reserve_purge( + &self, + _table_id: TableId, + _max_in_flight: usize, + ) -> Option { + panic!("try_reserve_purge should not be called in maintenance mode") + } + } + + #[derive(Default)] + struct SoftDropSchedulerCtx { + dropped_tables: StdMutex>, + purge_attempts: StdMutex>, + failed_table: StdMutex>, + never_complete_tables: StdMutex>, + in_flight_purges: Arc>>, + purge_scan_cursor: AtomicUsize, + region_gc_calls: AtomicUsize, + } + + #[async_trait::async_trait] + impl SchedulerCtx for SoftDropSchedulerCtx { + async fn get_table_to_region_stats(&self) -> Result>> { + self.region_gc_calls.fetch_add(1, Ordering::Relaxed); + Ok(HashMap::new()) + } + + async fn get_table_reparts(&self) -> Result> { + Ok(vec![]) + } + + async fn get_table_route( + &self, + _table_id: TableId, + ) -> Result<(TableId, PhysicalTableRouteValue)> { + unreachable!() + } + + async fn batch_get_table_route( + &self, + _table_ids: &[TableId], + ) -> Result> { + Ok(HashMap::new()) + } + + async fn gc_regions( + &self, + _region_ids: &[RegionId], + _full_file_listing: bool, + _timeout: Duration, + _region_routes_override: Region2Peers, + ) -> Result { + Ok(GcReport::default()) + } + + async fn list_dropped_tables(&self) -> Result> { + Ok(self.dropped_tables.lock().unwrap().clone()) + } + + async fn purge_dropped_table(&self, table_id: TableId) -> Result<()> { + self.purge_attempts.lock().unwrap().push(table_id); + if self + .never_complete_tables + .lock() + .unwrap() + .contains(&table_id) + { + std::future::pending().await + } + if *self.failed_table.lock().unwrap() == Some(table_id) { + return crate::error::UnexpectedSnafu { + violated: format!("mock purge failure for table {table_id}"), + } + .fail(); + } + Ok(()) + } + + fn try_reserve_purge( + &self, + table_id: TableId, + max_in_flight: usize, + ) -> Option { + PurgeReservation::try_new(self.in_flight_purges.clone(), table_id, max_in_flight) + } + + fn next_purge_scan_start(&self, table_count: usize) -> usize { + if table_count == 0 { + return 0; + } + self.purge_scan_cursor.fetch_add(1, Ordering::Relaxed) % table_count + } + } + + fn dropped_table(table_id: TableId, dropped_at: Option) -> DroppedTableName { + DroppedTableName { + table_id, + table_name: TableName::new("greptime", "public", format!("table_{table_id}")), + dropped_at, + retention_expires_at: None, + drop_generation: None, + } + } + + fn dropped_table_with_deadline(table_id: TableId, deadline: i64) -> DroppedTableName { + DroppedTableName { + retention_expires_at: Some(deadline), + ..dropped_table(table_id, Some(i64::MAX)) + } + } + + fn soft_drop_scheduler(ctx: Arc) -> GcScheduler { + let (tx, rx) = GcScheduler::channel(); + drop(tx); + GcScheduler { + ctx, + runtime_switch_manager: new_test_runtime_switch_manager(), + receiver: rx, + config: GcSchedulerOptions { + enable: true, + experimental_soft_drop: crate::gc::options::ExperimentalSoftDropGcOptions { + enable: true, + retention: Duration::from_millis(100), + }, + ..Default::default() + }, + region_gc_tracker: Arc::new(Mutex::new(HashMap::new())), + last_tracker_cleanup: Arc::new(Mutex::new(Instant::now())), + } + } + + async fn wait_for_purge_attempts(ctx: &SoftDropSchedulerCtx, expected: usize) { + tokio::time::timeout(Duration::from_secs(1), async { + while ctx.purge_attempts.lock().unwrap().len() < expected { + tokio::task::yield_now().await; + } + }) + .await + .expect("purge tasks should be scheduled"); + } + + #[tokio::test] + async fn test_purge_expired_soft_dropped_tables_filters_by_retention_and_timestamp() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![ + dropped_table(1, Some(899)), + dropped_table(2, Some(900)), + dropped_table(3, Some(901)), + dropped_table(4, None), + dropped_table(5, Some(i64::MAX)), + ]; + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 2).await; + + let mut attempts = ctx.purge_attempts.lock().unwrap().clone(); + attempts.sort_unstable(); + assert_eq!(vec![1, 2], attempts); + } + + #[tokio::test] + async fn test_purge_uses_fixed_deadline_before_legacy_retention_fallback() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![ + dropped_table_with_deadline(1, 1_000), + dropped_table_with_deadline(2, 1_001), + dropped_table(3, Some(900)), + ]; + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 2).await; + + let mut attempts = ctx.purge_attempts.lock().unwrap().clone(); + attempts.sort_unstable(); + assert_eq!(vec![1, 3], attempts); + } + + #[tokio::test] + async fn test_purge_saturates_legacy_retention_that_exceeds_i64() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![dropped_table(1, Some(0))]; + ctx.never_complete_tables.lock().unwrap().insert(1); + let mut scheduler = soft_drop_scheduler(ctx.clone()); + scheduler.config.experimental_soft_drop.retention = + Duration::from_millis(i64::MAX as u64 + 1); + + scheduler.purge_expired_soft_dropped_tables(0).await; + + assert!(!ctx.in_flight_purges.lock().unwrap().contains(&1)); + } + + #[tokio::test] + async fn test_purge_failure_does_not_prevent_other_purges() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![ + dropped_table(1, Some(0)), + dropped_table(2, Some(0)), + dropped_table(3, Some(0)), + ]; + *ctx.failed_table.lock().unwrap() = Some(2); + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 3).await; + + let mut attempts = ctx.purge_attempts.lock().unwrap().clone(); + attempts.sort_unstable(); + assert_eq!(vec![1, 2, 3], attempts); + } + + #[tokio::test] + async fn test_purge_scans_rotate_after_failure() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = + vec![dropped_table(1, Some(0)), dropped_table(2, Some(0))]; + *ctx.failed_table.lock().unwrap() = Some(1); + let mut scheduler = soft_drop_scheduler(ctx.clone()); + scheduler.config.max_concurrent_tables = 1; + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 1).await; + tokio::time::timeout(Duration::from_secs(1), async { + while !ctx.in_flight_purges.lock().unwrap().is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .expect("failed purge should release its in-flight slot"); + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 2).await; + + assert_eq!(vec![1, 2], *ctx.purge_attempts.lock().unwrap()); + } + + #[tokio::test] + async fn test_tick_ignores_experimental_soft_drop_config_and_still_runs_region_gc() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![dropped_table(1, Some(i64::MIN))]; + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.handle_tick().await.unwrap(); + + assert!(ctx.purge_attempts.lock().unwrap().is_empty()); + assert_eq!(1, ctx.region_gc_calls.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn test_tick_with_no_tombstones_makes_no_purge_submissions() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.handle_tick().await.unwrap(); + + assert!(ctx.purge_attempts.lock().unwrap().is_empty()); + assert_eq!(1, ctx.region_gc_calls.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn test_multiple_purge_scans_do_not_duplicate_in_flight_purges() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![ + dropped_table(1, Some(i64::MIN)), + dropped_table(2, Some(i64::MIN)), + ]; + ctx.never_complete_tables.lock().unwrap().insert(1); + *ctx.failed_table.lock().unwrap() = Some(2); + let mut scheduler = soft_drop_scheduler(ctx.clone()); + scheduler.config.max_concurrent_tables = 2; + + tokio::time::timeout( + Duration::from_millis(100), + scheduler.purge_expired_soft_dropped_tables(1_000), + ) + .await + .expect("first scan should not wait for purge completion"); + wait_for_purge_attempts(&ctx, 2).await; + tokio::time::timeout(Duration::from_secs(1), async { + while ctx.in_flight_purges.lock().unwrap().contains(&2) { + tokio::task::yield_now().await; + } + }) + .await + .expect("terminal failure should release its in-flight slot"); + tokio::time::timeout( + Duration::from_millis(100), + scheduler.purge_expired_soft_dropped_tables(1_000), + ) + .await + .expect("second scan should not wait for purge completion"); + wait_for_purge_attempts(&ctx, 3).await; + + let attempts = ctx.purge_attempts.lock().unwrap().clone(); + assert_eq!( + 1, + attempts.iter().filter(|&&table_id| table_id == 1).count() + ); + assert_eq!( + 2, + attempts.iter().filter(|&&table_id| table_id == 2).count() + ); + assert_eq!(0, ctx.region_gc_calls.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn test_purge_scan_caps_submissions() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![ + dropped_table(1, Some(i64::MIN)), + dropped_table(2, Some(i64::MIN)), + dropped_table(3, Some(i64::MIN)), + ]; + let mut scheduler = soft_drop_scheduler(ctx.clone()); + scheduler.config.max_concurrent_tables = 2; + ctx.never_complete_tables.lock().unwrap().extend([1, 2]); + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 2).await; + + assert_eq!(2, ctx.purge_attempts.lock().unwrap().len()); + assert_eq!(0, ctx.region_gc_calls.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn test_manual_gc_does_not_purge_soft_dropped_tables() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![dropped_table(1, Some(i64::MIN))]; + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.handle_manual_gc(None, None, None).await.unwrap(); + + assert!(ctx.purge_attempts.lock().unwrap().is_empty()); } struct ErrorMockSchedulerCtx; @@ -447,6 +856,26 @@ mod tests { } .fail() } + + async fn list_dropped_tables(&self) -> Result> { + Ok(vec![]) + } + + async fn purge_dropped_table(&self, _table_id: TableId) -> Result<()> { + Ok(()) + } + + fn try_reserve_purge( + &self, + table_id: TableId, + max_in_flight: usize, + ) -> Option { + PurgeReservation::try_new( + Arc::new(StdMutex::new(HashSet::new())), + table_id, + max_in_flight, + ) + } } #[tokio::test] diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 8a6f3147bfe..fd7be6ae6d4 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -58,7 +58,7 @@ use crate::cache_invalidator::MetasrvCacheInvalidator; use crate::cluster::MetaPeerClientRef; use crate::error::{self, BuildWalProviderSnafu, OtherSnafu, Result}; use crate::events::EventHandlerImpl; -use crate::gc::GcScheduler; +use crate::gc::{DefaultGcSchedulerCtx, EXPERIMENTAL_SOFT_DROP_ENABLED, GcScheduler}; use crate::greptimedb_telemetry::get_greptimedb_telemetry_task; use crate::handler::failure_handler::RegionFailureHandler; use crate::handler::flow_state_handler::FlowStateHandler; @@ -430,6 +430,7 @@ impl MetasrvBuilder { flow_metadata_allocator: flow_metadata_allocator.clone(), region_failure_detector_controller, soft_drop_enabled: ddl_soft_drop_enabled(&options), + soft_drop_retention: ddl_soft_drop_retention(&options), }; let procedure_manager_c = procedure_manager.clone(); let repartition_procedure_factory: RepartitionProcedureFactoryRef = if options.gc.enable { @@ -521,12 +522,16 @@ impl MetasrvBuilder { }; let gc_ticker = if options.gc.enable { - let (gc_scheduler, gc_ticker) = GcScheduler::new_with_config( + let gc_scheduler_ctx = DefaultGcSchedulerCtx::try_new( table_metadata_manager.clone(), procedure_manager.clone(), + ddl_manager.clone(), meta_peer_client.clone(), mailbox.clone(), options.grpc.server_addr.clone(), + )?; + let (gc_scheduler, gc_ticker) = GcScheduler::new_with_config( + gc_scheduler_ctx, runtime_switch_manager.clone(), options.gc.clone(), )?; @@ -687,10 +692,13 @@ fn build_procedure_manager( } /// Resolves if soft-drop is enabled from metasrv options. -fn ddl_soft_drop_enabled(_options: &MetasrvOptions) -> bool { - // TODO(hl): add a dedicated soft-drop cluster config - // when wiring the user-facing option. - false +fn ddl_soft_drop_enabled(options: &MetasrvOptions) -> bool { + EXPERIMENTAL_SOFT_DROP_ENABLED && options.gc.experimental_soft_drop.enable +} + +/// Returns soft-drop retention for recovering persisted procedures. +fn ddl_soft_drop_retention(options: &MetasrvOptions) -> Option { + Some(options.gc.experimental_soft_drop.retention) } impl Default for MetasrvBuilder { @@ -704,3 +712,49 @@ pub struct DdlManagerConfigureContext { pub kv_backend: KvBackendRef, pub meta_peer_client: MetaPeerClientRef, } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_ddl_soft_drop_gate_preserves_retention_for_recovery() { + let mut options = MetasrvOptions::default(); + options.gc.enable = true; + options.gc.experimental_soft_drop.enable = true; + options.gc.experimental_soft_drop.retention = Duration::from_secs(123); + + assert!(!ddl_soft_drop_enabled(&options)); + assert_eq!( + Some(Duration::from_secs(123)), + ddl_soft_drop_retention(&options) + ); + } + + #[test] + fn test_ddl_soft_drop_is_disabled_by_default() { + assert!(!ddl_soft_drop_enabled(&MetasrvOptions::default())); + } + + #[test] + fn test_ignored_soft_drop_options_are_not_validated() { + let mut options = MetasrvOptions::default(); + options.gc.experimental_soft_drop.enable = true; + options.gc.experimental_soft_drop.retention = Duration::ZERO; + + assert!(options.gc.validate().is_ok()); + } + + #[tokio::test] + async fn test_builder_skips_gc_validation_when_gc_is_disabled() { + let mut options = MetasrvOptions::default(); + options.gc.enable = false; + options.gc.max_concurrent_tables = 0; + + MetasrvBuilder::new() + .options(options) + .build() + .await + .unwrap(); + } +} diff --git a/src/meta-srv/src/metrics.rs b/src/meta-srv/src/metrics.rs index 879b1de2065..76283495c06 100644 --- a/src/meta-srv/src/metrics.rs +++ b/src/meta-srv/src/metrics.rs @@ -127,4 +127,12 @@ lazy_static! { "Total regions that failed GC", ) .unwrap(); + + /// Soft-drop purge submissions by status. + pub static ref METRIC_META_GC_SOFT_DROP_PURGES_TOTAL: IntCounterVec = register_int_counter_vec!( + "greptime_metasrv_gc_soft_drop_purges_total", + "Soft-drop purge submissions by status", + &["status"], + ) + .unwrap(); } diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index bcae558773a..81dfe56e670 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -547,6 +547,7 @@ pub mod test_data { leader_region_registry: Arc::new(LeaderRegionRegistry::default()), region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), soft_drop_enabled: false, + soft_drop_retention: None, } } } diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 122318d45f1..37ca78f66f3 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -251,6 +251,7 @@ impl GreptimeDbStandaloneBuilder { flow_metadata_allocator, region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), soft_drop_enabled: false, + soft_drop_retention: None, }, procedure_manager.clone(), repartition_procedure_factory, From 3ae8d7d46ae919334e177c21de67778f13730058 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Thu, 16 Jul 2026 18:56:27 +0800 Subject: [PATCH 47/48] refactor: reconcile OTLP trace schemas request-wide (#8485) * refactor: global trace schema Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: fmt Signed-off-by: shuiyisong * refactor: split files Signed-off-by: shuiyisong * chore: add comments and doc Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/frontend/src/instance/otlp.rs | 935 +------- src/frontend/src/instance/otlp/README.md | 182 ++ .../src/instance/otlp/trace_ingest.rs | 1965 +++++++++++++++++ .../src/instance/otlp/trace_ingest/tests.rs | 1212 ++++++++++ src/frontend/src/instance/otlp/trace_types.rs | 163 +- src/operator/src/insert.rs | 69 +- src/servers/src/otlp/coerce.rs | 67 +- src/servers/src/otlp/trace/v1.rs | 444 +++- src/servers/src/row_writer.rs | 31 +- tests-integration/tests/http.rs | 6 +- 10 files changed, 4033 insertions(+), 1041 deletions(-) create mode 100644 src/frontend/src/instance/otlp/README.md create mode 100644 src/frontend/src/instance/otlp/trace_ingest.rs create mode 100644 src/frontend/src/instance/otlp/trace_ingest/tests.rs diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 16e3b996064..c970ca144db 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -12,20 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod trace_ingest; +pub mod trace_semconv; +pub mod trace_types; + use std::sync::Arc; -use api::helper::ColumnDataTypeWrapper; -use api::v1::alter_table_expr::Kind; -use api::v1::{ - AlterTableExpr, ColumnDataType, ModifyColumnType, ModifyColumnTypes, RowInsertRequests, -}; use async_trait::async_trait; use auth::{PermissionChecker, PermissionCheckerRef, PermissionReq}; use client::Output; -use common_error::ext::{BoxedError, ErrorExt}; -use common_error::status_code::StatusCode; +use common_error::ext::BoxedError; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; -use common_telemetry::{tracing, warn}; +use common_telemetry::tracing; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; @@ -34,66 +32,20 @@ use servers::error::{self, AuthSnafu, Result as ServerResult}; use servers::http::prom_store::PHYSICAL_TABLE_PARAM; use servers::interceptor::{OpenTelemetryProtocolInterceptor, OpenTelemetryProtocolInterceptorRef}; use servers::otlp; -use servers::otlp::coerce::{coerce_value_data, trace_value_datatype}; -use servers::otlp::trace::TraceAuxData; -use servers::otlp::trace::span::{TraceSpan, TraceSpanGroup}; use servers::query_handler::{ OpenTelemetryProtocolHandler, PipelineHandlerRef, TraceIngestOutcome, }; use session::context::QueryContextRef; -use snafu::{IntoError, ResultExt}; +use snafu::ResultExt; use table::requests::{ OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM, SEMANTIC_PER_TABLE_INDEX_KEY, - SEMANTIC_PIPELINE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SEMANTIC_TRACE_CONVENTIONS, - SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, SIGNAL_TYPE_LOG, SIGNAL_TYPE_METRIC, - SIGNAL_TYPE_TRACE, SOURCE_OPENTELEMETRY, TABLE_DATA_MODEL_TRACE_V1, + SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_LOG, SIGNAL_TYPE_METRIC, + SOURCE_OPENTELEMETRY, }; +use self::trace_ingest::trace_conventions; use crate::instance::Instance; -use crate::instance::otlp::trace_semconv::trace_semconv_fixed_type; -use crate::instance::otlp::trace_types::{ - PendingTraceColumnRewrite, choose_trace_reconcile_decision, enrich_trace_reconcile_error, - is_trace_reconcile_candidate_type, push_observed_trace_type, validate_trace_column_rewrites, -}; -use crate::metrics::{ - OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_TRACES_FAILURE_COUNT, OTLP_TRACES_ROWS, -}; - -pub mod trace_semconv; -pub mod trace_types; - -const TRACE_FAILURE_MESSAGE_LIMIT: usize = 4; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ChunkFailureReaction { - RetryPerSpan, - DiscardChunk, - Propagate, -} - -impl ChunkFailureReaction { - fn as_metric_label(self) -> &'static str { - match self { - Self::RetryPerSpan => "retry_per_span", - Self::DiscardChunk => "discard_chunk", - Self::Propagate => "propagate_failure", - } - } -} - -struct TraceChunkIngestContext<'a> { - pipeline_handler: PipelineHandlerRef, - pipeline: &'a PipelineWay, - pipeline_params: &'a GreptimePipelineParams, - table_name: &'a str, - is_trace_v1_model: bool, -} - -struct TraceIngestState { - aux_data: TraceAuxData, - outcome: TraceIngestOutcome, - failure_messages: Vec, -} +use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS}; #[async_trait] impl OpenTelemetryProtocolHandler for Instance { @@ -265,868 +217,3 @@ impl OpenTelemetryProtocolHandler for Instance { Ok(outputs) } } - -impl Instance { - /// Ingest OTLP trace spans with chunk-level writes and span-level fallback on - /// deterministic chunk failures. - #[allow(clippy::too_many_arguments)] - async fn ingest_trace_spans( - &self, - pipeline_handler: PipelineHandlerRef, - pipeline: &PipelineWay, - pipeline_params: &GreptimePipelineParams, - table_name: String, - groups: Vec, - conventions: &str, - ctx: QueryContextRef, - ) -> ServerResult { - let is_trace_v1_model = matches!(pipeline, PipelineWay::OtlpTraceDirectV1); - - // Only the main span table gets the identity; the derived `_services` / - // `_operations` lookup tables keep the unstamped `ctx`. - let main_ctx = { - let mut c = (*ctx).clone(); - c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_TRACE); - c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY); - if is_trace_v1_model { - c.set_extension(SEMANTIC_PIPELINE, TABLE_DATA_MODEL_TRACE_V1); - c.set_extension(SEMANTIC_TRACE_CONVENTIONS, conventions); - } - Arc::new(c) - }; - - let ingest_ctx = TraceChunkIngestContext { - pipeline_handler, - pipeline, - pipeline_params, - table_name: &table_name, - is_trace_v1_model, - }; - let mut ingest_state = TraceIngestState { - aux_data: TraceAuxData::default(), - outcome: TraceIngestOutcome::default(), - failure_messages: Vec::new(), - }; - - for group in groups { - let chunks = chunk_owned(group.spans, self.trace_ingest_chunk_size); - for chunk in chunks { - self.ingest_trace_chunk(&ingest_ctx, chunk, main_ctx.clone(), &mut ingest_state) - .await?; - } - } - - OTLP_TRACES_ROWS.inc_by(ingest_state.outcome.accepted_spans as u64); - - if !ingest_state.aux_data.is_empty() { - // Auxiliary trace tables are derived from spans whose main-table - // writes are already confirmed, so they never create new accepted - // spans and they do not affect rejected span counts. - let (aux_requests, _) = otlp::trace::to_grpc_insert_requests_for_aux_tables( - std::mem::take(&mut ingest_state.aux_data), - ingest_ctx.pipeline, - ingest_ctx.table_name, - )?; - - if !aux_requests.inserts.is_empty() { - match self - .insert_trace_requests(aux_requests, ingest_ctx.is_trace_v1_model, ctx) - .await - { - Ok(output) => { - Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); - } - Err(err) => { - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - "aux_table_update_failed", - format!( - "Auxiliary trace tables were not fully updated ({})", - err.status_code().as_ref() - ), - ); - } - } - } - } - - ingest_state.outcome.error_message = Self::finish_trace_failure_message( - ingest_state.outcome.accepted_spans, - ingest_state.outcome.rejected_spans, - ingest_state.failure_messages, - ); - - Ok(ingest_state.outcome) - } - - /// Ingest one owned trace chunk so successful spans can be moved into the - /// accepted set without extra cloning. - async fn ingest_trace_chunk( - &self, - ingest_ctx: &TraceChunkIngestContext<'_>, - chunk: Vec, - ctx: QueryContextRef, - ingest_state: &mut TraceIngestState, - ) -> ServerResult<()> { - // Try the fast path first so healthy batches keep their original - // throughput and write amplification stays low. - let (requests, chunk_rows) = otlp::trace::to_grpc_insert_requests_from_spans( - &chunk, - ingest_ctx.pipeline, - ingest_ctx.pipeline_params, - ingest_ctx.table_name, - &ctx, - ingest_ctx.pipeline_handler.clone(), - )?; - - match self - .insert_trace_requests(requests, ingest_ctx.is_trace_v1_model, ctx.clone()) - .await - { - Ok(output) => { - Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); - ingest_state.outcome.accepted_spans += chunk_rows; - for span in &chunk { - ingest_state.aux_data.observe_span(span); - } - } - Err(err) => match Self::classify_trace_chunk_failure(err.status_code()) { - ChunkFailureReaction::RetryPerSpan => { - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - ChunkFailureReaction::RetryPerSpan.as_metric_label(), - format!("Chunk fallback triggered by {}", err.status_code().as_ref()), - ); - // Only deterministic failures are retried span by span. - // This includes schemaless table or column creation paths for - // trace ingestion. Ambiguous failures are handled below - // without retrying because the chunk may already have been - // ingested. - self.ingest_trace_chunk_span_by_span( - ingest_ctx, - chunk, - ctx.clone(), - ingest_state, - ) - .await?; - } - ChunkFailureReaction::DiscardChunk => { - ingest_state.outcome.rejected_spans += chunk.len(); - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - ChunkFailureReaction::DiscardChunk.as_metric_label(), - format!( - "Discarded {} spans after ambiguous chunk failure ({})", - chunk.len(), - err.status_code().as_ref() - ), - ); - // TODO(shuiyisong): Add an idempotent retry-safe recovery path for - // ambiguous chunk failures such as timeout-like errors. - } - // Retryable or ambiguous failures must fail the request instead of - // becoming partial success. This path is not retry-safe because the - // chunk may already have been committed before the error surfaced. - ChunkFailureReaction::Propagate => { - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - ChunkFailureReaction::Propagate.as_metric_label(), - format!( - "Propagating retryable chunk failure ({})", - err.status_code().as_ref() - ), - ); - return Err(err); - } - }, - } - - Ok(()) - } - - /// Retry spans one by one only after a deterministic chunk failure. - async fn ingest_trace_chunk_span_by_span( - &self, - ingest_ctx: &TraceChunkIngestContext<'_>, - chunk: Vec, - ctx: QueryContextRef, - ingest_state: &mut TraceIngestState, - ) -> ServerResult<()> { - for span in chunk { - let (requests, rows) = otlp::trace::to_grpc_insert_requests_from_spans( - std::slice::from_ref(&span), - ingest_ctx.pipeline, - ingest_ctx.pipeline_params, - ingest_ctx.table_name, - &ctx, - ingest_ctx.pipeline_handler.clone(), - )?; - - match self - .insert_trace_requests(requests, ingest_ctx.is_trace_v1_model, ctx.clone()) - .await - { - Ok(output) => { - Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); - ingest_state.outcome.accepted_spans += rows; - ingest_state.aux_data.observe_span(&span); - } - Err(err) => { - if Self::should_propagate_trace_span_failure(err.status_code()) { - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - ChunkFailureReaction::Propagate.as_metric_label(), - format!( - "Propagating retryable span failure for {}:{} ({})", - span.trace_id, - span.span_id, - err.status_code().as_ref() - ), - ); - return Err(err); - } - - ingest_state.outcome.rejected_spans += 1; - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - "span_rejected", - format!( - "Rejected span {}:{} ({})", - span.trace_id, - span.span_id, - err.status_code().as_ref() - ), - ); - } - } - } - - Ok(()) - } - - /// Reconcile and insert one trace request batch. - async fn insert_trace_requests( - &self, - mut requests: RowInsertRequests, - is_trace_v1_model: bool, - ctx: QueryContextRef, - ) -> ServerResult { - if is_trace_v1_model { - self.reconcile_trace_column_types(&mut requests, &ctx) - .await?; - self.handle_trace_inserts(requests, ctx) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu) - } else { - self.handle_log_inserts(requests, ctx) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu) - } - } - - fn classify_trace_chunk_failure(status: StatusCode) -> ChunkFailureReaction { - match status { - StatusCode::InvalidArguments - | StatusCode::InvalidSyntax - | StatusCode::Unsupported - | StatusCode::TableNotFound - | StatusCode::TableColumnNotFound => ChunkFailureReaction::RetryPerSpan, - StatusCode::DatabaseNotFound => ChunkFailureReaction::DiscardChunk, - StatusCode::Cancelled | StatusCode::DeadlineExceeded => ChunkFailureReaction::Propagate, - StatusCode::StorageUnavailable - | StatusCode::RuntimeResourcesExhausted - | StatusCode::Internal - | StatusCode::RegionNotReady - | StatusCode::TableUnavailable - | StatusCode::RegionBusy => ChunkFailureReaction::Propagate, - _ => ChunkFailureReaction::DiscardChunk, - } - } - - fn should_propagate_trace_span_failure(status: StatusCode) -> bool { - matches!( - Self::classify_trace_chunk_failure(status), - ChunkFailureReaction::Propagate - ) - } - - fn add_trace_write_cost(outcome: &mut TraceIngestOutcome, cost: usize) { - outcome.write_cost += cost; - } - - fn push_trace_failure_message(messages: &mut Vec, label: &str, message: String) { - OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).inc(); - - if messages.len() < TRACE_FAILURE_MESSAGE_LIMIT { - messages.push(message); - } else if messages.len() == TRACE_FAILURE_MESSAGE_LIMIT { - tracing::debug!( - label, - limit = TRACE_FAILURE_MESSAGE_LIMIT, - "Trace ingest failure message limit reached; suppressing additional failure details" - ); - } - } - - fn finish_trace_failure_message( - accepted_spans: usize, - rejected_spans: usize, - messages: Vec, - ) -> Option { - if rejected_spans == 0 && messages.is_empty() { - return None; - } - - let mut summary = format!( - "Accepted {} spans, rejected {} spans", - accepted_spans, rejected_spans - ); - - if !messages.is_empty() { - summary.push_str(": "); - summary.push_str(&messages.join("; ")); - } - - Some(summary) - } - - /// Widen existing trace table columns to Float64 before request rewrite. - async fn alter_trace_table_columns_to_float64( - &self, - ctx: &QueryContextRef, - table_name: &str, - column_names: &[String], - ) -> ServerResult<()> { - let catalog_name = ctx.current_catalog().to_string(); - let schema_name = ctx.current_schema(); - let alter_expr = AlterTableExpr { - catalog_name: catalog_name.clone(), - schema_name: schema_name.clone(), - table_name: table_name.to_string(), - kind: Some(Kind::ModifyColumnTypes(ModifyColumnTypes { - modify_column_types: column_names - .iter() - .map(|column_name| ModifyColumnType { - column_name: column_name.clone(), - target_type: ColumnDataType::Float64 as i32, - target_type_extension: None, - }) - .collect(), - })), - }; - - if let Err(err) = self - .statement_executor - .alter_table_inner(alter_expr, ctx.clone()) - .await - { - let table = self - .catalog_manager - .table(&catalog_name, &schema_name, table_name, None) - .await - .map_err(servers::error::Error::from)?; - let alter_already_applied = table - .map(|table| { - let table_schema = table.schema(); - column_names.iter().all(|column_name| { - table_schema - .column_schema_by_name(column_name) - .and_then(|table_col| { - ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) - .ok() - .map(|wrapper| wrapper.datatype()) - }) - == Some(ColumnDataType::Float64) - }) - }) - .unwrap_or(false); - - if alter_already_applied { - return Ok(()); - } - - warn!( - table_name, - columns = ?column_names, - error = %err, - "failed to widen trace columns before insert" - ); - - return Err(wrap_trace_alter_failure(err)); - } - - Ok(()) - } - - /// Coerce request column types and values to match the existing table schema - /// for compatible type pairs. Existing table schema wins when present; - /// otherwise the full request batch decides a stable target type. - async fn reconcile_trace_column_types( - &self, - requests: &mut RowInsertRequests, - ctx: &QueryContextRef, - ) -> ServerResult<()> { - let catalog = ctx.current_catalog(); - let schema = ctx.current_schema(); - - for req in &mut requests.inserts { - let table = self - .catalog_manager - .table(catalog, &schema, &req.table_name, None) - .await?; - - let Some(rows) = req.rows.as_mut() else { - continue; - }; - - let table_schema = table.map(|table| table.schema()); - let mut pending_rewrites = Vec::new(); - let mut pending_alter_columns = Vec::new(); - - for (col_idx, col_schema) in rows.schema.iter().enumerate() { - let Some(current_type) = ColumnDataType::try_from(col_schema.datatype).ok() else { - continue; - }; - - let mut observed_types = Vec::new(); - push_observed_trace_type(&mut observed_types, current_type); - - // Scan the full request first so the final type decision is not affected - // by row order inside the batch. - for row in &rows.rows { - let Some(value) = row - .values - .get(col_idx) - .and_then(|value| value.value_data.as_ref()) - else { - continue; - }; - - let Some(value_type) = trace_value_datatype(value) else { - continue; - }; - push_observed_trace_type(&mut observed_types, value_type); - } - - let existing_type = table_schema - .as_ref() - .and_then(|schema| schema.column_schema_by_name(&col_schema.column_name)) - .and_then(|table_col| { - ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) - .ok() - .map(|wrapper| wrapper.datatype()) - }); - let fixed_type = trace_semconv_fixed_type(&col_schema.column_name); - - if !observed_types - .iter() - .copied() - .any(is_trace_reconcile_candidate_type) - && existing_type - .map(|datatype| !is_trace_reconcile_candidate_type(datatype)) - .unwrap_or(true) - && fixed_type.is_none() - { - continue; - } - - // Decide the final type once per column, then rewrite all affected cells - // together in one row pass below. - let Some(decision) = choose_trace_reconcile_decision( - &col_schema.column_name, - &observed_types, - existing_type, - ) - .map_err(|_| { - enrich_trace_reconcile_error( - &req.table_name, - &col_schema.column_name, - &observed_types, - existing_type, - fixed_type, - ) - })? - else { - continue; - }; - let target_type = decision.target_type(); - - if !decision.requires_alter() - && observed_types - .iter() - .all(|observed| *observed == target_type) - && col_schema.datatype == target_type as i32 - { - continue; - } - - if decision.requires_alter() - && !pending_alter_columns.contains(&col_schema.column_name) - { - pending_alter_columns.push(col_schema.column_name.clone()); - } - - pending_rewrites.push(PendingTraceColumnRewrite { - col_idx, - target_type, - column_name: col_schema.column_name.clone(), - }); - } - - if pending_rewrites.is_empty() { - continue; - } - - validate_trace_column_rewrites(&rows.rows, &pending_rewrites, &req.table_name)?; - - if !pending_alter_columns.is_empty() { - self.alter_trace_table_columns_to_float64( - ctx, - &req.table_name, - &pending_alter_columns, - ) - .await?; - } - - // Update schema metadata before mutating row values so both stay in sync. - for pending_rewrite in &pending_rewrites { - rows.schema[pending_rewrite.col_idx].datatype = pending_rewrite.target_type as i32; - } - - // Apply all pending column rewrites in one row pass. - for row in &mut rows.rows { - for pending_rewrite in &pending_rewrites { - let Some(value) = row.values.get_mut(pending_rewrite.col_idx) else { - continue; - }; - let Some(request_type) = - value.value_data.as_ref().and_then(trace_value_datatype) - else { - continue; - }; - if request_type == pending_rewrite.target_type { - continue; - } - - value.value_data = coerce_value_data( - &value.value_data, - pending_rewrite.target_type, - request_type, - ) - .map_err(|_| { - error::InvalidParameterSnafu { - reason: format!( - "failed to coerce trace column '{}' in table '{}' from {:?} to {:?}", - pending_rewrite.column_name, - req.table_name, - request_type, - pending_rewrite.target_type - ), - } - .build() - })?; - } - } - } - - Ok(()) - } -} - -fn chunk_owned(items: Vec, chunk_size: usize) -> Vec> { - if items.is_empty() { - return Vec::new(); - } - - if chunk_size == 0 { - return vec![items]; - } - - let mut chunks = Vec::with_capacity(items.len().div_ceil(chunk_size)); - let mut iter = items.into_iter(); - while iter.len() > 0 { - chunks.push(iter.by_ref().take(chunk_size).collect()); - } - chunks -} - -/// Preserve the original alter failure status so chunk retry behavior stays correct. -fn wrap_trace_alter_failure(err: E) -> servers::error::Error -where - E: ErrorExt + Send + Sync + 'static, -{ - error::ExecuteGrpcQuerySnafu.into_error(BoxedError::new(err)) -} - -/// Derives `trace.conventions` from the request's resource/scope `schema_url`s. -/// A single distinct non-empty value is concrete; multiple distinct values are -/// `mixed`; none is `unknown`. `schema_url` is row-level in OTLP, so the -/// table-level value is best-effort per the RFC conflict rule. -fn trace_conventions(request: &ExportTraceServiceRequest) -> String { - let mut seen: Option<&str> = None; - let mut mixed = false; - - for resource_spans in &request.resource_spans { - let urls = std::iter::once(resource_spans.schema_url.as_str()).chain( - resource_spans - .scope_spans - .iter() - .map(|s| s.schema_url.as_str()), - ); - for url in urls { - if url.is_empty() { - continue; - } - match seen { - None => seen = Some(url), - Some(prev) if prev == url => {} - Some(_) => { - mixed = true; - break; - } - } - } - if mixed { - break; - } - } - - if mixed { - SEMANTIC_VALUE_MIXED.to_string() - } else { - seen.map(str::to_string) - .unwrap_or_else(|| SEMANTIC_VALUE_UNKNOWN.to_string()) - } -} - -#[cfg(test)] -mod tests { - use common_error::ext::ErrorExt; - use common_error::status_code::StatusCode; - use servers::query_handler::TraceIngestOutcome; - - use super::{ChunkFailureReaction, Instance, chunk_owned, wrap_trace_alter_failure}; - use crate::metrics::OTLP_TRACES_FAILURE_COUNT; - - #[test] - fn test_chunk_owned() { - let chunks = chunk_owned(vec![1, 2, 3], 2); - assert_eq!(chunks.iter().map(Vec::len).collect::>(), vec![2, 1]); - - let chunks = chunk_owned(vec![1, 2, 3], 0); - assert_eq!(chunks.len(), 1); - assert_eq!(chunks[0].len(), 3); - - assert!(chunk_owned::(Vec::new(), 0).is_empty()); - } - - #[test] - fn test_classify_trace_chunk_failure() { - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::InvalidArguments), - ChunkFailureReaction::RetryPerSpan - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::InvalidSyntax), - ChunkFailureReaction::RetryPerSpan - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::Unsupported), - ChunkFailureReaction::RetryPerSpan - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::TableColumnNotFound), - ChunkFailureReaction::RetryPerSpan - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::TableNotFound), - ChunkFailureReaction::RetryPerSpan - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::DatabaseNotFound), - ChunkFailureReaction::DiscardChunk - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::DeadlineExceeded), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::Cancelled), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::StorageUnavailable), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::Internal), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::RegionNotReady), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::TableUnavailable), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::RegionBusy), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::RuntimeResourcesExhausted), - ChunkFailureReaction::Propagate - ); - } - - #[test] - fn test_classify_trace_span_failure() { - assert!(Instance::should_propagate_trace_span_failure( - StatusCode::DeadlineExceeded - )); - assert!(Instance::should_propagate_trace_span_failure( - StatusCode::StorageUnavailable - )); - assert!(!Instance::should_propagate_trace_span_failure( - StatusCode::InvalidArguments - )); - } - - #[test] - fn test_add_trace_write_cost() { - let mut outcome = TraceIngestOutcome::default(); - Instance::add_trace_write_cost(&mut outcome, 3); - Instance::add_trace_write_cost(&mut outcome, 5); - assert_eq!(outcome.write_cost, 8); - } - - #[test] - fn test_finish_trace_failure_message() { - let message = Instance::finish_trace_failure_message( - 3, - 2, - vec!["Rejected span trace:span (InvalidArguments)".to_string()], - ) - .unwrap(); - assert!(message.contains("Accepted 3 spans, rejected 2 spans")); - assert!(message.contains("Rejected span trace:span")); - - assert_eq!(Instance::finish_trace_failure_message(2, 0, vec![]), None); - } - - #[test] - fn test_finish_trace_failure_message_without_detail_messages() { - assert_eq!( - Instance::finish_trace_failure_message(0, 2, vec![]), - Some("Accepted 0 spans, rejected 2 spans".to_string()) - ); - } - - #[test] - fn test_push_trace_failure_message_increments_labeled_counter() { - let label = "retry_per_span_counter_test"; - let initial = OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(); - let mut messages = Vec::new(); - - Instance::push_trace_failure_message( - &mut messages, - label, - "Chunk fallback triggered by InvalidArguments".to_string(), - ); - - assert_eq!(messages.len(), 1); - assert_eq!( - OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(), - initial + 1 - ); - } - - #[test] - fn test_push_trace_failure_message_caps_recorded_messages() { - let label = "retry_per_span_limit_test"; - let mut messages = Vec::new(); - - for idx in 0..=4 { - Instance::push_trace_failure_message(&mut messages, label, format!("failure-{idx}")); - } - - assert_eq!(messages.len(), 4); - assert_eq!( - messages, - vec![ - "failure-0".to_string(), - "failure-1".to_string(), - "failure-2".to_string(), - "failure-3".to_string() - ] - ); - } - - #[test] - fn test_classify_trace_chunk_failure_defaults_to_discard() { - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::Unknown), - ChunkFailureReaction::DiscardChunk - ); - } - - #[test] - fn test_wrap_trace_alter_failure_preserves_status_code() { - let err = wrap_trace_alter_failure( - servers::error::TableNotFoundSnafu { - catalog: "greptime".to_string(), - schema: "public".to_string(), - table: "trace_type_missing".to_string(), - } - .build(), - ); - - assert_eq!(err.status_code(), StatusCode::TableNotFound); - } - - use opentelemetry_proto::tonic::trace::v1::{ResourceSpans, ScopeSpans}; - - use super::{ExportTraceServiceRequest, trace_conventions}; - - fn resource_spans(resource_url: &str, scope_urls: &[&str]) -> ResourceSpans { - ResourceSpans { - schema_url: resource_url.to_string(), - scope_spans: scope_urls - .iter() - .map(|u| ScopeSpans { - schema_url: u.to_string(), - ..Default::default() - }) - .collect(), - ..Default::default() - } - } - - #[test] - fn test_trace_conventions() { - let unknown = ExportTraceServiceRequest::default(); - assert_eq!(trace_conventions(&unknown), "unknown"); - - let url = "https://opentelemetry.io/schemas/1.27.0"; - let single = ExportTraceServiceRequest { - resource_spans: vec![resource_spans("", &[url, url])], - }; - assert_eq!(trace_conventions(&single), url); - - let resource_level = ExportTraceServiceRequest { - resource_spans: vec![resource_spans(url, &[""])], - }; - assert_eq!(trace_conventions(&resource_level), url); - - let conflicting = ExportTraceServiceRequest { - resource_spans: vec![resource_spans( - "", - &[url, "https://opentelemetry.io/schemas/1.30.0"], - )], - }; - assert_eq!(trace_conventions(&conflicting), "mixed"); - } -} diff --git a/src/frontend/src/instance/otlp/README.md b/src/frontend/src/instance/otlp/README.md new file mode 100644 index 00000000000..6380d56b7ac --- /dev/null +++ b/src/frontend/src/instance/otlp/README.md @@ -0,0 +1,182 @@ +# OTLP trace ingestion + +This directory owns the frontend orchestration for OTLP trace ingestion. It +starts after the HTTP server has decoded an OTLP protobuf request and ends after +the main trace table and its derived lookup tables have been written. + +The responsibility is split across crates: + +- `servers` decodes OTLP requests, parses spans, and converts spans into row + insert requests; +- this module chunks requests, reconciles v1 schemas, writes rows, handles + fallback, and accounts for accepted and rejected spans; +- `operator` creates or extends trace tables and performs the inserts. + +## End-to-end path + +```mermaid +flowchart TD + A["POST /v1/otlp/v1/traces
OTLP protobuf"] --> B["HTTP handler
decode, select table and pipeline"] + B --> C["Instance::traces
permission and interceptor"] + C --> D["Derive trace conventions
parse resource/scope groups"] + D --> E["Chunk spans within each group"] + E --> F{"Trace pipeline"} + + F -->|greptime_trace_v0| G["v0 conversion
fixed columns and JSON attributes"] + F -->|greptime_trace_v1| H["v1 conversion
flattened dynamic attributes"] + H --> I["Collect request-wide schema observations"] + I --> J["Resolve types, prevalidate rewrites,
apply DDL, and recheck table schema"] + + G --> K["Write main-table chunks"] + J --> K + K --> L{"Chunk result"} + L -->|success| M["Accept spans and collect
service/operation data"] + L -->|deterministic failure| N["Retry spans individually"] + L -->|discard| O["Count rejected spans"] + L -->|propagate| P["Return a request error"] + N --> M + N --> O + + M --> Q["Write <table>_services and
<table>_operations"] + O --> R["Build TraceIngestOutcome"] + Q --> R + R --> S["HTTP full success,
partial success, or failure"] +``` + +### 1. HTTP ingress + +The route is registered as `POST /v1/otlp/v1/traces` in +[`servers/src/http.rs`](../../../../servers/src/http.rs). The handler in +[`servers/src/http/otlp.rs`](../../../../servers/src/http/otlp.rs): + +1. accepts protobuf OTLP payloads and rejects JSON content types; +2. selects the main table, defaulting to `opentelemetry_traces`; +3. resolves the requested pipeline; and +4. calls `OpenTelemetryProtocolHandler::traces`. + +The trace converter currently accepts the internal `greptime_trace_v0` and +`greptime_trace_v1` pipelines. Their `PipelineWay` variants are defined in +[`pipeline/src/manager.rs`](../../../../pipeline/src/manager.rs). + +### 2. Parse and chunk + +[`Instance::traces`](../otlp.rs) checks the OTLP permission, runs the protocol +interceptor, derives the semantic-convention version, and parses the request. +The parser in +[`servers/src/otlp/trace/span.rs`](../../../../servers/src/otlp/trace/span.rs) +produces one `TraceSpanGroup` per resource/scope pair. + +`ingest_trace_spans` in [`trace_ingest.rs`](trace_ingest.rs) splits the spans in +each group into owned chunks. `trace_ingest_chunk_size` defaults to 128; setting +it to 0 disables splitting. The option is defined in +[`service_config/otlp.rs`](../../service_config/otlp.rs). + +### 3. Convert the selected trace model + +| Model | Conversion and storage shape | Frontend behavior | +| --- | --- | --- | +| v0 | [`trace/v0.rs`](../../../../servers/src/otlp/trace/v0.rs) writes a fixed span schema and keeps span, scope, and resource attributes in JSON columns. | Converts and writes one chunk at a time through the legacy log insertion path. | +| v1 | [`trace/v1.rs`](../../../../servers/src/otlp/trace/v1.rs) flattens OTLP attributes into dynamic columns and records the value types seen in each chunk. | Scans chunks one at a time to build a request-wide schema plan, then materializes and writes one chunk at a time. | + +Consequently, v1 chunks from one request can have different schemas: attribute +keys are data-dependent. The request-wide plan reconciles compatible column +observations across all chunks. The v0 main-table schema is comparatively fixed. + +The v1 converter intentionally uses `write_column_unchecked` from +[`row_writer.rs`](../../../../servers/src/row_writer.rs) to preserve the original +dynamic values instead of making the first value in a chunk decide the column +type. `MultiTableData` later pads sparse rows to the completed chunk schema. The +frontend schema planner is the validation and reconciliation boundary. + +## Trace v1 schema reconciliation + +The v1 path performs these steps before the normal chunk writes: + +1. Convert each chunk into a temporary `RowInsertRequest` and a + `TraceBatchSchema` that records observed scalar types, sparse row presence, + and binary logical types. Drop the dense request after collecting its + lightweight schema observations. +2. Mark a chunk for span-only fallback if one column contains both raw OTLP + bytes and arrays or key-value lists encoded as JSONB. Both use protobuf + `BinaryValue`, but a Greptime column cannot represent both logical schemas. + The chunk bypasses request-wide schema planning and is reconciled span by + span. +3. Aggregate the compatible observations into `TraceRequestSchema` and + compare them with the existing table schema and the fixed types in + [`trace_semconv.rs`](trace_semconv.rs). +4. Select a stable target type. Existing table types are authoritative when the + incoming values can be coerced. New columns use all remaining request + observations. + The supported existing-column widening is `Int64` to `Float64`. +5. Use `prepare_trace_column_rewrites` in + [`trace_types.rs`](trace_types.rs) to precompute every coercion without + mutating rows. A chunk whose coercion cannot be prepared is split into + single-span planning units; only spans that still fail are removed from the + request-wide plan, so invalid data cannot contribute to persistent DDL. +6. Create the table or add missing columns with `ensure_trace_table_on_demand`, + and widen planned numeric columns. The implementation is in + [`operator/src/insert.rs`](../../../../operator/src/insert.rs). +7. Read and plan against the table schema again after DDL to account for + concurrent creates or alters. Then rematerialize, rewrite, and write one + chunk at a time. If the schema does not converge, affected chunks use the + per-request reconciliation path. + +This ordering keeps conversion failures atomic: rows are not partly rewritten, +and columns from an invalid chunk are not added to the table before that chunk +falls back. + +## Writes, fallback, and accounting + +Schema planning and request reconciliation happen before a main-table write. +`classify_trace_prewrite_failure` in [`trace_ingest.rs`](trace_ingest.rs) +determines the next action when that preparation fails: + +| Reaction | Behavior | +| --- | --- | +| `RetryPerSpan` | Prepare and write each span independently. Successful spans are accepted; deterministic per-span preparation failures are rejected. A v1 chunk with incompatible raw-binary and JSONB values starts here without attempting a chunk write. | +| `DiscardChunk` | Do not retry the chunk at span granularity. Count every span in the chunk as rejected. | +| `Propagate` | Return retryable or unclassified failures instead of reporting a partial success. Retry hints take precedence over status-code classification. | + +Span-specific invalid values are isolated before request-wide schema planning. +Therefore, a deterministic error that escapes request-wide preparation rejects +all remaining spans instead of repeating the same global failure per span. + +Once a main-table write is dispatched, any returned error is propagated. The +distributed insert may have committed writes on some peers before another peer +failed, so retrying or rejecting the whole chunk would risk duplicate data or +incorrect accounting. + +Accounting follows the main-table write: + +- a span is accepted only after its main-table insert is confirmed; +- only accepted spans contribute service and operation rows; +- auxiliary rows are deduplicated and written afterwards to + `_services` and `_operations`; +- auxiliary rows accumulated by earlier successful chunks are still flushed + before a later main-table error is returned; +- an auxiliary-table failure adds failure detail but does not change the + accepted or rejected span counts; and +- failure details are bounded before they are folded into `TraceIngestOutcome`. + +Finally, the HTTP handler returns: + +- full success when there are no rejected spans or failure details; +- OTLP partial success when spans were rejected or auxiliary/fallback details + were recorded; or +- HTTP `400 Bad Request` when no span was accepted and at least one was + rejected. + +A propagated error bypasses that outcome and fails the request. + +## Code map + +| Area | Entry point | +| --- | --- | +| HTTP route and response mapping | [`servers/src/http/otlp.rs`](../../../../servers/src/http/otlp.rs) | +| Frontend protocol handler | [`instance/otlp.rs`](../otlp.rs) | +| Chunking, v1 planning, writes, fallback, accounting | [`trace_ingest.rs`](trace_ingest.rs) | +| Reconciliation decisions and atomic value rewrites | [`trace_types.rs`](trace_types.rs) | +| Semantic-convention fixed types | [`trace_semconv.rs`](trace_semconv.rs) | +| OTLP span parsing | [`servers/src/otlp/trace/span.rs`](../../../../servers/src/otlp/trace/span.rs) | +| v0 and v1 row conversion | [`servers/src/otlp/trace`](../../../../servers/src/otlp/trace.rs) | +| Trace-aware table creation and inserts | [`operator/src/insert.rs`](../../../../operator/src/insert.rs) | diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs new file mode 100644 index 00000000000..0866189b571 --- /dev/null +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -0,0 +1,1965 @@ +// 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 std::collections::{BTreeMap, HashMap, HashSet}; +use std::sync::Arc; + +use api::helper::ColumnDataTypeWrapper; +use api::v1::alter_table_expr::Kind; +use api::v1::{ + AlterTableExpr, ColumnDataType, ColumnSchema, ModifyColumnType, ModifyColumnTypes, + RowInsertRequest, RowInsertRequests, Rows, Value, +}; +use client::Output; +use common_error::ext::{BoxedError, ErrorExt}; +use common_error::status_code::StatusCode; +use common_telemetry::warn; +use datatypes::prelude::ConcreteDataType; +use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; +use pipeline::{GreptimePipelineParams, PipelineWay}; +use servers::error::{self, Result as ServerResult}; +use servers::otlp; +use servers::otlp::coerce::{coerce_value_data, trace_value_datatype}; +use servers::otlp::trace::TraceAuxData; +use servers::otlp::trace::span::{TraceSpan, TraceSpanGroup}; +use servers::otlp::trace::v1::{TraceBatchSchema, TraceBinaryType, TraceRetryColumns}; +use servers::query_handler::{PipelineHandlerRef, TraceIngestOutcome}; +use session::context::QueryContextRef; +use snafu::{IntoError, ResultExt}; +use table::requests::{ + SEMANTIC_PIPELINE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SEMANTIC_TRACE_CONVENTIONS, + SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, SIGNAL_TYPE_TRACE, SOURCE_OPENTELEMETRY, + TABLE_DATA_MODEL_TRACE_V1, +}; + +use crate::instance::Instance; +use crate::instance::otlp::trace_semconv::trace_semconv_fixed_type; +use crate::instance::otlp::trace_types::{ + PendingTraceColumnRewrite, PreparedTraceColumnRewrites, TraceColumnRewriteError, + choose_trace_reconcile_decision, enrich_trace_reconcile_error, + is_trace_reconcile_candidate_type, prepare_trace_column_rewrites, push_observed_trace_type, +}; +use crate::metrics::{OTLP_TRACES_FAILURE_COUNT, OTLP_TRACES_ROWS}; + +const TRACE_FAILURE_MESSAGE_LIMIT: usize = 4; + +/// Determines how trace ingestion responds to a failure before a write is dispatched. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ChunkFailureReaction { + RetryPerSpan, + DiscardChunk, + Propagate, +} + +impl ChunkFailureReaction { + fn as_metric_label(self) -> &'static str { + match self { + Self::RetryPerSpan => "retry_per_span", + Self::DiscardChunk => "discard_chunk", + Self::Propagate => "propagate_failure", + } + } +} + +/// Shared dependencies and request metadata used while ingesting trace chunks. +struct TraceChunkIngestContext<'a> { + pipeline_handler: PipelineHandlerRef, + pipeline: &'a PipelineWay, + pipeline_params: &'a GreptimePipelineParams, + table_name: &'a str, + is_trace_v1_model: bool, +} + +/// Accumulates trace outcomes, auxiliary rows, and bounded failure details. +struct TraceIngestState { + aux_data: TraceAuxData, + outcome: TraceIngestOutcome, + failure_messages: Vec, +} + +/// How a v1 chunk should be reconciled when it is written. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TraceChunkSchemaState { + Prepared, + ReconcilePerChunk, + SpanFallbackOnly, +} + +impl TraceChunkSchemaState { + fn mark_for_reconcile(&mut self) { + if matches!(self, Self::Prepared) { + *self = Self::ReconcilePerChunk; + } + } +} + +/// Sparse representation retained for a v1 trace chunk. +struct TraceChunkRetry { + table_name: String, + schema: Vec, + rows: Vec, +} + +/// The projected retry values and identifying metadata for one span. +struct TraceRetryRow { + span_metadata: TraceSpanMetadata, + values: Vec, +} + +/// One retry cell and the logical type to restore before insertion. +struct TraceRetryValue { + column_index: usize, + value: Value, + schema_type: Option, +} + +/// The scalar or binary logical type recorded for a retry value. +enum TraceRetrySchemaType { + Scalar(ColumnDataType), + Binary(TraceBinaryType), +} + +/// Span identity and operation fields retained after request conversion. +struct TraceSpanMetadata { + trace_id: String, + span_id: String, + operation: Option<(String, String, String)>, +} + +impl From<&TraceSpan> for TraceSpanMetadata { + fn from(span: &TraceSpan) -> Self { + Self { + trace_id: span.trace_id.clone(), + span_id: span.span_id.clone(), + operation: span.service_name.as_ref().map(|service_name| { + ( + service_name.clone(), + span.span_name.clone(), + span.span_kind.clone(), + ) + }), + } + } +} + +impl TraceSpanMetadata { + fn add_to_aux_data(self, aux_data: &mut TraceAuxData) { + if let Some((service_name, span_name, span_kind)) = self.operation { + aux_data.services.insert(service_name.clone()); + aux_data + .operations + .insert((service_name, span_name, span_kind)); + } + } +} + +impl TraceChunkRetry { + fn try_new( + table_name: String, + rows: Rows, + span_metadata: Vec, + retry_columns: TraceRetryColumns, + ) -> ServerResult { + let Rows { schema, rows } = rows; + let row_count = rows.len(); + if row_count != span_metadata.len() { + return error::InternalSnafu { + err_msg: format!( + "trace row count {} does not match span metadata count {}", + row_count, + span_metadata.len() + ), + } + .fail(); + } + + let retry_columns_by_index = schema + .iter() + .map(|column| retry_columns.get(&column.column_name)) + .collect::>(); + let mut fixed_columns = Vec::with_capacity(schema.len()); + let mut sparse_columns = vec![Vec::new(); rows.len()]; + for (column_index, retry_column) in retry_columns_by_index.iter().enumerate() { + let Some(retry_column) = retry_column else { + fixed_columns.push(column_index); + continue; + }; + for present_row in &retry_column.present_rows { + if *present_row < sparse_columns.len() { + sparse_columns[*present_row].push(column_index); + } + } + } + + let mut retry_rows = Vec::with_capacity(rows.len()); + for (row_index, (mut row, span_metadata)) in rows.into_iter().zip(span_metadata).enumerate() + { + if row.values.len() > schema.len() { + return error::InternalSnafu { + err_msg: format!( + "trace row column count {} exceeds schema column count {}", + row.values.len(), + schema.len() + ), + } + .fail(); + } + let sparse_row = &mut sparse_columns[row_index]; + let mut selected_columns = Vec::with_capacity(fixed_columns.len() + sparse_row.len()); + selected_columns.extend_from_slice(&fixed_columns); + selected_columns.append(sparse_row); + selected_columns.sort_unstable(); + selected_columns.dedup(); + + let mut values = Vec::with_capacity(selected_columns.len()); + for column_index in selected_columns { + let column = &schema[column_index]; + let value = row + .values + .get_mut(column_index) + .map(std::mem::take) + .unwrap_or_default(); + let value_type = value.value_data.as_ref().and_then(trace_value_datatype); + let binary_type = if value_type == Some(ColumnDataType::Binary) { + retry_columns_by_index[column_index].and_then(|retry_column| { + retry_column + .binary_types + .binary_search_by_key(&row_index, |(row_index, _)| *row_index) + .ok() + .map(|index| retry_column.binary_types[index].1) + }) + } else { + None + }; + let schema_type = binary_type.map(TraceRetrySchemaType::Binary).or_else(|| { + value_type + .filter(|datatype| { + is_trace_reconcile_candidate_type(*datatype) + && ColumnDataType::try_from(column.datatype).ok() != Some(*datatype) + }) + .map(TraceRetrySchemaType::Scalar) + }); + values.push(TraceRetryValue { + column_index, + value, + schema_type, + }); + } + retry_rows.push(TraceRetryRow { + span_metadata, + values, + }); + } + + Ok(Self { + table_name, + schema, + rows: retry_rows, + }) + } + + fn to_request(&self) -> ServerResult { + let mut rows = Vec::with_capacity(self.rows.len()); + for retry_row in &self.rows { + let mut values = vec![Value::default(); self.schema.len()]; + for retry_value in &retry_row.values { + let Some(value) = values.get_mut(retry_value.column_index) else { + return error::InternalSnafu { + err_msg: format!( + "trace retry column index {} is out of bounds", + retry_value.column_index + ), + } + .fail(); + }; + *value = retry_value.value.clone(); + } + rows.push(api::v1::Row { values }); + } + + Ok(RowInsertRequest { + table_name: self.table_name.clone(), + rows: Some(Rows { + schema: self.schema.clone(), + rows, + }), + }) + } + + fn into_single_span_chunks(self) -> ServerResult> { + let Self { + table_name, + schema, + rows, + } = self; + rows.into_iter() + .map(|row| { + let (span_metadata, rows_for_insert) = row.into_rows(&schema)?; + Self::try_new( + table_name.clone(), + rows_for_insert, + vec![span_metadata], + TraceRetryColumns::default(), + ) + }) + .collect() + } + + fn add_to_aux_data(self, aux_data: &mut TraceAuxData) { + for row in self.rows { + row.span_metadata.add_to_aux_data(aux_data); + } + } +} + +impl TraceRetryRow { + fn into_rows(self, schema: &[ColumnSchema]) -> ServerResult<(TraceSpanMetadata, Rows)> { + let mut projected_schema = Vec::with_capacity(self.values.len()); + let mut projected_values = Vec::with_capacity(self.values.len()); + for retry_value in self.values { + let Some(mut column) = schema.get(retry_value.column_index).cloned() else { + return error::InternalSnafu { + err_msg: format!( + "trace fallback column index {} is out of bounds", + retry_value.column_index + ), + } + .fail(); + }; + match retry_value.schema_type { + Some(TraceRetrySchemaType::Scalar(datatype)) => { + column.datatype = datatype as i32; + column.datatype_extension = None; + } + Some(TraceRetrySchemaType::Binary(binary_type)) => { + binary_type.apply_to_schema(&mut column); + } + None => {} + } + projected_schema.push(column); + projected_values.push(retry_value.value); + } + + Ok(( + self.span_metadata, + Rows { + schema: projected_schema, + rows: vec![api::v1::Row { + values: projected_values, + }], + }, + )) + } +} + +/// Request-wide schema observations collected across all v1 trace chunks. +#[derive(Default)] +struct TraceRequestSchema { + columns: Vec, + column_indexes: HashMap, +} + +/// Per-column observations and the resolved target type for a trace request. +struct TraceColumnRequestSchema { + schema: ColumnSchema, + batches: BTreeMap, + target_type: Option, +} + +/// The schema and value types observed for one column in one chunk. +struct TraceBatchColumnObservation { + value_types: Vec, + schema_type: ColumnDataType, + concrete_type: ConcreteDataType, +} + +/// Maps incompatible columns to the chunk indexes that must use fallback. +type TraceSchemaExclusions = HashMap>; + +/// Table columns to add or widen before inserting the prepared chunks. +struct TraceTablePreAlter { + ensure_columns: Vec, + modify_float64_columns: Vec, + alter_existing: bool, +} + +/// Either a ready table-alter plan or the observations that cannot be unified. +enum TraceRequestSchemaPlan { + Ready(TraceTablePreAlter), + IncompatibleObservations(TraceSchemaExclusions), +} + +impl TraceTablePreAlter { + fn requires_ddl(&self) -> bool { + !self.ensure_columns.is_empty() || !self.modify_float64_columns.is_empty() + } +} + +impl TraceRequestSchema { + fn observe_batch_schema( + &mut self, + batch_index: usize, + schema: &[ColumnSchema], + batch_schema: &TraceBatchSchema, + ) { + for column in schema { + if batch_schema.has_incompatible_logical_types_for(&column.column_name) { + continue; + } + let observed_types = batch_schema.value_types(&column.column_name); + self.observe_trace_column(batch_index, column, None); + if let Some(observed_types) = observed_types { + for value_type in observed_types { + self.observe_trace_column(batch_index, column, Some(*value_type)); + } + } + } + } + + fn observe_retry_chunk( + &mut self, + batch_index: usize, + chunk: &TraceChunkRetry, + ) -> ServerResult<()> { + // Split chunks already carry row-projected schemas; untouched chunks keep + // their original batch schema while observations stay sparse. + for row in &chunk.rows { + for retry_value in &row.values { + let Some(column) = chunk.schema.get(retry_value.column_index) else { + return error::InternalSnafu { + err_msg: format!( + "trace retry column index {} is out of bounds", + retry_value.column_index + ), + } + .fail(); + }; + let value_type = retry_value + .value + .value_data + .as_ref() + .and_then(trace_value_datatype); + self.observe_trace_column(batch_index, column, value_type); + } + } + Ok(()) + } + + fn observe_trace_column( + &mut self, + batch_index: usize, + schema: &ColumnSchema, + value_type: Option, + ) { + let column_schema = self.observe_column(schema); + if let Ok(current_type) = ColumnDataType::try_from(schema.datatype) { + let observation = column_schema.batches.entry(batch_index).or_insert_with(|| { + let wrapper = + ColumnDataTypeWrapper::new(current_type, schema.datatype_extension.clone()); + TraceBatchColumnObservation { + value_types: Vec::new(), + schema_type: current_type, + concrete_type: ConcreteDataType::from(wrapper), + } + }); + push_observed_trace_type(&mut observation.value_types, current_type); + } + if let Some(value_type) = value_type { + column_schema.observe_type(batch_index, value_type); + } + } + + fn incompatible_schema_observations( + &self, + table_schema: Option<&datatypes::schema::Schema>, + ) -> TraceSchemaExclusions { + let mut exclusions = HashMap::new(); + for column in &self.columns { + let existing_type = table_schema + .and_then(|schema| schema.column_schema_by_name(&column.schema.column_name)) + .and_then(|table_col| { + ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) + .ok() + .map(|wrapper| { + let datatype = wrapper.datatype(); + (datatype, ConcreteDataType::from(wrapper)) + }) + }); + let incompatible_batches = column.incompatible_schema_batches(existing_type); + if !incompatible_batches.is_empty() { + exclusions.insert(column.schema.column_name.clone(), incompatible_batches); + } + } + exclusions + } + + fn resolve_table_schema( + &mut self, + table_schema: Option<&datatypes::schema::Schema>, + ) -> TraceRequestSchemaPlan { + let mut pre_alter = TraceTablePreAlter { + ensure_columns: Vec::new(), + modify_float64_columns: Vec::new(), + alter_existing: table_schema.is_some(), + }; + + for column in &mut self.columns { + let Some(current_type) = ColumnDataType::try_from(column.schema.datatype).ok() else { + continue; + }; + let observed_types = column.observed_types(); + + let existing_type = table_schema + .and_then(|schema| schema.column_schema_by_name(&column.schema.column_name)) + .and_then(|table_col| { + ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) + .ok() + .map(|wrapper| wrapper.datatype()) + }); + let fixed_type = trace_semconv_fixed_type(&column.schema.column_name); + + let needs_reconcile = observed_types + .iter() + .copied() + .any(is_trace_reconcile_candidate_type) + || existing_type + .map(is_trace_reconcile_candidate_type) + .unwrap_or(false) + || fixed_type.is_some(); + + let target_type = if needs_reconcile { + let decision = match choose_trace_reconcile_decision( + &column.schema.column_name, + &observed_types, + existing_type, + ) { + Ok(decision) => decision, + Err(_) => { + return TraceRequestSchemaPlan::IncompatibleObservations(HashMap::from([ + ( + column.schema.column_name.clone(), + column.batches.keys().copied().collect(), + ), + ])); + } + }; + decision + .map(|decision| { + if decision.requires_alter() { + pre_alter + .modify_float64_columns + .push(column.schema.column_name.clone()); + } + decision.target_type() + }) + .unwrap_or(current_type) + } else { + current_type + }; + + column.target_type = Some(target_type); + + if table_schema + .map(|schema| { + schema + .column_schema_by_name(&column.schema.column_name) + .is_none() + }) + .unwrap_or(false) + { + pre_alter.ensure_columns.push(column.target_schema()); + } + } + + if table_schema.is_none() { + pre_alter.ensure_columns = self + .columns + .iter() + .map(TraceColumnRequestSchema::target_schema) + .collect(); + } + + TraceRequestSchemaPlan::Ready(pre_alter) + } + + fn prepare_request_rewrite( + &self, + batch_index: usize, + request: &RowInsertRequest, + ) -> Result, TraceColumnRewriteError> { + let Some(rows) = request.rows.as_ref() else { + return Ok(None); + }; + let pending_rewrites = self.pending_batch_rewrites(batch_index, rows); + if pending_rewrites.is_empty() { + return Ok(None); + } + prepare_trace_column_rewrites(&rows.rows, pending_rewrites, &request.table_name).map(Some) + } + + fn apply_request_rewrite( + request: &mut RowInsertRequest, + prepared: PreparedTraceColumnRewrites, + ) { + if let Some(rows) = request.rows.as_mut() { + prepared.apply(rows); + } + } + + fn pending_batch_rewrites( + &self, + batch_index: usize, + rows: &api::v1::Rows, + ) -> Vec { + let mut pending_rewrites = Vec::new(); + for (col_idx, col_schema) in rows.schema.iter().enumerate() { + let Some(global_idx) = self.column_indexes.get(&col_schema.column_name).copied() else { + continue; + }; + let global_column = &self.columns[global_idx]; + let Some(target_type) = global_column.target_type() else { + continue; + }; + if !global_column + .batches + .get(&batch_index) + .is_some_and(|observation| { + observation + .value_types + .iter() + .any(|datatype| *datatype != target_type) + }) + { + continue; + } + pending_rewrites.push(PendingTraceColumnRewrite { + col_idx, + target_type, + column_name: global_column.schema.column_name.clone(), + }); + } + + pending_rewrites + } + + #[cfg(test)] + fn remove_batches(&mut self, batch_indexes: &HashSet) { + for column in &mut self.columns { + column + .batches + .retain(|batch_index, _| !batch_indexes.contains(batch_index)); + column.target_type = None; + } + self.columns.retain(|column| !column.batches.is_empty()); + + self.column_indexes.clear(); + for (index, column) in self.columns.iter().enumerate() { + self.column_indexes + .insert(column.schema.column_name.clone(), index); + } + } + + #[cfg(test)] + fn resolved_target_types(&self) -> Vec> { + self.columns + .iter() + .map(|column| column.target_type) + .collect() + } + + fn observe_column(&mut self, schema: &ColumnSchema) -> &mut TraceColumnRequestSchema { + if let Some(index) = self.column_indexes.get(&schema.column_name).copied() { + return &mut self.columns[index]; + } + + let index = self.columns.len(); + self.columns.push(TraceColumnRequestSchema { + schema: schema.clone(), + batches: BTreeMap::new(), + target_type: None, + }); + self.column_indexes + .insert(schema.column_name.clone(), index); + &mut self.columns[index] + } +} + +impl TraceColumnRequestSchema { + fn observe_type(&mut self, batch_index: usize, datatype: ColumnDataType) { + if let Some(observation) = self.batches.get_mut(&batch_index) { + push_observed_trace_type(&mut observation.value_types, datatype); + } + } + + fn observed_types(&self) -> Vec { + let mut observed_types = Vec::new(); + for datatype in self + .batches + .values() + .flat_map(|observation| &observation.value_types) + { + push_observed_trace_type(&mut observed_types, *datatype); + } + observed_types + } + + fn incompatible_schema_batches( + &self, + existing_type: Option<(ColumnDataType, ConcreteDataType)>, + ) -> HashSet { + let mut incompatible_batches = HashSet::new(); + if let Some((existing_datatype, existing_concrete_type)) = existing_type { + for (batch_index, observation) in &self.batches { + if trace_logical_types_incompatible( + observation.schema_type, + &observation.concrete_type, + existing_datatype, + &existing_concrete_type, + ) { + incompatible_batches.insert(*batch_index); + } + } + return incompatible_batches; + } + + let mut reconcile_candidates = Vec::<&TraceBatchColumnObservation>::new(); + for observation in self + .batches + .values() + .filter(|observation| is_trace_reconcile_candidate_type(observation.schema_type)) + { + if !reconcile_candidates.iter().any(|candidate| { + candidate.schema_type == observation.schema_type + && candidate.concrete_type == observation.concrete_type + }) { + reconcile_candidates.push(observation); + } + } + if !reconcile_candidates.is_empty() { + for (batch_index, observation) in &self.batches { + if reconcile_candidates.iter().any(|candidate| { + trace_logical_types_incompatible( + observation.schema_type, + &observation.concrete_type, + candidate.schema_type, + &candidate.concrete_type, + ) + }) { + incompatible_batches.insert(*batch_index); + } + } + return incompatible_batches; + } + + let mut schemas = Vec::<(ColumnDataType, ConcreteDataType, bool)>::new(); + for observation in self.batches.values() { + if schemas.iter().any(|(observed_datatype, observed_type, _)| { + *observed_datatype == observation.schema_type + && observed_type == &observation.concrete_type + }) { + continue; + } + + let mut conflicts = false; + for (observed_datatype, observed_type, observed_conflicts) in &mut schemas { + if trace_logical_types_incompatible( + *observed_datatype, + observed_type, + observation.schema_type, + &observation.concrete_type, + ) { + *observed_conflicts = true; + conflicts = true; + } + } + schemas.push(( + observation.schema_type, + observation.concrete_type.clone(), + conflicts, + )); + } + + for (batch_index, observation) in &self.batches { + if schemas + .iter() + .any(|(observed_datatype, observed_type, conflict)| { + *observed_datatype == observation.schema_type + && observed_type == &observation.concrete_type + && *conflict + }) + { + incompatible_batches.insert(*batch_index); + } + } + + incompatible_batches + } + + fn target_type(&self) -> Option { + self.target_type + .or_else(|| ColumnDataType::try_from(self.schema.datatype).ok()) + } + + fn target_schema(&self) -> ColumnSchema { + let mut schema = self.schema.clone(); + if let Some(target_type) = self.target_type() { + schema.datatype = target_type as i32; + } + schema + } +} + +impl Instance { + /// Ingest OTLP trace spans with chunk-level writes and span-level fallback on + /// deterministic chunk failures. + #[allow(clippy::too_many_arguments)] + pub(super) async fn ingest_trace_spans( + &self, + pipeline_handler: PipelineHandlerRef, + pipeline: &PipelineWay, + pipeline_params: &GreptimePipelineParams, + table_name: String, + groups: Vec, + conventions: &str, + ctx: QueryContextRef, + ) -> ServerResult { + let is_trace_v1_model = matches!(pipeline, PipelineWay::OtlpTraceDirectV1); + + // Only the main span table gets the identity; the derived `_services` / + // `_operations` lookup tables keep the unstamped `ctx`. + let main_ctx = { + let mut c = (*ctx).clone(); + c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_TRACE); + c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY); + if is_trace_v1_model { + c.set_extension(SEMANTIC_PIPELINE, TABLE_DATA_MODEL_TRACE_V1); + c.set_extension(SEMANTIC_TRACE_CONVENTIONS, conventions); + } + Arc::new(c) + }; + + let ingest_ctx = TraceChunkIngestContext { + pipeline_handler, + pipeline, + pipeline_params, + table_name: &table_name, + is_trace_v1_model, + }; + let mut ingest_state = TraceIngestState { + aux_data: TraceAuxData::default(), + outcome: TraceIngestOutcome::default(), + failure_messages: Vec::new(), + }; + + let main_result: ServerResult<()> = async { + if is_trace_v1_model { + let mut request_schema = TraceRequestSchema::default(); + let mut chunks = Vec::new(); + let mut chunk_states = Vec::new(); + // Convert each chunk once, then retain only its sparse values while + // request-wide schema planning runs. + for chunk in groups + .into_iter() + .flat_map(|group| chunk_owned(group.spans, self.trace_ingest_chunk_size)) + { + let batch_index = chunks.len(); + let span_metadata = chunk.iter().map(TraceSpanMetadata::from).collect(); + let (table_data, batch_schema) = + otlp::trace::v1::v1_to_main_table_data_with_schema(chunk)?; + let schema_state = if batch_schema.has_incompatible_logical_types() { + TraceChunkSchemaState::SpanFallbackOnly + } else { + TraceChunkSchemaState::Prepared + }; + let (schema, rows) = table_data.into_schema_and_rows(); + if schema_state == TraceChunkSchemaState::Prepared { + request_schema.observe_batch_schema(batch_index, &schema, &batch_schema); + } + chunks.push(TraceChunkRetry::try_new( + ingest_ctx.table_name.to_string(), + Rows { schema, rows }, + span_metadata, + batch_schema.into_retry_columns(), + )?); + chunk_states.push(schema_state); + } + + if !chunks.is_empty() { + self.ingest_trace_v1_prepared_chunks( + &ingest_ctx, + request_schema, + chunks, + chunk_states, + main_ctx.clone(), + &mut ingest_state, + ) + .await?; + } + } else { + for group in groups { + let chunks = chunk_owned(group.spans, self.trace_ingest_chunk_size); + for chunk in chunks { + self.ingest_trace_chunk( + &ingest_ctx, + chunk, + main_ctx.clone(), + &mut ingest_state, + ) + .await?; + } + } + } + + Ok(()) + } + .await; + + OTLP_TRACES_ROWS.inc_by(ingest_state.outcome.accepted_spans as u64); + + if !ingest_state.aux_data.is_empty() { + // Auxiliary trace tables are derived from spans whose main-table + // writes are already confirmed, so they never create new accepted + // spans and they do not affect rejected span counts. + let aux_requests = otlp::trace::to_grpc_insert_requests_for_aux_tables( + std::mem::take(&mut ingest_state.aux_data), + ingest_ctx.pipeline, + ingest_ctx.table_name, + ); + + match aux_requests { + Ok((aux_requests, _)) if !aux_requests.inserts.is_empty() => { + match self + .insert_trace_requests(aux_requests, ingest_ctx.is_trace_v1_model, ctx) + .await + { + Ok(output) => { + Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); + } + Err(err) => { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + "aux_table_update_failed", + format!( + "Auxiliary trace tables were not fully updated ({})", + err.status_code().as_ref() + ), + ); + } + } + } + // Preserve the existing conversion-error behavior after a + // successful main write, but never mask an earlier main error. + Err(err) if main_result.is_ok() => return Err(err), + Err(err) => Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + "aux_table_update_failed", + format!( + "Auxiliary trace tables were not fully updated ({})", + err.status_code().as_ref() + ), + ), + Ok(_) => {} + } + } + + main_result?; + + ingest_state.outcome.error_message = Self::finish_trace_failure_message( + ingest_state.outcome.accepted_spans, + ingest_state.outcome.rejected_spans, + ingest_state.failure_messages, + ); + + Ok(ingest_state.outcome) + } + + /// Ingest one owned trace chunk so successful spans can be moved into the + /// accepted set without extra cloning. + async fn ingest_trace_chunk( + &self, + ingest_ctx: &TraceChunkIngestContext<'_>, + chunk: Vec, + ctx: QueryContextRef, + ingest_state: &mut TraceIngestState, + ) -> ServerResult<()> { + // Try the fast path first so healthy batches keep their original + // throughput and write amplification stays low. + let (requests, chunk_rows) = otlp::trace::to_grpc_insert_requests_from_spans( + &chunk, + ingest_ctx.pipeline, + ingest_ctx.pipeline_params, + ingest_ctx.table_name, + &ctx, + ingest_ctx.pipeline_handler.clone(), + )?; + + let output = match self + .insert_trace_requests(requests, ingest_ctx.is_trace_v1_model, ctx) + .await + { + Ok(output) => output, + Err(err) => { + // A distributed insert can partially commit before returning an + // error, so retrying or rejecting the whole chunk is unsafe. + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::Propagate.as_metric_label(), + format!( + "Propagating chunk write failure ({})", + err.status_code().as_ref() + ), + ); + return Err(err); + } + }; + + Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); + ingest_state.outcome.accepted_spans += chunk_rows; + for span in &chunk { + ingest_state.aux_data.observe_span(span); + } + + Ok(()) + } + + async fn ingest_trace_v1_prepared_chunks( + &self, + ingest_ctx: &TraceChunkIngestContext<'_>, + mut request_schema: TraceRequestSchema, + mut chunks: Vec, + mut chunk_states: Vec, + ctx: QueryContextRef, + ingest_state: &mut TraceIngestState, + ) -> ServerResult<()> { + if let Err(err) = self + .prepare_trace_v1_request_schema( + ingest_ctx.table_name, + &mut request_schema, + &mut chunks, + &mut chunk_states, + &ctx, + ) + .await + { + if matches!( + Self::classify_trace_prewrite_failure(err.status_code(), err.is_retryable()), + ChunkFailureReaction::Propagate + ) { + return Err(err); + } + + // Span-specific conversion failures were already isolated during + // prevalidation. An error escaping request-wide schema preparation + // applies to every remaining span, so repeating it per span cannot + // recover valid data. + let span_count = chunks.iter().map(|chunk| chunk.rows.len()).sum::(); + ingest_state.outcome.rejected_spans += span_count; + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::DiscardChunk.as_metric_label(), + format!( + "Discarded {} spans after pre-write request failure ({})", + span_count, + err.status_code().as_ref() + ), + ); + return Ok(()); + } + + for (batch_index, (retry, schema_state)) in chunks.into_iter().zip(chunk_states).enumerate() + { + self.ingest_trace_v1_chunk( + &request_schema, + batch_index, + retry, + schema_state, + ctx.clone(), + ingest_state, + ) + .await?; + } + + Ok(()) + } + + async fn ingest_trace_v1_chunk( + &self, + request_schema: &TraceRequestSchema, + batch_index: usize, + retry: TraceChunkRetry, + mut schema_state: TraceChunkSchemaState, + ctx: QueryContextRef, + ingest_state: &mut TraceIngestState, + ) -> ServerResult<()> { + if schema_state == TraceChunkSchemaState::SpanFallbackOnly { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::RetryPerSpan.as_metric_label(), + "Chunk fallback triggered by incompatible binary and JSON values".to_string(), + ); + return self + .ingest_trace_v1_rows_span_by_span(retry, ctx, ingest_state) + .await; + } + + let mut request = retry.to_request()?; + let _ = Self::prepare_trace_v1_chunk_rewrite( + request_schema, + batch_index, + &mut request, + &mut schema_state, + )?; + let span_count = retry.rows.len(); + let mut requests = RowInsertRequests { + inserts: vec![request], + }; + if schema_state == TraceChunkSchemaState::ReconcilePerChunk + && let Err(err) = self.reconcile_trace_column_types(&mut requests, &ctx).await + { + match Self::classify_trace_prewrite_failure(err.status_code(), err.is_retryable()) { + ChunkFailureReaction::RetryPerSpan => { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::RetryPerSpan.as_metric_label(), + format!("Chunk fallback triggered by {}", err.status_code().as_ref()), + ); + return self + .ingest_trace_v1_rows_span_by_span(retry, ctx, ingest_state) + .await; + } + ChunkFailureReaction::DiscardChunk => { + ingest_state.outcome.rejected_spans += span_count; + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::DiscardChunk.as_metric_label(), + format!( + "Discarded {} spans after pre-write chunk failure ({})", + span_count, + err.status_code().as_ref() + ), + ); + return Ok(()); + } + ChunkFailureReaction::Propagate => { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::Propagate.as_metric_label(), + format!( + "Propagating pre-write chunk failure ({})", + err.status_code().as_ref() + ), + ); + return Err(err); + } + } + } + + let output = match self + .handle_trace_inserts(requests, ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + { + Ok(output) => output, + Err(err) => { + // Do not retry after dispatch: writes to other peers may already + // have committed even when the aggregate request returns an error. + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::Propagate.as_metric_label(), + format!( + "Propagating chunk write failure ({})", + err.status_code().as_ref() + ), + ); + return Err(err); + } + }; + + Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); + ingest_state.outcome.accepted_spans += span_count; + retry.add_to_aux_data(&mut ingest_state.aux_data); + + Ok(()) + } + + async fn ingest_trace_v1_rows_span_by_span( + &self, + retry: TraceChunkRetry, + ctx: QueryContextRef, + ingest_state: &mut TraceIngestState, + ) -> ServerResult<()> { + for row in retry.rows { + let (span, rows) = row.into_rows(&retry.schema)?; + let mut requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: retry.table_name.clone(), + rows: Some(rows), + }], + }; + + if let Err(err) = self.reconcile_trace_column_types(&mut requests, &ctx).await { + if matches!( + Self::classify_trace_prewrite_failure(err.status_code(), err.is_retryable()), + ChunkFailureReaction::Propagate + ) { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::Propagate.as_metric_label(), + format!( + "Propagating pre-write span failure for {}:{} ({})", + span.trace_id, + span.span_id, + err.status_code().as_ref() + ), + ); + return Err(err); + } + + ingest_state.outcome.rejected_spans += 1; + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + "span_rejected", + format!( + "Rejected span {}:{} ({})", + span.trace_id, + span.span_id, + err.status_code().as_ref() + ), + ); + continue; + } + + let output = match self + .handle_trace_inserts(requests, ctx.clone()) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + { + Ok(output) => output, + Err(err) => { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::Propagate.as_metric_label(), + format!( + "Propagating span write failure for {}:{} ({})", + span.trace_id, + span.span_id, + err.status_code().as_ref() + ), + ); + return Err(err); + } + }; + + Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); + ingest_state.outcome.accepted_spans += 1; + span.add_to_aux_data(&mut ingest_state.aux_data); + } + + Ok(()) + } + + /// Reconcile and insert one trace request batch. + async fn insert_trace_requests( + &self, + mut requests: RowInsertRequests, + is_trace_v1_model: bool, + ctx: QueryContextRef, + ) -> ServerResult { + if is_trace_v1_model { + self.reconcile_trace_column_types(&mut requests, &ctx) + .await?; + self.handle_trace_inserts(requests, ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } else { + self.handle_log_inserts(requests, ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } + } + + async fn plan_trace_v1_request_schema( + &self, + table_name: &str, + request_schema: &mut TraceRequestSchema, + ctx: &QueryContextRef, + ) -> ServerResult { + let catalog = ctx.current_catalog(); + let schema = ctx.current_schema(); + let table = self + .catalog_manager + .table(catalog, &schema, table_name, None) + .await?; + let table_schema = table.as_ref().map(|table| table.schema()); + let exclusions = request_schema.incompatible_schema_observations(table_schema.as_deref()); + if !exclusions.is_empty() { + return Ok(TraceRequestSchemaPlan::IncompatibleObservations(exclusions)); + } + + Ok(request_schema.resolve_table_schema(table_schema.as_deref())) + } + + async fn prepare_trace_v1_request_schema( + &self, + table_name: &str, + request_schema: &mut TraceRequestSchema, + chunks: &mut Vec, + chunk_states: &mut Vec, + ctx: &QueryContextRef, + ) -> ServerResult<()> { + loop { + // Fixed semconv targets are known before schema planning. Isolate invalid + // values before they can affect conflict resolution or DDL. + let exclusions = Self::prevalidate_trace_v1_fixed_columns(chunks, chunk_states)?; + if !exclusions.is_empty() { + Self::exclude_trace_v1_schema_observations( + request_schema, + chunks, + chunk_states, + &exclusions, + )?; + continue; + } + + if request_schema.columns.is_empty() { + return Ok(()); + } + + let pre_alter = match self + .plan_trace_v1_request_schema(table_name, request_schema, ctx) + .await? + { + TraceRequestSchemaPlan::Ready(pre_alter) => pre_alter, + TraceRequestSchemaPlan::IncompatibleObservations(exclusions) => { + Self::exclude_trace_v1_schema_observations( + request_schema, + chunks, + chunk_states, + &exclusions, + )?; + continue; + } + }; + if pre_alter.requires_ddl() { + let exclusions = Self::prevalidate_trace_v1_chunk_rewrites( + request_schema, + chunks, + chunk_states, + )?; + if !exclusions.is_empty() { + Self::exclude_trace_v1_schema_observations( + request_schema, + chunks, + chunk_states, + &exclusions, + )?; + continue; + } + + self.apply_trace_v1_pre_alter(ctx, table_name, pre_alter) + .await?; + let final_plan = self + .plan_trace_v1_request_schema(table_name, request_schema, ctx) + .await?; + let schema_converged = matches!( + final_plan, + TraceRequestSchemaPlan::Ready(final_plan) if !final_plan.requires_ddl() + ); + if !schema_converged { + for state in chunk_states.iter_mut() { + state.mark_for_reconcile(); + } + return Ok(()); + } + + return Ok(()); + } + + return Ok(()); + } + } + + fn prevalidate_trace_v1_fixed_columns( + chunks: &[TraceChunkRetry], + chunk_states: &[TraceChunkSchemaState], + ) -> ServerResult { + let mut exclusions = HashMap::>::new(); + for (batch_index, (chunk, state)) in chunks.iter().zip(chunk_states).enumerate() { + if *state != TraceChunkSchemaState::Prepared { + continue; + } + + 'rows: for row in &chunk.rows { + for retry_value in &row.values { + let Some(column) = chunk.schema.get(retry_value.column_index) else { + return error::InternalSnafu { + err_msg: format!( + "trace retry column index {} is out of bounds", + retry_value.column_index + ), + } + .fail(); + }; + let Some(target_type) = trace_semconv_fixed_type(&column.column_name) else { + continue; + }; + let Some(request_type) = retry_value + .value + .value_data + .as_ref() + .and_then(trace_value_datatype) + else { + continue; + }; + if request_type != target_type + && coerce_value_data( + &retry_value.value.value_data, + target_type, + request_type, + ) + .is_err() + { + exclusions + .entry(column.column_name.clone()) + .or_default() + .insert(batch_index); + break 'rows; + } + } + } + } + Ok(exclusions) + } + + fn prevalidate_trace_v1_chunk_rewrites( + request_schema: &TraceRequestSchema, + chunks: &[TraceChunkRetry], + chunk_states: &mut [TraceChunkSchemaState], + ) -> ServerResult { + let mut exclusions = HashMap::>::new(); + for (batch_index, (chunk, state)) in chunks.iter().zip(chunk_states).enumerate() { + if *state != TraceChunkSchemaState::Prepared { + continue; + } + let mut request = chunk.to_request()?; + if let Some(column_name) = Self::prepare_trace_v1_chunk_rewrite( + request_schema, + batch_index, + &mut request, + state, + )? { + exclusions + .entry(column_name) + .or_default() + .insert(batch_index); + } + } + Ok(exclusions) + } + + fn prepare_trace_v1_chunk_rewrite( + request_schema: &TraceRequestSchema, + batch_index: usize, + request: &mut RowInsertRequest, + state: &mut TraceChunkSchemaState, + ) -> ServerResult> { + if *state != TraceChunkSchemaState::Prepared { + return Ok(None); + } + + match request_schema.prepare_request_rewrite(batch_index, request) { + Ok(Some(rewrite)) => { + TraceRequestSchema::apply_request_rewrite(request, rewrite); + Ok(None) + } + Ok(None) => Ok(None), + Err(failure) => { + if !matches!( + Self::classify_trace_prewrite_failure( + failure.error.status_code(), + failure.error.is_retryable(), + ), + ChunkFailureReaction::RetryPerSpan + ) { + return Err(failure.error); + } + state.mark_for_reconcile(); + Ok(Some(failure.column_name)) + } + } + } + + fn exclude_trace_v1_schema_observations( + request_schema: &mut TraceRequestSchema, + chunks: &mut Vec, + chunk_states: &mut Vec, + exclusions: &TraceSchemaExclusions, + ) -> ServerResult<()> { + if chunks.len() != chunk_states.len() { + return error::InternalSnafu { + err_msg: format!( + "trace chunk count {} does not match state count {}", + chunks.len(), + chunk_states.len() + ), + } + .fail(); + } + + let batch_indexes = exclusions + .values() + .flatten() + .copied() + .collect::>(); + let previous_chunks = std::mem::take(chunks); + let previous_states = std::mem::take(chunk_states); + for (batch_index, (chunk, mut state)) in + previous_chunks.into_iter().zip(previous_states).enumerate() + { + if !batch_indexes.contains(&batch_index) { + chunks.push(chunk); + chunk_states.push(state); + } else if chunk.rows.len() <= 1 { + state.mark_for_reconcile(); + chunks.push(chunk); + chunk_states.push(state); + } else { + let split_chunks = chunk.into_single_span_chunks()?; + chunk_states.extend(std::iter::repeat_n( + TraceChunkSchemaState::Prepared, + split_chunks.len(), + )); + chunks.extend(split_chunks); + } + } + + *request_schema = TraceRequestSchema::default(); + for (batch_index, (chunk, state)) in chunks.iter().zip(chunk_states.iter()).enumerate() { + if *state == TraceChunkSchemaState::Prepared { + request_schema.observe_retry_chunk(batch_index, chunk)?; + } + } + + Ok(()) + } + + async fn apply_trace_v1_pre_alter( + &self, + ctx: &QueryContextRef, + table_name: &str, + pre_alter: TraceTablePreAlter, + ) -> ServerResult<()> { + let TraceTablePreAlter { + ensure_columns, + modify_float64_columns, + alter_existing, + } = pre_alter; + + if !ensure_columns.is_empty() { + self.inserter + .ensure_trace_table_on_demand( + table_name, + ensure_columns, + alter_existing, + ctx, + &self.statement_executor, + ) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; + } + + if !modify_float64_columns.is_empty() { + self.alter_trace_table_columns_to_float64(ctx, table_name, &modify_float64_columns) + .await?; + } + + Ok(()) + } + + fn classify_trace_prewrite_failure( + status: StatusCode, + retryable: bool, + ) -> ChunkFailureReaction { + if retryable { + return ChunkFailureReaction::Propagate; + } + + match status { + StatusCode::InvalidArguments + | StatusCode::InvalidSyntax + | StatusCode::Unsupported + | StatusCode::TableNotFound + | StatusCode::TableColumnNotFound => ChunkFailureReaction::RetryPerSpan, + StatusCode::DatabaseNotFound => ChunkFailureReaction::DiscardChunk, + _ => ChunkFailureReaction::Propagate, + } + } + + fn add_trace_write_cost(outcome: &mut TraceIngestOutcome, cost: usize) { + outcome.write_cost += cost; + } + + fn push_trace_failure_message(messages: &mut Vec, label: &str, message: String) { + OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).inc(); + + if messages.len() < TRACE_FAILURE_MESSAGE_LIMIT { + messages.push(message); + } else if messages.len() == TRACE_FAILURE_MESSAGE_LIMIT { + tracing::debug!( + label, + limit = TRACE_FAILURE_MESSAGE_LIMIT, + "Trace ingest failure message limit reached; suppressing additional failure details" + ); + } + } + + fn finish_trace_failure_message( + accepted_spans: usize, + rejected_spans: usize, + messages: Vec, + ) -> Option { + if rejected_spans == 0 && messages.is_empty() { + return None; + } + + let mut summary = format!( + "Accepted {} spans, rejected {} spans", + accepted_spans, rejected_spans + ); + + if !messages.is_empty() { + summary.push_str(": "); + summary.push_str(&messages.join("; ")); + } + + Some(summary) + } + + /// Widen existing trace table columns to Float64 before request rewrite. + async fn alter_trace_table_columns_to_float64( + &self, + ctx: &QueryContextRef, + table_name: &str, + column_names: &[String], + ) -> ServerResult<()> { + let catalog_name = ctx.current_catalog().to_string(); + let schema_name = ctx.current_schema(); + let alter_expr = AlterTableExpr { + catalog_name: catalog_name.clone(), + schema_name: schema_name.clone(), + table_name: table_name.to_string(), + kind: Some(Kind::ModifyColumnTypes(ModifyColumnTypes { + modify_column_types: column_names + .iter() + .map(|column_name| ModifyColumnType { + column_name: column_name.clone(), + target_type: ColumnDataType::Float64 as i32, + target_type_extension: None, + }) + .collect(), + })), + }; + + if let Err(err) = self + .statement_executor + .alter_table_inner(alter_expr, ctx.clone()) + .await + { + let table = self + .catalog_manager + .table(&catalog_name, &schema_name, table_name, None) + .await + .map_err(servers::error::Error::from)?; + let alter_already_applied = table + .map(|table| { + let table_schema = table.schema(); + column_names.iter().all(|column_name| { + table_schema + .column_schema_by_name(column_name) + .and_then(|table_col| { + ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) + .ok() + .map(|wrapper| wrapper.datatype()) + }) + == Some(ColumnDataType::Float64) + }) + }) + .unwrap_or(false); + + if alter_already_applied { + return Ok(()); + } + + warn!( + table_name, + columns = ?column_names, + error = %err, + "failed to widen trace columns before insert" + ); + + return Err(wrap_trace_alter_failure(err)); + } + + Ok(()) + } + + /// Coerce request column types and values to match the existing table schema + /// for compatible type pairs. Existing table schema wins when present; + /// otherwise the full request batch decides a stable target type. + async fn reconcile_trace_column_types( + &self, + requests: &mut RowInsertRequests, + ctx: &QueryContextRef, + ) -> ServerResult<()> { + let catalog = ctx.current_catalog(); + let schema = ctx.current_schema(); + + for req in &mut requests.inserts { + let table = self + .catalog_manager + .table(catalog, &schema, &req.table_name, None) + .await?; + + let Some(rows) = req.rows.as_mut() else { + continue; + }; + + let table_schema = table.map(|table| table.schema()); + let mut pending_rewrites = Vec::new(); + let mut pending_alter_columns = Vec::new(); + + for (col_idx, col_schema) in rows.schema.iter().enumerate() { + let Some(current_type) = ColumnDataType::try_from(col_schema.datatype).ok() else { + continue; + }; + + let mut observed_types = Vec::new(); + push_observed_trace_type(&mut observed_types, current_type); + + // Scan the full request first so the final type decision is not affected + // by row order inside the batch. + for row in &rows.rows { + let Some(value) = row + .values + .get(col_idx) + .and_then(|value| value.value_data.as_ref()) + else { + continue; + }; + + let Some(value_type) = trace_value_datatype(value) else { + continue; + }; + push_observed_trace_type(&mut observed_types, value_type); + } + + let existing_schema_type = table_schema + .as_ref() + .and_then(|schema| schema.column_schema_by_name(&col_schema.column_name)) + .and_then(|table_col| { + ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) + .ok() + .map(|wrapper| { + let datatype = wrapper.datatype(); + (datatype, ConcreteDataType::from(wrapper)) + }) + }); + let request_concrete_type = ConcreteDataType::from(ColumnDataTypeWrapper::new( + current_type, + col_schema.datatype_extension.clone(), + )); + if let Some((existing_datatype, existing_concrete_type)) = + existing_schema_type.as_ref() + && trace_logical_types_incompatible( + current_type, + &request_concrete_type, + *existing_datatype, + existing_concrete_type, + ) + { + return error::InvalidParameterSnafu { + reason: format!( + "incompatible logical types for trace column '{}' in table '{}'", + col_schema.column_name, req.table_name + ), + } + .fail(); + } + let existing_type = existing_schema_type.map(|(datatype, _)| datatype); + let fixed_type = trace_semconv_fixed_type(&col_schema.column_name); + + if !observed_types + .iter() + .copied() + .any(is_trace_reconcile_candidate_type) + && existing_type + .map(|datatype| !is_trace_reconcile_candidate_type(datatype)) + .unwrap_or(true) + && fixed_type.is_none() + { + continue; + } + + // Decide the final type once per column, then rewrite all affected cells + // together in one row pass below. + let Some(decision) = choose_trace_reconcile_decision( + &col_schema.column_name, + &observed_types, + existing_type, + ) + .map_err(|_| { + enrich_trace_reconcile_error( + &req.table_name, + &col_schema.column_name, + &observed_types, + existing_type, + fixed_type, + ) + })? + else { + continue; + }; + let target_type = decision.target_type(); + + if !decision.requires_alter() + && observed_types + .iter() + .all(|observed| *observed == target_type) + && col_schema.datatype == target_type as i32 + { + continue; + } + + if decision.requires_alter() + && !pending_alter_columns.contains(&col_schema.column_name) + { + pending_alter_columns.push(col_schema.column_name.clone()); + } + + pending_rewrites.push(PendingTraceColumnRewrite { + col_idx, + target_type, + column_name: col_schema.column_name.clone(), + }); + } + + if pending_rewrites.is_empty() { + continue; + } + + let prepared_rewrites = + prepare_trace_column_rewrites(&rows.rows, pending_rewrites, &req.table_name) + .map_err(|failure| failure.error)?; + + if !pending_alter_columns.is_empty() { + self.alter_trace_table_columns_to_float64( + ctx, + &req.table_name, + &pending_alter_columns, + ) + .await?; + } + + prepared_rewrites.apply(rows); + } + + Ok(()) + } +} + +fn trace_logical_types_incompatible( + left_datatype: ColumnDataType, + left_concrete_type: &ConcreteDataType, + right_datatype: ColumnDataType, + right_concrete_type: &ConcreteDataType, +) -> bool { + left_concrete_type != right_concrete_type + && (left_datatype == right_datatype + || !is_trace_reconcile_candidate_type(left_datatype) + || !is_trace_reconcile_candidate_type(right_datatype)) +} + +fn chunk_owned(items: Vec, chunk_size: usize) -> Vec> { + if items.is_empty() { + return Vec::new(); + } + + if chunk_size == 0 { + return vec![items]; + } + + let mut chunks = Vec::with_capacity(items.len().div_ceil(chunk_size)); + let mut iter = items.into_iter(); + while iter.len() > 0 { + chunks.push(iter.by_ref().take(chunk_size).collect()); + } + chunks +} + +/// Preserve the original alter failure status so chunk retry behavior stays correct. +fn wrap_trace_alter_failure(err: E) -> servers::error::Error +where + E: ErrorExt + Send + Sync + 'static, +{ + error::ExecuteGrpcQuerySnafu.into_error(BoxedError::new(err)) +} + +/// Derives `trace.conventions` from the request's resource/scope `schema_url`s. +/// A single distinct non-empty value is concrete; multiple distinct values are +/// `mixed`; none is `unknown`. `schema_url` is row-level in OTLP, so the +/// table-level value is best-effort per the RFC conflict rule. +pub(super) fn trace_conventions(request: &ExportTraceServiceRequest) -> String { + let mut seen: Option<&str> = None; + let mut mixed = false; + + for resource_spans in &request.resource_spans { + let urls = std::iter::once(resource_spans.schema_url.as_str()).chain( + resource_spans + .scope_spans + .iter() + .map(|s| s.schema_url.as_str()), + ); + for url in urls { + if url.is_empty() { + continue; + } + match seen { + None => seen = Some(url), + Some(prev) if prev == url => {} + Some(_) => { + mixed = true; + break; + } + } + } + if mixed { + break; + } + } + + if mixed { + SEMANTIC_VALUE_MIXED.to_string() + } else { + seen.map(str::to_string) + .unwrap_or_else(|| SEMANTIC_VALUE_UNKNOWN.to_string()) + } +} + +#[cfg(test)] +mod tests; diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs new file mode 100644 index 00000000000..949ab1535ac --- /dev/null +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -0,0 +1,1212 @@ +// 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 std::collections::{HashMap, HashSet}; + +use api::v1::column_data_type_extension::TypeExt; +use api::v1::value::ValueData; +use api::v1::{ + ColumnDataType, ColumnDataTypeExtension, ColumnSchema, JsonTypeExtension, Row, + RowInsertRequest, Rows, SemanticType, Value, +}; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ + ColumnSchema as DatatypesColumnSchema, SchemaBuilder as DatatypesSchemaBuilder, +}; +use opentelemetry_proto::tonic::common::v1::any_value::Value as OtlpValue; +use opentelemetry_proto::tonic::common::v1::{AnyValue, ArrayValue, KeyValue}; +use servers::otlp::trace::SERVICE_NAME_COLUMN; +use servers::otlp::trace::attributes::Attributes; +use servers::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; +use servers::otlp::trace::v1::{TraceBinaryType, TraceRetryColumn}; +use servers::query_handler::TraceIngestOutcome; + +use super::{ + ChunkFailureReaction, Instance, TraceChunkRetry, TraceChunkSchemaState, TraceRequestSchema, + TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, + wrap_trace_alter_failure, +}; +use crate::metrics::OTLP_TRACES_FAILURE_COUNT; + +#[test] +fn test_chunk_owned() { + let chunks = chunk_owned(vec![1, 2, 3], 2); + assert_eq!(chunks.iter().map(Vec::len).collect::>(), vec![2, 1]); + + let chunks = chunk_owned(vec![1, 2, 3], 0); + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].len(), 3); + + assert!(chunk_owned::(Vec::new(), 0).is_empty()); +} + +#[test] +fn test_classify_trace_prewrite_failure() { + for status in [ + StatusCode::InvalidArguments, + StatusCode::InvalidSyntax, + StatusCode::Unsupported, + StatusCode::TableColumnNotFound, + StatusCode::TableNotFound, + ] { + assert_eq!( + Instance::classify_trace_prewrite_failure(status, false), + ChunkFailureReaction::RetryPerSpan + ); + } + + assert_eq!( + Instance::classify_trace_prewrite_failure(StatusCode::DatabaseNotFound, false), + ChunkFailureReaction::DiscardChunk + ); + + for status in [ + StatusCode::DeadlineExceeded, + StatusCode::Cancelled, + StatusCode::StorageUnavailable, + StatusCode::Internal, + StatusCode::RegionNotReady, + StatusCode::TableUnavailable, + StatusCode::RegionBusy, + StatusCode::RuntimeResourcesExhausted, + StatusCode::Unknown, + ] { + assert_eq!( + Instance::classify_trace_prewrite_failure(status, false), + ChunkFailureReaction::Propagate + ); + } + + // Retry hints take precedence over status-based partial success handling. + assert_eq!( + Instance::classify_trace_prewrite_failure(StatusCode::InvalidArguments, true), + ChunkFailureReaction::Propagate + ); + assert_eq!( + Instance::classify_trace_prewrite_failure(StatusCode::DatabaseNotFound, true), + ChunkFailureReaction::Propagate + ); +} + +#[test] +fn test_add_trace_write_cost() { + let mut outcome = TraceIngestOutcome::default(); + Instance::add_trace_write_cost(&mut outcome, 3); + Instance::add_trace_write_cost(&mut outcome, 5); + assert_eq!(outcome.write_cost, 8); +} + +#[test] +fn test_finish_trace_failure_message() { + let message = Instance::finish_trace_failure_message( + 3, + 2, + vec!["Rejected span trace:span (InvalidArguments)".to_string()], + ) + .unwrap(); + assert!(message.contains("Accepted 3 spans, rejected 2 spans")); + assert!(message.contains("Rejected span trace:span")); + + assert_eq!(Instance::finish_trace_failure_message(2, 0, vec![]), None); +} + +#[test] +fn test_finish_trace_failure_message_without_detail_messages() { + assert_eq!( + Instance::finish_trace_failure_message(0, 2, vec![]), + Some("Accepted 0 spans, rejected 2 spans".to_string()) + ); +} + +#[test] +fn test_push_trace_failure_message_increments_labeled_counter() { + let label = "retry_per_span_counter_test"; + let initial = OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(); + let mut messages = Vec::new(); + + Instance::push_trace_failure_message( + &mut messages, + label, + "Chunk fallback triggered by InvalidArguments".to_string(), + ); + + assert_eq!(messages.len(), 1); + assert_eq!( + OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(), + initial + 1 + ); +} + +#[test] +fn test_push_trace_failure_message_caps_recorded_messages() { + let label = "retry_per_span_limit_test"; + let mut messages = Vec::new(); + + for idx in 0..=4 { + Instance::push_trace_failure_message(&mut messages, label, format!("failure-{idx}")); + } + + assert_eq!(messages.len(), 4); + assert_eq!( + messages, + vec![ + "failure-0".to_string(), + "failure-1".to_string(), + "failure-2".to_string(), + "failure-3".to_string() + ] + ); +} + +#[test] +fn test_wrap_trace_alter_failure_preserves_status_code() { + let err = wrap_trace_alter_failure( + servers::error::TableNotFoundSnafu { + catalog: "greptime".to_string(), + schema: "public".to_string(), + table: "trace_type_missing".to_string(), + } + .build(), + ); + + assert_eq!(err.status_code(), StatusCode::TableNotFound); +} + +fn field_schema(name: &str, datatype: ColumnDataType) -> ColumnSchema { + ColumnSchema { + column_name: name.to_string(), + datatype: datatype as i32, + semantic_type: SemanticType::Field as i32, + ..Default::default() + } +} + +fn tag_schema(name: &str, datatype: ColumnDataType) -> ColumnSchema { + ColumnSchema { + semantic_type: SemanticType::Tag as i32, + ..field_schema(name, datatype) + } +} + +fn json_field_schema(name: &str) -> ColumnSchema { + ColumnSchema { + datatype_extension: Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), + }), + ..field_schema(name, ColumnDataType::Binary) + } +} + +fn row(values: Vec>) -> Row { + Row { + values: values + .into_iter() + .map(|value_data| Value { value_data }) + .collect(), + } +} + +fn row_insert_request( + table_name: &str, + schema: Vec, + rows: Vec, +) -> RowInsertRequest { + RowInsertRequest { + table_name: table_name.to_string(), + rows: Some(Rows { schema, rows }), + } +} + +fn trace_attribute(key: &str, value: OtlpValue) -> KeyValue { + KeyValue { + key: key.to_string(), + value: Some(AnyValue { value: Some(value) }), + } +} + +fn trace_span(span_id: &str, span_attributes: Vec) -> TraceSpan { + TraceSpan { + service_name: Some("service".to_string()), + trace_id: "trace-id".to_string(), + span_id: span_id.to_string(), + parent_span_id: None, + resource_attributes: Attributes::from(vec![]), + scope_name: "scope".to_string(), + scope_version: "v1".to_string(), + scope_attributes: Attributes::from(vec![]), + trace_state: String::new(), + span_name: "operation".to_string(), + span_kind: "SPAN_KIND_SERVER".to_string(), + span_status_code: "STATUS_CODE_UNSET".to_string(), + span_status_message: String::new(), + span_attributes: Attributes::from(span_attributes), + span_events: SpanEvents::from(vec![]), + span_links: SpanLinks::from(vec![]), + start_in_nanosecond: 1, + end_in_nanosecond: 2, + } +} + +fn ready_plan(plan: TraceRequestSchemaPlan) -> TraceTablePreAlter { + let TraceRequestSchemaPlan::Ready(plan) = plan else { + panic!("expected a ready trace schema plan"); + }; + plan +} + +fn trace_chunk_retry(request: &RowInsertRequest, span_ids: &[&str]) -> TraceChunkRetry { + TraceChunkRetry::try_new( + request.table_name.clone(), + request.rows.clone().unwrap(), + span_ids + .iter() + .map(|span_id| TraceSpanMetadata { + trace_id: "trace-id".to_string(), + span_id: (*span_id).to_string(), + operation: None, + }) + .collect(), + HashMap::new(), + ) + .unwrap() +} + +fn prepare_and_apply( + request_schema: &TraceRequestSchema, + batch_index: usize, + request: &mut RowInsertRequest, +) { + if let Some(prepared) = request_schema + .prepare_request_rewrite(batch_index, request) + .unwrap() + { + TraceRequestSchema::apply_request_rewrite(request, prepared); + } +} + +#[test] +fn test_trace_chunk_retry_drops_absent_sparse_columns() { + let schema = vec![ + field_schema("parent_span_id", ColumnDataType::String), + field_schema("service_name", ColumnDataType::String), + field_schema("span_attributes.absent", ColumnDataType::Int64), + field_schema("scope_attributes.present", ColumnDataType::Boolean), + ]; + let requests = row_insert_request( + "traces", + schema, + vec![ + row(vec![None, None, None, Some(ValueData::BoolValue(true))]), + row(vec![ + None, + Some(ValueData::StringValue("service".to_string())), + Some(ValueData::I64Value(1)), + None, + ]), + ], + ); + let retry = TraceChunkRetry::try_new( + requests.table_name, + requests.rows.unwrap(), + ["span-1", "span-2"] + .into_iter() + .map(|span_id| TraceSpanMetadata { + trace_id: "trace-id".to_string(), + span_id: span_id.to_string(), + operation: None, + }) + .collect(), + HashMap::from([ + ( + "service_name".to_string(), + TraceRetryColumn { + present_rows: vec![1], + binary_types: Vec::new(), + }, + ), + ( + "span_attributes.absent".to_string(), + TraceRetryColumn { + present_rows: vec![1], + binary_types: Vec::new(), + }, + ), + ( + "scope_attributes.present".to_string(), + TraceRetryColumn { + present_rows: vec![0], + binary_types: Vec::new(), + }, + ), + ]), + ) + .unwrap(); + assert_eq!(retry.rows[0].values.len(), 2); + + let mut rows = retry.rows.into_iter(); + let (_, projected) = rows.next().unwrap().into_rows(&retry.schema).unwrap(); + + assert_eq!( + projected + .schema + .iter() + .map(|column| column.column_name.as_str()) + .collect::>(), + ["parent_span_id", "scope_attributes.present"] + ); + assert_eq!(projected.rows[0].values.len(), 2); + assert_eq!(projected.rows[0].values[0].value_data, None); + assert_eq!( + projected.rows[0].values[1].value_data, + Some(ValueData::BoolValue(true)) + ); +} + +#[test] +fn test_trace_chunk_retry_restores_row_logical_types() { + let column_name = "span_attributes.payload"; + let requests = row_insert_request( + "traces", + vec![field_schema(column_name, ColumnDataType::Binary)], + vec![ + row(vec![Some(ValueData::BinaryValue(vec![1, 2, 3]))]), + row(vec![Some(ValueData::BinaryValue(vec![4, 5, 6]))]), + row(vec![Some(ValueData::StringValue("value".to_string()))]), + ], + ); + let metadata = ["span-1", "span-2", "span-3"] + .into_iter() + .map(|span_id| TraceSpanMetadata { + trace_id: "trace-id".to_string(), + span_id: span_id.to_string(), + operation: None, + }) + .collect(); + let retry = TraceChunkRetry::try_new( + requests.table_name, + requests.rows.unwrap(), + metadata, + HashMap::from([( + column_name.to_string(), + TraceRetryColumn { + present_rows: vec![0, 1, 2], + binary_types: vec![(0, TraceBinaryType::Binary), (1, TraceBinaryType::Json)], + }, + )]), + ) + .unwrap(); + + let projected = retry + .rows + .into_iter() + .map(|row| row.into_rows(&retry.schema).unwrap().1) + .collect::>(); + + assert!(projected[0].schema[0].datatype_extension.is_none()); + assert!(matches!( + projected[1].schema[0] + .datatype_extension + .as_ref() + .and_then(|extension| extension.type_ext.as_ref()), + Some(TypeExt::JsonType(_)) + )); + assert_eq!( + projected[2].schema[0].datatype, + ColumnDataType::String as i32 + ); + assert!(projected[2].schema[0].datatype_extension.is_none()); +} + +#[test] +fn test_trace_chunk_retry_reconstructs_converter_request() { + let spans = vec![ + trace_span( + "span-1", + vec![ + trace_attribute("first", OtlpValue::IntValue(1)), + trace_attribute("payload", OtlpValue::BytesValue(vec![1, 2, 3])), + ], + ), + trace_span( + "span-2", + vec![ + trace_attribute("second", OtlpValue::DoubleValue(2.5)), + trace_attribute( + "payload", + OtlpValue::ArrayValue(ArrayValue { + values: vec![AnyValue { + value: Some(OtlpValue::StringValue("value".to_string())), + }], + }), + ), + ], + ), + ]; + let metadata = spans.iter().map(TraceSpanMetadata::from).collect(); + let (expected_table_data, _) = + servers::otlp::trace::v1::v1_to_main_table_data_with_schema(spans.clone()).unwrap(); + let (expected_schema, mut expected_rows) = expected_table_data.into_schema_and_rows(); + for row in &mut expected_rows { + row.values.resize(expected_schema.len(), Value::default()); + } + let expected = row_insert_request("traces", expected_schema, expected_rows); + let (table_data, batch_schema) = + servers::otlp::trace::v1::v1_to_main_table_data_with_schema(spans).unwrap(); + let (schema, rows) = table_data.into_schema_and_rows(); + let retry = TraceChunkRetry::try_new( + "traces".to_string(), + Rows { schema, rows }, + metadata, + batch_schema.into_retry_columns(), + ) + .unwrap(); + + assert_eq!(retry.to_request().unwrap(), expected); +} + +#[test] +fn test_trace_request_schema_detects_datatype_extension_conflicts() { + let column_name = "span_attributes.payload"; + let binary_schema = field_schema(column_name, ColumnDataType::Binary); + let json_schema = json_field_schema(column_name); + let legacy_json_schema = ColumnSchema { + datatype: ColumnDataType::Json as i32, + ..json_schema.clone() + }; + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &binary_schema, None); + request_schema.observe_trace_column(1, &json_schema, None); + assert_eq!( + request_schema.incompatible_schema_observations(None), + HashMap::from([(column_name.to_string(), HashSet::from([0, 1]))]) + ); + + let mut equivalent_json_request_schema = TraceRequestSchema::default(); + equivalent_json_request_schema.observe_trace_column(0, &json_schema, None); + equivalent_json_request_schema.observe_trace_column(1, &legacy_json_schema, None); + assert!( + equivalent_json_request_schema + .incompatible_schema_observations(None) + .is_empty() + ); + + let mut incompatible_json_request_schema = TraceRequestSchema::default(); + incompatible_json_request_schema.observe_trace_column(0, &binary_schema, None); + incompatible_json_request_schema.observe_trace_column(1, &legacy_json_schema, None); + assert_eq!( + incompatible_json_request_schema.incompatible_schema_observations(None), + HashMap::from([(column_name.to_string(), HashSet::from([0, 1]))]) + ); + + let existing_json_schema = + DatatypesSchemaBuilder::try_from_columns(vec![DatatypesColumnSchema::new( + column_name, + ConcreteDataType::json_datatype(), + true, + )]) + .unwrap() + .build() + .unwrap(); + let mut binary_request_schema = TraceRequestSchema::default(); + binary_request_schema.observe_trace_column(0, &binary_schema, None); + assert_eq!( + binary_request_schema.incompatible_schema_observations(Some(&existing_json_schema)), + HashMap::from([(column_name.to_string(), HashSet::from([0]))]) + ); + + let mut json_request_schema = TraceRequestSchema::default(); + json_request_schema.observe_trace_column(0, &json_schema, None); + assert!( + json_request_schema + .incompatible_schema_observations(Some(&existing_json_schema)) + .is_empty() + ); +} + +#[test] +fn test_trace_request_schema_isolates_non_candidate_batch() { + let column_name = "span_attributes.payload"; + let string_schema = field_schema(column_name, ColumnDataType::String); + let binary_schema = field_schema(column_name, ColumnDataType::Binary); + let existing_string_schema = + DatatypesSchemaBuilder::try_from_columns(vec![DatatypesColumnSchema::new( + column_name, + ConcreteDataType::string_datatype(), + true, + )]) + .unwrap() + .build() + .unwrap(); + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &string_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(1, &binary_schema, Some(ColumnDataType::Binary)); + + assert_eq!( + request_schema.incompatible_schema_observations(Some(&existing_string_schema)), + HashMap::from([(column_name.to_string(), HashSet::from([1]))]) + ); +} + +#[test] +fn test_trace_request_schema_isolates_non_candidate_new_column_batch() { + let column_name = "span_attributes.payload"; + for datatypes in [ + [ + ColumnDataType::String, + ColumnDataType::Int64, + ColumnDataType::Binary, + ], + [ + ColumnDataType::Binary, + ColumnDataType::Int64, + ColumnDataType::String, + ], + ] { + let binary_batch = datatypes + .iter() + .position(|datatype| *datatype == ColumnDataType::Binary) + .unwrap(); + let mut request_schema = TraceRequestSchema::default(); + for (batch_index, datatype) in datatypes.into_iter().enumerate() { + request_schema.observe_trace_column( + batch_index, + &field_schema(column_name, datatype), + Some(datatype), + ); + } + + assert_eq!( + request_schema.incompatible_schema_observations(None), + HashMap::from([(column_name.to_string(), HashSet::from([binary_batch]),)]) + ); + + request_schema.remove_batches(&HashSet::from([binary_batch])); + let plan = ready_plan(request_schema.resolve_table_schema(None)); + assert_eq!(plan.ensure_columns.len(), 1); + assert_eq!( + plan.ensure_columns[0].datatype, + ColumnDataType::Int64 as i32 + ); + } +} + +#[test] +fn test_trace_request_schema_isolates_unsupported_type_batches() { + let mixed_column = "span_attributes.mixed"; + let unrelated_column = "span_attributes.unrelated"; + let int_schema = field_schema(mixed_column, ColumnDataType::Int64); + let bool_schema = field_schema(mixed_column, ColumnDataType::Boolean); + let string_schema = field_schema(unrelated_column, ColumnDataType::String); + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &int_schema, Some(ColumnDataType::Int64)); + request_schema.observe_trace_column(0, &string_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(1, &bool_schema, Some(ColumnDataType::Boolean)); + request_schema.observe_trace_column(1, &string_schema, Some(ColumnDataType::String)); + + let TraceRequestSchemaPlan::IncompatibleObservations(exclusions) = + request_schema.resolve_table_schema(None) + else { + panic!("expected incompatible trace schema batches"); + }; + assert_eq!( + exclusions, + HashMap::from([(mixed_column.to_string(), HashSet::from([0, 1]))]) + ); + + let batch_indexes = exclusions + .values() + .flatten() + .copied() + .collect::>(); + request_schema.remove_batches(&batch_indexes); + let plan = ready_plan(request_schema.resolve_table_schema(None)); + assert!(!plan.requires_ddl()); +} + +#[test] +fn test_trace_request_schema_drops_failed_batch_observations() { + let table_name = "trace_failed_batch"; + let column_name = "span_attributes.value"; + let failed_only_column = "span_attributes.server.port"; + let safe_column = "span_attributes.safe"; + let string_schema = field_schema(column_name, ColumnDataType::String); + let float_schema = field_schema(column_name, ColumnDataType::Float64); + let failed_only_schema = field_schema(failed_only_column, ColumnDataType::String); + let service_schema = tag_schema(SERVICE_NAME_COLUMN, ColumnDataType::String); + let safe_schema = field_schema(safe_column, ColumnDataType::Boolean); + + let mut good_request = row_insert_request( + table_name, + vec![string_schema.clone()], + vec![row(vec![Some(ValueData::StringValue("1.5".to_string()))])], + ); + let mut bad_request = row_insert_request( + table_name, + vec![ + string_schema.clone(), + failed_only_schema.clone(), + service_schema.clone(), + safe_schema.clone(), + ], + vec![row(vec![ + Some(ValueData::StringValue("not-a-number".to_string())), + Some(ValueData::StringValue("invalid-port".to_string())), + Some(ValueData::StringValue("frontend".to_string())), + Some(ValueData::BoolValue(true)), + ])], + ); + let original_bad_request = bad_request.clone(); + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &string_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(0, &float_schema, Some(ColumnDataType::Float64)); + request_schema.observe_trace_column(1, &string_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(1, &failed_only_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(1, &service_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(1, &safe_schema, Some(ColumnDataType::Boolean)); + ready_plan(request_schema.resolve_table_schema(None)); + + let mut chunk_states = vec![ + TraceChunkSchemaState::Prepared, + TraceChunkSchemaState::Prepared, + ]; + assert_eq!( + Instance::prepare_trace_v1_chunk_rewrite( + &request_schema, + 0, + &mut good_request, + &mut chunk_states[0], + ) + .unwrap(), + None + ); + let failed_column = Instance::prepare_trace_v1_chunk_rewrite( + &request_schema, + 1, + &mut bad_request, + &mut chunk_states[1], + ) + .unwrap(); + let exclusions = HashMap::from([(failed_column.unwrap(), HashSet::from([1]))]); + assert_eq!( + exclusions, + HashMap::from([(column_name.to_string(), HashSet::from([1]))]) + ); + assert_eq!(chunk_states[0], TraceChunkSchemaState::Prepared); + assert_eq!(chunk_states[1], TraceChunkSchemaState::ReconcilePerChunk); + + let mut chunks = vec![ + trace_chunk_retry(&good_request, &["good-span"]), + trace_chunk_retry(&bad_request, &["bad-span"]), + ]; + Instance::exclude_trace_v1_schema_observations( + &mut request_schema, + &mut chunks, + &mut chunk_states, + &exclusions, + ) + .unwrap(); + assert!( + !request_schema + .column_indexes + .contains_key(failed_only_column) + ); + assert!( + !request_schema + .column_indexes + .contains_key(SERVICE_NAME_COLUMN) + ); + assert!(!request_schema.column_indexes.contains_key(safe_column)); + let plan = ready_plan(request_schema.resolve_table_schema(None)); + assert_eq!(plan.ensure_columns.len(), 1); + assert_eq!(plan.ensure_columns[0].column_name, column_name); + assert_eq!( + Instance::prepare_trace_v1_chunk_rewrite( + &request_schema, + 1, + &mut bad_request, + &mut chunk_states[1], + ) + .unwrap(), + None + ); + + let good_rows = good_request.rows.as_ref().unwrap(); + assert_eq!(good_rows.schema[0].datatype, ColumnDataType::Float64 as i32); + assert_eq!( + good_rows.rows[0].values[0].value_data, + Some(ValueData::F64Value(1.5)) + ); + assert_eq!(bad_request, original_bad_request); +} + +#[test] +fn test_trace_request_schema_splits_failed_chunk_before_exclusion() { + let table_name = "trace_failed_chunk"; + let value_column = "span_attributes.value"; + let port_column = "span_attributes.server.port"; + let boolean_schema = field_schema(value_column, ColumnDataType::Boolean); + let int_schema = field_schema(value_column, ColumnDataType::Int64); + let string_schema = field_schema(value_column, ColumnDataType::String); + let port_string_schema = field_schema(port_column, ColumnDataType::String); + let port_int_schema = field_schema(port_column, ColumnDataType::Int64); + let requests = [ + row_insert_request( + table_name, + vec![boolean_schema, port_string_schema], + vec![ + row(vec![ + Some(ValueData::BoolValue(true)), + Some(ValueData::StringValue("1".to_string())), + ]), + row(vec![ + Some(ValueData::BoolValue(false)), + Some(ValueData::StringValue("bad".to_string())), + ]), + ], + ), + row_insert_request( + table_name, + vec![int_schema, port_int_schema.clone()], + vec![row(vec![ + Some(ValueData::I64Value(2)), + Some(ValueData::I64Value(2)), + ])], + ), + row_insert_request( + table_name, + vec![string_schema, port_int_schema], + vec![row(vec![ + Some(ValueData::StringValue("3".to_string())), + Some(ValueData::I64Value(3)), + ])], + ), + ]; + let mut chunks = vec![ + trace_chunk_retry(&requests[0], &["good-span", "bad-span"]), + trace_chunk_retry(&requests[1], &["int-span"]), + trace_chunk_retry(&requests[2], &["string-span"]), + ]; + let mut chunk_states = vec![TraceChunkSchemaState::Prepared; chunks.len()]; + let mut request_schema = TraceRequestSchema::default(); + for (batch_index, chunk) in chunks.iter().enumerate() { + request_schema + .observe_retry_chunk(batch_index, chunk) + .unwrap(); + } + + ready_plan(request_schema.resolve_table_schema(None)); + let value_index = request_schema.column_indexes[value_column]; + assert_eq!( + request_schema.columns[value_index].target_type, + Some(ColumnDataType::String) + ); + + let exclusions = + Instance::prevalidate_trace_v1_chunk_rewrites(&request_schema, &chunks, &mut chunk_states) + .unwrap(); + assert_eq!( + exclusions, + HashMap::from([(port_column.to_string(), HashSet::from([0]))]) + ); + Instance::exclude_trace_v1_schema_observations( + &mut request_schema, + &mut chunks, + &mut chunk_states, + &exclusions, + ) + .unwrap(); + assert_eq!(chunks.len(), 4); + assert!(chunks.iter().all(|chunk| chunk.rows.len() == 1)); + + ready_plan(request_schema.resolve_table_schema(None)); + let exclusions = + Instance::prevalidate_trace_v1_chunk_rewrites(&request_schema, &chunks, &mut chunk_states) + .unwrap(); + assert_eq!( + exclusions, + HashMap::from([(port_column.to_string(), HashSet::from([1]))]) + ); + Instance::exclude_trace_v1_schema_observations( + &mut request_schema, + &mut chunks, + &mut chunk_states, + &exclusions, + ) + .unwrap(); + + ready_plan(request_schema.resolve_table_schema(None)); + let value_index = request_schema.column_indexes[value_column]; + assert_eq!( + request_schema.columns[value_index].target_type, + Some(ColumnDataType::String) + ); + assert_eq!(chunk_states[0], TraceChunkSchemaState::Prepared); + assert_eq!(chunk_states[1], TraceChunkSchemaState::ReconcilePerChunk); +} + +#[test] +fn test_trace_request_schema_prevalidates_fixed_columns_before_conflicts() { + let table_name = "trace_fixed_column_conflict"; + let conflict_column = "span_attributes.conflict"; + let value_column = "span_attributes.value"; + let port_column = "span_attributes.server.port"; + let requests = [ + row_insert_request( + table_name, + vec![ + field_schema(conflict_column, ColumnDataType::String), + field_schema(value_column, ColumnDataType::Boolean), + ], + vec![row(vec![ + Some(ValueData::StringValue("text".to_string())), + Some(ValueData::BoolValue(true)), + ])], + ), + row_insert_request( + table_name, + vec![ + field_schema(conflict_column, ColumnDataType::Binary), + field_schema(port_column, ColumnDataType::String), + ], + vec![row(vec![ + Some(ValueData::BinaryValue(vec![1])), + Some(ValueData::StringValue("bad".to_string())), + ])], + ), + row_insert_request( + table_name, + vec![field_schema(value_column, ColumnDataType::Int64)], + vec![row(vec![Some(ValueData::I64Value(2))])], + ), + row_insert_request( + table_name, + vec![field_schema(value_column, ColumnDataType::String)], + vec![row(vec![Some(ValueData::StringValue("3".to_string()))])], + ), + ]; + let mut chunks = requests + .iter() + .map(|request| trace_chunk_retry(request, &["span"])) + .collect::>(); + let mut chunk_states = vec![TraceChunkSchemaState::Prepared; chunks.len()]; + let mut request_schema = TraceRequestSchema::default(); + for (batch_index, chunk) in chunks.iter().enumerate() { + request_schema + .observe_retry_chunk(batch_index, chunk) + .unwrap(); + } + + assert_eq!( + request_schema.incompatible_schema_observations(None), + HashMap::from([(conflict_column.to_string(), HashSet::from([1]),)]) + ); + let exclusions = Instance::prevalidate_trace_v1_fixed_columns(&chunks, &chunk_states).unwrap(); + assert_eq!( + exclusions, + HashMap::from([(port_column.to_string(), HashSet::from([1]))]) + ); + Instance::exclude_trace_v1_schema_observations( + &mut request_schema, + &mut chunks, + &mut chunk_states, + &exclusions, + ) + .unwrap(); + + assert!( + request_schema + .incompatible_schema_observations(None) + .is_empty() + ); + ready_plan(request_schema.resolve_table_schema(None)); + let value_index = request_schema.column_indexes[value_column]; + assert_eq!( + request_schema.columns[value_index].target_type, + Some(ColumnDataType::String) + ); + assert_eq!(chunk_states[0], TraceChunkSchemaState::Prepared); + assert_eq!(chunk_states[1], TraceChunkSchemaState::ReconcilePerChunk); +} + +#[test] +fn test_trace_request_schema_uses_global_columns_and_types() { + let table_name = "trace_global_schema"; + let attr_num = "span_attributes.attr_num"; + let attr_later = "span_attributes.attr_later"; + let mut first_chunk = row_insert_request( + table_name, + vec![field_schema(attr_num, ColumnDataType::Int64)], + vec![row(vec![Some(ValueData::I64Value(1))])], + ); + let mut second_chunk = row_insert_request( + table_name, + vec![ + field_schema(attr_num, ColumnDataType::Float64), + field_schema(attr_later, ColumnDataType::Boolean), + ], + vec![row(vec![ + Some(ValueData::F64Value(2.5)), + Some(ValueData::BoolValue(true)), + ])], + ); + + let mut request_schema = TraceRequestSchema::default(); + let attr_num_i64 = field_schema(attr_num, ColumnDataType::Int64); + let attr_num_f64 = field_schema(attr_num, ColumnDataType::Float64); + let attr_later_bool = field_schema(attr_later, ColumnDataType::Boolean); + request_schema.observe_trace_column(0, &attr_num_i64, Some(ColumnDataType::Int64)); + request_schema.observe_trace_column(1, &attr_num_f64, Some(ColumnDataType::Float64)); + request_schema.observe_trace_column(1, &attr_later_bool, Some(ColumnDataType::Boolean)); + + let existing_schema = + DatatypesSchemaBuilder::try_from_columns(vec![DatatypesColumnSchema::new( + attr_num, + ConcreteDataType::int64_datatype(), + true, + )]) + .unwrap() + .build() + .unwrap(); + let pre_alter = ready_plan(request_schema.resolve_table_schema(Some(&existing_schema))); + assert_eq!(pre_alter.modify_float64_columns, vec![attr_num.to_string()]); + assert_eq!(pre_alter.ensure_columns.len(), 1); + assert_eq!(pre_alter.ensure_columns[0].column_name, attr_later); + + prepare_and_apply(&request_schema, 0, &mut first_chunk); + prepare_and_apply(&request_schema, 1, &mut second_chunk); + + let rows = first_chunk.rows.as_ref().unwrap(); + assert_eq!(rows.schema.len(), 1); + assert_eq!(rows.schema[0].column_name, attr_num); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Float64 as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::F64Value(1.0)) + ); + + let rows = second_chunk.rows.as_ref().unwrap(); + assert_eq!(rows.schema.len(), 2); + assert_eq!(rows.schema[0].column_name, attr_num); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Float64 as i32); + assert_eq!(rows.schema[1].column_name, attr_later); + assert_eq!(rows.schema[1].datatype, ColumnDataType::Boolean as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::F64Value(2.5)) + ); + assert_eq!( + rows.rows[0].values[1].value_data, + Some(ValueData::BoolValue(true)) + ); +} + +#[test] +fn test_trace_request_schema_resolves_three_compatible_types() { + let column_name = "span_attributes.value"; + let mut request_schema = TraceRequestSchema::default(); + for (batch_index, datatype) in [ + ColumnDataType::String, + ColumnDataType::Int64, + ColumnDataType::Float64, + ] + .into_iter() + .enumerate() + { + let schema = field_schema(column_name, datatype); + request_schema.observe_trace_column(batch_index, &schema, Some(datatype)); + } + + let plan = ready_plan(request_schema.resolve_table_schema(None)); + assert_eq!(plan.ensure_columns.len(), 1); + assert_eq!( + plan.ensure_columns[0].datatype, + ColumnDataType::Float64 as i32 + ); +} + +#[test] +fn test_trace_request_schema_creates_table_without_padding_batches() { + let table_name = "trace_create_schema"; + let attr_num = "span_attributes.attr_num"; + let attr_later = "span_attributes.attr_later"; + let mut first_chunk = row_insert_request( + table_name, + vec![field_schema(attr_num, ColumnDataType::Int64)], + vec![row(vec![Some(ValueData::I64Value(1))])], + ); + let mut second_chunk = row_insert_request( + table_name, + vec![ + field_schema(attr_num, ColumnDataType::Float64), + field_schema(attr_later, ColumnDataType::Boolean), + ], + vec![row(vec![ + Some(ValueData::F64Value(2.5)), + Some(ValueData::BoolValue(true)), + ])], + ); + + let mut request_schema = TraceRequestSchema::default(); + let attr_num_i64 = field_schema(attr_num, ColumnDataType::Int64); + let attr_num_f64 = field_schema(attr_num, ColumnDataType::Float64); + let attr_later_bool = field_schema(attr_later, ColumnDataType::Boolean); + request_schema.observe_trace_column(0, &attr_num_i64, Some(ColumnDataType::Int64)); + request_schema.observe_trace_column(1, &attr_num_f64, Some(ColumnDataType::Float64)); + request_schema.observe_trace_column(1, &attr_later_bool, Some(ColumnDataType::Boolean)); + + let pre_alter = ready_plan(request_schema.resolve_table_schema(None)); + assert!(pre_alter.modify_float64_columns.is_empty()); + assert_eq!(pre_alter.ensure_columns.len(), 2); + assert_eq!(pre_alter.ensure_columns[0].column_name, attr_num); + assert_eq!( + pre_alter.ensure_columns[0].datatype, + ColumnDataType::Float64 as i32 + ); + assert_eq!(pre_alter.ensure_columns[1].column_name, attr_later); + + prepare_and_apply(&request_schema, 0, &mut first_chunk); + prepare_and_apply(&request_schema, 1, &mut second_chunk); + + let rows = first_chunk.rows.as_ref().unwrap(); + assert_eq!(rows.schema.len(), 1); + assert_eq!(rows.schema[0].column_name, attr_num); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Float64 as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::F64Value(1.0)) + ); + assert_eq!(rows.rows[0].values.len(), 1); + + let rows = second_chunk.rows.as_ref().unwrap(); + assert_eq!(rows.schema.len(), 2); + assert_eq!(rows.schema[0].column_name, attr_num); + assert_eq!(rows.schema[1].column_name, attr_later); +} + +#[test] +fn test_trace_request_schema_re_resolves_after_concurrent_create() { + let table_name = "trace_concurrent_create"; + let column_name = "span_attributes.value"; + let string_schema = field_schema(column_name, ColumnDataType::String); + let mut requests = row_insert_request( + table_name, + vec![string_schema.clone()], + vec![row(vec![Some(ValueData::StringValue("42".to_string()))])], + ); + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &string_schema, Some(ColumnDataType::String)); + let initial_plan = ready_plan(request_schema.resolve_table_schema(None)); + assert!(!initial_plan.ensure_columns.is_empty()); + assert!(!initial_plan.alter_existing); + let initial_targets = request_schema.resolved_target_types(); + + let concurrent_schema = + DatatypesSchemaBuilder::try_from_columns(vec![DatatypesColumnSchema::new( + column_name, + ConcreteDataType::int64_datatype(), + true, + )]) + .unwrap() + .build() + .unwrap(); + let final_plan = ready_plan(request_schema.resolve_table_schema(Some(&concurrent_schema))); + assert!(!final_plan.requires_ddl()); + assert!(final_plan.alter_existing); + assert_ne!(initial_targets, request_schema.resolved_target_types()); + + prepare_and_apply(&request_schema, 0, &mut requests); + let rows = requests.rows.as_ref().unwrap(); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Int64 as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::I64Value(42)) + ); +} + +#[test] +fn test_trace_request_schema_keeps_request_intact_on_coercion_failure() { + let table_name = "trace_atomic_rewrite"; + let column_name = "span_attributes.value"; + let string_schema = field_schema(column_name, ColumnDataType::String); + let float_schema = field_schema(column_name, ColumnDataType::Float64); + let requests = row_insert_request( + table_name, + vec![string_schema.clone()], + vec![row(vec![Some(ValueData::StringValue( + "not-a-number".to_string(), + ))])], + ); + let original = requests.clone(); + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &string_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(0, &float_schema, Some(ColumnDataType::Float64)); + ready_plan(request_schema.resolve_table_schema(None)); + + assert!( + request_schema + .prepare_request_rewrite(0, &requests) + .is_err() + ); + assert_eq!(requests, original); +} + +use opentelemetry_proto::tonic::trace::v1::{ResourceSpans, ScopeSpans}; + +use super::{ExportTraceServiceRequest, trace_conventions}; + +fn resource_spans(resource_url: &str, scope_urls: &[&str]) -> ResourceSpans { + ResourceSpans { + schema_url: resource_url.to_string(), + scope_spans: scope_urls + .iter() + .map(|u| ScopeSpans { + schema_url: u.to_string(), + ..Default::default() + }) + .collect(), + ..Default::default() + } +} + +#[test] +fn test_trace_conventions() { + let unknown = ExportTraceServiceRequest::default(); + assert_eq!(trace_conventions(&unknown), "unknown"); + + let url = "https://opentelemetry.io/schemas/1.27.0"; + let single = ExportTraceServiceRequest { + resource_spans: vec![resource_spans("", &[url, url])], + }; + assert_eq!(trace_conventions(&single), url); + + let resource_level = ExportTraceServiceRequest { + resource_spans: vec![resource_spans(url, &[""])], + }; + assert_eq!(trace_conventions(&resource_level), url); + + let conflicting = ExportTraceServiceRequest { + resource_spans: vec![resource_spans( + "", + &[url, "https://opentelemetry.io/schemas/1.30.0"], + )], + }; + assert_eq!(trace_conventions(&conflicting), "mixed"); +} diff --git a/src/frontend/src/instance/otlp/trace_types.rs b/src/frontend/src/instance/otlp/trace_types.rs index e46e16fbd5c..0e2df25053e 100644 --- a/src/frontend/src/instance/otlp/trace_types.rs +++ b/src/frontend/src/instance/otlp/trace_types.rs @@ -12,7 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::{ColumnDataType, Row}; +use api::v1::value::ValueData; +use api::v1::{ColumnDataType, Row, Rows}; use servers::error::{self, Result as ServerResult}; use servers::otlp::coerce::{ coerce_value_data, is_supported_trace_coercion, resolve_new_trace_column_type, @@ -42,12 +43,54 @@ impl TraceReconcileDecision { } } +/// Describes a column rewrite before its row values have been validated. +#[derive(Debug)] pub(super) struct PendingTraceColumnRewrite { pub(super) col_idx: usize, pub(super) target_type: ColumnDataType, pub(super) column_name: String, } +/// Holds the schema and value rewrites prepared for atomic application. +#[derive(Debug)] +pub(super) struct PreparedTraceColumnRewrites { + columns: Vec, + values: Vec, +} + +/// Reports the column whose trace value could not be rewritten. +#[derive(Debug)] +pub(super) struct TraceColumnRewriteError { + pub(super) error: servers::error::Error, + pub(super) column_name: String, +} + +/// Updates one request column to its reconciled datatype. +#[derive(Debug)] +struct PreparedTraceColumnRewrite { + col_idx: usize, + target_type: ColumnDataType, +} + +/// Replaces one trace value with its precomputed coerced value. +#[derive(Debug)] +struct PreparedTraceValueRewrite { + row_idx: usize, + col_idx: usize, + value_data: Option, +} + +impl PreparedTraceColumnRewrites { + pub(super) fn apply(self, rows: &mut Rows) { + for column in self.columns { + rows.schema[column.col_idx].datatype = column.target_type as i32; + } + for value in self.values { + rows.rows[value.row_idx].values[value.col_idx].value_data = value.value_data; + } + } +} + /// Picks the reconciliation action for one trace column. /// /// Existing table schema is authoritative unless the only incompatible case is @@ -129,14 +172,22 @@ fn choose_fixed_trace_reconcile_decision( .fail() } -/// Validate all pending trace column rewrites before any schema mutation happens. -pub(super) fn validate_trace_column_rewrites( +/// Prepares an atomic rewrite plan without mutating the input rows. +/// +/// For each pending column rewrite, this precomputes every required value +/// coercion. Missing, null, and already-correct values are skipped. If any +/// coercion fails, it returns the failing column and leaves all rows unchanged. +/// +/// Target types must already have been selected by `TraceRequestSchema`. +/// Call [`PreparedTraceColumnRewrites::apply`] to update the schema and values. +pub(super) fn prepare_trace_column_rewrites( rows: &[Row], - pending_rewrites: &[PendingTraceColumnRewrite], + pending_rewrites: Vec, table_name: &str, -) -> ServerResult<()> { - for row in rows { - for pending_rewrite in pending_rewrites { +) -> Result { + let mut values = Vec::new(); + for (row_idx, row) in rows.iter().enumerate() { + for pending_rewrite in &pending_rewrites { let Some(value) = row.values.get(pending_rewrite.col_idx) else { continue; }; @@ -148,23 +199,39 @@ pub(super) fn validate_trace_column_rewrites( continue; } - coerce_value_data(&value.value_data, pending_rewrite.target_type, request_type) - .map_err(|_| { - error::InvalidParameterSnafu { - reason: format!( - "failed to coerce trace column '{}' in table '{}' from {:?} to {:?}", - pending_rewrite.column_name, - table_name, - request_type, - pending_rewrite.target_type - ), - } - .build() - })?; + let value_data = + coerce_value_data(&value.value_data, pending_rewrite.target_type, request_type) + .map_err(|_| { + TraceColumnRewriteError { + error: error::InvalidParameterSnafu { + reason: format!( + "failed to coerce trace column '{}' in table '{}' from {:?} to {:?}", + pending_rewrite.column_name, + table_name, + request_type, + pending_rewrite.target_type + ), + } + .build(), + column_name: pending_rewrite.column_name.clone(), + } + })?; + values.push(PreparedTraceValueRewrite { + row_idx, + col_idx: pending_rewrite.col_idx, + value_data, + }); } } - Ok(()) + let columns = pending_rewrites + .into_iter() + .map(|rewrite| PreparedTraceColumnRewrite { + col_idx: rewrite.col_idx, + target_type: rewrite.target_type, + }) + .collect(); + Ok(PreparedTraceColumnRewrites { columns, values }) } pub(super) fn enrich_trace_reconcile_error( @@ -228,14 +295,14 @@ pub(super) fn push_observed_trace_type( #[cfg(test)] mod tests { use api::v1::value::ValueData; - use api::v1::{ColumnDataType, Row, Value}; + use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, Value}; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use super::{ PendingTraceColumnRewrite, TraceReconcileDecision, choose_trace_reconcile_decision, - enrich_trace_reconcile_error, is_trace_reconcile_candidate_type, push_observed_trace_type, - validate_trace_column_rewrites, + enrich_trace_reconcile_error, is_trace_reconcile_candidate_type, + prepare_trace_column_rewrites, push_observed_trace_type, }; #[test] @@ -375,7 +442,7 @@ mod tests { } #[test] - fn test_validate_trace_column_rewrites_rejects_invalid_string_parse() { + fn test_prepare_trace_column_rewrites_rejects_invalid_string_parse() { let rows = vec![Row { values: vec![Value { value_data: Some(ValueData::StringValue("not_a_number".to_string())), @@ -387,29 +454,49 @@ mod tests { column_name: "span_attributes.attr_int".to_string(), }]; - let err = validate_trace_column_rewrites(&rows, &pending_rewrites, "trace_type_atomicity") + let err = prepare_trace_column_rewrites(&rows, pending_rewrites, "trace_type_atomicity") .unwrap_err(); - assert_eq!(err.status_code(), StatusCode::InvalidArguments); + assert_eq!(err.error.status_code(), StatusCode::InvalidArguments); + assert_eq!(err.column_name, "span_attributes.attr_int"); } #[test] - fn test_validate_trace_column_rewrites_whitelisted_values_validate_against_fixed_type() { - let rows = vec![Row { - values: vec![Value { - value_data: Some(ValueData::StringValue("503".to_string())), + fn test_prepare_trace_column_rewrites_applies_prepared_values() { + let mut rows = Rows { + schema: vec![ColumnSchema { + datatype: ColumnDataType::String as i32, + ..Default::default() }], - }]; + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::StringValue("503".to_string())), + }], + }], + }; let pending_rewrites = vec![PendingTraceColumnRewrite { col_idx: 0, target_type: ColumnDataType::Int64, column_name: "span_attributes.http.response.status_code".to_string(), }]; - validate_trace_column_rewrites(&rows, &pending_rewrites, "trace_type_atomicity").unwrap(); + let prepared = + prepare_trace_column_rewrites(&rows.rows, pending_rewrites, "trace_type_atomicity") + .unwrap(); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::StringValue("503".to_string())) + ); + + prepared.apply(&mut rows); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Int64 as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::I64Value(503)) + ); } #[test] - fn test_validate_trace_column_rewrites_whitelisted_boolean_rejects_invalid_string_parse() { + fn test_prepare_trace_column_rewrites_boolean_rejects_invalid_string_parse() { let rows = vec![Row { values: vec![Value { value_data: Some(ValueData::StringValue("not_a_bool".to_string())), @@ -421,9 +508,13 @@ mod tests { column_name: "span_attributes.messaging.destination.temporary".to_string(), }]; - let err = validate_trace_column_rewrites(&rows, &pending_rewrites, "trace_type_atomicity") + let err = prepare_trace_column_rewrites(&rows, pending_rewrites, "trace_type_atomicity") .unwrap_err(); - assert_eq!(err.status_code(), StatusCode::InvalidArguments); + assert_eq!(err.error.status_code(), StatusCode::InvalidArguments); + assert_eq!( + err.column_name, + "span_attributes.messaging.destination.temporary" + ); } #[test] diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index a433602a36a..d2fbe20d034 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -107,7 +107,7 @@ pub enum AutoCreateTableType { /// A table that merges rows by `last_non_null` strategy. LastNonNull, /// Create table that build index and default partition rules on trace_id - Trace, + Trace { alter_existing: bool }, } impl AutoCreateTableType { @@ -117,9 +117,18 @@ impl AutoCreateTableType { AutoCreateTableType::Physical => "physical", AutoCreateTableType::Log => "log", AutoCreateTableType::LastNonNull => "last_non_null", - AutoCreateTableType::Trace => "trace", + AutoCreateTableType::Trace { .. } => "trace", } } + + fn alter_existing(&self) -> bool { + !matches!( + self, + Self::Trace { + alter_existing: false + } + ) + } } /// Split insert requests into normal and instant requests. @@ -214,7 +223,9 @@ impl Inserter { requests, ctx, statement_executor, - AutoCreateTableType::Trace, + AutoCreateTableType::Trace { + alter_existing: true, + }, false, false, ) @@ -503,6 +514,39 @@ impl Inserter { }) } + /// Ensures a trace table has the request-global schema without requiring a + /// padded data row to drive on-demand creation or alteration. When + /// `alter_existing` is false, a table created after planning is left for the + /// caller to re-plan. + pub async fn ensure_trace_table_on_demand( + &self, + table_name: &str, + request_schema: Vec, + alter_existing: bool, + ctx: &QueryContextRef, + statement_executor: &StatementExecutor, + ) -> Result<()> { + let mut requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: table_name.to_string(), + rows: Some(api::v1::Rows { + schema: request_schema, + rows: Vec::new(), + }), + }], + }; + self.create_or_alter_tables_on_demand( + &mut requests, + ctx, + AutoCreateTableType::Trace { alter_existing }, + statement_executor, + false, + false, + ) + .await?; + Ok(()) + } + /// Creates or alter tables on demand: /// - if table does not exist, create table by inferred CreateExpr /// - if table exist, check if schema matches. If any new column found, alter table by inferred `AlterExpr` @@ -575,6 +619,7 @@ impl Inserter { ctx, accommodate_existing_schema, is_single_value, + auto_create_table_type.alter_existing(), )? { alter_tables.push(alter_expr); need_refresh_table_infos.insert(( @@ -639,7 +684,7 @@ impl Inserter { } } - AutoCreateTableType::Trace => { + AutoCreateTableType::Trace { .. } => { let trace_table_name = ctx .extension(TRACE_TABLE_NAME_SESSION_KEY) .unwrap_or(TRACE_TABLE_NAME); @@ -904,7 +949,12 @@ impl Inserter { ctx: &QueryContextRef, accommodate_existing_schema: bool, is_single_value: bool, + alter_existing: bool, ) -> Result> { + if !alter_existing { + return Ok(None); + } + let catalog_name = ctx.current_catalog(); let schema_name = ctx.current_schema(); let table_name = table.table_info().name.clone(); @@ -1187,7 +1237,7 @@ pub fn fill_table_options_for_create( table_options.insert(MERGE_MODE_KEY.to_string(), "last_non_null".to_string()); } } - AutoCreateTableType::Trace => { + AutoCreateTableType::Trace { .. } => { table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string()); } } @@ -1473,8 +1523,15 @@ mod tests { )), true, ); + // Do not apply an absent-table plan to a table that appeared concurrently. + assert!( + inserter + .get_alter_table_expr_on_demand(&mut req, &table, &ctx, false, false, false) + .unwrap() + .is_none() + ); let alter_expr = inserter - .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true) + .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true, true) .unwrap(); assert!(alter_expr.is_none()); diff --git a/src/servers/src/otlp/coerce.rs b/src/servers/src/otlp/coerce.rs index febec0fda9b..2d87744c5c2 100644 --- a/src/servers/src/otlp/coerce.rs +++ b/src/servers/src/otlp/coerce.rs @@ -108,38 +108,35 @@ pub fn trace_value_datatype(value: &ValueData) -> Option { pub fn resolve_new_trace_column_type( observed_types: impl IntoIterator, ) -> Result, TraceCoerceError> { - let mut observed_types = observed_types.into_iter().collect::>(); - observed_types.dedup(); - - match observed_types.as_slice() { - [] => Ok(None), - [datatype] => Ok(Some(*datatype)), - [_, _] - if observed_types.contains(&ColumnDataType::String) - && observed_types.contains(&ColumnDataType::Boolean) => - { - Ok(Some(ColumnDataType::Boolean)) + let mut observed = Vec::new(); + for datatype in observed_types { + if !observed.contains(&datatype) { + observed.push(datatype); } - [_, _] - if observed_types.contains(&ColumnDataType::String) - && observed_types.contains(&ColumnDataType::Int64) => - { - Ok(Some(ColumnDataType::Int64)) - } - [_, _] - if observed_types.contains(&ColumnDataType::String) - && observed_types.contains(&ColumnDataType::Float64) => - { - Ok(Some(ColumnDataType::Float64)) - } - [_, _] - if observed_types.contains(&ColumnDataType::Int64) - && observed_types.contains(&ColumnDataType::Float64) => - { - Ok(Some(ColumnDataType::Float64)) - } - _ => Err(TraceCoerceError::Unsupported), } + + if observed.is_empty() { + return Ok(None); + } + if observed.len() == 1 { + return Ok(observed.first().copied()); + } + + [ + ColumnDataType::Boolean, + ColumnDataType::Int64, + ColumnDataType::Float64, + ColumnDataType::String, + ] + .into_iter() + .find(|target| { + observed.contains(target) + && observed + .iter() + .all(|source| source == target || is_supported_trace_coercion(*source, *target)) + }) + .map(Some) + .ok_or(TraceCoerceError::Unsupported) } #[cfg(test)] @@ -337,7 +334,15 @@ mod tests { ColumnDataType::Int64, ColumnDataType::Float64, ]), - Err(TraceCoerceError::Unsupported) + Ok(Some(ColumnDataType::Float64)) + ); + assert_eq!( + resolve_new_trace_column_type([ + ColumnDataType::Float64, + ColumnDataType::String, + ColumnDataType::Int64, + ]), + Ok(Some(ColumnDataType::Float64)) ); } } diff --git a/src/servers/src/otlp/trace/v1.rs b/src/servers/src/otlp/trace/v1.rs index e04d344e54e..eedc598ffba 100644 --- a/src/servers/src/otlp/trace/v1.rs +++ b/src/servers/src/otlp/trace/v1.rs @@ -12,10 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; +use api::v1::column_data_type_extension::TypeExt; use api::v1::value::ValueData; -use api::v1::{ColumnDataType, RowInsertRequests, Value}; +use api::v1::{ + ColumnDataType, ColumnDataTypeExtension, ColumnSchema, JsonTypeExtension, RowInsertRequests, + Value, +}; use common_catalog::consts::{trace_operations_table_name, trace_services_table_name}; use common_grpc::precision::Precision; use opentelemetry_proto::tonic::common::v1::any_value::Value as OtlpValue; @@ -40,6 +44,155 @@ const APPROXIMATE_COLUMN_COUNT: usize = 30; // Use a timestamp(2100-01-01 00:00:00) as large as possible. const MAX_TIMESTAMP: i64 = 4102444800000000000; +/// Distinguishes raw bytes from JSONB values that both use a binary protobuf value. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TraceBinaryType { + /// Raw OTLP bytes with no datatype extension. + Binary, + /// An OTLP array or key-value list encoded with the JSONB extension. + Json, +} + +impl TraceBinaryType { + /// Applies this logical binary type to a column schema. + pub fn apply_to_schema(self, schema: &mut ColumnSchema) { + schema.datatype = ColumnDataType::Binary as i32; + schema.datatype_extension = match self { + Self::Binary => None, + Self::Json => Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), + }), + }; + } +} + +/// Per-column observations collected while building one trace chunk. +#[derive(Default)] +struct TraceBatchColumnSchema { + value_types: Vec, + present_rows: Vec, + binary_types: Vec<(usize, TraceBinaryType)>, +} + +impl TraceBatchColumnSchema { + /// Records a row once while preserving encounter order. + fn observe_row(&mut self, row_index: usize) { + if self.present_rows.last() != Some(&row_index) { + self.present_rows.push(row_index); + } + } + + /// Returns whether raw binary and JSONB values share this column. + fn has_incompatible_logical_types(&self) -> bool { + let Some((_, first_type)) = self.binary_types.first() else { + return false; + }; + self.binary_types + .iter() + .any(|(_, binary_type)| binary_type != first_type) + } +} + +/// Sparse column metadata retained for span-level fallback. +pub struct TraceRetryColumn { + /// Row indexes that contain this dynamic column. + pub present_rows: Vec, + /// Row-specific binary kinds needed to restore the logical schema. + pub binary_types: Vec<(usize, TraceBinaryType)>, +} + +/// Retry metadata keyed by dynamic trace column name. +pub type TraceRetryColumns = HashMap; + +/// Schema observations for dynamic columns in one converted trace chunk. +#[derive(Default)] +pub struct TraceBatchSchema { + columns: HashMap, +} + +impl TraceBatchSchema { + /// Records a scalar value type and the row where the column is present. + fn observe_value_type(&mut self, name: &str, row_index: usize, value_type: ColumnDataType) { + let column = self.columns.entry(name.to_string()).or_default(); + column.observe_row(row_index); + if !column.value_types.contains(&value_type) { + column.value_types.push(value_type); + } + } + + /// Records the logical kind of a binary protobuf value for one row. + fn observe_binary_type(&mut self, name: &str, row_index: usize, binary_type: TraceBinaryType) { + let column = self.columns.entry(name.to_string()).or_default(); + column.observe_row(row_index); + if !column.value_types.contains(&ColumnDataType::Binary) { + column.value_types.push(ColumnDataType::Binary); + } + if let Some((last_row_index, last_type)) = column.binary_types.last_mut() + && *last_row_index == row_index + { + *last_type = binary_type; + } else { + column.binary_types.push((row_index, binary_type)); + } + } + + /// Records a sparse column occurrence that has no dynamic value type. + fn observe_present_column(&mut self, name: &str, row_index: usize) { + self.columns + .entry(name.to_string()) + .or_default() + .observe_row(row_index); + } + + /// Returns the distinct value types in first-seen order for a column. + pub fn value_types(&self, name: &str) -> Option<&[ColumnDataType]> { + self.columns + .get(name) + .map(|column| column.value_types.as_slice()) + } + + /// Returns whether any column mixes raw binary and JSONB values. + pub fn has_incompatible_logical_types(&self) -> bool { + self.columns + .values() + .any(TraceBatchColumnSchema::has_incompatible_logical_types) + } + + /// Returns whether the named column mixes raw binary and JSONB values. + pub fn has_incompatible_logical_types_for(&self, name: &str) -> bool { + self.columns + .get(name) + .is_some_and(TraceBatchColumnSchema::has_incompatible_logical_types) + } + + /// Converts batch observations into sparse metadata for per-span retries. + pub fn into_retry_columns(self) -> TraceRetryColumns { + self.columns + .into_iter() + .filter(|(_, column)| !column.present_rows.is_empty()) + .map(|(name, column)| { + let binary_types = if column.has_incompatible_logical_types() + || column + .value_types + .iter() + .any(|datatype| *datatype != ColumnDataType::Binary) + { + column.binary_types + } else { + Vec::new() + }; + ( + name, + TraceRetryColumn { + present_rows: column.present_rows, + binary_types, + }, + ) + }) + .collect() + } +} + /// Converts trace spans into row insert requests for the main v1 trace table. /// /// Auxiliary service and operation table writes are built separately so the @@ -52,20 +205,58 @@ pub fn v1_to_grpc_main_insert_requests( _query_ctx: &QueryContextRef, _pipeline_handler: PipelineHandlerRef, ) -> Result<(RowInsertRequests, usize)> { + let requests = v1_to_grpc_main_insert_requests_from_iter(spans.iter().cloned(), table_name)?; + Ok((requests, spans.len())) +} + +/// Converts owned spans into unpadded main-table rows and schema observations. +pub fn v1_to_main_table_data_with_schema( + spans: Vec, +) -> Result<(TableData, TraceBatchSchema)> { + build_trace_table_data_with_schema(spans.into_iter()) +} + +/// Builds the main-table request without collecting batch schema observations. +fn v1_to_grpc_main_insert_requests_from_iter( + spans: impl ExactSizeIterator, + table_name: &str, +) -> Result { let mut multi_table_writer = MultiTableData::default(); - let trace_writer = build_trace_table_data(spans)?; + let trace_writer = build_trace_table_data_from_iter(spans, None)?; multi_table_writer.add_table_data(table_name, trace_writer); - Ok(multi_table_writer.into_row_insert_requests()) + Ok(multi_table_writer.into_row_insert_requests().0) } /// Builds the row-oriented payload for the main v1 trace table. pub fn build_trace_table_data(spans: &[TraceSpan]) -> Result { - let mut trace_writer = TableData::new(APPROXIMATE_COLUMN_COUNT, spans.len()); - for span in spans.iter().cloned() { - write_span_to_row(&mut trace_writer, span)?; - } + build_trace_table_data_from_iter(spans.iter().cloned(), None) +} +/// Builds trace rows while collecting dynamic column observations. +fn build_trace_table_data_with_schema( + spans: impl ExactSizeIterator, +) -> Result<(TableData, TraceBatchSchema)> { + let mut batch_schema = TraceBatchSchema::default(); + let trace_writer = build_trace_table_data_from_iter(spans, Some(&mut batch_schema))?; + Ok((trace_writer, batch_schema)) +} + +/// Shared row builder with optional batch schema observation. +fn build_trace_table_data_from_iter( + spans: impl ExactSizeIterator, + mut batch_schema: Option<&mut TraceBatchSchema>, +) -> Result { + let mut trace_writer = TableData::new(APPROXIMATE_COLUMN_COUNT, spans.len()); + for span in spans { + let row_index = trace_writer.num_rows(); + write_span_to_row_inner( + &mut trace_writer, + span, + row_index, + batch_schema.as_deref_mut(), + )?; + } Ok(trace_writer) } @@ -91,6 +282,17 @@ pub fn build_aux_table_requests( } pub fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> { + let row_index = writer.num_rows(); + write_span_to_row_inner(writer, span, row_index, None) +} + +/// Writes one span and optionally records its dynamic columns for reconciliation. +fn write_span_to_row_inner( + writer: &mut TableData, + span: TraceSpan, + row_index: usize, + mut batch_schema: Option<&mut TraceBatchSchema>, +) -> Result<()> { let mut row = writer.alloc_one_row(); // write ts @@ -132,6 +334,9 @@ pub fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> row_writer::write_fields(writer, fields.into_iter(), &mut row)?; if let Some(service_name) = span.service_name { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_present_column(SERVICE_NAME_COLUMN, row_index); + } row_writer::write_tags( writer, std::iter::once((SERVICE_NAME_COLUMN.to_string(), service_name)), @@ -139,13 +344,29 @@ pub fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> )?; } - write_attributes(writer, "span_attributes", span.span_attributes, &mut row)?; - write_attributes(writer, "scope_attributes", span.scope_attributes, &mut row)?; - write_attributes( + write_attributes_with_schema( + writer, + "span_attributes", + span.span_attributes, + &mut row, + row_index, + batch_schema.as_deref_mut(), + )?; + write_attributes_with_schema( + writer, + "scope_attributes", + span.scope_attributes, + &mut row, + row_index, + batch_schema.as_deref_mut(), + )?; + write_attributes_with_schema( writer, "resource_attributes", span.resource_attributes, &mut row, + row_index, + batch_schema, )?; row_writer::write_json( @@ -217,11 +438,25 @@ fn write_trace_operations_to_row( Ok(()) } +#[cfg(test)] pub(crate) fn write_attributes( writer: &mut TableData, prefix: &str, attributes: Attributes, row: &mut Vec, +) -> Result<()> { + let row_index = writer.num_rows(); + write_attributes_with_schema(writer, prefix, attributes, row, row_index, None) +} + +/// Writes flattened attributes without coercion and optionally records their actual types. +fn write_attributes_with_schema( + writer: &mut TableData, + prefix: &str, + attributes: Attributes, + row: &mut Vec, + row_index: usize, + mut batch_schema: Option<&mut TraceBatchSchema>, ) -> Result<()> { for attr in attributes.take().into_iter() { let key_suffix = attr.key; @@ -234,6 +469,9 @@ pub(crate) fn write_attributes( let key = format!("{}.{}", prefix, key_suffix); match attr.value.and_then(|v| v.value) { Some(OtlpValue::StringValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::String); + } // Keep the raw request value here. Mixed trace types are reconciled later // in the frontend once we can also see the existing table schema. writer.write_field_unchecked( @@ -244,6 +482,9 @@ pub(crate) fn write_attributes( ); } Some(OtlpValue::BoolValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::Boolean); + } // Do not coerce or promote types while building the request-local rows. writer.write_field_unchecked( &key, @@ -253,6 +494,9 @@ pub(crate) fn write_attributes( ); } Some(OtlpValue::IntValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::Int64); + } // Preserving the original value avoids order-dependent behavior inside one batch. writer.write_field_unchecked( &key, @@ -262,6 +506,9 @@ pub(crate) fn write_attributes( ); } Some(OtlpValue::DoubleValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::Float64); + } writer.write_field_unchecked( &key, ColumnDataType::Float64, @@ -269,28 +516,40 @@ pub(crate) fn write_attributes( row, ); } - Some(OtlpValue::ArrayValue(v)) => row_writer::write_json( - writer, - key, - any_value_to_jsonb(OtlpValue::ArrayValue(v)), - row, - )?, - Some(OtlpValue::KvlistValue(v)) => row_writer::write_json( - writer, - key, - any_value_to_jsonb(OtlpValue::KvlistValue(v)), - row, - )?, - Some(OtlpValue::BytesValue(v)) => { - row_writer::write_fields( - writer, - std::iter::once(make_column_data( - &key, - ColumnDataType::Binary, - Some(ValueData::BinaryValue(v)), + Some(OtlpValue::ArrayValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Json); + } + writer.write_column_unchecked( + row_writer::build_json_column_schema(key), + Some(ValueData::BinaryValue( + any_value_to_jsonb(OtlpValue::ArrayValue(v)).to_vec(), )), row, - )?; + ); + } + Some(OtlpValue::KvlistValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Json); + } + writer.write_column_unchecked( + row_writer::build_json_column_schema(key), + Some(ValueData::BinaryValue( + any_value_to_jsonb(OtlpValue::KvlistValue(v)).to_vec(), + )), + row, + ); + } + Some(OtlpValue::BytesValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Binary); + } + writer.write_field_unchecked( + key, + ColumnDataType::Binary, + Some(ValueData::BinaryValue(v)), + row, + ); } None => {} } @@ -341,6 +600,129 @@ mod tests { } } + #[test] + fn test_batch_schema_preserves_column_order_and_attribute_types() { + let mut span1 = make_span("svc-a", "trace-a", "span-a"); + span1.span_attributes = Attributes::from(vec![make_kv("val", OtlpValue::IntValue(10))]); + let mut span2 = make_span("svc-a", "trace-a", "span-b"); + span2.span_attributes = Attributes::from(vec![make_kv("val", OtlpValue::DoubleValue(1.5))]); + + let (table_data, batch_schema) = + build_trace_table_data_with_schema(vec![span1, span2].into_iter()).unwrap(); + + assert_eq!( + batch_schema.value_types("span_attributes.val").unwrap(), + [ColumnDataType::Int64, ColumnDataType::Float64] + ); + let timestamp_index = table_data + .columns() + .iter() + .position(|column| column.column_name == TIMESTAMP_COLUMN) + .unwrap(); + let attribute_index = table_data + .columns() + .iter() + .position(|column| column.column_name == "span_attributes.val") + .unwrap(); + assert!(timestamp_index < attribute_index); + } + + #[test] + fn test_optional_batch_schema_observation_preserves_rows() { + let mut span = make_span("svc-a", "trace-a", "span-a"); + span.span_attributes = Attributes::from(vec![make_kv("val", OtlpValue::IntValue(10))]); + + let without_schema = build_trace_table_data(std::slice::from_ref(&span)).unwrap(); + let (with_schema, batch_schema) = + build_trace_table_data_with_schema(vec![span].into_iter()).unwrap(); + + assert_eq!( + without_schema.into_schema_and_rows(), + with_schema.into_schema_and_rows() + ); + assert_eq!( + batch_schema.value_types("span_attributes.val").unwrap(), + [ColumnDataType::Int64] + ); + let retry_columns = batch_schema.into_retry_columns(); + let retry_column = &retry_columns["span_attributes.val"]; + assert_eq!(retry_column.present_rows, [0]); + assert!(retry_column.binary_types.is_empty()); + } + + #[test] + fn test_batch_schema_tracks_binary_and_json_rows() { + let mut binary_span = make_span("svc-a", "trace-a", "span-a"); + binary_span.span_attributes = Attributes::from(vec![make_kv( + "val", + OtlpValue::BytesValue(vec![1_u8, 2, 3]), + )]); + let mut json_span = make_span("svc-a", "trace-a", "span-b"); + json_span.span_attributes = Attributes::from(vec![make_kv( + "val", + OtlpValue::ArrayValue(ArrayValue { + values: vec![AnyValue { + value: Some(OtlpValue::IntValue(1)), + }], + }), + )]); + + let (_, batch_schema) = + build_trace_table_data_with_schema(vec![binary_span, json_span].into_iter()).unwrap(); + + assert!(batch_schema.has_incompatible_logical_types()); + assert!(batch_schema.has_incompatible_logical_types_for("span_attributes.val")); + let retry_columns = batch_schema.into_retry_columns(); + let retry_column = &retry_columns["span_attributes.val"]; + assert_eq!(retry_column.present_rows, [0, 1]); + assert_eq!( + retry_column.binary_types, + [(0, TraceBinaryType::Binary), (1, TraceBinaryType::Json)] + ); + } + + #[test] + fn test_batch_schema_preserves_scalar_then_binary_values() { + let mut scalar_span = make_span("svc-a", "trace-a", "span-a"); + scalar_span.span_attributes = Attributes::from(vec![ + make_kv("bytes", OtlpValue::StringValue("text".to_string())), + make_kv("json", OtlpValue::StringValue("text".to_string())), + ]); + let mut binary_span = make_span("svc-a", "trace-a", "span-b"); + binary_span.span_attributes = Attributes::from(vec![ + make_kv("bytes", OtlpValue::BytesValue(vec![1_u8, 2, 3])), + make_kv( + "json", + OtlpValue::ArrayValue(ArrayValue { + values: vec![AnyValue { + value: Some(OtlpValue::IntValue(1)), + }], + }), + ), + ]); + + let (_, batch_schema) = + build_trace_table_data_with_schema(vec![scalar_span, binary_span].into_iter()).unwrap(); + + assert_eq!( + batch_schema.value_types("span_attributes.bytes").unwrap(), + [ColumnDataType::String, ColumnDataType::Binary] + ); + assert_eq!( + batch_schema.value_types("span_attributes.json").unwrap(), + [ColumnDataType::String, ColumnDataType::Binary] + ); + let retry_columns = batch_schema.into_retry_columns(); + assert_eq!( + retry_columns["span_attributes.bytes"].binary_types, + [(1, TraceBinaryType::Binary)] + ); + assert_eq!( + retry_columns["span_attributes.json"].binary_types, + [(1, TraceBinaryType::Json)] + ); + } + #[test] fn test_keep_mixed_numeric_values_until_frontend_reconciliation() { let mut writer = TableData::new(4, 2); diff --git a/src/servers/src/row_writer.rs b/src/servers/src/row_writer.rs index fb3351323c9..5b7abe151bf 100644 --- a/src/servers/src/row_writer.rs +++ b/src/servers/src/row_writer.rs @@ -96,17 +96,30 @@ impl TableData { value: Option, one_row: &mut Vec, ) { - let name = name.to_string(); - if let Some(index) = self.column_indexes.get(&name).copied() { - one_row[index].value_data = value; - } else { - let index = self.schema.len(); - self.schema.push(ColumnSchema { - column_name: name.clone(), + self.write_column_unchecked( + ColumnSchema { + column_name: name.to_string(), datatype: datatype as i32, semantic_type: SemanticType::Field as i32, ..Default::default() - }); + }, + value, + one_row, + ); + } + + pub fn write_column_unchecked( + &mut self, + column_schema: ColumnSchema, + value: Option, + one_row: &mut Vec, + ) { + if let Some(index) = self.column_indexes.get(&column_schema.column_name).copied() { + one_row[index].value_data = value; + } else { + let index = self.schema.len(); + let name = column_schema.column_name.clone(); + self.schema.push(column_schema); self.column_indexes.insert(name, index); one_row.push(Value { value_data: value }); } @@ -235,7 +248,7 @@ pub fn write_f64( ) } -fn build_json_column_schema(name: impl ToString) -> ColumnSchema { +pub(crate) fn build_json_column_schema(name: impl ToString) -> ColumnSchema { ColumnSchema { column_name: name.to_string(), datatype: ColumnDataType::Binary as i32, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index ba6710da34b..5f4e5886d8a 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -6997,10 +6997,8 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { "unexpected error body: {body:?}" ); assert!( - body.message.contains("Chunk fallback triggered by") - || body - .message - .contains("Discarded 2 spans after ambiguous chunk failure"), + body.message + .contains("Discarded 2 spans after pre-write request failure"), "unexpected error body: {body:?}" ); From c969fb0a5ea132a2a3ba761b6643960d8a4b9748 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Thu, 16 Jul 2026 19:20:13 +0800 Subject: [PATCH 48/48] perf: reduce parquet metadata cache footprint (#8527) * perf(mito2): reuse parquet metadata size accounting Signed-off-by: Ruihang Xia * perf(mito2): skip unused parquet column indexes Signed-off-by: Ruihang Xia * perf(mito2): cache compact parquet metadata Signed-off-by: Ruihang Xia * fix(mito2): tolerate metadata cache encoding failures Signed-off-by: Ruihang Xia * fix(mito2): skip disabled metadata cache encoding Signed-off-by: Ruihang Xia * refactor(mito2): avoid nested metadata cache task Signed-off-by: Ruihang Xia --------- Signed-off-by: Ruihang Xia --- Cargo.lock | 1 + src/mito2/Cargo.toml | 1 + src/mito2/src/access_layer.rs | 48 +- src/mito2/src/cache.rs | 661 +++++++++++++++++++++--- src/mito2/src/cache/file_cache.rs | 69 ++- src/mito2/src/cache/write_cache.rs | 3 +- src/mito2/src/region/opener.rs | 98 +++- src/mito2/src/sst/parquet/file_range.rs | 2 +- src/mito2/src/sst/parquet/metadata.rs | 53 +- src/mito2/src/sst/parquet/reader.rs | 47 +- 10 files changed, 836 insertions(+), 147 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 29640722ba1..87827905fe6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8519,6 +8519,7 @@ dependencies = [ "toml 0.8.23", "tracing", "uuid", + "zstd", ] [[package]] diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 7ae38da9f2a..e48a5758eab 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -93,6 +93,7 @@ tokio-stream.workspace = true tokio-util.workspace = true tracing.workspace = true uuid.workspace = true +zstd.workspace = true [dev-dependencies] common-function.workspace = true diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index b6adb6eb595..5622acacefa 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -16,11 +16,13 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use async_stream::try_stream; +use common_telemetry::warn; use common_time::Timestamp; use futures::{Stream, TryStreamExt}; use object_store::services::Fs; use object_store::util::{join_dir, with_instrument_layers}; use object_store::{ATOMIC_WRITE_DIR, ErrorKind, OLD_ATOMIC_WRITE_DIR, ObjectStore}; +use parquet::file::metadata::PageIndexPolicy; use smallvec::SmallVec; use snafu::ResultExt; use store_api::metadata::RegionMetadataRef; @@ -28,9 +30,9 @@ use store_api::region_request::PathType; use store_api::sst_entry::StorageSstEntry; use store_api::storage::{FileId, RegionId, SequenceNumber}; -use crate::cache::CacheManagerRef; use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; use crate::cache::write_cache::SstUploadRequest; +use crate::cache::{CacheManagerRef, SstMetaPreparation, prepare_sst_meta_sync}; use crate::config::{BloomFilterConfig, FulltextIndexConfig, IndexConfig, InvertedIndexConfig}; use crate::error::{ CleanDirSnafu, DeleteIndexSnafu, DeleteIndexesSnafu, DeleteSstsSnafu, OpenDalSnafu, Result, @@ -410,14 +412,46 @@ impl AccessLayer { }; // Put parquet metadata to cache manager. - if !sst_info.is_empty() { + if !sst_info.is_empty() && cache_manager.sst_meta_cache_enabled() { for sst in &sst_info { if let Some(parquet_metadata) = &sst.file_metadata { - cache_manager.put_parquet_meta_data( - RegionFileId::new(region_id, sst.file_id), - parquet_metadata.clone(), - Some(region_metadata.clone()), - ) + let file_id = RegionFileId::new(region_id, sst.file_id); + let file_path = format!( + "region_id={}, file_id={}", + file_id.region_id(), + file_id.file_id() + ); + let page_index_policy = if parquet_metadata.offset_index().is_some() { + PageIndexPolicy::Optional + } else { + PageIndexPolicy::Skip + }; + let parquet_metadata = parquet_metadata.clone(); + let region_metadata = region_metadata.clone(); + let cache_manager = cache_manager.clone(); + // Compact cache preparation is best-effort. Run the entire operation in one + // detached blocking task so it neither blocks an async worker nor delays the + // SST write. + common_runtime::spawn_blocking_global(move || { + match prepare_sst_meta_sync( + &file_path, + Arc::unwrap_or_clone(parquet_metadata), + Some(region_metadata), + page_index_policy, + ) { + Ok(SstMetaPreparation::Prepared(metadata)) => { + cache_manager.put_prepared_sst_meta(file_id, metadata, true); + } + Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!( + encoding_error; + "Failed to encode parquet metadata for cache, file: {}", + file_path + ), + Err(err) => { + warn!(err; "Failed to cache parquet metadata for {}", file_path); + } + } + }); } } } diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 5915aa32083..27db7daa5ab 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -26,10 +26,11 @@ pub(crate) mod write_cache; use std::collections::{BTreeMap, HashMap}; use std::mem; use std::ops::Range; -use std::sync::{Arc, RwLock}; +use std::sync::{Arc, RwLock, Weak}; use bytes::Bytes; use common_base::readable_size::ReadableSize; +use common_datasource::compression::CompressionType; use common_telemetry::warn; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::record_batch::RecordBatch; @@ -41,11 +42,13 @@ use moka::notification::RemovalCause; use moka::sync::Cache; use object_store::ObjectStore; use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; -use parquet::file::metadata::{FileMetaData, PageIndexPolicy, ParquetMetaData}; +use parquet::file::metadata::{ + FileMetaData, PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, ParquetMetaDataWriter, +}; use puffin::puffin_manager::cache::{PuffinMetadataCache, PuffinMetadataCacheRef}; use smallvec::SmallVec; use snafu::{OptionExt, ResultExt}; -use store_api::metadata::RegionMetadataRef; +use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::{ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector}; use crate::cache::cache_size::parquet_meta_size; @@ -54,7 +57,10 @@ use crate::cache::index::inverted_index::{InvertedIndexCache, InvertedIndexCache #[cfg(feature = "vector_index")] use crate::cache::index::vector_index::{VectorIndexCache, VectorIndexCacheRef}; use crate::cache::write_cache::WriteCacheRef; -use crate::error::{InvalidMetadataSnafu, InvalidParquetSnafu, Result, UnexpectedSnafu}; +use crate::error::{ + CompressObjectSnafu, DecompressObjectSnafu, InvalidMetadataSnafu, InvalidParquetSnafu, + JoinSnafu, ReadParquetSnafu, Result, UnexpectedSnafu, WriteParquetSnafu, +}; use crate::memtable::record_batch_estimated_size; use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, CACHE_HIT, CACHE_MISS}; use crate::read::Batch; @@ -66,6 +72,8 @@ use crate::sst::parquet::reader::MetadataCacheMetrics; /// Metrics type key for sst meta. const SST_META_TYPE: &str = "sst_meta"; +/// Metrics type key for the optional decoded SST metadata acceleration tier. +const SST_META_DECODED_TYPE: &str = "sst_meta_decoded"; /// Metrics type key for vector. const VECTOR_TYPE: &str = "vector"; /// Metrics type key for pages. @@ -152,11 +160,235 @@ impl RangeResultMemoryLimiter { #[derive(Debug)] pub(crate) struct CachedSstMeta { parquet_metadata: Arc, + parquet_metadata_size: usize, region_metadata: RegionMetadataRef, region_metadata_weight: usize, page_index_policy: PageIndexPolicy, } +/// Compact, authoritative form of one SST's metadata. +/// +/// The entry contains a zstd-compressed, self-contained Parquet metadata stream. The decoded +/// representation is held weakly to coalesce concurrent decodes without retaining memory outside +/// the cache capacity. +#[derive(Debug)] +pub(crate) struct CompactSstMeta { + encoded_metadata: Bytes, + decoded_size: usize, + region_metadata: Weak, + page_index_policy: PageIndexPolicy, + decoded: tokio::sync::Mutex>, +} + +/// Both forms produced by decoding metadata after a cache miss. +#[derive(Debug)] +pub(crate) struct PreparedSstMeta { + compact: Arc, + decoded: Arc, +} + +impl PreparedSstMeta { + pub(crate) fn decoded(&self) -> Arc { + self.decoded.clone() + } +} + +/// Result of decoding SST metadata and attempting to encode its compact cache entry. +#[derive(Debug)] +pub(crate) enum SstMetaPreparation { + /// Both cache representations are ready for admission. + Prepared(PreparedSstMeta), + /// The metadata is usable by the reader, but its compact cache encoding failed. + DecodedOnly { + decoded: Arc, + encoding_error: crate::error::Error, + }, +} + +impl SstMetaPreparation { + pub(crate) fn decoded(&self) -> Arc { + match self { + Self::Prepared(metadata) => metadata.decoded(), + Self::DecodedOnly { decoded, .. } => decoded.clone(), + } + } +} + +impl CompactSstMeta { + async fn decode(&self) -> Result> { + let mut decoded_guard = self.decoded.lock().await; + if let Some(decoded) = decoded_guard.upgrade() { + return Ok(decoded); + } + + let encoded_metadata = self.encoded_metadata.clone(); + let decoded_size = self.decoded_size; + let region_metadata = self.region_metadata.upgrade(); + let page_index_policy = self.page_index_policy; + let decoded = common_runtime::spawn_blocking_global(move || { + let bytes = zstd::bulk::decompress(&encoded_metadata, decoded_size).context( + DecompressObjectSnafu { + compress_type: CompressionType::Zstd, + path: "cached SST metadata", + }, + )?; + let bytes = Bytes::from(bytes); + let mut reader = ParquetMetaDataReader::new() + .with_column_index_policy(PageIndexPolicy::Skip) + .with_offset_index_policy(page_index_policy); + reader.try_parse(&bytes).context(ReadParquetSnafu { + path: "cached SST metadata", + })?; + let metadata = reader.finish().context(ReadParquetSnafu { + path: "cached SST metadata", + })?; + CachedSstMeta::try_new_with_page_index_policy( + "cached SST metadata", + metadata, + region_metadata, + page_index_policy, + ) + .map(Arc::new) + }) + .await + .context(JoinSnafu)??; + + *decoded_guard = Arc::downgrade(&decoded); + Ok(decoded) + } + + fn satisfies_page_index_policy(&self, requested: PageIndexPolicy) -> bool { + satisfies_page_index_policy(self.page_index_policy, requested) + } +} + +/// Decodes SST metadata without preparing a compact cache entry. +pub(crate) async fn decode_sst_meta( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, +) -> Result> { + let file_path = file_path.to_string(); + common_runtime::spawn_blocking_global(move || { + let parquet_metadata = strip_column_indexes(parquet_metadata); + CachedSstMeta::try_new_with_page_index_policy( + &file_path, + parquet_metadata, + region_metadata, + page_index_policy, + ) + .map(Arc::new) + }) + .await + .context(JoinSnafu)? +} + +/// Decodes SST metadata and attempts to encode both cache representations on the blocking runtime. +pub(crate) async fn prepare_sst_meta( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, +) -> Result { + let file_path = file_path.to_string(); + common_runtime::spawn_blocking_global(move || { + prepare_sst_meta_sync( + &file_path, + parquet_metadata, + region_metadata, + page_index_policy, + ) + }) + .await + .context(JoinSnafu)? +} + +/// Synchronously prepares SST metadata. Callers must run this on a blocking runtime. +pub(crate) fn prepare_sst_meta_sync( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, +) -> Result { + let parquet_metadata = strip_column_indexes(parquet_metadata); + let cache_encoding = encode_compact_sst_meta(file_path, &parquet_metadata); + finish_sst_meta_preparation( + file_path, + parquet_metadata, + region_metadata, + page_index_policy, + cache_encoding, + ) +} + +fn strip_column_indexes(parquet_metadata: ParquetMetaData) -> ParquetMetaData { + // Defensively discard column indexes supplied by external metadata producers. Mito only + // consumes offset indexes. + let mut builder = parquet_metadata.into_builder(); + builder.take_column_index(); + builder.build() +} + +fn encode_compact_sst_meta( + file_path: &str, + parquet_metadata: &ParquetMetaData, +) -> Result<(Bytes, usize)> { + let mut encoded = Vec::new(); + ParquetMetaDataWriter::new(&mut encoded, parquet_metadata) + .finish() + .context(WriteParquetSnafu)?; + let decoded_size = encoded.len(); + let encoded_metadata = zstd::bulk::compress(&encoded, 3).context(CompressObjectSnafu { + compress_type: CompressionType::Zstd, + path: file_path, + })?; + + // `zstd::bulk::compress` allocates for the compression upper bound and leaves the excess + // capacity in its `Vec`. Convert through a boxed slice so the retained allocation matches the + // cache weight. + Ok(( + Bytes::from(encoded_metadata.into_boxed_slice()), + decoded_size, + )) +} + +fn finish_sst_meta_preparation( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, + cache_encoding: Result<(Bytes, usize)>, +) -> Result { + let decoded = Arc::new(CachedSstMeta::try_new_with_page_index_policy( + file_path, + parquet_metadata, + region_metadata, + page_index_policy, + )?); + let (encoded_metadata, decoded_size) = match cache_encoding { + Ok(encoded) => encoded, + Err(encoding_error) => { + return Ok(SstMetaPreparation::DecodedOnly { + decoded, + encoding_error, + }); + } + }; + let compact = Arc::new(CompactSstMeta { + encoded_metadata, + decoded_size, + region_metadata: Arc::downgrade(&decoded.region_metadata), + page_index_policy, + decoded: tokio::sync::Mutex::new(Arc::downgrade(&decoded)), + }); + + Ok(SstMetaPreparation::Prepared(PreparedSstMeta { + compact, + decoded, + })) +} + impl CachedSstMeta { #[cfg(test)] pub(crate) fn try_new(file_path: &str, parquet_metadata: ParquetMetaData) -> Result { @@ -215,9 +447,11 @@ impl CachedSstMeta { // Keep the previous JSON-byte floor and charge the decoded structures as well. let region_metadata_weight = region_metadata.estimated_size().max(json.len()); let parquet_metadata = Arc::new(strip_region_metadata_from_parquet(parquet_metadata)); + let parquet_metadata_size = parquet_meta_size(&parquet_metadata); Ok(Self { parquet_metadata, + parquet_metadata_size, region_metadata, region_metadata_weight, page_index_policy, @@ -228,21 +462,30 @@ impl CachedSstMeta { self.parquet_metadata.clone() } + /// Returns the immutable parquet metadata size computed when it was decoded. + pub(crate) fn parquet_metadata_size(&self) -> usize { + self.parquet_metadata_size + } + pub(crate) fn region_metadata(&self) -> RegionMetadataRef { self.region_metadata.clone() } fn satisfies_page_index_policy(&self, requested: PageIndexPolicy) -> bool { - match requested { - PageIndexPolicy::Skip => true, - PageIndexPolicy::Optional => self.page_index_policy != PageIndexPolicy::Skip, - PageIndexPolicy::Required => self.page_index_policy == PageIndexPolicy::Required, - } + satisfies_page_index_policy(self.page_index_policy, requested) + } +} + +fn satisfies_page_index_policy(cached: PageIndexPolicy, requested: PageIndexPolicy) -> bool { + match requested { + PageIndexPolicy::Skip => true, + PageIndexPolicy::Optional => cached != PageIndexPolicy::Skip, + PageIndexPolicy::Required => cached == PageIndexPolicy::Required, } } fn infer_loaded_page_index_policy(parquet_metadata: &ParquetMetaData) -> PageIndexPolicy { - if parquet_metadata.column_index().is_some() || parquet_metadata.offset_index().is_some() { + if parquet_metadata.offset_index().is_some() { PageIndexPolicy::Optional } else { PageIndexPolicy::Skip @@ -270,12 +513,10 @@ fn strip_region_metadata_from_parquet(parquet_metadata: ParquetMetaData) -> Parq let mut builder = parquet_metadata.into_builder(); let row_groups = builder.take_row_groups(); - let column_index = builder.take_column_index(); let offset_index = builder.take_offset_index(); parquet::file::metadata::ParquetMetaDataBuilder::new(stripped_file_metadata) .set_row_groups(row_groups) - .set_column_index(column_index) .set_offset_index(offset_index) .build() } @@ -407,6 +648,16 @@ pub enum CacheStrategy { } impl CacheStrategy { + /// Returns whether the SST metadata cache is enabled for this strategy. + pub(crate) fn sst_meta_cache_enabled(&self) -> bool { + match self { + CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { + cache_manager.sst_meta_cache_enabled() + } + CacheStrategy::Disabled => false, + } + } + /// Gets fused SST metadata with cache metrics tracking. pub(crate) async fn get_sst_meta_data( &self, @@ -450,11 +701,16 @@ impl CacheStrategy { .map(|metadata| metadata.parquet_metadata()) } - /// Calls [CacheManager::put_sst_meta_data()]. - pub(crate) fn put_sst_meta_data(&self, file_id: RegionFileId, metadata: Arc) { + /// Puts compact and decoded forms of SST metadata into the shared cache capacity. + pub(crate) fn put_prepared_sst_meta( + &self, + file_id: RegionFileId, + metadata: PreparedSstMeta, + retain_decoded: bool, + ) { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { - cache_manager.put_sst_meta_data(file_id, metadata); + cache_manager.put_prepared_sst_meta(file_id, metadata, retain_decoded); } CacheStrategy::Disabled => {} } @@ -730,8 +986,10 @@ impl CacheStrategy { /// All caches are disabled by default. #[derive(Default)] pub struct CacheManager { - /// Cache for SST metadata. + /// Cache for compact, authoritative SST metadata. sst_meta_cache: Option, + /// Cache for decoded SST metadata, used only as an acceleration tier. + sst_decoded_meta_cache: Option, /// Cache for vectors. vector_cache: Option, /// Cache for SST byte ranges. @@ -777,40 +1035,82 @@ impl CacheManager { metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, ) -> Option> { - if let Some(metadata) = self.get_sst_meta_data_from_mem_cache(file_id, page_index_policy) { + let cache_key = SstMetaKey(file_id.region_id(), file_id.file_id()); + let compact = self + .get_compact_sst_meta(&cache_key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + let decoded = self + .get_decoded_sst_meta(&cache_key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + + if let Some(compact) = compact { + CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); metrics.mem_cache_hit += 1; - return Some(metadata); + if let Some(decoded) = decoded { + CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(); + return Some(decoded); + } + + CACHE_MISS.with_label_values(&[SST_META_DECODED_TYPE]).inc(); + match compact.decode().await { + Ok(decoded) => { + self.put_sst_meta_data(file_id, decoded.clone()); + return Some(decoded); + } + Err(err) => { + warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); + self.remove_parquet_meta_data(file_id); + } + } + } else if let Some(decoded) = decoded { + // Metadata produced by a new SST writer can enter the decoded tier before its compact + // representation is prepared. + CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); + CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(); + metrics.mem_cache_hit += 1; + return Some(decoded); + } else { + CACHE_MISS.with_label_values(&[SST_META_TYPE]).inc(); } let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); - if let Some(write_cache) = &self.write_cache - && let Some(metadata) = write_cache - .file_cache() - .get_sst_meta_data(key, metrics, page_index_policy) + if let Some(write_cache) = &self.write_cache { + let file_cache = write_cache.file_cache(); + if self.sst_meta_cache_enabled() { + if let Some(metadata) = file_cache + .get_sst_meta_data(key, metrics, page_index_policy) + .await + { + metrics.file_cache_hit += 1; + let decoded = metadata.decoded(); + match metadata { + SstMetaPreparation::Prepared(metadata) => { + self.put_prepared_sst_meta(file_id, metadata, true); + } + SstMetaPreparation::DecodedOnly { encoding_error, .. } => { + warn!( + encoding_error; + "Failed to encode file-cached SST metadata for memory cache, region_id: {}, file_id: {}", + file_id.region_id(), + file_id.file_id() + ); + } + } + return Some(decoded); + } + } else if let Some(decoded) = file_cache + .get_decoded_sst_meta_data(key, metrics, page_index_policy) .await - { - metrics.file_cache_hit += 1; - self.put_sst_meta_data(file_id, metadata.clone()); - return Some(metadata); + { + metrics.file_cache_hit += 1; + return Some(decoded); + } } metrics.cache_miss += 1; None } - /// Gets cached [ParquetMetaData] with metrics tracking. - /// Tries in-memory cache first, then file cache, updating metrics accordingly. - pub(crate) async fn get_parquet_meta_data( - &self, - file_id: RegionFileId, - metrics: &mut MetadataCacheMetrics, - page_index_policy: PageIndexPolicy, - ) -> Option> { - self.get_sst_meta_data(file_id, metrics, page_index_policy) - .await - .map(|metadata| metadata.parquet_metadata()) - } - /// Gets cached fused SST metadata from in-memory cache. /// This method does not perform I/O. pub(crate) fn get_sst_meta_data_from_mem_cache( @@ -818,12 +1118,11 @@ impl CacheManager { file_id: RegionFileId, page_index_policy: PageIndexPolicy, ) -> Option> { - self.sst_meta_cache.as_ref().and_then(|sst_meta_cache| { - let value = sst_meta_cache.get(&SstMetaKey(file_id.region_id(), file_id.file_id())); - let value = - value.filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); - update_hit_miss(value, SST_META_TYPE) - }) + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); + let value = self + .get_decoded_sst_meta(&key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + update_hit_miss(value, SST_META_DECODED_TYPE) } /// Gets cached [ParquetMetaData] from in-memory cache. @@ -838,15 +1137,65 @@ impl CacheManager { /// Puts fused SST metadata into the cache. pub(crate) fn put_sst_meta_data(&self, file_id: RegionFileId, metadata: Arc) { - if let Some(cache) = &self.sst_meta_cache { + if let Some(cache) = &self.sst_decoded_meta_cache { let key = SstMetaKey(file_id.region_id(), file_id.file_id()); CACHE_BYTES - .with_label_values(&[SST_META_TYPE]) - .add(meta_cache_weight(&key, &metadata).into()); + .with_label_values(&[SST_META_DECODED_TYPE]) + .add(decoded_meta_cache_weight(&key, &metadata).into()); cache.insert(key, metadata); } } + /// Puts a compact metadata entry and optionally retains its decoded acceleration entry. + pub(crate) fn put_prepared_sst_meta( + &self, + file_id: RegionFileId, + metadata: PreparedSstMeta, + retain_decoded: bool, + ) { + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); + if let Some(cache) = &self.sst_meta_cache { + CACHE_BYTES + .with_label_values(&[SST_META_TYPE]) + .add(meta_cache_weight(&key, &metadata.compact).into()); + cache.insert(key.clone(), metadata.compact); + } + if retain_decoded { + self.put_sst_meta_data(file_id, metadata.decoded); + } + } + + fn get_compact_sst_meta(&self, key: &SstMetaKey) -> Option> { + self.sst_meta_cache.as_ref()?.get(key) + } + + fn get_decoded_sst_meta(&self, key: &SstMetaKey) -> Option> { + self.sst_decoded_meta_cache.as_ref()?.get(key) + } + + /// Gets and decodes only the compact tier without promoting into the decoded cache. + /// + /// This is used by startup preloading so inspecting already-cached metadata cannot consume the + /// decoded reservation and prematurely stop compact preloading. + pub(crate) async fn get_compact_sst_meta_data( + &self, + file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + ) -> Option> { + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); + let compact = self + .get_compact_sst_meta(&key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy))?; + match compact.decode().await { + Ok(metadata) => Some(metadata), + Err(err) => { + warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); + self.remove_parquet_meta_data(file_id); + None + } + } + } + /// Puts [ParquetMetaData] into the cache. pub fn put_parquet_meta_data( &self, @@ -854,7 +1203,7 @@ impl CacheManager { metadata: Arc, region_metadata: Option, ) { - if self.sst_meta_cache.is_some() { + if self.sst_decoded_meta_cache.is_some() { let file_path = format!( "region_id={}, file_id={}", file_id.region_id(), @@ -877,17 +1226,24 @@ impl CacheManager { /// Removes [ParquetMetaData] from the cache. pub fn remove_parquet_meta_data(&self, file_id: RegionFileId) { + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); if let Some(cache) = &self.sst_meta_cache { - cache.remove(&SstMetaKey(file_id.region_id(), file_id.file_id())); + cache.remove(&key); + } + if let Some(cache) = &self.sst_decoded_meta_cache { + cache.remove(&key); } } - /// Returns the total weighted size of the in-memory SST meta cache. - pub(crate) fn sst_meta_cache_weighted_size(&self) -> u64 { - self.sst_meta_cache - .as_ref() - .map(|cache| cache.weighted_size()) - .unwrap_or(0) + /// Returns whether the authoritative SST metadata tier has reached its reservation. + pub(crate) fn sst_meta_cache_is_full(&self) -> bool { + let Some(cache) = &self.sst_meta_cache else { + return true; + }; + cache + .policy() + .max_capacity() + .is_some_and(|capacity| cache.weighted_size() >= capacity) } /// Returns true if the in-memory SST meta cache is enabled. @@ -1190,9 +1546,14 @@ impl CacheManagerBuilder { /// Builds the [CacheManager]. pub fn build(self) -> CacheManager { - let sst_meta_cache = (self.sst_meta_cache_size != 0).then(|| { + // Reserve half the configured capacity for the compact authoritative tier. This prevents + // decoded acceleration entries from evicting metadata that would require storage I/O to + // recover. Both reservations together equal the configured limit. + let compact_meta_capacity = self.sst_meta_cache_size.div_ceil(2); + let decoded_meta_capacity = self.sst_meta_cache_size / 2; + let sst_meta_cache = (compact_meta_capacity != 0).then(|| { Cache::builder() - .max_capacity(self.sst_meta_cache_size) + .max_capacity(compact_meta_capacity) .weigher(meta_cache_weight) .eviction_listener(|k, v, cause| { let size = meta_cache_weight(&k, &v); @@ -1205,6 +1566,21 @@ impl CacheManagerBuilder { }) .build() }); + let sst_decoded_meta_cache = (decoded_meta_capacity != 0).then(|| { + Cache::builder() + .max_capacity(decoded_meta_capacity) + .weigher(decoded_meta_cache_weight) + .eviction_listener(|k, v, cause| { + let size = decoded_meta_cache_weight(&k, &v); + CACHE_BYTES + .with_label_values(&[SST_META_DECODED_TYPE]) + .sub(size.into()); + CACHE_EVICTION + .with_label_values(&[SST_META_DECODED_TYPE, removal_cause_str(cause)]) + .inc(); + }) + .build() + }); let vector_cache = (self.vector_cache_size != 0).then(|| { Cache::builder() .max_capacity(self.vector_cache_size) @@ -1274,6 +1650,7 @@ impl CacheManagerBuilder { }); CacheManager { sst_meta_cache, + sst_decoded_meta_cache, vector_cache, page_cache, write_cache: self.write_cache, @@ -1295,10 +1672,15 @@ impl CacheManagerBuilder { } } -fn meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { - // We ignore the size of `Arc`. - let size = - k.estimated_size() + parquet_meta_size(&v.parquet_metadata) + v.region_metadata_weight; +fn meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { + // We ignore the size of `Arc`. Region metadata is already present in the compressed Parquet + // stream and is materialized only in the decoded acceleration tier. + let size = k.estimated_size() + mem::size_of::() + v.encoded_metadata.len(); + u32::try_from(size).unwrap_or(u32::MAX) +} + +fn decoded_meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { + let size = k.estimated_size() + v.parquet_metadata_size + v.region_metadata_weight; u32::try_from(size).unwrap_or(u32::MAX) } @@ -1677,7 +2059,9 @@ impl SelectorResultValue { } /// Maps (region id, file id) to fused SST metadata. -type SstMetaCache = Cache>; +type SstMetaCache = Cache>; +/// Maps (region id, file id) to decoded SST metadata retained as an acceleration tier. +type SstDecodedMetaCache = Cache>; /// Maps [Value] to a vector that holds this value repeatedly. /// /// e.g. `"hello" => ["hello", "hello", "hello"]` @@ -1695,6 +2079,7 @@ mod tests { use api::v1::index::{BloomFilterMeta, InvertedIndexMetas}; use datatypes::schema::ColumnSchema; use datatypes::vectors::Int64Vector; + use parquet::file::page_index::offset_index::{OffsetIndexMetaData, PageLocation}; use puffin::file_metadata::FileMetadata; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::storage::ColumnId; @@ -1715,6 +2100,7 @@ mod tests { async fn test_disable_cache() { let cache = CacheManager::default(); assert!(cache.sst_meta_cache.is_none()); + assert!(cache.sst_decoded_meta_cache.is_none()); assert!(cache.vector_cache.is_none()); assert!(cache.page_cache.is_none()); @@ -1725,7 +2111,7 @@ mod tests { cache.put_parquet_meta_data(file_id, metadata, None); assert!( cache - .get_parquet_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data(file_id, &mut metrics, Default::default()) .await .is_none() ); @@ -1754,6 +2140,30 @@ mod tests { assert!(cache.write_cache().is_none()); } + #[test] + fn test_sst_meta_cache_splits_capacity() { + let cache = CacheManager::builder().sst_meta_cache_size(101).build(); + + assert_eq!( + Some(51), + cache + .sst_meta_cache + .as_ref() + .unwrap() + .policy() + .max_capacity() + ); + assert_eq!( + Some(50), + cache + .sst_decoded_meta_cache + .as_ref() + .unwrap() + .policy() + .max_capacity() + ); + } + #[tokio::test] async fn test_parquet_meta_cache() { let cache = CacheManager::builder().sst_meta_cache_size(2000).build(); @@ -1762,7 +2172,7 @@ mod tests { let file_id = RegionFileId::new(region_id, FileId::random()); assert!( cache - .get_parquet_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data(file_id, &mut metrics, Default::default()) .await .is_none() ); @@ -1787,7 +2197,7 @@ mod tests { cache.remove_parquet_meta_data(file_id); assert!( cache - .get_parquet_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data(file_id, &mut metrics, Default::default()) .await .is_none() ); @@ -1810,6 +2220,109 @@ mod tests { assert!(Arc::ptr_eq(®ion_metadata, &cached.region_metadata())); } + #[tokio::test] + async fn test_compact_sst_meta_round_trip() { + let cache = CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .build(); + let region_id = RegionId::new(1, 1); + let file_id = RegionFileId::new(region_id, FileId::random()); + let (metadata, expected_region_metadata) = sst_parquet_meta(); + let metadata = Arc::unwrap_or_clone(metadata); + let offset_indexes = metadata + .row_groups() + .iter() + .map(|row_group| { + row_group + .columns() + .iter() + .map(|_| OffsetIndexMetaData { + page_locations: vec![PageLocation { + offset: 42, + compressed_page_size: 128, + first_row_index: 0, + }], + unencoded_byte_array_data_bytes: None, + }) + .collect() + }) + .collect(); + let metadata = metadata + .into_builder() + .set_offset_index(Some(offset_indexes)) + .build(); + let expected_rows = metadata.file_metadata().num_rows(); + let prepared = prepare_sst_meta("test.parquet", metadata, None, PageIndexPolicy::Required) + .await + .unwrap(); + let SstMetaPreparation::Prepared(prepared) = prepared else { + panic!("valid metadata should produce a compact cache entry"); + }; + + // Retain only the compact form so the lookup must exercise decompression and decoding. + cache.put_prepared_sst_meta(file_id, prepared, false); + let mut metrics = MetadataCacheMetrics::default(); + let cached = cache + .get_sst_meta_data(file_id, &mut metrics, PageIndexPolicy::Required) + .await + .unwrap(); + + assert_eq!(1, metrics.mem_cache_hit); + assert_eq!(0, metrics.cache_miss); + assert_eq!( + expected_rows, + cached.parquet_metadata.file_metadata().num_rows() + ); + assert!(cached.parquet_metadata.offset_index().is_some()); + assert_eq!(expected_region_metadata, cached.region_metadata()); + assert!( + cached + .parquet_metadata + .file_metadata() + .key_value_metadata() + .is_none_or(|key_values| key_values + .iter() + .all(|key_value| key_value.key != PARQUET_METADATA_KEY)) + ); + } + + #[test] + fn test_cache_encoding_failure_preserves_decoded_metadata() { + let (metadata, expected_region_metadata) = sst_parquet_meta(); + let expected_rows = metadata.file_metadata().num_rows(); + let cache_encoding: Result<(Bytes, usize)> = Err(UnexpectedSnafu { + reason: "injected compact metadata encoding failure", + } + .build()); + + let preparation = finish_sst_meta_preparation( + "test.parquet", + Arc::unwrap_or_clone(metadata), + None, + PageIndexPolicy::Skip, + cache_encoding, + ) + .unwrap(); + let SstMetaPreparation::DecodedOnly { + decoded, + encoding_error, + } = preparation + else { + panic!("cache encoding failure should return decoded-only metadata"); + }; + + assert_eq!( + expected_rows, + decoded.parquet_metadata.file_metadata().num_rows() + ); + assert_eq!(expected_region_metadata, decoded.region_metadata()); + assert!( + encoding_error + .to_string() + .contains("injected compact metadata encoding failure") + ); + } + #[tokio::test] async fn test_parquet_meta_cache_respects_page_index_policy() { let cache = CacheManager::builder().sst_meta_cache_size(2000).build(); @@ -1868,7 +2381,7 @@ mod tests { } #[test] - fn test_meta_cache_weight_accounts_for_decoded_region_metadata() { + fn test_decoded_meta_cache_weight_accounts_for_region_metadata() { let region_metadata = Arc::new(wide_region_metadata(128)); let json_len = region_metadata.to_json().unwrap().len(); let metadata = sst_parquet_meta_with_region_metadata(region_metadata.clone()); @@ -1879,15 +2392,17 @@ mod tests { assert!(cached.region_metadata_weight > json_len); assert_eq!( - meta_cache_weight(&key, &cached) as usize, - key.estimated_size() - + parquet_meta_size(&cached.parquet_metadata) - + cached.region_metadata_weight + decoded_meta_cache_weight(&key, &cached) as usize, + key.estimated_size() + cached.parquet_metadata_size + cached.region_metadata_weight + ); + assert_eq!( + cached.parquet_metadata_size, + parquet_meta_size(&cached.parquet_metadata) ); } #[test] - fn test_meta_cache_weight_saturates_on_overflow() { + fn test_decoded_meta_cache_weight_saturates_on_overflow() { let region_metadata = Arc::new(wide_region_metadata(1)); let metadata = sst_parquet_meta_with_region_metadata(region_metadata.clone()); let mut cached = @@ -1896,7 +2411,7 @@ mod tests { let cached = Arc::new(cached); let key = SstMetaKey(region_metadata.region_id, FileId::random()); - assert_eq!(u32::MAX, meta_cache_weight(&key, &cached)); + assert_eq!(u32::MAX, decoded_meta_cache_weight(&key, &cached)); } #[test] diff --git a/src/mito2/src/cache/file_cache.rs b/src/mito2/src/cache/file_cache.rs index 681b2f90e3f..64f2fad6f1b 100644 --- a/src/mito2/src/cache/file_cache.rs +++ b/src/mito2/src/cache/file_cache.rs @@ -34,7 +34,9 @@ use store_api::storage::{FileId, RegionId}; use tokio::sync::mpsc::{Sender, UnboundedReceiver}; use crate::access_layer::TempFileCleaner; -use crate::cache::{CachedSstMeta, FILE_TYPE, INDEX_TYPE}; +use crate::cache::{ + CachedSstMeta, FILE_TYPE, INDEX_TYPE, SstMetaPreparation, decode_sst_meta, prepare_sst_meta, +}; use crate::error::{self, OpenDalSnafu, Result}; use crate::metrics::{ CACHE_BYTES, CACHE_HIT, CACHE_MISS, WRITE_CACHE_DOWNLOAD_BYTES_TOTAL, @@ -619,35 +621,56 @@ impl FileCache { /// Get fused SST metadata from the file cache. /// If the file is not in the cache, or metadata loading/decoding fails, return None. + /// Compact cache encoding failures return decoded-only metadata to the caller. pub(crate) async fn get_sst_meta_data( &self, key: IndexKey, cache_metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, + ) -> Option { + let file_path = self.inner.cache_file_path(key); + let metadata = self + .get_parquet_meta_data(key, cache_metrics, page_index_policy) + .await?; + match prepare_sst_meta(&file_path, metadata, None, page_index_policy).await { + Ok(metadata) => Some(metadata), + Err(err) => { + CACHE_MISS + .with_label_values(&[key.file_type.metric_label()]) + .inc(); + warn!( + err; "Failed to prepare cached parquet metadata for key {:?}", + key + ); + None + } + } + } + + /// Gets decoded SST metadata without preparing an in-memory cache entry. + pub(crate) async fn get_decoded_sst_meta_data( + &self, + key: IndexKey, + cache_metrics: &mut MetadataCacheMetrics, + page_index_policy: PageIndexPolicy, ) -> Option> { let file_path = self.inner.cache_file_path(key); - self.get_parquet_meta_data(key, cache_metrics, page_index_policy) - .await - .and_then(|metadata| { - match CachedSstMeta::try_new_with_page_index_policy( - &file_path, - metadata, - None, - page_index_policy, - ) { - Ok(metadata) => Some(Arc::new(metadata)), - Err(err) => { - CACHE_MISS - .with_label_values(&[key.file_type.metric_label()]) - .inc(); - warn!( - err; "Failed to decode cached parquet metadata for key {:?}", - key - ); - None - } - } - }) + let metadata = self + .get_parquet_meta_data(key, cache_metrics, page_index_policy) + .await?; + match decode_sst_meta(&file_path, metadata, None, page_index_policy).await { + Ok(metadata) => Some(metadata), + Err(err) => { + CACHE_MISS + .with_label_values(&[key.file_type.metric_label()]) + .inc(); + warn!( + err; "Failed to decode cached parquet metadata for key {:?}", + key + ); + None + } + } } async fn get_reader(&self, file_path: &str) -> object_store::Result> { diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index 6b02a92811a..d186307c9b1 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -638,6 +638,7 @@ mod tests { .write_cache(Some(write_cache.clone())) .build(), ); + assert!(!cache_manager.sst_meta_cache_enabled()); // Create source let metadata = Arc::new(sst_region_metadata()); @@ -682,7 +683,7 @@ mod tests { let sst_info = sst_infos.remove(0); let write_parquet_metadata = sst_info.file_metadata.unwrap(); - // Read metadata from write cache + // Read metadata from write cache without preparing an in-memory metadata cache entry. let handle = sst_file_handle_with_file_id(sst_info.file_id, 0, 1000); let builder = ParquetReaderBuilder::new( data_home, diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index ad59e4fb94f..9388f943ceb 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -43,8 +43,8 @@ use store_api::storage::{ColumnId, RegionId}; use tokio::sync::Semaphore; use crate::access_layer::AccessLayer; -use crate::cache::CacheManagerRef; use crate::cache::file_cache::{FileCache, FileType, IndexKey}; +use crate::cache::{CacheManagerRef, SstMetaPreparation, prepare_sst_meta}; use crate::config::MitoConfig; use crate::engine::region_hook::RegionHookRef; use crate::error; @@ -1092,7 +1092,7 @@ async fn preload_parquet_meta_cache_for_files( ) -> usize { if !cache_manager.sst_meta_cache_enabled() || sst_meta_cache_capacity == 0 - || cache_manager.sst_meta_cache_weighted_size() >= sst_meta_cache_capacity + || cache_manager.sst_meta_cache_is_full() { return 0; } @@ -1104,30 +1104,60 @@ async fn preload_parquet_meta_cache_for_files( let mut loaded = 0usize; for file_handle in files { - // Stop when the shared SST meta cache is full. - if cache_manager.sst_meta_cache_weighted_size() >= sst_meta_cache_capacity { + // Stop when the protected authoritative tier is full. + if cache_manager.sst_meta_cache_is_full() { break; } let file_id = file_handle.file_id(); let mut cache_metrics = MetadataCacheMetrics::default(); if let Some(metadata) = cache_manager - .get_parquet_meta_data(file_id, &mut cache_metrics, PageIndexPolicy::Optional) + .get_compact_sst_meta_data(file_id, PageIndexPolicy::Optional) .await { if file_handle.primary_key_range().is_none() - && let Some(primary_key_range) = - extract_primary_key_range(&metadata, ®ion_metadata) + && let Some(primary_key_range) = extract_primary_key_range( + metadata.parquet_metadata().as_ref(), + ®ion_metadata, + ) { file_handle.set_primary_key_range(primary_key_range); } - // Metadata is either already in memory or loaded from file cache. - if cache_metrics.mem_cache_hit == 0 { - loaded += 1; - } continue; } + if let Some(write_cache) = cache_manager.write_cache() { + let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); + if let Some(metadata) = write_cache + .file_cache() + .get_sst_meta_data(key, &mut cache_metrics, PageIndexPolicy::Optional) + .await + { + let decoded = metadata.decoded(); + if file_handle.primary_key_range().is_none() + && let Some(primary_key_range) = extract_primary_key_range( + decoded.parquet_metadata().as_ref(), + ®ion_metadata, + ) + { + file_handle.set_primary_key_range(primary_key_range); + } + match metadata { + SstMetaPreparation::Prepared(metadata) => { + cache_manager.put_prepared_sst_meta(file_id, metadata, false); + loaded += 1; + } + SstMetaPreparation::DecodedOnly { encoding_error, .. } => warn!( + encoding_error; + "Failed to encode file-cached SST metadata during preload, region: {}, file: {}", + region_id, + file_id.file_id() + ), + } + continue; + } + } + if !allow_direct_load { continue; } @@ -1143,8 +1173,35 @@ async fn preload_parquet_meta_cache_for_files( { file_handle.set_primary_key_range(primary_key_range); } - cache_manager.put_parquet_meta_data(file_id, Arc::new(metadata), None); - loaded += 1; + match prepare_sst_meta( + &file_path, + metadata, + // The SST can predate schema changes, so decode its embedded region metadata + // instead of substituting the region's current schema. + None, + PageIndexPolicy::Optional, + ) + .await + { + Ok(SstMetaPreparation::Prepared(metadata)) => { + // Preload the compact canonical entry only. Query traffic promotes decoded + // entries into the separate acceleration tier according to actual demand. + cache_manager.put_prepared_sst_meta(file_id, metadata, false); + loaded += 1; + } + Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!( + encoding_error; + "Failed to encode preloaded parquet metadata, region: {}, file: {}", + region_id, + file_path + ), + Err(err) => { + warn!( + err; "Failed to prepare preloaded parquet metadata, region: {}, file: {}", + region_id, file_path + ); + } + } } Err(err) => { // Preloading is best-effort. Failure shouldn't affect region open. @@ -1257,7 +1314,7 @@ mod tests { use object_store::ObjectStore; use object_store::services::{Fs, Memory, S3}; use parquet::arrow::ArrowWriter; - use parquet::file::metadata::KeyValue; + use parquet::file::metadata::{KeyValue, PageIndexPolicy}; use parquet::file::properties::WriterProperties; use store_api::region_request::PathType; use store_api::storage::{FileId, RegionId}; @@ -1490,17 +1547,11 @@ mod tests { assert_eq!(loaded, 1); assert!( cache_manager - .get_parquet_meta_data_from_mem_cache(region_file_id) - .is_some() - ); - // The cached entry must carry the page index so that later `Optional` queries hit - // the in-memory cache instead of reloading metadata on demand. - assert!( - cache_manager - .get_sst_meta_data_from_mem_cache( + .get_compact_sst_meta_data( region_file_id, parquet::file::metadata::PageIndexPolicy::Optional, ) + .await .is_some() ); assert!(file_handle.primary_key_range().is_some()); @@ -1645,7 +1696,8 @@ mod tests { assert_eq!(loaded, 1); assert!( cache_manager - .get_parquet_meta_data_from_mem_cache(region_file_id) + .get_compact_sst_meta_data(region_file_id, PageIndexPolicy::Optional) + .await .is_some() ); } diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index db2ce371991..52c74ef70a6 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -332,7 +332,7 @@ impl FileRangeContext { /// Returns the estimated memory size of this context. /// Mainly accounts for the parquet metadata size. pub(crate) fn memory_size(&self) -> usize { - crate::cache::cache_size::parquet_meta_size(self.reader_builder.parquet_metadata()) + self.reader_builder.parquet_metadata_size() } } diff --git a/src/mito2/src/sst/parquet/metadata.rs b/src/mito2/src/sst/parquet/metadata.rs index 441be373b0c..9d38d6e2167 100644 --- a/src/mito2/src/sst/parquet/metadata.rs +++ b/src/mito2/src/sst/parquet/metadata.rs @@ -81,7 +81,10 @@ impl<'a> MetadataLoader<'a> { let file_size = self.get_file_size().await?; let reader = ParquetMetaDataReader::new() .with_prefetch_hint(Some(DEFAULT_PREFETCH_SIZE as usize)) - .with_page_index_policy(self.page_index_policy); + // Mito uses offset indexes to translate row selections into byte ranges. It does not + // consume Parquet column indexes, so decoding them only bloats the metadata cache. + .with_column_index_policy(PageIndexPolicy::Skip) + .with_offset_index_policy(self.page_index_policy); let num_reads = AtomicUsize::new(0); let bytes_read = AtomicU64::new(0); @@ -190,6 +193,7 @@ mod tests { }; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use datatypes::arrow::record_batch::RecordBatch; + use object_store::services::Memory; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use parquet::file::metadata::{KeyValue, ParquetMetaData}; @@ -199,12 +203,12 @@ mod tests { use crate::sst::parquet::PARQUET_METADATA_KEY; use crate::test_util::sst_util::sst_region_metadata; - fn build_test_metadata( + fn build_test_parquet_bytes( include_primary_key: bool, primary_keys: &[&[u8]], row_group_sizes: &[usize], stats_enabled: EnabledStatistics, - ) -> ParquetMetaData { + ) -> Vec { let total_rows = row_group_sizes.iter().sum::(); let mut fields = vec![Field::new("field", ArrowDataType::Int64, true)]; let mut columns: Vec = @@ -253,11 +257,44 @@ mod tests { } writer.close().unwrap(); - ParquetRecordBatchReaderBuilder::try_new(Bytes::from(parquet_bytes)) - .unwrap() - .metadata() - .as_ref() - .clone() + parquet_bytes + } + + fn build_test_metadata( + include_primary_key: bool, + primary_keys: &[&[u8]], + row_group_sizes: &[usize], + stats_enabled: EnabledStatistics, + ) -> ParquetMetaData { + ParquetRecordBatchReaderBuilder::try_new(Bytes::from(build_test_parquet_bytes( + include_primary_key, + primary_keys, + row_group_sizes, + stats_enabled, + ))) + .unwrap() + .metadata() + .as_ref() + .clone() + } + + #[tokio::test] + async fn test_metadata_loader_only_loads_offset_indexes() { + let parquet_bytes = build_test_parquet_bytes(false, &[], &[4], EnabledStatistics::Page); + let file_size = parquet_bytes.len() as u64; + let file_path = "test.parquet"; + let object_store = ObjectStore::new(Memory::default()).unwrap().finish(); + object_store.write(file_path, parquet_bytes).await.unwrap(); + + let mut loader = MetadataLoader::new(object_store, file_path, file_size); + loader.with_page_index_policy(PageIndexPolicy::Required); + let metadata = loader + .load(&mut MetadataCacheMetrics::default()) + .await + .unwrap(); + + assert!(metadata.column_index().is_none()); + assert!(metadata.offset_index().is_some()); } #[test] diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 6459a9877e8..87c521181c1 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -52,7 +52,7 @@ use table::predicate::Predicate; use self::stream::{NestedSchemaAligner, ProjectedRecordBatchStream}; use crate::cache::index::result_cache::PredicateKey; -use crate::cache::{CacheStrategy, CachedSstMeta}; +use crate::cache::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst_meta}; #[cfg(feature = "vector_index")] use crate::error::ApplyVectorIndexSnafu; use crate::error::{ @@ -393,6 +393,7 @@ impl ParquetReaderBuilder { ) .await?; let mut parquet_meta = sst_meta.parquet_metadata(); + let mut parquet_metadata_size = sst_meta.parquet_metadata_size(); let region_meta = sst_meta.region_metadata(); let region_partition_expr_str = self .expected_metadata @@ -504,6 +505,7 @@ impl ParquetReaderBuilder { ) .await?; parquet_meta = sst_meta.parquet_metadata(); + parquet_metadata_size = sst_meta.parquet_metadata_size(); cache_miss |= page_index_cache_miss; } @@ -551,6 +553,7 @@ impl ParquetReaderBuilder { file_handle: self.file_handle.clone(), file_path, parquet_meta, + parquet_metadata_size, arrow_metadata, output_schema, object_store: self.object_store.clone(), @@ -695,18 +698,34 @@ impl ParquetReaderBuilder { metadata_loader.with_page_index_policy(page_index_policy); let metadata = metadata_loader.load(cache_metrics).await?; - let metadata = Arc::new(CachedSstMeta::try_new_with_page_index_policy( - file_path, - metadata, - None, - page_index_policy, - )?); - // Cache the metadata. - self.cache_strategy - .put_sst_meta_data(file_id, metadata.clone()); + let decoded = if self.cache_strategy.sst_meta_cache_enabled() { + let metadata = prepare_sst_meta(file_path, metadata, None, page_index_policy).await?; + let decoded = metadata.decoded(); + match metadata { + SstMetaPreparation::Prepared(metadata) => { + self.cache_strategy + .put_prepared_sst_meta(file_id, metadata, true); + } + SstMetaPreparation::DecodedOnly { encoding_error, .. } => { + warn!( + encoding_error; + "Failed to encode SST metadata for cache, using decoded metadata for {}", + file_path + ); + } + } + decoded + } else { + Arc::new(CachedSstMeta::try_new_with_page_index_policy( + file_path, + metadata, + None, + page_index_policy, + )?) + }; cache_metrics.metadata_load_cost += start.elapsed(); - Ok((metadata, true)) + Ok((decoded, true)) } /// Computes row groups to read, along with their respective row selections. @@ -1735,6 +1754,8 @@ pub(crate) struct RowGroupReaderBuilder { file_path: String, /// Metadata of the parquet file. parquet_meta: Arc, + /// Immutable metadata size, computed once when the footer is decoded. + parquet_metadata_size: usize, /// Arrow reader metadata for building async stream. arrow_metadata: ArrowReaderMetadata, /// Projected output schema aligned with `projection.projected_root_presence`. @@ -1777,6 +1798,10 @@ impl RowGroupReaderBuilder { &self.parquet_meta } + pub(crate) fn parquet_metadata_size(&self) -> usize { + self.parquet_metadata_size + } + pub(crate) fn cache_strategy(&self) -> &CacheStrategy { &self.cache_strategy }