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, METRICS_EXEMPLARS_TABLE, PARENT_SPAN_ID_COLUMN,
43 SERVICE_NAME_COLUMN, SPAN_ID_COLUMN, TIMESTAMP_END_COLUMN, TRACE_ID_COLUMN, TRACE_TABLE_NAME,
44 TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table,
45 trace_operations_table_name, trace_services_table_name,
46};
47use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME;
48use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME;
49use common_grpc_expr::util::ColumnExpr;
50use common_meta::cache::TableFlownodeSetCacheRef;
51use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME;
52use common_meta::node_manager::NodeManagerRef;
53use common_meta::peer::Peer;
54use common_meta::rpc::ddl::TriggerReason;
55use common_query::Output;
56use common_query::native_histogram::{is_native_histogram_value_type, native_histogram_value_type};
57use common_query::prelude::{
58 EXEMPLAR_LABELS_COLUMN, EXEMPLAR_METRIC_COLUMN, greptime_timestamp, 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: format!(
2212 "{METRICS_EXEMPLARS_TABLE} requires Mito, append_mode=true and primary_key_encoding=sparse"
2213 )
2214 }
2215 );
2216 let columns = meta.schema.column_schemas();
2217 let lookup = |name: &str| {
2218 columns
2219 .iter()
2220 .enumerate()
2221 .find(|(_, c)| c.name == name)
2222 .with_context(|| InvalidInsertRequestSnafu {
2223 reason: format!("missing exemplar column {name}"),
2224 })
2225 };
2226 let (_, timestamp) = lookup(greptime_timestamp())?;
2227 for (position, name, id, data_type) in [
2228 (
2229 0,
2230 DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
2231 ReservedColumnId::table_id(),
2232 datatypes::prelude::ConcreteDataType::uint32_datatype(),
2233 ),
2234 (
2235 1,
2236 DATA_SCHEMA_TSID_COLUMN_NAME,
2237 ReservedColumnId::tsid(),
2238 datatypes::prelude::ConcreteDataType::uint64_datatype(),
2239 ),
2240 ] {
2241 let (index, column) = lookup(name)?;
2242 ensure!(
2243 meta.primary_key_indices.get(position) == Some(&index)
2244 && meta.column_ids.get(index) == Some(&id)
2245 && column.data_type == data_type
2246 && !column.is_nullable()
2247 && column.default_constraint().is_none(),
2248 InvalidInsertRequestSnafu {
2249 reason: format!("invalid exemplar source identity column {name}"),
2250 }
2251 );
2252 }
2253 ensure!(
2254 timestamp.is_time_index() && timestamp.data_type.is_timestamp(),
2255 InvalidInsertRequestSnafu {
2256 reason: "invalid exemplar time index"
2257 }
2258 );
2259 let (metric_index, metric) = lookup(EXEMPLAR_METRIC_COLUMN)?;
2260 ensure!(
2261 metric.data_type.is_string() && meta.primary_key_indices.contains(&metric_index),
2262 InvalidInsertRequestSnafu {
2263 reason: "exemplar metric name must be a string tag"
2264 }
2265 );
2266 let (value_index, value) = lookup(greptime_value())?;
2267 ensure!(
2268 value.data_type == datatypes::prelude::ConcreteDataType::float64_datatype()
2269 && !meta.primary_key_indices.contains(&value_index),
2270 InvalidInsertRequestSnafu {
2271 reason: "exemplar value must be a DOUBLE field"
2272 }
2273 );
2274 let (labels_index, labels) = lookup(EXEMPLAR_LABELS_COLUMN)?;
2275 ensure!(
2276 labels.data_type.is_json2() && !meta.primary_key_indices.contains(&labels_index),
2277 InvalidInsertRequestSnafu {
2278 reason: "exemplar_labels must be a JSON2 field"
2279 }
2280 );
2281 if let Some(extension) = labels
2282 .extension_type::<datatypes::extension::json::Json2ExtensionType>()
2283 .map_err(|error| {
2284 InvalidInsertRequestSnafu {
2285 reason: error.to_string(),
2286 }
2287 .build()
2288 })?
2289 {
2290 ensure!(
2291 extension
2292 .metadata()
2293 .json_settings()
2294 .type_hints()
2295 .iter()
2296 .all(|hint| hint.path.len() == 1 && hint.data_type.is_string()),
2297 InvalidInsertRequestSnafu {
2298 reason: "exemplar JSON2 hints must preserve string labels"
2299 }
2300 );
2301 }
2302 for (index, column) in columns.iter().enumerate() {
2303 if [
2304 greptime_timestamp(),
2305 greptime_value(),
2306 EXEMPLAR_METRIC_COLUMN,
2307 EXEMPLAR_LABELS_COLUMN,
2308 DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
2309 DATA_SCHEMA_TSID_COLUMN_NAME,
2310 ]
2311 .contains(&column.name.as_str())
2312 {
2313 continue;
2314 }
2315 ensure!(
2316 meta.primary_key_indices.contains(&index)
2317 && column.data_type.is_string()
2318 && column.is_nullable(),
2319 InvalidInsertRequestSnafu {
2320 reason: "inherited exemplar labels must be nullable string tags"
2321 }
2322 );
2323 ensure!(
2325 matches!(
2326 column.default_constraint(),
2327 None | Some(ColumnDefaultConstraint::Value(
2328 datatypes::value::Value::Null
2329 ))
2330 ),
2331 InvalidInsertRequestSnafu {
2332 reason: "inherited exemplar labels must not have non-NULL defaults"
2333 }
2334 );
2335 }
2336 Ok(())
2337}
2338
2339pub fn validate_trace_table_model(table_info: &TableInfo, ctx: &QueryContextRef) -> Result<()> {
2342 let Some(expected @ (TABLE_DATA_MODEL_TRACE_V1 | TABLE_DATA_MODEL_TRACE_V2)) =
2343 ctx.extension(SEMANTIC_PIPELINE)
2344 else {
2345 return Ok(());
2346 };
2347 if let Some(actual) = table_info.meta.options.data_model() {
2348 ensure!(
2349 actual == expected,
2350 InvalidInsertRequestSnafu {
2351 reason: format!(
2352 "Trace table `{}` uses {actual}, but the request uses {expected}",
2353 table_info.name,
2354 ),
2355 }
2356 );
2357 }
2358 Ok(())
2359}
2360
2361pub fn fill_table_options_for_create(
2363 table_options: &mut std::collections::HashMap<String, String>,
2364 create_type: &AutoCreateTableType,
2365 ctx: &QueryContextRef,
2366) {
2367 for key in VALID_TABLE_OPTION_KEYS {
2368 if let Some(value) = ctx.extension(key) {
2369 table_options.insert(key.to_string(), value.to_string());
2370 }
2371 }
2372
2373 if !matches!(create_type, AutoCreateTableType::Logical(_)) {
2374 let mut has_twcs_option = false;
2375 for key in COMPACTION_OPTION_KEYS {
2376 if let Some(value) = ctx.extension(key) {
2377 has_twcs_option |= key.starts_with(COMPACTION_TWCS_PREFIX);
2378 table_options.insert(key.to_string(), value.to_string());
2379 }
2380 }
2381 if has_twcs_option {
2382 table_options
2383 .entry(COMPACTION_TYPE.to_string())
2384 .or_insert_with(|| COMPACTION_TYPE_TWCS.to_string());
2385 }
2386 }
2387
2388 for (key, value) in ctx.extensions() {
2390 if is_semantic_option_key(&key) && validate_semantic_option(&key, &value) {
2391 table_options.insert(key, value);
2392 }
2393 }
2394
2395 match create_type {
2396 AutoCreateTableType::Logical(physical_table) => {
2397 table_options.insert(
2398 LOGICAL_TABLE_METADATA_KEY.to_string(),
2399 physical_table.clone(),
2400 );
2401 }
2402 AutoCreateTableType::Exemplar => {
2403 table_options.retain(|key, _| {
2404 !is_semantic_option_key(key)
2405 && key != SEMANTIC_PER_TABLE_INDEX_KEY
2406 && key != LOGICAL_TABLE_METADATA_KEY
2407 && key != PHYSICAL_TABLE_METADATA_KEY
2408 && key != MERGE_MODE_KEY
2409 && key != "otlp_metric_compat"
2410 });
2411 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2412 table_options.insert(
2413 store_api::metric_engine_consts::PRIMARY_KEY_ENCODING.to_string(),
2414 "sparse".to_string(),
2415 );
2416 }
2417 AutoCreateTableType::Physical => {
2418 if let Some(append_mode) = ctx.extension(APPEND_MODE_KEY) {
2419 table_options.insert(APPEND_MODE_KEY.to_string(), append_mode.to_string());
2420 }
2421 if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
2422 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
2423 }
2424 }
2425 AutoCreateTableType::Log => {
2428 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2429 }
2430 AutoCreateTableType::LastNonNull => {
2431 if ctx
2432 .extension(APPEND_MODE_KEY)
2433 .is_some_and(|value| value.eq_ignore_ascii_case("true"))
2434 {
2435 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2436 table_options.insert(MERGE_MODE_KEY.to_string(), "last_row".to_string());
2437 } else if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
2438 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
2439 } else {
2440 table_options.insert(MERGE_MODE_KEY.to_string(), "last_non_null".to_string());
2441 }
2442 }
2443 AutoCreateTableType::Trace { .. } => {
2444 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
2445 }
2446 }
2447}
2448
2449pub type PerTableSemanticIndex = BTreeMap<String, BTreeMap<String, BTreeMap<String, String>>>;
2455
2456pub fn parse_per_table_semantic_index(ctx: &QueryContextRef) -> Option<PerTableSemanticIndex> {
2461 let raw = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY)?;
2462 match serde_json::from_str(raw) {
2463 Ok(index) => Some(index),
2464 Err(_) => {
2465 warn!("failed to parse semantic per-table index, skipping per-table options");
2466 None
2467 }
2468 }
2469}
2470
2471pub fn apply_per_table_semantic_options(
2479 table_options: &mut std::collections::HashMap<String, String>,
2480 index: Option<&PerTableSemanticIndex>,
2481 schema: &str,
2482 table_name: &str,
2483) {
2484 let Some(entry) = index
2485 .and_then(|index| index.get(schema))
2486 .and_then(|tables| tables.get(table_name))
2487 else {
2488 return;
2489 };
2490 for (key, value) in entry {
2491 if is_semantic_option_key(key) && validate_semantic_option(key, value) {
2492 table_options.insert(key.clone(), value.clone());
2493 }
2494 }
2495}
2496
2497pub fn build_create_table_expr(
2498 table: &TableReference,
2499 request_schema: &[ColumnSchema],
2500 engine: &str,
2501) -> Result<CreateTableExpr> {
2502 expr_helper::create_table_expr_by_column_schemas(table, request_schema, engine, None)
2503}
2504
2505pub const SPLUNK_PK_METADATA_ORDER_KEY: &str = "splunk_pk_metadata_order";
2509
2510fn reorder_splunk_primary_keys(primary_keys: &mut [String]) {
2513 const LEAD: [&str; 3] = ["host", "source", "sourcetype"];
2514 primary_keys.sort_by_key(|name| {
2517 LEAD.iter()
2518 .position(|&lead| lead == name.as_str())
2519 .unwrap_or(LEAD.len())
2520 });
2521}
2522
2523struct CreateAlterTableResult {
2525 instant_table_ids: HashSet<TableId>,
2527 table_infos: HashMap<TableId, Arc<TableInfo>>,
2529}
2530
2531const MAX_MIRROR_PENDING_ROWS: usize = 1_000_000;
2539
2540const MIRROR_LOG_INTERVAL: Duration = Duration::from_secs(10);
2542const MIRROR_LOG_NEVER_REPORTED: u64 = u64::MAX;
2543
2544struct MirrorLog {
2545 last_log_millis: AtomicU64,
2546 events: AtomicU64,
2547}
2548
2549impl MirrorLog {
2550 const fn new() -> Self {
2551 Self {
2552 last_log_millis: AtomicU64::new(MIRROR_LOG_NEVER_REPORTED),
2553 events: AtomicU64::new(0),
2554 }
2555 }
2556
2557 fn claim_report(&self, now_millis: u64) -> Option<u64> {
2560 self.events.fetch_add(1, Ordering::Relaxed);
2561 let last = self.last_log_millis.load(Ordering::Relaxed);
2562 if last != MIRROR_LOG_NEVER_REPORTED
2563 && now_millis.saturating_sub(last) < MIRROR_LOG_INTERVAL.as_millis() as u64
2564 {
2565 return None;
2566 }
2567 if self
2568 .last_log_millis
2569 .compare_exchange(last, now_millis, Ordering::Relaxed, Ordering::Relaxed)
2570 .is_err()
2571 {
2572 return None;
2573 }
2574 Some(self.events.swap(0, Ordering::Relaxed))
2575 }
2576}
2577
2578static MIRROR_DROP_LOG: MirrorLog = MirrorLog::new();
2579static MIRROR_FAILURE_LOG: MirrorLog = MirrorLog::new();
2580
2581fn mirror_drop_log_millis() -> u64 {
2584 static START: LazyLock<Instant> = LazyLock::new(Instant::now);
2585 START.elapsed().as_millis() as u64
2586}
2587
2588struct FlowMirrorTask {
2589 requests: HashMap<Peer, RegionInsertRequests>,
2590}
2591
2592impl FlowMirrorTask {
2593 async fn new(
2594 cache: &TableFlownodeSetCacheRef,
2595 requests: impl Iterator<Item = &RegionInsertRequest>,
2596 ) -> Result<Self> {
2597 let mut src_table_reqs: HashMap<TableId, Option<(Vec<Peer>, RegionInsertRequests)>> =
2598 HashMap::new();
2599
2600 for req in requests {
2601 let table_id = RegionId::from_u64(req.region_id).table_id();
2602 match src_table_reqs.get_mut(&table_id) {
2603 Some(Some((_peers, reqs))) => reqs.requests.push(req.clone()),
2604 Some(None) => continue,
2606 _ => {
2607 let peers = cache
2609 .get(table_id)
2610 .await
2611 .context(RequestInsertsSnafu)?
2612 .unwrap_or_default()
2613 .values()
2614 .cloned()
2615 .collect::<HashSet<_>>()
2616 .into_iter()
2617 .collect::<Vec<_>>();
2618
2619 if !peers.is_empty() {
2620 let mut reqs = RegionInsertRequests::default();
2621 reqs.requests.push(req.clone());
2622 src_table_reqs.insert(table_id, Some((peers, reqs)));
2623 } else {
2624 src_table_reqs.insert(table_id, None);
2626 }
2627 }
2628 }
2629 }
2630
2631 let mut inserts: HashMap<Peer, RegionInsertRequests> = HashMap::new();
2632
2633 for (_table_id, (peers, reqs)) in src_table_reqs
2634 .into_iter()
2635 .filter_map(|(k, v)| v.map(|v| (k, v)))
2636 {
2637 if peers.len() == 1 {
2638 inserts
2640 .entry(peers[0].clone())
2641 .or_default()
2642 .requests
2643 .extend(reqs.requests);
2644 continue;
2645 } else {
2646 for flownode in peers {
2648 inserts
2649 .entry(flownode.clone())
2650 .or_default()
2651 .requests
2652 .extend(reqs.requests.clone());
2653 }
2654 }
2655 }
2656
2657 Ok(Self { requests: inserts })
2658 }
2659
2660 fn pending_rows(&self) -> u64 {
2669 self.requests.values().map(region_inserts_rows).sum()
2670 }
2671
2672 fn detach(
2673 self,
2674 node_manager: NodeManagerRef,
2675 mirror_pending_rows: Arc<AtomicU64>,
2676 has_instant_rows: bool,
2677 ) -> Result<()> {
2678 let num_rows = self.pending_rows();
2681 if num_rows == 0 {
2682 return Ok(());
2683 }
2684 if has_instant_rows && num_rows > MAX_MIRROR_PENDING_ROWS as u64 {
2685 return InvalidInsertRequestSnafu {
2686 reason: format!(
2687 "flow mirror batch has {num_rows} peer-fan-out rows, exceeding the limit of {}; reduce the batch size",
2688 MAX_MIRROR_PENDING_ROWS
2689 ),
2690 }
2691 .fail();
2692 }
2693
2694 let pending = reserve_mirror_pending_rows(&mirror_pending_rows, num_rows);
2695 if pending > MAX_MIRROR_PENDING_ROWS as u64 {
2696 release_mirror_pending_rows(&mirror_pending_rows, num_rows);
2697 if has_instant_rows {
2698 return Err(meter_core::collect::WriteRejected::new(
2699 "flow mirror pending rows limit exceeded for instant-TTL insert",
2700 ))
2701 .context(WriteRejectedSnafu);
2702 }
2703 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.inc_by(num_rows);
2704 if let Some(events) = MIRROR_DROP_LOG.claim_report(mirror_drop_log_millis()) {
2705 warn!(
2706 "Flow mirror write dropped: {} pending rows exceeds limit {}, {} drop events since last report",
2707 pending, MAX_MIRROR_PENDING_ROWS, events
2708 );
2709 }
2710 return Ok(());
2711 }
2712
2713 for (peer, inserts) in self.requests {
2714 let peer_rows = region_inserts_rows(&inserts);
2718 let node_manager = node_manager.clone();
2719 let mirror_pending_rows = mirror_pending_rows.clone();
2720 common_runtime::spawn_global(async move {
2721 let result = node_manager
2722 .flownode(&peer)
2723 .await
2724 .handle_inserts(inserts)
2725 .await
2726 .context(RequestInsertsSnafu);
2727
2728 match result {
2729 Ok(resp) => {
2730 let affected_rows = resp.affected_rows;
2731 crate::metrics::DIST_MIRROR_ROW_COUNT.inc_by(affected_rows);
2732 }
2733 Err(err) => {
2734 if let Some(events) =
2735 MIRROR_FAILURE_LOG.claim_report(mirror_drop_log_millis())
2736 {
2737 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);
2738 }
2739 }
2740 }
2741 release_mirror_pending_rows(&mirror_pending_rows, peer_rows);
2744 });
2745 }
2746
2747 Ok(())
2748 }
2749}
2750
2751fn region_inserts_rows(inserts: &RegionInsertRequests) -> u64 {
2754 inserts
2755 .requests
2756 .iter()
2757 .filter_map(|req| req.rows.as_ref())
2758 .map(|rows| rows.rows.len() as u64)
2759 .sum()
2760}
2761
2762fn reserve_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) -> u64 {
2767 let pending = mirror_pending_rows.fetch_add(rows, Ordering::Relaxed) + rows;
2768 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(rows as i64);
2771 pending
2772}
2773
2774fn release_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) {
2780 let prev = mirror_pending_rows
2781 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |pending| {
2782 Some(pending.saturating_sub(rows))
2783 })
2784 .unwrap_or_else(|prev| prev);
2785 let released = rows.min(prev);
2789 if released > 0 {
2790 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(released as i64);
2791 }
2792}
2793
2794#[cfg(test)]
2795mod tests {
2796 use std::future::Future;
2797 use std::sync::Arc;
2798 use std::sync::atomic::{AtomicU64, Ordering};
2799 use std::task::Poll;
2800
2801 use api::helper::ColumnDataTypeWrapper;
2802 use api::region::RegionResponse;
2803 use api::v1::flow::FlowResponse;
2804 use api::v1::helper::{field_column_schema, time_index_column_schema};
2805 use api::v1::{Row, RowInsertRequest, Rows, Value};
2806 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
2807 use common_error::ext::{ErrorExt, RetryHint};
2808 use common_error::mock::MockError;
2809 use common_error::status_code::StatusCode;
2810 use common_meta::cache::new_table_flownode_set_cache;
2811 use common_meta::cache_invalidator::DummyCacheInvalidator;
2812 use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler;
2813 use common_meta::node_manager::{
2814 Datanode, DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef,
2815 };
2816 use common_meta::test_util::{MockDatanodeManager, MockFlownodeHandler, MockFlownodeManager};
2817 use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD;
2818 use common_query::prelude::{greptime_native_histogram, set_default_prefix};
2819 use datatypes::data_type::ConcreteDataType;
2820 use datatypes::schema::ColumnSchema;
2821 use moka::future::Cache;
2822 use session::context::QueryContext;
2823 use table::TableRef;
2824 use table::dist_table::DummyDataSource;
2825 use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType};
2826 use tokio::sync::{Semaphore, mpsc};
2827
2828 use crate::insert::*;
2829 use crate::test_util::{
2830 create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
2831 };
2832 use crate::write_dispatcher::{TrackedWrite, WriteRetryOptions};
2833
2834 fn make_table_ref_with_schema(
2835 ts_name: &str,
2836 field_name: &str,
2837 field_type: ConcreteDataType,
2838 ) -> TableRef {
2839 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2840 ColumnSchema::new(
2841 ts_name,
2842 ConcreteDataType::timestamp_millisecond_datatype(),
2843 false,
2844 )
2845 .with_time_index(true),
2846 ColumnSchema::new(field_name, field_type, true),
2847 ])
2848 .unwrap()
2849 .build()
2850 .unwrap();
2851 let meta = TableMetaBuilder::empty()
2852 .schema(Arc::new(schema))
2853 .primary_key_indices(vec![])
2854 .value_indices(vec![1])
2855 .engine("mito")
2856 .next_column_id(0)
2857 .options(Default::default())
2858 .created_on(Default::default())
2859 .build()
2860 .unwrap();
2861 let info = Arc::new(
2862 TableInfoBuilder::default()
2863 .table_id(1)
2864 .table_version(0)
2865 .name("test_table")
2866 .schema_name(DEFAULT_SCHEMA_NAME)
2867 .catalog_name(DEFAULT_CATALOG_NAME)
2868 .desc(None)
2869 .table_type(TableType::Base)
2870 .meta(meta)
2871 .build()
2872 .unwrap(),
2873 );
2874 Arc::new(table::Table::new(
2875 info,
2876 table::metadata::FilterPushDownType::Unsupported,
2877 Arc::new(DummyDataSource),
2878 ))
2879 }
2880
2881 fn make_metric_physical_table_ref_with_time_unit(unit: TimeUnit) -> TableRef {
2882 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2883 ColumnSchema::new(
2884 greptime_timestamp(),
2885 ConcreteDataType::timestamp_datatype(unit),
2886 false,
2887 )
2888 .with_time_index(true),
2889 ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
2890 ])
2891 .unwrap()
2892 .build()
2893 .unwrap();
2894 let meta = TableMetaBuilder::empty()
2895 .schema(Arc::new(schema))
2896 .primary_key_indices(vec![])
2897 .value_indices(vec![1])
2898 .engine("metric")
2899 .next_column_id(0)
2900 .options(Default::default())
2901 .created_on(Default::default())
2902 .build()
2903 .unwrap();
2904 let info = Arc::new(
2905 TableInfoBuilder::default()
2906 .table_id(1)
2907 .table_version(0)
2908 .name("greptime_physical_table")
2909 .schema_name(DEFAULT_SCHEMA_NAME)
2910 .catalog_name(DEFAULT_CATALOG_NAME)
2911 .desc(None)
2912 .table_type(TableType::Base)
2913 .meta(meta)
2914 .build()
2915 .unwrap(),
2916 );
2917 Arc::new(table::Table::new(
2918 info,
2919 table::metadata::FilterPushDownType::Unsupported,
2920 Arc::new(DummyDataSource),
2921 ))
2922 }
2923
2924 fn ms_row_insert_request(timestamp_ms: i64) -> RowInsertRequest {
2925 ms_row_insert_request_named("my_metric", timestamp_ms)
2926 }
2927
2928 fn ms_row_insert_request_named(table: &str, timestamp_ms: i64) -> RowInsertRequest {
2929 RowInsertRequest {
2930 table_name: table.to_string(),
2931 rows: Some(Rows {
2932 schema: vec![
2933 time_index_column_schema(
2934 greptime_timestamp(),
2935 ColumnDataType::TimestampMillisecond,
2936 ),
2937 field_column_schema(greptime_value(), ColumnDataType::Float64),
2938 ],
2939 rows: vec![Row {
2940 values: vec![
2941 Value {
2942 value_data: Some(ValueData::TimestampMillisecondValue(timestamp_ms)),
2943 },
2944 Value {
2945 value_data: Some(ValueData::F64Value(1.0)),
2946 },
2947 ],
2948 }],
2949 }),
2950 }
2951 }
2952
2953 fn convert_row_insert_requests_time_unit(
2955 requests: &mut RowInsertRequests,
2956 target_unit: TimeUnit,
2957 ) -> Result<()> {
2958 for request in &mut requests.inserts {
2959 let Some(rows) = request.rows.as_mut() else {
2960 continue;
2961 };
2962 convert_rows_time_unit(rows, target_unit)?;
2963 }
2964 Ok(())
2965 }
2966
2967 #[test]
2968 fn test_convert_row_insert_requests_time_unit_noop_when_matching() {
2969 let mut requests = RowInsertRequests {
2970 inserts: vec![ms_row_insert_request(123)],
2971 };
2972 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Millisecond).unwrap();
2973 let rows = requests.inserts[0].rows.as_ref().unwrap();
2974 assert_eq!(
2975 rows.schema[0].datatype,
2976 ColumnDataType::TimestampMillisecond as i32
2977 );
2978 assert!(matches!(
2979 rows.rows[0].values[0].value_data,
2980 Some(ValueData::TimestampMillisecondValue(123))
2981 ));
2982 }
2983
2984 #[test]
2985 fn test_convert_row_insert_requests_time_unit_widens_losslessly() {
2986 let mut requests = RowInsertRequests {
2987 inserts: vec![ms_row_insert_request(123)],
2988 };
2989 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Microsecond).unwrap();
2990 let rows = requests.inserts[0].rows.as_ref().unwrap();
2991 assert_eq!(
2992 rows.schema[0].datatype,
2993 ColumnDataType::TimestampMicrosecond as i32
2994 );
2995 assert!(matches!(
2996 rows.rows[0].values[0].value_data,
2997 Some(ValueData::TimestampMicrosecondValue(123_000))
2998 ));
2999 }
3000
3001 #[test]
3002 fn test_convert_row_insert_requests_time_unit_truncates_on_narrowing() {
3003 let requests = |value_ns: i64| RowInsertRequests {
3006 inserts: vec![RowInsertRequest {
3007 table_name: "my_metric".to_string(),
3008 rows: Some(Rows {
3009 schema: vec![
3010 time_index_column_schema(
3011 greptime_timestamp(),
3012 ColumnDataType::TimestampNanosecond,
3013 ),
3014 field_column_schema(greptime_value(), ColumnDataType::Float64),
3015 ],
3016 rows: vec![Row {
3017 values: vec![
3018 Value {
3019 value_data: Some(ValueData::TimestampNanosecondValue(value_ns)),
3020 },
3021 Value {
3022 value_data: Some(ValueData::F64Value(1.0)),
3023 },
3024 ],
3025 }],
3026 }),
3027 }],
3028 };
3029
3030 let mut reqs = requests(123_456_789);
3031 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
3032 let rows = reqs.inserts[0].rows.as_ref().unwrap();
3033 assert!(matches!(
3034 rows.rows[0].values[0].value_data,
3035 Some(ValueData::TimestampMicrosecondValue(123_456))
3036 ));
3037
3038 let mut reqs = requests(123_456_789);
3039 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Millisecond).unwrap();
3040 let rows = reqs.inserts[0].rows.as_ref().unwrap();
3041 assert!(matches!(
3042 rows.rows[0].values[0].value_data,
3043 Some(ValueData::TimestampMillisecondValue(123))
3044 ));
3045
3046 let mut reqs = requests(-123_456_789);
3047 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
3048 let rows = reqs.inserts[0].rows.as_ref().unwrap();
3049 assert!(matches!(
3050 rows.rows[0].values[0].value_data,
3051 Some(ValueData::TimestampMicrosecondValue(-123_457))
3052 ));
3053 }
3054
3055 #[test]
3056 fn test_convert_row_insert_requests_time_unit_overflow() {
3057 let mut requests = RowInsertRequests {
3058 inserts: vec![ms_row_insert_request(i64::MAX)],
3059 };
3060 let err =
3061 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Nanosecond).unwrap_err();
3062 assert!(err.to_string().contains("overflows"), "{err}");
3063 }
3064
3065 #[test]
3066 fn test_table_time_index_unit() {
3067 assert_eq!(
3068 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
3069 TimeUnit::Microsecond
3070 )),
3071 Some(TimeUnit::Microsecond)
3072 );
3073 assert_eq!(
3074 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
3075 TimeUnit::Millisecond
3076 )),
3077 Some(TimeUnit::Millisecond)
3078 );
3079 }
3080
3081 #[tokio::test]
3082 async fn test_accommodate_existing_schema_and_reject_kind_changes() {
3083 let ts_name = "my_ts";
3084 let field_name = "my_field";
3085 let table =
3086 make_table_ref_with_schema(ts_name, field_name, ConcreteDataType::float64_datatype());
3087
3088 let mut req = RowInsertRequest {
3090 table_name: "test_table".to_string(),
3091 rows: Some(Rows {
3092 schema: vec![
3093 time_index_column_schema("ts_wrong", ColumnDataType::TimestampMillisecond),
3094 field_column_schema("field_wrong", ColumnDataType::Float64),
3095 ],
3096 rows: vec![api::v1::Row {
3097 values: vec![Value::default(), Value::default()],
3098 }],
3099 }),
3100 };
3101 let ctx = Arc::new(QueryContext::with(
3102 DEFAULT_CATALOG_NAME,
3103 DEFAULT_SCHEMA_NAME,
3104 ));
3105
3106 let kv_backend = prepare_mocked_backend().await;
3107 let inserter = Inserter::new(
3108 catalog::memory::MemoryCatalogManager::new(),
3109 create_partition_rule_manager(kv_backend.clone()).await,
3110 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3111 Arc::new(new_table_flownode_set_cache(
3112 String::new(),
3113 Cache::new(100),
3114 kv_backend.clone(),
3115 )),
3116 true,
3117 );
3118 assert!(
3120 inserter
3121 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, false, false, false)
3122 .unwrap()
3123 .is_none()
3124 );
3125 let alter_expr = inserter
3126 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true, true)
3127 .unwrap();
3128 assert!(alter_expr.is_none());
3129
3130 let req_schema = req.rows.as_ref().unwrap().schema.clone();
3132 assert_eq!(req_schema[0].column_name, ts_name);
3133 assert_eq!(req_schema[1].column_name, field_name);
3134
3135 let (datatype, datatype_extension) =
3136 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
3137 .unwrap()
3138 .into_parts();
3139 let mut histogram_req = RowInsertRequest {
3140 table_name: "test_table".to_string(),
3141 rows: Some(Rows {
3142 schema: vec![
3143 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
3144 api::v1::ColumnSchema {
3145 column_name: greptime_native_histogram().to_string(),
3146 datatype: datatype as i32,
3147 semantic_type: SemanticType::Field as i32,
3148 datatype_extension,
3149 options: None,
3150 },
3151 ],
3152 rows: vec![],
3153 }),
3154 };
3155 let error = inserter
3156 .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, false, true, true)
3157 .unwrap_err();
3158 assert!(
3159 error
3160 .to_string()
3161 .contains("cannot mix native histogram and float sample fields")
3162 );
3163
3164 let histogram_table = make_table_ref_with_schema(
3165 "ts",
3166 greptime_native_histogram(),
3167 native_histogram_value_type().clone(),
3168 );
3169 let mut sample_req = RowInsertRequest {
3170 table_name: "test_table".to_string(),
3171 rows: Some(Rows {
3172 schema: vec![
3173 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
3174 field_column_schema(greptime_value(), ColumnDataType::Float64),
3175 ],
3176 rows: vec![],
3177 }),
3178 };
3179 let error = inserter
3180 .get_alter_table_expr_on_demand(
3181 &mut sample_req,
3182 &histogram_table,
3183 &ctx,
3184 false,
3185 true,
3186 true,
3187 )
3188 .unwrap_err();
3189 assert!(
3190 error
3191 .to_string()
3192 .contains("cannot mix native histogram and float sample fields")
3193 );
3194 }
3195
3196 #[test]
3197 fn test_native_histogram_detection_survives_prefix_change() {
3198 set_default_prefix(Some("custom")).unwrap();
3199 let table = make_table_ref_with_schema(
3200 "custom_timestamp",
3201 NATIVE_HISTOGRAM_FIELD,
3202 native_histogram_value_type().clone(),
3203 );
3204 let (datatype, datatype_extension) =
3205 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
3206 .unwrap()
3207 .into_parts();
3208 let request_schema = [api::v1::ColumnSchema {
3209 column_name: greptime_native_histogram().to_string(),
3210 datatype: datatype as i32,
3211 semantic_type: SemanticType::Field as i32,
3212 datatype_extension,
3213 options: None,
3214 }];
3215
3216 assert!(request_is_native_histogram(&request_schema));
3217 assert!(table_is_native_histogram(&table));
3218 }
3219
3220 #[derive(Clone)]
3221 struct GatedDatanodeHandler {
3222 started: mpsc::UnboundedSender<u32>,
3223 completed: mpsc::UnboundedSender<u32>,
3224 release: Arc<Semaphore>,
3225 }
3226
3227 #[async_trait::async_trait]
3228 impl Datanode for GatedDatanodeHandler {
3229 async fn handle_request(
3230 &self,
3231 request: Arc<api::v1::region::RegionRequest>,
3232 ) -> common_meta::error::Result<RegionResponse> {
3233 let Some(api::v1::region::region_request::Body::Inserts(requests)) = &request.body
3234 else {
3235 panic!("expected inserts");
3236 };
3237 let region = RegionId::from(requests.requests[0].region_id).region_number();
3238 self.started.send(region).unwrap();
3239 self.release.acquire().await.unwrap().forget();
3240 self.completed.send(region).unwrap();
3241 if region == 2 {
3242 Err(common_meta::error::Error::retry_later(MockError::new(
3243 StatusCode::RegionBusy,
3244 )))
3245 } else {
3246 Ok(RegionResponse::new(1))
3247 }
3248 }
3249
3250 async fn handle_query(
3251 &self,
3252 _request: common_query::request::QueryRequest,
3253 ) -> common_meta::error::Result<common_recordbatch::SendableRecordBatchStream> {
3254 unreachable!()
3255 }
3256 }
3257
3258 #[async_trait::async_trait]
3259 impl DatanodeManager for GatedDatanodeHandler {
3260 async fn datanode(&self, _peer: &Peer) -> DatanodeRef {
3261 Arc::new(self.clone())
3262 }
3263 }
3264
3265 #[async_trait::async_trait]
3266 impl FlownodeManager for GatedDatanodeHandler {
3267 async fn flownode(&self, _peer: &Peer) -> FlownodeRef {
3268 unreachable!()
3269 }
3270 }
3271
3272 #[tokio::test]
3273 async fn test_insert_cancellation_drains_inflight_writes() {
3274 for confirmed in [false, true] {
3275 for max_retries in [0, 2] {
3276 let mut inserter = batcher_test_inserter().await;
3277 let (started, mut starts) = mpsc::unbounded_channel();
3278 let (completed, mut completions) = mpsc::unbounded_channel();
3279 let release = Arc::new(Semaphore::new(0));
3280 let node_manager = Arc::new(GatedDatanodeHandler {
3281 started,
3282 completed,
3283 release: release.clone(),
3284 });
3285 inserter.node_manager = node_manager.clone();
3286 let dispatcher = Arc::new(
3287 WriteDispatcher::new(node_manager, inserter.partition_manager.clone())
3288 .with_retry(
3289 WriteRetryOptions {
3290 max_retries,
3291 initial_backoff: Duration::ZERO,
3292 max_backoff: Duration::ZERO,
3293 ..Default::default()
3294 },
3295 inserter.catalog_manager.clone(),
3296 Arc::new(DummyCacheInvalidator),
3297 ),
3298 );
3299 let inserter = inserter.with_write_dispatcher(dispatcher);
3300 let table_infos = HashMap::from_iter([(
3301 1,
3302 Arc::new(new_test_table_info(1, "table_1", [1].into_iter())),
3303 )]);
3304 let counter = Arc::new(session::write_progress::WriteProgress::default());
3305 let mut ctx = QueryContext::arc().fork();
3306 ctx.set_write_progress(counter.clone());
3307 let ctx = Arc::new(ctx);
3308 let ctx = if confirmed {
3309 Arc::new(ctx.with_confirmed_writes())
3310 } else {
3311 ctx
3312 };
3313 let requests = InstantAndNormalInsertRequests {
3314 normal_requests: RegionInsertRequests {
3315 requests: [1, 2]
3316 .map(|region| RegionInsertRequest {
3317 region_id: RegionId::new(1, region).as_u64(),
3318 rows: Some(Rows {
3319 schema: vec![],
3320 rows: vec![Row { values: vec![] }],
3321 }),
3322 ..Default::default()
3323 })
3324 .to_vec(),
3325 },
3326 instant_requests: RegionInsertRequests::default(),
3327 };
3328 let caller =
3329 tokio::spawn(
3330 async move { inserter.do_request(requests, &table_infos, &ctx).await },
3331 );
3332 for _ in 0..2 {
3333 tokio::time::timeout(Duration::from_secs(5), starts.recv())
3334 .await
3335 .unwrap()
3336 .unwrap();
3337 }
3338 caller.abort();
3339 assert!(caller.await.unwrap_err().is_cancelled());
3340 release.add_permits(2);
3341 let mut completed = tokio::time::timeout(Duration::from_secs(5), async {
3343 let mut regions = Vec::new();
3344 while let Some(region) = completions.recv().await {
3345 regions.push(region);
3346 }
3347 regions
3348 })
3349 .await
3350 .unwrap();
3351 completed.sort_unstable();
3352 assert_eq!(
3353 completed,
3354 vec![1, 2],
3355 "confirmed={confirmed}, retries={max_retries}"
3356 );
3357 assert!(
3358 starts.try_recv().is_err(),
3359 "cancelled caller started another retry"
3360 );
3361 assert_eq!(counter.written(), 1);
3362 }
3363 }
3364 }
3365
3366 #[tokio::test]
3367 async fn test_tracked_cancellation_stops_backoff() {
3368 let inserter = batcher_test_inserter().await;
3369 let (started, mut starts) = mpsc::unbounded_channel();
3370 let (completed, mut completions) = mpsc::unbounded_channel();
3371 let node_manager = Arc::new(GatedDatanodeHandler {
3372 started,
3373 completed,
3374 release: Arc::new(Semaphore::new(2)),
3375 });
3376 let dispatcher = WriteDispatcher::new(node_manager, inserter.partition_manager.clone())
3377 .with_retry(
3378 WriteRetryOptions {
3379 max_retries: 2,
3380 initial_backoff: Duration::from_secs(60),
3381 max_backoff: Duration::from_secs(60),
3382 ..Default::default()
3383 },
3384 inserter.catalog_manager.clone(),
3385 Arc::new(DummyCacheInvalidator),
3386 );
3387 let factory = RegionRequestFactory::new(RegionRequestHeader::default());
3388 let requests = [1, 2]
3389 .into_iter()
3390 .map(|region| TrackedWrite {
3391 request: factory.build_insert(RegionInsertRequests {
3392 requests: vec![RegionInsertRequest {
3393 region_id: RegionId::new(1, region).as_u64(),
3394 ..Default::default()
3395 }],
3396 }),
3397 owners: vec![(region - 1) as usize],
3398 })
3399 .collect();
3400 let (writes, progress) = dispatcher.prepare_tracked(requests, 2).await;
3401 let cancellation = CancellationToken::new();
3402 let execute = dispatcher.execute_tracked(writes, progress, None, Some(&cancellation));
3403 tokio::pin!(execute);
3404 tokio::time::timeout(
3406 Duration::from_secs(1),
3407 future::poll_fn(|cx| {
3408 assert!(execute.as_mut().poll(cx).is_pending());
3409 if completions.len() == 2 {
3410 Poll::Ready(())
3411 } else {
3412 Poll::Pending
3413 }
3414 }),
3415 )
3416 .await
3417 .unwrap();
3418 for region in [1, 2] {
3419 assert_eq!(starts.try_recv().unwrap(), region);
3420 assert_eq!(completions.try_recv().unwrap(), region);
3421 }
3422 cancellation.cancel();
3423 let progress = tokio::time::timeout(Duration::from_secs(1), execute)
3424 .await
3425 .unwrap();
3426 assert_eq!(progress[0].affected_rows, 1);
3427 assert!(progress[0].error.is_none());
3428 assert_eq!(progress[1].affected_rows, 0);
3429 assert_eq!(
3430 progress[1].error.as_ref().unwrap().retry_hint(),
3431 RetryHint::Retryable
3432 );
3433 assert!(starts.try_recv().is_err());
3434 }
3435
3436 #[tokio::test]
3438 async fn test_write_meter_admission() {
3439 use std::cell::Cell;
3440 use std::sync::Mutex;
3441 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
3442
3443 use api::region::RegionResponse;
3444 use api::v1::region::region_request::Body;
3445 use arrow::array::{Int32Array, TimestampMillisecondArray};
3446 use arrow::record_batch::RecordBatch;
3447 use bytes::Bytes;
3448 use common_error::ext::{ErrorExt, RetryHint};
3449 use common_error::status_code::StatusCode;
3450 use common_grpc::flight::{FlightEncoder, FlightMessage};
3451 use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher;
3452 use futures::future::BoxFuture;
3453 use meter_core::ItemCalculator;
3454 use meter_core::collect::{Collect, WriteRejected};
3455 use meter_core::data::MeterRecord;
3456 use meter_core::global::global_registry;
3457 use session::context::Channel;
3458
3459 const CATALOG: &str = "write_meter_test";
3460
3461 #[derive(Default)]
3462 struct Meter {
3463 reject: AtomicBool,
3464 attempts: Mutex<Vec<MeterRecord>>,
3465 accepted_value: AtomicU64,
3466 }
3467
3468 impl Collect for Meter {
3469 fn on_write(
3470 &self,
3471 record: MeterRecord,
3472 ) -> BoxFuture<'_, std::result::Result<(), WriteRejected>> {
3473 Box::pin(async move {
3474 if record.catalog != CATALOG {
3475 return Ok(());
3476 }
3477 let value = record.value;
3478 self.attempts.lock().unwrap().push(record);
3479 if self.reject.load(Ordering::Relaxed) {
3480 return Err(WriteRejected::new("database row quota exhausted"));
3481 }
3482 self.accepted_value.fetch_add(value, Ordering::Relaxed);
3483 Ok(())
3484 })
3485 }
3486
3487 fn on_read(&self, _: MeterRecord) {}
3488 }
3489
3490 impl ItemCalculator<InstantAndNormalInsertRequests> for Meter {
3491 fn calc(&self, _: &InstantAndNormalInsertRequests) -> u64 {
3492 17
3493 }
3494 }
3495
3496 let kv_backend = prepare_mocked_backend().await;
3497 let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
3498 let (sender, mut dispatched) = tokio::sync::mpsc::channel(16);
3499 let watcher = DatanodeWatcher::new(sender).with_handler(|_, request| {
3500 let rows = match request.body.unwrap() {
3501 Body::Inserts(requests) => requests
3502 .requests
3503 .iter()
3504 .filter_map(|request| request.rows.as_ref())
3505 .map(|rows| rows.rows.len())
3506 .sum(),
3507 Body::BulkInsert(_) => 2,
3509 body => panic!("unexpected request: {body:?}"),
3510 };
3511 Ok(RegionResponse::new(rows))
3512 });
3513 let flow_cache = Cache::new(100);
3514 let inserter = Inserter::new(
3515 catalog::memory::MemoryCatalogManager::new(),
3516 partition_manager,
3517 Arc::new(MockDatanodeManager::new(watcher)),
3518 Arc::new(new_table_flownode_set_cache(
3519 String::new(),
3520 flow_cache.clone(),
3521 kv_backend,
3522 )),
3523 true,
3524 );
3525 let mut table_info = new_test_table_info(1, "table_1", [1].into_iter());
3526 table_info.catalog_name = CATALOG.to_string();
3527 table_info.schema_name = "target_db".to_string();
3528 let table_info = Arc::new(table_info);
3529 let table_infos = HashMap::from_iter([(1, table_info.clone())]);
3530 let ctx = Arc::new(QueryContext::with_channel(
3531 DEFAULT_CATALOG_NAME,
3532 DEFAULT_SCHEMA_NAME,
3533 Channel::Postgres,
3534 ));
3535 let rows_request = || {
3536 let build = |num_rows| RegionInsertRequests {
3537 requests: vec![RegionInsertRequest {
3538 region_id: RegionId::new(1, 1).as_u64(),
3539 rows: Some(Rows {
3540 schema: vec![],
3541 rows: vec![api::v1::Row { values: vec![] }; num_rows],
3542 }),
3543 ..Default::default()
3544 }],
3545 };
3546 InstantAndNormalInsertRequests {
3547 normal_requests: build(3),
3548 instant_requests: build(2),
3549 }
3550 };
3551
3552 let output = inserter
3554 .do_request(rows_request(), &table_infos, &ctx)
3555 .await
3556 .unwrap();
3557 assert_eq!(output.meta.cost, 0);
3558 assert!(matches!(output.data, OutputData::AffectedRows(3)));
3559 dispatched.try_recv().unwrap();
3560
3561 let meter = Arc::new(Meter::default());
3562 global_registry().set_collector(meter.clone());
3563 global_registry().register_calculator(meter.clone());
3564 let enabled = Cell::new(false);
3566 write_meter!({
3567 enabled.set(true);
3568 MeterRecord::new("probe".into(), "probe".into(), 0, 0, 0)
3569 })
3570 .await
3571 .unwrap();
3572 let enabled = enabled.get();
3573
3574 let output = inserter
3575 .do_request(rows_request(), &table_infos, &ctx)
3576 .await
3577 .unwrap();
3578 assert_eq!(output.meta.cost, if enabled { 17 } else { 0 });
3579 assert!(matches!(output.data, OutputData::AffectedRows(3)));
3580 dispatched.try_recv().unwrap();
3581 assert!(flow_cache.contains_key(&1));
3582 flow_cache.invalidate_all();
3583 meter.reject.store(true, Ordering::Relaxed);
3584 let result = inserter
3585 .do_request(rows_request(), &table_infos, &ctx)
3586 .await;
3587 if enabled {
3588 let error = result.unwrap_err();
3589 assert_eq!(error.status_code(), StatusCode::RateLimited);
3590 assert_eq!(error.retry_hint(), RetryHint::Retryable);
3591 assert!(error.to_string().contains("database row quota exhausted"));
3592 assert!(dispatched.try_recv().is_err());
3593 assert!(
3594 !flow_cache.contains_key(&1),
3595 "rejected write reached flow mirroring"
3596 );
3597 assert_eq!(meter.accepted_value.load(Ordering::Relaxed), 17);
3598 let attempts = meter.attempts.lock().unwrap();
3599 assert_eq!(attempts.len(), 2);
3600 for record in attempts.iter() {
3601 assert_eq!(record.catalog, CATALOG);
3602 assert_eq!(record.schema, "target_db");
3603 assert_eq!(
3604 (record.rows, record.value, record.source),
3605 (5, 17, Channel::Postgres as u8)
3606 );
3607 }
3608 } else {
3609 assert_eq!(result.unwrap().meta.cost, 0);
3610 dispatched.try_recv().unwrap();
3611 assert!(meter.attempts.lock().unwrap().is_empty());
3612 }
3613 meter.attempts.lock().unwrap().clear();
3614
3615 if enabled {
3617 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3618 let rows = Rows {
3619 schema: vec![
3620 api::v1::helper::tag_column_schema("a", ColumnDataType::Int32),
3621 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
3622 field_column_schema("b", ColumnDataType::Int32),
3623 ],
3624 rows: vec![api::v1::Row {
3625 values: vec![
3626 api::v1::value::ValueData::I32Value(60).into(),
3627 Value {
3628 value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue(
3629 0,
3630 )),
3631 },
3632 api::v1::value::ValueData::I32Value(0).into(),
3633 ],
3634 }],
3635 };
3636 let error = inserter
3637 .submit_table_rows(rows, table_info.clone(), ctx.clone(), &batcher)
3638 .await
3639 .unwrap_err();
3640 assert_eq!(error.status_code(), StatusCode::RateLimited);
3641 let mut attempts = meter.attempts.lock().unwrap();
3642 assert_eq!(attempts.len(), 1);
3643 assert_eq!(attempts[0].catalog, CATALOG);
3644 assert_eq!(attempts[0].schema, "target_db");
3645 assert_eq!(attempts[0].rows, 1);
3646 attempts.clear();
3647 }
3648
3649 let table = Arc::new(table::Table::new(
3650 table_info.clone(),
3651 table::metadata::FilterPushDownType::Unsupported,
3652 Arc::new(DummyDataSource),
3653 ));
3654 let batch = RecordBatch::try_new(
3655 table_info.meta.schema.arrow_schema().clone(),
3656 vec![
3657 Arc::new(Int32Array::from(vec![60, 70])),
3658 Arc::new(TimestampMillisecondArray::from(vec![0, 1])),
3659 Arc::new(Int32Array::from(vec![0, 0])),
3660 ],
3661 )
3662 .unwrap();
3663 let bulk_insert = |batch: RecordBatch| {
3664 let flight_data = FlightEncoder::default()
3665 .encode(FlightMessage::RecordBatch(batch.clone()))
3666 .into_iter()
3667 .next()
3668 .unwrap();
3669 inserter.handle_bulk_insert(
3670 table.clone(),
3671 flight_data,
3672 batch,
3673 Bytes::new(),
3674 false,
3675 Channel::Grpc,
3676 )
3677 };
3678 assert_eq!(bulk_insert(batch.slice(0, 0)).await.unwrap(), 0);
3680 assert!(meter.attempts.lock().unwrap().is_empty());
3681 assert!(dispatched.try_recv().is_err());
3682
3683 for reject in [false, true, false] {
3685 meter.reject.store(reject, Ordering::Relaxed);
3686 let result = bulk_insert(batch.clone()).await;
3687 if enabled && reject {
3688 let error = result.unwrap_err();
3689 assert_eq!(error.status_code(), StatusCode::RateLimited);
3690 assert_eq!(error.retry_hint(), RetryHint::Retryable);
3691 assert!(dispatched.try_recv().is_err());
3692 } else {
3693 assert_eq!(result.unwrap(), 2);
3694 dispatched.try_recv().unwrap();
3695 }
3696 }
3697 {
3698 let attempts = meter.attempts.lock().unwrap();
3699 assert_eq!(attempts.len(), if enabled { 3 } else { 0 });
3700 for record in attempts.iter() {
3701 assert_eq!(record.catalog, CATALOG);
3702 assert_eq!(record.schema, "target_db");
3703 assert_eq!(
3704 (record.rows, record.value, record.source),
3705 (2, 0, Channel::Grpc as u8)
3706 );
3707 }
3708 }
3709 assert_eq!(
3710 meter.accepted_value.load(Ordering::Relaxed),
3711 if enabled { 17 } else { 0 }
3712 );
3713
3714 meter.attempts.lock().unwrap().clear();
3717 let original = Arc::new(QueryContext::with_channel(
3718 CATALOG,
3719 "a",
3720 Channel::Prometheus,
3721 ));
3722 let mut batches = ["a", "b", "a"].map(|schema| {
3723 let ctx = if schema == "a" {
3724 original.clone()
3725 } else {
3726 Arc::new(QueryContext::with_channel(
3727 CATALOG,
3728 schema,
3729 Channel::Prometheus,
3730 ))
3731 };
3732 (
3733 ctx,
3734 RowInsertRequests {
3735 inserts: vec![RowInsertRequest {
3736 table_name: "data".into(),
3737 rows: Some(Rows {
3738 schema: vec![],
3739 rows: vec![api::v1::Row::default(); 2],
3740 }),
3741 }],
3742 },
3743 )
3744 });
3745 admit_row_insert_batches(&mut batches).await.unwrap();
3746 admit_row_insert_batches(&mut batches).await.unwrap();
3747 assert_eq!(original.write_rows_to_admit(CATALOG, "a", 4), 4);
3748 for (ctx, _) in &batches {
3749 assert_eq!(
3750 ctx.write_rows_to_admit(CATALOG, &ctx.current_schema(), 2),
3751 0
3752 );
3753 assert_eq!(ctx.channel(), Channel::Prometheus);
3754 }
3755 {
3756 let attempts = meter.attempts.lock().unwrap();
3757 let totals = attempts
3758 .iter()
3759 .map(|r| (r.schema.as_str(), r.rows, r.value))
3760 .collect::<Vec<_>>();
3761 assert_eq!(
3762 totals,
3763 if enabled {
3764 vec![("a", 4, 0), ("b", 2, 0)]
3765 } else {
3766 vec![]
3767 }
3768 );
3769 }
3770 meter.attempts.lock().unwrap().clear();
3771 meter.reject.store(false, Ordering::Relaxed);
3772 let ctx = Arc::new(QueryContext::with_channel(
3773 CATALOG,
3774 "logical",
3775 Channel::Otlp,
3776 ));
3777 let admitted = admit_write(2, &ctx).await.unwrap();
3778 let mut requests = RowInsertRequests {
3779 inserts: vec![RowInsertRequest {
3780 table_name: "metric".to_string(),
3781 rows: Some(Rows {
3782 schema: vec![],
3783 rows: vec![api::v1::Row::default(); 2],
3784 }),
3785 }],
3786 };
3787 let original = requests.clone();
3788 let cost = Inserter::meter_row_inserts(&mut requests, &admitted)
3789 .await
3790 .unwrap();
3791 assert_eq!(cost, if enabled { 17 } else { 0 });
3792 assert_eq!(requests, original);
3793 let records = meter
3794 .attempts
3795 .lock()
3796 .unwrap()
3797 .iter()
3798 .map(|record| (record.rows, record.value))
3799 .collect::<Vec<_>>();
3800 assert_eq!(
3801 records,
3802 if enabled {
3803 vec![(2, 0), (0, 17)]
3804 } else {
3805 vec![]
3806 }
3807 );
3808 meter.reject.store(true, Ordering::Relaxed);
3809 let result = Inserter::meter_row_inserts(&mut requests, &admitted).await;
3810 if enabled {
3811 assert_eq!(result.unwrap_err().status_code(), StatusCode::RateLimited);
3812 } else {
3813 assert_eq!(result.unwrap(), 0);
3814 }
3815 assert_eq!(requests, original);
3816 }
3817
3818 #[test]
3819 fn test_trace_create_options_copy_compaction_hints() {
3820 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3821 ctx.set_extension("compaction.twcs.time_window", "1h");
3822 ctx.set_extension("compaction.twcs.max_output_file_size", "256MB");
3823 ctx.set_extension("compaction.twcs.remote_compaction", "true");
3824 ctx.set_extension("compaction.unknown", "ignored");
3825 let ctx = Arc::new(ctx);
3826 let mut table_options = Default::default();
3827
3828 fill_table_options_for_create(
3829 &mut table_options,
3830 &AutoCreateTableType::Trace {
3831 alter_existing: true,
3832 },
3833 &ctx,
3834 );
3835
3836 assert_eq!(
3837 Some("twcs"),
3838 table_options.get(COMPACTION_TYPE).map(String::as_str)
3839 );
3840 assert_eq!(
3841 Some("1h"),
3842 table_options
3843 .get("compaction.twcs.time_window")
3844 .map(String::as_str)
3845 );
3846 assert_eq!(
3847 Some("256MB"),
3848 table_options
3849 .get("compaction.twcs.max_output_file_size")
3850 .map(String::as_str)
3851 );
3852 assert_eq!(
3853 Some("true"),
3854 table_options
3855 .get("compaction.twcs.remote_compaction")
3856 .map(String::as_str)
3857 );
3858 assert!(!table_options.contains_key("compaction.unknown"));
3859 }
3860
3861 #[test]
3862 fn test_exemplar_create_options_isolate_metric_context() {
3863 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3864 for (key, value) in [
3865 (APPEND_MODE_KEY, "false"),
3866 (MERGE_MODE_KEY, "last_non_null"),
3867 (PHYSICAL_TABLE_METADATA_KEY, "true"),
3868 (LOGICAL_TABLE_METADATA_KEY, "physical_metrics"),
3869 ("otlp_metric_compat", "prometheus"),
3870 (table::requests::SEMANTIC_SIGNAL_TYPE, "metric"),
3871 (SEMANTIC_PER_TABLE_INDEX_KEY, "{}"),
3872 ("ttl", "7d"),
3873 ("compaction.twcs.time_window", "1h"),
3874 ] {
3875 ctx.set_extension(key, value);
3876 }
3877 let mut options = Default::default();
3878 fill_table_options_for_create(&mut options, &AutoCreateTableType::Exemplar, &Arc::new(ctx));
3879 assert_eq!(
3880 options,
3881 std::collections::HashMap::from([
3882 (APPEND_MODE_KEY.into(), "true".into()),
3883 (
3884 store_api::metric_engine_consts::PRIMARY_KEY_ENCODING.into(),
3885 "sparse".into()
3886 ),
3887 ("ttl".into(), "7d".into()),
3888 (COMPACTION_TYPE.into(), "twcs".into()),
3889 ("compaction.twcs.time_window".into(), "1h".into()),
3890 ])
3891 );
3892 }
3893
3894 #[test]
3895 fn test_skip_wal_does_not_change_table_options() {
3896 check_skip_wal_does_not_change_table_options(false);
3897 check_skip_wal_does_not_change_table_options(true);
3898 }
3899
3900 fn check_skip_wal_does_not_change_table_options(skip_wal: bool) {
3901 let ctx = Arc::new(QueryContext::with(
3902 DEFAULT_CATALOG_NAME,
3903 DEFAULT_SCHEMA_NAME,
3904 ));
3905 ctx.set_skip_wal(skip_wal);
3906 let mut options = Default::default();
3907 fill_table_options_for_create(&mut options, &AutoCreateTableType::Physical, &ctx);
3908 assert!(!options.contains_key(session::hints::INSERT_SKIP_WAL_HINT));
3909 assert!(!options.contains_key("skip_wal"));
3910 }
3911
3912 #[test]
3913 fn test_last_non_null_create_options_preserve_default_without_append_mode() {
3914 let ctx = Arc::new(QueryContext::with(
3915 DEFAULT_CATALOG_NAME,
3916 DEFAULT_SCHEMA_NAME,
3917 ));
3918 let mut table_options = Default::default();
3919
3920 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3921
3922 assert_eq!(
3923 Some("last_non_null"),
3924 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3925 );
3926 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3927 }
3928
3929 #[test]
3930 fn test_fill_table_options_copies_semantic_extensions() {
3931 use table::requests::{
3932 SEMANTIC_METRIC_TYPE, SEMANTIC_PER_TABLE_INDEX_KEY, SEMANTIC_SIGNAL_TYPE,
3933 SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, SOURCE_OPENTELEMETRY,
3934 };
3935
3936 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3937 ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
3938 ctx.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
3939 ctx.set_extension(SEMANTIC_SOURCE_VERSION, "2.0");
3940 ctx.set_extension(SEMANTIC_METRIC_TYPE, "bogus");
3941 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, "{}");
3943 let ctx = Arc::new(ctx);
3944 let mut table_options = Default::default();
3945
3946 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::Physical, &ctx);
3947
3948 assert_eq!(
3949 Some(SIGNAL_TYPE_METRIC),
3950 table_options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str)
3951 );
3952 assert_eq!(
3953 Some(SOURCE_OPENTELEMETRY),
3954 table_options.get(SEMANTIC_SOURCE).map(String::as_str)
3955 );
3956 assert_eq!(
3957 Some("2.0"),
3958 table_options
3959 .get(SEMANTIC_SOURCE_VERSION)
3960 .map(String::as_str)
3961 );
3962 assert!(!table_options.contains_key(SEMANTIC_METRIC_TYPE));
3963 assert!(!table_options.contains_key(SEMANTIC_PER_TABLE_INDEX_KEY));
3964 }
3965
3966 #[test]
3967 fn test_apply_per_table_semantic_options() {
3968 use table::requests::{
3969 SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, SEMANTIC_PER_TABLE_INDEX_KEY,
3970 };
3971
3972 let index = format!(
3973 r#"{{
3974 "{DEFAULT_SCHEMA_NAME}": {{
3975 "http_requests_total": {{
3976 "greptime.semantic.metric.type": "counter",
3977 "greptime.semantic.metric.unit": "By",
3978 "greptime.semantic.metric.type_BOGUS": "x"
3979 }},
3980 "other_table": {{
3981 "greptime.semantic.metric.type": "gauge"
3982 }}
3983 }},
3984 "other_schema": {{
3985 "http_requests_total": {{
3986 "greptime.semantic.metric.type": "gauge"
3987 }}
3988 }}
3989 }}"#
3990 );
3991 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3992 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index);
3993 let ctx = Arc::new(ctx);
3994
3995 let index = parse_per_table_semantic_index(&ctx);
3996 assert!(index.is_some());
3997 let index = index.as_ref();
3998
3999 let mut table_options = std::collections::HashMap::new();
4000 apply_per_table_semantic_options(
4001 &mut table_options,
4002 index,
4003 DEFAULT_SCHEMA_NAME,
4004 "http_requests_total",
4005 );
4006 assert_eq!(
4008 table_options.get(SEMANTIC_METRIC_TYPE).map(String::as_str),
4009 Some("counter")
4010 );
4011 assert_eq!(
4012 table_options.get(SEMANTIC_METRIC_UNIT).map(String::as_str),
4013 Some("By")
4014 );
4015 assert!(!table_options.contains_key("greptime.semantic.metric.type_BOGUS"));
4018 assert_eq!(table_options.len(), 2);
4019
4020 let mut empty = std::collections::HashMap::new();
4021 apply_per_table_semantic_options(&mut empty, index, DEFAULT_SCHEMA_NAME, "not_in_index");
4022 assert!(empty.is_empty());
4023
4024 let mut opts = std::collections::HashMap::new();
4027 apply_per_table_semantic_options(
4028 &mut opts,
4029 index,
4030 "schema_without_entry",
4031 "http_requests_total",
4032 );
4033 assert!(opts.is_empty());
4034
4035 let bare = Arc::new(QueryContext::with(
4037 DEFAULT_CATALOG_NAME,
4038 DEFAULT_SCHEMA_NAME,
4039 ));
4040 assert!(parse_per_table_semantic_index(&bare).is_none());
4041 let mut opts = std::collections::HashMap::new();
4042 apply_per_table_semantic_options(
4043 &mut opts,
4044 None,
4045 DEFAULT_SCHEMA_NAME,
4046 "http_requests_total",
4047 );
4048 assert!(opts.is_empty());
4049 }
4050
4051 #[test]
4052 fn test_last_non_null_create_options_preserve_default_with_append_mode_false() {
4053 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
4054 ctx.set_extension(APPEND_MODE_KEY, "false");
4055 let ctx = Arc::new(ctx);
4056 let mut table_options = Default::default();
4057
4058 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
4059
4060 assert!(!table_options.contains_key(APPEND_MODE_KEY));
4061 assert_eq!(
4062 Some("last_non_null"),
4063 table_options.get(MERGE_MODE_KEY).map(String::as_str)
4064 );
4065 }
4066
4067 #[test]
4068 fn test_last_non_null_create_options_use_configured_merge_mode() {
4069 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
4070 ctx.set_extension(MERGE_MODE_KEY, "last_row");
4071 let ctx = Arc::new(ctx);
4072 let mut table_options = Default::default();
4073
4074 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
4075
4076 assert_eq!(
4077 Some("last_row"),
4078 table_options.get(MERGE_MODE_KEY).map(String::as_str)
4079 );
4080 assert!(!table_options.contains_key(APPEND_MODE_KEY));
4081 }
4082
4083 #[test]
4084 fn test_last_non_null_create_options_use_last_row_with_append_mode_true() {
4085 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
4086 ctx.set_extension(APPEND_MODE_KEY, "true");
4087 let ctx = Arc::new(ctx);
4088 let mut table_options = Default::default();
4089
4090 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
4091
4092 assert_eq!(
4093 Some("true"),
4094 table_options.get(APPEND_MODE_KEY).map(String::as_str)
4095 );
4096 assert_eq!(
4097 Some("last_row"),
4098 table_options.get(MERGE_MODE_KEY).map(String::as_str)
4099 );
4100 }
4101
4102 struct UnexpectedBatcher;
4103
4104 #[async_trait::async_trait]
4105 impl PendingRowsBatcher for UnexpectedBatcher {
4106 async fn acquire(&self) -> Result<Arc<tokio::sync::OwnedSemaphorePermit>> {
4107 panic!("empty writes must not acquire batch admission")
4108 }
4109
4110 async fn submit(
4111 &self,
4112 _table_info: TableInfoRef,
4113 _batch: arrow::record_batch::RecordBatch,
4114 _ctx: QueryContextRef,
4115 _permit: Arc<tokio::sync::OwnedSemaphorePermit>,
4116 ) -> Result<Output> {
4117 panic!("empty writes must not submit a batch")
4118 }
4119 }
4120
4121 async fn batcher_test_inserter() -> Inserter {
4122 let kv_backend = prepare_mocked_backend().await;
4123 Inserter::new(
4124 catalog::memory::MemoryCatalogManager::new(),
4125 create_partition_rule_manager(kv_backend.clone()).await,
4126 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
4127 Arc::new(new_table_flownode_set_cache(
4128 String::new(),
4129 Cache::new(100),
4130 kv_backend,
4131 )),
4132 true,
4133 )
4134 }
4135
4136 #[tokio::test]
4137 async fn test_logical_batcher_eligibility() {
4138 use catalog::RegisterTableRequest;
4139 use catalog::memory::MemoryCatalogManager;
4140 use common_meta::instruction::{CacheIdent, CreateFlow};
4141 use common_meta::kv_backend::KvBackendRef;
4142 use common_meta::kv_backend::memory::MemoryKvBackend;
4143 use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder};
4144 let requests = RowInsertRequests {
4145 inserts: vec![RowInsertRequest {
4146 table_name: "test_table".to_string(),
4147 rows: None,
4148 }],
4149 };
4150 let original =
4151 make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
4152 .table_info();
4153 for case in [
4154 "eligible",
4155 "physical",
4156 "ordinary",
4157 "instant",
4158 "disabled",
4159 "hint",
4160 "flow",
4161 "required_tag",
4162 "default",
4163 ] {
4164 let mut info = (*original).clone();
4165 info.meta.engine = METRIC_ENGINE_NAME.to_string();
4166 info.meta.options.extra_options.insert(
4167 LOGICAL_TABLE_METADATA_KEY.to_string(),
4168 "physical".to_string(),
4169 );
4170 let mut ctx = QueryContext::arc().fork();
4171 match case {
4172 "physical" => {
4173 info.meta
4174 .options
4175 .extra_options
4176 .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "other".to_string());
4177 }
4178 "ordinary" => info.meta.engine = "mito".to_string(),
4179 "instant" => info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant),
4180 "hint" => ctx.set_extension(AUTO_CREATE_TABLE_KEY, "false"),
4181 "required_tag" | "default" => {
4182 let mut columns = info.meta.schema.column_schemas().to_vec();
4183 let mut tag = ColumnSchema::new(
4184 "tag",
4185 ConcreteDataType::string_datatype(),
4186 case != "required_tag",
4187 );
4188 if case == "default" {
4189 tag = tag
4190 .with_default_constraint(Some(ColumnDefaultConstraint::null_value()))
4191 .unwrap();
4192 }
4193 columns.push(tag);
4194 info.meta.schema = Arc::new(
4195 SchemaBuilder::try_from_columns(columns)
4196 .unwrap()
4197 .build()
4198 .unwrap(),
4199 );
4200 info.meta.primary_key_indices = vec![2];
4201 }
4202 _ => {}
4203 }
4204 let catalog = MemoryCatalogManager::with_default_setup();
4205 let table = Arc::new(table::Table::new(
4206 Arc::new(info),
4207 table::metadata::FilterPushDownType::Unsupported,
4208 Arc::new(DummyDataSource),
4209 ));
4210 catalog
4211 .register_table_sync(RegisterTableRequest {
4212 catalog: DEFAULT_CATALOG_NAME.to_string(),
4213 schema: DEFAULT_SCHEMA_NAME.to_string(),
4214 table_name: "test_table".to_string(),
4215 table_id: 1,
4216 table,
4217 })
4218 .unwrap();
4219 let mut inserter = batcher_test_inserter().await;
4220 inserter.catalog_manager = catalog;
4221 inserter.auto_create_table = case != "disabled";
4222 let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::default());
4223 inserter.table_flownode_set_cache = Arc::new(new_table_flownode_set_cache(
4224 String::new(),
4225 Cache::new(10),
4226 kv_backend,
4227 ));
4228 if case == "flow" {
4229 inserter
4230 .table_flownode_set_cache
4231 .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
4232 flow_id: 1,
4233 source_table_ids: vec![1],
4234 partition_to_peer_mapping: vec![(0, Peer::empty(1))],
4235 })])
4236 .await
4237 .unwrap();
4238 }
4239 assert_eq!(
4240 inserter
4241 .can_batch_metric_rows(&requests, &Arc::new(ctx), "physical")
4242 .await
4243 .unwrap(),
4244 case == "eligible",
4245 "{case}"
4246 );
4247 }
4248 }
4249
4250 #[tokio::test]
4251 async fn test_batcher_meter_preserves_request() {
4252 let mut requests = RowInsertRequests {
4253 inserts: vec![RowInsertRequest {
4254 table_name: "sample".to_string(),
4255 rows: Some(Rows {
4256 schema: vec![],
4257 rows: vec![api::v1::Row {
4258 values: vec![Value {
4259 value_data: Some(api::v1::value::ValueData::F64Value(1.5)),
4260 }],
4261 }],
4262 }),
4263 }],
4264 };
4265 let expected = requests.clone();
4266 Inserter::meter_row_inserts(&mut requests, &QueryContext::arc())
4267 .await
4268 .unwrap();
4269 assert_eq!(requests, expected);
4270 }
4271
4272 #[tokio::test]
4273 async fn test_instant_table_bypasses_batcher() {
4274 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
4275 let inserter = batcher_test_inserter()
4276 .await
4277 .with_pending_rows_batcher(Some(batcher));
4278 let mut ctx = session::context::QueryContextBuilder::default().build();
4279 ctx.set_batching_enabled(true);
4280 let ctx = Arc::new(ctx);
4281 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
4282 .table_info();
4283 assert!(inserter.table_batcher(&table, &ctx).is_some());
4284 let mut instant = (*table).clone();
4285 instant.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
4286 assert!(inserter.table_batcher(&Arc::new(instant), &ctx).is_none());
4287 }
4288
4289 #[tokio::test]
4290 async fn test_empty_prepared_rows_skip_batcher() {
4291 let inserter = batcher_test_inserter().await;
4292 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
4293 .table_info();
4294 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
4295 let ctx = QueryContext::arc();
4296 let output = inserter
4297 .submit_table_rows(
4298 Rows {
4299 schema: vec![],
4300 rows: vec![],
4301 },
4302 table.clone(),
4303 ctx.clone(),
4304 &batcher,
4305 )
4306 .await
4307 .unwrap();
4308 assert!(matches!(output.data, OutputData::AffectedRows(0)));
4309 let output = inserter
4310 .submit_pending_rows(
4311 RowInsertRequests {
4312 inserts: vec![RowInsertRequest {
4313 table_name: table.name.clone(),
4314 rows: None,
4315 }],
4316 },
4317 HashMap::from_iter([(table.table_id(), table)]),
4318 ctx,
4319 &batcher,
4320 )
4321 .await
4322 .unwrap();
4323 assert!(matches!(output.data, OutputData::AffectedRows(0)));
4324 }
4325
4326 #[derive(Clone)]
4330 struct GatedFlownodeHandler {
4331 gate: Arc<tokio::sync::Semaphore>,
4332 affected_rows: u64,
4333 }
4334
4335 #[async_trait::async_trait]
4336 impl MockFlownodeHandler for GatedFlownodeHandler {
4337 async fn handle_inserts(
4338 &self,
4339 _peer: &Peer,
4340 _requests: api::v1::region::InsertRequests,
4341 ) -> common_meta::error::Result<FlowResponse> {
4342 let _permit = self.gate.acquire().await.unwrap();
4343 Ok(FlowResponse {
4344 affected_rows: self.affected_rows,
4345 ..Default::default()
4346 })
4347 }
4348 }
4349
4350 fn flownode_peer() -> Peer {
4351 Peer {
4352 id: 1,
4353 addr: "127.0.0.1:4001".to_string(),
4354 }
4355 }
4356
4357 #[derive(Clone)]
4358 struct RecordingNodeManager {
4359 datanode_dispatch: tokio::sync::mpsc::UnboundedSender<usize>,
4360 flownode_dispatch: tokio::sync::mpsc::UnboundedSender<Peer>,
4361 gates: Arc<std::sync::Mutex<HashMap<u64, Arc<tokio::sync::Semaphore>>>>,
4362 fail_flownode: bool,
4363 }
4364
4365 #[async_trait::async_trait]
4366 impl common_meta::node_manager::DatanodeManager for RecordingNodeManager {
4367 async fn datanode(&self, peer: &Peer) -> common_meta::node_manager::DatanodeRef {
4368 Arc::new(RecordingNode {
4369 peer: peer.clone(),
4370 manager: self.clone(),
4371 })
4372 }
4373 }
4374
4375 #[async_trait::async_trait]
4376 impl common_meta::node_manager::FlownodeManager for RecordingNodeManager {
4377 async fn flownode(&self, peer: &Peer) -> common_meta::node_manager::FlownodeRef {
4378 Arc::new(RecordingNode {
4379 peer: peer.clone(),
4380 manager: self.clone(),
4381 })
4382 }
4383 }
4384
4385 struct RecordingNode {
4386 peer: Peer,
4387 manager: RecordingNodeManager,
4388 }
4389
4390 #[async_trait::async_trait]
4391 impl common_meta::node_manager::Datanode for RecordingNode {
4392 async fn handle_request(
4393 &self,
4394 _request: Arc<api::v1::region::RegionRequest>,
4395 ) -> common_meta::error::Result<api::region::RegionResponse> {
4396 let rows = match _request.body.as_ref() {
4397 Some(api::v1::region::region_request::Body::Inserts(requests)) => requests
4398 .requests
4399 .iter()
4400 .filter_map(|request| request.rows.as_ref())
4401 .map(|rows| rows.rows.len())
4402 .sum(),
4403 _ => 0,
4404 };
4405 let _ = self.manager.datanode_dispatch.send(rows);
4406 Ok(api::region::RegionResponse::new(rows))
4407 }
4408
4409 async fn handle_query(
4410 &self,
4411 _request: common_query::request::QueryRequest,
4412 ) -> common_meta::error::Result<common_recordbatch::SendableRecordBatchStream> {
4413 unreachable!()
4414 }
4415 }
4416
4417 #[async_trait::async_trait]
4418 impl common_meta::node_manager::Flownode for RecordingNode {
4419 async fn handle(
4420 &self,
4421 _request: api::v1::flow::FlowRequest,
4422 ) -> common_meta::error::Result<FlowResponse> {
4423 unreachable!()
4424 }
4425
4426 async fn handle_inserts(
4427 &self,
4428 _request: api::v1::region::InsertRequests,
4429 ) -> common_meta::error::Result<FlowResponse> {
4430 let _ = self.manager.flownode_dispatch.send(self.peer.clone());
4431 if self.manager.fail_flownode {
4432 return Err(common_meta::error::UnexpectedSnafu {
4433 err_msg: "test flownode failure".to_string(),
4434 }
4435 .build());
4436 }
4437 let gate = self
4438 .manager
4439 .gates
4440 .lock()
4441 .unwrap()
4442 .get(&self.peer.id)
4443 .cloned();
4444 if let Some(gate) = gate {
4445 let _permit = gate.acquire().await.unwrap();
4446 }
4447 Ok(FlowResponse::default())
4448 }
4449
4450 async fn handle_mark_window_dirty(
4451 &self,
4452 _request: api::v1::flow::DirtyWindowRequests,
4453 ) -> common_meta::error::Result<FlowResponse> {
4454 unreachable!()
4455 }
4456 }
4457
4458 fn mirror_requests(peer: &Peer, num_rows: usize) -> HashMap<Peer, RegionInsertRequests> {
4459 HashMap::from_iter([(
4460 peer.clone(),
4461 RegionInsertRequests {
4462 requests: vec![RegionInsertRequest {
4463 region_id: RegionId::new(1, 1).as_u64(),
4464 rows: Some(Rows {
4465 schema: vec![],
4466 rows: vec![api::v1::Row { values: vec![] }; num_rows],
4467 }),
4468 ..Default::default()
4469 }],
4470 },
4471 )])
4472 }
4473
4474 static MIRROR_METRIC_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
4479
4480 async fn lock_mirror_metrics() -> tokio::sync::MutexGuard<'static, ()> {
4481 MIRROR_METRIC_TEST_LOCK.lock().await
4482 }
4483
4484 #[tokio::test]
4488 async fn flow_mirror_saturation_preserves_instant_and_normal_contracts() {
4489 use common_error::ext::{ErrorExt, RetryHint};
4490 use common_error::status_code::StatusCode;
4491
4492 let _guard = lock_mirror_metrics().await;
4493 let kv_backend = prepare_mocked_backend().await;
4494 let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
4495 let (datanode_tx, mut datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4496 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4497 let node_manager = Arc::new(RecordingNodeManager {
4498 datanode_dispatch: datanode_tx,
4499 flownode_dispatch: flownode_tx,
4500 gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
4501 fail_flownode: false,
4502 });
4503 let flow_cache = Cache::new(10);
4504 let flow_cache_backend = prepare_mocked_backend().await;
4505 let inserter = Inserter::new(
4506 catalog::memory::MemoryCatalogManager::new(),
4507 partition_manager,
4508 node_manager,
4509 Arc::new(new_table_flownode_set_cache(
4510 String::new(),
4511 flow_cache.clone(),
4512 flow_cache_backend,
4513 )),
4514 true,
4515 );
4516 let peer = flownode_peer();
4517 inserter
4518 .table_flownode_set_cache
4519 .invalidate(&[common_meta::instruction::CacheIdent::CreateFlow(
4520 common_meta::instruction::CreateFlow {
4521 flow_id: 1,
4522 source_table_ids: vec![1, 2],
4523 partition_to_peer_mapping: vec![(0, peer)],
4524 },
4525 )])
4526 .await
4527 .unwrap();
4528 let mut normal_info = new_test_table_info(1, "normal_table", [1].into_iter());
4529 normal_info.catalog_name = DEFAULT_CATALOG_NAME.to_string();
4530 normal_info.schema_name = DEFAULT_SCHEMA_NAME.to_string();
4531 let mut instant_info = new_test_table_info(2, "instant_table", [1].into_iter());
4532 instant_info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
4533 let table_infos =
4534 HashMap::from_iter([(1, Arc::new(normal_info)), (2, Arc::new(instant_info))]);
4535 let ctx = Arc::new(QueryContext::with(
4536 DEFAULT_CATALOG_NAME,
4537 DEFAULT_SCHEMA_NAME,
4538 ));
4539 let insert = |table_id, num_rows| RegionInsertRequests {
4540 requests: vec![RegionInsertRequest {
4541 region_id: RegionId::new(table_id, 1).as_u64(),
4542 rows: Some(Rows {
4543 schema: vec![],
4544 rows: vec![api::v1::Row::default(); num_rows],
4545 }),
4546 ..Default::default()
4547 }],
4548 };
4549 let make_request = |normal, instant| InstantAndNormalInsertRequests {
4550 normal_requests: if normal == 0 {
4551 RegionInsertRequests::default()
4552 } else {
4553 insert(1, normal)
4554 },
4555 instant_requests: if instant == 0 {
4556 RegionInsertRequests::default()
4557 } else {
4558 insert(2, instant)
4559 },
4560 };
4561
4562 let pending = inserter.mirror_pending_rows.clone();
4563 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4564 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
4565 pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
4566 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4567 for request in [make_request(0, 1), make_request(1, 1)] {
4568 let error = inserter
4569 .do_request(request, &table_infos, &ctx)
4570 .await
4571 .unwrap_err();
4572 assert_eq!(error.status_code(), StatusCode::RateLimited);
4573 assert_eq!(error.retry_hint(), RetryHint::Retryable);
4574 assert!(datanode_rx.try_recv().is_err());
4575 assert!(flownode_rx.try_recv().is_err());
4576 assert_eq!(
4577 pending.load(Ordering::Relaxed),
4578 MAX_MIRROR_PENDING_ROWS as u64
4579 );
4580 }
4581 assert_eq!(
4582 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4583 dropped_before
4584 );
4585
4586 for pending_before in [0, MAX_MIRROR_PENDING_ROWS as u64] {
4587 if pending_before == 0 {
4588 pending.store(0, Ordering::Relaxed);
4589 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
4590 }
4591 let gauge_before_rejection = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4592 for (normal, instant) in [
4593 (0, MAX_MIRROR_PENDING_ROWS + 1),
4594 (MAX_MIRROR_PENDING_ROWS / 2, MAX_MIRROR_PENDING_ROWS / 2 + 1),
4595 ] {
4596 let error = inserter
4597 .do_request(make_request(normal, instant), &table_infos, &ctx)
4598 .await
4599 .unwrap_err();
4600 assert_eq!(error.status_code(), StatusCode::InvalidArguments);
4601 assert_eq!(error.retry_hint(), RetryHint::NonRetryable);
4602 let message = error.to_string();
4603 assert!(message.contains("1000001"), "{message}");
4604 assert!(message.contains("1000000"), "{message}");
4605 assert!(message.contains("reduce the batch size"), "{message}");
4606 assert!(datanode_rx.try_recv().is_err());
4607 assert!(flownode_rx.try_recv().is_err());
4608 assert_eq!(pending.load(Ordering::Relaxed), pending_before);
4609 assert_eq!(
4610 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4611 gauge_before_rejection
4612 );
4613 assert_eq!(
4614 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4615 dropped_before
4616 );
4617 }
4618 if pending_before == 0 {
4619 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
4620 pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
4621 }
4622 }
4623
4624 let result = inserter
4625 .do_request(make_request(1, 0), &table_infos, &ctx)
4626 .await
4627 .unwrap();
4628 assert!(matches!(result.data, OutputData::AffectedRows(1)));
4629 assert_eq!(1, datanode_rx.try_recv().unwrap());
4630 assert!(flownode_rx.try_recv().is_err());
4631 assert_eq!(
4632 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4633 dropped_before + 1
4634 );
4635 pending.store(0, Ordering::Relaxed);
4636 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
4637 assert_eq!(
4638 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4639 gauge_before
4640 );
4641 }
4642
4643 #[tokio::test]
4644 async fn flow_mirror_new_counts_regions_and_peer_clones() {
4645 use common_meta::instruction::{CacheIdent, CreateFlow};
4646
4647 let peer1 = flownode_peer();
4648 let peer2 = Peer {
4649 id: 2,
4650 addr: "127.0.0.1:4002".to_string(),
4651 };
4652 let cache_backend = prepare_mocked_backend().await;
4653 let cache = Arc::new(new_table_flownode_set_cache(
4654 String::new(),
4655 Cache::new(10),
4656 cache_backend,
4657 ));
4658 cache
4659 .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
4660 flow_id: 1,
4661 source_table_ids: vec![1],
4662 partition_to_peer_mapping: vec![(0, peer1), (1, peer2)],
4663 })])
4664 .await
4665 .unwrap();
4666 let request = |region_id| RegionInsertRequest {
4667 region_id,
4668 rows: Some(Rows {
4669 schema: vec![],
4670 rows: vec![api::v1::Row { values: vec![] }; 3],
4671 }),
4672 ..Default::default()
4673 };
4674 let requests = [
4675 request(RegionId::new(1, 0).as_u64()),
4676 request(RegionId::new(1, 1).as_u64()),
4677 ];
4678 let task = FlowMirrorTask::new(&cache, requests.iter()).await.unwrap();
4679 assert_eq!(12, task.pending_rows());
4680 assert_eq!(2, task.requests.len());
4681 }
4682
4683 #[tokio::test]
4684 async fn flow_mirror_dropped_when_pending_exceeds_limit() {
4685 let _guard = lock_mirror_metrics().await;
4686 let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64));
4687 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4688 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4689 let num_rows = 100;
4690
4691 let task = FlowMirrorTask {
4692 requests: mirror_requests(&flownode_peer(), num_rows),
4693 };
4694 task.detach(
4696 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
4697 pending.clone(),
4698 false,
4699 )
4700 .unwrap();
4701
4702 assert_eq!(
4703 num_rows as u64,
4704 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get() - dropped_before
4705 );
4706 assert_eq!(
4708 MAX_MIRROR_PENDING_ROWS as u64,
4709 pending.load(Ordering::Relaxed)
4710 );
4711 assert_eq!(
4714 gauge_before,
4715 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4716 "a dropped batch must not change the pending gauge"
4717 );
4718
4719 let task = FlowMirrorTask {
4720 requests: mirror_requests(&flownode_peer(), num_rows),
4721 };
4722 let error = task
4723 .detach(
4724 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
4725 pending.clone(),
4726 true,
4727 )
4728 .err()
4729 .unwrap();
4730 use common_error::ext::{ErrorExt, RetryHint};
4731 use common_error::status_code::StatusCode;
4732 assert_eq!(error.status_code(), StatusCode::RateLimited);
4733 assert_eq!(error.retry_hint(), RetryHint::Retryable);
4734 assert_eq!(
4735 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4736 dropped_before + num_rows as u64,
4737 "rejected instant rows are not successful best-effort drops"
4738 );
4739 assert_eq!(
4740 MAX_MIRROR_PENDING_ROWS as u64,
4741 pending.load(Ordering::Relaxed)
4742 );
4743 assert_eq!(
4744 gauge_before,
4745 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get()
4746 );
4747 }
4748
4749 #[tokio::test]
4753 async fn flow_mirror_reserves_and_releases_per_peer_shares() {
4754 let _guard = lock_mirror_metrics().await;
4755 let per_peer_rows = 10;
4756 let total_rows = (2 * per_peer_rows) as u64;
4757 let start = MAX_MIRROR_PENDING_ROWS as u64 - total_rows;
4759 let pending = Arc::new(AtomicU64::new(start));
4760 let gate = Arc::new(tokio::sync::Semaphore::new(0));
4761 let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4762 gate: gate.clone(),
4763 affected_rows: per_peer_rows as u64,
4764 }));
4765
4766 let mut requests = mirror_requests(&flownode_peer(), per_peer_rows);
4767 requests.extend(mirror_requests(
4768 &Peer {
4769 id: 2,
4770 addr: "127.0.0.1:4002".to_string(),
4771 },
4772 per_peer_rows,
4773 ));
4774
4775 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4776 FlowMirrorTask { requests }
4777 .detach(node_manager, pending.clone(), false)
4778 .unwrap();
4779
4780 assert_eq!(
4783 MAX_MIRROR_PENDING_ROWS as u64,
4784 pending.load(Ordering::Relaxed)
4785 );
4786
4787 gate.add_permits(2);
4793 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4794 while pending.load(Ordering::Relaxed) != start
4795 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4796 {
4797 assert!(
4798 std::time::Instant::now() < deadline,
4799 "the mirror tasks must release exactly their own shares"
4800 );
4801 std::thread::sleep(std::time::Duration::from_millis(1));
4802 }
4803 }
4804
4805 #[tokio::test]
4806 async fn flow_mirror_fast_peer_releases_only_its_share_before_drop() {
4807 let _guard = lock_mirror_metrics().await;
4808 let peer1 = flownode_peer();
4809 let peer2 = Peer {
4810 id: 2,
4811 addr: "127.0.0.1:4002".to_string(),
4812 };
4813 let gate1 = Arc::new(tokio::sync::Semaphore::new(0));
4814 let gate2 = Arc::new(tokio::sync::Semaphore::new(0));
4815 let gates = Arc::new(std::sync::Mutex::new(HashMap::from_iter([
4816 (peer1.id, gate1.clone()),
4817 (peer2.id, gate2.clone()),
4818 ])));
4819 let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4820 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4821 let node_manager = Arc::new(RecordingNodeManager {
4822 datanode_dispatch: datanode_tx,
4823 flownode_dispatch: flownode_tx,
4824 gates,
4825 fail_flownode: false,
4826 });
4827 let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64 - 20));
4828 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4829 let mut requests = mirror_requests(&peer1, 10);
4830 requests.extend(mirror_requests(&peer2, 10));
4831 let task = FlowMirrorTask { requests };
4832 task.detach(node_manager.clone(), pending.clone(), false)
4833 .unwrap();
4834 let first = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4835 .await
4836 .unwrap()
4837 .unwrap();
4838 let second = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4839 .await
4840 .unwrap()
4841 .unwrap();
4842 assert_ne!(first.id, second.id);
4843 let fast_gate = if first.id == peer1.id {
4844 gate1.clone()
4845 } else {
4846 gate2.clone()
4847 };
4848 let slow_gate = if first.id == peer1.id {
4849 gate2.clone()
4850 } else {
4851 gate1.clone()
4852 };
4853 fast_gate.add_permits(1);
4854 let deadline = std::time::Instant::now() + Duration::from_secs(10);
4855 while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 10 {
4856 assert!(
4857 std::time::Instant::now() < deadline,
4858 "fast peer did not release"
4859 );
4860 std::thread::sleep(Duration::from_millis(1));
4861 }
4862 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4863 FlowMirrorTask {
4864 requests: mirror_requests(&peer1, 11),
4865 }
4866 .detach(node_manager, pending.clone(), false)
4867 .unwrap();
4868 assert_eq!(
4869 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4870 dropped_before + 11
4871 );
4872 assert!(flownode_rx.try_recv().is_err());
4873 slow_gate.add_permits(1);
4874 while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 20
4875 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4876 {
4877 assert!(
4878 std::time::Instant::now() < deadline,
4879 "slow peer did not release"
4880 );
4881 std::thread::sleep(Duration::from_millis(1));
4882 }
4883 }
4884
4885 #[tokio::test]
4886 async fn flow_mirror_failure_releases_budget_and_gauge() {
4887 let _guard = lock_mirror_metrics().await;
4888 let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4889 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4890 let node_manager = Arc::new(RecordingNodeManager {
4891 datanode_dispatch: datanode_tx,
4892 flownode_dispatch: flownode_tx,
4893 gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
4894 fail_flownode: true,
4895 });
4896 let pending = Arc::new(AtomicU64::new(0));
4897 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4898 FlowMirrorTask {
4899 requests: mirror_requests(&flownode_peer(), 5),
4900 }
4901 .detach(node_manager, pending.clone(), false)
4902 .unwrap();
4903 assert_eq!(
4904 tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4905 .await
4906 .unwrap()
4907 .unwrap()
4908 .id,
4909 flownode_peer().id
4910 );
4911 let deadline = std::time::Instant::now() + Duration::from_secs(10);
4912 while pending.load(Ordering::Relaxed) != 0
4913 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4914 {
4915 assert!(
4916 std::time::Instant::now() < deadline,
4917 "failed task leaked reservation"
4918 );
4919 std::thread::sleep(Duration::from_millis(1));
4920 }
4921 }
4922
4923 #[tokio::test]
4924 async fn flow_mirror_pending_counter_consistent() {
4925 let _guard = lock_mirror_metrics().await;
4926 let pending = Arc::new(AtomicU64::new(0));
4927 let gate = Arc::new(tokio::sync::Semaphore::new(0));
4928 let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4929 gate: gate.clone(),
4930 affected_rows: 7,
4931 }));
4932
4933 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4935 FlowMirrorTask {
4936 requests: HashMap::new(),
4937 }
4938 .detach(node_manager.clone(), pending.clone(), false)
4939 .unwrap();
4940 assert_eq!(0, pending.load(Ordering::Relaxed));
4941 assert_eq!(
4942 gauge_before,
4943 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4944 "an empty batch must not touch the pending gauge"
4945 );
4946
4947 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4948 let num_rows = 7;
4949
4950 FlowMirrorTask {
4951 requests: mirror_requests(&flownode_peer(), num_rows),
4952 }
4953 .detach(node_manager, pending.clone(), false)
4954 .unwrap();
4955
4956 assert_eq!(num_rows as u64, pending.load(Ordering::Relaxed));
4959 assert_eq!(
4960 num_rows as i64,
4961 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() - gauge_before,
4962 "the pending gauge must track the pending rows"
4963 );
4964 assert_eq!(
4965 dropped_before,
4966 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4967 "a batch within the limit must not be counted as dropped"
4968 );
4969
4970 gate.add_permits(1);
4973 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4974 while pending.load(Ordering::Relaxed) != 0
4975 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4976 {
4977 assert!(
4978 std::time::Instant::now() < deadline,
4979 "the mirror task did not release its pending reservation"
4980 );
4981 std::thread::sleep(std::time::Duration::from_millis(1));
4982 }
4983 }
4984
4985 #[test]
4986 fn mirror_logs_are_rate_limited() {
4987 let log = Arc::new(MirrorLog::new());
4988 assert_eq!(Some(1), log.claim_report(1_000));
4989 assert_eq!(None, log.claim_report(1_500));
4990 assert_eq!(None, log.claim_report(10_999));
4991 assert_eq!(Some(3), log.claim_report(11_000));
4992 assert_eq!(None, log.claim_report(11_001));
4993 assert_eq!(Some(2), log.claim_report(21_000));
4994
4995 let concurrent = Arc::new(MirrorLog::new());
4996 assert_eq!(Some(1), concurrent.claim_report(50_000));
4997 let barrier = Arc::new(std::sync::Barrier::new(8));
4998 let threads = (0..8)
4999 .map(|_| {
5000 let log = concurrent.clone();
5001 let barrier = barrier.clone();
5002 std::thread::spawn(move || {
5003 barrier.wait();
5004 log.claim_report(60_000)
5005 })
5006 })
5007 .collect::<Vec<_>>();
5008 let winners = threads
5009 .into_iter()
5010 .filter_map(|thread| thread.join().unwrap())
5011 .collect::<Vec<_>>();
5012 assert_eq!(winners.len(), 1);
5013 let residual = concurrent.events.swap(0, Ordering::Relaxed);
5014 assert_eq!(winners[0] + residual, 8);
5015 }
5016}