1mod otlp;
16
17pub use otlp::MetricInsertPreflight;
18
19use std::collections::BTreeMap;
20use std::sync::atomic::{AtomicU64, Ordering};
21use std::sync::{Arc, LazyLock};
22use std::time::{Duration, Instant};
23
24use ahash::{HashMap, HashMapExt, HashSet, HashSetExt};
25use api::v1::alter_table_expr::Kind;
26use api::v1::column_def::{options_from_compression, options_from_skipping, try_as_column_def};
27use api::v1::region::{
28 InsertRequest as RegionInsertRequest, InsertRequests as RegionInsertRequests,
29 RegionRequestHeader,
30};
31use api::v1::value::ValueData;
32use api::v1::{
33 AlterTableExpr, ColumnDataType, ColumnSchema, CreateTableExpr, InsertRequests,
34 RowInsertRequest, RowInsertRequests, Rows, SemanticType,
35};
36use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema};
37use arrow_schema::extension::ExtensionType;
38use catalog::CatalogManagerRef;
39use client::{OutputData, OutputMeta};
40use common_base::readable_size::ReadableSize;
41use common_catalog::consts::{
42 DEFAULT_PRIVATE_SCHEMA_NAME, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, SPAN_ID_COLUMN,
43 TIMESTAMP_END_COLUMN, TRACE_ID_COLUMN, TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY,
44 default_engine, is_ddl_reserved_table, trace_operations_table_name, trace_services_table_name,
45};
46use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME;
47use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME;
48use common_grpc_expr::util::ColumnExpr;
49use common_meta::cache::TableFlownodeSetCacheRef;
50use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME;
51use common_meta::node_manager::NodeManagerRef;
52use common_meta::peer::Peer;
53use common_meta::rpc::ddl::TriggerReason;
54use common_query::Output;
55use common_query::native_histogram::{is_native_histogram_value_type, native_histogram_value_type};
56use common_query::prelude::{
57 EXEMPLAR_LABELS_COLUMN, EXEMPLAR_METRIC_COLUMN, METRICS_EXEMPLARS_TABLE, greptime_timestamp,
58 greptime_value,
59};
60use common_telemetry::tracing_context::TracingContext;
61use common_telemetry::{debug, error, warn};
62use common_time::Timestamp;
63use common_time::timestamp::TimeUnit;
64use datatypes::schema::{
65 ColumnDefaultConstraint, ColumnEncoding, CompressionOptions, SkippingIndexOptions,
66};
67use futures_util::future;
68use meter_core::data::MeterRecord;
69use meter_macros::write_meter;
70use partition::manager::PartitionRuleManagerRef;
71use session::context::QueryContextRef;
72use snafu::ResultExt;
73use snafu::prelude::*;
74use sql::partition::partition_rule_for_hexstring;
75use sql::statements::create::Partitions;
76use sql::statements::insert::Insert;
77use store_api::metric_engine_consts::{
78 DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY,
79 METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
80};
81use store_api::mito_engine_options::{
82 APPEND_MODE_KEY, COMPACTION_OPTION_KEYS, COMPACTION_TWCS_PREFIX, COMPACTION_TYPE,
83 COMPACTION_TYPE_TWCS, MERGE_MODE_KEY, TTL_KEY,
84};
85use store_api::storage::consts::ReservedColumnId;
86use store_api::storage::{RegionId, TableId};
87use table::TableRef;
88use table::metadata::{TableInfo, TableInfoRef};
89use table::requests::{
90 AUTO_CREATE_TABLE_KEY, InsertRequest as TableInsertRequest, SEMANTIC_PER_TABLE_INDEX_KEY,
91 SEMANTIC_PIPELINE, TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1, TABLE_DATA_MODEL_TRACE_V2,
92 TRACE_TABLE_PARTITIONS_HINT_KEY, VALID_TABLE_OPTION_KEYS, is_semantic_option_key,
93 validate_semantic_option,
94};
95use table::table_reference::TableReference;
96use tokio_util::sync::CancellationToken;
97
98use crate::batcher::PendingRowsBatcher;
99use crate::error::{
100 CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, InvalidInsertRequestSnafu,
101 JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu, WriteRejectedSnafu,
102};
103use crate::expr_helper;
104use crate::region_req_factory::RegionRequestFactory;
105use crate::req_convert::common::preprocess_row_insert_requests;
106use crate::req_convert::insert::{
107 ColumnToRow, ImpureDefaultFiller, RowToRegion, StatementToRegion, TableToRegion,
108 fill_reqs_with_impure_default, rows_to_record_batch,
109};
110use crate::statement::StatementExecutor;
111use crate::write_dispatcher::{WriteDispatcher, WriteDispatcherRef};
112
113pub struct Inserter {
114 catalog_manager: CatalogManagerRef,
115 pub(crate) partition_manager: PartitionRuleManagerRef,
116 pub(crate) node_manager: NodeManagerRef,
117 pub(crate) write_dispatcher: WriteDispatcherRef,
118 pub(crate) table_flownode_set_cache: TableFlownodeSetCacheRef,
119 auto_create_table: bool,
123 pending_rows_batcher: Option<Arc<dyn PendingRowsBatcher>>,
124 mirror_pending_rows: Arc<AtomicU64>,
128}
129
130pub type InserterRef = Arc<Inserter>;
131
132#[derive(Clone)]
134pub enum AutoCreateTableType {
135 Logical(String),
137 Physical,
139 Exemplar,
141 Log,
143 LastNonNull,
145 Trace { alter_existing: bool },
147}
148
149impl AutoCreateTableType {
150 pub fn as_str(&self) -> &'static str {
151 match self {
152 AutoCreateTableType::Logical(_) => "logical",
153 AutoCreateTableType::Physical => "physical",
154 AutoCreateTableType::Exemplar => "exemplar",
155 AutoCreateTableType::Log => "log",
156 AutoCreateTableType::LastNonNull => "last_non_null",
157 AutoCreateTableType::Trace { .. } => "trace",
158 }
159 }
160
161 fn alter_existing(&self) -> bool {
162 !matches!(
163 self,
164 Self::Trace {
165 alter_existing: false
166 }
167 )
168 }
169}
170
171#[derive(Clone, Default)]
178pub struct InstantAndNormalInsertRequests {
179 pub normal_requests: RegionInsertRequests,
181 pub instant_requests: RegionInsertRequests,
184}
185
186#[derive(Default)]
188pub struct PreparedMetricInserts {
189 requests: InstantAndNormalInsertRequests,
190 table_infos: HashMap<TableId, TableInfoRef>,
191}
192
193impl PreparedMetricInserts {
194 pub fn source_tables(&self) -> impl Iterator<Item = &TableInfoRef> {
196 self.table_infos.values()
197 }
198}
199
200impl Inserter {
201 pub async fn can_batch_metric_rows(
205 &self,
206 requests: &RowInsertRequests,
207 ctx: &QueryContextRef,
208 physical_table: &str,
209 ) -> Result<bool> {
210 if self.auto_create_disabled_reason(ctx)?.is_some() || ctx.extension(TTL_KEY).is_some() {
211 return Ok(false);
212 }
213 for request in &requests.inserts {
214 if request.rows.as_ref().is_some_and(|rows| {
219 rows.schema.iter().any(|column| {
220 column.datatype_extension.is_some()
221 || !matches!(
222 ColumnDataType::try_from(column.datatype),
223 Ok(ColumnDataType::TimestampSecond
224 | ColumnDataType::TimestampMillisecond
225 | ColumnDataType::TimestampMicrosecond
226 | ColumnDataType::TimestampNanosecond
227 | ColumnDataType::Float64
228 | ColumnDataType::String)
229 )
230 })
231 }) {
232 return Ok(false);
233 }
234
235 let Some(table) = self
236 .get_table(
237 ctx.current_catalog(),
238 &ctx.current_schema(),
239 &request.table_name,
240 )
241 .await?
242 else {
243 continue;
244 };
245 let info = table.table_info();
246 if info.meta.engine != METRIC_ENGINE_NAME
247 || info.is_ttl_instant_table()
248 || info
249 .meta
250 .options
251 .extra_options
252 .get(LOGICAL_TABLE_METADATA_KEY)
253 .map(String::as_str)
254 != Some(physical_table)
255 || info
256 .meta
257 .schema
258 .column_schemas()
259 .iter()
260 .any(|column| column.default_constraint().is_some())
261 {
262 return Ok(false);
263 }
264 if info
267 .meta
268 .primary_key_indices
269 .iter()
270 .any(|&index| !info.meta.schema.column_schemas()[index].is_nullable())
271 {
272 return Ok(false);
273 }
274 match self.table_flownode_set_cache.get(info.table_id()).await {
277 Ok(None) => {}
278 Ok(Some(flows)) if flows.is_empty() => {}
279 _ => return Ok(false),
280 }
281 }
282 Ok(true)
283 }
284
285 pub async fn meter_row_inserts(
288 requests: &mut RowInsertRequests,
289 ctx: &QueryContextRef,
290 ) -> Result<u64> {
291 let metered = InstantAndNormalInsertRequests {
292 normal_requests: RegionInsertRequests {
293 requests: requests
294 .inserts
295 .iter_mut()
296 .map(|request| RegionInsertRequest {
297 rows: request.rows.take(),
298 ..Default::default()
299 })
300 .collect(),
301 },
302 instant_requests: RegionInsertRequests::default(),
303 };
304 let cost = write_meter!(
305 ctx.current_catalog(),
306 ctx.current_schema(),
307 metered,
308 ctx.write_rows_to_admit(
309 ctx.current_catalog(),
310 &ctx.current_schema(),
311 count_insert_rows(&metered)?
312 ),
313 ctx.channel() as u8
314 )
315 .await
316 .context(WriteRejectedSnafu);
317 for (request, region) in requests
318 .inserts
319 .iter_mut()
320 .zip(metered.normal_requests.requests)
321 {
322 request.rows = region.rows;
323 }
324 cost
325 }
326
327 pub fn new(
329 catalog_manager: CatalogManagerRef,
330 partition_manager: PartitionRuleManagerRef,
331 node_manager: NodeManagerRef,
332 table_flownode_set_cache: TableFlownodeSetCacheRef,
333 auto_create_table: bool,
334 ) -> Self {
335 Self {
336 write_dispatcher: Arc::new(WriteDispatcher::new(
337 node_manager.clone(),
338 partition_manager.clone(),
339 )),
340 catalog_manager,
341 partition_manager,
342 node_manager,
343 table_flownode_set_cache,
344 auto_create_table,
345 pending_rows_batcher: None,
346 mirror_pending_rows: Arc::new(AtomicU64::new(0)),
347 }
348 }
349
350 pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self {
352 self.write_dispatcher = dispatcher;
353 self
354 }
355
356 pub fn with_pending_rows_batcher(
358 mut self,
359 batcher: Option<Arc<dyn PendingRowsBatcher>>,
360 ) -> Self {
361 self.pending_rows_batcher = batcher;
362 self
363 }
364
365 pub async fn handle_column_inserts(
366 &self,
367 requests: InsertRequests,
368 ctx: QueryContextRef,
369 statement_executor: &StatementExecutor,
370 ) -> Result<Output> {
371 let row_inserts = ColumnToRow::convert(requests)?;
372 self.handle_row_inserts(row_inserts, ctx, statement_executor, false, false)
373 .await
374 }
375
376 pub async fn handle_row_inserts(
378 &self,
379 mut requests: RowInsertRequests,
380 ctx: QueryContextRef,
381 statement_executor: &StatementExecutor,
382 accommodate_existing_schema: bool,
383 is_single_value: bool,
384 ) -> Result<Output> {
385 preprocess_row_insert_requests(&mut requests.inserts)?;
386 self.handle_row_inserts_with_create_type(
387 requests,
388 ctx,
389 statement_executor,
390 AutoCreateTableType::Physical,
391 accommodate_existing_schema,
392 is_single_value,
393 )
394 .await
395 }
396
397 pub async fn handle_exemplar_inserts(
399 &self,
400 mut requests: RowInsertRequests,
401 ctx: QueryContextRef,
402 statement_executor: &StatementExecutor,
403 sources: &[TableInfoRef],
404 ) -> Result<Output> {
405 ensure!(
406 requests
407 .inserts
408 .iter()
409 .all(|r| r.table_name == METRICS_EXEMPLARS_TABLE),
410 InvalidInsertRequestSnafu {
411 reason: "invalid exemplar destination"
412 }
413 );
414 self.bind_exemplar_table_ids(&mut requests, &ctx, sources)
415 .await?;
416 preprocess_row_insert_requests(&mut requests.inserts)?;
417 self.handle_row_inserts_with_create_type(
418 requests,
419 ctx,
420 statement_executor,
421 AutoCreateTableType::Exemplar,
422 false,
423 false,
424 )
425 .await
426 }
427
428 async fn bind_exemplar_table_ids(
429 &self,
430 requests: &mut RowInsertRequests,
431 ctx: &QueryContextRef,
432 sources: &[TableInfoRef],
433 ) -> Result<()> {
434 validate_column_count_match(requests)?;
435 let mut tables = HashMap::new();
436 for info in sources.iter().filter(|info| {
437 info.catalog_name == ctx.current_catalog() && info.schema_name == ctx.current_schema()
438 }) {
439 if let Some(previous) = tables.insert(info.name.clone(), info.clone()) {
440 ensure!(
441 previous.table_id() == info.table_id(),
442 InvalidInsertRequestSnafu {
443 reason: format!(
444 "exemplar source {} changed during write preparation",
445 info.name
446 ),
447 }
448 );
449 }
450 }
451 for request in &mut requests.inserts {
452 let Some(rows) = request.rows.as_mut() else {
453 continue;
454 };
455 ensure!(
456 !rows
457 .schema
458 .iter()
459 .any(|c| c.column_name == DATA_SCHEMA_TABLE_ID_COLUMN_NAME),
460 InvalidInsertRequestSnafu {
461 reason: "exemplar table IDs must be resolved from their source tables",
462 }
463 );
464 let metric_index = rows
465 .schema
466 .iter()
467 .position(|c| c.column_name == EXEMPLAR_METRIC_COLUMN)
468 .context(InvalidInsertRequestSnafu {
469 reason: "missing exemplar source metric",
470 })?;
471 for row in &mut rows.rows {
472 let Some(ValueData::StringValue(name)) = &row.values[metric_index].value_data
473 else {
474 return InvalidInsertRequestSnafu {
475 reason: "invalid exemplar source metric",
476 }
477 .fail();
478 };
479 if !tables.contains_key(name) {
480 let table = self.get_table(ctx.current_catalog(), &ctx.current_schema(), name).await?
481 .with_context(|| InvalidInsertRequestSnafu {
482 reason: format!("exemplar source {name} must be an existing logical metric-engine table"),
483 })?;
484 tables.insert(name.clone(), table.table_info());
485 }
486 let info = &tables[name];
487 ensure!(
488 info.meta.engine == METRIC_ENGINE_NAME
489 && info
490 .meta
491 .options
492 .extra_options
493 .contains_key(LOGICAL_TABLE_METADATA_KEY),
494 InvalidInsertRequestSnafu {
495 reason: format!(
496 "exemplar source {name} must be an existing logical metric-engine table"
497 ),
498 }
499 );
500 row.values.push(ValueData::U32Value(info.table_id()).into());
501 }
502 rows.schema.push(api::v1::helper::tag_column_schema(
503 DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
504 ColumnDataType::Uint32,
505 ));
506 }
507 Ok(())
508 }
509
510 pub async fn handle_log_inserts(
512 &self,
513 requests: RowInsertRequests,
514 ctx: QueryContextRef,
515 statement_executor: &StatementExecutor,
516 ) -> Result<Output> {
517 self.handle_row_inserts_with_create_type(
518 requests,
519 ctx,
520 statement_executor,
521 AutoCreateTableType::Log,
522 false,
523 false,
524 )
525 .await
526 }
527
528 pub async fn handle_trace_inserts(
529 &self,
530 requests: RowInsertRequests,
531 ctx: QueryContextRef,
532 statement_executor: &StatementExecutor,
533 ) -> Result<Output> {
534 self.handle_row_inserts_with_create_type(
535 requests,
536 ctx,
537 statement_executor,
538 AutoCreateTableType::Trace {
539 alter_existing: true,
540 },
541 false,
542 false,
543 )
544 .await
545 }
546
547 pub async fn handle_last_non_null_inserts(
549 &self,
550 requests: RowInsertRequests,
551 ctx: QueryContextRef,
552 statement_executor: &StatementExecutor,
553 accommodate_existing_schema: bool,
554 is_single_value: bool,
555 ) -> Result<Output> {
556 self.handle_row_inserts_with_create_type(
557 requests,
558 ctx,
559 statement_executor,
560 AutoCreateTableType::LastNonNull,
561 accommodate_existing_schema,
562 is_single_value,
563 )
564 .await
565 }
566
567 async fn handle_row_inserts_with_create_type(
569 &self,
570 mut requests: RowInsertRequests,
571 ctx: QueryContextRef,
572 statement_executor: &StatementExecutor,
573 create_type: AutoCreateTableType,
574 accommodate_existing_schema: bool,
575 is_single_value: bool,
576 ) -> Result<Output> {
577 let skip_wal = ctx.skip_wal();
578 let align_time_index_unit =
579 matches!(create_type, AutoCreateTableType::Exemplar).then_some(TimeUnit::Nanosecond);
580
581 let batcher = self
582 .pending_rows_batcher
583 .as_ref()
584 .filter(|_| ctx.batching_enabled());
585
586 requests.inserts.retain(|req| {
588 req.rows
589 .as_ref()
590 .map(|r| !r.rows.is_empty())
591 .unwrap_or_default()
592 });
593 validate_column_count_match(&requests)?;
594
595 let CreateAlterTableResult {
596 instant_table_ids,
597 table_infos,
598 } = self
599 .create_or_alter_tables_on_demand(
600 &mut requests,
601 &ctx,
602 create_type,
603 statement_executor,
604 accommodate_existing_schema,
605 is_single_value,
606 align_time_index_unit,
607 )
608 .await?;
609
610 if let Some(batcher) = batcher
613 && instant_table_ids.is_empty()
614 && table_infos
615 .values()
616 .all(|info| info.meta.engine == default_engine())
617 {
618 return self
619 .submit_pending_rows(requests, table_infos, ctx, batcher)
620 .await;
621 }
622
623 let name_to_info = table_infos
624 .values()
625 .map(|info| (info.name.clone(), info.clone()))
626 .collect::<HashMap<_, _>>();
627 let inserts = RowToRegion::new(
628 name_to_info,
629 instant_table_ids,
630 self.partition_manager.as_ref(),
631 )
632 .convert(requests, skip_wal)
633 .await?;
634
635 self.do_request(inserts, &table_infos, &ctx).await
636 }
637
638 async fn submit_pending_rows(
639 &self,
640 mut requests: RowInsertRequests,
641 table_infos: HashMap<TableId, Arc<TableInfo>>,
642 ctx: QueryContextRef,
643 batcher: &Arc<dyn PendingRowsBatcher>,
644 ) -> Result<Output> {
645 requests.inserts.retain(|request| {
648 request
649 .rows
650 .as_ref()
651 .is_some_and(|rows| !rows.rows.is_empty())
652 });
653 let by_name = table_infos
654 .values()
655 .map(|info| (info.name.as_str(), info))
656 .collect::<HashMap<_, _>>();
657 let mut prepared = Vec::with_capacity(requests.inserts.len());
658 for request in &mut requests.inserts {
659 let table_info =
660 by_name
661 .get(request.table_name.as_str())
662 .context(TableNotFoundSnafu {
663 table_name: &request.table_name,
664 })?;
665 let Some(rows) = &mut request.rows else {
666 continue;
667 };
668 ImpureDefaultFiller::new((*table_info).clone())?.fill_rows(rows);
669 let batch = rows_to_record_batch(rows, table_info)?;
670 prepared.push(((*table_info).clone(), batch));
671 }
672 let metered = InstantAndNormalInsertRequests {
675 normal_requests: RegionInsertRequests {
676 requests: requests
677 .inserts
678 .into_iter()
679 .map(|request| RegionInsertRequest {
680 rows: request.rows,
681 ..Default::default()
682 })
683 .collect(),
684 },
685 instant_requests: RegionInsertRequests::default(),
686 };
687 let table_info = table_infos.values().next();
688 let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str());
689 let schema =
690 table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone());
691 let write_cost = write_meter!(
692 catalog,
693 &schema,
694 metered,
695 ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&metered)?),
696 ctx.channel() as u8
697 )
698 .await
699 .context(WriteRejectedSnafu)?;
700 prepared.retain(|(_, batch)| batch.num_rows() != 0);
701 let results = if prepared.is_empty() {
702 Vec::new()
703 } else {
704 let permit = batcher.acquire().await?;
706 let submissions = prepared.into_iter().map(|(info, batch)| {
707 let mut target_ctx = ctx.fork();
709 target_ctx.set_current_catalog(&info.catalog_name);
710 target_ctx.set_current_schema(&info.schema_name);
711 batcher.submit(info, batch, Arc::new(target_ctx), permit.clone())
712 });
713 future::join_all(submissions).await
715 };
716 let mut affected_rows = 0;
717 let mut meta = OutputMeta::new_with_cost(write_cost as _);
718 let mut first_error = None;
719 for result in results {
720 match result {
721 Ok(output) => {
722 affected_rows += output.extract_rows_and_cost().0;
723 meta.write_completions.extend(output.meta.write_completions);
724 }
725 Err(error) if ctx.confirmed_writes() => {
726 meta.write_completions
727 .push(common_query::WriteCompletion::with_progress(async move {
728 common_query::WriteProgress::failed(common_error::ext::BoxedError::new(
729 error,
730 ))
731 }));
732 }
733 Err(error) => {
734 first_error.get_or_insert(error);
735 }
736 }
737 }
738 if let Some(error) = first_error {
739 for completion in meta.write_completions {
740 let _ = completion.wait().await;
741 }
742 return Err(error);
743 }
744 Ok(Output::new(OutputData::AffectedRows(affected_rows), meta))
745 }
746
747 pub async fn handle_metric_row_inserts(
752 &self,
753 requests: RowInsertRequests,
754 ctx: QueryContextRef,
755 statement_executor: &StatementExecutor,
756 physical_table: String,
757 ) -> Result<Output> {
758 let prepared = self
759 .prepare_metric_row_inserts(requests, ctx.clone(), statement_executor, physical_table)
760 .await?;
761 self.handle_prepared_metric_inserts(prepared, ctx).await
762 }
763
764 pub async fn prepare_metric_row_inserts(
766 &self,
767 mut requests: RowInsertRequests,
768 ctx: QueryContextRef,
769 statement_executor: &StatementExecutor,
770 physical_table: String,
771 ) -> Result<PreparedMetricInserts> {
772 let skip_wal = ctx.skip_wal();
773
774 requests.inserts.retain(|req| {
776 req.rows
777 .as_ref()
778 .map(|r| !r.rows.is_empty())
779 .unwrap_or_default()
780 });
781 validate_column_count_match(&requests)?;
782
783 let mut has_new_tables = false;
785 for request in &requests.inserts {
786 if self
787 .get_table(
788 ctx.current_catalog(),
789 &ctx.current_schema(),
790 &request.table_name,
791 )
792 .await?
793 .is_none()
794 {
795 has_new_tables = true;
796 break;
797 }
798 }
799 let new_table_unit = if has_new_tables {
800 let physical = self
801 .create_physical_table_on_demand(&ctx, physical_table.clone(), statement_executor)
802 .await?;
803 table_time_index_unit(&physical)
804 } else {
805 None
806 };
807 let CreateAlterTableResult {
808 instant_table_ids,
809 table_infos,
810 } = self
811 .create_or_alter_tables_on_demand(
812 &mut requests,
813 &ctx,
814 AutoCreateTableType::Logical(physical_table.clone()),
815 statement_executor,
816 true,
817 true,
818 new_table_unit,
819 )
820 .await?;
821 let name_to_info = table_infos
822 .values()
823 .map(|info| (info.name.clone(), info.clone()))
824 .collect::<HashMap<_, _>>();
825 let inserts = RowToRegion::new(name_to_info, instant_table_ids, &self.partition_manager)
826 .convert(requests, skip_wal)
827 .await?;
828
829 Ok(PreparedMetricInserts {
830 requests: inserts,
831 table_infos,
832 })
833 }
834
835 pub async fn handle_prepared_metric_inserts(
837 &self,
838 prepared: PreparedMetricInserts,
839 ctx: QueryContextRef,
840 ) -> Result<Output> {
841 self.do_request(prepared.requests, &prepared.table_infos, &ctx)
842 .await
843 }
844
845 fn table_batcher(
846 &self,
847 table_info: &TableInfoRef,
848 ctx: &QueryContextRef,
849 ) -> Option<&Arc<dyn PendingRowsBatcher>> {
850 self.pending_rows_batcher.as_ref().filter(|_| {
851 ctx.batching_enabled()
852 && !table_info.is_ttl_instant_table()
853 && table_info.meta.engine == default_engine()
854 })
855 }
856
857 async fn submit_table_rows(
858 &self,
859 rows: Rows,
860 table_info: TableInfoRef,
861 ctx: QueryContextRef,
862 batcher: &Arc<dyn PendingRowsBatcher>,
863 ) -> Result<Output> {
864 let requests = RowInsertRequests {
865 inserts: vec![RowInsertRequest {
866 table_name: table_info.name.clone(),
867 rows: Some(rows),
868 }],
869 };
870 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info)]);
871 self.submit_pending_rows(requests, table_infos, ctx, batcher)
872 .await
873 }
874
875 pub async fn handle_table_insert(
876 &self,
877 request: TableInsertRequest,
878 ctx: QueryContextRef,
879 ) -> Result<Output> {
880 let catalog = request.catalog_name.as_str();
881 let schema = request.schema_name.as_str();
882 let table_name = request.table_name.as_str();
883 let table = self.get_table(catalog, schema, table_name).await?;
884 let table = table.with_context(|| TableNotFoundSnafu {
885 table_name: common_catalog::format_full_table_name(catalog, schema, table_name),
886 })?;
887 let table_info = table.table_info();
888
889 let converter = TableToRegion::new(&table_info, &self.partition_manager);
890 let skip_wal = request.skip_wal;
891 let rows = converter.prepare(request)?;
892 if let Some(batcher) = self.table_batcher(&table_info, &ctx) {
893 return self.submit_table_rows(rows, table_info, ctx, batcher).await;
894 }
895 let inserts = converter.partition(rows, skip_wal).await?;
896
897 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]);
898
899 self.do_request(inserts, &table_infos, &ctx).await
900 }
901
902 pub async fn handle_statement_insert(
903 &self,
904 insert: &Insert,
905 ctx: &QueryContextRef,
906 ) -> Result<Output> {
907 let converter =
908 StatementToRegion::new(self.catalog_manager.as_ref(), &self.partition_manager, ctx);
909 let (rows, table_info) = converter.prepare(insert, ctx).await?;
910 if let Some(batcher) = self.table_batcher(&table_info, ctx) {
911 return self
912 .submit_table_rows(rows, table_info, ctx.clone(), batcher)
913 .await;
914 }
915 let inserts = converter.partition(rows, table_info.clone(), ctx).await?;
916
917 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]);
918
919 self.do_request(inserts, &table_infos, ctx).await
920 }
921}
922
923pub async fn admit_write(rows: u64, ctx: &QueryContextRef) -> Result<QueryContextRef> {
926 write_meter!(MeterRecord::new(
929 ctx.current_catalog().to_string(),
930 ctx.current_schema(),
931 0,
932 ctx.write_rows_to_admit(ctx.current_catalog(), &ctx.current_schema(), rows),
933 ctx.channel() as u8,
934 ))
935 .await
936 .context(WriteRejectedSnafu)?;
937 Ok(Arc::new(ctx.with_write_admission()))
938}
939
940pub async fn admit_row_insert_batches(
943 batches: &mut [(QueryContextRef, RowInsertRequests)],
944) -> Result<()> {
945 let mut totals = BTreeMap::<_, (QueryContextRef, u64)>::new();
946 for (ctx, requests) in batches.iter() {
947 let catalog = ctx.current_catalog();
948 let schema = ctx.current_schema();
949
950 if ctx.write_rows_to_admit(catalog, &schema, 1) == 0 {
951 continue;
952 }
953 let (_, total) = totals
954 .entry((catalog.to_string(), schema.clone()))
955 .or_insert_with(|| (ctx.clone(), 0));
956 for rows in requests.inserts.iter().filter_map(|r| r.rows.as_ref()) {
957 *total =
958 total
959 .checked_add(rows.rows.len() as u64)
960 .context(InvalidInsertRequestSnafu {
961 reason: "Insert row count exceeds u64::MAX",
962 })?;
963 }
964 }
965 for (ctx, rows) in totals.values() {
966 admit_write(*rows, ctx).await?;
967 }
968 for (ctx, _) in batches {
969 *ctx = Arc::new(ctx.with_write_admission());
970 }
971 Ok(())
972}
973
974fn count_insert_rows(requests: &InstantAndNormalInsertRequests) -> Result<u64> {
975 requests
976 .normal_requests
977 .requests
978 .iter()
979 .chain(&requests.instant_requests.requests)
980 .filter_map(|request| request.rows.as_ref())
981 .try_fold(0u64, |total, rows| {
982 total
983 .checked_add(rows.rows.len() as u64)
984 .context(InvalidInsertRequestSnafu {
985 reason: "Insert row count exceeds u64::MAX",
986 })
987 })
988}
989
990impl Inserter {
991 async fn do_request(
992 &self,
993 requests: InstantAndNormalInsertRequests,
994 table_infos: &HashMap<TableId, Arc<TableInfo>>,
995 ctx: &QueryContextRef,
996 ) -> Result<Output> {
997 let requests = fill_reqs_with_impure_default(table_infos, requests)?;
999
1000 let table_info = table_infos.values().next();
1003 let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str());
1004 let schema =
1005 table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone());
1006 let write_cost = write_meter!(
1007 catalog,
1008 schema.clone(),
1009 requests,
1010 ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&requests)?),
1011 ctx.channel() as u8
1012 )
1013 .await
1014 .context(WriteRejectedSnafu)?;
1015 let request_factory = RegionRequestFactory::new(RegionRequestHeader {
1016 tracing_context: TracingContext::from_current_span().to_w3c(),
1017 dbname: ctx.get_db_string(),
1018 ..Default::default()
1019 });
1020
1021 let InstantAndNormalInsertRequests {
1022 normal_requests,
1023 instant_requests,
1024 } = requests;
1025
1026 let flow_mirror_task = FlowMirrorTask::new(
1028 &self.table_flownode_set_cache,
1029 normal_requests
1030 .requests
1031 .iter()
1032 .chain(instant_requests.requests.iter()),
1033 )
1034 .await?;
1035 let has_instant_rows = instant_requests.requests.iter().any(|request| {
1036 request
1037 .rows
1038 .as_ref()
1039 .is_some_and(|rows| !rows.rows.is_empty())
1040 });
1041 flow_mirror_task.detach(
1042 self.node_manager.clone(),
1043 self.mirror_pending_rows.clone(),
1044 has_instant_rows,
1045 )?;
1046
1047 let dispatcher = self.write_dispatcher.clone();
1048 if ctx.confirmed_writes() {
1049 let writes = normal_requests
1050 .requests
1051 .into_iter()
1052 .map(|request| {
1053 let rows = request.rows.as_ref().map_or(0, |rows| rows.rows.len());
1054 crate::write_dispatcher::TrackedWrite {
1055 request: request_factory.build_insert(RegionInsertRequests {
1056 requests: vec![request],
1057 }),
1058 owners: vec![0; rows],
1059 }
1060 })
1061 .collect();
1062 let (writes, progress) = dispatcher.prepare_tracked(writes, 1).await;
1063 let cancellation = CancellationToken::new();
1064 let _guard = cancellation.clone().drop_guard();
1065 let counter = ctx.write_progress().cloned();
1066 let progress = common_runtime::spawn_global(async move {
1067 let progress = dispatcher
1068 .execute_tracked(writes, progress, None, Some(&cancellation))
1069 .await;
1070 if let Some(counter) = counter {
1071 counter.add_written(progress.iter().map(|part| part.affected_rows).sum());
1072 }
1073 progress
1074 })
1075 .await
1076 .context(JoinTaskSnafu)?
1077 .into_iter()
1078 .fold(
1079 common_query::WriteProgress::default(),
1080 |mut result, part| {
1081 result.merge(part);
1082 result
1083 },
1084 );
1085 crate::metrics::DIST_INGEST_ROW_COUNT
1086 .with_label_values(&[ctx.get_db_string().as_str()])
1087 .inc_by(progress.affected_rows as u64);
1088 let mut meta = OutputMeta::new_with_cost(write_cost as _);
1089 let rows = progress.affected_rows;
1090 meta.write_completions
1091 .push(common_query::WriteCompletion::with_progress(async move {
1092 progress
1093 }));
1094 return Ok(Output::new(OutputData::AffectedRows(rows), meta));
1095 }
1096 let request = request_factory.build_insert(normal_requests);
1097 let writes = dispatcher.prepare(vec![request]).await?;
1098 let cancellation = CancellationToken::new();
1100 let _guard = cancellation.clone().drop_guard();
1101 let progress = ctx.write_progress().cloned();
1102 let affected_rows = common_runtime::spawn_global(async move {
1103 dispatcher
1104 .execute_with_progress(writes, None, Some(&cancellation), progress)
1105 .await
1106 })
1107 .await
1108 .context(JoinTaskSnafu)?
1109 .context(RequestInsertsSnafu)?
1110 .affected_rows;
1111 crate::metrics::DIST_INGEST_ROW_COUNT
1112 .with_label_values(&[ctx.get_db_string().as_str()])
1113 .inc_by(affected_rows as u64);
1114 Ok(Output::new(
1115 OutputData::AffectedRows(affected_rows),
1116 OutputMeta::new_with_cost(write_cost as _),
1117 ))
1118 }
1119
1120 fn auto_create_disabled_reason(&self, ctx: &QueryContextRef) -> Result<Option<&'static str>> {
1124 let auto_create_table_hint = ctx
1125 .extension(AUTO_CREATE_TABLE_KEY)
1126 .map(|v| v.parse::<bool>())
1127 .transpose()
1128 .map_err(|_| {
1129 InvalidInsertRequestSnafu {
1130 reason: "`auto_create_table` hint must be a boolean",
1131 }
1132 .build()
1133 })?
1134 .unwrap_or(true);
1135 Ok(if !self.auto_create_table {
1136 Some("auto-create table is disabled by frontend config")
1137 } else if !auto_create_table_hint {
1138 Some("`auto_create_table` hint is disabled")
1139 } else {
1140 None
1141 })
1142 }
1143
1144 fn is_auto_create_exempt_private_table(schema: &str, table: &str) -> bool {
1147 schema == DEFAULT_PRIVATE_SCHEMA_NAME
1148 && matches!(
1149 table,
1150 DEFAULT_EVENTS_TABLE_NAME | SLOW_QUERY_TABLE_NAME | REGION_STATS_HISTORY_TABLE_NAME
1151 )
1152 }
1153
1154 pub async fn ensure_bulk_insert_schema(
1158 &self,
1159 table: TableRef,
1160 request_schema: &ArrowSchema,
1161 ctx: &QueryContextRef,
1162 statement_executor: &StatementExecutor,
1163 ) -> Result<TableRef> {
1164 let table_info = table.table_info();
1165 if self.auto_create_disabled_reason(ctx)?.is_some()
1166 && !Self::is_auto_create_exempt_private_table(&table_info.schema_name, &table_info.name)
1167 {
1168 return Ok(table);
1169 }
1170
1171 let table_schema = table.schema();
1172 let schema = request_schema
1173 .fields()
1174 .iter()
1175 .filter(|field| table_schema.column_schema_by_name(field.name()).is_none())
1176 .map(|field| {
1177 let data_type = field.data_type();
1178 ensure!(
1181 !data_type.is_nested() && !matches!(data_type, ArrowDataType::Dictionary(..)),
1182 crate::error::NotSupportedSnafu {
1183 feat: format!(
1184 "automatically adding bulk insert column '{}' with type {:?}",
1185 field.name(),
1186 data_type
1187 ),
1188 }
1189 );
1190 let column = datatypes::schema::ColumnSchema::try_from(field.as_ref())
1191 .context(crate::error::ConvertSchemaSnafu)?;
1192 let column_def =
1195 try_as_column_def(&column, false).context(crate::error::ColumnDataTypeSnafu)?;
1196 Ok(ColumnSchema {
1197 column_name: column_def.name,
1198 datatype: column_def.data_type,
1199 semantic_type: column_def.semantic_type,
1200 datatype_extension: column_def.datatype_extension,
1201 options: column_def.options,
1202 })
1203 })
1204 .collect::<Result<Vec<_>>>()?;
1205 let mut request = RowInsertRequest {
1206 table_name: table_info.name.clone(),
1207 rows: Some(Rows {
1208 schema,
1209 rows: Vec::new(),
1210 }),
1211 };
1212 let Some(alter_expr) =
1213 self.get_alter_table_expr_on_demand(&mut request, &table, ctx, false, false, true)?
1214 else {
1215 return Ok(table);
1216 };
1217
1218 statement_executor
1219 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1220 .await?;
1221 self.get_table(
1222 &table_info.catalog_name,
1223 &table_info.schema_name,
1224 &table_info.name,
1225 )
1226 .await?
1227 .with_context(|| TableNotFoundSnafu {
1228 table_name: table_info.full_table_name(),
1229 })
1230 }
1231
1232 pub async fn ensure_trace_table_on_demand(
1237 &self,
1238 table_name: &str,
1239 request_schema: Vec<ColumnSchema>,
1240 alter_existing: bool,
1241 ctx: &QueryContextRef,
1242 statement_executor: &StatementExecutor,
1243 ) -> Result<()> {
1244 let mut requests = RowInsertRequests {
1245 inserts: vec![RowInsertRequest {
1246 table_name: table_name.to_string(),
1247 rows: Some(api::v1::Rows {
1248 schema: request_schema,
1249 rows: Vec::new(),
1250 }),
1251 }],
1252 };
1253 self.create_or_alter_tables_on_demand(
1254 &mut requests,
1255 ctx,
1256 AutoCreateTableType::Trace { alter_existing },
1257 statement_executor,
1258 false,
1259 false,
1260 None,
1261 )
1262 .await?;
1263 Ok(())
1264 }
1265
1266 #[allow(clippy::too_many_arguments)]
1284 async fn create_or_alter_tables_on_demand(
1285 &self,
1286 requests: &mut RowInsertRequests,
1287 ctx: &QueryContextRef,
1288 auto_create_table_type: AutoCreateTableType,
1289 statement_executor: &StatementExecutor,
1290 accommodate_existing_schema: bool,
1291 is_single_value: bool,
1292 align_time_index_unit: Option<TimeUnit>,
1293 ) -> Result<CreateAlterTableResult> {
1294 let is_exemplar = matches!(&auto_create_table_type, AutoCreateTableType::Exemplar);
1295 let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND
1296 .with_label_values(&[auto_create_table_type.as_str()])
1297 .start_timer();
1298 let catalog = ctx.current_catalog();
1299 let schema = ctx.current_schema();
1300
1301 let auto_create_disabled_reason = self.auto_create_disabled_reason(ctx)?;
1302 let has_auto_create_exempt_table = auto_create_disabled_reason.is_some()
1304 && requests
1305 .inserts
1306 .iter()
1307 .any(|req| Self::is_auto_create_exempt_private_table(&schema, &req.table_name));
1308 let mut table_infos = HashMap::new();
1309 if let Some(disabled_reason) = auto_create_disabled_reason
1311 && !has_auto_create_exempt_table
1312 {
1313 let mut instant_table_ids = HashSet::new();
1314 for req in &mut requests.inserts {
1315 let table = match self.get_table(catalog, &schema, &req.table_name).await? {
1316 Some(table) => table,
1317 None if is_ddl_reserved_table(&schema, &req.table_name) => {
1320 statement_executor
1321 .create_declared_relationships_table(catalog, ctx.clone())
1322 .await?
1323 }
1324 None => {
1325 return InvalidInsertRequestSnafu {
1326 reason: format!(
1327 "Table `{}` does not exist, and {}",
1328 req.table_name, disabled_reason
1329 ),
1330 }
1331 .fail();
1332 }
1333 };
1334 if matches!(
1336 auto_create_table_type,
1337 AutoCreateTableType::Logical(_) | AutoCreateTableType::Exemplar
1338 ) && let Some(rows) = req.rows.as_mut()
1339 && let Some(target_unit) = table_time_index_unit(&table)
1340 {
1341 convert_rows_time_unit(rows, target_unit)?;
1342 }
1343 let table_info = table.table_info();
1344 if matches!(auto_create_table_type, AutoCreateTableType::Exemplar) {
1345 validate_exemplar_table(&table_info)?;
1346 }
1347 if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1348 validate_trace_table_model(&table_info, ctx)?;
1349 }
1350 if table_info.is_ttl_instant_table() {
1351 instant_table_ids.insert(table_info.table_id());
1352 }
1353 table_infos.insert(table_info.table_id(), table.table_info());
1354 }
1355 let ret = CreateAlterTableResult {
1356 instant_table_ids,
1357 table_infos,
1358 };
1359 return Ok(ret);
1360 }
1361
1362 let mut create_tables = vec![];
1363 let mut alter_tables = vec![];
1364 let mut need_refresh_table_infos = HashSet::new();
1365 let mut instant_table_ids = HashSet::new();
1366 let mut per_table_semantics: Option<Option<PerTableSemanticIndex>> = None;
1367
1368 for req in &mut requests.inserts {
1369 let auto_create_allowed = auto_create_disabled_reason.is_none()
1371 || Self::is_auto_create_exempt_private_table(&schema, &req.table_name);
1372 match self.get_table(catalog, &schema, &req.table_name).await? {
1373 Some(table) => {
1374 let table_info = table.table_info();
1375 if matches!(auto_create_table_type, AutoCreateTableType::Exemplar) {
1376 validate_exemplar_table(&table_info)?;
1377 }
1378 if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1379 validate_trace_table_model(&table_info, ctx)?;
1380 }
1381 if table_info.is_ttl_instant_table() {
1382 instant_table_ids.insert(table_info.table_id());
1383 }
1384 if matches!(auto_create_table_type, AutoCreateTableType::Logical(_))
1388 && let Some(rows) = req.rows.as_mut()
1389 && let Some(target_unit) = table_time_index_unit(&table)
1390 {
1391 convert_rows_time_unit(rows, target_unit)?;
1392 }
1393 if auto_create_allowed
1394 && let Some(alter_expr) = self.get_alter_table_expr_on_demand(
1395 req,
1396 &table,
1397 ctx,
1398 accommodate_existing_schema,
1399 is_single_value,
1400 auto_create_table_type.alter_existing(),
1401 )?
1402 {
1403 alter_tables.push(alter_expr);
1404 need_refresh_table_infos.insert((
1405 catalog.to_string(),
1406 schema.clone(),
1407 req.table_name.clone(),
1408 ));
1409 } else {
1410 table_infos.insert(table_info.table_id(), table.table_info());
1411 }
1412 }
1413 None if is_ddl_reserved_table(&schema, &req.table_name) => {
1417 let table = statement_executor
1418 .create_declared_relationships_table(catalog, ctx.clone())
1419 .await?;
1420 let table_info = table.table_info();
1421 if table_info.is_ttl_instant_table() {
1422 instant_table_ids.insert(table_info.table_id());
1423 }
1424 table_infos.insert(table_info.table_id(), table_info);
1425 }
1426 None if !auto_create_allowed
1427 && let Some(disabled_reason) = auto_create_disabled_reason =>
1428 {
1429 return InvalidInsertRequestSnafu {
1430 reason: format!(
1431 "Table `{}` does not exist, and {}",
1432 req.table_name, disabled_reason,
1433 ),
1434 }
1435 .fail();
1436 }
1437 None => {
1438 if let Some(physical_unit) = align_time_index_unit
1442 && let Some(rows) = req.rows.as_mut()
1443 {
1444 convert_rows_time_unit(rows, physical_unit)?;
1445 }
1446 let semantic_index = per_table_semantics
1447 .get_or_insert_with(|| parse_per_table_semantic_index(ctx))
1448 .as_ref();
1449 let create_expr = self.get_create_table_expr_on_demand(
1450 req,
1451 &auto_create_table_type,
1452 ctx,
1453 semantic_index,
1454 )?;
1455 create_tables.push(create_expr);
1456 }
1457 }
1458 }
1459
1460 match auto_create_table_type {
1461 AutoCreateTableType::Logical(_) => {
1462 if !create_tables.is_empty() {
1463 let tables = self
1465 .create_logical_tables(create_tables, ctx, statement_executor)
1466 .await?;
1467
1468 for table in tables {
1469 let table_info = table.table_info();
1470 if table_info.is_ttl_instant_table() {
1471 instant_table_ids.insert(table_info.table_id());
1472 }
1473 table_infos.insert(table_info.table_id(), table.table_info());
1474 }
1475 }
1476 if !alter_tables.is_empty() {
1477 statement_executor
1479 .alter_logical_tables(alter_tables, ctx.clone(), TriggerReason::AutoAlter)
1480 .await?;
1481 }
1482 }
1483 AutoCreateTableType::Physical
1484 | AutoCreateTableType::Exemplar
1485 | AutoCreateTableType::Log
1486 | AutoCreateTableType::LastNonNull => {
1487 for create_table in create_tables {
1490 let table = self
1491 .create_physical_table(create_table, None, ctx, statement_executor)
1492 .await?;
1493 let table = if is_exemplar {
1495 let name = table.table_info().name.clone();
1496 self.get_table(catalog, &schema, &name)
1497 .await?
1498 .context(TableNotFoundSnafu { table_name: name })?
1499 } else {
1500 table
1501 };
1502 if matches!(auto_create_table_type, AutoCreateTableType::Exemplar) {
1503 validate_exemplar_table(&table.table_info())?;
1504 if let Some(req) = requests
1505 .inserts
1506 .iter_mut()
1507 .find(|r| r.table_name == table.table_info().name)
1508 && let Some(alter) = self.get_alter_table_expr_on_demand(
1509 req, &table, ctx, false, false, true,
1510 )?
1511 {
1512 alter_tables.push(alter);
1513 need_refresh_table_infos.insert((
1514 catalog.to_string(),
1515 schema.clone(),
1516 req.table_name.clone(),
1517 ));
1518 }
1519 }
1520 let table_info = table.table_info();
1521 if table_info.is_ttl_instant_table() {
1522 instant_table_ids.insert(table_info.table_id());
1523 }
1524 table_infos.insert(table_info.table_id(), table.table_info());
1525 }
1526 for alter_expr in alter_tables.into_iter() {
1527 statement_executor
1528 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1529 .await?;
1530 }
1531 }
1532
1533 AutoCreateTableType::Trace { .. } => {
1534 let trace_table_name = ctx
1535 .extension(TRACE_TABLE_NAME_SESSION_KEY)
1536 .unwrap_or(TRACE_TABLE_NAME);
1537
1538 let trace_table_partitions = if let Some(trace_table_partitions) =
1539 ctx.extension(TRACE_TABLE_PARTITIONS_HINT_KEY)
1540 {
1541 let p = trace_table_partitions.parse::<u32>().map_err(|_| {
1542 InvalidInsertRequestSnafu {
1543 reason: format!(
1544 "Failed to parse trace_table_partitions: {}",
1545 trace_table_partitions
1546 ),
1547 }
1548 .build()
1549 })?;
1550 Some(p)
1551 } else {
1552 None
1553 };
1554
1555 for mut create_table in create_tables {
1558 if create_table.table_name == trace_services_table_name(trace_table_name)
1559 || create_table.table_name == trace_operations_table_name(trace_table_name)
1560 {
1561 create_table
1563 .table_options
1564 .insert(APPEND_MODE_KEY.to_string(), "false".to_string());
1565 create_table.table_options.remove(TTL_KEY);
1567
1568 let table = self
1569 .create_physical_table(create_table, None, ctx, statement_executor)
1570 .await?;
1571 let table_info = table.table_info();
1572 if table_info.is_ttl_instant_table() {
1573 instant_table_ids.insert(table_info.table_id());
1574 }
1575 table_infos.insert(table_info.table_id(), table.table_info());
1576 } else {
1577 let partitions = if matches!(trace_table_partitions, Some(0) | Some(1)) {
1580 None
1582 } else {
1583 let p = partition_rule_for_hexstring(
1584 TRACE_ID_COLUMN,
1585 trace_table_partitions,
1586 )
1587 .context(CreatePartitionRulesSnafu)?;
1588 Some(p)
1589 };
1590
1591 let index_columns =
1596 [TRACE_ID_COLUMN, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN];
1597 for index_column in index_columns {
1598 if let Some(col) = create_table
1599 .column_defs
1600 .iter_mut()
1601 .find(|c| c.name == index_column)
1602 {
1603 col.options =
1604 options_from_skipping(&SkippingIndexOptions::default())
1605 .context(ColumnOptionsSnafu)?;
1606 } else {
1607 warn!(
1608 "Column {} not found when creating index for trace table: {}.",
1609 index_column, create_table.table_name
1610 );
1611 }
1612 }
1613
1614 let id_compression = CompressionOptions {
1618 encoding: Some(ColumnEncoding::Plain),
1619 page_size: Some(ReadableSize::kb(64)),
1620 ..Default::default()
1621 };
1622 let compressed_columns = [
1623 (SPAN_ID_COLUMN, id_compression),
1624 (PARENT_SPAN_ID_COLUMN, id_compression),
1625 (
1626 TIMESTAMP_END_COLUMN,
1627 CompressionOptions {
1628 encoding: Some(ColumnEncoding::DeltaBinaryPacked),
1629 ..Default::default()
1630 },
1631 ),
1632 ];
1633 for (name, compression) in compressed_columns {
1634 if let Some(col) =
1635 create_table.column_defs.iter_mut().find(|c| c.name == name)
1636 {
1637 let compression = options_from_compression(&compression)
1638 .context(ColumnOptionsSnafu)?;
1639 col.options
1640 .get_or_insert_with(Default::default)
1641 .options
1642 .extend(compression.options);
1643 }
1644 }
1645
1646 create_table.table_options.insert(
1648 TABLE_DATA_MODEL.to_string(),
1649 ctx.extension(SEMANTIC_PIPELINE)
1650 .unwrap_or(TABLE_DATA_MODEL_TRACE_V1)
1651 .to_string(),
1652 );
1653
1654 let table = self
1655 .create_physical_table(
1656 create_table,
1657 partitions,
1658 ctx,
1659 statement_executor,
1660 )
1661 .await?;
1662 let table_info = table.table_info();
1663 if table_info.is_ttl_instant_table() {
1664 instant_table_ids.insert(table_info.table_id());
1665 }
1666 table_infos.insert(table_info.table_id(), table.table_info());
1667 }
1668 }
1669 for alter_expr in alter_tables.into_iter() {
1670 statement_executor
1671 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1672 .await?;
1673 }
1674 }
1675 }
1676
1677 for (catalog, schema, table_name) in need_refresh_table_infos {
1679 let table = self
1680 .get_table(&catalog, &schema, &table_name)
1681 .await?
1682 .context(TableNotFoundSnafu {
1683 table_name: common_catalog::format_full_table_name(
1684 &catalog,
1685 &schema,
1686 &table_name,
1687 ),
1688 })?;
1689 let table_info = table.table_info();
1690 if is_exemplar {
1691 validate_exemplar_table(&table_info)?;
1692 }
1693 table_infos.insert(table_info.table_id(), table.table_info());
1694 }
1695
1696 if is_exemplar {
1697 for info in table_infos.values() {
1698 if let Some(unit) = info
1699 .meta
1700 .schema
1701 .timestamp_column()
1702 .and_then(|column| column.data_type.as_timestamp())
1703 .map(|timestamp| timestamp.unit())
1704 {
1705 for request in &mut requests.inserts {
1706 if request.table_name == info.name
1707 && let Some(rows) = request.rows.as_mut()
1708 {
1709 convert_rows_time_unit(rows, unit)?;
1710 }
1711 }
1712 }
1713 }
1714 }
1715
1716 Ok(CreateAlterTableResult {
1717 instant_table_ids,
1718 table_infos,
1719 })
1720 }
1721
1722 async fn create_physical_table_on_demand(
1723 &self,
1724 ctx: &QueryContextRef,
1725 physical_table: String,
1726 statement_executor: &StatementExecutor,
1727 ) -> Result<TableRef> {
1728 let catalog_name = ctx.current_catalog();
1729 let schema_name = ctx.current_schema();
1730
1731 if let Some(table) = self
1733 .get_table(catalog_name, &schema_name, &physical_table)
1734 .await?
1735 {
1736 return Ok(table);
1737 }
1738
1739 if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? {
1741 return InvalidInsertRequestSnafu {
1742 reason: format!(
1743 "Physical table `{physical_table}` does not exist, and {disabled_reason}"
1744 ),
1745 }
1746 .fail();
1747 }
1748
1749 let table_reference = TableReference::full(catalog_name, &schema_name, &physical_table);
1750 debug!("Ensuring physical metric table `{table_reference}` exists for insert");
1751
1752 let default_schema = vec![
1754 ColumnSchema {
1755 column_name: greptime_timestamp().to_string(),
1756 datatype: ColumnDataType::TimestampMillisecond as _,
1757 semantic_type: SemanticType::Timestamp as _,
1758 datatype_extension: None,
1759 options: None,
1760 },
1761 ColumnSchema {
1762 column_name: greptime_value().to_string(),
1763 datatype: ColumnDataType::Float64 as _,
1764 semantic_type: SemanticType::Field as _,
1765 datatype_extension: None,
1766 options: None,
1767 },
1768 ];
1769 let create_table_expr =
1770 &mut build_create_table_expr(&table_reference, &default_schema, default_engine())?;
1771
1772 create_table_expr.engine = METRIC_ENGINE_NAME.to_string();
1773 create_table_expr
1774 .table_options
1775 .insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), "true".to_string());
1776
1777 let res = statement_executor
1779 .create_table_inner(
1780 create_table_expr,
1781 None,
1782 ctx.clone(),
1783 TriggerReason::AutoCreate,
1784 )
1785 .await;
1786
1787 match res {
1788 Ok(table) => Ok(table),
1789 Err(err) => {
1790 error!(err; "Failed to create table {table_reference}");
1791 Err(err)
1792 }
1793 }
1794 }
1795
1796 async fn get_table(
1797 &self,
1798 catalog: &str,
1799 schema: &str,
1800 table: &str,
1801 ) -> Result<Option<TableRef>> {
1802 self.catalog_manager
1803 .table(catalog, schema, table, None)
1804 .await
1805 .context(CatalogSnafu)
1806 }
1807
1808 fn get_create_table_expr_on_demand(
1809 &self,
1810 req: &RowInsertRequest,
1811 create_type: &AutoCreateTableType,
1812 ctx: &QueryContextRef,
1813 semantic_index: Option<&PerTableSemanticIndex>,
1814 ) -> Result<CreateTableExpr> {
1815 let schema = ctx.current_schema();
1816 let mut table_options = std::collections::HashMap::with_capacity(4);
1817 fill_table_options_for_create(&mut table_options, create_type, ctx);
1818 if !matches!(create_type, AutoCreateTableType::Exemplar) {
1819 apply_per_table_semantic_options(
1820 &mut table_options,
1821 semantic_index,
1822 ctx.current_schema().as_str(),
1823 &req.table_name,
1824 );
1825 }
1826
1827 let engine_name = if let AutoCreateTableType::Logical(_) = create_type {
1828 METRIC_ENGINE_NAME
1830 } else {
1831 default_engine()
1832 };
1833
1834 let table_ref = TableReference::full(ctx.current_catalog(), &schema, &req.table_name);
1835 let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1837 let mut create_table_expr =
1838 build_create_table_expr(&table_ref, request_schema, engine_name)?;
1839
1840 if matches!(create_type, AutoCreateTableType::Exemplar) {
1841 create_table_expr
1843 .column_defs
1844 .sort_by_key(|column| column.name != DATA_SCHEMA_TABLE_ID_COLUMN_NAME);
1845 create_table_expr
1846 .primary_keys
1847 .sort_by_key(|name| name != DATA_SCHEMA_TABLE_ID_COLUMN_NAME);
1848 for column in &mut create_table_expr.column_defs {
1849 if [
1850 DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
1851 DATA_SCHEMA_TSID_COLUMN_NAME,
1852 ]
1853 .contains(&column.name.as_str())
1854 {
1855 column.is_nullable = false;
1856 }
1857 }
1858 }
1859
1860 if ctx.extension(SPLUNK_PK_METADATA_ORDER_KEY).is_some() {
1862 reorder_splunk_primary_keys(&mut create_table_expr.primary_keys);
1863 }
1864
1865 debug!("Ensuring table `{table_ref}` exists for insert");
1866 create_table_expr.table_options.extend(table_options);
1867 Ok(create_table_expr)
1868 }
1869
1870 fn get_alter_table_expr_on_demand(
1879 &self,
1880 req: &mut RowInsertRequest,
1881 table: &TableRef,
1882 ctx: &QueryContextRef,
1883 accommodate_existing_schema: bool,
1884 is_single_value: bool,
1885 alter_existing: bool,
1886 ) -> Result<Option<AlterTableExpr>> {
1887 if !alter_existing {
1888 return Ok(None);
1889 }
1890
1891 let catalog_name = ctx.current_catalog();
1892 let schema_name = ctx.current_schema();
1893 let table_name = table.table_info().name.clone();
1894
1895 if is_ddl_reserved_table(&schema_name, &table_name) {
1898 return Ok(None);
1899 }
1900
1901 let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1902 let request_field_count = request_schema
1903 .iter()
1904 .filter(|col| col.semantic_type == SemanticType::Field as i32)
1905 .count();
1906 if is_single_value {
1907 let request_is_native_histogram = request_is_native_histogram(request_schema);
1908 let table_is_native_histogram = table_is_native_histogram(table);
1909 ensure!(
1910 request_is_native_histogram == table_is_native_histogram,
1911 InvalidInsertRequestSnafu {
1912 reason: format!(
1913 "Table `{table_name}` cannot mix native histogram and float sample fields"
1914 ),
1915 }
1916 );
1917 }
1918
1919 let column_exprs = ColumnExpr::from_column_schemas(request_schema);
1920 let add_columns = expr_helper::extract_add_columns_expr(&table.schema(), column_exprs)?;
1921 let Some(mut add_columns) = add_columns else {
1922 return Ok(None);
1923 };
1924
1925 if accommodate_existing_schema {
1927 let table_schema = table.schema();
1928 let ts_col_name = table_schema.timestamp_column().map(|c| c.name.clone());
1930 let mut field_col_name = None;
1932 if is_single_value && request_field_count <= 1 {
1933 let mut multiple_field_cols = false;
1934 table.field_columns().for_each(|col| {
1935 if field_col_name.is_none() {
1936 field_col_name = Some(col.name.clone());
1937 } else {
1938 multiple_field_cols = true;
1939 }
1940 });
1941 if multiple_field_cols {
1942 field_col_name = None;
1943 }
1944 }
1945
1946 if let Some(rows) = req.rows.as_mut() {
1948 for col in &mut rows.schema {
1949 match col.semantic_type {
1950 x if x == SemanticType::Timestamp as i32 => {
1951 if let Some(ref ts_name) = ts_col_name
1952 && col.column_name != *ts_name
1953 {
1954 col.column_name = ts_name.clone();
1955 }
1956 }
1957 x if x == SemanticType::Field as i32 => {
1958 if let Some(ref field_name) = field_col_name
1959 && col.column_name != *field_name
1960 {
1961 col.column_name = field_name.clone();
1962 }
1963 }
1964 _ => {}
1965 }
1966 }
1967 }
1968
1969 add_columns.add_columns.retain(|col| {
1971 let def = col.column_def.as_ref().unwrap();
1972 def.semantic_type == SemanticType::Tag as i32
1973 || (def.semantic_type == SemanticType::Field as i32 && field_col_name.is_none())
1974 });
1975
1976 if add_columns.add_columns.is_empty() {
1977 return Ok(None);
1978 }
1979 }
1980
1981 Ok(Some(AlterTableExpr {
1982 catalog_name: catalog_name.to_string(),
1983 schema_name: schema_name.clone(),
1984 table_name: table_name.clone(),
1985 kind: Some(Kind::AddColumns(add_columns)),
1986 }))
1987 }
1988
1989 async fn create_physical_table(
1991 &self,
1992 mut create_table_expr: CreateTableExpr,
1993 partitions: Option<Partitions>,
1994 ctx: &QueryContextRef,
1995 statement_executor: &StatementExecutor,
1996 ) -> Result<TableRef> {
1997 let res = statement_executor
1998 .create_table_inner(
1999 &mut create_table_expr,
2000 partitions,
2001 ctx.clone(),
2002 TriggerReason::AutoCreate,
2003 )
2004 .await;
2005
2006 let table_ref = TableReference::full(
2007 &create_table_expr.catalog_name,
2008 &create_table_expr.schema_name,
2009 &create_table_expr.table_name,
2010 );
2011
2012 match res {
2013 Ok(table) => {
2014 validate_trace_table_model(&table.table_info(), ctx)?;
2015 Ok(table)
2016 }
2017 Err(err) => {
2018 error!(err; "Failed to create table {}", table_ref);
2019 Err(err)
2020 }
2021 }
2022 }
2023
2024 async fn create_logical_tables(
2025 &self,
2026 create_table_exprs: Vec<CreateTableExpr>,
2027 ctx: &QueryContextRef,
2028 statement_executor: &StatementExecutor,
2029 ) -> Result<Vec<TableRef>> {
2030 let res = statement_executor
2031 .create_logical_tables(&create_table_exprs, ctx.clone(), TriggerReason::AutoCreate)
2032 .await;
2033
2034 match res {
2035 Ok(res) => Ok(res),
2036 Err(err) => {
2037 let failed_tables = create_table_exprs
2038 .into_iter()
2039 .map(|expr| {
2040 format!(
2041 "{}.{}.{}",
2042 expr.catalog_name, expr.schema_name, expr.table_name
2043 )
2044 })
2045 .collect::<Vec<_>>();
2046 error!(
2047 err;
2048 "Failed to create logical tables {:?}",
2049 failed_tables
2050 );
2051 Err(err)
2052 }
2053 }
2054 }
2055
2056 pub fn write_dispatcher(&self) -> &WriteDispatcherRef {
2058 &self.write_dispatcher
2059 }
2060
2061 pub fn node_manager(&self) -> &NodeManagerRef {
2062 &self.node_manager
2063 }
2064
2065 pub fn partition_manager(&self) -> &PartitionRuleManagerRef {
2066 &self.partition_manager
2067 }
2068
2069 pub fn table_flownode_set_cache(&self) -> &TableFlownodeSetCacheRef {
2070 &self.table_flownode_set_cache
2071 }
2072}
2073
2074fn request_is_native_histogram(request_schema: &[ColumnSchema]) -> bool {
2075 let mut fields = request_schema
2076 .iter()
2077 .filter(|col| col.semantic_type == SemanticType::Field as i32);
2078 let Some(col) = fields.next() else {
2079 return false;
2080 };
2081
2082 fields.next().is_none()
2083 && api::helper::is_column_type_value_eq(
2084 col.datatype,
2085 col.datatype_extension.clone(),
2086 native_histogram_value_type(),
2087 )
2088}
2089
2090fn table_time_index_unit(table: &TableRef) -> Option<TimeUnit> {
2094 table
2095 .table_info()
2096 .meta
2097 .schema
2098 .timestamp_column()
2099 .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit()))
2100}
2101
2102fn convert_rows_time_unit(rows: &mut Rows, target_unit: TimeUnit) -> Result<()> {
2103 let Some(ts_index) = rows
2104 .schema
2105 .iter()
2106 .position(|col| col.semantic_type == SemanticType::Timestamp as i32)
2107 else {
2108 return Ok(());
2109 };
2110 let Some(source_unit) = ColumnDataType::try_from(rows.schema[ts_index].datatype)
2111 .ok()
2112 .and_then(api::helper::timestamp_unit)
2113 else {
2114 return Ok(());
2115 };
2116 if source_unit == target_unit {
2117 return Ok(());
2118 }
2119
2120 rows.schema[ts_index].datatype = api::helper::timestamp_datatype(target_unit) as i32;
2121 rows.schema[ts_index].datatype_extension = None;
2123
2124 for row in &mut rows.rows {
2132 debug_assert_eq!(row.values.len(), rows.schema.len());
2133 let value = &mut row.values[ts_index];
2134 let Some(value_data) = value.value_data.take() else {
2135 continue;
2136 };
2137 value.value_data =
2138 convert_timestamp_value_data(value_data, source_unit, target_unit, ts_index)?;
2139 }
2140 Ok(())
2141}
2142
2143fn convert_timestamp_value_data(
2144 value_data: ValueData,
2145 source_unit: TimeUnit,
2146 target_unit: TimeUnit,
2147 column_index: usize,
2148) -> Result<Option<ValueData>> {
2149 let timestamp = match value_data {
2150 ValueData::TimestampSecondValue(v) => Timestamp::new_second(v),
2151 ValueData::TimestampMillisecondValue(v) => Timestamp::new_millisecond(v),
2152 ValueData::TimestampMicrosecondValue(v) => Timestamp::new_microsecond(v),
2153 ValueData::TimestampNanosecondValue(v) => Timestamp::new_nanosecond(v),
2154 other => return Ok(Some(other)),
2156 };
2157 let converted = timestamp
2158 .convert_to(target_unit)
2159 .with_context(|| InvalidInsertRequestSnafu {
2160 reason: format!(
2161 "timestamp column {column_index} value {} in unit {source_unit:?} overflows when converting to unit {target_unit:?}",
2162 timestamp.value()
2163 ),
2164 })?;
2165 Ok(api::helper::to_grpc_value(datatypes::value::Value::Timestamp(converted)).value_data)
2166}
2167
2168fn table_is_native_histogram(table: &TableRef) -> bool {
2169 let mut fields = table.field_columns();
2170 let Some(col) = fields.next() else {
2171 return false;
2172 };
2173
2174 fields.next().is_none() && is_native_histogram_value_type(&col.data_type)
2175}
2176
2177fn validate_column_count_match(requests: &RowInsertRequests) -> Result<()> {
2178 for request in &requests.inserts {
2179 let rows = request.rows.as_ref().unwrap();
2180 let column_count = rows.schema.len();
2181 rows.rows.iter().try_for_each(|r| {
2182 ensure!(
2183 r.values.len() == column_count,
2184 InvalidInsertRequestSnafu {
2185 reason: format!(
2186 "column count mismatch, columns: {}, values: {}",
2187 column_count,
2188 r.values.len()
2189 )
2190 }
2191 );
2192 Ok(())
2193 })?;
2194 }
2195 Ok(())
2196}
2197
2198pub fn validate_exemplar_table(table: &TableInfo) -> Result<()> {
2200 let meta = &table.meta;
2201 let options = &meta.options.extra_options;
2202 ensure!(
2203 meta.engine == default_engine()
2204 && options
2205 .get(APPEND_MODE_KEY)
2206 .is_some_and(|v| v.eq_ignore_ascii_case("true"))
2207 && options
2208 .get(store_api::metric_engine_consts::PRIMARY_KEY_ENCODING)
2209 .is_some_and(|v| v.eq_ignore_ascii_case("sparse")),
2210 InvalidInsertRequestSnafu {
2211 reason: "metrics_exemplars requires Mito, append_mode=true and primary_key_encoding=sparse"
2212 }
2213 );
2214 let columns = meta.schema.column_schemas();
2215 let lookup = |name: &str| {
2216 columns
2217 .iter()
2218 .enumerate()
2219 .find(|(_, c)| c.name == name)
2220 .with_context(|| InvalidInsertRequestSnafu {
2221 reason: format!("missing exemplar column {name}"),
2222 })
2223 };
2224 let (_, timestamp) = lookup(greptime_timestamp())?;
2225 for (position, name, id, data_type) in [
2226 (
2227 0,
2228 DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
2229 ReservedColumnId::table_id(),
2230 datatypes::prelude::ConcreteDataType::uint32_datatype(),
2231 ),
2232 (
2233 1,
2234 DATA_SCHEMA_TSID_COLUMN_NAME,
2235 ReservedColumnId::tsid(),
2236 datatypes::prelude::ConcreteDataType::uint64_datatype(),
2237 ),
2238 ] {
2239 let (index, column) = lookup(name)?;
2240 ensure!(
2241 meta.primary_key_indices.get(position) == Some(&index)
2242 && meta.column_ids.get(index) == Some(&id)
2243 && column.data_type == data_type
2244 && !column.is_nullable()
2245 && column.default_constraint().is_none(),
2246 InvalidInsertRequestSnafu {
2247 reason: format!("invalid exemplar source identity column {name}"),
2248 }
2249 );
2250 }
2251 ensure!(
2252 timestamp.is_time_index() && timestamp.data_type.is_timestamp(),
2253 InvalidInsertRequestSnafu {
2254 reason: "invalid exemplar time index"
2255 }
2256 );
2257 let (metric_index, metric) = lookup(EXEMPLAR_METRIC_COLUMN)?;
2258 ensure!(
2259 metric.data_type.is_string() && meta.primary_key_indices.contains(&metric_index),
2260 InvalidInsertRequestSnafu {
2261 reason: "exemplar metric name must be a string tag"
2262 }
2263 );
2264 let (value_index, value) = lookup(greptime_value())?;
2265 ensure!(
2266 value.data_type == datatypes::prelude::ConcreteDataType::float64_datatype()
2267 && !meta.primary_key_indices.contains(&value_index),
2268 InvalidInsertRequestSnafu {
2269 reason: "exemplar value must be a DOUBLE field"
2270 }
2271 );
2272 let (labels_index, labels) = lookup(EXEMPLAR_LABELS_COLUMN)?;
2273 ensure!(
2274 labels.data_type.is_json2() && !meta.primary_key_indices.contains(&labels_index),
2275 InvalidInsertRequestSnafu {
2276 reason: "exemplar_labels must be a JSON2 field"
2277 }
2278 );
2279 if let Some(extension) = labels
2280 .extension_type::<datatypes::extension::json::Json2ExtensionType>()
2281 .map_err(|error| {
2282 InvalidInsertRequestSnafu {
2283 reason: error.to_string(),
2284 }
2285 .build()
2286 })?
2287 {
2288 ensure!(
2289 extension
2290 .metadata()
2291 .json_settings()
2292 .type_hints()
2293 .iter()
2294 .all(|hint| hint.path.len() == 1 && hint.data_type.is_string()),
2295 InvalidInsertRequestSnafu {
2296 reason: "exemplar JSON2 hints must preserve string labels"
2297 }
2298 );
2299 }
2300 for (index, column) in columns.iter().enumerate() {
2301 if [
2302 greptime_timestamp(),
2303 greptime_value(),
2304 EXEMPLAR_METRIC_COLUMN,
2305 EXEMPLAR_LABELS_COLUMN,
2306 DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
2307 DATA_SCHEMA_TSID_COLUMN_NAME,
2308 ]
2309 .contains(&column.name.as_str())
2310 {
2311 continue;
2312 }
2313 ensure!(
2314 meta.primary_key_indices.contains(&index)
2315 && column.data_type.is_string()
2316 && column.is_nullable(),
2317 InvalidInsertRequestSnafu {
2318 reason: "inherited exemplar labels must be nullable string tags"
2319 }
2320 );
2321 ensure!(
2323 matches!(
2324 column.default_constraint(),
2325 None | Some(ColumnDefaultConstraint::Value(
2326 datatypes::value::Value::Null
2327 ))
2328 ),
2329 InvalidInsertRequestSnafu {
2330 reason: "inherited exemplar labels must not have non-NULL defaults"
2331 }
2332 );
2333 }
2334 Ok(())
2335}
2336
2337pub fn validate_trace_table_model(table_info: &TableInfo, ctx: &QueryContextRef) -> Result<()> {
2340 let Some(expected @ (TABLE_DATA_MODEL_TRACE_V1 | TABLE_DATA_MODEL_TRACE_V2)) =
2341 ctx.extension(SEMANTIC_PIPELINE)
2342 else {
2343 return Ok(());
2344 };
2345 if let Some(actual) = table_info.meta.options.data_model() {
2346 ensure!(
2347 actual == expected,
2348 InvalidInsertRequestSnafu {
2349 reason: format!(
2350 "Trace table `{}` uses {actual}, but the request uses {expected}",
2351 table_info.name,
2352 ),
2353 }
2354 );
2355 }
2356 Ok(())
2357}
2358
2359pub fn fill_table_options_for_create(
2361 table_options: &mut std::collections::HashMap<String, String>,
2362 create_type: &AutoCreateTableType,
2363 ctx: &QueryContextRef,
2364) {
2365 for key in VALID_TABLE_OPTION_KEYS {
2366 if let Some(value) = ctx.extension(key) {
2367 table_options.insert(key.to_string(), value.to_string());
2368 }
2369 }
2370
2371 if !matches!(create_type, AutoCreateTableType::Logical(_)) {
2372 let mut has_twcs_option = false;
2373 for key in COMPACTION_OPTION_KEYS {
2374 if let Some(value) = ctx.extension(key) {
2375 has_twcs_option |= key.starts_with(COMPACTION_TWCS_PREFIX);
2376 table_options.insert(key.to_string(), value.to_string());
2377 }
2378 }
2379 if has_twcs_option {
2380 table_options
2381 .entry(COMPACTION_TYPE.to_string())
2382 .or_insert_with(|| COMPACTION_TYPE_TWCS.to_string());
2383 }
2384 }
2385
2386 for (key, value) in ctx.extensions() {
2388 if is_semantic_option_key(&key) && validate_semantic_option(&key, &value) {
2389 table_options.insert(key, value);
2390 }
2391 }
2392
2393 match create_type {
2394 AutoCreateTableType::Logical(physical_table) => {
2395 table_options.insert(
2396 LOGICAL_TABLE_METADATA_KEY.to_string(),
2397 physical_table.clone(),
2398 );
2399 }
2400 AutoCreateTableType::Exemplar => {
2401 table_options.retain(|key, _| {
2402 !is_semantic_option_key(key)
2403 && key != SEMANTIC_PER_TABLE_INDEX_KEY
2404 && key != LOGICAL_TABLE_METADATA_KEY
2405 && key != PHYSICAL_TABLE_METADATA_KEY
2406 && key != MERGE_MODE_KEY
2407 && key != "otlp_metric_compat"
2408 });
2409 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2410 table_options.insert(
2411 store_api::metric_engine_consts::PRIMARY_KEY_ENCODING.to_string(),
2412 "sparse".to_string(),
2413 );
2414 }
2415 AutoCreateTableType::Physical => {
2416 if let Some(append_mode) = ctx.extension(APPEND_MODE_KEY) {
2417 table_options.insert(APPEND_MODE_KEY.to_string(), append_mode.to_string());
2418 }
2419 if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
2420 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
2421 }
2422 }
2423 AutoCreateTableType::Log => {
2426 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2427 }
2428 AutoCreateTableType::LastNonNull => {
2429 if ctx
2430 .extension(APPEND_MODE_KEY)
2431 .is_some_and(|value| value.eq_ignore_ascii_case("true"))
2432 {
2433 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2434 table_options.insert(MERGE_MODE_KEY.to_string(), "last_row".to_string());
2435 } else if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
2436 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
2437 } else {
2438 table_options.insert(MERGE_MODE_KEY.to_string(), "last_non_null".to_string());
2439 }
2440 }
2441 AutoCreateTableType::Trace { .. } => {
2442 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2443 }
2444 }
2445}
2446
2447pub type PerTableSemanticIndex = BTreeMap<String, BTreeMap<String, BTreeMap<String, String>>>;
2453
2454pub fn parse_per_table_semantic_index(ctx: &QueryContextRef) -> Option<PerTableSemanticIndex> {
2459 let raw = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY)?;
2460 match serde_json::from_str(raw) {
2461 Ok(index) => Some(index),
2462 Err(_) => {
2463 warn!("failed to parse semantic per-table index, skipping per-table options");
2464 None
2465 }
2466 }
2467}
2468
2469pub fn apply_per_table_semantic_options(
2477 table_options: &mut std::collections::HashMap<String, String>,
2478 index: Option<&PerTableSemanticIndex>,
2479 schema: &str,
2480 table_name: &str,
2481) {
2482 let Some(entry) = index
2483 .and_then(|index| index.get(schema))
2484 .and_then(|tables| tables.get(table_name))
2485 else {
2486 return;
2487 };
2488 for (key, value) in entry {
2489 if is_semantic_option_key(key) && validate_semantic_option(key, value) {
2490 table_options.insert(key.clone(), value.clone());
2491 }
2492 }
2493}
2494
2495pub fn build_create_table_expr(
2496 table: &TableReference,
2497 request_schema: &[ColumnSchema],
2498 engine: &str,
2499) -> Result<CreateTableExpr> {
2500 expr_helper::create_table_expr_by_column_schemas(table, request_schema, engine, None)
2501}
2502
2503pub const SPLUNK_PK_METADATA_ORDER_KEY: &str = "splunk_pk_metadata_order";
2507
2508fn reorder_splunk_primary_keys(primary_keys: &mut [String]) {
2511 const LEAD: [&str; 3] = ["host", "source", "sourcetype"];
2512 primary_keys.sort_by_key(|name| {
2515 LEAD.iter()
2516 .position(|&lead| lead == name.as_str())
2517 .unwrap_or(LEAD.len())
2518 });
2519}
2520
2521struct CreateAlterTableResult {
2523 instant_table_ids: HashSet<TableId>,
2525 table_infos: HashMap<TableId, Arc<TableInfo>>,
2527}
2528
2529const MAX_MIRROR_PENDING_ROWS: usize = 1_000_000;
2537
2538const MIRROR_LOG_INTERVAL: Duration = Duration::from_secs(10);
2540const MIRROR_LOG_NEVER_REPORTED: u64 = u64::MAX;
2541
2542struct MirrorLog {
2543 last_log_millis: AtomicU64,
2544 events: AtomicU64,
2545}
2546
2547impl MirrorLog {
2548 const fn new() -> Self {
2549 Self {
2550 last_log_millis: AtomicU64::new(MIRROR_LOG_NEVER_REPORTED),
2551 events: AtomicU64::new(0),
2552 }
2553 }
2554
2555 fn claim_report(&self, now_millis: u64) -> Option<u64> {
2558 self.events.fetch_add(1, Ordering::Relaxed);
2559 let last = self.last_log_millis.load(Ordering::Relaxed);
2560 if last != MIRROR_LOG_NEVER_REPORTED
2561 && now_millis.saturating_sub(last) < MIRROR_LOG_INTERVAL.as_millis() as u64
2562 {
2563 return None;
2564 }
2565 if self
2566 .last_log_millis
2567 .compare_exchange(last, now_millis, Ordering::Relaxed, Ordering::Relaxed)
2568 .is_err()
2569 {
2570 return None;
2571 }
2572 Some(self.events.swap(0, Ordering::Relaxed))
2573 }
2574}
2575
2576static MIRROR_DROP_LOG: MirrorLog = MirrorLog::new();
2577static MIRROR_FAILURE_LOG: MirrorLog = MirrorLog::new();
2578
2579fn mirror_drop_log_millis() -> u64 {
2582 static START: LazyLock<Instant> = LazyLock::new(Instant::now);
2583 START.elapsed().as_millis() as u64
2584}
2585
2586struct FlowMirrorTask {
2587 requests: HashMap<Peer, RegionInsertRequests>,
2588}
2589
2590impl FlowMirrorTask {
2591 async fn new(
2592 cache: &TableFlownodeSetCacheRef,
2593 requests: impl Iterator<Item = &RegionInsertRequest>,
2594 ) -> Result<Self> {
2595 let mut src_table_reqs: HashMap<TableId, Option<(Vec<Peer>, RegionInsertRequests)>> =
2596 HashMap::new();
2597
2598 for req in requests {
2599 let table_id = RegionId::from_u64(req.region_id).table_id();
2600 match src_table_reqs.get_mut(&table_id) {
2601 Some(Some((_peers, reqs))) => reqs.requests.push(req.clone()),
2602 Some(None) => continue,
2604 _ => {
2605 let peers = cache
2607 .get(table_id)
2608 .await
2609 .context(RequestInsertsSnafu)?
2610 .unwrap_or_default()
2611 .values()
2612 .cloned()
2613 .collect::<HashSet<_>>()
2614 .into_iter()
2615 .collect::<Vec<_>>();
2616
2617 if !peers.is_empty() {
2618 let mut reqs = RegionInsertRequests::default();
2619 reqs.requests.push(req.clone());
2620 src_table_reqs.insert(table_id, Some((peers, reqs)));
2621 } else {
2622 src_table_reqs.insert(table_id, None);
2624 }
2625 }
2626 }
2627 }
2628
2629 let mut inserts: HashMap<Peer, RegionInsertRequests> = HashMap::new();
2630
2631 for (_table_id, (peers, reqs)) in src_table_reqs
2632 .into_iter()
2633 .filter_map(|(k, v)| v.map(|v| (k, v)))
2634 {
2635 if peers.len() == 1 {
2636 inserts
2638 .entry(peers[0].clone())
2639 .or_default()
2640 .requests
2641 .extend(reqs.requests);
2642 continue;
2643 } else {
2644 for flownode in peers {
2646 inserts
2647 .entry(flownode.clone())
2648 .or_default()
2649 .requests
2650 .extend(reqs.requests.clone());
2651 }
2652 }
2653 }
2654
2655 Ok(Self { requests: inserts })
2656 }
2657
2658 fn pending_rows(&self) -> u64 {
2667 self.requests.values().map(region_inserts_rows).sum()
2668 }
2669
2670 fn detach(
2671 self,
2672 node_manager: NodeManagerRef,
2673 mirror_pending_rows: Arc<AtomicU64>,
2674 has_instant_rows: bool,
2675 ) -> Result<()> {
2676 let num_rows = self.pending_rows();
2679 if num_rows == 0 {
2680 return Ok(());
2681 }
2682 if has_instant_rows && num_rows > MAX_MIRROR_PENDING_ROWS as u64 {
2683 return InvalidInsertRequestSnafu {
2684 reason: format!(
2685 "flow mirror batch has {num_rows} peer-fan-out rows, exceeding the limit of {}; reduce the batch size",
2686 MAX_MIRROR_PENDING_ROWS
2687 ),
2688 }
2689 .fail();
2690 }
2691
2692 let pending = reserve_mirror_pending_rows(&mirror_pending_rows, num_rows);
2693 if pending > MAX_MIRROR_PENDING_ROWS as u64 {
2694 release_mirror_pending_rows(&mirror_pending_rows, num_rows);
2695 if has_instant_rows {
2696 return Err(meter_core::collect::WriteRejected::new(
2697 "flow mirror pending rows limit exceeded for instant-TTL insert",
2698 ))
2699 .context(WriteRejectedSnafu);
2700 }
2701 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.inc_by(num_rows);
2702 if let Some(events) = MIRROR_DROP_LOG.claim_report(mirror_drop_log_millis()) {
2703 warn!(
2704 "Flow mirror write dropped: {} pending rows exceeds limit {}, {} drop events since last report",
2705 pending, MAX_MIRROR_PENDING_ROWS, events
2706 );
2707 }
2708 return Ok(());
2709 }
2710
2711 for (peer, inserts) in self.requests {
2712 let peer_rows = region_inserts_rows(&inserts);
2716 let node_manager = node_manager.clone();
2717 let mirror_pending_rows = mirror_pending_rows.clone();
2718 common_runtime::spawn_global(async move {
2719 let result = node_manager
2720 .flownode(&peer)
2721 .await
2722 .handle_inserts(inserts)
2723 .await
2724 .context(RequestInsertsSnafu);
2725
2726 match result {
2727 Ok(resp) => {
2728 let affected_rows = resp.affected_rows;
2729 crate::metrics::DIST_MIRROR_ROW_COUNT.inc_by(affected_rows);
2730 }
2731 Err(err) => {
2732 if let Some(events) =
2733 MIRROR_FAILURE_LOG.claim_report(mirror_drop_log_millis())
2734 {
2735 error!(err; flownode_id = peer.id, flownode_addr = %peer.addr, "Failed to insert data into flownode ({} total mirror failures across peers since last report)", events);
2736 }
2737 }
2738 }
2739 release_mirror_pending_rows(&mirror_pending_rows, peer_rows);
2742 });
2743 }
2744
2745 Ok(())
2746 }
2747}
2748
2749fn region_inserts_rows(inserts: &RegionInsertRequests) -> u64 {
2752 inserts
2753 .requests
2754 .iter()
2755 .filter_map(|req| req.rows.as_ref())
2756 .map(|rows| rows.rows.len() as u64)
2757 .sum()
2758}
2759
2760fn reserve_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) -> u64 {
2765 let pending = mirror_pending_rows.fetch_add(rows, Ordering::Relaxed) + rows;
2766 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(rows as i64);
2769 pending
2770}
2771
2772fn release_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) {
2778 let prev = mirror_pending_rows
2779 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |pending| {
2780 Some(pending.saturating_sub(rows))
2781 })
2782 .unwrap_or_else(|prev| prev);
2783 let released = rows.min(prev);
2787 if released > 0 {
2788 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(released as i64);
2789 }
2790}
2791
2792#[cfg(test)]
2793mod tests {
2794 use std::future::Future;
2795 use std::sync::Arc;
2796 use std::sync::atomic::{AtomicU64, Ordering};
2797 use std::task::Poll;
2798
2799 use api::helper::ColumnDataTypeWrapper;
2800 use api::region::RegionResponse;
2801 use api::v1::flow::FlowResponse;
2802 use api::v1::helper::{field_column_schema, time_index_column_schema};
2803 use api::v1::{Row, RowInsertRequest, Rows, Value};
2804 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
2805 use common_error::ext::{ErrorExt, RetryHint};
2806 use common_error::mock::MockError;
2807 use common_error::status_code::StatusCode;
2808 use common_meta::cache::new_table_flownode_set_cache;
2809 use common_meta::cache_invalidator::DummyCacheInvalidator;
2810 use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler;
2811 use common_meta::node_manager::{
2812 Datanode, DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef,
2813 };
2814 use common_meta::test_util::{MockDatanodeManager, MockFlownodeHandler, MockFlownodeManager};
2815 use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD;
2816 use common_query::prelude::{greptime_native_histogram, set_default_prefix};
2817 use datatypes::data_type::ConcreteDataType;
2818 use datatypes::schema::ColumnSchema;
2819 use moka::future::Cache;
2820 use session::context::QueryContext;
2821 use table::TableRef;
2822 use table::dist_table::DummyDataSource;
2823 use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType};
2824 use tokio::sync::{Semaphore, mpsc};
2825
2826 use crate::insert::*;
2827 use crate::test_util::{
2828 create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
2829 };
2830 use crate::write_dispatcher::{TrackedWrite, WriteRetryOptions};
2831
2832 fn make_table_ref_with_schema(
2833 ts_name: &str,
2834 field_name: &str,
2835 field_type: ConcreteDataType,
2836 ) -> TableRef {
2837 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2838 ColumnSchema::new(
2839 ts_name,
2840 ConcreteDataType::timestamp_millisecond_datatype(),
2841 false,
2842 )
2843 .with_time_index(true),
2844 ColumnSchema::new(field_name, field_type, true),
2845 ])
2846 .unwrap()
2847 .build()
2848 .unwrap();
2849 let meta = TableMetaBuilder::empty()
2850 .schema(Arc::new(schema))
2851 .primary_key_indices(vec![])
2852 .value_indices(vec![1])
2853 .engine("mito")
2854 .next_column_id(0)
2855 .options(Default::default())
2856 .created_on(Default::default())
2857 .build()
2858 .unwrap();
2859 let info = Arc::new(
2860 TableInfoBuilder::default()
2861 .table_id(1)
2862 .table_version(0)
2863 .name("test_table")
2864 .schema_name(DEFAULT_SCHEMA_NAME)
2865 .catalog_name(DEFAULT_CATALOG_NAME)
2866 .desc(None)
2867 .table_type(TableType::Base)
2868 .meta(meta)
2869 .build()
2870 .unwrap(),
2871 );
2872 Arc::new(table::Table::new(
2873 info,
2874 table::metadata::FilterPushDownType::Unsupported,
2875 Arc::new(DummyDataSource),
2876 ))
2877 }
2878
2879 fn make_metric_physical_table_ref_with_time_unit(unit: TimeUnit) -> TableRef {
2880 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2881 ColumnSchema::new(
2882 greptime_timestamp(),
2883 ConcreteDataType::timestamp_datatype(unit),
2884 false,
2885 )
2886 .with_time_index(true),
2887 ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
2888 ])
2889 .unwrap()
2890 .build()
2891 .unwrap();
2892 let meta = TableMetaBuilder::empty()
2893 .schema(Arc::new(schema))
2894 .primary_key_indices(vec![])
2895 .value_indices(vec![1])
2896 .engine("metric")
2897 .next_column_id(0)
2898 .options(Default::default())
2899 .created_on(Default::default())
2900 .build()
2901 .unwrap();
2902 let info = Arc::new(
2903 TableInfoBuilder::default()
2904 .table_id(1)
2905 .table_version(0)
2906 .name("greptime_physical_table")
2907 .schema_name(DEFAULT_SCHEMA_NAME)
2908 .catalog_name(DEFAULT_CATALOG_NAME)
2909 .desc(None)
2910 .table_type(TableType::Base)
2911 .meta(meta)
2912 .build()
2913 .unwrap(),
2914 );
2915 Arc::new(table::Table::new(
2916 info,
2917 table::metadata::FilterPushDownType::Unsupported,
2918 Arc::new(DummyDataSource),
2919 ))
2920 }
2921
2922 fn ms_row_insert_request(timestamp_ms: i64) -> RowInsertRequest {
2923 ms_row_insert_request_named("my_metric", timestamp_ms)
2924 }
2925
2926 fn ms_row_insert_request_named(table: &str, timestamp_ms: i64) -> RowInsertRequest {
2927 RowInsertRequest {
2928 table_name: table.to_string(),
2929 rows: Some(Rows {
2930 schema: vec![
2931 time_index_column_schema(
2932 greptime_timestamp(),
2933 ColumnDataType::TimestampMillisecond,
2934 ),
2935 field_column_schema(greptime_value(), ColumnDataType::Float64),
2936 ],
2937 rows: vec![Row {
2938 values: vec![
2939 Value {
2940 value_data: Some(ValueData::TimestampMillisecondValue(timestamp_ms)),
2941 },
2942 Value {
2943 value_data: Some(ValueData::F64Value(1.0)),
2944 },
2945 ],
2946 }],
2947 }),
2948 }
2949 }
2950
2951 fn convert_row_insert_requests_time_unit(
2953 requests: &mut RowInsertRequests,
2954 target_unit: TimeUnit,
2955 ) -> Result<()> {
2956 for request in &mut requests.inserts {
2957 let Some(rows) = request.rows.as_mut() else {
2958 continue;
2959 };
2960 convert_rows_time_unit(rows, target_unit)?;
2961 }
2962 Ok(())
2963 }
2964
2965 #[test]
2966 fn test_convert_row_insert_requests_time_unit_noop_when_matching() {
2967 let mut requests = RowInsertRequests {
2968 inserts: vec![ms_row_insert_request(123)],
2969 };
2970 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Millisecond).unwrap();
2971 let rows = requests.inserts[0].rows.as_ref().unwrap();
2972 assert_eq!(
2973 rows.schema[0].datatype,
2974 ColumnDataType::TimestampMillisecond as i32
2975 );
2976 assert!(matches!(
2977 rows.rows[0].values[0].value_data,
2978 Some(ValueData::TimestampMillisecondValue(123))
2979 ));
2980 }
2981
2982 #[test]
2983 fn test_convert_row_insert_requests_time_unit_widens_losslessly() {
2984 let mut requests = RowInsertRequests {
2985 inserts: vec![ms_row_insert_request(123)],
2986 };
2987 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Microsecond).unwrap();
2988 let rows = requests.inserts[0].rows.as_ref().unwrap();
2989 assert_eq!(
2990 rows.schema[0].datatype,
2991 ColumnDataType::TimestampMicrosecond as i32
2992 );
2993 assert!(matches!(
2994 rows.rows[0].values[0].value_data,
2995 Some(ValueData::TimestampMicrosecondValue(123_000))
2996 ));
2997 }
2998
2999 #[test]
3000 fn test_convert_row_insert_requests_time_unit_truncates_on_narrowing() {
3001 let requests = |value_ns: i64| RowInsertRequests {
3004 inserts: vec![RowInsertRequest {
3005 table_name: "my_metric".to_string(),
3006 rows: Some(Rows {
3007 schema: vec![
3008 time_index_column_schema(
3009 greptime_timestamp(),
3010 ColumnDataType::TimestampNanosecond,
3011 ),
3012 field_column_schema(greptime_value(), ColumnDataType::Float64),
3013 ],
3014 rows: vec![Row {
3015 values: vec![
3016 Value {
3017 value_data: Some(ValueData::TimestampNanosecondValue(value_ns)),
3018 },
3019 Value {
3020 value_data: Some(ValueData::F64Value(1.0)),
3021 },
3022 ],
3023 }],
3024 }),
3025 }],
3026 };
3027
3028 let mut reqs = requests(123_456_789);
3029 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
3030 let rows = reqs.inserts[0].rows.as_ref().unwrap();
3031 assert!(matches!(
3032 rows.rows[0].values[0].value_data,
3033 Some(ValueData::TimestampMicrosecondValue(123_456))
3034 ));
3035
3036 let mut reqs = requests(123_456_789);
3037 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Millisecond).unwrap();
3038 let rows = reqs.inserts[0].rows.as_ref().unwrap();
3039 assert!(matches!(
3040 rows.rows[0].values[0].value_data,
3041 Some(ValueData::TimestampMillisecondValue(123))
3042 ));
3043
3044 let mut reqs = requests(-123_456_789);
3045 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
3046 let rows = reqs.inserts[0].rows.as_ref().unwrap();
3047 assert!(matches!(
3048 rows.rows[0].values[0].value_data,
3049 Some(ValueData::TimestampMicrosecondValue(-123_457))
3050 ));
3051 }
3052
3053 #[test]
3054 fn test_convert_row_insert_requests_time_unit_overflow() {
3055 let mut requests = RowInsertRequests {
3056 inserts: vec![ms_row_insert_request(i64::MAX)],
3057 };
3058 let err =
3059 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Nanosecond).unwrap_err();
3060 assert!(err.to_string().contains("overflows"), "{err}");
3061 }
3062
3063 #[test]
3064 fn test_table_time_index_unit() {
3065 assert_eq!(
3066 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
3067 TimeUnit::Microsecond
3068 )),
3069 Some(TimeUnit::Microsecond)
3070 );
3071 assert_eq!(
3072 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
3073 TimeUnit::Millisecond
3074 )),
3075 Some(TimeUnit::Millisecond)
3076 );
3077 }
3078
3079 #[tokio::test]
3080 async fn test_accommodate_existing_schema_and_reject_kind_changes() {
3081 let ts_name = "my_ts";
3082 let field_name = "my_field";
3083 let table =
3084 make_table_ref_with_schema(ts_name, field_name, ConcreteDataType::float64_datatype());
3085
3086 let mut req = RowInsertRequest {
3088 table_name: "test_table".to_string(),
3089 rows: Some(Rows {
3090 schema: vec![
3091 time_index_column_schema("ts_wrong", ColumnDataType::TimestampMillisecond),
3092 field_column_schema("field_wrong", ColumnDataType::Float64),
3093 ],
3094 rows: vec![api::v1::Row {
3095 values: vec![Value::default(), Value::default()],
3096 }],
3097 }),
3098 };
3099 let ctx = Arc::new(QueryContext::with(
3100 DEFAULT_CATALOG_NAME,
3101 DEFAULT_SCHEMA_NAME,
3102 ));
3103
3104 let kv_backend = prepare_mocked_backend().await;
3105 let inserter = Inserter::new(
3106 catalog::memory::MemoryCatalogManager::new(),
3107 create_partition_rule_manager(kv_backend.clone()).await,
3108 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3109 Arc::new(new_table_flownode_set_cache(
3110 String::new(),
3111 Cache::new(100),
3112 kv_backend.clone(),
3113 )),
3114 true,
3115 );
3116 assert!(
3118 inserter
3119 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, false, false, false)
3120 .unwrap()
3121 .is_none()
3122 );
3123 let alter_expr = inserter
3124 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true, true)
3125 .unwrap();
3126 assert!(alter_expr.is_none());
3127
3128 let req_schema = req.rows.as_ref().unwrap().schema.clone();
3130 assert_eq!(req_schema[0].column_name, ts_name);
3131 assert_eq!(req_schema[1].column_name, field_name);
3132
3133 let (datatype, datatype_extension) =
3134 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
3135 .unwrap()
3136 .into_parts();
3137 let mut histogram_req = RowInsertRequest {
3138 table_name: "test_table".to_string(),
3139 rows: Some(Rows {
3140 schema: vec![
3141 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
3142 api::v1::ColumnSchema {
3143 column_name: greptime_native_histogram().to_string(),
3144 datatype: datatype as i32,
3145 semantic_type: SemanticType::Field as i32,
3146 datatype_extension,
3147 options: None,
3148 },
3149 ],
3150 rows: vec![],
3151 }),
3152 };
3153 let error = inserter
3154 .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, false, true, true)
3155 .unwrap_err();
3156 assert!(
3157 error
3158 .to_string()
3159 .contains("cannot mix native histogram and float sample fields")
3160 );
3161
3162 let histogram_table = make_table_ref_with_schema(
3163 "ts",
3164 greptime_native_histogram(),
3165 native_histogram_value_type().clone(),
3166 );
3167 let mut sample_req = RowInsertRequest {
3168 table_name: "test_table".to_string(),
3169 rows: Some(Rows {
3170 schema: vec![
3171 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
3172 field_column_schema(greptime_value(), ColumnDataType::Float64),
3173 ],
3174 rows: vec![],
3175 }),
3176 };
3177 let error = inserter
3178 .get_alter_table_expr_on_demand(
3179 &mut sample_req,
3180 &histogram_table,
3181 &ctx,
3182 false,
3183 true,
3184 true,
3185 )
3186 .unwrap_err();
3187 assert!(
3188 error
3189 .to_string()
3190 .contains("cannot mix native histogram and float sample fields")
3191 );
3192 }
3193
3194 #[test]
3195 fn test_native_histogram_detection_survives_prefix_change() {
3196 set_default_prefix(Some("custom")).unwrap();
3197 let table = make_table_ref_with_schema(
3198 "custom_timestamp",
3199 NATIVE_HISTOGRAM_FIELD,
3200 native_histogram_value_type().clone(),
3201 );
3202 let (datatype, datatype_extension) =
3203 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
3204 .unwrap()
3205 .into_parts();
3206 let request_schema = [api::v1::ColumnSchema {
3207 column_name: greptime_native_histogram().to_string(),
3208 datatype: datatype as i32,
3209 semantic_type: SemanticType::Field as i32,
3210 datatype_extension,
3211 options: None,
3212 }];
3213
3214 assert!(request_is_native_histogram(&request_schema));
3215 assert!(table_is_native_histogram(&table));
3216 }
3217
3218 #[derive(Clone)]
3219 struct GatedDatanodeHandler {
3220 started: mpsc::UnboundedSender<u32>,
3221 completed: mpsc::UnboundedSender<u32>,
3222 release: Arc<Semaphore>,
3223 }
3224
3225 #[async_trait::async_trait]
3226 impl Datanode for GatedDatanodeHandler {
3227 async fn handle_request(
3228 &self,
3229 request: Arc<api::v1::region::RegionRequest>,
3230 ) -> common_meta::error::Result<RegionResponse> {
3231 let Some(api::v1::region::region_request::Body::Inserts(requests)) = &request.body
3232 else {
3233 panic!("expected inserts");
3234 };
3235 let region = RegionId::from(requests.requests[0].region_id).region_number();
3236 self.started.send(region).unwrap();
3237 self.release.acquire().await.unwrap().forget();
3238 self.completed.send(region).unwrap();
3239 if region == 2 {
3240 Err(common_meta::error::Error::retry_later(MockError::new(
3241 StatusCode::RegionBusy,
3242 )))
3243 } else {
3244 Ok(RegionResponse::new(1))
3245 }
3246 }
3247
3248 async fn handle_query(
3249 &self,
3250 _request: common_query::request::QueryRequest,
3251 ) -> common_meta::error::Result<common_recordbatch::SendableRecordBatchStream> {
3252 unreachable!()
3253 }
3254 }
3255
3256 #[async_trait::async_trait]
3257 impl DatanodeManager for GatedDatanodeHandler {
3258 async fn datanode(&self, _peer: &Peer) -> DatanodeRef {
3259 Arc::new(self.clone())
3260 }
3261 }
3262
3263 #[async_trait::async_trait]
3264 impl FlownodeManager for GatedDatanodeHandler {
3265 async fn flownode(&self, _peer: &Peer) -> FlownodeRef {
3266 unreachable!()
3267 }
3268 }
3269
3270 #[tokio::test]
3271 async fn test_insert_cancellation_drains_inflight_writes() {
3272 for confirmed in [false, true] {
3273 for max_retries in [0, 2] {
3274 let mut inserter = batcher_test_inserter().await;
3275 let (started, mut starts) = mpsc::unbounded_channel();
3276 let (completed, mut completions) = mpsc::unbounded_channel();
3277 let release = Arc::new(Semaphore::new(0));
3278 let node_manager = Arc::new(GatedDatanodeHandler {
3279 started,
3280 completed,
3281 release: release.clone(),
3282 });
3283 inserter.node_manager = node_manager.clone();
3284 let dispatcher = Arc::new(
3285 WriteDispatcher::new(node_manager, inserter.partition_manager.clone())
3286 .with_retry(
3287 WriteRetryOptions {
3288 max_retries,
3289 initial_backoff: Duration::ZERO,
3290 max_backoff: Duration::ZERO,
3291 ..Default::default()
3292 },
3293 inserter.catalog_manager.clone(),
3294 Arc::new(DummyCacheInvalidator),
3295 ),
3296 );
3297 let inserter = inserter.with_write_dispatcher(dispatcher);
3298 let table_infos = HashMap::from_iter([(
3299 1,
3300 Arc::new(new_test_table_info(1, "table_1", [1].into_iter())),
3301 )]);
3302 let counter = Arc::new(session::write_progress::WriteProgress::default());
3303 let mut ctx = QueryContext::arc().fork();
3304 ctx.set_write_progress(counter.clone());
3305 let ctx = Arc::new(ctx);
3306 let ctx = if confirmed {
3307 Arc::new(ctx.with_confirmed_writes())
3308 } else {
3309 ctx
3310 };
3311 let requests = InstantAndNormalInsertRequests {
3312 normal_requests: RegionInsertRequests {
3313 requests: [1, 2]
3314 .map(|region| RegionInsertRequest {
3315 region_id: RegionId::new(1, region).as_u64(),
3316 rows: Some(Rows {
3317 schema: vec![],
3318 rows: vec![Row { values: vec![] }],
3319 }),
3320 ..Default::default()
3321 })
3322 .to_vec(),
3323 },
3324 instant_requests: RegionInsertRequests::default(),
3325 };
3326 let caller =
3327 tokio::spawn(
3328 async move { inserter.do_request(requests, &table_infos, &ctx).await },
3329 );
3330 for _ in 0..2 {
3331 tokio::time::timeout(Duration::from_secs(5), starts.recv())
3332 .await
3333 .unwrap()
3334 .unwrap();
3335 }
3336 caller.abort();
3337 assert!(caller.await.unwrap_err().is_cancelled());
3338 release.add_permits(2);
3339 let mut completed = tokio::time::timeout(Duration::from_secs(5), async {
3341 let mut regions = Vec::new();
3342 while let Some(region) = completions.recv().await {
3343 regions.push(region);
3344 }
3345 regions
3346 })
3347 .await
3348 .unwrap();
3349 completed.sort_unstable();
3350 assert_eq!(
3351 completed,
3352 vec![1, 2],
3353 "confirmed={confirmed}, retries={max_retries}"
3354 );
3355 assert!(
3356 starts.try_recv().is_err(),
3357 "cancelled caller started another retry"
3358 );
3359 assert_eq!(counter.written(), 1);
3360 }
3361 }
3362 }
3363
3364 #[tokio::test]
3365 async fn test_tracked_cancellation_stops_backoff() {
3366 let inserter = batcher_test_inserter().await;
3367 let (started, mut starts) = mpsc::unbounded_channel();
3368 let (completed, mut completions) = mpsc::unbounded_channel();
3369 let node_manager = Arc::new(GatedDatanodeHandler {
3370 started,
3371 completed,
3372 release: Arc::new(Semaphore::new(2)),
3373 });
3374 let dispatcher = WriteDispatcher::new(node_manager, inserter.partition_manager.clone())
3375 .with_retry(
3376 WriteRetryOptions {
3377 max_retries: 2,
3378 initial_backoff: Duration::from_secs(60),
3379 max_backoff: Duration::from_secs(60),
3380 ..Default::default()
3381 },
3382 inserter.catalog_manager.clone(),
3383 Arc::new(DummyCacheInvalidator),
3384 );
3385 let factory = RegionRequestFactory::new(RegionRequestHeader::default());
3386 let requests = [1, 2]
3387 .into_iter()
3388 .map(|region| TrackedWrite {
3389 request: factory.build_insert(RegionInsertRequests {
3390 requests: vec![RegionInsertRequest {
3391 region_id: RegionId::new(1, region).as_u64(),
3392 ..Default::default()
3393 }],
3394 }),
3395 owners: vec![(region - 1) as usize],
3396 })
3397 .collect();
3398 let (writes, progress) = dispatcher.prepare_tracked(requests, 2).await;
3399 let cancellation = CancellationToken::new();
3400 let execute = dispatcher.execute_tracked(writes, progress, None, Some(&cancellation));
3401 tokio::pin!(execute);
3402 tokio::time::timeout(
3404 Duration::from_secs(1),
3405 future::poll_fn(|cx| {
3406 assert!(execute.as_mut().poll(cx).is_pending());
3407 if completions.len() == 2 {
3408 Poll::Ready(())
3409 } else {
3410 Poll::Pending
3411 }
3412 }),
3413 )
3414 .await
3415 .unwrap();
3416 for region in [1, 2] {
3417 assert_eq!(starts.try_recv().unwrap(), region);
3418 assert_eq!(completions.try_recv().unwrap(), region);
3419 }
3420 cancellation.cancel();
3421 let progress = tokio::time::timeout(Duration::from_secs(1), execute)
3422 .await
3423 .unwrap();
3424 assert_eq!(progress[0].affected_rows, 1);
3425 assert!(progress[0].error.is_none());
3426 assert_eq!(progress[1].affected_rows, 0);
3427 assert_eq!(
3428 progress[1].error.as_ref().unwrap().retry_hint(),
3429 RetryHint::Retryable
3430 );
3431 assert!(starts.try_recv().is_err());
3432 }
3433
3434 #[tokio::test]
3436 async fn test_write_meter_admission() {
3437 use std::cell::Cell;
3438 use std::sync::Mutex;
3439 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
3440
3441 use api::region::RegionResponse;
3442 use api::v1::region::region_request::Body;
3443 use arrow::array::{Int32Array, TimestampMillisecondArray};
3444 use arrow::record_batch::RecordBatch;
3445 use bytes::Bytes;
3446 use common_error::ext::{ErrorExt, RetryHint};
3447 use common_error::status_code::StatusCode;
3448 use common_grpc::flight::{FlightEncoder, FlightMessage};
3449 use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher;
3450 use futures::future::BoxFuture;
3451 use meter_core::ItemCalculator;
3452 use meter_core::collect::{Collect, WriteRejected};
3453 use meter_core::data::MeterRecord;
3454 use meter_core::global::global_registry;
3455 use session::context::Channel;
3456
3457 const CATALOG: &str = "write_meter_test";
3458
3459 #[derive(Default)]
3460 struct Meter {
3461 reject: AtomicBool,
3462 attempts: Mutex<Vec<MeterRecord>>,
3463 accepted_value: AtomicU64,
3464 }
3465
3466 impl Collect for Meter {
3467 fn on_write(
3468 &self,
3469 record: MeterRecord,
3470 ) -> BoxFuture<'_, std::result::Result<(), WriteRejected>> {
3471 Box::pin(async move {
3472 if record.catalog != CATALOG {
3473 return Ok(());
3474 }
3475 let value = record.value;
3476 self.attempts.lock().unwrap().push(record);
3477 if self.reject.load(Ordering::Relaxed) {
3478 return Err(WriteRejected::new("database row quota exhausted"));
3479 }
3480 self.accepted_value.fetch_add(value, Ordering::Relaxed);
3481 Ok(())
3482 })
3483 }
3484
3485 fn on_read(&self, _: MeterRecord) {}
3486 }
3487
3488 impl ItemCalculator<InstantAndNormalInsertRequests> for Meter {
3489 fn calc(&self, _: &InstantAndNormalInsertRequests) -> u64 {
3490 17
3491 }
3492 }
3493
3494 let kv_backend = prepare_mocked_backend().await;
3495 let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
3496 let (sender, mut dispatched) = tokio::sync::mpsc::channel(16);
3497 let watcher = DatanodeWatcher::new(sender).with_handler(|_, request| {
3498 let rows = match request.body.unwrap() {
3499 Body::Inserts(requests) => requests
3500 .requests
3501 .iter()
3502 .filter_map(|request| request.rows.as_ref())
3503 .map(|rows| rows.rows.len())
3504 .sum(),
3505 Body::BulkInsert(_) => 2,
3507 body => panic!("unexpected request: {body:?}"),
3508 };
3509 Ok(RegionResponse::new(rows))
3510 });
3511 let flow_cache = Cache::new(100);
3512 let inserter = Inserter::new(
3513 catalog::memory::MemoryCatalogManager::new(),
3514 partition_manager,
3515 Arc::new(MockDatanodeManager::new(watcher)),
3516 Arc::new(new_table_flownode_set_cache(
3517 String::new(),
3518 flow_cache.clone(),
3519 kv_backend,
3520 )),
3521 true,
3522 );
3523 let mut table_info = new_test_table_info(1, "table_1", [1].into_iter());
3524 table_info.catalog_name = CATALOG.to_string();
3525 table_info.schema_name = "target_db".to_string();
3526 let table_info = Arc::new(table_info);
3527 let table_infos = HashMap::from_iter([(1, table_info.clone())]);
3528 let ctx = Arc::new(QueryContext::with_channel(
3529 DEFAULT_CATALOG_NAME,
3530 DEFAULT_SCHEMA_NAME,
3531 Channel::Postgres,
3532 ));
3533 let rows_request = || {
3534 let build = |num_rows| RegionInsertRequests {
3535 requests: vec![RegionInsertRequest {
3536 region_id: RegionId::new(1, 1).as_u64(),
3537 rows: Some(Rows {
3538 schema: vec![],
3539 rows: vec![api::v1::Row { values: vec![] }; num_rows],
3540 }),
3541 ..Default::default()
3542 }],
3543 };
3544 InstantAndNormalInsertRequests {
3545 normal_requests: build(3),
3546 instant_requests: build(2),
3547 }
3548 };
3549
3550 let output = inserter
3552 .do_request(rows_request(), &table_infos, &ctx)
3553 .await
3554 .unwrap();
3555 assert_eq!(output.meta.cost, 0);
3556 assert!(matches!(output.data, OutputData::AffectedRows(3)));
3557 dispatched.try_recv().unwrap();
3558
3559 let meter = Arc::new(Meter::default());
3560 global_registry().set_collector(meter.clone());
3561 global_registry().register_calculator(meter.clone());
3562 let enabled = Cell::new(false);
3564 write_meter!({
3565 enabled.set(true);
3566 MeterRecord::new("probe".into(), "probe".into(), 0, 0, 0)
3567 })
3568 .await
3569 .unwrap();
3570 let enabled = enabled.get();
3571
3572 let output = inserter
3573 .do_request(rows_request(), &table_infos, &ctx)
3574 .await
3575 .unwrap();
3576 assert_eq!(output.meta.cost, if enabled { 17 } else { 0 });
3577 assert!(matches!(output.data, OutputData::AffectedRows(3)));
3578 dispatched.try_recv().unwrap();
3579 assert!(flow_cache.contains_key(&1));
3580 flow_cache.invalidate_all();
3581 meter.reject.store(true, Ordering::Relaxed);
3582 let result = inserter
3583 .do_request(rows_request(), &table_infos, &ctx)
3584 .await;
3585 if enabled {
3586 let error = result.unwrap_err();
3587 assert_eq!(error.status_code(), StatusCode::RateLimited);
3588 assert_eq!(error.retry_hint(), RetryHint::Retryable);
3589 assert!(error.to_string().contains("database row quota exhausted"));
3590 assert!(dispatched.try_recv().is_err());
3591 assert!(
3592 !flow_cache.contains_key(&1),
3593 "rejected write reached flow mirroring"
3594 );
3595 assert_eq!(meter.accepted_value.load(Ordering::Relaxed), 17);
3596 let attempts = meter.attempts.lock().unwrap();
3597 assert_eq!(attempts.len(), 2);
3598 for record in attempts.iter() {
3599 assert_eq!(record.catalog, CATALOG);
3600 assert_eq!(record.schema, "target_db");
3601 assert_eq!(
3602 (record.rows, record.value, record.source),
3603 (5, 17, Channel::Postgres as u8)
3604 );
3605 }
3606 } else {
3607 assert_eq!(result.unwrap().meta.cost, 0);
3608 dispatched.try_recv().unwrap();
3609 assert!(meter.attempts.lock().unwrap().is_empty());
3610 }
3611 meter.attempts.lock().unwrap().clear();
3612
3613 if enabled {
3615 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3616 let rows = Rows {
3617 schema: vec![
3618 api::v1::helper::tag_column_schema("a", ColumnDataType::Int32),
3619 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
3620 field_column_schema("b", ColumnDataType::Int32),
3621 ],
3622 rows: vec![api::v1::Row {
3623 values: vec![
3624 api::v1::value::ValueData::I32Value(60).into(),
3625 Value {
3626 value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue(
3627 0,
3628 )),
3629 },
3630 api::v1::value::ValueData::I32Value(0).into(),
3631 ],
3632 }],
3633 };
3634 let error = inserter
3635 .submit_table_rows(rows, table_info.clone(), ctx.clone(), &batcher)
3636 .await
3637 .unwrap_err();
3638 assert_eq!(error.status_code(), StatusCode::RateLimited);
3639 let mut attempts = meter.attempts.lock().unwrap();
3640 assert_eq!(attempts.len(), 1);
3641 assert_eq!(attempts[0].catalog, CATALOG);
3642 assert_eq!(attempts[0].schema, "target_db");
3643 assert_eq!(attempts[0].rows, 1);
3644 attempts.clear();
3645 }
3646
3647 let table = Arc::new(table::Table::new(
3648 table_info.clone(),
3649 table::metadata::FilterPushDownType::Unsupported,
3650 Arc::new(DummyDataSource),
3651 ));
3652 let batch = RecordBatch::try_new(
3653 table_info.meta.schema.arrow_schema().clone(),
3654 vec![
3655 Arc::new(Int32Array::from(vec![60, 70])),
3656 Arc::new(TimestampMillisecondArray::from(vec![0, 1])),
3657 Arc::new(Int32Array::from(vec![0, 0])),
3658 ],
3659 )
3660 .unwrap();
3661 let bulk_insert = |batch: RecordBatch| {
3662 let flight_data = FlightEncoder::default()
3663 .encode(FlightMessage::RecordBatch(batch.clone()))
3664 .into_iter()
3665 .next()
3666 .unwrap();
3667 inserter.handle_bulk_insert(
3668 table.clone(),
3669 flight_data,
3670 batch,
3671 Bytes::new(),
3672 false,
3673 Channel::Grpc,
3674 )
3675 };
3676 assert_eq!(bulk_insert(batch.slice(0, 0)).await.unwrap(), 0);
3678 assert!(meter.attempts.lock().unwrap().is_empty());
3679 assert!(dispatched.try_recv().is_err());
3680
3681 for reject in [false, true, false] {
3683 meter.reject.store(reject, Ordering::Relaxed);
3684 let result = bulk_insert(batch.clone()).await;
3685 if enabled && reject {
3686 let error = result.unwrap_err();
3687 assert_eq!(error.status_code(), StatusCode::RateLimited);
3688 assert_eq!(error.retry_hint(), RetryHint::Retryable);
3689 assert!(dispatched.try_recv().is_err());
3690 } else {
3691 assert_eq!(result.unwrap(), 2);
3692 dispatched.try_recv().unwrap();
3693 }
3694 }
3695 {
3696 let attempts = meter.attempts.lock().unwrap();
3697 assert_eq!(attempts.len(), if enabled { 3 } else { 0 });
3698 for record in attempts.iter() {
3699 assert_eq!(record.catalog, CATALOG);
3700 assert_eq!(record.schema, "target_db");
3701 assert_eq!(
3702 (record.rows, record.value, record.source),
3703 (2, 0, Channel::Grpc as u8)
3704 );
3705 }
3706 }
3707 assert_eq!(
3708 meter.accepted_value.load(Ordering::Relaxed),
3709 if enabled { 17 } else { 0 }
3710 );
3711
3712 meter.attempts.lock().unwrap().clear();
3715 let original = Arc::new(QueryContext::with_channel(
3716 CATALOG,
3717 "a",
3718 Channel::Prometheus,
3719 ));
3720 let mut batches = ["a", "b", "a"].map(|schema| {
3721 let ctx = if schema == "a" {
3722 original.clone()
3723 } else {
3724 Arc::new(QueryContext::with_channel(
3725 CATALOG,
3726 schema,
3727 Channel::Prometheus,
3728 ))
3729 };
3730 (
3731 ctx,
3732 RowInsertRequests {
3733 inserts: vec![RowInsertRequest {
3734 table_name: "data".into(),
3735 rows: Some(Rows {
3736 schema: vec![],
3737 rows: vec![api::v1::Row::default(); 2],
3738 }),
3739 }],
3740 },
3741 )
3742 });
3743 admit_row_insert_batches(&mut batches).await.unwrap();
3744 admit_row_insert_batches(&mut batches).await.unwrap();
3745 assert_eq!(original.write_rows_to_admit(CATALOG, "a", 4), 4);
3746 for (ctx, _) in &batches {
3747 assert_eq!(
3748 ctx.write_rows_to_admit(CATALOG, &ctx.current_schema(), 2),
3749 0
3750 );
3751 assert_eq!(ctx.channel(), Channel::Prometheus);
3752 }
3753 {
3754 let attempts = meter.attempts.lock().unwrap();
3755 let totals = attempts
3756 .iter()
3757 .map(|r| (r.schema.as_str(), r.rows, r.value))
3758 .collect::<Vec<_>>();
3759 assert_eq!(
3760 totals,
3761 if enabled {
3762 vec![("a", 4, 0), ("b", 2, 0)]
3763 } else {
3764 vec![]
3765 }
3766 );
3767 }
3768 meter.attempts.lock().unwrap().clear();
3769 meter.reject.store(false, Ordering::Relaxed);
3770 let ctx = Arc::new(QueryContext::with_channel(
3771 CATALOG,
3772 "logical",
3773 Channel::Otlp,
3774 ));
3775 let admitted = admit_write(2, &ctx).await.unwrap();
3776 let mut requests = RowInsertRequests {
3777 inserts: vec![RowInsertRequest {
3778 table_name: "metric".to_string(),
3779 rows: Some(Rows {
3780 schema: vec![],
3781 rows: vec![api::v1::Row::default(); 2],
3782 }),
3783 }],
3784 };
3785 let original = requests.clone();
3786 let cost = Inserter::meter_row_inserts(&mut requests, &admitted)
3787 .await
3788 .unwrap();
3789 assert_eq!(cost, if enabled { 17 } else { 0 });
3790 assert_eq!(requests, original);
3791 let records = meter
3792 .attempts
3793 .lock()
3794 .unwrap()
3795 .iter()
3796 .map(|record| (record.rows, record.value))
3797 .collect::<Vec<_>>();
3798 assert_eq!(
3799 records,
3800 if enabled {
3801 vec![(2, 0), (0, 17)]
3802 } else {
3803 vec![]
3804 }
3805 );
3806 meter.reject.store(true, Ordering::Relaxed);
3807 let result = Inserter::meter_row_inserts(&mut requests, &admitted).await;
3808 if enabled {
3809 assert_eq!(result.unwrap_err().status_code(), StatusCode::RateLimited);
3810 } else {
3811 assert_eq!(result.unwrap(), 0);
3812 }
3813 assert_eq!(requests, original);
3814 }
3815
3816 #[test]
3817 fn test_trace_create_options_copy_compaction_hints() {
3818 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3819 ctx.set_extension("compaction.twcs.time_window", "1h");
3820 ctx.set_extension("compaction.twcs.max_output_file_size", "256MB");
3821 ctx.set_extension("compaction.twcs.remote_compaction", "true");
3822 ctx.set_extension("compaction.unknown", "ignored");
3823 let ctx = Arc::new(ctx);
3824 let mut table_options = Default::default();
3825
3826 fill_table_options_for_create(
3827 &mut table_options,
3828 &AutoCreateTableType::Trace {
3829 alter_existing: true,
3830 },
3831 &ctx,
3832 );
3833
3834 assert_eq!(
3835 Some("twcs"),
3836 table_options.get(COMPACTION_TYPE).map(String::as_str)
3837 );
3838 assert_eq!(
3839 Some("1h"),
3840 table_options
3841 .get("compaction.twcs.time_window")
3842 .map(String::as_str)
3843 );
3844 assert_eq!(
3845 Some("256MB"),
3846 table_options
3847 .get("compaction.twcs.max_output_file_size")
3848 .map(String::as_str)
3849 );
3850 assert_eq!(
3851 Some("true"),
3852 table_options
3853 .get("compaction.twcs.remote_compaction")
3854 .map(String::as_str)
3855 );
3856 assert!(!table_options.contains_key("compaction.unknown"));
3857 }
3858
3859 #[test]
3860 fn test_exemplar_create_options_isolate_metric_context() {
3861 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3862 for (key, value) in [
3863 (APPEND_MODE_KEY, "false"),
3864 (MERGE_MODE_KEY, "last_non_null"),
3865 (PHYSICAL_TABLE_METADATA_KEY, "true"),
3866 (LOGICAL_TABLE_METADATA_KEY, "physical_metrics"),
3867 ("otlp_metric_compat", "prometheus"),
3868 (table::requests::SEMANTIC_SIGNAL_TYPE, "metric"),
3869 (SEMANTIC_PER_TABLE_INDEX_KEY, "{}"),
3870 ("ttl", "7d"),
3871 ("compaction.twcs.time_window", "1h"),
3872 ] {
3873 ctx.set_extension(key, value);
3874 }
3875 let mut options = Default::default();
3876 fill_table_options_for_create(&mut options, &AutoCreateTableType::Exemplar, &Arc::new(ctx));
3877 assert_eq!(
3878 options,
3879 std::collections::HashMap::from([
3880 (APPEND_MODE_KEY.into(), "true".into()),
3881 (
3882 store_api::metric_engine_consts::PRIMARY_KEY_ENCODING.into(),
3883 "sparse".into()
3884 ),
3885 ("ttl".into(), "7d".into()),
3886 (COMPACTION_TYPE.into(), "twcs".into()),
3887 ("compaction.twcs.time_window".into(), "1h".into()),
3888 ])
3889 );
3890 }
3891
3892 #[test]
3893 fn test_skip_wal_does_not_change_table_options() {
3894 check_skip_wal_does_not_change_table_options(false);
3895 check_skip_wal_does_not_change_table_options(true);
3896 }
3897
3898 fn check_skip_wal_does_not_change_table_options(skip_wal: bool) {
3899 let ctx = Arc::new(QueryContext::with(
3900 DEFAULT_CATALOG_NAME,
3901 DEFAULT_SCHEMA_NAME,
3902 ));
3903 ctx.set_skip_wal(skip_wal);
3904 let mut options = Default::default();
3905 fill_table_options_for_create(&mut options, &AutoCreateTableType::Physical, &ctx);
3906 assert!(!options.contains_key(session::hints::INSERT_SKIP_WAL_HINT));
3907 assert!(!options.contains_key("skip_wal"));
3908 }
3909
3910 #[test]
3911 fn test_last_non_null_create_options_preserve_default_without_append_mode() {
3912 let ctx = Arc::new(QueryContext::with(
3913 DEFAULT_CATALOG_NAME,
3914 DEFAULT_SCHEMA_NAME,
3915 ));
3916 let mut table_options = Default::default();
3917
3918 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3919
3920 assert_eq!(
3921 Some("last_non_null"),
3922 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3923 );
3924 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3925 }
3926
3927 #[test]
3928 fn test_fill_table_options_copies_semantic_extensions() {
3929 use table::requests::{
3930 SEMANTIC_METRIC_TYPE, SEMANTIC_PER_TABLE_INDEX_KEY, SEMANTIC_SIGNAL_TYPE,
3931 SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, SOURCE_OPENTELEMETRY,
3932 };
3933
3934 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3935 ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
3936 ctx.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
3937 ctx.set_extension(SEMANTIC_SOURCE_VERSION, "2.0");
3938 ctx.set_extension(SEMANTIC_METRIC_TYPE, "bogus");
3939 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, "{}");
3941 let ctx = Arc::new(ctx);
3942 let mut table_options = Default::default();
3943
3944 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::Physical, &ctx);
3945
3946 assert_eq!(
3947 Some(SIGNAL_TYPE_METRIC),
3948 table_options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str)
3949 );
3950 assert_eq!(
3951 Some(SOURCE_OPENTELEMETRY),
3952 table_options.get(SEMANTIC_SOURCE).map(String::as_str)
3953 );
3954 assert_eq!(
3955 Some("2.0"),
3956 table_options
3957 .get(SEMANTIC_SOURCE_VERSION)
3958 .map(String::as_str)
3959 );
3960 assert!(!table_options.contains_key(SEMANTIC_METRIC_TYPE));
3961 assert!(!table_options.contains_key(SEMANTIC_PER_TABLE_INDEX_KEY));
3962 }
3963
3964 #[test]
3965 fn test_apply_per_table_semantic_options() {
3966 use table::requests::{
3967 SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, SEMANTIC_PER_TABLE_INDEX_KEY,
3968 };
3969
3970 let index = format!(
3971 r#"{{
3972 "{DEFAULT_SCHEMA_NAME}": {{
3973 "http_requests_total": {{
3974 "greptime.semantic.metric.type": "counter",
3975 "greptime.semantic.metric.unit": "By",
3976 "greptime.semantic.metric.type_BOGUS": "x"
3977 }},
3978 "other_table": {{
3979 "greptime.semantic.metric.type": "gauge"
3980 }}
3981 }},
3982 "other_schema": {{
3983 "http_requests_total": {{
3984 "greptime.semantic.metric.type": "gauge"
3985 }}
3986 }}
3987 }}"#
3988 );
3989 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3990 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index);
3991 let ctx = Arc::new(ctx);
3992
3993 let index = parse_per_table_semantic_index(&ctx);
3994 assert!(index.is_some());
3995 let index = index.as_ref();
3996
3997 let mut table_options = std::collections::HashMap::new();
3998 apply_per_table_semantic_options(
3999 &mut table_options,
4000 index,
4001 DEFAULT_SCHEMA_NAME,
4002 "http_requests_total",
4003 );
4004 assert_eq!(
4006 table_options.get(SEMANTIC_METRIC_TYPE).map(String::as_str),
4007 Some("counter")
4008 );
4009 assert_eq!(
4010 table_options.get(SEMANTIC_METRIC_UNIT).map(String::as_str),
4011 Some("By")
4012 );
4013 assert!(!table_options.contains_key("greptime.semantic.metric.type_BOGUS"));
4016 assert_eq!(table_options.len(), 2);
4017
4018 let mut empty = std::collections::HashMap::new();
4019 apply_per_table_semantic_options(&mut empty, index, DEFAULT_SCHEMA_NAME, "not_in_index");
4020 assert!(empty.is_empty());
4021
4022 let mut opts = std::collections::HashMap::new();
4025 apply_per_table_semantic_options(
4026 &mut opts,
4027 index,
4028 "schema_without_entry",
4029 "http_requests_total",
4030 );
4031 assert!(opts.is_empty());
4032
4033 let bare = Arc::new(QueryContext::with(
4035 DEFAULT_CATALOG_NAME,
4036 DEFAULT_SCHEMA_NAME,
4037 ));
4038 assert!(parse_per_table_semantic_index(&bare).is_none());
4039 let mut opts = std::collections::HashMap::new();
4040 apply_per_table_semantic_options(
4041 &mut opts,
4042 None,
4043 DEFAULT_SCHEMA_NAME,
4044 "http_requests_total",
4045 );
4046 assert!(opts.is_empty());
4047 }
4048
4049 #[test]
4050 fn test_last_non_null_create_options_preserve_default_with_append_mode_false() {
4051 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
4052 ctx.set_extension(APPEND_MODE_KEY, "false");
4053 let ctx = Arc::new(ctx);
4054 let mut table_options = Default::default();
4055
4056 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
4057
4058 assert!(!table_options.contains_key(APPEND_MODE_KEY));
4059 assert_eq!(
4060 Some("last_non_null"),
4061 table_options.get(MERGE_MODE_KEY).map(String::as_str)
4062 );
4063 }
4064
4065 #[test]
4066 fn test_last_non_null_create_options_use_configured_merge_mode() {
4067 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
4068 ctx.set_extension(MERGE_MODE_KEY, "last_row");
4069 let ctx = Arc::new(ctx);
4070 let mut table_options = Default::default();
4071
4072 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
4073
4074 assert_eq!(
4075 Some("last_row"),
4076 table_options.get(MERGE_MODE_KEY).map(String::as_str)
4077 );
4078 assert!(!table_options.contains_key(APPEND_MODE_KEY));
4079 }
4080
4081 #[test]
4082 fn test_last_non_null_create_options_use_last_row_with_append_mode_true() {
4083 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
4084 ctx.set_extension(APPEND_MODE_KEY, "true");
4085 let ctx = Arc::new(ctx);
4086 let mut table_options = Default::default();
4087
4088 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
4089
4090 assert_eq!(
4091 Some("true"),
4092 table_options.get(APPEND_MODE_KEY).map(String::as_str)
4093 );
4094 assert_eq!(
4095 Some("last_row"),
4096 table_options.get(MERGE_MODE_KEY).map(String::as_str)
4097 );
4098 }
4099
4100 struct UnexpectedBatcher;
4101
4102 #[async_trait::async_trait]
4103 impl PendingRowsBatcher for UnexpectedBatcher {
4104 async fn acquire(&self) -> Result<Arc<tokio::sync::OwnedSemaphorePermit>> {
4105 panic!("empty writes must not acquire batch admission")
4106 }
4107
4108 async fn submit(
4109 &self,
4110 _table_info: TableInfoRef,
4111 _batch: arrow::record_batch::RecordBatch,
4112 _ctx: QueryContextRef,
4113 _permit: Arc<tokio::sync::OwnedSemaphorePermit>,
4114 ) -> Result<Output> {
4115 panic!("empty writes must not submit a batch")
4116 }
4117 }
4118
4119 async fn batcher_test_inserter() -> Inserter {
4120 let kv_backend = prepare_mocked_backend().await;
4121 Inserter::new(
4122 catalog::memory::MemoryCatalogManager::new(),
4123 create_partition_rule_manager(kv_backend.clone()).await,
4124 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
4125 Arc::new(new_table_flownode_set_cache(
4126 String::new(),
4127 Cache::new(100),
4128 kv_backend,
4129 )),
4130 true,
4131 )
4132 }
4133
4134 #[tokio::test]
4135 async fn test_logical_batcher_eligibility() {
4136 use catalog::RegisterTableRequest;
4137 use catalog::memory::MemoryCatalogManager;
4138 use common_meta::instruction::{CacheIdent, CreateFlow};
4139 use common_meta::kv_backend::KvBackendRef;
4140 use common_meta::kv_backend::memory::MemoryKvBackend;
4141 use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder};
4142 let requests = RowInsertRequests {
4143 inserts: vec![RowInsertRequest {
4144 table_name: "test_table".to_string(),
4145 rows: None,
4146 }],
4147 };
4148 let original =
4149 make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
4150 .table_info();
4151 for case in [
4152 "eligible",
4153 "physical",
4154 "ordinary",
4155 "instant",
4156 "disabled",
4157 "hint",
4158 "flow",
4159 "required_tag",
4160 "default",
4161 ] {
4162 let mut info = (*original).clone();
4163 info.meta.engine = METRIC_ENGINE_NAME.to_string();
4164 info.meta.options.extra_options.insert(
4165 LOGICAL_TABLE_METADATA_KEY.to_string(),
4166 "physical".to_string(),
4167 );
4168 let mut ctx = QueryContext::arc().fork();
4169 match case {
4170 "physical" => {
4171 info.meta
4172 .options
4173 .extra_options
4174 .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "other".to_string());
4175 }
4176 "ordinary" => info.meta.engine = "mito".to_string(),
4177 "instant" => info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant),
4178 "hint" => ctx.set_extension(AUTO_CREATE_TABLE_KEY, "false"),
4179 "required_tag" | "default" => {
4180 let mut columns = info.meta.schema.column_schemas().to_vec();
4181 let mut tag = ColumnSchema::new(
4182 "tag",
4183 ConcreteDataType::string_datatype(),
4184 case != "required_tag",
4185 );
4186 if case == "default" {
4187 tag = tag
4188 .with_default_constraint(Some(ColumnDefaultConstraint::null_value()))
4189 .unwrap();
4190 }
4191 columns.push(tag);
4192 info.meta.schema = Arc::new(
4193 SchemaBuilder::try_from_columns(columns)
4194 .unwrap()
4195 .build()
4196 .unwrap(),
4197 );
4198 info.meta.primary_key_indices = vec![2];
4199 }
4200 _ => {}
4201 }
4202 let catalog = MemoryCatalogManager::with_default_setup();
4203 let table = Arc::new(table::Table::new(
4204 Arc::new(info),
4205 table::metadata::FilterPushDownType::Unsupported,
4206 Arc::new(DummyDataSource),
4207 ));
4208 catalog
4209 .register_table_sync(RegisterTableRequest {
4210 catalog: DEFAULT_CATALOG_NAME.to_string(),
4211 schema: DEFAULT_SCHEMA_NAME.to_string(),
4212 table_name: "test_table".to_string(),
4213 table_id: 1,
4214 table,
4215 })
4216 .unwrap();
4217 let mut inserter = batcher_test_inserter().await;
4218 inserter.catalog_manager = catalog;
4219 inserter.auto_create_table = case != "disabled";
4220 let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::default());
4221 inserter.table_flownode_set_cache = Arc::new(new_table_flownode_set_cache(
4222 String::new(),
4223 Cache::new(10),
4224 kv_backend,
4225 ));
4226 if case == "flow" {
4227 inserter
4228 .table_flownode_set_cache
4229 .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
4230 flow_id: 1,
4231 source_table_ids: vec![1],
4232 partition_to_peer_mapping: vec![(0, Peer::empty(1))],
4233 })])
4234 .await
4235 .unwrap();
4236 }
4237 assert_eq!(
4238 inserter
4239 .can_batch_metric_rows(&requests, &Arc::new(ctx), "physical")
4240 .await
4241 .unwrap(),
4242 case == "eligible",
4243 "{case}"
4244 );
4245 }
4246 }
4247
4248 #[tokio::test]
4249 async fn test_batcher_meter_preserves_request() {
4250 let mut requests = RowInsertRequests {
4251 inserts: vec![RowInsertRequest {
4252 table_name: "sample".to_string(),
4253 rows: Some(Rows {
4254 schema: vec![],
4255 rows: vec![api::v1::Row {
4256 values: vec![Value {
4257 value_data: Some(api::v1::value::ValueData::F64Value(1.5)),
4258 }],
4259 }],
4260 }),
4261 }],
4262 };
4263 let expected = requests.clone();
4264 Inserter::meter_row_inserts(&mut requests, &QueryContext::arc())
4265 .await
4266 .unwrap();
4267 assert_eq!(requests, expected);
4268 }
4269
4270 #[tokio::test]
4271 async fn test_instant_table_bypasses_batcher() {
4272 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
4273 let inserter = batcher_test_inserter()
4274 .await
4275 .with_pending_rows_batcher(Some(batcher));
4276 let mut ctx = session::context::QueryContextBuilder::default().build();
4277 ctx.set_batching_enabled(true);
4278 let ctx = Arc::new(ctx);
4279 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
4280 .table_info();
4281 assert!(inserter.table_batcher(&table, &ctx).is_some());
4282 let mut instant = (*table).clone();
4283 instant.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
4284 assert!(inserter.table_batcher(&Arc::new(instant), &ctx).is_none());
4285 }
4286
4287 #[tokio::test]
4288 async fn test_empty_prepared_rows_skip_batcher() {
4289 let inserter = batcher_test_inserter().await;
4290 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
4291 .table_info();
4292 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
4293 let ctx = QueryContext::arc();
4294 let output = inserter
4295 .submit_table_rows(
4296 Rows {
4297 schema: vec![],
4298 rows: vec![],
4299 },
4300 table.clone(),
4301 ctx.clone(),
4302 &batcher,
4303 )
4304 .await
4305 .unwrap();
4306 assert!(matches!(output.data, OutputData::AffectedRows(0)));
4307 let output = inserter
4308 .submit_pending_rows(
4309 RowInsertRequests {
4310 inserts: vec![RowInsertRequest {
4311 table_name: table.name.clone(),
4312 rows: None,
4313 }],
4314 },
4315 HashMap::from_iter([(table.table_id(), table)]),
4316 ctx,
4317 &batcher,
4318 )
4319 .await
4320 .unwrap();
4321 assert!(matches!(output.data, OutputData::AffectedRows(0)));
4322 }
4323
4324 #[derive(Clone)]
4328 struct GatedFlownodeHandler {
4329 gate: Arc<tokio::sync::Semaphore>,
4330 affected_rows: u64,
4331 }
4332
4333 #[async_trait::async_trait]
4334 impl MockFlownodeHandler for GatedFlownodeHandler {
4335 async fn handle_inserts(
4336 &self,
4337 _peer: &Peer,
4338 _requests: api::v1::region::InsertRequests,
4339 ) -> common_meta::error::Result<FlowResponse> {
4340 let _permit = self.gate.acquire().await.unwrap();
4341 Ok(FlowResponse {
4342 affected_rows: self.affected_rows,
4343 ..Default::default()
4344 })
4345 }
4346 }
4347
4348 fn flownode_peer() -> Peer {
4349 Peer {
4350 id: 1,
4351 addr: "127.0.0.1:4001".to_string(),
4352 }
4353 }
4354
4355 #[derive(Clone)]
4356 struct RecordingNodeManager {
4357 datanode_dispatch: tokio::sync::mpsc::UnboundedSender<usize>,
4358 flownode_dispatch: tokio::sync::mpsc::UnboundedSender<Peer>,
4359 gates: Arc<std::sync::Mutex<HashMap<u64, Arc<tokio::sync::Semaphore>>>>,
4360 fail_flownode: bool,
4361 }
4362
4363 #[async_trait::async_trait]
4364 impl common_meta::node_manager::DatanodeManager for RecordingNodeManager {
4365 async fn datanode(&self, peer: &Peer) -> common_meta::node_manager::DatanodeRef {
4366 Arc::new(RecordingNode {
4367 peer: peer.clone(),
4368 manager: self.clone(),
4369 })
4370 }
4371 }
4372
4373 #[async_trait::async_trait]
4374 impl common_meta::node_manager::FlownodeManager for RecordingNodeManager {
4375 async fn flownode(&self, peer: &Peer) -> common_meta::node_manager::FlownodeRef {
4376 Arc::new(RecordingNode {
4377 peer: peer.clone(),
4378 manager: self.clone(),
4379 })
4380 }
4381 }
4382
4383 struct RecordingNode {
4384 peer: Peer,
4385 manager: RecordingNodeManager,
4386 }
4387
4388 #[async_trait::async_trait]
4389 impl common_meta::node_manager::Datanode for RecordingNode {
4390 async fn handle_request(
4391 &self,
4392 _request: Arc<api::v1::region::RegionRequest>,
4393 ) -> common_meta::error::Result<api::region::RegionResponse> {
4394 let rows = match _request.body.as_ref() {
4395 Some(api::v1::region::region_request::Body::Inserts(requests)) => requests
4396 .requests
4397 .iter()
4398 .filter_map(|request| request.rows.as_ref())
4399 .map(|rows| rows.rows.len())
4400 .sum(),
4401 _ => 0,
4402 };
4403 let _ = self.manager.datanode_dispatch.send(rows);
4404 Ok(api::region::RegionResponse::new(rows))
4405 }
4406
4407 async fn handle_query(
4408 &self,
4409 _request: common_query::request::QueryRequest,
4410 ) -> common_meta::error::Result<common_recordbatch::SendableRecordBatchStream> {
4411 unreachable!()
4412 }
4413 }
4414
4415 #[async_trait::async_trait]
4416 impl common_meta::node_manager::Flownode for RecordingNode {
4417 async fn handle(
4418 &self,
4419 _request: api::v1::flow::FlowRequest,
4420 ) -> common_meta::error::Result<FlowResponse> {
4421 unreachable!()
4422 }
4423
4424 async fn handle_inserts(
4425 &self,
4426 _request: api::v1::region::InsertRequests,
4427 ) -> common_meta::error::Result<FlowResponse> {
4428 let _ = self.manager.flownode_dispatch.send(self.peer.clone());
4429 if self.manager.fail_flownode {
4430 return Err(common_meta::error::UnexpectedSnafu {
4431 err_msg: "test flownode failure".to_string(),
4432 }
4433 .build());
4434 }
4435 let gate = self
4436 .manager
4437 .gates
4438 .lock()
4439 .unwrap()
4440 .get(&self.peer.id)
4441 .cloned();
4442 if let Some(gate) = gate {
4443 let _permit = gate.acquire().await.unwrap();
4444 }
4445 Ok(FlowResponse::default())
4446 }
4447
4448 async fn handle_mark_window_dirty(
4449 &self,
4450 _request: api::v1::flow::DirtyWindowRequests,
4451 ) -> common_meta::error::Result<FlowResponse> {
4452 unreachable!()
4453 }
4454 }
4455
4456 fn mirror_requests(peer: &Peer, num_rows: usize) -> HashMap<Peer, RegionInsertRequests> {
4457 HashMap::from_iter([(
4458 peer.clone(),
4459 RegionInsertRequests {
4460 requests: vec![RegionInsertRequest {
4461 region_id: RegionId::new(1, 1).as_u64(),
4462 rows: Some(Rows {
4463 schema: vec![],
4464 rows: vec![api::v1::Row { values: vec![] }; num_rows],
4465 }),
4466 ..Default::default()
4467 }],
4468 },
4469 )])
4470 }
4471
4472 static MIRROR_METRIC_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
4477
4478 async fn lock_mirror_metrics() -> tokio::sync::MutexGuard<'static, ()> {
4479 MIRROR_METRIC_TEST_LOCK.lock().await
4480 }
4481
4482 #[tokio::test]
4486 async fn flow_mirror_saturation_preserves_instant_and_normal_contracts() {
4487 use common_error::ext::{ErrorExt, RetryHint};
4488 use common_error::status_code::StatusCode;
4489
4490 let _guard = lock_mirror_metrics().await;
4491 let kv_backend = prepare_mocked_backend().await;
4492 let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
4493 let (datanode_tx, mut datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4494 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4495 let node_manager = Arc::new(RecordingNodeManager {
4496 datanode_dispatch: datanode_tx,
4497 flownode_dispatch: flownode_tx,
4498 gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
4499 fail_flownode: false,
4500 });
4501 let flow_cache = Cache::new(10);
4502 let flow_cache_backend = prepare_mocked_backend().await;
4503 let inserter = Inserter::new(
4504 catalog::memory::MemoryCatalogManager::new(),
4505 partition_manager,
4506 node_manager,
4507 Arc::new(new_table_flownode_set_cache(
4508 String::new(),
4509 flow_cache.clone(),
4510 flow_cache_backend,
4511 )),
4512 true,
4513 );
4514 let peer = flownode_peer();
4515 inserter
4516 .table_flownode_set_cache
4517 .invalidate(&[common_meta::instruction::CacheIdent::CreateFlow(
4518 common_meta::instruction::CreateFlow {
4519 flow_id: 1,
4520 source_table_ids: vec![1, 2],
4521 partition_to_peer_mapping: vec![(0, peer)],
4522 },
4523 )])
4524 .await
4525 .unwrap();
4526 let mut normal_info = new_test_table_info(1, "normal_table", [1].into_iter());
4527 normal_info.catalog_name = DEFAULT_CATALOG_NAME.to_string();
4528 normal_info.schema_name = DEFAULT_SCHEMA_NAME.to_string();
4529 let mut instant_info = new_test_table_info(2, "instant_table", [1].into_iter());
4530 instant_info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
4531 let table_infos =
4532 HashMap::from_iter([(1, Arc::new(normal_info)), (2, Arc::new(instant_info))]);
4533 let ctx = Arc::new(QueryContext::with(
4534 DEFAULT_CATALOG_NAME,
4535 DEFAULT_SCHEMA_NAME,
4536 ));
4537 let insert = |table_id, num_rows| RegionInsertRequests {
4538 requests: vec![RegionInsertRequest {
4539 region_id: RegionId::new(table_id, 1).as_u64(),
4540 rows: Some(Rows {
4541 schema: vec![],
4542 rows: vec![api::v1::Row::default(); num_rows],
4543 }),
4544 ..Default::default()
4545 }],
4546 };
4547 let make_request = |normal, instant| InstantAndNormalInsertRequests {
4548 normal_requests: if normal == 0 {
4549 RegionInsertRequests::default()
4550 } else {
4551 insert(1, normal)
4552 },
4553 instant_requests: if instant == 0 {
4554 RegionInsertRequests::default()
4555 } else {
4556 insert(2, instant)
4557 },
4558 };
4559
4560 let pending = inserter.mirror_pending_rows.clone();
4561 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4562 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
4563 pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
4564 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4565 for request in [make_request(0, 1), make_request(1, 1)] {
4566 let error = inserter
4567 .do_request(request, &table_infos, &ctx)
4568 .await
4569 .unwrap_err();
4570 assert_eq!(error.status_code(), StatusCode::RateLimited);
4571 assert_eq!(error.retry_hint(), RetryHint::Retryable);
4572 assert!(datanode_rx.try_recv().is_err());
4573 assert!(flownode_rx.try_recv().is_err());
4574 assert_eq!(
4575 pending.load(Ordering::Relaxed),
4576 MAX_MIRROR_PENDING_ROWS as u64
4577 );
4578 }
4579 assert_eq!(
4580 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4581 dropped_before
4582 );
4583
4584 for pending_before in [0, MAX_MIRROR_PENDING_ROWS as u64] {
4585 if pending_before == 0 {
4586 pending.store(0, Ordering::Relaxed);
4587 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
4588 }
4589 let gauge_before_rejection = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4590 for (normal, instant) in [
4591 (0, MAX_MIRROR_PENDING_ROWS + 1),
4592 (MAX_MIRROR_PENDING_ROWS / 2, MAX_MIRROR_PENDING_ROWS / 2 + 1),
4593 ] {
4594 let error = inserter
4595 .do_request(make_request(normal, instant), &table_infos, &ctx)
4596 .await
4597 .unwrap_err();
4598 assert_eq!(error.status_code(), StatusCode::InvalidArguments);
4599 assert_eq!(error.retry_hint(), RetryHint::NonRetryable);
4600 let message = error.to_string();
4601 assert!(message.contains("1000001"), "{message}");
4602 assert!(message.contains("1000000"), "{message}");
4603 assert!(message.contains("reduce the batch size"), "{message}");
4604 assert!(datanode_rx.try_recv().is_err());
4605 assert!(flownode_rx.try_recv().is_err());
4606 assert_eq!(pending.load(Ordering::Relaxed), pending_before);
4607 assert_eq!(
4608 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4609 gauge_before_rejection
4610 );
4611 assert_eq!(
4612 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4613 dropped_before
4614 );
4615 }
4616 if pending_before == 0 {
4617 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
4618 pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
4619 }
4620 }
4621
4622 let result = inserter
4623 .do_request(make_request(1, 0), &table_infos, &ctx)
4624 .await
4625 .unwrap();
4626 assert!(matches!(result.data, OutputData::AffectedRows(1)));
4627 assert_eq!(1, datanode_rx.try_recv().unwrap());
4628 assert!(flownode_rx.try_recv().is_err());
4629 assert_eq!(
4630 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4631 dropped_before + 1
4632 );
4633 pending.store(0, Ordering::Relaxed);
4634 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
4635 assert_eq!(
4636 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4637 gauge_before
4638 );
4639 }
4640
4641 #[tokio::test]
4642 async fn flow_mirror_new_counts_regions_and_peer_clones() {
4643 use common_meta::instruction::{CacheIdent, CreateFlow};
4644
4645 let peer1 = flownode_peer();
4646 let peer2 = Peer {
4647 id: 2,
4648 addr: "127.0.0.1:4002".to_string(),
4649 };
4650 let cache_backend = prepare_mocked_backend().await;
4651 let cache = Arc::new(new_table_flownode_set_cache(
4652 String::new(),
4653 Cache::new(10),
4654 cache_backend,
4655 ));
4656 cache
4657 .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
4658 flow_id: 1,
4659 source_table_ids: vec![1],
4660 partition_to_peer_mapping: vec![(0, peer1), (1, peer2)],
4661 })])
4662 .await
4663 .unwrap();
4664 let request = |region_id| RegionInsertRequest {
4665 region_id,
4666 rows: Some(Rows {
4667 schema: vec![],
4668 rows: vec![api::v1::Row { values: vec![] }; 3],
4669 }),
4670 ..Default::default()
4671 };
4672 let requests = [
4673 request(RegionId::new(1, 0).as_u64()),
4674 request(RegionId::new(1, 1).as_u64()),
4675 ];
4676 let task = FlowMirrorTask::new(&cache, requests.iter()).await.unwrap();
4677 assert_eq!(12, task.pending_rows());
4678 assert_eq!(2, task.requests.len());
4679 }
4680
4681 #[tokio::test]
4682 async fn flow_mirror_dropped_when_pending_exceeds_limit() {
4683 let _guard = lock_mirror_metrics().await;
4684 let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64));
4685 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4686 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4687 let num_rows = 100;
4688
4689 let task = FlowMirrorTask {
4690 requests: mirror_requests(&flownode_peer(), num_rows),
4691 };
4692 task.detach(
4694 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
4695 pending.clone(),
4696 false,
4697 )
4698 .unwrap();
4699
4700 assert_eq!(
4701 num_rows as u64,
4702 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get() - dropped_before
4703 );
4704 assert_eq!(
4706 MAX_MIRROR_PENDING_ROWS as u64,
4707 pending.load(Ordering::Relaxed)
4708 );
4709 assert_eq!(
4712 gauge_before,
4713 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4714 "a dropped batch must not change the pending gauge"
4715 );
4716
4717 let task = FlowMirrorTask {
4718 requests: mirror_requests(&flownode_peer(), num_rows),
4719 };
4720 let error = task
4721 .detach(
4722 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
4723 pending.clone(),
4724 true,
4725 )
4726 .err()
4727 .unwrap();
4728 use common_error::ext::{ErrorExt, RetryHint};
4729 use common_error::status_code::StatusCode;
4730 assert_eq!(error.status_code(), StatusCode::RateLimited);
4731 assert_eq!(error.retry_hint(), RetryHint::Retryable);
4732 assert_eq!(
4733 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4734 dropped_before + num_rows as u64,
4735 "rejected instant rows are not successful best-effort drops"
4736 );
4737 assert_eq!(
4738 MAX_MIRROR_PENDING_ROWS as u64,
4739 pending.load(Ordering::Relaxed)
4740 );
4741 assert_eq!(
4742 gauge_before,
4743 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get()
4744 );
4745 }
4746
4747 #[tokio::test]
4751 async fn flow_mirror_reserves_and_releases_per_peer_shares() {
4752 let _guard = lock_mirror_metrics().await;
4753 let per_peer_rows = 10;
4754 let total_rows = (2 * per_peer_rows) as u64;
4755 let start = MAX_MIRROR_PENDING_ROWS as u64 - total_rows;
4757 let pending = Arc::new(AtomicU64::new(start));
4758 let gate = Arc::new(tokio::sync::Semaphore::new(0));
4759 let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4760 gate: gate.clone(),
4761 affected_rows: per_peer_rows as u64,
4762 }));
4763
4764 let mut requests = mirror_requests(&flownode_peer(), per_peer_rows);
4765 requests.extend(mirror_requests(
4766 &Peer {
4767 id: 2,
4768 addr: "127.0.0.1:4002".to_string(),
4769 },
4770 per_peer_rows,
4771 ));
4772
4773 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4774 FlowMirrorTask { requests }
4775 .detach(node_manager, pending.clone(), false)
4776 .unwrap();
4777
4778 assert_eq!(
4781 MAX_MIRROR_PENDING_ROWS as u64,
4782 pending.load(Ordering::Relaxed)
4783 );
4784
4785 gate.add_permits(2);
4791 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4792 while pending.load(Ordering::Relaxed) != start
4793 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4794 {
4795 assert!(
4796 std::time::Instant::now() < deadline,
4797 "the mirror tasks must release exactly their own shares"
4798 );
4799 std::thread::sleep(std::time::Duration::from_millis(1));
4800 }
4801 }
4802
4803 #[tokio::test]
4804 async fn flow_mirror_fast_peer_releases_only_its_share_before_drop() {
4805 let _guard = lock_mirror_metrics().await;
4806 let peer1 = flownode_peer();
4807 let peer2 = Peer {
4808 id: 2,
4809 addr: "127.0.0.1:4002".to_string(),
4810 };
4811 let gate1 = Arc::new(tokio::sync::Semaphore::new(0));
4812 let gate2 = Arc::new(tokio::sync::Semaphore::new(0));
4813 let gates = Arc::new(std::sync::Mutex::new(HashMap::from_iter([
4814 (peer1.id, gate1.clone()),
4815 (peer2.id, gate2.clone()),
4816 ])));
4817 let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4818 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4819 let node_manager = Arc::new(RecordingNodeManager {
4820 datanode_dispatch: datanode_tx,
4821 flownode_dispatch: flownode_tx,
4822 gates,
4823 fail_flownode: false,
4824 });
4825 let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64 - 20));
4826 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4827 let mut requests = mirror_requests(&peer1, 10);
4828 requests.extend(mirror_requests(&peer2, 10));
4829 let task = FlowMirrorTask { requests };
4830 task.detach(node_manager.clone(), pending.clone(), false)
4831 .unwrap();
4832 let first = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4833 .await
4834 .unwrap()
4835 .unwrap();
4836 let second = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4837 .await
4838 .unwrap()
4839 .unwrap();
4840 assert_ne!(first.id, second.id);
4841 let fast_gate = if first.id == peer1.id {
4842 gate1.clone()
4843 } else {
4844 gate2.clone()
4845 };
4846 let slow_gate = if first.id == peer1.id {
4847 gate2.clone()
4848 } else {
4849 gate1.clone()
4850 };
4851 fast_gate.add_permits(1);
4852 let deadline = std::time::Instant::now() + Duration::from_secs(10);
4853 while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 10 {
4854 assert!(
4855 std::time::Instant::now() < deadline,
4856 "fast peer did not release"
4857 );
4858 std::thread::sleep(Duration::from_millis(1));
4859 }
4860 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4861 FlowMirrorTask {
4862 requests: mirror_requests(&peer1, 11),
4863 }
4864 .detach(node_manager, pending.clone(), false)
4865 .unwrap();
4866 assert_eq!(
4867 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4868 dropped_before + 11
4869 );
4870 assert!(flownode_rx.try_recv().is_err());
4871 slow_gate.add_permits(1);
4872 while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 20
4873 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4874 {
4875 assert!(
4876 std::time::Instant::now() < deadline,
4877 "slow peer did not release"
4878 );
4879 std::thread::sleep(Duration::from_millis(1));
4880 }
4881 }
4882
4883 #[tokio::test]
4884 async fn flow_mirror_failure_releases_budget_and_gauge() {
4885 let _guard = lock_mirror_metrics().await;
4886 let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4887 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4888 let node_manager = Arc::new(RecordingNodeManager {
4889 datanode_dispatch: datanode_tx,
4890 flownode_dispatch: flownode_tx,
4891 gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
4892 fail_flownode: true,
4893 });
4894 let pending = Arc::new(AtomicU64::new(0));
4895 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4896 FlowMirrorTask {
4897 requests: mirror_requests(&flownode_peer(), 5),
4898 }
4899 .detach(node_manager, pending.clone(), false)
4900 .unwrap();
4901 assert_eq!(
4902 tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4903 .await
4904 .unwrap()
4905 .unwrap()
4906 .id,
4907 flownode_peer().id
4908 );
4909 let deadline = std::time::Instant::now() + Duration::from_secs(10);
4910 while pending.load(Ordering::Relaxed) != 0
4911 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4912 {
4913 assert!(
4914 std::time::Instant::now() < deadline,
4915 "failed task leaked reservation"
4916 );
4917 std::thread::sleep(Duration::from_millis(1));
4918 }
4919 }
4920
4921 #[tokio::test]
4922 async fn flow_mirror_pending_counter_consistent() {
4923 let _guard = lock_mirror_metrics().await;
4924 let pending = Arc::new(AtomicU64::new(0));
4925 let gate = Arc::new(tokio::sync::Semaphore::new(0));
4926 let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4927 gate: gate.clone(),
4928 affected_rows: 7,
4929 }));
4930
4931 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4933 FlowMirrorTask {
4934 requests: HashMap::new(),
4935 }
4936 .detach(node_manager.clone(), pending.clone(), false)
4937 .unwrap();
4938 assert_eq!(0, pending.load(Ordering::Relaxed));
4939 assert_eq!(
4940 gauge_before,
4941 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4942 "an empty batch must not touch the pending gauge"
4943 );
4944
4945 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4946 let num_rows = 7;
4947
4948 FlowMirrorTask {
4949 requests: mirror_requests(&flownode_peer(), num_rows),
4950 }
4951 .detach(node_manager, pending.clone(), false)
4952 .unwrap();
4953
4954 assert_eq!(num_rows as u64, pending.load(Ordering::Relaxed));
4957 assert_eq!(
4958 num_rows as i64,
4959 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() - gauge_before,
4960 "the pending gauge must track the pending rows"
4961 );
4962 assert_eq!(
4963 dropped_before,
4964 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4965 "a batch within the limit must not be counted as dropped"
4966 );
4967
4968 gate.add_permits(1);
4971 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4972 while pending.load(Ordering::Relaxed) != 0
4973 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4974 {
4975 assert!(
4976 std::time::Instant::now() < deadline,
4977 "the mirror task did not release its pending reservation"
4978 );
4979 std::thread::sleep(std::time::Duration::from_millis(1));
4980 }
4981 }
4982
4983 #[test]
4984 fn mirror_logs_are_rate_limited() {
4985 let log = Arc::new(MirrorLog::new());
4986 assert_eq!(Some(1), log.claim_report(1_000));
4987 assert_eq!(None, log.claim_report(1_500));
4988 assert_eq!(None, log.claim_report(10_999));
4989 assert_eq!(Some(3), log.claim_report(11_000));
4990 assert_eq!(None, log.claim_report(11_001));
4991 assert_eq!(Some(2), log.claim_report(21_000));
4992
4993 let concurrent = Arc::new(MirrorLog::new());
4994 assert_eq!(Some(1), concurrent.claim_report(50_000));
4995 let barrier = Arc::new(std::sync::Barrier::new(8));
4996 let threads = (0..8)
4997 .map(|_| {
4998 let log = concurrent.clone();
4999 let barrier = barrier.clone();
5000 std::thread::spawn(move || {
5001 barrier.wait();
5002 log.claim_report(60_000)
5003 })
5004 })
5005 .collect::<Vec<_>>();
5006 let winners = threads
5007 .into_iter()
5008 .filter_map(|thread| thread.join().unwrap())
5009 .collect::<Vec<_>>();
5010 assert_eq!(winners.len(), 1);
5011 let residual = concurrent.events.swap(0, Ordering::Relaxed);
5012 assert_eq!(winners[0] + residual, 8);
5013 }
5014}