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