1#[cfg(feature = "enterprise")]
16mod flow_write;
17
18use std::collections::BTreeMap;
19use std::sync::atomic::{AtomicU64, Ordering};
20use std::sync::{Arc, LazyLock};
21use std::time::{Duration, Instant};
22
23use ahash::{HashMap, HashMapExt, HashSet, HashSetExt};
24use api::v1::alter_table_expr::Kind;
25use api::v1::column_def::{options_from_compression, options_from_skipping, try_as_column_def};
26use api::v1::region::{
27 InsertRequest as RegionInsertRequest, InsertRequests as RegionInsertRequests,
28 RegionRequestHeader,
29};
30use api::v1::value::ValueData;
31use api::v1::{
32 AlterTableExpr, ColumnDataType, ColumnSchema, CreateTableExpr, InsertRequests,
33 RowInsertRequest, RowInsertRequests, Rows, SemanticType,
34};
35use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema};
36use catalog::CatalogManagerRef;
37use client::{OutputData, OutputMeta};
38use common_base::readable_size::ReadableSize;
39use common_catalog::consts::{
40 DEFAULT_PRIVATE_SCHEMA_NAME, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, SPAN_ID_COLUMN,
41 TIMESTAMP_END_COLUMN, TRACE_ID_COLUMN, TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY,
42 default_engine, is_ddl_reserved_table, trace_operations_table_name, trace_services_table_name,
43};
44use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME;
45use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME;
46use common_grpc_expr::util::ColumnExpr;
47use common_meta::cache::TableFlownodeSetCacheRef;
48use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME;
49use common_meta::node_manager::NodeManagerRef;
50use common_meta::peer::Peer;
51use common_meta::rpc::ddl::TriggerReason;
52use common_query::Output;
53use common_query::native_histogram::{is_native_histogram_value_type, native_histogram_value_type};
54use common_query::prelude::{greptime_timestamp, greptime_value};
55use common_telemetry::tracing_context::TracingContext;
56use common_telemetry::{debug, error, warn};
57use common_time::Timestamp;
58use common_time::timestamp::TimeUnit;
59use datatypes::schema::{ColumnEncoding, CompressionOptions, SkippingIndexOptions};
60use futures_util::future;
61use meter_core::data::MeterRecord;
62use meter_macros::write_meter;
63use partition::manager::PartitionRuleManagerRef;
64use session::context::QueryContextRef;
65use snafu::ResultExt;
66use snafu::prelude::*;
67use sql::partition::partition_rule_for_hexstring;
68use sql::statements::create::Partitions;
69use sql::statements::insert::Insert;
70use store_api::metric_engine_consts::{
71 LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
72};
73use store_api::mito_engine_options::{
74 APPEND_MODE_KEY, COMPACTION_OPTION_KEYS, COMPACTION_TWCS_PREFIX, COMPACTION_TYPE,
75 COMPACTION_TYPE_TWCS, MERGE_MODE_KEY, TTL_KEY,
76};
77use store_api::storage::{RegionId, TableId};
78use table::TableRef;
79use table::metadata::{TableInfo, TableInfoRef};
80use table::requests::{
81 AUTO_CREATE_TABLE_KEY, InsertRequest as TableInsertRequest, SEMANTIC_PER_TABLE_INDEX_KEY,
82 SEMANTIC_PIPELINE, TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1, TABLE_DATA_MODEL_TRACE_V2,
83 TRACE_TABLE_PARTITIONS_HINT_KEY, VALID_TABLE_OPTION_KEYS, is_semantic_option_key,
84 validate_semantic_option,
85};
86use table::table_reference::TableReference;
87use tokio_util::sync::CancellationToken;
88
89use crate::batcher::PendingRowsBatcher;
90use crate::error::{
91 CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, InvalidInsertRequestSnafu,
92 JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu, WriteRejectedSnafu,
93};
94use crate::expr_helper;
95use crate::region_req_factory::RegionRequestFactory;
96use crate::req_convert::common::preprocess_row_insert_requests;
97use crate::req_convert::insert::{
98 ColumnToRow, ImpureDefaultFiller, RowToRegion, StatementToRegion, TableToRegion,
99 fill_reqs_with_impure_default, rows_to_record_batch,
100};
101use crate::statement::StatementExecutor;
102use crate::write_dispatcher::{WriteDispatcher, WriteDispatcherRef};
103
104pub struct Inserter {
105 catalog_manager: CatalogManagerRef,
106 pub(crate) partition_manager: PartitionRuleManagerRef,
107 pub(crate) node_manager: NodeManagerRef,
108 pub(crate) write_dispatcher: WriteDispatcherRef,
109 pub(crate) table_flownode_set_cache: TableFlownodeSetCacheRef,
110 auto_create_table: bool,
114 pending_rows_batcher: Option<Arc<dyn PendingRowsBatcher>>,
115 mirror_pending_rows: Arc<AtomicU64>,
119}
120
121pub type InserterRef = Arc<Inserter>;
122
123#[derive(Clone)]
125pub enum AutoCreateTableType {
126 Logical(String),
128 Physical,
130 Log,
132 LastNonNull,
134 Trace { alter_existing: bool },
136}
137
138impl AutoCreateTableType {
139 pub fn as_str(&self) -> &'static str {
140 match self {
141 AutoCreateTableType::Logical(_) => "logical",
142 AutoCreateTableType::Physical => "physical",
143 AutoCreateTableType::Log => "log",
144 AutoCreateTableType::LastNonNull => "last_non_null",
145 AutoCreateTableType::Trace { .. } => "trace",
146 }
147 }
148
149 fn alter_existing(&self) -> bool {
150 !matches!(
151 self,
152 Self::Trace {
153 alter_existing: false
154 }
155 )
156 }
157}
158
159#[derive(Clone)]
166pub struct InstantAndNormalInsertRequests {
167 pub normal_requests: RegionInsertRequests,
169 pub instant_requests: RegionInsertRequests,
172}
173
174impl Inserter {
175 pub async fn can_batch_metric_rows(
179 &self,
180 requests: &RowInsertRequests,
181 ctx: &QueryContextRef,
182 physical_table: &str,
183 ) -> Result<bool> {
184 if self.auto_create_disabled_reason(ctx)?.is_some() || ctx.extension(TTL_KEY).is_some() {
185 return Ok(false);
186 }
187 for request in &requests.inserts {
188 if request.rows.as_ref().is_some_and(|rows| {
193 rows.schema.iter().any(|column| {
194 column.datatype_extension.is_some()
195 || !matches!(
196 ColumnDataType::try_from(column.datatype),
197 Ok(ColumnDataType::TimestampSecond
198 | ColumnDataType::TimestampMillisecond
199 | ColumnDataType::TimestampMicrosecond
200 | ColumnDataType::TimestampNanosecond
201 | ColumnDataType::Float64
202 | ColumnDataType::String)
203 )
204 })
205 }) {
206 return Ok(false);
207 }
208
209 let Some(table) = self
210 .get_table(
211 ctx.current_catalog(),
212 &ctx.current_schema(),
213 &request.table_name,
214 )
215 .await?
216 else {
217 continue;
218 };
219 let info = table.table_info();
220 if info.meta.engine != METRIC_ENGINE_NAME
221 || info.is_ttl_instant_table()
222 || info
223 .meta
224 .options
225 .extra_options
226 .get(LOGICAL_TABLE_METADATA_KEY)
227 .map(String::as_str)
228 != Some(physical_table)
229 || info
230 .meta
231 .schema
232 .column_schemas()
233 .iter()
234 .any(|column| column.default_constraint().is_some())
235 {
236 return Ok(false);
237 }
238 if info
241 .meta
242 .primary_key_indices
243 .iter()
244 .any(|&index| !info.meta.schema.column_schemas()[index].is_nullable())
245 {
246 return Ok(false);
247 }
248 match self.table_flownode_set_cache.get(info.table_id()).await {
251 Ok(None) => {}
252 Ok(Some(flows)) if flows.is_empty() => {}
253 _ => return Ok(false),
254 }
255 }
256 Ok(true)
257 }
258
259 pub async fn meter_row_inserts(
262 requests: &mut RowInsertRequests,
263 ctx: &QueryContextRef,
264 ) -> Result<u64> {
265 let metered = InstantAndNormalInsertRequests {
266 normal_requests: RegionInsertRequests {
267 requests: requests
268 .inserts
269 .iter_mut()
270 .map(|request| RegionInsertRequest {
271 rows: request.rows.take(),
272 ..Default::default()
273 })
274 .collect(),
275 },
276 instant_requests: RegionInsertRequests::default(),
277 };
278 let cost = write_meter!(
279 ctx.current_catalog(),
280 ctx.current_schema(),
281 metered,
282 ctx.write_rows_to_admit(
283 ctx.current_catalog(),
284 &ctx.current_schema(),
285 count_insert_rows(&metered)?
286 ),
287 ctx.channel() as u8
288 )
289 .await
290 .context(WriteRejectedSnafu);
291 for (request, region) in requests
292 .inserts
293 .iter_mut()
294 .zip(metered.normal_requests.requests)
295 {
296 request.rows = region.rows;
297 }
298 cost
299 }
300
301 pub fn new(
303 catalog_manager: CatalogManagerRef,
304 partition_manager: PartitionRuleManagerRef,
305 node_manager: NodeManagerRef,
306 table_flownode_set_cache: TableFlownodeSetCacheRef,
307 auto_create_table: bool,
308 ) -> Self {
309 Self {
310 write_dispatcher: Arc::new(WriteDispatcher::new(
311 node_manager.clone(),
312 partition_manager.clone(),
313 )),
314 catalog_manager,
315 partition_manager,
316 node_manager,
317 table_flownode_set_cache,
318 auto_create_table,
319 pending_rows_batcher: None,
320 mirror_pending_rows: Arc::new(AtomicU64::new(0)),
321 }
322 }
323
324 pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self {
326 self.write_dispatcher = dispatcher;
327 self
328 }
329
330 pub fn with_pending_rows_batcher(
332 mut self,
333 batcher: Option<Arc<dyn PendingRowsBatcher>>,
334 ) -> Self {
335 self.pending_rows_batcher = batcher;
336 self
337 }
338
339 pub async fn handle_column_inserts(
340 &self,
341 requests: InsertRequests,
342 ctx: QueryContextRef,
343 statement_executor: &StatementExecutor,
344 ) -> Result<Output> {
345 let row_inserts = ColumnToRow::convert(requests)?;
346 self.handle_row_inserts(row_inserts, ctx, statement_executor, false, false)
347 .await
348 }
349
350 pub async fn handle_row_inserts(
352 &self,
353 mut requests: RowInsertRequests,
354 ctx: QueryContextRef,
355 statement_executor: &StatementExecutor,
356 accommodate_existing_schema: bool,
357 is_single_value: bool,
358 ) -> Result<Output> {
359 preprocess_row_insert_requests(&mut requests.inserts)?;
360 self.handle_row_inserts_with_create_type(
361 requests,
362 ctx,
363 statement_executor,
364 AutoCreateTableType::Physical,
365 accommodate_existing_schema,
366 is_single_value,
367 )
368 .await
369 }
370
371 pub async fn handle_log_inserts(
373 &self,
374 requests: RowInsertRequests,
375 ctx: QueryContextRef,
376 statement_executor: &StatementExecutor,
377 ) -> Result<Output> {
378 self.handle_row_inserts_with_create_type(
379 requests,
380 ctx,
381 statement_executor,
382 AutoCreateTableType::Log,
383 false,
384 false,
385 )
386 .await
387 }
388
389 pub async fn handle_trace_inserts(
390 &self,
391 requests: RowInsertRequests,
392 ctx: QueryContextRef,
393 statement_executor: &StatementExecutor,
394 ) -> Result<Output> {
395 self.handle_row_inserts_with_create_type(
396 requests,
397 ctx,
398 statement_executor,
399 AutoCreateTableType::Trace {
400 alter_existing: true,
401 },
402 false,
403 false,
404 )
405 .await
406 }
407
408 pub async fn handle_last_non_null_inserts(
410 &self,
411 requests: RowInsertRequests,
412 ctx: QueryContextRef,
413 statement_executor: &StatementExecutor,
414 accommodate_existing_schema: bool,
415 is_single_value: bool,
416 ) -> Result<Output> {
417 self.handle_row_inserts_with_create_type(
418 requests,
419 ctx,
420 statement_executor,
421 AutoCreateTableType::LastNonNull,
422 accommodate_existing_schema,
423 is_single_value,
424 )
425 .await
426 }
427
428 async fn handle_row_inserts_with_create_type(
430 &self,
431 mut requests: RowInsertRequests,
432 ctx: QueryContextRef,
433 statement_executor: &StatementExecutor,
434 create_type: AutoCreateTableType,
435 accommodate_existing_schema: bool,
436 is_single_value: bool,
437 ) -> Result<Output> {
438 let skip_wal = ctx.skip_wal();
439
440 let batcher = self
441 .pending_rows_batcher
442 .as_ref()
443 .filter(|_| ctx.batching_enabled());
444
445 requests.inserts.retain(|req| {
447 req.rows
448 .as_ref()
449 .map(|r| !r.rows.is_empty())
450 .unwrap_or_default()
451 });
452 validate_column_count_match(&requests)?;
453
454 let CreateAlterTableResult {
455 instant_table_ids,
456 table_infos,
457 } = self
458 .create_or_alter_tables_on_demand(
459 &mut requests,
460 &ctx,
461 create_type,
462 statement_executor,
463 accommodate_existing_schema,
464 is_single_value,
465 None,
466 )
467 .await?;
468
469 if let Some(batcher) = batcher
472 && instant_table_ids.is_empty()
473 && table_infos
474 .values()
475 .all(|info| info.meta.engine == default_engine())
476 {
477 return self
478 .submit_pending_rows(requests, table_infos, ctx, batcher)
479 .await;
480 }
481
482 let name_to_info = table_infos
483 .values()
484 .map(|info| (info.name.clone(), info.clone()))
485 .collect::<HashMap<_, _>>();
486 let inserts = RowToRegion::new(
487 name_to_info,
488 instant_table_ids,
489 self.partition_manager.as_ref(),
490 )
491 .convert(requests, skip_wal)
492 .await?;
493
494 self.do_request(inserts, &table_infos, &ctx).await
495 }
496
497 async fn submit_pending_rows(
498 &self,
499 mut requests: RowInsertRequests,
500 table_infos: HashMap<TableId, Arc<TableInfo>>,
501 ctx: QueryContextRef,
502 batcher: &Arc<dyn PendingRowsBatcher>,
503 ) -> Result<Output> {
504 requests.inserts.retain(|request| {
507 request
508 .rows
509 .as_ref()
510 .is_some_and(|rows| !rows.rows.is_empty())
511 });
512 let by_name = table_infos
513 .values()
514 .map(|info| (info.name.as_str(), info))
515 .collect::<HashMap<_, _>>();
516 let mut prepared = Vec::with_capacity(requests.inserts.len());
517 for request in &mut requests.inserts {
518 let table_info =
519 by_name
520 .get(request.table_name.as_str())
521 .context(TableNotFoundSnafu {
522 table_name: &request.table_name,
523 })?;
524 let Some(rows) = &mut request.rows else {
525 continue;
526 };
527 ImpureDefaultFiller::new((*table_info).clone())?.fill_rows(rows);
528 let batch = rows_to_record_batch(rows, table_info)?;
529 prepared.push(((*table_info).clone(), batch));
530 }
531 let metered = InstantAndNormalInsertRequests {
534 normal_requests: RegionInsertRequests {
535 requests: requests
536 .inserts
537 .into_iter()
538 .map(|request| RegionInsertRequest {
539 rows: request.rows,
540 ..Default::default()
541 })
542 .collect(),
543 },
544 instant_requests: RegionInsertRequests::default(),
545 };
546 let table_info = table_infos.values().next();
547 let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str());
548 let schema =
549 table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone());
550 let write_cost = write_meter!(
551 catalog,
552 &schema,
553 metered,
554 ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&metered)?),
555 ctx.channel() as u8
556 )
557 .await
558 .context(WriteRejectedSnafu)?;
559 prepared.retain(|(_, batch)| batch.num_rows() != 0);
560 let results = if prepared.is_empty() {
561 Vec::new()
562 } else {
563 let permit = batcher.acquire().await?;
565 let submissions = prepared.into_iter().map(|(info, batch)| {
566 let mut target_ctx = ctx.fork();
568 target_ctx.set_current_catalog(&info.catalog_name);
569 target_ctx.set_current_schema(&info.schema_name);
570 batcher.submit(info, batch, Arc::new(target_ctx), permit.clone())
571 });
572 future::join_all(submissions).await
574 };
575 let mut affected_rows = 0;
576 let mut meta = OutputMeta::new_with_cost(write_cost as _);
577 for result in results {
578 let output = result?;
579 affected_rows += output.extract_rows_and_cost().0;
580 meta.write_completions.extend(output.meta.write_completions);
581 }
582 Ok(Output::new(OutputData::AffectedRows(affected_rows), meta))
583 }
584
585 pub async fn handle_metric_row_inserts(
587 &self,
588 mut requests: RowInsertRequests,
589 ctx: QueryContextRef,
590 statement_executor: &StatementExecutor,
591 physical_table: String,
592 ) -> Result<Output> {
593 let skip_wal = ctx.skip_wal();
594
595 requests.inserts.retain(|req| {
597 req.rows
598 .as_ref()
599 .map(|r| !r.rows.is_empty())
600 .unwrap_or_default()
601 });
602 validate_column_count_match(&requests)?;
603
604 let physical_table_ref = self
606 .create_physical_table_on_demand(&ctx, physical_table.clone(), statement_executor)
607 .await?;
608
609 let CreateAlterTableResult {
621 instant_table_ids,
622 table_infos,
623 } = self
624 .create_or_alter_tables_on_demand(
625 &mut requests,
626 &ctx,
627 AutoCreateTableType::Logical(physical_table.clone()),
628 statement_executor,
629 true,
630 true,
631 table_time_index_unit(&physical_table_ref),
632 )
633 .await?;
634 let name_to_info = table_infos
635 .values()
636 .map(|info| (info.name.clone(), info.clone()))
637 .collect::<HashMap<_, _>>();
638 let inserts = RowToRegion::new(name_to_info, instant_table_ids, &self.partition_manager)
639 .convert(requests, skip_wal)
640 .await?;
641
642 self.do_request(inserts, &table_infos, &ctx).await
643 }
644
645 fn table_batcher(
646 &self,
647 table_info: &TableInfoRef,
648 ctx: &QueryContextRef,
649 ) -> Option<&Arc<dyn PendingRowsBatcher>> {
650 self.pending_rows_batcher.as_ref().filter(|_| {
651 ctx.batching_enabled()
652 && !table_info.is_ttl_instant_table()
653 && table_info.meta.engine == default_engine()
654 })
655 }
656
657 async fn submit_table_rows(
658 &self,
659 rows: Rows,
660 table_info: TableInfoRef,
661 ctx: QueryContextRef,
662 batcher: &Arc<dyn PendingRowsBatcher>,
663 ) -> Result<Output> {
664 let requests = RowInsertRequests {
665 inserts: vec![RowInsertRequest {
666 table_name: table_info.name.clone(),
667 rows: Some(rows),
668 }],
669 };
670 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info)]);
671 self.submit_pending_rows(requests, table_infos, ctx, batcher)
672 .await
673 }
674
675 pub async fn handle_table_insert(
676 &self,
677 request: TableInsertRequest,
678 ctx: QueryContextRef,
679 ) -> Result<Output> {
680 let catalog = request.catalog_name.as_str();
681 let schema = request.schema_name.as_str();
682 let table_name = request.table_name.as_str();
683 let table = self.get_table(catalog, schema, table_name).await?;
684 let table = table.with_context(|| TableNotFoundSnafu {
685 table_name: common_catalog::format_full_table_name(catalog, schema, table_name),
686 })?;
687 let table_info = table.table_info();
688
689 let converter = TableToRegion::new(&table_info, &self.partition_manager);
690 let skip_wal = request.skip_wal;
691 let rows = converter.prepare(request)?;
692 if let Some(batcher) = self.table_batcher(&table_info, &ctx) {
693 return self.submit_table_rows(rows, table_info, ctx, batcher).await;
694 }
695 let inserts = converter.partition(rows, skip_wal).await?;
696
697 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]);
698
699 self.do_request(inserts, &table_infos, &ctx).await
700 }
701
702 pub async fn handle_statement_insert(
703 &self,
704 insert: &Insert,
705 ctx: &QueryContextRef,
706 ) -> Result<Output> {
707 let converter =
708 StatementToRegion::new(self.catalog_manager.as_ref(), &self.partition_manager, ctx);
709 let (rows, table_info) = converter.prepare(insert, ctx).await?;
710 if let Some(batcher) = self.table_batcher(&table_info, ctx) {
711 return self
712 .submit_table_rows(rows, table_info, ctx.clone(), batcher)
713 .await;
714 }
715 let inserts = converter.partition(rows, table_info.clone(), ctx).await?;
716
717 let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]);
718
719 self.do_request(inserts, &table_infos, ctx).await
720 }
721}
722
723pub async fn admit_write(rows: u64, ctx: &QueryContextRef) -> Result<QueryContextRef> {
726 write_meter!(MeterRecord::new(
729 ctx.current_catalog().to_string(),
730 ctx.current_schema(),
731 0,
732 ctx.write_rows_to_admit(ctx.current_catalog(), &ctx.current_schema(), rows),
733 ctx.channel() as u8,
734 ))
735 .await
736 .context(WriteRejectedSnafu)?;
737 Ok(Arc::new(ctx.with_write_admission()))
738}
739
740pub async fn admit_row_insert_batches(
743 batches: &mut [(QueryContextRef, RowInsertRequests)],
744) -> Result<()> {
745 let mut totals = BTreeMap::<_, (QueryContextRef, u64)>::new();
746 for (ctx, requests) in batches.iter() {
747 let catalog = ctx.current_catalog();
748 let schema = ctx.current_schema();
749 if ctx.write_rows_to_admit(catalog, &schema, 1) == 0 {
750 continue;
751 }
752 let (_, total) = totals
753 .entry((catalog.to_string(), schema.clone()))
754 .or_insert_with(|| (ctx.clone(), 0));
755 for rows in requests.inserts.iter().filter_map(|r| r.rows.as_ref()) {
756 *total =
757 total
758 .checked_add(rows.rows.len() as u64)
759 .context(InvalidInsertRequestSnafu {
760 reason: "Insert row count exceeds u64::MAX",
761 })?;
762 }
763 }
764 for (ctx, rows) in totals.values() {
765 admit_write(*rows, ctx).await?;
766 }
767 for (ctx, _) in batches {
768 *ctx = Arc::new(ctx.with_write_admission());
769 }
770 Ok(())
771}
772
773fn count_insert_rows(requests: &InstantAndNormalInsertRequests) -> Result<u64> {
774 requests
775 .normal_requests
776 .requests
777 .iter()
778 .chain(&requests.instant_requests.requests)
779 .filter_map(|request| request.rows.as_ref())
780 .try_fold(0u64, |total, rows| {
781 total
782 .checked_add(rows.rows.len() as u64)
783 .context(InvalidInsertRequestSnafu {
784 reason: "Insert row count exceeds u64::MAX",
785 })
786 })
787}
788
789impl Inserter {
790 async fn do_request(
791 &self,
792 requests: InstantAndNormalInsertRequests,
793 table_infos: &HashMap<TableId, Arc<TableInfo>>,
794 ctx: &QueryContextRef,
795 ) -> Result<Output> {
796 let requests = fill_reqs_with_impure_default(table_infos, requests)?;
798
799 let table_info = table_infos.values().next();
802 let catalog = table_info.map_or(ctx.current_catalog(), |info| info.catalog_name.as_str());
803 let schema =
804 table_info.map_or_else(|| ctx.current_schema(), |info| info.schema_name.clone());
805 let write_cost = write_meter!(
806 catalog,
807 schema.clone(),
808 requests,
809 ctx.write_rows_to_admit(catalog, &schema, count_insert_rows(&requests)?),
810 ctx.channel() as u8
811 )
812 .await
813 .context(WriteRejectedSnafu)?;
814 let request_factory = RegionRequestFactory::new(RegionRequestHeader {
815 tracing_context: TracingContext::from_current_span().to_w3c(),
816 dbname: ctx.get_db_string(),
817 ..Default::default()
818 });
819
820 let InstantAndNormalInsertRequests {
821 normal_requests,
822 instant_requests,
823 } = requests;
824
825 let flow_mirror_task = FlowMirrorTask::new(
827 &self.table_flownode_set_cache,
828 normal_requests
829 .requests
830 .iter()
831 .chain(instant_requests.requests.iter()),
832 )
833 .await?;
834 let has_instant_rows = instant_requests.requests.iter().any(|request| {
835 request
836 .rows
837 .as_ref()
838 .is_some_and(|rows| !rows.rows.is_empty())
839 });
840 flow_mirror_task.detach(
841 self.node_manager.clone(),
842 self.mirror_pending_rows.clone(),
843 has_instant_rows,
844 )?;
845
846 #[cfg(feature = "enterprise")]
849 if flow_write::is_child(ctx) {
850 let dispatcher = self.write_dispatcher.clone();
851 let request = request_factory.build_insert(normal_requests);
852 let writes = dispatcher.prepare(vec![request]).await?;
853 let tracked_tasks = writes.into_iter().map(|(peer, mut request)| {
854 let mut query_context = crate::utils::to_meta_query_context(ctx.clone());
855 query_context.extensions.insert(
856 common_grpc::flow_write::CHILD_KEY.to_string(),
857 "1".to_string(),
858 );
859 if let Some(header) = request.header.as_mut() {
860 header.query_context = Some(query_context.into());
861 }
862 let node_manager = self.node_manager.clone();
863 async move { flow_write::dispatch(&node_manager, &peer, request).await }
864 });
865 let results = future::join_all(tracked_tasks).await;
866 let mut affected_rows = 0;
867 let mut first_error = None;
868 for result in results {
869 match result {
870 Ok(response) => affected_rows += response.affected_rows,
871 Err(error) if first_error.is_none() => first_error = Some(error),
872 Err(_) => {}
873 }
874 }
875 if let Some(error) = first_error {
876 return Err(error);
877 }
878 crate::metrics::DIST_INGEST_ROW_COUNT
879 .with_label_values(&[ctx.get_db_string().as_str()])
880 .inc_by(affected_rows as u64);
881 return Ok(Output::new(
882 OutputData::AffectedRows(affected_rows),
883 OutputMeta::new_with_cost(write_cost as _),
884 ));
885 }
886
887 let dispatcher = self.write_dispatcher.clone();
888 let request = request_factory.build_insert(normal_requests);
889 let writes = dispatcher.prepare(vec![request]).await?;
890 let cancellation = CancellationToken::new();
892 let _guard = cancellation.clone().drop_guard();
893 let affected_rows = common_runtime::spawn_global(async move {
894 dispatcher.execute(writes, None, Some(&cancellation)).await
895 })
896 .await
897 .context(JoinTaskSnafu)?
898 .context(RequestInsertsSnafu)?
899 .affected_rows;
900 crate::metrics::DIST_INGEST_ROW_COUNT
901 .with_label_values(&[ctx.get_db_string().as_str()])
902 .inc_by(affected_rows as u64);
903 Ok(Output::new(
904 OutputData::AffectedRows(affected_rows),
905 OutputMeta::new_with_cost(write_cost as _),
906 ))
907 }
908
909 fn auto_create_disabled_reason(&self, ctx: &QueryContextRef) -> Result<Option<&'static str>> {
913 let auto_create_table_hint = ctx
914 .extension(AUTO_CREATE_TABLE_KEY)
915 .map(|v| v.parse::<bool>())
916 .transpose()
917 .map_err(|_| {
918 InvalidInsertRequestSnafu {
919 reason: "`auto_create_table` hint must be a boolean",
920 }
921 .build()
922 })?
923 .unwrap_or(true);
924 Ok(if !self.auto_create_table {
925 Some("auto-create table is disabled by frontend config")
926 } else if !auto_create_table_hint {
927 Some("`auto_create_table` hint is disabled")
928 } else {
929 None
930 })
931 }
932
933 fn is_auto_create_exempt_private_table(schema: &str, table: &str) -> bool {
936 schema == DEFAULT_PRIVATE_SCHEMA_NAME
937 && matches!(
938 table,
939 DEFAULT_EVENTS_TABLE_NAME | SLOW_QUERY_TABLE_NAME | REGION_STATS_HISTORY_TABLE_NAME
940 )
941 }
942
943 pub async fn ensure_bulk_insert_schema(
947 &self,
948 table: TableRef,
949 request_schema: &ArrowSchema,
950 ctx: &QueryContextRef,
951 statement_executor: &StatementExecutor,
952 ) -> Result<TableRef> {
953 let table_info = table.table_info();
954 if self.auto_create_disabled_reason(ctx)?.is_some()
955 && !Self::is_auto_create_exempt_private_table(&table_info.schema_name, &table_info.name)
956 {
957 return Ok(table);
958 }
959
960 let table_schema = table.schema();
961 let schema = request_schema
962 .fields()
963 .iter()
964 .filter(|field| table_schema.column_schema_by_name(field.name()).is_none())
965 .map(|field| {
966 let data_type = field.data_type();
967 ensure!(
970 !data_type.is_nested() && !matches!(data_type, ArrowDataType::Dictionary(..)),
971 crate::error::NotSupportedSnafu {
972 feat: format!(
973 "automatically adding bulk insert column '{}' with type {:?}",
974 field.name(),
975 data_type
976 ),
977 }
978 );
979 let column = datatypes::schema::ColumnSchema::try_from(field.as_ref())
980 .context(crate::error::ConvertSchemaSnafu)?;
981 let column_def =
984 try_as_column_def(&column, false).context(crate::error::ColumnDataTypeSnafu)?;
985 Ok(ColumnSchema {
986 column_name: column_def.name,
987 datatype: column_def.data_type,
988 semantic_type: column_def.semantic_type,
989 datatype_extension: column_def.datatype_extension,
990 options: column_def.options,
991 })
992 })
993 .collect::<Result<Vec<_>>>()?;
994 let mut request = RowInsertRequest {
995 table_name: table_info.name.clone(),
996 rows: Some(Rows {
997 schema,
998 rows: Vec::new(),
999 }),
1000 };
1001 let Some(alter_expr) =
1002 self.get_alter_table_expr_on_demand(&mut request, &table, ctx, false, false, true)?
1003 else {
1004 return Ok(table);
1005 };
1006
1007 statement_executor
1008 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1009 .await?;
1010 self.get_table(
1011 &table_info.catalog_name,
1012 &table_info.schema_name,
1013 &table_info.name,
1014 )
1015 .await?
1016 .with_context(|| TableNotFoundSnafu {
1017 table_name: table_info.full_table_name(),
1018 })
1019 }
1020
1021 pub async fn ensure_trace_table_on_demand(
1026 &self,
1027 table_name: &str,
1028 request_schema: Vec<ColumnSchema>,
1029 alter_existing: bool,
1030 ctx: &QueryContextRef,
1031 statement_executor: &StatementExecutor,
1032 ) -> Result<()> {
1033 let mut requests = RowInsertRequests {
1034 inserts: vec![RowInsertRequest {
1035 table_name: table_name.to_string(),
1036 rows: Some(api::v1::Rows {
1037 schema: request_schema,
1038 rows: Vec::new(),
1039 }),
1040 }],
1041 };
1042 self.create_or_alter_tables_on_demand(
1043 &mut requests,
1044 ctx,
1045 AutoCreateTableType::Trace { alter_existing },
1046 statement_executor,
1047 false,
1048 false,
1049 None,
1050 )
1051 .await?;
1052 Ok(())
1053 }
1054
1055 #[allow(clippy::too_many_arguments)]
1074 async fn create_or_alter_tables_on_demand(
1075 &self,
1076 requests: &mut RowInsertRequests,
1077 ctx: &QueryContextRef,
1078 auto_create_table_type: AutoCreateTableType,
1079 statement_executor: &StatementExecutor,
1080 accommodate_existing_schema: bool,
1081 is_single_value: bool,
1082 align_time_index_unit: Option<TimeUnit>,
1083 ) -> Result<CreateAlterTableResult> {
1084 let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND
1085 .with_label_values(&[auto_create_table_type.as_str()])
1086 .start_timer();
1087 let catalog = ctx.current_catalog();
1088 let schema = ctx.current_schema();
1089
1090 let auto_create_disabled_reason = self.auto_create_disabled_reason(ctx)?;
1091 let has_auto_create_exempt_table = auto_create_disabled_reason.is_some()
1093 && requests
1094 .inserts
1095 .iter()
1096 .any(|req| Self::is_auto_create_exempt_private_table(&schema, &req.table_name));
1097 let mut table_infos = HashMap::new();
1098 if let Some(disabled_reason) = auto_create_disabled_reason
1100 && !has_auto_create_exempt_table
1101 {
1102 let mut instant_table_ids = HashSet::new();
1103 for req in &mut requests.inserts {
1104 let table = match self.get_table(catalog, &schema, &req.table_name).await? {
1105 Some(table) => table,
1106 None if is_ddl_reserved_table(&schema, &req.table_name) => {
1109 statement_executor
1110 .create_declared_relationships_table(catalog, ctx.clone())
1111 .await?
1112 }
1113 None => {
1114 return InvalidInsertRequestSnafu {
1115 reason: format!(
1116 "Table `{}` does not exist, and {}",
1117 req.table_name, disabled_reason
1118 ),
1119 }
1120 .fail();
1121 }
1122 };
1123 if align_time_index_unit.is_some()
1127 && let Some(rows) = req.rows.as_mut()
1128 && let Some(target_unit) = table_time_index_unit(&table)
1129 {
1130 convert_rows_time_unit(rows, target_unit)?;
1131 }
1132 let table_info = table.table_info();
1133 if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1134 validate_trace_table_model(&table_info, ctx)?;
1135 }
1136 if table_info.is_ttl_instant_table() {
1137 instant_table_ids.insert(table_info.table_id());
1138 }
1139 table_infos.insert(table_info.table_id(), table.table_info());
1140 }
1141 let ret = CreateAlterTableResult {
1142 instant_table_ids,
1143 table_infos,
1144 };
1145 return Ok(ret);
1146 }
1147
1148 let mut create_tables = vec![];
1149 let mut alter_tables = vec![];
1150 let mut need_refresh_table_infos = HashSet::new();
1151 let mut instant_table_ids = HashSet::new();
1152 let mut per_table_semantics: Option<Option<PerTableSemanticIndex>> = None;
1153
1154 for req in &mut requests.inserts {
1155 let auto_create_allowed = auto_create_disabled_reason.is_none()
1157 || Self::is_auto_create_exempt_private_table(&schema, &req.table_name);
1158 match self.get_table(catalog, &schema, &req.table_name).await? {
1159 Some(table) => {
1160 let table_info = table.table_info();
1161 if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) {
1162 validate_trace_table_model(&table_info, ctx)?;
1163 }
1164 if table_info.is_ttl_instant_table() {
1165 instant_table_ids.insert(table_info.table_id());
1166 }
1167 if align_time_index_unit.is_some()
1171 && let Some(rows) = req.rows.as_mut()
1172 && let Some(target_unit) = table_time_index_unit(&table)
1173 {
1174 convert_rows_time_unit(rows, target_unit)?;
1175 }
1176 if auto_create_allowed
1177 && let Some(alter_expr) = self.get_alter_table_expr_on_demand(
1178 req,
1179 &table,
1180 ctx,
1181 accommodate_existing_schema,
1182 is_single_value,
1183 auto_create_table_type.alter_existing(),
1184 )?
1185 {
1186 alter_tables.push(alter_expr);
1187 need_refresh_table_infos.insert((
1188 catalog.to_string(),
1189 schema.clone(),
1190 req.table_name.clone(),
1191 ));
1192 } else {
1193 table_infos.insert(table_info.table_id(), table.table_info());
1194 }
1195 }
1196 None if is_ddl_reserved_table(&schema, &req.table_name) => {
1200 let table = statement_executor
1201 .create_declared_relationships_table(catalog, ctx.clone())
1202 .await?;
1203 let table_info = table.table_info();
1204 if table_info.is_ttl_instant_table() {
1205 instant_table_ids.insert(table_info.table_id());
1206 }
1207 table_infos.insert(table_info.table_id(), table_info);
1208 }
1209 None if !auto_create_allowed
1210 && let Some(disabled_reason) = auto_create_disabled_reason =>
1211 {
1212 return InvalidInsertRequestSnafu {
1213 reason: format!(
1214 "Table `{}` does not exist, and {}",
1215 req.table_name, disabled_reason,
1216 ),
1217 }
1218 .fail();
1219 }
1220 None => {
1221 if let Some(physical_unit) = align_time_index_unit
1225 && let Some(rows) = req.rows.as_mut()
1226 {
1227 convert_rows_time_unit(rows, physical_unit)?;
1228 }
1229 let semantic_index = per_table_semantics
1230 .get_or_insert_with(|| parse_per_table_semantic_index(ctx))
1231 .as_ref();
1232 let create_expr = self.get_create_table_expr_on_demand(
1233 req,
1234 &auto_create_table_type,
1235 ctx,
1236 semantic_index,
1237 )?;
1238 create_tables.push(create_expr);
1239 }
1240 }
1241 }
1242
1243 match auto_create_table_type {
1244 AutoCreateTableType::Logical(_) => {
1245 if !create_tables.is_empty() {
1246 let tables = self
1248 .create_logical_tables(create_tables, ctx, statement_executor)
1249 .await?;
1250
1251 for table in tables {
1252 let table_info = table.table_info();
1253 if table_info.is_ttl_instant_table() {
1254 instant_table_ids.insert(table_info.table_id());
1255 }
1256 table_infos.insert(table_info.table_id(), table.table_info());
1257 }
1258 }
1259 if !alter_tables.is_empty() {
1260 statement_executor
1262 .alter_logical_tables(alter_tables, ctx.clone(), TriggerReason::AutoAlter)
1263 .await?;
1264 }
1265 }
1266 AutoCreateTableType::Physical
1267 | AutoCreateTableType::Log
1268 | AutoCreateTableType::LastNonNull => {
1269 for create_table in create_tables {
1272 let table = self
1273 .create_physical_table(create_table, None, ctx, statement_executor)
1274 .await?;
1275 let table_info = table.table_info();
1276 if table_info.is_ttl_instant_table() {
1277 instant_table_ids.insert(table_info.table_id());
1278 }
1279 table_infos.insert(table_info.table_id(), table.table_info());
1280 }
1281 for alter_expr in alter_tables.into_iter() {
1282 statement_executor
1283 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1284 .await?;
1285 }
1286 }
1287
1288 AutoCreateTableType::Trace { .. } => {
1289 let trace_table_name = ctx
1290 .extension(TRACE_TABLE_NAME_SESSION_KEY)
1291 .unwrap_or(TRACE_TABLE_NAME);
1292
1293 let trace_table_partitions = if let Some(trace_table_partitions) =
1294 ctx.extension(TRACE_TABLE_PARTITIONS_HINT_KEY)
1295 {
1296 let p = trace_table_partitions.parse::<u32>().map_err(|_| {
1297 InvalidInsertRequestSnafu {
1298 reason: format!(
1299 "Failed to parse trace_table_partitions: {}",
1300 trace_table_partitions
1301 ),
1302 }
1303 .build()
1304 })?;
1305 Some(p)
1306 } else {
1307 None
1308 };
1309
1310 for mut create_table in create_tables {
1313 if create_table.table_name == trace_services_table_name(trace_table_name)
1314 || create_table.table_name == trace_operations_table_name(trace_table_name)
1315 {
1316 create_table
1318 .table_options
1319 .insert(APPEND_MODE_KEY.to_string(), "false".to_string());
1320 create_table.table_options.remove(TTL_KEY);
1322
1323 let table = self
1324 .create_physical_table(create_table, None, ctx, statement_executor)
1325 .await?;
1326 let table_info = table.table_info();
1327 if table_info.is_ttl_instant_table() {
1328 instant_table_ids.insert(table_info.table_id());
1329 }
1330 table_infos.insert(table_info.table_id(), table.table_info());
1331 } else {
1332 let partitions = if matches!(trace_table_partitions, Some(0) | Some(1)) {
1335 None
1337 } else {
1338 let p = partition_rule_for_hexstring(
1339 TRACE_ID_COLUMN,
1340 trace_table_partitions,
1341 )
1342 .context(CreatePartitionRulesSnafu)?;
1343 Some(p)
1344 };
1345
1346 let index_columns =
1351 [TRACE_ID_COLUMN, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN];
1352 for index_column in index_columns {
1353 if let Some(col) = create_table
1354 .column_defs
1355 .iter_mut()
1356 .find(|c| c.name == index_column)
1357 {
1358 col.options =
1359 options_from_skipping(&SkippingIndexOptions::default())
1360 .context(ColumnOptionsSnafu)?;
1361 } else {
1362 warn!(
1363 "Column {} not found when creating index for trace table: {}.",
1364 index_column, create_table.table_name
1365 );
1366 }
1367 }
1368
1369 let id_compression = CompressionOptions {
1373 encoding: Some(ColumnEncoding::Plain),
1374 page_size: Some(ReadableSize::kb(64)),
1375 ..Default::default()
1376 };
1377 let compressed_columns = [
1378 (SPAN_ID_COLUMN, id_compression),
1379 (PARENT_SPAN_ID_COLUMN, id_compression),
1380 (
1381 TIMESTAMP_END_COLUMN,
1382 CompressionOptions {
1383 encoding: Some(ColumnEncoding::DeltaBinaryPacked),
1384 ..Default::default()
1385 },
1386 ),
1387 ];
1388 for (name, compression) in compressed_columns {
1389 if let Some(col) =
1390 create_table.column_defs.iter_mut().find(|c| c.name == name)
1391 {
1392 let compression = options_from_compression(&compression)
1393 .context(ColumnOptionsSnafu)?;
1394 col.options
1395 .get_or_insert_with(Default::default)
1396 .options
1397 .extend(compression.options);
1398 }
1399 }
1400
1401 create_table.table_options.insert(
1403 TABLE_DATA_MODEL.to_string(),
1404 ctx.extension(SEMANTIC_PIPELINE)
1405 .unwrap_or(TABLE_DATA_MODEL_TRACE_V1)
1406 .to_string(),
1407 );
1408
1409 let table = self
1410 .create_physical_table(
1411 create_table,
1412 partitions,
1413 ctx,
1414 statement_executor,
1415 )
1416 .await?;
1417 let table_info = table.table_info();
1418 if table_info.is_ttl_instant_table() {
1419 instant_table_ids.insert(table_info.table_id());
1420 }
1421 table_infos.insert(table_info.table_id(), table.table_info());
1422 }
1423 }
1424 for alter_expr in alter_tables.into_iter() {
1425 statement_executor
1426 .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter)
1427 .await?;
1428 }
1429 }
1430 }
1431
1432 for (catalog, schema, table_name) in need_refresh_table_infos {
1434 let table = self
1435 .get_table(&catalog, &schema, &table_name)
1436 .await?
1437 .context(TableNotFoundSnafu {
1438 table_name: common_catalog::format_full_table_name(
1439 &catalog,
1440 &schema,
1441 &table_name,
1442 ),
1443 })?;
1444 let table_info = table.table_info();
1445 table_infos.insert(table_info.table_id(), table.table_info());
1446 }
1447
1448 Ok(CreateAlterTableResult {
1449 instant_table_ids,
1450 table_infos,
1451 })
1452 }
1453
1454 async fn create_physical_table_on_demand(
1455 &self,
1456 ctx: &QueryContextRef,
1457 physical_table: String,
1458 statement_executor: &StatementExecutor,
1459 ) -> Result<TableRef> {
1460 let catalog_name = ctx.current_catalog();
1461 let schema_name = ctx.current_schema();
1462
1463 if let Some(table) = self
1465 .get_table(catalog_name, &schema_name, &physical_table)
1466 .await?
1467 {
1468 return Ok(table);
1469 }
1470
1471 if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? {
1473 return InvalidInsertRequestSnafu {
1474 reason: format!(
1475 "Physical table `{physical_table}` does not exist, and {disabled_reason}"
1476 ),
1477 }
1478 .fail();
1479 }
1480
1481 let table_reference = TableReference::full(catalog_name, &schema_name, &physical_table);
1482 debug!("Ensuring physical metric table `{table_reference}` exists for insert");
1483
1484 let default_schema = vec![
1486 ColumnSchema {
1487 column_name: greptime_timestamp().to_string(),
1488 datatype: ColumnDataType::TimestampMillisecond as _,
1489 semantic_type: SemanticType::Timestamp as _,
1490 datatype_extension: None,
1491 options: None,
1492 },
1493 ColumnSchema {
1494 column_name: greptime_value().to_string(),
1495 datatype: ColumnDataType::Float64 as _,
1496 semantic_type: SemanticType::Field as _,
1497 datatype_extension: None,
1498 options: None,
1499 },
1500 ];
1501 let create_table_expr =
1502 &mut build_create_table_expr(&table_reference, &default_schema, default_engine())?;
1503
1504 create_table_expr.engine = METRIC_ENGINE_NAME.to_string();
1505 create_table_expr
1506 .table_options
1507 .insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), "true".to_string());
1508
1509 let res = statement_executor
1511 .create_table_inner(
1512 create_table_expr,
1513 None,
1514 ctx.clone(),
1515 TriggerReason::AutoCreate,
1516 )
1517 .await;
1518
1519 match res {
1520 Ok(table) => Ok(table),
1521 Err(err) => {
1522 error!(err; "Failed to create table {table_reference}");
1523 Err(err)
1524 }
1525 }
1526 }
1527
1528 async fn get_table(
1529 &self,
1530 catalog: &str,
1531 schema: &str,
1532 table: &str,
1533 ) -> Result<Option<TableRef>> {
1534 self.catalog_manager
1535 .table(catalog, schema, table, None)
1536 .await
1537 .context(CatalogSnafu)
1538 }
1539
1540 fn get_create_table_expr_on_demand(
1541 &self,
1542 req: &RowInsertRequest,
1543 create_type: &AutoCreateTableType,
1544 ctx: &QueryContextRef,
1545 semantic_index: Option<&PerTableSemanticIndex>,
1546 ) -> Result<CreateTableExpr> {
1547 let schema = ctx.current_schema();
1548 let mut table_options = std::collections::HashMap::with_capacity(4);
1549 fill_table_options_for_create(&mut table_options, create_type, ctx);
1550 apply_per_table_semantic_options(
1551 &mut table_options,
1552 semantic_index,
1553 ctx.current_schema().as_str(),
1554 &req.table_name,
1555 );
1556
1557 let engine_name = if let AutoCreateTableType::Logical(_) = create_type {
1558 METRIC_ENGINE_NAME
1560 } else {
1561 default_engine()
1562 };
1563
1564 let table_ref = TableReference::full(ctx.current_catalog(), &schema, &req.table_name);
1565 let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1567 let mut create_table_expr =
1568 build_create_table_expr(&table_ref, request_schema, engine_name)?;
1569
1570 if ctx.extension(SPLUNK_PK_METADATA_ORDER_KEY).is_some() {
1572 reorder_splunk_primary_keys(&mut create_table_expr.primary_keys);
1573 }
1574
1575 debug!("Ensuring table `{table_ref}` exists for insert");
1576 create_table_expr.table_options.extend(table_options);
1577 Ok(create_table_expr)
1578 }
1579
1580 fn get_alter_table_expr_on_demand(
1588 &self,
1589 req: &mut RowInsertRequest,
1590 table: &TableRef,
1591 ctx: &QueryContextRef,
1592 accommodate_existing_schema: bool,
1593 is_single_value: bool,
1594 alter_existing: bool,
1595 ) -> Result<Option<AlterTableExpr>> {
1596 if !alter_existing {
1597 return Ok(None);
1598 }
1599
1600 let catalog_name = ctx.current_catalog();
1601 let schema_name = ctx.current_schema();
1602 let table_name = table.table_info().name.clone();
1603
1604 if is_ddl_reserved_table(&schema_name, &table_name) {
1607 return Ok(None);
1608 }
1609
1610 let request_schema = req.rows.as_ref().unwrap().schema.as_slice();
1611 let request_field_count = request_schema
1612 .iter()
1613 .filter(|col| col.semantic_type == SemanticType::Field as i32)
1614 .count();
1615 let column_exprs = ColumnExpr::from_column_schemas(request_schema);
1616 let add_columns = expr_helper::extract_add_columns_expr(&table.schema(), column_exprs)?;
1617 let Some(mut add_columns) = add_columns else {
1618 return Ok(None);
1619 };
1620
1621 if is_single_value {
1622 let request_is_native_histogram = request_is_native_histogram(request_schema);
1623 let table_is_native_histogram = table_is_native_histogram(table);
1624 ensure!(
1625 request_is_native_histogram == table_is_native_histogram,
1626 InvalidInsertRequestSnafu {
1627 reason: format!(
1628 "Table `{table_name}` cannot mix native histogram and float sample fields"
1629 ),
1630 }
1631 );
1632 }
1633
1634 if accommodate_existing_schema {
1636 let table_schema = table.schema();
1637 let ts_col_name = table_schema.timestamp_column().map(|c| c.name.clone());
1639 let mut field_col_name = None;
1641 if is_single_value && request_field_count <= 1 {
1642 let mut multiple_field_cols = false;
1643 table.field_columns().for_each(|col| {
1644 if field_col_name.is_none() {
1645 field_col_name = Some(col.name.clone());
1646 } else {
1647 multiple_field_cols = true;
1648 }
1649 });
1650 if multiple_field_cols {
1651 field_col_name = None;
1652 }
1653 }
1654
1655 if let Some(rows) = req.rows.as_mut() {
1657 for col in &mut rows.schema {
1658 match col.semantic_type {
1659 x if x == SemanticType::Timestamp as i32 => {
1660 if let Some(ref ts_name) = ts_col_name
1661 && col.column_name != *ts_name
1662 {
1663 col.column_name = ts_name.clone();
1664 }
1665 }
1666 x if x == SemanticType::Field as i32 => {
1667 if let Some(ref field_name) = field_col_name
1668 && col.column_name != *field_name
1669 {
1670 col.column_name = field_name.clone();
1671 }
1672 }
1673 _ => {}
1674 }
1675 }
1676 }
1677
1678 add_columns.add_columns.retain(|col| {
1680 let def = col.column_def.as_ref().unwrap();
1681 def.semantic_type == SemanticType::Tag as i32
1682 || (def.semantic_type == SemanticType::Field as i32 && field_col_name.is_none())
1683 });
1684
1685 if add_columns.add_columns.is_empty() {
1686 return Ok(None);
1687 }
1688 }
1689
1690 Ok(Some(AlterTableExpr {
1691 catalog_name: catalog_name.to_string(),
1692 schema_name: schema_name.clone(),
1693 table_name: table_name.clone(),
1694 kind: Some(Kind::AddColumns(add_columns)),
1695 }))
1696 }
1697
1698 async fn create_physical_table(
1700 &self,
1701 mut create_table_expr: CreateTableExpr,
1702 partitions: Option<Partitions>,
1703 ctx: &QueryContextRef,
1704 statement_executor: &StatementExecutor,
1705 ) -> Result<TableRef> {
1706 let res = statement_executor
1707 .create_table_inner(
1708 &mut create_table_expr,
1709 partitions,
1710 ctx.clone(),
1711 TriggerReason::AutoCreate,
1712 )
1713 .await;
1714
1715 let table_ref = TableReference::full(
1716 &create_table_expr.catalog_name,
1717 &create_table_expr.schema_name,
1718 &create_table_expr.table_name,
1719 );
1720
1721 match res {
1722 Ok(table) => {
1723 validate_trace_table_model(&table.table_info(), ctx)?;
1724 Ok(table)
1725 }
1726 Err(err) => {
1727 error!(err; "Failed to create table {}", table_ref);
1728 Err(err)
1729 }
1730 }
1731 }
1732
1733 async fn create_logical_tables(
1734 &self,
1735 create_table_exprs: Vec<CreateTableExpr>,
1736 ctx: &QueryContextRef,
1737 statement_executor: &StatementExecutor,
1738 ) -> Result<Vec<TableRef>> {
1739 let res = statement_executor
1740 .create_logical_tables(&create_table_exprs, ctx.clone(), TriggerReason::AutoCreate)
1741 .await;
1742
1743 match res {
1744 Ok(res) => Ok(res),
1745 Err(err) => {
1746 let failed_tables = create_table_exprs
1747 .into_iter()
1748 .map(|expr| {
1749 format!(
1750 "{}.{}.{}",
1751 expr.catalog_name, expr.schema_name, expr.table_name
1752 )
1753 })
1754 .collect::<Vec<_>>();
1755 error!(
1756 err;
1757 "Failed to create logical tables {:?}",
1758 failed_tables
1759 );
1760 Err(err)
1761 }
1762 }
1763 }
1764
1765 pub fn write_dispatcher(&self) -> &WriteDispatcherRef {
1767 &self.write_dispatcher
1768 }
1769
1770 pub fn node_manager(&self) -> &NodeManagerRef {
1771 &self.node_manager
1772 }
1773
1774 pub fn partition_manager(&self) -> &PartitionRuleManagerRef {
1775 &self.partition_manager
1776 }
1777
1778 pub fn table_flownode_set_cache(&self) -> &TableFlownodeSetCacheRef {
1779 &self.table_flownode_set_cache
1780 }
1781}
1782
1783fn request_is_native_histogram(request_schema: &[ColumnSchema]) -> bool {
1784 let mut fields = request_schema
1785 .iter()
1786 .filter(|col| col.semantic_type == SemanticType::Field as i32);
1787 let Some(col) = fields.next() else {
1788 return false;
1789 };
1790
1791 fields.next().is_none()
1792 && api::helper::is_column_type_value_eq(
1793 col.datatype,
1794 col.datatype_extension.clone(),
1795 native_histogram_value_type(),
1796 )
1797}
1798
1799fn table_time_index_unit(table: &TableRef) -> Option<TimeUnit> {
1803 table
1804 .table_info()
1805 .meta
1806 .schema
1807 .timestamp_column()
1808 .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit()))
1809}
1810
1811fn convert_rows_time_unit(rows: &mut Rows, target_unit: TimeUnit) -> Result<()> {
1812 let Some(ts_index) = rows
1813 .schema
1814 .iter()
1815 .position(|col| col.semantic_type == SemanticType::Timestamp as i32)
1816 else {
1817 return Ok(());
1818 };
1819 let Some(source_unit) = ColumnDataType::try_from(rows.schema[ts_index].datatype)
1820 .ok()
1821 .and_then(api::helper::timestamp_unit)
1822 else {
1823 return Ok(());
1824 };
1825 if source_unit == target_unit {
1826 return Ok(());
1827 }
1828
1829 rows.schema[ts_index].datatype = api::helper::timestamp_datatype(target_unit) as i32;
1830 rows.schema[ts_index].datatype_extension = None;
1832
1833 for row in &mut rows.rows {
1841 debug_assert_eq!(row.values.len(), rows.schema.len());
1842 let value = &mut row.values[ts_index];
1843 let Some(value_data) = value.value_data.take() else {
1844 continue;
1845 };
1846 value.value_data =
1847 convert_timestamp_value_data(value_data, source_unit, target_unit, ts_index)?;
1848 }
1849 Ok(())
1850}
1851
1852fn convert_timestamp_value_data(
1853 value_data: ValueData,
1854 source_unit: TimeUnit,
1855 target_unit: TimeUnit,
1856 column_index: usize,
1857) -> Result<Option<ValueData>> {
1858 let timestamp = match value_data {
1859 ValueData::TimestampSecondValue(v) => Timestamp::new_second(v),
1860 ValueData::TimestampMillisecondValue(v) => Timestamp::new_millisecond(v),
1861 ValueData::TimestampMicrosecondValue(v) => Timestamp::new_microsecond(v),
1862 ValueData::TimestampNanosecondValue(v) => Timestamp::new_nanosecond(v),
1863 other => return Ok(Some(other)),
1865 };
1866 let converted = timestamp
1867 .convert_to(target_unit)
1868 .with_context(|| InvalidInsertRequestSnafu {
1869 reason: format!(
1870 "timestamp column {column_index} value {} in unit {source_unit:?} overflows when converting to unit {target_unit:?}",
1871 timestamp.value()
1872 ),
1873 })?;
1874 Ok(api::helper::to_grpc_value(datatypes::value::Value::Timestamp(converted)).value_data)
1875}
1876
1877fn table_is_native_histogram(table: &TableRef) -> bool {
1878 let mut fields = table.field_columns();
1879 let Some(col) = fields.next() else {
1880 return false;
1881 };
1882
1883 fields.next().is_none() && is_native_histogram_value_type(&col.data_type)
1884}
1885
1886fn validate_column_count_match(requests: &RowInsertRequests) -> Result<()> {
1887 for request in &requests.inserts {
1888 let rows = request.rows.as_ref().unwrap();
1889 let column_count = rows.schema.len();
1890 rows.rows.iter().try_for_each(|r| {
1891 ensure!(
1892 r.values.len() == column_count,
1893 InvalidInsertRequestSnafu {
1894 reason: format!(
1895 "column count mismatch, columns: {}, values: {}",
1896 column_count,
1897 r.values.len()
1898 )
1899 }
1900 );
1901 Ok(())
1902 })?;
1903 }
1904 Ok(())
1905}
1906
1907pub fn validate_trace_table_model(table_info: &TableInfo, ctx: &QueryContextRef) -> Result<()> {
1910 let Some(expected @ (TABLE_DATA_MODEL_TRACE_V1 | TABLE_DATA_MODEL_TRACE_V2)) =
1911 ctx.extension(SEMANTIC_PIPELINE)
1912 else {
1913 return Ok(());
1914 };
1915 if let Some(actual) = table_info.meta.options.data_model() {
1916 ensure!(
1917 actual == expected,
1918 InvalidInsertRequestSnafu {
1919 reason: format!(
1920 "Trace table `{}` uses {actual}, but the request uses {expected}",
1921 table_info.name,
1922 ),
1923 }
1924 );
1925 }
1926 Ok(())
1927}
1928
1929pub fn fill_table_options_for_create(
1931 table_options: &mut std::collections::HashMap<String, String>,
1932 create_type: &AutoCreateTableType,
1933 ctx: &QueryContextRef,
1934) {
1935 for key in VALID_TABLE_OPTION_KEYS {
1936 if let Some(value) = ctx.extension(key) {
1937 table_options.insert(key.to_string(), value.to_string());
1938 }
1939 }
1940
1941 if !matches!(create_type, AutoCreateTableType::Logical(_)) {
1942 let mut has_twcs_option = false;
1943 for key in COMPACTION_OPTION_KEYS {
1944 if let Some(value) = ctx.extension(key) {
1945 has_twcs_option |= key.starts_with(COMPACTION_TWCS_PREFIX);
1946 table_options.insert(key.to_string(), value.to_string());
1947 }
1948 }
1949 if has_twcs_option {
1950 table_options
1951 .entry(COMPACTION_TYPE.to_string())
1952 .or_insert_with(|| COMPACTION_TYPE_TWCS.to_string());
1953 }
1954 }
1955
1956 for (key, value) in ctx.extensions() {
1958 if is_semantic_option_key(&key) && validate_semantic_option(&key, &value) {
1959 table_options.insert(key, value);
1960 }
1961 }
1962
1963 match create_type {
1964 AutoCreateTableType::Logical(physical_table) => {
1965 table_options.insert(
1966 LOGICAL_TABLE_METADATA_KEY.to_string(),
1967 physical_table.clone(),
1968 );
1969 }
1970 AutoCreateTableType::Physical => {
1971 if let Some(append_mode) = ctx.extension(APPEND_MODE_KEY) {
1972 table_options.insert(APPEND_MODE_KEY.to_string(), append_mode.to_string());
1973 }
1974 if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
1975 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
1976 }
1977 }
1978 AutoCreateTableType::Log => {
1981 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1982 }
1983 AutoCreateTableType::LastNonNull => {
1984 if ctx
1985 .extension(APPEND_MODE_KEY)
1986 .is_some_and(|value| value.eq_ignore_ascii_case("true"))
1987 {
1988 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1989 table_options.insert(MERGE_MODE_KEY.to_string(), "last_row".to_string());
1990 } else if let Some(merge_mode) = ctx.extension(MERGE_MODE_KEY) {
1991 table_options.insert(MERGE_MODE_KEY.to_string(), merge_mode.to_string());
1992 } else {
1993 table_options.insert(MERGE_MODE_KEY.to_string(), "last_non_null".to_string());
1994 }
1995 }
1996 AutoCreateTableType::Trace { .. } => {
1997 table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string());
1998 }
1999 }
2000}
2001
2002pub type PerTableSemanticIndex = BTreeMap<String, BTreeMap<String, BTreeMap<String, String>>>;
2008
2009pub fn parse_per_table_semantic_index(ctx: &QueryContextRef) -> Option<PerTableSemanticIndex> {
2014 let raw = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY)?;
2015 match serde_json::from_str(raw) {
2016 Ok(index) => Some(index),
2017 Err(_) => {
2018 warn!("failed to parse semantic per-table index, skipping per-table options");
2019 None
2020 }
2021 }
2022}
2023
2024pub fn apply_per_table_semantic_options(
2032 table_options: &mut std::collections::HashMap<String, String>,
2033 index: Option<&PerTableSemanticIndex>,
2034 schema: &str,
2035 table_name: &str,
2036) {
2037 let Some(entry) = index
2038 .and_then(|index| index.get(schema))
2039 .and_then(|tables| tables.get(table_name))
2040 else {
2041 return;
2042 };
2043 for (key, value) in entry {
2044 if is_semantic_option_key(key) && validate_semantic_option(key, value) {
2045 table_options.insert(key.clone(), value.clone());
2046 }
2047 }
2048}
2049
2050pub fn build_create_table_expr(
2051 table: &TableReference,
2052 request_schema: &[ColumnSchema],
2053 engine: &str,
2054) -> Result<CreateTableExpr> {
2055 expr_helper::create_table_expr_by_column_schemas(table, request_schema, engine, None)
2056}
2057
2058pub const SPLUNK_PK_METADATA_ORDER_KEY: &str = "splunk_pk_metadata_order";
2062
2063fn reorder_splunk_primary_keys(primary_keys: &mut [String]) {
2066 const LEAD: [&str; 3] = ["host", "source", "sourcetype"];
2067 primary_keys.sort_by_key(|name| {
2070 LEAD.iter()
2071 .position(|&lead| lead == name.as_str())
2072 .unwrap_or(LEAD.len())
2073 });
2074}
2075
2076struct CreateAlterTableResult {
2078 instant_table_ids: HashSet<TableId>,
2080 table_infos: HashMap<TableId, Arc<TableInfo>>,
2082}
2083
2084const MAX_MIRROR_PENDING_ROWS: usize = 1_000_000;
2092
2093const MIRROR_LOG_INTERVAL: Duration = Duration::from_secs(10);
2095const MIRROR_LOG_NEVER_REPORTED: u64 = u64::MAX;
2096
2097struct MirrorLog {
2098 last_log_millis: AtomicU64,
2099 events: AtomicU64,
2100}
2101
2102impl MirrorLog {
2103 const fn new() -> Self {
2104 Self {
2105 last_log_millis: AtomicU64::new(MIRROR_LOG_NEVER_REPORTED),
2106 events: AtomicU64::new(0),
2107 }
2108 }
2109
2110 fn claim_report(&self, now_millis: u64) -> Option<u64> {
2113 self.events.fetch_add(1, Ordering::Relaxed);
2114 let last = self.last_log_millis.load(Ordering::Relaxed);
2115 if last != MIRROR_LOG_NEVER_REPORTED
2116 && now_millis.saturating_sub(last) < MIRROR_LOG_INTERVAL.as_millis() as u64
2117 {
2118 return None;
2119 }
2120 if self
2121 .last_log_millis
2122 .compare_exchange(last, now_millis, Ordering::Relaxed, Ordering::Relaxed)
2123 .is_err()
2124 {
2125 return None;
2126 }
2127 Some(self.events.swap(0, Ordering::Relaxed))
2128 }
2129}
2130
2131static MIRROR_DROP_LOG: MirrorLog = MirrorLog::new();
2132static MIRROR_FAILURE_LOG: MirrorLog = MirrorLog::new();
2133
2134fn mirror_drop_log_millis() -> u64 {
2137 static START: LazyLock<Instant> = LazyLock::new(Instant::now);
2138 START.elapsed().as_millis() as u64
2139}
2140
2141struct FlowMirrorTask {
2142 requests: HashMap<Peer, RegionInsertRequests>,
2143}
2144
2145impl FlowMirrorTask {
2146 async fn new(
2147 cache: &TableFlownodeSetCacheRef,
2148 requests: impl Iterator<Item = &RegionInsertRequest>,
2149 ) -> Result<Self> {
2150 let mut src_table_reqs: HashMap<TableId, Option<(Vec<Peer>, RegionInsertRequests)>> =
2151 HashMap::new();
2152
2153 for req in requests {
2154 let table_id = RegionId::from_u64(req.region_id).table_id();
2155 match src_table_reqs.get_mut(&table_id) {
2156 Some(Some((_peers, reqs))) => reqs.requests.push(req.clone()),
2157 Some(None) => continue,
2159 _ => {
2160 let peers = cache
2162 .get(table_id)
2163 .await
2164 .context(RequestInsertsSnafu)?
2165 .unwrap_or_default()
2166 .values()
2167 .cloned()
2168 .collect::<HashSet<_>>()
2169 .into_iter()
2170 .collect::<Vec<_>>();
2171
2172 if !peers.is_empty() {
2173 let mut reqs = RegionInsertRequests::default();
2174 reqs.requests.push(req.clone());
2175 src_table_reqs.insert(table_id, Some((peers, reqs)));
2176 } else {
2177 src_table_reqs.insert(table_id, None);
2179 }
2180 }
2181 }
2182 }
2183
2184 let mut inserts: HashMap<Peer, RegionInsertRequests> = HashMap::new();
2185
2186 for (_table_id, (peers, reqs)) in src_table_reqs
2187 .into_iter()
2188 .filter_map(|(k, v)| v.map(|v| (k, v)))
2189 {
2190 if peers.len() == 1 {
2191 inserts
2193 .entry(peers[0].clone())
2194 .or_default()
2195 .requests
2196 .extend(reqs.requests);
2197 continue;
2198 } else {
2199 for flownode in peers {
2201 inserts
2202 .entry(flownode.clone())
2203 .or_default()
2204 .requests
2205 .extend(reqs.requests.clone());
2206 }
2207 }
2208 }
2209
2210 Ok(Self { requests: inserts })
2211 }
2212
2213 fn pending_rows(&self) -> u64 {
2222 self.requests.values().map(region_inserts_rows).sum()
2223 }
2224
2225 fn detach(
2226 self,
2227 node_manager: NodeManagerRef,
2228 mirror_pending_rows: Arc<AtomicU64>,
2229 has_instant_rows: bool,
2230 ) -> Result<()> {
2231 let num_rows = self.pending_rows();
2234 if num_rows == 0 {
2235 return Ok(());
2236 }
2237 if has_instant_rows && num_rows > MAX_MIRROR_PENDING_ROWS as u64 {
2238 return InvalidInsertRequestSnafu {
2239 reason: format!(
2240 "flow mirror batch has {num_rows} peer-fan-out rows, exceeding the limit of {}; reduce the batch size",
2241 MAX_MIRROR_PENDING_ROWS
2242 ),
2243 }
2244 .fail();
2245 }
2246
2247 let pending = reserve_mirror_pending_rows(&mirror_pending_rows, num_rows);
2248 if pending > MAX_MIRROR_PENDING_ROWS as u64 {
2249 release_mirror_pending_rows(&mirror_pending_rows, num_rows);
2250 if has_instant_rows {
2251 return Err(meter_core::collect::WriteRejected::new(
2252 "flow mirror pending rows limit exceeded for instant-TTL insert",
2253 ))
2254 .context(WriteRejectedSnafu);
2255 }
2256 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.inc_by(num_rows);
2257 if let Some(events) = MIRROR_DROP_LOG.claim_report(mirror_drop_log_millis()) {
2258 warn!(
2259 "Flow mirror write dropped: {} pending rows exceeds limit {}, {} drop events since last report",
2260 pending, MAX_MIRROR_PENDING_ROWS, events
2261 );
2262 }
2263 return Ok(());
2264 }
2265
2266 for (peer, inserts) in self.requests {
2267 let peer_rows = region_inserts_rows(&inserts);
2271 let node_manager = node_manager.clone();
2272 let mirror_pending_rows = mirror_pending_rows.clone();
2273 common_runtime::spawn_global(async move {
2274 let result = node_manager
2275 .flownode(&peer)
2276 .await
2277 .handle_inserts(inserts)
2278 .await
2279 .context(RequestInsertsSnafu);
2280
2281 match result {
2282 Ok(resp) => {
2283 let affected_rows = resp.affected_rows;
2284 crate::metrics::DIST_MIRROR_ROW_COUNT.inc_by(affected_rows);
2285 }
2286 Err(err) => {
2287 if let Some(events) =
2288 MIRROR_FAILURE_LOG.claim_report(mirror_drop_log_millis())
2289 {
2290 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);
2291 }
2292 }
2293 }
2294 release_mirror_pending_rows(&mirror_pending_rows, peer_rows);
2297 });
2298 }
2299
2300 Ok(())
2301 }
2302}
2303
2304fn region_inserts_rows(inserts: &RegionInsertRequests) -> u64 {
2307 inserts
2308 .requests
2309 .iter()
2310 .filter_map(|req| req.rows.as_ref())
2311 .map(|rows| rows.rows.len() as u64)
2312 .sum()
2313}
2314
2315fn reserve_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) -> u64 {
2320 let pending = mirror_pending_rows.fetch_add(rows, Ordering::Relaxed) + rows;
2321 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(rows as i64);
2324 pending
2325}
2326
2327fn release_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) {
2333 let prev = mirror_pending_rows
2334 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |pending| {
2335 Some(pending.saturating_sub(rows))
2336 })
2337 .unwrap_or_else(|prev| prev);
2338 let released = rows.min(prev);
2342 if released > 0 {
2343 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(released as i64);
2344 }
2345}
2346
2347#[cfg(test)]
2348mod tests {
2349 use std::sync::Arc;
2350 use std::sync::atomic::{AtomicU64, Ordering};
2351
2352 use api::helper::ColumnDataTypeWrapper;
2353 use api::v1::flow::FlowResponse;
2354 use api::v1::helper::{field_column_schema, time_index_column_schema};
2355 use api::v1::{Row, RowInsertRequest, Rows, Value};
2356 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
2357 use common_meta::cache::new_table_flownode_set_cache;
2358 use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler;
2359 use common_meta::test_util::{MockDatanodeManager, MockFlownodeHandler, MockFlownodeManager};
2360 use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD;
2361 use common_query::prelude::{greptime_native_histogram, set_default_prefix};
2362 use datatypes::data_type::ConcreteDataType;
2363 use datatypes::schema::ColumnSchema;
2364 use moka::future::Cache;
2365 use session::context::QueryContext;
2366 use table::TableRef;
2367 use table::dist_table::DummyDataSource;
2368 use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType};
2369
2370 use crate::insert::*;
2371 use crate::test_util::{
2372 create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
2373 };
2374
2375 fn make_table_ref_with_schema(
2376 ts_name: &str,
2377 field_name: &str,
2378 field_type: ConcreteDataType,
2379 ) -> TableRef {
2380 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2381 ColumnSchema::new(
2382 ts_name,
2383 ConcreteDataType::timestamp_millisecond_datatype(),
2384 false,
2385 )
2386 .with_time_index(true),
2387 ColumnSchema::new(field_name, field_type, true),
2388 ])
2389 .unwrap()
2390 .build()
2391 .unwrap();
2392 let meta = TableMetaBuilder::empty()
2393 .schema(Arc::new(schema))
2394 .primary_key_indices(vec![])
2395 .value_indices(vec![1])
2396 .engine("mito")
2397 .next_column_id(0)
2398 .options(Default::default())
2399 .created_on(Default::default())
2400 .build()
2401 .unwrap();
2402 let info = Arc::new(
2403 TableInfoBuilder::default()
2404 .table_id(1)
2405 .table_version(0)
2406 .name("test_table")
2407 .schema_name(DEFAULT_SCHEMA_NAME)
2408 .catalog_name(DEFAULT_CATALOG_NAME)
2409 .desc(None)
2410 .table_type(TableType::Base)
2411 .meta(meta)
2412 .build()
2413 .unwrap(),
2414 );
2415 Arc::new(table::Table::new(
2416 info,
2417 table::metadata::FilterPushDownType::Unsupported,
2418 Arc::new(DummyDataSource),
2419 ))
2420 }
2421
2422 fn make_metric_physical_table_ref_with_time_unit(unit: TimeUnit) -> TableRef {
2423 let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![
2424 ColumnSchema::new(
2425 greptime_timestamp(),
2426 ConcreteDataType::timestamp_datatype(unit),
2427 false,
2428 )
2429 .with_time_index(true),
2430 ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
2431 ])
2432 .unwrap()
2433 .build()
2434 .unwrap();
2435 let meta = TableMetaBuilder::empty()
2436 .schema(Arc::new(schema))
2437 .primary_key_indices(vec![])
2438 .value_indices(vec![1])
2439 .engine("metric")
2440 .next_column_id(0)
2441 .options(Default::default())
2442 .created_on(Default::default())
2443 .build()
2444 .unwrap();
2445 let info = Arc::new(
2446 TableInfoBuilder::default()
2447 .table_id(1)
2448 .table_version(0)
2449 .name("greptime_physical_table")
2450 .schema_name(DEFAULT_SCHEMA_NAME)
2451 .catalog_name(DEFAULT_CATALOG_NAME)
2452 .desc(None)
2453 .table_type(TableType::Base)
2454 .meta(meta)
2455 .build()
2456 .unwrap(),
2457 );
2458 Arc::new(table::Table::new(
2459 info,
2460 table::metadata::FilterPushDownType::Unsupported,
2461 Arc::new(DummyDataSource),
2462 ))
2463 }
2464
2465 fn ms_row_insert_request(timestamp_ms: i64) -> RowInsertRequest {
2466 ms_row_insert_request_named("my_metric", timestamp_ms)
2467 }
2468
2469 fn ms_row_insert_request_named(table: &str, timestamp_ms: i64) -> RowInsertRequest {
2470 RowInsertRequest {
2471 table_name: table.to_string(),
2472 rows: Some(Rows {
2473 schema: vec![
2474 time_index_column_schema(
2475 greptime_timestamp(),
2476 ColumnDataType::TimestampMillisecond,
2477 ),
2478 field_column_schema(greptime_value(), ColumnDataType::Float64),
2479 ],
2480 rows: vec![Row {
2481 values: vec![
2482 Value {
2483 value_data: Some(ValueData::TimestampMillisecondValue(timestamp_ms)),
2484 },
2485 Value {
2486 value_data: Some(ValueData::F64Value(1.0)),
2487 },
2488 ],
2489 }],
2490 }),
2491 }
2492 }
2493
2494 fn convert_row_insert_requests_time_unit(
2496 requests: &mut RowInsertRequests,
2497 target_unit: TimeUnit,
2498 ) -> Result<()> {
2499 for request in &mut requests.inserts {
2500 let Some(rows) = request.rows.as_mut() else {
2501 continue;
2502 };
2503 convert_rows_time_unit(rows, target_unit)?;
2504 }
2505 Ok(())
2506 }
2507
2508 #[test]
2509 fn test_convert_row_insert_requests_time_unit_noop_when_matching() {
2510 let mut requests = RowInsertRequests {
2511 inserts: vec![ms_row_insert_request(123)],
2512 };
2513 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Millisecond).unwrap();
2514 let rows = requests.inserts[0].rows.as_ref().unwrap();
2515 assert_eq!(
2516 rows.schema[0].datatype,
2517 ColumnDataType::TimestampMillisecond as i32
2518 );
2519 assert!(matches!(
2520 rows.rows[0].values[0].value_data,
2521 Some(ValueData::TimestampMillisecondValue(123))
2522 ));
2523 }
2524
2525 #[test]
2526 fn test_convert_row_insert_requests_time_unit_widens_losslessly() {
2527 let mut requests = RowInsertRequests {
2528 inserts: vec![ms_row_insert_request(123)],
2529 };
2530 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Microsecond).unwrap();
2531 let rows = requests.inserts[0].rows.as_ref().unwrap();
2532 assert_eq!(
2533 rows.schema[0].datatype,
2534 ColumnDataType::TimestampMicrosecond as i32
2535 );
2536 assert!(matches!(
2537 rows.rows[0].values[0].value_data,
2538 Some(ValueData::TimestampMicrosecondValue(123_000))
2539 ));
2540 }
2541
2542 #[test]
2543 fn test_convert_row_insert_requests_time_unit_truncates_on_narrowing() {
2544 let requests = |value_ns: i64| RowInsertRequests {
2547 inserts: vec![RowInsertRequest {
2548 table_name: "my_metric".to_string(),
2549 rows: Some(Rows {
2550 schema: vec![
2551 time_index_column_schema(
2552 greptime_timestamp(),
2553 ColumnDataType::TimestampNanosecond,
2554 ),
2555 field_column_schema(greptime_value(), ColumnDataType::Float64),
2556 ],
2557 rows: vec![Row {
2558 values: vec![
2559 Value {
2560 value_data: Some(ValueData::TimestampNanosecondValue(value_ns)),
2561 },
2562 Value {
2563 value_data: Some(ValueData::F64Value(1.0)),
2564 },
2565 ],
2566 }],
2567 }),
2568 }],
2569 };
2570
2571 let mut reqs = requests(123_456_789);
2572 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
2573 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2574 assert!(matches!(
2575 rows.rows[0].values[0].value_data,
2576 Some(ValueData::TimestampMicrosecondValue(123_456))
2577 ));
2578
2579 let mut reqs = requests(123_456_789);
2580 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Millisecond).unwrap();
2581 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2582 assert!(matches!(
2583 rows.rows[0].values[0].value_data,
2584 Some(ValueData::TimestampMillisecondValue(123))
2585 ));
2586
2587 let mut reqs = requests(-123_456_789);
2588 convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap();
2589 let rows = reqs.inserts[0].rows.as_ref().unwrap();
2590 assert!(matches!(
2591 rows.rows[0].values[0].value_data,
2592 Some(ValueData::TimestampMicrosecondValue(-123_457))
2593 ));
2594 }
2595
2596 #[test]
2597 fn test_convert_row_insert_requests_time_unit_overflow() {
2598 let mut requests = RowInsertRequests {
2599 inserts: vec![ms_row_insert_request(i64::MAX)],
2600 };
2601 let err =
2602 convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Nanosecond).unwrap_err();
2603 assert!(err.to_string().contains("overflows"), "{err}");
2604 }
2605
2606 #[test]
2607 fn test_table_time_index_unit() {
2608 assert_eq!(
2609 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
2610 TimeUnit::Microsecond
2611 )),
2612 Some(TimeUnit::Microsecond)
2613 );
2614 assert_eq!(
2615 table_time_index_unit(&make_metric_physical_table_ref_with_time_unit(
2616 TimeUnit::Millisecond
2617 )),
2618 Some(TimeUnit::Millisecond)
2619 );
2620 }
2621
2622 #[tokio::test]
2623 async fn test_accommodate_existing_schema_and_reject_kind_changes() {
2624 let ts_name = "my_ts";
2625 let field_name = "my_field";
2626 let table =
2627 make_table_ref_with_schema(ts_name, field_name, ConcreteDataType::float64_datatype());
2628
2629 let mut req = RowInsertRequest {
2631 table_name: "test_table".to_string(),
2632 rows: Some(Rows {
2633 schema: vec![
2634 time_index_column_schema("ts_wrong", ColumnDataType::TimestampMillisecond),
2635 field_column_schema("field_wrong", ColumnDataType::Float64),
2636 ],
2637 rows: vec![api::v1::Row {
2638 values: vec![Value::default(), Value::default()],
2639 }],
2640 }),
2641 };
2642 let ctx = Arc::new(QueryContext::with(
2643 DEFAULT_CATALOG_NAME,
2644 DEFAULT_SCHEMA_NAME,
2645 ));
2646
2647 let kv_backend = prepare_mocked_backend().await;
2648 let inserter = Inserter::new(
2649 catalog::memory::MemoryCatalogManager::new(),
2650 create_partition_rule_manager(kv_backend.clone()).await,
2651 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
2652 Arc::new(new_table_flownode_set_cache(
2653 String::new(),
2654 Cache::new(100),
2655 kv_backend.clone(),
2656 )),
2657 true,
2658 );
2659 assert!(
2661 inserter
2662 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, false, false, false)
2663 .unwrap()
2664 .is_none()
2665 );
2666 let alter_expr = inserter
2667 .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true, true)
2668 .unwrap();
2669 assert!(alter_expr.is_none());
2670
2671 let req_schema = req.rows.as_ref().unwrap().schema.clone();
2673 assert_eq!(req_schema[0].column_name, ts_name);
2674 assert_eq!(req_schema[1].column_name, field_name);
2675
2676 let (datatype, datatype_extension) =
2677 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
2678 .unwrap()
2679 .into_parts();
2680 let mut histogram_req = RowInsertRequest {
2681 table_name: "test_table".to_string(),
2682 rows: Some(Rows {
2683 schema: vec![
2684 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2685 api::v1::ColumnSchema {
2686 column_name: greptime_native_histogram().to_string(),
2687 datatype: datatype as i32,
2688 semantic_type: SemanticType::Field as i32,
2689 datatype_extension,
2690 options: None,
2691 },
2692 ],
2693 rows: vec![],
2694 }),
2695 };
2696 let error = inserter
2697 .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, false, true, true)
2698 .unwrap_err();
2699 assert!(
2700 error
2701 .to_string()
2702 .contains("cannot mix native histogram and float sample fields")
2703 );
2704
2705 let histogram_table = make_table_ref_with_schema(
2706 "ts",
2707 greptime_native_histogram(),
2708 native_histogram_value_type().clone(),
2709 );
2710 let mut sample_req = RowInsertRequest {
2711 table_name: "test_table".to_string(),
2712 rows: Some(Rows {
2713 schema: vec![
2714 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2715 field_column_schema(greptime_value(), ColumnDataType::Float64),
2716 ],
2717 rows: vec![],
2718 }),
2719 };
2720 let error = inserter
2721 .get_alter_table_expr_on_demand(
2722 &mut sample_req,
2723 &histogram_table,
2724 &ctx,
2725 false,
2726 true,
2727 true,
2728 )
2729 .unwrap_err();
2730 assert!(
2731 error
2732 .to_string()
2733 .contains("cannot mix native histogram and float sample fields")
2734 );
2735 }
2736
2737 #[test]
2738 fn test_native_histogram_detection_survives_prefix_change() {
2739 set_default_prefix(Some("custom")).unwrap();
2740 let table = make_table_ref_with_schema(
2741 "custom_timestamp",
2742 NATIVE_HISTOGRAM_FIELD,
2743 native_histogram_value_type().clone(),
2744 );
2745 let (datatype, datatype_extension) =
2746 ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone())
2747 .unwrap()
2748 .into_parts();
2749 let request_schema = [api::v1::ColumnSchema {
2750 column_name: greptime_native_histogram().to_string(),
2751 datatype: datatype as i32,
2752 semantic_type: SemanticType::Field as i32,
2753 datatype_extension,
2754 options: None,
2755 }];
2756
2757 assert!(request_is_native_histogram(&request_schema));
2758 assert!(table_is_native_histogram(&table));
2759 }
2760
2761 #[tokio::test]
2763 async fn test_write_meter_admission() {
2764 use std::cell::Cell;
2765 use std::sync::Mutex;
2766 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
2767
2768 use api::region::RegionResponse;
2769 use api::v1::region::region_request::Body;
2770 use arrow::array::{Int32Array, TimestampMillisecondArray};
2771 use arrow::record_batch::RecordBatch;
2772 use bytes::Bytes;
2773 use common_error::ext::{ErrorExt, RetryHint};
2774 use common_error::status_code::StatusCode;
2775 use common_grpc::flight::{FlightEncoder, FlightMessage};
2776 use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher;
2777 use futures::future::BoxFuture;
2778 use meter_core::ItemCalculator;
2779 use meter_core::collect::{Collect, WriteRejected};
2780 use meter_core::data::MeterRecord;
2781 use meter_core::global::global_registry;
2782 use session::context::Channel;
2783
2784 const CATALOG: &str = "write_meter_test";
2785
2786 #[derive(Default)]
2787 struct Meter {
2788 reject: AtomicBool,
2789 attempts: Mutex<Vec<MeterRecord>>,
2790 accepted_value: AtomicU64,
2791 }
2792
2793 impl Collect for Meter {
2794 fn on_write(
2795 &self,
2796 record: MeterRecord,
2797 ) -> BoxFuture<'_, std::result::Result<(), WriteRejected>> {
2798 Box::pin(async move {
2799 if record.catalog != CATALOG {
2800 return Ok(());
2801 }
2802 let value = record.value;
2803 self.attempts.lock().unwrap().push(record);
2804 if self.reject.load(Ordering::Relaxed) {
2805 return Err(WriteRejected::new("database row quota exhausted"));
2806 }
2807 self.accepted_value.fetch_add(value, Ordering::Relaxed);
2808 Ok(())
2809 })
2810 }
2811
2812 fn on_read(&self, _: MeterRecord) {}
2813 }
2814
2815 impl ItemCalculator<InstantAndNormalInsertRequests> for Meter {
2816 fn calc(&self, _: &InstantAndNormalInsertRequests) -> u64 {
2817 17
2818 }
2819 }
2820
2821 let kv_backend = prepare_mocked_backend().await;
2822 let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
2823 let (sender, mut dispatched) = tokio::sync::mpsc::channel(16);
2824 let watcher = DatanodeWatcher::new(sender).with_handler(|_, request| {
2825 let rows = match request.body.unwrap() {
2826 Body::Inserts(requests) => requests
2827 .requests
2828 .iter()
2829 .filter_map(|request| request.rows.as_ref())
2830 .map(|rows| rows.rows.len())
2831 .sum(),
2832 Body::BulkInsert(_) => 2,
2834 body => panic!("unexpected request: {body:?}"),
2835 };
2836 Ok(RegionResponse::new(rows))
2837 });
2838 let flow_cache = Cache::new(100);
2839 let inserter = Inserter::new(
2840 catalog::memory::MemoryCatalogManager::new(),
2841 partition_manager,
2842 Arc::new(MockDatanodeManager::new(watcher)),
2843 Arc::new(new_table_flownode_set_cache(
2844 String::new(),
2845 flow_cache.clone(),
2846 kv_backend,
2847 )),
2848 true,
2849 );
2850 let mut table_info = new_test_table_info(1, "table_1", [1].into_iter());
2851 table_info.catalog_name = CATALOG.to_string();
2852 table_info.schema_name = "target_db".to_string();
2853 let table_info = Arc::new(table_info);
2854 let table_infos = HashMap::from_iter([(1, table_info.clone())]);
2855 let ctx = Arc::new(QueryContext::with_channel(
2856 DEFAULT_CATALOG_NAME,
2857 DEFAULT_SCHEMA_NAME,
2858 Channel::Postgres,
2859 ));
2860 let rows_request = || {
2861 let build = |num_rows| RegionInsertRequests {
2862 requests: vec![RegionInsertRequest {
2863 region_id: RegionId::new(1, 1).as_u64(),
2864 rows: Some(Rows {
2865 schema: vec![],
2866 rows: vec![api::v1::Row { values: vec![] }; num_rows],
2867 }),
2868 ..Default::default()
2869 }],
2870 };
2871 InstantAndNormalInsertRequests {
2872 normal_requests: build(3),
2873 instant_requests: build(2),
2874 }
2875 };
2876
2877 let output = inserter
2879 .do_request(rows_request(), &table_infos, &ctx)
2880 .await
2881 .unwrap();
2882 assert_eq!(output.meta.cost, 0);
2883 assert!(matches!(output.data, OutputData::AffectedRows(3)));
2884 dispatched.try_recv().unwrap();
2885
2886 let meter = Arc::new(Meter::default());
2887 global_registry().set_collector(meter.clone());
2888 global_registry().register_calculator(meter.clone());
2889 let enabled = Cell::new(false);
2891 write_meter!({
2892 enabled.set(true);
2893 MeterRecord::new("probe".into(), "probe".into(), 0, 0, 0)
2894 })
2895 .await
2896 .unwrap();
2897 let enabled = enabled.get();
2898
2899 let output = inserter
2900 .do_request(rows_request(), &table_infos, &ctx)
2901 .await
2902 .unwrap();
2903 assert_eq!(output.meta.cost, if enabled { 17 } else { 0 });
2904 assert!(matches!(output.data, OutputData::AffectedRows(3)));
2905 dispatched.try_recv().unwrap();
2906 assert!(flow_cache.contains_key(&1));
2907 flow_cache.invalidate_all();
2908 meter.reject.store(true, Ordering::Relaxed);
2909 let result = inserter
2910 .do_request(rows_request(), &table_infos, &ctx)
2911 .await;
2912 if enabled {
2913 let error = result.unwrap_err();
2914 assert_eq!(error.status_code(), StatusCode::RateLimited);
2915 assert_eq!(error.retry_hint(), RetryHint::Retryable);
2916 assert!(error.to_string().contains("database row quota exhausted"));
2917 assert!(dispatched.try_recv().is_err());
2918 assert!(
2919 !flow_cache.contains_key(&1),
2920 "rejected write reached flow mirroring"
2921 );
2922 assert_eq!(meter.accepted_value.load(Ordering::Relaxed), 17);
2923 let attempts = meter.attempts.lock().unwrap();
2924 assert_eq!(attempts.len(), 2);
2925 for record in attempts.iter() {
2926 assert_eq!(record.catalog, CATALOG);
2927 assert_eq!(record.schema, "target_db");
2928 assert_eq!(
2929 (record.rows, record.value, record.source),
2930 (5, 17, Channel::Postgres as u8)
2931 );
2932 }
2933 } else {
2934 assert_eq!(result.unwrap().meta.cost, 0);
2935 dispatched.try_recv().unwrap();
2936 assert!(meter.attempts.lock().unwrap().is_empty());
2937 }
2938 meter.attempts.lock().unwrap().clear();
2939
2940 if enabled {
2942 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
2943 let rows = Rows {
2944 schema: vec![
2945 api::v1::helper::tag_column_schema("a", ColumnDataType::Int32),
2946 time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
2947 field_column_schema("b", ColumnDataType::Int32),
2948 ],
2949 rows: vec![api::v1::Row {
2950 values: vec![
2951 api::v1::value::ValueData::I32Value(60).into(),
2952 Value {
2953 value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue(
2954 0,
2955 )),
2956 },
2957 api::v1::value::ValueData::I32Value(0).into(),
2958 ],
2959 }],
2960 };
2961 let error = inserter
2962 .submit_table_rows(rows, table_info.clone(), ctx.clone(), &batcher)
2963 .await
2964 .unwrap_err();
2965 assert_eq!(error.status_code(), StatusCode::RateLimited);
2966 let mut attempts = meter.attempts.lock().unwrap();
2967 assert_eq!(attempts.len(), 1);
2968 assert_eq!(attempts[0].catalog, CATALOG);
2969 assert_eq!(attempts[0].schema, "target_db");
2970 assert_eq!(attempts[0].rows, 1);
2971 attempts.clear();
2972 }
2973
2974 let table = Arc::new(table::Table::new(
2975 table_info.clone(),
2976 table::metadata::FilterPushDownType::Unsupported,
2977 Arc::new(DummyDataSource),
2978 ));
2979 let batch = RecordBatch::try_new(
2980 table_info.meta.schema.arrow_schema().clone(),
2981 vec![
2982 Arc::new(Int32Array::from(vec![60, 70])),
2983 Arc::new(TimestampMillisecondArray::from(vec![0, 1])),
2984 Arc::new(Int32Array::from(vec![0, 0])),
2985 ],
2986 )
2987 .unwrap();
2988 let bulk_insert = |batch: RecordBatch| {
2989 let flight_data = FlightEncoder::default()
2990 .encode(FlightMessage::RecordBatch(batch.clone()))
2991 .into_iter()
2992 .next()
2993 .unwrap();
2994 inserter.handle_bulk_insert(
2995 table.clone(),
2996 flight_data,
2997 batch,
2998 Bytes::new(),
2999 false,
3000 Channel::Grpc,
3001 )
3002 };
3003 assert_eq!(bulk_insert(batch.slice(0, 0)).await.unwrap(), 0);
3005 assert!(meter.attempts.lock().unwrap().is_empty());
3006 assert!(dispatched.try_recv().is_err());
3007
3008 for reject in [false, true, false] {
3010 meter.reject.store(reject, Ordering::Relaxed);
3011 let result = bulk_insert(batch.clone()).await;
3012 if enabled && reject {
3013 let error = result.unwrap_err();
3014 assert_eq!(error.status_code(), StatusCode::RateLimited);
3015 assert_eq!(error.retry_hint(), RetryHint::Retryable);
3016 assert!(dispatched.try_recv().is_err());
3017 } else {
3018 assert_eq!(result.unwrap(), 2);
3019 dispatched.try_recv().unwrap();
3020 }
3021 }
3022 {
3023 let attempts = meter.attempts.lock().unwrap();
3024 assert_eq!(attempts.len(), if enabled { 3 } else { 0 });
3025 for record in attempts.iter() {
3026 assert_eq!(record.catalog, CATALOG);
3027 assert_eq!(record.schema, "target_db");
3028 assert_eq!(
3029 (record.rows, record.value, record.source),
3030 (2, 0, Channel::Grpc as u8)
3031 );
3032 }
3033 }
3034 assert_eq!(
3035 meter.accepted_value.load(Ordering::Relaxed),
3036 if enabled { 17 } else { 0 }
3037 );
3038
3039 meter.attempts.lock().unwrap().clear();
3042 let original = Arc::new(QueryContext::with_channel(
3043 CATALOG,
3044 "a",
3045 Channel::Prometheus,
3046 ));
3047 let mut batches = ["a", "b", "a"].map(|schema| {
3048 let ctx = if schema == "a" {
3049 original.clone()
3050 } else {
3051 Arc::new(QueryContext::with_channel(
3052 CATALOG,
3053 schema,
3054 Channel::Prometheus,
3055 ))
3056 };
3057 (
3058 ctx,
3059 RowInsertRequests {
3060 inserts: vec![RowInsertRequest {
3061 table_name: "data".into(),
3062 rows: Some(Rows {
3063 schema: vec![],
3064 rows: vec![api::v1::Row::default(); 2],
3065 }),
3066 }],
3067 },
3068 )
3069 });
3070 admit_row_insert_batches(&mut batches).await.unwrap();
3071 admit_row_insert_batches(&mut batches).await.unwrap();
3072 assert_eq!(original.write_rows_to_admit(CATALOG, "a", 4), 4);
3073 for (ctx, _) in &batches {
3074 assert_eq!(
3075 ctx.write_rows_to_admit(CATALOG, &ctx.current_schema(), 2),
3076 0
3077 );
3078 assert_eq!(ctx.channel(), Channel::Prometheus);
3079 }
3080 {
3081 let attempts = meter.attempts.lock().unwrap();
3082 let totals = attempts
3083 .iter()
3084 .map(|r| (r.schema.as_str(), r.rows, r.value))
3085 .collect::<Vec<_>>();
3086 assert_eq!(
3087 totals,
3088 if enabled {
3089 vec![("a", 4, 0), ("b", 2, 0)]
3090 } else {
3091 vec![]
3092 }
3093 );
3094 }
3095 meter.attempts.lock().unwrap().clear();
3096 meter.reject.store(false, Ordering::Relaxed);
3097 let ctx = Arc::new(QueryContext::with_channel(
3098 CATALOG,
3099 "logical",
3100 Channel::Otlp,
3101 ));
3102 let admitted = admit_write(2, &ctx).await.unwrap();
3103 let mut requests = RowInsertRequests {
3104 inserts: vec![RowInsertRequest {
3105 table_name: "metric".to_string(),
3106 rows: Some(Rows {
3107 schema: vec![],
3108 rows: vec![api::v1::Row::default(); 2],
3109 }),
3110 }],
3111 };
3112 let original = requests.clone();
3113 let cost = Inserter::meter_row_inserts(&mut requests, &admitted)
3114 .await
3115 .unwrap();
3116 assert_eq!(cost, if enabled { 17 } else { 0 });
3117 assert_eq!(requests, original);
3118 let records = meter
3119 .attempts
3120 .lock()
3121 .unwrap()
3122 .iter()
3123 .map(|record| (record.rows, record.value))
3124 .collect::<Vec<_>>();
3125 assert_eq!(
3126 records,
3127 if enabled {
3128 vec![(2, 0), (0, 17)]
3129 } else {
3130 vec![]
3131 }
3132 );
3133 meter.reject.store(true, Ordering::Relaxed);
3134 let result = Inserter::meter_row_inserts(&mut requests, &admitted).await;
3135 if enabled {
3136 assert_eq!(result.unwrap_err().status_code(), StatusCode::RateLimited);
3137 } else {
3138 assert_eq!(result.unwrap(), 0);
3139 }
3140 assert_eq!(requests, original);
3141 }
3142
3143 #[test]
3144 fn test_trace_create_options_copy_compaction_hints() {
3145 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3146 ctx.set_extension("compaction.twcs.time_window", "1h");
3147 ctx.set_extension("compaction.twcs.max_output_file_size", "256MB");
3148 ctx.set_extension("compaction.twcs.remote_compaction", "true");
3149 ctx.set_extension("compaction.unknown", "ignored");
3150 let ctx = Arc::new(ctx);
3151 let mut table_options = Default::default();
3152
3153 fill_table_options_for_create(
3154 &mut table_options,
3155 &AutoCreateTableType::Trace {
3156 alter_existing: true,
3157 },
3158 &ctx,
3159 );
3160
3161 assert_eq!(
3162 Some("twcs"),
3163 table_options.get(COMPACTION_TYPE).map(String::as_str)
3164 );
3165 assert_eq!(
3166 Some("1h"),
3167 table_options
3168 .get("compaction.twcs.time_window")
3169 .map(String::as_str)
3170 );
3171 assert_eq!(
3172 Some("256MB"),
3173 table_options
3174 .get("compaction.twcs.max_output_file_size")
3175 .map(String::as_str)
3176 );
3177 assert_eq!(
3178 Some("true"),
3179 table_options
3180 .get("compaction.twcs.remote_compaction")
3181 .map(String::as_str)
3182 );
3183 assert!(!table_options.contains_key("compaction.unknown"));
3184 }
3185
3186 #[test]
3187 fn test_skip_wal_does_not_change_table_options() {
3188 check_skip_wal_does_not_change_table_options(false);
3189 check_skip_wal_does_not_change_table_options(true);
3190 }
3191
3192 fn check_skip_wal_does_not_change_table_options(skip_wal: bool) {
3193 let ctx = Arc::new(QueryContext::with(
3194 DEFAULT_CATALOG_NAME,
3195 DEFAULT_SCHEMA_NAME,
3196 ));
3197 ctx.set_skip_wal(skip_wal);
3198 let mut options = Default::default();
3199 fill_table_options_for_create(&mut options, &AutoCreateTableType::Physical, &ctx);
3200 assert!(!options.contains_key(session::hints::INSERT_SKIP_WAL_HINT));
3201 assert!(!options.contains_key("skip_wal"));
3202 }
3203
3204 #[test]
3205 fn test_last_non_null_create_options_preserve_default_without_append_mode() {
3206 let ctx = Arc::new(QueryContext::with(
3207 DEFAULT_CATALOG_NAME,
3208 DEFAULT_SCHEMA_NAME,
3209 ));
3210 let mut table_options = Default::default();
3211
3212 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3213
3214 assert_eq!(
3215 Some("last_non_null"),
3216 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3217 );
3218 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3219 }
3220
3221 #[test]
3222 fn test_fill_table_options_copies_semantic_extensions() {
3223 use table::requests::{
3224 SEMANTIC_METRIC_TYPE, SEMANTIC_PER_TABLE_INDEX_KEY, SEMANTIC_SIGNAL_TYPE,
3225 SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, SOURCE_OPENTELEMETRY,
3226 };
3227
3228 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3229 ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
3230 ctx.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
3231 ctx.set_extension(SEMANTIC_SOURCE_VERSION, "2.0");
3232 ctx.set_extension(SEMANTIC_METRIC_TYPE, "bogus");
3233 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, "{}");
3235 let ctx = Arc::new(ctx);
3236 let mut table_options = Default::default();
3237
3238 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::Physical, &ctx);
3239
3240 assert_eq!(
3241 Some(SIGNAL_TYPE_METRIC),
3242 table_options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str)
3243 );
3244 assert_eq!(
3245 Some(SOURCE_OPENTELEMETRY),
3246 table_options.get(SEMANTIC_SOURCE).map(String::as_str)
3247 );
3248 assert_eq!(
3249 Some("2.0"),
3250 table_options
3251 .get(SEMANTIC_SOURCE_VERSION)
3252 .map(String::as_str)
3253 );
3254 assert!(!table_options.contains_key(SEMANTIC_METRIC_TYPE));
3255 assert!(!table_options.contains_key(SEMANTIC_PER_TABLE_INDEX_KEY));
3256 }
3257
3258 #[test]
3259 fn test_apply_per_table_semantic_options() {
3260 use table::requests::{
3261 SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, SEMANTIC_PER_TABLE_INDEX_KEY,
3262 };
3263
3264 let index = format!(
3265 r#"{{
3266 "{DEFAULT_SCHEMA_NAME}": {{
3267 "http_requests_total": {{
3268 "greptime.semantic.metric.type": "counter",
3269 "greptime.semantic.metric.unit": "By",
3270 "greptime.semantic.metric.type_BOGUS": "x"
3271 }},
3272 "other_table": {{
3273 "greptime.semantic.metric.type": "gauge"
3274 }}
3275 }},
3276 "other_schema": {{
3277 "http_requests_total": {{
3278 "greptime.semantic.metric.type": "gauge"
3279 }}
3280 }}
3281 }}"#
3282 );
3283 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3284 ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index);
3285 let ctx = Arc::new(ctx);
3286
3287 let index = parse_per_table_semantic_index(&ctx);
3288 assert!(index.is_some());
3289 let index = index.as_ref();
3290
3291 let mut table_options = std::collections::HashMap::new();
3292 apply_per_table_semantic_options(
3293 &mut table_options,
3294 index,
3295 DEFAULT_SCHEMA_NAME,
3296 "http_requests_total",
3297 );
3298 assert_eq!(
3300 table_options.get(SEMANTIC_METRIC_TYPE).map(String::as_str),
3301 Some("counter")
3302 );
3303 assert_eq!(
3304 table_options.get(SEMANTIC_METRIC_UNIT).map(String::as_str),
3305 Some("By")
3306 );
3307 assert!(!table_options.contains_key("greptime.semantic.metric.type_BOGUS"));
3310 assert_eq!(table_options.len(), 2);
3311
3312 let mut empty = std::collections::HashMap::new();
3313 apply_per_table_semantic_options(&mut empty, index, DEFAULT_SCHEMA_NAME, "not_in_index");
3314 assert!(empty.is_empty());
3315
3316 let mut opts = std::collections::HashMap::new();
3319 apply_per_table_semantic_options(
3320 &mut opts,
3321 index,
3322 "schema_without_entry",
3323 "http_requests_total",
3324 );
3325 assert!(opts.is_empty());
3326
3327 let bare = Arc::new(QueryContext::with(
3329 DEFAULT_CATALOG_NAME,
3330 DEFAULT_SCHEMA_NAME,
3331 ));
3332 assert!(parse_per_table_semantic_index(&bare).is_none());
3333 let mut opts = std::collections::HashMap::new();
3334 apply_per_table_semantic_options(
3335 &mut opts,
3336 None,
3337 DEFAULT_SCHEMA_NAME,
3338 "http_requests_total",
3339 );
3340 assert!(opts.is_empty());
3341 }
3342
3343 #[test]
3344 fn test_last_non_null_create_options_preserve_default_with_append_mode_false() {
3345 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3346 ctx.set_extension(APPEND_MODE_KEY, "false");
3347 let ctx = Arc::new(ctx);
3348 let mut table_options = Default::default();
3349
3350 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3351
3352 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3353 assert_eq!(
3354 Some("last_non_null"),
3355 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3356 );
3357 }
3358
3359 #[test]
3360 fn test_last_non_null_create_options_use_configured_merge_mode() {
3361 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3362 ctx.set_extension(MERGE_MODE_KEY, "last_row");
3363 let ctx = Arc::new(ctx);
3364 let mut table_options = Default::default();
3365
3366 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3367
3368 assert_eq!(
3369 Some("last_row"),
3370 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3371 );
3372 assert!(!table_options.contains_key(APPEND_MODE_KEY));
3373 }
3374
3375 #[test]
3376 fn test_last_non_null_create_options_use_last_row_with_append_mode_true() {
3377 let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
3378 ctx.set_extension(APPEND_MODE_KEY, "true");
3379 let ctx = Arc::new(ctx);
3380 let mut table_options = Default::default();
3381
3382 fill_table_options_for_create(&mut table_options, &AutoCreateTableType::LastNonNull, &ctx);
3383
3384 assert_eq!(
3385 Some("true"),
3386 table_options.get(APPEND_MODE_KEY).map(String::as_str)
3387 );
3388 assert_eq!(
3389 Some("last_row"),
3390 table_options.get(MERGE_MODE_KEY).map(String::as_str)
3391 );
3392 }
3393
3394 struct UnexpectedBatcher;
3395
3396 #[async_trait::async_trait]
3397 impl PendingRowsBatcher for UnexpectedBatcher {
3398 async fn acquire(&self) -> Result<Arc<tokio::sync::OwnedSemaphorePermit>> {
3399 panic!("empty writes must not acquire batch admission")
3400 }
3401
3402 async fn submit(
3403 &self,
3404 _table_info: TableInfoRef,
3405 _batch: arrow::record_batch::RecordBatch,
3406 _ctx: QueryContextRef,
3407 _permit: Arc<tokio::sync::OwnedSemaphorePermit>,
3408 ) -> Result<Output> {
3409 panic!("empty writes must not submit a batch")
3410 }
3411 }
3412
3413 async fn batcher_test_inserter() -> Inserter {
3414 let kv_backend = prepare_mocked_backend().await;
3415 Inserter::new(
3416 catalog::memory::MemoryCatalogManager::new(),
3417 create_partition_rule_manager(kv_backend.clone()).await,
3418 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3419 Arc::new(new_table_flownode_set_cache(
3420 String::new(),
3421 Cache::new(100),
3422 kv_backend,
3423 )),
3424 true,
3425 )
3426 }
3427
3428 #[tokio::test]
3429 async fn test_logical_batcher_eligibility() {
3430 use catalog::RegisterTableRequest;
3431 use catalog::memory::MemoryCatalogManager;
3432 use common_meta::instruction::{CacheIdent, CreateFlow};
3433 use common_meta::kv_backend::KvBackendRef;
3434 use common_meta::kv_backend::memory::MemoryKvBackend;
3435 use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder};
3436 let requests = RowInsertRequests {
3437 inserts: vec![RowInsertRequest {
3438 table_name: "test_table".to_string(),
3439 rows: None,
3440 }],
3441 };
3442 let original =
3443 make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3444 .table_info();
3445 for case in [
3446 "eligible",
3447 "physical",
3448 "ordinary",
3449 "instant",
3450 "disabled",
3451 "hint",
3452 "flow",
3453 "required_tag",
3454 "default",
3455 ] {
3456 let mut info = (*original).clone();
3457 info.meta.engine = METRIC_ENGINE_NAME.to_string();
3458 info.meta.options.extra_options.insert(
3459 LOGICAL_TABLE_METADATA_KEY.to_string(),
3460 "physical".to_string(),
3461 );
3462 let mut ctx = QueryContext::arc().fork();
3463 match case {
3464 "physical" => {
3465 info.meta
3466 .options
3467 .extra_options
3468 .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "other".to_string());
3469 }
3470 "ordinary" => info.meta.engine = "mito".to_string(),
3471 "instant" => info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant),
3472 "hint" => ctx.set_extension(AUTO_CREATE_TABLE_KEY, "false"),
3473 "required_tag" | "default" => {
3474 let mut columns = info.meta.schema.column_schemas().to_vec();
3475 let mut tag = ColumnSchema::new(
3476 "tag",
3477 ConcreteDataType::string_datatype(),
3478 case != "required_tag",
3479 );
3480 if case == "default" {
3481 tag = tag
3482 .with_default_constraint(Some(ColumnDefaultConstraint::null_value()))
3483 .unwrap();
3484 }
3485 columns.push(tag);
3486 info.meta.schema = Arc::new(
3487 SchemaBuilder::try_from_columns(columns)
3488 .unwrap()
3489 .build()
3490 .unwrap(),
3491 );
3492 info.meta.primary_key_indices = vec![2];
3493 }
3494 _ => {}
3495 }
3496 let catalog = MemoryCatalogManager::with_default_setup();
3497 let table = Arc::new(table::Table::new(
3498 Arc::new(info),
3499 table::metadata::FilterPushDownType::Unsupported,
3500 Arc::new(DummyDataSource),
3501 ));
3502 catalog
3503 .register_table_sync(RegisterTableRequest {
3504 catalog: DEFAULT_CATALOG_NAME.to_string(),
3505 schema: DEFAULT_SCHEMA_NAME.to_string(),
3506 table_name: "test_table".to_string(),
3507 table_id: 1,
3508 table,
3509 })
3510 .unwrap();
3511 let mut inserter = batcher_test_inserter().await;
3512 inserter.catalog_manager = catalog;
3513 inserter.auto_create_table = case != "disabled";
3514 let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::default());
3515 inserter.table_flownode_set_cache = Arc::new(new_table_flownode_set_cache(
3516 String::new(),
3517 Cache::new(10),
3518 kv_backend,
3519 ));
3520 if case == "flow" {
3521 inserter
3522 .table_flownode_set_cache
3523 .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
3524 flow_id: 1,
3525 source_table_ids: vec![1],
3526 partition_to_peer_mapping: vec![(0, Peer::empty(1))],
3527 })])
3528 .await
3529 .unwrap();
3530 }
3531 assert_eq!(
3532 inserter
3533 .can_batch_metric_rows(&requests, &Arc::new(ctx), "physical")
3534 .await
3535 .unwrap(),
3536 case == "eligible",
3537 "{case}"
3538 );
3539 }
3540 }
3541
3542 #[tokio::test]
3543 async fn test_batcher_meter_preserves_request() {
3544 let mut requests = RowInsertRequests {
3545 inserts: vec![RowInsertRequest {
3546 table_name: "sample".to_string(),
3547 rows: Some(Rows {
3548 schema: vec![],
3549 rows: vec![api::v1::Row {
3550 values: vec![Value {
3551 value_data: Some(api::v1::value::ValueData::F64Value(1.5)),
3552 }],
3553 }],
3554 }),
3555 }],
3556 };
3557 let expected = requests.clone();
3558 Inserter::meter_row_inserts(&mut requests, &QueryContext::arc())
3559 .await
3560 .unwrap();
3561 assert_eq!(requests, expected);
3562 }
3563
3564 #[tokio::test]
3565 async fn test_instant_table_bypasses_batcher() {
3566 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3567 let inserter = batcher_test_inserter()
3568 .await
3569 .with_pending_rows_batcher(Some(batcher));
3570 let mut ctx = session::context::QueryContextBuilder::default().build();
3571 ctx.set_batching_enabled(true);
3572 let ctx = Arc::new(ctx);
3573 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3574 .table_info();
3575 assert!(inserter.table_batcher(&table, &ctx).is_some());
3576 let mut instant = (*table).clone();
3577 instant.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
3578 assert!(inserter.table_batcher(&Arc::new(instant), &ctx).is_none());
3579 }
3580
3581 #[tokio::test]
3582 async fn test_empty_prepared_rows_skip_batcher() {
3583 let inserter = batcher_test_inserter().await;
3584 let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype())
3585 .table_info();
3586 let batcher: Arc<dyn PendingRowsBatcher> = Arc::new(UnexpectedBatcher);
3587 let ctx = QueryContext::arc();
3588 let output = inserter
3589 .submit_table_rows(
3590 Rows {
3591 schema: vec![],
3592 rows: vec![],
3593 },
3594 table.clone(),
3595 ctx.clone(),
3596 &batcher,
3597 )
3598 .await
3599 .unwrap();
3600 assert!(matches!(output.data, OutputData::AffectedRows(0)));
3601 let output = inserter
3602 .submit_pending_rows(
3603 RowInsertRequests {
3604 inserts: vec![RowInsertRequest {
3605 table_name: table.name.clone(),
3606 rows: None,
3607 }],
3608 },
3609 HashMap::from_iter([(table.table_id(), table)]),
3610 ctx,
3611 &batcher,
3612 )
3613 .await
3614 .unwrap();
3615 assert!(matches!(output.data, OutputData::AffectedRows(0)));
3616 }
3617
3618 #[derive(Clone)]
3622 struct GatedFlownodeHandler {
3623 gate: Arc<tokio::sync::Semaphore>,
3624 affected_rows: u64,
3625 }
3626
3627 #[async_trait::async_trait]
3628 impl MockFlownodeHandler for GatedFlownodeHandler {
3629 async fn handle_inserts(
3630 &self,
3631 _peer: &Peer,
3632 _requests: api::v1::region::InsertRequests,
3633 ) -> common_meta::error::Result<FlowResponse> {
3634 let _permit = self.gate.acquire().await.unwrap();
3635 Ok(FlowResponse {
3636 affected_rows: self.affected_rows,
3637 ..Default::default()
3638 })
3639 }
3640 }
3641
3642 fn flownode_peer() -> Peer {
3643 Peer {
3644 id: 1,
3645 addr: "127.0.0.1:4001".to_string(),
3646 }
3647 }
3648
3649 #[derive(Clone)]
3650 struct RecordingNodeManager {
3651 datanode_dispatch: tokio::sync::mpsc::UnboundedSender<usize>,
3652 flownode_dispatch: tokio::sync::mpsc::UnboundedSender<Peer>,
3653 gates: Arc<std::sync::Mutex<HashMap<u64, Arc<tokio::sync::Semaphore>>>>,
3654 fail_flownode: bool,
3655 }
3656
3657 #[async_trait::async_trait]
3658 impl common_meta::node_manager::DatanodeManager for RecordingNodeManager {
3659 async fn datanode(&self, peer: &Peer) -> common_meta::node_manager::DatanodeRef {
3660 Arc::new(RecordingNode {
3661 peer: peer.clone(),
3662 manager: self.clone(),
3663 })
3664 }
3665 }
3666
3667 #[async_trait::async_trait]
3668 impl common_meta::node_manager::FlownodeManager for RecordingNodeManager {
3669 async fn flownode(&self, peer: &Peer) -> common_meta::node_manager::FlownodeRef {
3670 Arc::new(RecordingNode {
3671 peer: peer.clone(),
3672 manager: self.clone(),
3673 })
3674 }
3675 }
3676
3677 struct RecordingNode {
3678 peer: Peer,
3679 manager: RecordingNodeManager,
3680 }
3681
3682 #[async_trait::async_trait]
3683 impl common_meta::node_manager::Datanode for RecordingNode {
3684 async fn handle_request(
3685 &self,
3686 _request: Arc<api::v1::region::RegionRequest>,
3687 ) -> common_meta::error::Result<api::region::RegionResponse> {
3688 let rows = match _request.body.as_ref() {
3689 Some(api::v1::region::region_request::Body::Inserts(requests)) => requests
3690 .requests
3691 .iter()
3692 .filter_map(|request| request.rows.as_ref())
3693 .map(|rows| rows.rows.len())
3694 .sum(),
3695 _ => 0,
3696 };
3697 let _ = self.manager.datanode_dispatch.send(rows);
3698 Ok(api::region::RegionResponse::new(rows))
3699 }
3700
3701 async fn handle_query(
3702 &self,
3703 _request: common_query::request::QueryRequest,
3704 ) -> common_meta::error::Result<common_recordbatch::SendableRecordBatchStream> {
3705 unreachable!()
3706 }
3707 }
3708
3709 #[async_trait::async_trait]
3710 impl common_meta::node_manager::Flownode for RecordingNode {
3711 async fn handle(
3712 &self,
3713 _request: api::v1::flow::FlowRequest,
3714 ) -> common_meta::error::Result<FlowResponse> {
3715 unreachable!()
3716 }
3717
3718 async fn handle_inserts(
3719 &self,
3720 _request: api::v1::region::InsertRequests,
3721 ) -> common_meta::error::Result<FlowResponse> {
3722 let _ = self.manager.flownode_dispatch.send(self.peer.clone());
3723 if self.manager.fail_flownode {
3724 return Err(common_meta::error::UnexpectedSnafu {
3725 err_msg: "test flownode failure".to_string(),
3726 }
3727 .build());
3728 }
3729 let gate = self
3730 .manager
3731 .gates
3732 .lock()
3733 .unwrap()
3734 .get(&self.peer.id)
3735 .cloned();
3736 if let Some(gate) = gate {
3737 let _permit = gate.acquire().await.unwrap();
3738 }
3739 Ok(FlowResponse::default())
3740 }
3741
3742 async fn handle_mark_window_dirty(
3743 &self,
3744 _request: api::v1::flow::DirtyWindowRequests,
3745 ) -> common_meta::error::Result<FlowResponse> {
3746 unreachable!()
3747 }
3748 }
3749
3750 fn mirror_requests(peer: &Peer, num_rows: usize) -> HashMap<Peer, RegionInsertRequests> {
3751 HashMap::from_iter([(
3752 peer.clone(),
3753 RegionInsertRequests {
3754 requests: vec![RegionInsertRequest {
3755 region_id: RegionId::new(1, 1).as_u64(),
3756 rows: Some(Rows {
3757 schema: vec![],
3758 rows: vec![api::v1::Row { values: vec![] }; num_rows],
3759 }),
3760 ..Default::default()
3761 }],
3762 },
3763 )])
3764 }
3765
3766 static MIRROR_METRIC_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
3771
3772 async fn lock_mirror_metrics() -> tokio::sync::MutexGuard<'static, ()> {
3773 MIRROR_METRIC_TEST_LOCK.lock().await
3774 }
3775
3776 #[tokio::test]
3780 async fn flow_mirror_saturation_preserves_instant_and_normal_contracts() {
3781 use common_error::ext::{ErrorExt, RetryHint};
3782 use common_error::status_code::StatusCode;
3783
3784 let _guard = lock_mirror_metrics().await;
3785 let kv_backend = prepare_mocked_backend().await;
3786 let partition_manager = create_partition_rule_manager(kv_backend.clone()).await;
3787 let (datanode_tx, mut datanode_rx) = tokio::sync::mpsc::unbounded_channel();
3788 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
3789 let node_manager = Arc::new(RecordingNodeManager {
3790 datanode_dispatch: datanode_tx,
3791 flownode_dispatch: flownode_tx,
3792 gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
3793 fail_flownode: false,
3794 });
3795 let flow_cache = Cache::new(10);
3796 let flow_cache_backend = prepare_mocked_backend().await;
3797 let inserter = Inserter::new(
3798 catalog::memory::MemoryCatalogManager::new(),
3799 partition_manager,
3800 node_manager,
3801 Arc::new(new_table_flownode_set_cache(
3802 String::new(),
3803 flow_cache.clone(),
3804 flow_cache_backend,
3805 )),
3806 true,
3807 );
3808 let peer = flownode_peer();
3809 inserter
3810 .table_flownode_set_cache
3811 .invalidate(&[common_meta::instruction::CacheIdent::CreateFlow(
3812 common_meta::instruction::CreateFlow {
3813 flow_id: 1,
3814 source_table_ids: vec![1, 2],
3815 partition_to_peer_mapping: vec![(0, peer)],
3816 },
3817 )])
3818 .await
3819 .unwrap();
3820 let mut normal_info = new_test_table_info(1, "normal_table", [1].into_iter());
3821 normal_info.catalog_name = DEFAULT_CATALOG_NAME.to_string();
3822 normal_info.schema_name = DEFAULT_SCHEMA_NAME.to_string();
3823 let mut instant_info = new_test_table_info(2, "instant_table", [1].into_iter());
3824 instant_info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant);
3825 let table_infos =
3826 HashMap::from_iter([(1, Arc::new(normal_info)), (2, Arc::new(instant_info))]);
3827 let ctx = Arc::new(QueryContext::with(
3828 DEFAULT_CATALOG_NAME,
3829 DEFAULT_SCHEMA_NAME,
3830 ));
3831 let insert = |table_id, num_rows| RegionInsertRequests {
3832 requests: vec![RegionInsertRequest {
3833 region_id: RegionId::new(table_id, 1).as_u64(),
3834 rows: Some(Rows {
3835 schema: vec![],
3836 rows: vec![api::v1::Row::default(); num_rows],
3837 }),
3838 ..Default::default()
3839 }],
3840 };
3841 let make_request = |normal, instant| InstantAndNormalInsertRequests {
3842 normal_requests: if normal == 0 {
3843 RegionInsertRequests::default()
3844 } else {
3845 insert(1, normal)
3846 },
3847 instant_requests: if instant == 0 {
3848 RegionInsertRequests::default()
3849 } else {
3850 insert(2, instant)
3851 },
3852 };
3853
3854 let pending = inserter.mirror_pending_rows.clone();
3855 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
3856 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
3857 pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
3858 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
3859 for request in [make_request(0, 1), make_request(1, 1)] {
3860 let error = inserter
3861 .do_request(request, &table_infos, &ctx)
3862 .await
3863 .unwrap_err();
3864 assert_eq!(error.status_code(), StatusCode::RateLimited);
3865 assert_eq!(error.retry_hint(), RetryHint::Retryable);
3866 assert!(datanode_rx.try_recv().is_err());
3867 assert!(flownode_rx.try_recv().is_err());
3868 assert_eq!(
3869 pending.load(Ordering::Relaxed),
3870 MAX_MIRROR_PENDING_ROWS as u64
3871 );
3872 }
3873 assert_eq!(
3874 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
3875 dropped_before
3876 );
3877
3878 for pending_before in [0, MAX_MIRROR_PENDING_ROWS as u64] {
3879 if pending_before == 0 {
3880 pending.store(0, Ordering::Relaxed);
3881 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
3882 }
3883 let gauge_before_rejection = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
3884 for (normal, instant) in [
3885 (0, MAX_MIRROR_PENDING_ROWS + 1),
3886 (MAX_MIRROR_PENDING_ROWS / 2, MAX_MIRROR_PENDING_ROWS / 2 + 1),
3887 ] {
3888 let error = inserter
3889 .do_request(make_request(normal, instant), &table_infos, &ctx)
3890 .await
3891 .unwrap_err();
3892 assert_eq!(error.status_code(), StatusCode::InvalidArguments);
3893 assert_eq!(error.retry_hint(), RetryHint::NonRetryable);
3894 let message = error.to_string();
3895 assert!(message.contains("1000001"), "{message}");
3896 assert!(message.contains("1000000"), "{message}");
3897 assert!(message.contains("reduce the batch size"), "{message}");
3898 assert!(datanode_rx.try_recv().is_err());
3899 assert!(flownode_rx.try_recv().is_err());
3900 assert_eq!(pending.load(Ordering::Relaxed), pending_before);
3901 assert_eq!(
3902 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
3903 gauge_before_rejection
3904 );
3905 assert_eq!(
3906 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
3907 dropped_before
3908 );
3909 }
3910 if pending_before == 0 {
3911 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64);
3912 pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed);
3913 }
3914 }
3915
3916 let result = inserter
3917 .do_request(make_request(1, 0), &table_infos, &ctx)
3918 .await
3919 .unwrap();
3920 assert!(matches!(result.data, OutputData::AffectedRows(1)));
3921 assert_eq!(1, datanode_rx.try_recv().unwrap());
3922 assert!(flownode_rx.try_recv().is_err());
3923 assert_eq!(
3924 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
3925 dropped_before + 1
3926 );
3927 pending.store(0, Ordering::Relaxed);
3928 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64);
3929 assert_eq!(
3930 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
3931 gauge_before
3932 );
3933 }
3934
3935 #[tokio::test]
3936 async fn flow_mirror_new_counts_regions_and_peer_clones() {
3937 use common_meta::instruction::{CacheIdent, CreateFlow};
3938
3939 let peer1 = flownode_peer();
3940 let peer2 = Peer {
3941 id: 2,
3942 addr: "127.0.0.1:4002".to_string(),
3943 };
3944 let cache_backend = prepare_mocked_backend().await;
3945 let cache = Arc::new(new_table_flownode_set_cache(
3946 String::new(),
3947 Cache::new(10),
3948 cache_backend,
3949 ));
3950 cache
3951 .invalidate(&[CacheIdent::CreateFlow(CreateFlow {
3952 flow_id: 1,
3953 source_table_ids: vec![1],
3954 partition_to_peer_mapping: vec![(0, peer1), (1, peer2)],
3955 })])
3956 .await
3957 .unwrap();
3958 let request = |region_id| RegionInsertRequest {
3959 region_id,
3960 rows: Some(Rows {
3961 schema: vec![],
3962 rows: vec![api::v1::Row { values: vec![] }; 3],
3963 }),
3964 ..Default::default()
3965 };
3966 let requests = [
3967 request(RegionId::new(1, 0).as_u64()),
3968 request(RegionId::new(1, 1).as_u64()),
3969 ];
3970 let task = FlowMirrorTask::new(&cache, requests.iter()).await.unwrap();
3971 assert_eq!(12, task.pending_rows());
3972 assert_eq!(2, task.requests.len());
3973 }
3974
3975 #[tokio::test]
3976 async fn flow_mirror_dropped_when_pending_exceeds_limit() {
3977 let _guard = lock_mirror_metrics().await;
3978 let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64));
3979 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
3980 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
3981 let num_rows = 100;
3982
3983 let task = FlowMirrorTask {
3984 requests: mirror_requests(&flownode_peer(), num_rows),
3985 };
3986 task.detach(
3988 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
3989 pending.clone(),
3990 false,
3991 )
3992 .unwrap();
3993
3994 assert_eq!(
3995 num_rows as u64,
3996 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get() - dropped_before
3997 );
3998 assert_eq!(
4000 MAX_MIRROR_PENDING_ROWS as u64,
4001 pending.load(Ordering::Relaxed)
4002 );
4003 assert_eq!(
4006 gauge_before,
4007 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4008 "a dropped batch must not change the pending gauge"
4009 );
4010
4011 let task = FlowMirrorTask {
4012 requests: mirror_requests(&flownode_peer(), num_rows),
4013 };
4014 let error = task
4015 .detach(
4016 Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
4017 pending.clone(),
4018 true,
4019 )
4020 .err()
4021 .unwrap();
4022 use common_error::ext::{ErrorExt, RetryHint};
4023 use common_error::status_code::StatusCode;
4024 assert_eq!(error.status_code(), StatusCode::RateLimited);
4025 assert_eq!(error.retry_hint(), RetryHint::Retryable);
4026 assert_eq!(
4027 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4028 dropped_before + num_rows as u64,
4029 "rejected instant rows are not successful best-effort drops"
4030 );
4031 assert_eq!(
4032 MAX_MIRROR_PENDING_ROWS as u64,
4033 pending.load(Ordering::Relaxed)
4034 );
4035 assert_eq!(
4036 gauge_before,
4037 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get()
4038 );
4039 }
4040
4041 #[tokio::test]
4045 async fn flow_mirror_reserves_and_releases_per_peer_shares() {
4046 let _guard = lock_mirror_metrics().await;
4047 let per_peer_rows = 10;
4048 let total_rows = (2 * per_peer_rows) as u64;
4049 let start = MAX_MIRROR_PENDING_ROWS as u64 - total_rows;
4051 let pending = Arc::new(AtomicU64::new(start));
4052 let gate = Arc::new(tokio::sync::Semaphore::new(0));
4053 let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4054 gate: gate.clone(),
4055 affected_rows: per_peer_rows as u64,
4056 }));
4057
4058 let mut requests = mirror_requests(&flownode_peer(), per_peer_rows);
4059 requests.extend(mirror_requests(
4060 &Peer {
4061 id: 2,
4062 addr: "127.0.0.1:4002".to_string(),
4063 },
4064 per_peer_rows,
4065 ));
4066
4067 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4068 FlowMirrorTask { requests }
4069 .detach(node_manager, pending.clone(), false)
4070 .unwrap();
4071
4072 assert_eq!(
4075 MAX_MIRROR_PENDING_ROWS as u64,
4076 pending.load(Ordering::Relaxed)
4077 );
4078
4079 gate.add_permits(2);
4085 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4086 while pending.load(Ordering::Relaxed) != start
4087 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4088 {
4089 assert!(
4090 std::time::Instant::now() < deadline,
4091 "the mirror tasks must release exactly their own shares"
4092 );
4093 std::thread::sleep(std::time::Duration::from_millis(1));
4094 }
4095 }
4096
4097 #[tokio::test]
4098 async fn flow_mirror_fast_peer_releases_only_its_share_before_drop() {
4099 let _guard = lock_mirror_metrics().await;
4100 let peer1 = flownode_peer();
4101 let peer2 = Peer {
4102 id: 2,
4103 addr: "127.0.0.1:4002".to_string(),
4104 };
4105 let gate1 = Arc::new(tokio::sync::Semaphore::new(0));
4106 let gate2 = Arc::new(tokio::sync::Semaphore::new(0));
4107 let gates = Arc::new(std::sync::Mutex::new(HashMap::from_iter([
4108 (peer1.id, gate1.clone()),
4109 (peer2.id, gate2.clone()),
4110 ])));
4111 let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4112 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4113 let node_manager = Arc::new(RecordingNodeManager {
4114 datanode_dispatch: datanode_tx,
4115 flownode_dispatch: flownode_tx,
4116 gates,
4117 fail_flownode: false,
4118 });
4119 let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64 - 20));
4120 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4121 let mut requests = mirror_requests(&peer1, 10);
4122 requests.extend(mirror_requests(&peer2, 10));
4123 let task = FlowMirrorTask { requests };
4124 task.detach(node_manager.clone(), pending.clone(), false)
4125 .unwrap();
4126 let first = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4127 .await
4128 .unwrap()
4129 .unwrap();
4130 let second = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4131 .await
4132 .unwrap()
4133 .unwrap();
4134 assert_ne!(first.id, second.id);
4135 let fast_gate = if first.id == peer1.id {
4136 gate1.clone()
4137 } else {
4138 gate2.clone()
4139 };
4140 let slow_gate = if first.id == peer1.id {
4141 gate2.clone()
4142 } else {
4143 gate1.clone()
4144 };
4145 fast_gate.add_permits(1);
4146 let deadline = std::time::Instant::now() + Duration::from_secs(10);
4147 while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 10 {
4148 assert!(
4149 std::time::Instant::now() < deadline,
4150 "fast peer did not release"
4151 );
4152 std::thread::sleep(Duration::from_millis(1));
4153 }
4154 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4155 FlowMirrorTask {
4156 requests: mirror_requests(&peer1, 11),
4157 }
4158 .detach(node_manager, pending.clone(), false)
4159 .unwrap();
4160 assert_eq!(
4161 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4162 dropped_before + 11
4163 );
4164 assert!(flownode_rx.try_recv().is_err());
4165 slow_gate.add_permits(1);
4166 while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 20
4167 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4168 {
4169 assert!(
4170 std::time::Instant::now() < deadline,
4171 "slow peer did not release"
4172 );
4173 std::thread::sleep(Duration::from_millis(1));
4174 }
4175 }
4176
4177 #[tokio::test]
4178 async fn flow_mirror_failure_releases_budget_and_gauge() {
4179 let _guard = lock_mirror_metrics().await;
4180 let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel();
4181 let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel();
4182 let node_manager = Arc::new(RecordingNodeManager {
4183 datanode_dispatch: datanode_tx,
4184 flownode_dispatch: flownode_tx,
4185 gates: Arc::new(std::sync::Mutex::new(HashMap::new())),
4186 fail_flownode: true,
4187 });
4188 let pending = Arc::new(AtomicU64::new(0));
4189 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4190 FlowMirrorTask {
4191 requests: mirror_requests(&flownode_peer(), 5),
4192 }
4193 .detach(node_manager, pending.clone(), false)
4194 .unwrap();
4195 assert_eq!(
4196 tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv())
4197 .await
4198 .unwrap()
4199 .unwrap()
4200 .id,
4201 flownode_peer().id
4202 );
4203 let deadline = std::time::Instant::now() + Duration::from_secs(10);
4204 while pending.load(Ordering::Relaxed) != 0
4205 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4206 {
4207 assert!(
4208 std::time::Instant::now() < deadline,
4209 "failed task leaked reservation"
4210 );
4211 std::thread::sleep(Duration::from_millis(1));
4212 }
4213 }
4214
4215 #[tokio::test]
4216 async fn flow_mirror_pending_counter_consistent() {
4217 let _guard = lock_mirror_metrics().await;
4218 let pending = Arc::new(AtomicU64::new(0));
4219 let gate = Arc::new(tokio::sync::Semaphore::new(0));
4220 let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler {
4221 gate: gate.clone(),
4222 affected_rows: 7,
4223 }));
4224
4225 let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get();
4227 FlowMirrorTask {
4228 requests: HashMap::new(),
4229 }
4230 .detach(node_manager.clone(), pending.clone(), false)
4231 .unwrap();
4232 assert_eq!(0, pending.load(Ordering::Relaxed));
4233 assert_eq!(
4234 gauge_before,
4235 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(),
4236 "an empty batch must not touch the pending gauge"
4237 );
4238
4239 let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get();
4240 let num_rows = 7;
4241
4242 FlowMirrorTask {
4243 requests: mirror_requests(&flownode_peer(), num_rows),
4244 }
4245 .detach(node_manager, pending.clone(), false)
4246 .unwrap();
4247
4248 assert_eq!(num_rows as u64, pending.load(Ordering::Relaxed));
4251 assert_eq!(
4252 num_rows as i64,
4253 crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() - gauge_before,
4254 "the pending gauge must track the pending rows"
4255 );
4256 assert_eq!(
4257 dropped_before,
4258 crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(),
4259 "a batch within the limit must not be counted as dropped"
4260 );
4261
4262 gate.add_permits(1);
4265 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
4266 while pending.load(Ordering::Relaxed) != 0
4267 || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before
4268 {
4269 assert!(
4270 std::time::Instant::now() < deadline,
4271 "the mirror task did not release its pending reservation"
4272 );
4273 std::thread::sleep(std::time::Duration::from_millis(1));
4274 }
4275 }
4276
4277 #[test]
4278 fn mirror_logs_are_rate_limited() {
4279 let log = Arc::new(MirrorLog::new());
4280 assert_eq!(Some(1), log.claim_report(1_000));
4281 assert_eq!(None, log.claim_report(1_500));
4282 assert_eq!(None, log.claim_report(10_999));
4283 assert_eq!(Some(3), log.claim_report(11_000));
4284 assert_eq!(None, log.claim_report(11_001));
4285 assert_eq!(Some(2), log.claim_report(21_000));
4286
4287 let concurrent = Arc::new(MirrorLog::new());
4288 assert_eq!(Some(1), concurrent.claim_report(50_000));
4289 let barrier = Arc::new(std::sync::Barrier::new(8));
4290 let threads = (0..8)
4291 .map(|_| {
4292 let log = concurrent.clone();
4293 let barrier = barrier.clone();
4294 std::thread::spawn(move || {
4295 barrier.wait();
4296 log.claim_report(60_000)
4297 })
4298 })
4299 .collect::<Vec<_>>();
4300 let winners = threads
4301 .into_iter()
4302 .filter_map(|thread| thread.join().unwrap())
4303 .collect::<Vec<_>>();
4304 assert_eq!(winners.len(), 1);
4305 let residual = concurrent.events.swap(0, Ordering::Relaxed);
4306 assert_eq!(winners[0] + residual, 8);
4307 }
4308}