From a907cbfbe4197512d568c811c3daed99b19e436d Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Sat, 10 Oct 2026 02:45:43 +0000 Subject: [PATCH] perf: cache OTLP trace service and operation writes (#9432) * perf: cache OTLP trace service and operation writes Signed-off-by: shuiyisong * chore: fix CR issues for tests Signed-off-by: shuiyisong * fix: bound asynchronous trace auxiliary writes Signed-off-by: shuiyisong * perf: bound trace auxiliary cache by estimated memory Signed-off-by: shuiyisong * chore: size trace auxiliary cache from host or pod memory Signed-off-by: shuiyisong * chore: add metrics for cache hit and miss Signed-off-by: shuiyisong * fix: test Signed-off-by: shuiyisong * chore: unify frontend cache metrics Signed-off-by: shuiyisong * test: merge and reduce Signed-off-by: shuiyisong * chore: add init for the metrics Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- Cargo.lock | 6 + config/config.md | 2 + config/frontend.example.toml | 7 + config/standalone.example.toml | 7 + src/cmd/tests/load_config_test.rs | 23 + src/common/frontend/Cargo.toml | 2 + src/common/frontend/src/lib.rs | 1 + src/common/frontend/src/metrics.rs | 117 ++++ src/common/query/src/lib.rs | 39 +- src/frontend/AGENTS.md | 14 + src/frontend/Cargo.toml | 1 + src/frontend/src/instance.rs | 34 +- src/frontend/src/instance/builder.rs | 17 +- src/frontend/src/instance/otlp.rs | 1 + src/frontend/src/instance/otlp/README.md | 52 ++ .../src/instance/otlp/trace_ingest.rs | 346 ++++++++++- .../src/instance/otlp/trace_ingest/tests.rs | 250 +++++++- src/frontend/src/metrics.rs | 7 + src/frontend/src/service_config/otlp.rs | 44 ++ src/operator/AGENTS.md | 2 + src/operator/src/batcher.rs | 4 +- src/operator/src/insert.rs | 15 +- src/pipeline/Cargo.toml | 1 + src/pipeline/src/manager/pipeline_cache.rs | 218 +++++-- src/pipeline/src/manager/pipeline_operator.rs | 5 +- src/servers/AGENTS.md | 2 + src/servers/Cargo.toml | 1 + src/servers/src/batcher/table.rs | 74 ++- src/servers/src/mysql/handler.rs | 23 +- tests-integration/Cargo.toml | 1 + tests-integration/src/otlp.rs | 577 +++++++++++++++++- tests-integration/tests/http.rs | 1 + 32 files changed, 1784 insertions(+), 110 deletions(-) create mode 100644 src/common/frontend/src/metrics.rs diff --git a/Cargo.lock b/Cargo.lock index c02c7fdcc60..9b447169fbc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2459,7 +2459,9 @@ dependencies = [ "common-macro", "common-meta", "greptime-proto", + "lazy_static", "meta-client", + "prometheus 0.14.0", "serde", "serde_json", "snafu 0.8.6", @@ -5604,6 +5606,7 @@ dependencies = [ "lazy_static", "log-query", "meta-client", + "moka", "opentelemetry-proto", "opentelemetry-semantic-conventions 0.31.0", "operator", @@ -10860,6 +10863,7 @@ dependencies = [ "chrono-tz", "common-catalog", "common-error", + "common-frontend", "common-macro", "common-meta", "common-query", @@ -13522,6 +13526,7 @@ dependencies = [ "common-catalog", "common-decimal", "common-error", + "common-frontend", "common-grpc", "common-macro", "common-mem-prof", @@ -15122,6 +15127,7 @@ dependencies = [ "store-api", "substrait 1.4.0-alpha.0", "table", + "temp-env", "tempfile", "time", "tokio", diff --git a/config/config.md b/config/config.md index c538559e7cf..612a5cd43c0 100644 --- a/config/config.md +++ b/config/config.md @@ -95,6 +95,7 @@ | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.trace_aux_cache_size` | String | Auto | Estimated memory budget for cached trace service/operation keys per frontend,
shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching.
Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB.
Uses 32MiB if memory detection is unavailable.
Excludes cache and allocator overhead; shared table names are charged per entry. | | `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | @@ -368,6 +369,7 @@ | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.trace_aux_cache_size` | String | Auto | Estimated memory budget for cached trace service/operation keys per frontend,
shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching.
Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB.
Uses 32MiB if memory detection is unavailable.
Excludes cache and allocator overhead; shared table names are charged per entry. | | `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 21e7bbb40f3..e4ef6e36935 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -310,6 +310,13 @@ enable = true enable = true ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Estimated memory budget for cached trace service/operation keys per frontend, +## shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching. +## Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB. +## Uses 32MiB if memory detection is unavailable. +## Excludes cache and allocator overhead; shared table names are charged per entry. +## @toml2docs:none-default="Auto" +#+ trace_aux_cache_size = "32MiB" ## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric ## resource attributes, so metrics-only services reach the semantic graph. experimental_enable_resource_info = false diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 7e1f8dcf6b0..c656049c5d2 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -282,6 +282,13 @@ enable = true enable = true ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Estimated memory budget for cached trace service/operation keys per frontend, +## shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching. +## Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB. +## Uses 32MiB if memory detection is unavailable. +## Excludes cache and allocator overhead; shared table names are charged per entry. +## @toml2docs:none-default="Auto" +#+ trace_aux_cache_size = "32MiB" ## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric ## resource attributes, so metrics-only services reach the semantic graph. experimental_enable_resource_info = false diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 1f4c69c8938..00cdafd9b8a 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -731,6 +731,29 @@ fn test_example_configs_document_actual_defaults() { assert!(errors.is_empty(), "{}", errors.join("\n")); } +#[test] +fn test_load_trace_aux_cache_size() { + let default = FrontendOptions::default().otlp.trace_aux_cache_size; + for (setting, size) in [ + ("", default), + ("trace_aux_cache_size = 0", ReadableSize(0)), + ("trace_aux_cache_size = \"1MiB\"", ReadableSize::mb(1)), + ] { + let config = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(config.path(), format!("[otlp]\n{setting}\n")).unwrap(); + + let frontend = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + let standalone = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + for options in [frontend.component, standalone.component.frontend_options()] { + assert_eq!(options.otlp.trace_aux_cache_size, size); + } + } +} + #[test] fn test_load_removed_histogram_options() { for enabled in [false, true] { diff --git a/src/common/frontend/Cargo.toml b/src/common/frontend/Cargo.toml index 26092ae13f7..84a8373dcbb 100644 --- a/src/common/frontend/Cargo.toml +++ b/src/common/frontend/Cargo.toml @@ -13,7 +13,9 @@ common-grpc.workspace = true common-macro.workspace = true common-meta.workspace = true greptime-proto.workspace = true +lazy_static.workspace = true meta-client.workspace = true +prometheus.workspace = true serde.workspace = true serde_json.workspace = true snafu.workspace = true diff --git a/src/common/frontend/src/lib.rs b/src/common/frontend/src/lib.rs index dacffd19968..b6ec72da1a3 100644 --- a/src/common/frontend/src/lib.rs +++ b/src/common/frontend/src/lib.rs @@ -18,6 +18,7 @@ use std::str::FromStr; use snafu::OptionExt; pub mod error; +pub mod metrics; pub mod selector; pub mod slow_query_event; diff --git a/src/common/frontend/src/metrics.rs b/src/common/frontend/src/metrics.rs new file mode 100644 index 00000000000..7a2c152fa1f --- /dev/null +++ b/src/common/frontend/src/metrics.rs @@ -0,0 +1,117 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Cache lookup counters shared by frontend, pipeline, and protocol handlers. +//! +//! The `type` label identifies `otlp_trace_aux`, `otlp_metrics_legacy`, +//! `pipeline_table`, `pipeline`, `pipeline_original`, `pipeline_failover`, or +//! `mysql_prepared_stmt`. Counters aggregate across catalogs and connections. +//! Trace lookups count distinct service/operation keys once per request, excluding +//! rechecks; legacy metric lookups count each requested metric name. Pipeline +//! loaders count successful coalesced waiters as hits and failed loads as misses. +//! Failover counts one outcome after resolving all candidate schemas; an ambiguous +//! result is a miss. Table and prepared-statement caches count each lookup. +//! Both outcomes for all known caches are exported from the first metric update. + +use lazy_static::lazy_static; +use prometheus::{IntCounterVec, register_int_counter_vec}; + +const CACHE_TYPES: [&str; 7] = [ + "otlp_trace_aux", + "otlp_metrics_legacy", + "pipeline_table", + "pipeline", + "pipeline_original", + "pipeline_failover", + "mysql_prepared_stmt", +]; + +lazy_static! { + static ref CACHE_COUNTERS: (IntCounterVec, IntCounterVec) = { + let hits = register_int_counter_vec!( + "greptime_frontend_cache_hit", + "frontend cache hit", + &["type"] + ) + .unwrap(); + let misses = register_int_counter_vec!( + "greptime_frontend_cache_miss", + "frontend cache miss", + &["type"] + ) + .unwrap(); + // Register both outcomes together so an all-hit or all-miss cache has both series. + for cache_type in CACHE_TYPES { + hits.with_label_values(&[cache_type]); + misses.with_label_values(&[cache_type]); + } + (hits, misses) + }; + + /// Frontend cache hits by cache type. + pub static ref CACHE_HIT: IntCounterVec = CACHE_COUNTERS.0.clone(); + /// Frontend cache misses by cache type. + pub static ref CACHE_MISS: IntCounterVec = CACHE_COUNTERS.1.clone(); +} + +/// Records one completed cache lookup. +pub fn record_cache_lookup(cache_type: &str, hit: bool) { + if hit { + CACHE_HIT.with_label_values(&[cache_type]).inc(); + } else { + CACHE_MISS.with_label_values(&[cache_type]).inc(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_cache_lookup_counters() { + record_cache_lookup("pipeline", false); + // Scrape before accessing CACHE_HIT, which would initialize it and hide a missing series. + let families = prometheus::gather(); + for (name, expected) in [ + ("greptime_frontend_cache_hit", 0.0), + ("greptime_frontend_cache_miss", 1.0), + ] { + let family = families + .iter() + .find(|family| family.name() == name) + .unwrap(); + assert_eq!(family.get_metric().len(), CACHE_TYPES.len()); + for cache_type in CACHE_TYPES { + let metric = family + .get_metric() + .iter() + .find(|metric| metric.get_label()[0].value() == cache_type) + .unwrap(); + assert_eq!(metric.get_label()[0].name(), "type"); + assert_eq!( + metric.get_counter().value(), + if cache_type == "pipeline" { + expected + } else { + 0.0 + } + ); + } + } + record_cache_lookup("pipeline", true); + record_cache_lookup("pipeline", false); + assert_eq!(CACHE_HIT.with_label_values(&["pipeline"]).get(), 1); + assert_eq!(CACHE_MISS.with_label_values(&["pipeline"]).get(), 2); + } +} diff --git a/src/common/query/src/lib.rs b/src/common/query/src/lib.rs index 8ce028ef9fe..8151dfc5731 100644 --- a/src/common/query/src/lib.rs +++ b/src/common/query/src/lib.rs @@ -25,10 +25,13 @@ pub mod stream; pub mod test_util; use std::fmt::{Debug, Display, Formatter}; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use api::greptime_proto::v1::AddColumnLocation as Location; use api::greptime_proto::v1::add_column_location::LocationType; +use common_error::ext::BoxedError; use common_recordbatch::{ RecordBatches, SendableRecordBatchMapper, SendableRecordBatchStream, map_dictionary_to_values, map_dictionary_to_values_schema, @@ -74,6 +77,30 @@ pub struct OutputMeta { /// May exist for query output. One can retrieve execution metrics from this plan. pub plan: Option>, pub cost: OutputCost, + /// Process-local storage results attached by asynchronous ordinary-table batching. + /// Callers needing confirmed writes must await these. Dropping them does not cancel writes. + pub write_completions: Vec, +} + +/// Actual storage completion of an already accepted asynchronous write. +pub struct WriteCompletion(Pin> + Send>>); + +impl WriteCompletion { + /// Wraps the storage result without changing how the write is scheduled. + pub fn new(completion: impl Future> + Send + 'static) -> Self { + Self(Box::pin(completion)) + } + + /// Waits for storage success or failure, independently of queue admission. + pub async fn wait(self) -> Result<(), BoxedError> { + self.0.await + } +} + +impl Debug for WriteCompletion { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WriteCompletion").finish_non_exhaustive() + } } impl Output { @@ -163,18 +190,26 @@ impl Debug for OutputData { impl OutputMeta { pub fn new(plan: Option>, cost: usize) -> Self { - Self { plan, cost } + Self { + plan, + cost, + ..Default::default() + } } pub fn new_with_plan(plan: Arc) -> Self { Self { plan: Some(plan), cost: 0, + ..Default::default() } } pub fn new_with_cost(cost: usize) -> Self { - Self { plan: None, cost } + Self { + cost, + ..Default::default() + } } } diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index bab27676655..fdaa5dbcfe4 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -34,6 +34,9 @@ remote datanodes via `operator`/`client`. | `heartbeat` | `src/frontend/src/heartbeat.rs` | Heartbeat to metasrv; handles suspend / cache invalidation | | `service_config` | `src/frontend/src/service_config/` | Per-protocol option structs | +Shared frontend cache counters live in `src/common/frontend/src/metrics.rs`, +including caches owned by pipeline and protocol handlers. + ## Request lifecycles - **SQL query** (`instance.rs`): `do_query_inner` handles parsing, interceptors, @@ -58,6 +61,17 @@ remote datanodes via `operator`/`client`. the lazy schema handshake; checks permissions and reconciles missing columns through `Inserter` once per stream, then reuses the refreshed table. +- **OTLP trace lookup tables** (`instance/otlp/trace_ingest.rs`): `Instance` owns + a shared cache of confirmed service/operation writes, which bypass batching. + `otlp.trace_aux_cache_size` sets its estimated byte budget at startup; zero disables it. + Async main writes retain completion results; a frontend task writes auxiliary + rows only for confirmed chunks before populating the cache. + Cold async trace requests reserve a shared auxiliary slot before main writes, + bounded separately by the table batcher's `max_inflight_requests`; confirmed + cache hits bypass this admission. The slot lasts through auxiliary completion. + These tables are ingestion-managed; manual mutation requires restarting + serving frontends. + - **Logical-table batching** (`instance/logical_batcher.rs`): `Services` initializes one shared batcher for opted-in HTTP Prom and nonlegacy OTLP metric-engine writes. Downstream routers can initialize it when enabling replacement endpoints. diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index 1aa057d0fcd..92803933a8b 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -58,6 +58,7 @@ humantime-serde.workspace = true lazy_static.workspace = true log-query.workspace = true meta-client.workspace = true +moka = { workspace = true, features = ["sync"] } opentelemetry-proto.workspace = true opentelemetry-semantic-conventions = { version = "0.31", features = ["semconv_experimental"] } operator.workspace = true diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index d3e6f2addfc..b91290d2a55 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -48,8 +48,10 @@ use catalog::process_manager::{ use client::OutputData; use common_base::Plugins; use common_base::cancellation::CancellableFuture; +use common_batcher::request_limiter::RequestLimiter; use common_error::ext::{BoxedError, ErrorExt}; use common_event_recorder::EventRecorderRef; +use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::key::TableMetadataManagerRef; @@ -141,6 +143,8 @@ pub struct Instance { slow_query_options: SlowQueryOptions, influxdb_default_merge_mode: InfluxdbMergeMode, trace_ingest_chunk_size: usize, + trace_aux_cache: otlp::TraceAuxCache, + trace_aux_limiter: Option, otlp_resource_info: bool, suspend: Arc, @@ -148,7 +152,7 @@ pub struct Instance { // first layer key: db-string // key: direct input metric name // value: if runs in legacy mode - otlp_metrics_table_legacy_cache: DashMap>, + otlp_metrics_table_legacy_cache: Arc>>, } impl Instance { @@ -570,6 +574,12 @@ fn fast_legacy_check( .iter() .filter_map(|name| cache.get(name)) .collect::>(); + CACHE_HIT + .with_label_values(&["otlp_metrics_legacy"]) + .inc_by(hit_cache.len() as u64); + CACHE_MISS + .with_label_values(&["otlp_metrics_legacy"]) + .inc_by((names.len() - hit_cache.len()) as u64); if !hit_cache.is_empty() { let hit_legacy = hit_cache.iter().any(|en| *en.value()); let hit_prom = hit_cache.iter().any(|en| !*en.value()); @@ -3418,12 +3428,16 @@ mod tests { } #[tokio::test] - async fn test_event_recorder_is_exposed() -> TestResult<()> { + async fn test_instance_shared_state() -> TestResult<()> { let instance = test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?) .await?; let _event_recorder = instance.event_recorder(); + assert!(Arc::ptr_eq( + &instance.otlp_metrics_table_legacy_cache, + &instance.clone().otlp_metrics_table_legacy_cache, + )); Ok(()) } @@ -3554,11 +3568,20 @@ mod tests { })?; let ctx = test_query_ctx(1); let handler: PipelineHandlerRef = Arc::new(instance.clone()); + let hits = CACHE_HIT.with_label_values(&["pipeline_table"]); + let misses = CACHE_MISS.with_label_values(&["pipeline_table"]); + let before = (hits.get(), misses.get()); handler .get_pipeline("pipeline", None, ctx.clone()) .await .unwrap(); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 1)); + handler + .get_pipeline("pipeline", None, ctx.clone()) + .await + .unwrap(); + assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1)); assert_permission_denied( PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await, ); @@ -3877,11 +3900,15 @@ mod tests { #[test] fn test_fast_legacy_check_is_read_only() { + let hits = CACHE_HIT.with_label_values(&["otlp_metrics_legacy"]); + let misses = CACHE_MISS.with_label_values(&["otlp_metrics_legacy"]); + let before = (hits.get(), misses.get()); let cache = DashMap::new(); cache.insert("metric1".to_string(), true); let names = vec!["metric1".to_string(), "metric2".to_string()]; assert_eq!(Some(true), fast_legacy_check(&cache, &names).unwrap()); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 1)); assert!(!cache.contains_key("metric2")); cache_legacy_mode(&cache, &names, true).unwrap(); @@ -3893,6 +3920,9 @@ mod tests { cache_incompatible.insert("metric1".to_string(), true); cache_incompatible.insert("metric2".to_string(), false); assert!(fast_legacy_check(&cache_incompatible, &names).is_err()); + assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1)); + assert_eq!(None, fast_legacy_check(&cache, &[]).unwrap()); + assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1)); } #[test] diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index d261daed48e..032722499f9 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -22,6 +22,7 @@ use catalog::kvbackend::KvBackendCatalogManager; use catalog::process_manager::ProcessManagerRef; use catalog::system_schema::semantic_graph::EntityGraphProviderRef; use common_base::Plugins; +use common_batcher::request_limiter::RequestLimiter; use common_datasource::object_store::LocalFileAccess; use common_event_recorder::{EventRecorderImpl, EventRecorderRef}; use common_meta::cache::{LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableRouteCacheRef}; @@ -50,8 +51,8 @@ use partition::manager::PartitionRuleManager; use pipeline::pipeline_operator::PipelineOperator; use query::QueryEngineFactory; use query::region_query::RegionQueryHandlerFactoryRef; -use servers::batcher::BatchingProtocol; use servers::batcher::table::TablePendingRowsBatcher; +use servers::batcher::{BatchingProtocol, pending_rows_batch_sync_enabled}; use snafu::{OptionExt, ResultExt}; use crate::error::{self, DataFusionSnafu, ExternalSnafu, Result}; @@ -60,6 +61,7 @@ use crate::frontend::FrontendOptions; use crate::heartbeat::frontend_peer_addr; use crate::instance::Instance; use crate::instance::entity_graph::EntityGraphProviderImpl; +use crate::instance::otlp::TraceAuxCache; use crate::instance::region_query::FrontendRegionQueryHandler; use crate::service_config::BatcherOptions; @@ -385,6 +387,15 @@ impl FrontendBuilder { )); admin_event_recorder.install(&event_recorder); + let batcher_options = self.options.table_batcher_options(); + let trace_aux_limiter = if batcher_options.pending_rows_batching_enabled() + && !pending_rows_batch_sync_enabled() + { + RequestLimiter::try_new(batcher_options.max_inflight_requests) + } else { + None + }; + Ok(Instance { logical_batcher: Default::default(), frontend_peer_addr, @@ -400,10 +411,12 @@ impl FrontendBuilder { event_recorder, slow_query_recorder, process_manager, - otlp_metrics_table_legacy_cache: DashMap::new(), + otlp_metrics_table_legacy_cache: Arc::new(DashMap::new()), slow_query_options: self.options.slow_query.clone(), influxdb_default_merge_mode: self.options.influxdb.default_merge_mode, trace_ingest_chunk_size: self.options.otlp.trace_ingest_chunk_size, + trace_aux_cache: TraceAuxCache::new(self.options.otlp.trace_aux_cache_size.as_bytes()), + trace_aux_limiter, otlp_resource_info: self.options.otlp.experimental_enable_resource_info, suspend: Arc::new(AtomicBool::new(false)), }) diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 5d96138319a..e4f395db457 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -52,6 +52,7 @@ use table::requests::{ SOURCE_OPENTELEMETRY, }; +pub(crate) use self::trace_ingest::TraceAuxCache; use self::trace_ingest::trace_conventions; use crate::instance::Instance; use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS}; diff --git a/src/frontend/src/instance/otlp/README.md b/src/frontend/src/instance/otlp/README.md index c123f1caae1..e308ebec184 100644 --- a/src/frontend/src/instance/otlp/README.md +++ b/src/frontend/src/instance/otlp/README.md @@ -175,6 +175,58 @@ Accounting follows the main-table write: accepted or rejected span counts; and - failure details are bounded before they are folded into `TraceIngestOutcome`. +A cache shared by each frontend's `Instance` clones skips previously confirmed +auxiliary writes for v0, v1, and v2. `otlp.trace_aux_cache_size` sets an estimated +memory budget across all catalogs, schemas, and trace tables on each frontend. +By default, the budget is 1/128 of the pod's cgroup memory limit or total host +memory, with a minimum of `32MiB`. If memory detection is unavailable, it uses +`32MiB`. An explicit size overrides this default; `0` disables caching. +The estimate includes key structs and allocated string capacity, charging shared +table names per entry. It excludes cache and allocator overhead, so it is not a +strict process-memory limit. Changes take effect on frontend or standalone +restart. There is no time-based expiry. +Services and operations are cached independently. Only a successful +auxiliary write populates the cache. Auxiliary writes always bypass batching, +including when main spans use asynchronous batching. Failures, eviction, +frontend restarts, and concurrent misses can cause repeat writes. Write cost +reflects only the writes actually performed. + +`greptime_frontend_cache_hit{type="otlp_trace_aux"}` and +`greptime_frontend_cache_miss{type="otlp_trace_aux"}` count distinct service and +operation keys at the first cache check in each request. +Requests with auxiliary admission count before acquiring a slot; other requests +count when preparing auxiliary writes. Admission and write rechecks do not count. +When the cache is disabled, all checked keys count as misses. + +Successful auxiliary writes with `skip_wal` enabled are cached too. If a +datanode crashes before flushing those rows while the frontend survives, the +cache can suppress their recreation, even by later WAL-enabled requests, until +eviction or frontend restart. Missing service/operation discovery rows in this +case are an accepted limitation; the cache does not guarantee durability beyond +the write's WAL policy. + +With asynchronous main-table batching, the response still acknowledges queue +admission. A frontend task waits for the storage results and writes auxiliary +rows only for successful chunks. Failed chunks do not populate the auxiliary +tables or cache. Deferred auxiliary failures are logged; they cannot change an +already returned response. Deferred writes retain the original request's row +admission, and their write cost is not included in the early response. + +Async trace requests with cache misses acquire one shared auxiliary-work slot +before submitting main chunks. Each frontend limits these requests separately +from main-table batching, using `pending_rows_batcher.max_inflight_requests`. +The slot is held through auxiliary completion, including failures. Requests +whose service/operation entries are all cached bypass this admission and need +no deferred task. Cache lookups and updates hold no lock across storage writes. + +The auxiliary tables are ingestion-managed. When caching is enabled, manual +`DROP`, `TRUNCATE`, `DELETE`, or other mutations require clearing the affected +caches or restarting all serving frontends before relying on ingestion to +populate the lookup rows again. +There is no automatic invalidation or public cache-clear command. In v0, fewer +duplicate rows are appended; distinct service/operation lookup results remain +the same. + A rejection detail carries the failing cause, not just a status code, so an unusable attribute value names its column, source value, and target type. Repeats of the same `(site, cause)` collapse into one entry with an occurrence count, and diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs index e5b16c7d191..8b90a19ba49 100644 --- a/src/frontend/src/instance/otlp/trace_ingest.rs +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::{BTreeMap, HashMap, HashSet}; +use std::mem; use std::sync::Arc; use api::helper::ColumnDataTypeWrapper; @@ -22,13 +23,17 @@ use api::v1::{ RowInsertRequest, RowInsertRequests, Rows, Value, }; use client::Output; +use common_batcher::request_limiter::RequestLimiter; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; use common_meta::rpc::ddl::TriggerReason; +use common_query::WriteCompletion; use common_telemetry::{debug, warn}; use datatypes::prelude::ConcreteDataType; +use moka::sync::Cache; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use pipeline::PipelineWay; +use prometheus::IntCounter; use servers::error::{self, Result as ServerResult}; use servers::otlp; use servers::otlp::coerce::{coerce_value_data, is_supported_trace_coercion, trace_value_datatype}; @@ -43,6 +48,8 @@ use table::requests::{ SEMANTIC_TRACE_CONVENTIONS, SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, SIGNAL_TYPE_TRACE, SOURCE_OPENTELEMETRY, TABLE_DATA_MODEL_TRACE_V1, TABLE_DATA_MODEL_TRACE_V2, }; +use table::table_name::TableName; +use tokio::sync::OwnedSemaphorePermit; use crate::instance::Instance; use crate::instance::otlp::trace_semconv::trace_semconv_fixed_type; @@ -52,7 +59,131 @@ use crate::instance::otlp::trace_types::{ is_trace_reconcile_candidate_type, prepare_trace_column_rewrites, push_observed_trace_type, truncate_for_diagnostics, }; -use crate::metrics::{OTLP_TRACES_FAILURE_COUNT, OTLP_TRACES_ROWS}; +use crate::metrics::{ + OTLP_TRACE_AUX_CACHE_HIT, OTLP_TRACE_AUX_CACHE_MISS, OTLP_TRACES_FAILURE_COUNT, + OTLP_TRACES_ROWS, +}; + +/// Confirmed auxiliary writes, shared across requests and instance clones. +/// +/// ponytail: these tables are ingestion-managed; add invalidation if external +/// mutations must be supported without restarting every serving frontend. +#[derive(Clone)] +pub(crate) struct TraceAuxCache { + entries: Cache, +} + +#[derive(Clone, Hash, PartialEq, Eq)] +struct TraceAuxCacheKey { + table: Arc, + entry: TraceAuxEntry, +} + +impl TraceAuxCacheKey { + /// Estimated key memory, excluding cache and allocator overhead. + fn size(&self) -> usize { + let entry_size = match &self.entry { + TraceAuxEntry::Service(service) => service.capacity(), + TraceAuxEntry::Operation(service, span, kind) => { + service.capacity() + span.capacity() + kind.capacity() + } + }; + // Charge the shared table name per entry to keep the estimate conservative. + mem::size_of::() + + mem::size_of::() + + self.table.catalog_name.capacity() + + self.table.schema_name.capacity() + + self.table.table_name.capacity() + + entry_size + } +} + +#[derive(Clone, Hash, PartialEq, Eq)] +enum TraceAuxEntry { + Service(String), + Operation(String, String, String), +} + +impl TraceAuxCache { + /// Creates a cache with an estimated byte budget; zero disables caching. + pub(crate) fn new(max_size: u64) -> Self { + Self { + entries: Cache::builder() + .max_capacity(max_size) + .weigher(|key: &TraceAuxCacheKey, _: &()| { + u32::try_from(key.size()).unwrap_or(u32::MAX) + }) + .build(), + } + } + + /// Removes confirmed writes and returns keys to remember after insertion. + /// Only the first cache check in a request supplies metrics. + fn filter( + &self, + table: TableName, + aux_data: &mut TraceAuxData, + metrics: Option<(&IntCounter, &IntCounter)>, + ) -> Vec { + let lookups = aux_data.services.len() + aux_data.operations.len(); + let table = Arc::new(table); + let mut pending = Vec::new(); + let mut retain = |entry| { + let key = TraceAuxCacheKey { + table: table.clone(), + entry, + }; + if self.entries.get(&key).is_some() { + false + } else { + pending.push(key); + true + } + }; + aux_data + .services + .retain(|service| retain(TraceAuxEntry::Service(service.clone()))); + aux_data.operations.retain(|(service, span, kind)| { + retain(TraceAuxEntry::Operation( + service.clone(), + span.clone(), + kind.clone(), + )) + }); + if let Some((hits, misses)) = metrics { + hits.inc_by((lookups - pending.len()) as u64); + misses.inc_by(pending.len() as u64); + } + pending + } + + fn record(&self, keys: Vec) { + for key in keys { + self.entries.insert(key, ()); + } + } + + /// Reserves background capacity only for misses, without holding a cache lock. + async fn acquire_for_misses( + &self, + limiter: &RequestLimiter, + table: TableName, + mut aux_data: TraceAuxData, + metrics: (&IntCounter, &IntCounter), + ) -> ServerResult>> { + self.filter(table.clone(), &mut aux_data, Some(metrics)); + if aux_data.is_empty() { + return Ok(None); + } + let permit = limiter + .acquire() + .await + .map_err(|_| error::BatcherChannelClosedSnafu.build())?; + // Another request may have confirmed these entries while admission waited. + self.filter(table, &mut aux_data, None); + Ok((!aux_data.is_empty()).then_some(permit)) + } +} /// Merge converted V2 rows only when their fixed schemas agree. fn merge_trace_v2_rows(rows: &mut Rows, batch: Rows) -> ServerResult<()> { @@ -101,11 +232,52 @@ struct TraceChunkIngestContext<'a> { /// Accumulates trace outcomes, auxiliary rows, and bounded failure details. struct TraceIngestState { + needs_aux_write: bool, aux_data: TraceAuxData, + pending_aux_data: Vec<(TraceAuxData, Vec)>, outcome: TraceIngestOutcome, failure_messages: TraceFailureMessages, } +impl TraceIngestState { + fn add_aux_data(&mut self, aux_data: TraceAuxData, completions: Vec) { + if completions.is_empty() { + self.aux_data.services.extend(aux_data.services); + self.aux_data.operations.extend(aux_data.operations); + } else if !aux_data.is_empty() { + self.pending_aux_data.push((aux_data, completions)); + } + } +} + +/// Keeps confirmed chunks even if another chunk fails or loses its completion sender. +async fn confirmed_trace_aux_data( + pending: Vec<(TraceAuxData, Vec)>, +) -> TraceAuxData { + let results = futures::future::join_all(pending.into_iter().map( + |(aux_data, completions)| async move { + for completion in completions { + completion.wait().await?; + } + Ok::<_, BoxedError>(aux_data) + }, + )) + .await; + let mut confirmed = TraceAuxData::default(); + for result in results { + match result { + Ok(aux_data) => { + confirmed.services.extend(aux_data.services); + confirmed.operations.extend(aux_data.operations); + } + Err(err) => { + warn!(err; "Skipping auxiliary trace rows after asynchronous main write failure") + } + } + } + confirmed +} + /// Bounded, deduplicated failure details for one trace request. /// /// Occurrences past [`TRACE_FAILURE_MESSAGE_LIMIT`] distinct failures are @@ -876,6 +1048,32 @@ impl Instance { conventions: &str, ctx: QueryContextRef, ) -> ServerResult { + let aux_limiter = self + .trace_aux_limiter + .as_ref() + .filter(|_| ctx.batching_enabled()); + // Reserve once per request, before any main chunks are submitted. Main + // batcher permits remain independent so multi-chunk requests can progress. + let aux_permit = if let Some(limiter) = aux_limiter { + let mut aux_data = TraceAuxData::default(); + for span in groups.iter().flat_map(|group| &group.spans) { + aux_data.observe_span(span); + } + self.trace_aux_cache + .acquire_for_misses( + limiter, + TableName::new(ctx.current_catalog(), ctx.current_schema(), &table_name), + aux_data, + (&OTLP_TRACE_AUX_CACHE_HIT, &OTLP_TRACE_AUX_CACHE_MISS), + ) + .await? + } else { + None + }; + // Requests without auxiliary admission first check the cache when preparing writes. + let aux_cache_metrics = aux_limiter + .is_none() + .then_some((&*OTLP_TRACE_AUX_CACHE_HIT, &*OTLP_TRACE_AUX_CACHE_MISS)); let is_trace_v1_model = matches!(pipeline, PipelineWay::OtlpTraceDirectV1); let data_model = match pipeline { PipelineWay::OtlpTraceDirectV1 => Some(TABLE_DATA_MODEL_TRACE_V1), @@ -918,7 +1116,11 @@ impl Instance { table_name: &table_name, }; let mut ingest_state = TraceIngestState { + // Confirmed hits stay satisfied for this request even if evicted + // later: cache eviction does not delete the stored lookup rows. + needs_aux_write: aux_limiter.is_none() || aux_permit.is_some(), aux_data: TraceAuxData::default(), + pending_aux_data: Vec::new(), outcome: TraceIngestOutcome::default(), failure_messages: TraceFailureMessages::default(), }; @@ -988,20 +1190,64 @@ impl Instance { OTLP_TRACES_ROWS.inc_by(ingest_state.outcome.accepted_spans as u64); + if !ingest_state.pending_aux_data.is_empty() { + let pending = std::mem::take(&mut ingest_state.pending_aux_data); + let instance = self.clone(); + let pipeline = match pipeline { + PipelineWay::OtlpTraceDirectV1 => PipelineWay::OtlpTraceDirectV1, + PipelineWay::OtlpTraceDirectV2 => PipelineWay::OtlpTraceDirectV2, + _ => PipelineWay::OtlpTraceDirectV0, + }; + let table_name = table_name.clone(); + let ctx = ctx.clone(); + common_runtime::spawn_ingest(async move { + let _permit = aux_permit; + let aux_data = confirmed_trace_aux_data(pending).await; + let result: ServerResult<()> = async { + let (requests, keys) = instance.prepare_trace_aux_write( + aux_data, + &pipeline, + &table_name, + &ctx, + aux_cache_metrics, + )?; + if !requests.inserts.is_empty() { + instance + .insert_trace_aux_requests(requests, keys, &pipeline, ctx) + .await?; + } + Ok(()) + } + .await; + if let Err(err) = result { + OTLP_TRACES_FAILURE_COUNT + .with_label_values(&["aux_table_update_failed"]) + .inc(); + warn!(table_name, err = ?err, "Failed to write auxiliary trace rows after asynchronous main write"); + } + }); + } + if !ingest_state.aux_data.is_empty() { - // Auxiliary trace tables are derived from spans whose main-table - // writes are already confirmed, so they never create new accepted - // spans and they do not affect rejected span counts. - let aux_requests = otlp::trace::to_grpc_insert_requests_for_aux_tables( + // Auxiliary trace tables are derived from accepted spans, so they + // never create new accepted spans or affect rejected span counts. + let aux_requests = self.prepare_trace_aux_write( std::mem::take(&mut ingest_state.aux_data), ingest_ctx.pipeline, ingest_ctx.table_name, + &ctx, + aux_cache_metrics, ); match aux_requests { - Ok((aux_requests, _)) if !aux_requests.inserts.is_empty() => { + Ok((aux_requests, pending_aux_keys)) if !aux_requests.inserts.is_empty() => { match self - .insert_trace_requests(aux_requests, ingest_ctx.pipeline, ctx) + .insert_trace_aux_requests( + aux_requests, + pending_aux_keys, + ingest_ctx.pipeline, + ctx, + ) .await { Ok(output) => { @@ -1076,6 +1322,72 @@ impl Instance { Ok(ingest_state.outcome) } + fn prepare_trace_aux_write( + &self, + mut aux_data: TraceAuxData, + pipeline: &PipelineWay, + table_name: &str, + ctx: &QueryContextRef, + metrics: Option<(&IntCounter, &IntCounter)>, + ) -> ServerResult<(RowInsertRequests, Vec)> { + let keys = self.trace_aux_cache.filter( + TableName::new(ctx.current_catalog(), ctx.current_schema(), table_name), + &mut aux_data, + metrics, + ); + if aux_data.is_empty() { + return Ok((RowInsertRequests::default(), keys)); + } + let (mut requests, _) = + otlp::trace::to_grpc_insert_requests_for_aux_tables(aux_data, pipeline, table_name)?; + // Remove fully cached tables before v1 reconciliation performs catalog lookups. + requests.inserts.retain(|request| { + request + .rows + .as_ref() + .is_some_and(|rows| !rows.rows.is_empty()) + }); + Ok((requests, keys)) + } + + /// Checks deferred cache population without confusing visible rows with completion. + #[cfg(feature = "testing")] + pub fn trace_aux_data_cached_for_test( + &self, + table_name: &str, + mut aux_data: TraceAuxData, + ctx: &QueryContextRef, + ) -> bool { + self.trace_aux_cache + .filter( + TableName::new(ctx.current_catalog(), ctx.current_schema(), table_name), + &mut aux_data, + None, + ) + .is_empty() + } + + async fn insert_trace_aux_requests( + &self, + requests: RowInsertRequests, + keys: Vec, + pipeline: &PipelineWay, + ctx: QueryContextRef, + ) -> ServerResult { + // Cache entries require a completed write, even when main spans use + // batching with early acknowledgement. + let mut aux_ctx = ctx.fork(); + aux_ctx.set_batching_enabled(false); + let output = self + .insert_trace_requests(requests, pipeline, Arc::new(aux_ctx)) + .await?; + // ponytail: cache skip_wal writes too. After datanode data loss, lookup + // rows may stay missing until eviction/restart; require durable writes + // before caching if that limitation becomes unacceptable. + self.trace_aux_cache.record(keys); + Ok(output) + } + /// Ingest one owned trace chunk so successful spans can be moved into the /// accepted set without extra cloning. async fn ingest_trace_chunk( @@ -1167,8 +1479,12 @@ impl Instance { Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); ingest_state.outcome.accepted_spans += chunk_rows; - for span in &chunk { - ingest_state.aux_data.observe_span(span); + if ingest_state.needs_aux_write { + let mut aux_data = TraceAuxData::default(); + for span in &chunk { + aux_data.observe_span(span); + } + ingest_state.add_aux_data(aux_data, output.meta.write_completions); } Ok(()) @@ -1345,7 +1661,11 @@ impl Instance { Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); ingest_state.outcome.accepted_spans += span_count; - retry.add_to_aux_data(&mut ingest_state.aux_data); + if ingest_state.needs_aux_write { + let mut aux_data = TraceAuxData::default(); + retry.add_to_aux_data(&mut aux_data); + ingest_state.add_aux_data(aux_data, output.meta.write_completions); + } Ok(()) } @@ -1428,7 +1748,11 @@ impl Instance { Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); ingest_state.outcome.accepted_spans += 1; - span.add_to_aux_data(&mut ingest_state.aux_data); + if ingest_state.needs_aux_write { + let mut aux_data = TraceAuxData::default(); + span.add_to_aux_data(&mut aux_data); + ingest_state.add_aux_data(aux_data, output.meta.write_completions); + } } Ok(()) diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index 1459119efdd..d497db83e13 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::{HashMap, HashSet}; +use std::sync::Arc; use api::v1::column_data_type_extension::TypeExt; use api::v1::value::ValueData; @@ -20,26 +21,269 @@ use api::v1::{ ColumnDataType, ColumnDataTypeExtension, ColumnSchema, JsonTypeExtension, Row, RowInsertRequest, Rows, SemanticType, Value, }; +use common_batcher::request_limiter::RequestLimiter; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; +use common_query::WriteCompletion; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ ColumnSchema as DatatypesColumnSchema, SchemaBuilder as DatatypesSchemaBuilder, }; use opentelemetry_proto::tonic::common::v1::any_value::Value as OtlpValue; use opentelemetry_proto::tonic::common::v1::{AnyValue, ArrayValue, KeyValue}; -use servers::otlp::trace::SERVICE_NAME_COLUMN; +use prometheus::IntCounter; use servers::otlp::trace::attributes::Attributes; use servers::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; use servers::otlp::trace::v1::{TraceBinaryType, TraceRetryColumn}; +use servers::otlp::trace::{SERVICE_NAME_COLUMN, TraceAuxData}; +use table::table_name::TableName; use super::{ - ChunkFailureReaction, Instance, TraceChunkRetry, TraceChunkSchemaState, TraceFailureMessages, - TraceRequestSchema, TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, + ChunkFailureReaction, Instance, TraceAuxCache, TraceAuxCacheKey, TraceAuxEntry, + TraceChunkRetry, TraceChunkSchemaState, TraceFailureMessages, TraceRequestSchema, + TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, wrap_trace_alter_failure, }; use crate::metrics::OTLP_TRACES_FAILURE_COUNT; +fn trace_aux_data(service: &str, span: &str, kind: &str) -> TraceAuxData { + TraceAuxData { + services: HashSet::from([service.to_string()]), + operations: HashSet::from([(service.to_string(), span.to_string(), kind.to_string())]), + } +} + +fn trace_aux_cache_metrics() -> (IntCounter, IntCounter) { + ( + IntCounter::new("test_trace_aux_cache_hit", "Cache hits").unwrap(), + IntCounter::new("test_trace_aux_cache_miss", "Cache misses").unwrap(), + ) +} + +#[tokio::test] +async fn test_trace_aux_waits_for_successful_chunk_completions() { + let (success_tx, success_rx) = tokio::sync::oneshot::channel(); + let (failure_tx, failure_rx) = tokio::sync::oneshot::channel(); + let completion = |receiver: tokio::sync::oneshot::Receiver| { + WriteCompletion::new(async move { + if receiver.await.unwrap_or(false) { + Ok(()) + } else { + Err(common_error::ext::BoxedError::new( + servers::error::InternalSnafu { + err_msg: "main write failed", + } + .build(), + )) + } + }) + }; + let mut confirmed = Box::pin(super::confirmed_trace_aux_data(vec![ + ( + trace_aux_data("good", "span", "server"), + vec![completion(success_rx)], + ), + ( + trace_aux_data("bad", "span", "server"), + vec![completion(failure_rx)], + ), + ])); + assert!(futures::poll!(confirmed.as_mut()).is_pending()); + success_tx.send(true).unwrap(); + assert!(futures::poll!(confirmed.as_mut()).is_pending()); + // A lost completion sender must not count as successful storage. + drop(failure_tx); + let aux = confirmed.await; + assert_eq!(aux.services, HashSet::from(["good".to_string()])); + assert_eq!( + aux.operations, + trace_aux_data("good", "span", "server").operations + ); +} + +#[test] +fn test_trace_aux_cache_confirmed_writes() { + let cache = TraceAuxCache::new(1024); + let (hits, misses) = trace_aux_cache_metrics(); + let metrics = Some((&hits, &misses)); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + let pending = cache.filter(table.clone(), &mut data, metrics); + assert_eq!(pending.len(), 2); + assert_eq!((hits.get(), misses.get()), (0, 2)); + + // Preparing a write must not suppress a retry after failure or cancellation. + let mut retry = trace_aux_data("svc", "span", "server"); + assert_eq!(cache.filter(table.clone(), &mut retry, None).len(), 2); + assert_eq!((hits.get(), misses.get()), (0, 2)); + cache.record(pending); + + let mut repeated = trace_aux_data("svc", "span", "server"); + assert!( + cache + .clone() + .filter(table.clone(), &mut repeated, metrics) + .is_empty() + ); + assert!(repeated.is_empty()); + assert_eq!((hits.get(), misses.get()), (2, 2)); + cache.filter(table.clone(), &mut retry, None); + assert_eq!((hits.get(), misses.get()), (2, 2)); + + for (span, kind) in [("new_span", "server"), ("span", "client")] { + let mut data = trace_aux_data("svc", span, kind); + assert_eq!(cache.filter(table.clone(), &mut data, metrics).len(), 1); + assert!(data.services.is_empty()); + assert_eq!(data.operations.len(), 1); + } + assert_eq!((hits.get(), misses.get()), (4, 4)); + cache.filter(table.clone(), &mut TraceAuxData::default(), metrics); + assert_eq!((hits.get(), misses.get()), (4, 4)); + + let mut other_service = trace_aux_data("other", "span", "server"); + assert_eq!(cache.filter(table, &mut other_service, None).len(), 2); + + for other_table in [ + TableName::new("other", "public", "traces"), + TableName::new("greptime", "other", "traces"), + TableName::new("greptime", "public", "other"), + ] { + let mut data = trace_aux_data("svc", "span", "server"); + assert_eq!(cache.filter(other_table, &mut data, None).len(), 2); + } + + // An operation hit must not suppress a missing service entry. + let cache = TraceAuxCache::new(1024); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + data.services.clear(); + cache.record(cache.filter(table.clone(), &mut data, None)); + + let mut data = trace_aux_data("svc", "span", "server"); + assert_eq!(cache.filter(table, &mut data, None).len(), 1); + assert_eq!(data.services.len(), 1); + assert!(data.operations.is_empty()); +} + +#[test] +fn test_trace_aux_cache_zero_size_disables_caching() { + let cache = TraceAuxCache::new(0); + let (hits, misses) = trace_aux_cache_metrics(); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + cache.record(cache.filter(table.clone(), &mut data, Some((&hits, &misses)))); + + let mut repeated = trace_aux_data("svc", "span", "server"); + assert_eq!( + cache + .filter(table, &mut repeated, Some((&hits, &misses))) + .len(), + 2 + ); + assert_eq!(repeated.services.len(), 1); + assert_eq!(repeated.operations.len(), 1); + assert_eq!(hits.get(), 0); + assert_eq!(misses.get(), 4); +} + +#[test] +fn test_trace_aux_cache_accounts_for_key_memory() { + let cache = TraceAuxCache::new(1024); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + cache.record(cache.filter(table.clone(), &mut data, None)); + cache.entries.run_pending_tasks(); + assert_eq!(cache.entries.entry_count(), 2); + assert!(cache.entries.weighted_size() > 2); + assert!(cache.entries.weighted_size() <= 1024); + + // Spare string capacity counts even when the string content is empty. + let large = || String::with_capacity(2048); + let service = || TraceAuxEntry::Service("svc".to_string()); + for (table, entry) in [ + (TableName::new(large(), "public", "traces"), service()), + (TableName::new("greptime", large(), "traces"), service()), + (TableName::new("greptime", "public", large()), service()), + (table.clone(), TraceAuxEntry::Service(large())), + ( + table.clone(), + TraceAuxEntry::Operation(large(), String::new(), String::new()), + ), + ( + table.clone(), + TraceAuxEntry::Operation(String::new(), large(), String::new()), + ), + ( + table, + TraceAuxEntry::Operation(String::new(), String::new(), large()), + ), + ] { + let cache = TraceAuxCache::new(1024); + cache.record(vec![TraceAuxCacheKey { + table: Arc::new(table), + entry, + }]); + cache.entries.run_pending_tasks(); + assert_eq!(cache.entries.entry_count(), 0); + } + + // Empty strings still retain the key and table structs. + let cache = TraceAuxCache::new(1); + cache.record(vec![TraceAuxCacheKey { + table: Arc::new(TableName::new("", "", "")), + entry: TraceAuxEntry::Service(String::new()), + }]); + cache.entries.run_pending_tasks(); + assert_eq!(cache.entries.entry_count(), 0); +} + +#[tokio::test] +async fn test_trace_aux_admission_rechecks_cache() { + let cache = TraceAuxCache::new(1024); + let (hits, misses) = trace_aux_cache_metrics(); + let limiter = RequestLimiter::try_new(1).unwrap(); + let table = TableName::new("greptime", "public", "traces"); + let data = || trace_aux_data("svc", "span", "server"); + let permit = cache + .acquire_for_misses(&limiter, table.clone(), data(), (&hits, &misses)) + .await + .unwrap() + .unwrap(); + let cloned_cache = cache.clone(); + let cloned_limiter = limiter.clone(); + let mut waiting = Box::pin(cloned_cache.acquire_for_misses( + &cloned_limiter, + table.clone(), + data(), + (&hits, &misses), + )); + assert!(futures::poll!(waiting.as_mut()).is_pending()); + assert_eq!((hits.get(), misses.get()), (0, 4)); + + // Successful storage populates the shared cache before releasing admission. + cache.record(cache.filter(table.clone(), &mut data(), None)); + assert!( + cloned_cache + .acquire_for_misses(&cloned_limiter, table.clone(), data(), (&hits, &misses)) + .await + .unwrap() + .is_none() + ); + drop(permit); + assert!(waiting.await.unwrap().is_none()); + assert_eq!((hits.get(), misses.get()), (2, 4)); + + // Rechecking a now-warm request must release its briefly acquired slot. + let other = cache.acquire_for_misses( + &limiter, + table, + trace_aux_data("svc", "other", "server"), + (&hits, &misses), + ); + assert!(futures::poll!(Box::pin(other).as_mut()).is_ready()); + assert_eq!((hits.get(), misses.get()), (3, 5)); +} + #[test] fn test_chunk_owned() { let chunks = chunk_owned(vec![1, 2, 3], 2); diff --git a/src/frontend/src/metrics.rs b/src/frontend/src/metrics.rs index f649db0b46f..74574e56930 100644 --- a/src/frontend/src/metrics.rs +++ b/src/frontend/src/metrics.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; use lazy_static::lazy_static; use prometheus::*; @@ -68,6 +69,12 @@ lazy_static! { ) .unwrap(); + /// Initial trace auxiliary-cache lookups, excluding admission and write rechecks. + pub static ref OTLP_TRACE_AUX_CACHE_HIT: IntCounter = + CACHE_HIT.with_label_values(&["otlp_trace_aux"]); + pub static ref OTLP_TRACE_AUX_CACHE_MISS: IntCounter = + CACHE_MISS.with_label_values(&["otlp_trace_aux"]); + /// The number of OpenTelemetry logs send by frontend node. pub static ref OTLP_LOGS_ROWS: IntCounter = register_int_counter!( "greptime_frontend_otlp_logs_rows", diff --git a/src/frontend/src/service_config/otlp.rs b/src/frontend/src/service_config/otlp.rs index e67f4b5bcba..7a6d3f75b91 100644 --- a/src/frontend/src/service_config/otlp.rs +++ b/src/frontend/src/service_config/otlp.rs @@ -12,9 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_base::readable_size::ReadableSize; +use common_stat::get_total_memory_readable; use serde::{Deserialize, Serialize}; const DEFAULT_TRACE_INGEST_CHUNK_SIZE: usize = 512; +const MIN_TRACE_AUX_CACHE_SIZE: ReadableSize = ReadableSize::mb(32); #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] @@ -22,6 +25,12 @@ pub struct OtlpOptions { pub enable: bool, /// Maximum spans per trace ingest chunk. Set to 0 to disable splitting. pub trace_ingest_chunk_size: usize, + /// Estimated memory budget for cached trace service/operation keys per frontend, + /// shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching. + /// Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB. + /// Uses 32MiB if memory detection is unavailable. + /// Excludes cache and allocator overhead; shared table names are charged per entry. + pub trace_aux_cache_size: ReadableSize, /// Whether to synthesize the `greptime_otel_resource_info` descriptor /// table from the resource attributes of OTLP metrics, so metrics-only /// services reach the semantic graph. Off by default: it creates and @@ -34,24 +43,51 @@ impl Default for OtlpOptions { Self { enable: true, trace_ingest_chunk_size: DEFAULT_TRACE_INGEST_CHUNK_SIZE, + trace_aux_cache_size: default_trace_aux_cache_size(get_total_memory_readable()), experimental_enable_resource_info: false, } } } +fn default_trace_aux_cache_size(memory: Option) -> ReadableSize { + memory + .map(|size| size / 128) + .unwrap_or_default() + .max(MIN_TRACE_AUX_CACHE_SIZE) +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn test_default_trace_aux_cache_size() { + for (memory, expected) in [ + (None, ReadableSize::mb(32)), + (Some(ReadableSize(0)), ReadableSize::mb(32)), + (Some(ReadableSize::gb(1)), ReadableSize::mb(32)), + (Some(ReadableSize::gb(4)), ReadableSize::mb(32)), + (Some(ReadableSize::gb(8)), ReadableSize::mb(64)), + (Some(ReadableSize::gb(64)), ReadableSize::mb(512)), + ] { + assert_eq!(default_trace_aux_cache_size(memory), expected); + } + } + #[test] fn test_otlp_options() { let default = OtlpOptions::default(); assert!(default.enable); assert_eq!(default.trace_ingest_chunk_size, 512); + assert_eq!( + default.trace_aux_cache_size, + default_trace_aux_cache_size(get_total_memory_readable()) + ); assert!(!default.experimental_enable_resource_info); let options: OtlpOptions = toml::from_str("enable = false").unwrap(); assert!(!options.enable); + assert_eq!(options.trace_aux_cache_size, default.trace_aux_cache_size); assert_eq!( options.trace_ingest_chunk_size, DEFAULT_TRACE_INGEST_CHUNK_SIZE @@ -63,5 +99,13 @@ mod tests { let serialized = toml::to_string(&options).unwrap(); assert_eq!(toml::from_str::(&serialized).unwrap(), options); + + for (value, size) in [("0", ReadableSize(0)), ("\"1MiB\"", ReadableSize::mb(1))] { + let options: OtlpOptions = + toml::from_str(&format!("trace_aux_cache_size = {value}")).unwrap(); + assert_eq!(options.trace_aux_cache_size, size); + let serialized = toml::to_string(&options).unwrap(); + assert_eq!(toml::from_str::(&serialized).unwrap(), options); + } } } diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md index 62853a18f61..49cbd617567 100644 --- a/src/operator/AGENTS.md +++ b/src/operator/AGENTS.md @@ -46,3 +46,5 @@ the default `meter-macros/noop` build and active metering when changing this pat Finite requests split internally use `admit_write` / `admit_row_insert_batches` before dispatch. Their database-scoped `QueryContext` admission prevents a second row debit while `do_request` still records WCU for the actual writes. +Async table-batcher outputs carry storage completions in `OutputMeta`; keep +them when combining submissions so callers can distinguish admission from success. diff --git a/src/operator/src/batcher.rs b/src/operator/src/batcher.rs index 677524d53f3..bc9ef86b34e 100644 --- a/src/operator/src/batcher.rs +++ b/src/operator/src/batcher.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use arrow::record_batch::RecordBatch; use async_trait::async_trait; +use common_query::Output; use session::context::QueryContextRef; use table::metadata::TableInfoRef; use tokio::sync::OwnedSemaphorePermit; @@ -32,11 +33,12 @@ pub trait PendingRowsBatcher: Send + Sync { /// Submits rows according to the acknowledgement policy, retaining the slot until /// writing completes even when the response acknowledges queue admission only. /// Cancelling the response wait does not retract an already enqueued submission. + /// Early acknowledgements retain actual storage results in the output metadata. async fn submit( &self, table_info: TableInfoRef, batch: RecordBatch, ctx: QueryContextRef, permit: Arc, - ) -> Result; + ) -> Result; } diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 80866b45b66..c7f60c0883c 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -555,11 +555,14 @@ impl Inserter { // Observe every table completion even when another table fails. future::join_all(submissions).await }; - let affected_rows = results.into_iter().sum::>()?; - Ok(Output::new( - OutputData::AffectedRows(affected_rows), - OutputMeta::new_with_cost(write_cost as _), - )) + let mut affected_rows = 0; + let mut meta = OutputMeta::new_with_cost(write_cost as _); + for result in results { + let output = result?; + affected_rows += output.extract_rows_and_cost().0; + meta.write_completions.extend(output.meta.write_completions); + } + Ok(Output::new(OutputData::AffectedRows(affected_rows), meta)) } /// Handles row inserts request with metric engine. @@ -3302,7 +3305,7 @@ mod tests { _batch: arrow::record_batch::RecordBatch, _ctx: QueryContextRef, _permit: Arc, - ) -> Result { + ) -> Result { panic!("empty writes must not submit a batch") } } diff --git a/src/pipeline/Cargo.toml b/src/pipeline/Cargo.toml index 5a9cb3e9aa6..9ebe51ecf24 100644 --- a/src/pipeline/Cargo.toml +++ b/src/pipeline/Cargo.toml @@ -19,6 +19,7 @@ chrono.workspace = true chrono-tz.workspace = true common-catalog.workspace = true common-error.workspace = true +common-frontend.workspace = true common-macro.workspace = true common-meta.workspace = true common-query.workspace = true diff --git a/src/pipeline/src/manager/pipeline_cache.rs b/src/pipeline/src/manager/pipeline_cache.rs index e60963e7f84..f0e9ca52970 100644 --- a/src/pipeline/src/manager/pipeline_cache.rs +++ b/src/pipeline/src/manager/pipeline_cache.rs @@ -16,6 +16,7 @@ use std::future::Future; use std::sync::Arc; use std::time::Duration; +use common_frontend::metrics; use datatypes::timestamp::TimestampNanosecond; use moka::future::Cache; @@ -32,7 +33,7 @@ const PIPELINES_CACHE_SIZE: u64 = 10000; /// to encapsulate inner cache. Only public methods are exposed. /// /// `pipelines` and `original_pipelines` are keyed by the *requested* schema so -/// a lookup is a single key probe, as [`Cache::try_get_with`] requires; +/// a lookup is a single key probe through [`Cache::entry`]; /// resolving it to a stored schema is the loader's job. `failover_cache` has no /// loader and keeps the stored-schema key. pub(crate) struct PipelineCache { @@ -71,7 +72,7 @@ impl PipelineCache { } } - /// Concurrent misses on the same key share one `init` call. + /// Concurrent misses share one `init` call; successful waiters count as hits. pub(crate) async fn get_pipeline_with( &self, schema: &str, @@ -80,13 +81,17 @@ impl PipelineCache { init: impl Future>>, ) -> Result> { let key = generate_pipeline_cache_key(schema, name, version); - self.pipelines - .try_get_with(key, init) - .await + let entry = self.pipelines.entry(key).or_try_insert_with(init).await; + metrics::record_cache_lookup( + "pipeline", + entry.as_ref().is_ok_and(|entry| !entry.is_fresh()), + ); + entry + .map(|entry| entry.into_value()) .map_err(|error| CacheLoadSnafu { error }.build()) } - /// Concurrent misses on the same key share one `init` call. + /// Concurrent misses share one `init` call; successful waiters count as hits. pub(crate) async fn get_pipeline_str_with( &self, schema: &str, @@ -95,9 +100,17 @@ impl PipelineCache { init: impl Future>, ) -> Result { let key = generate_pipeline_cache_key(schema, name, version); - self.original_pipelines - .try_get_with(key, init) - .await + let entry = self + .original_pipelines + .entry(key) + .or_try_insert_with(init) + .await; + metrics::record_cache_lookup( + "pipeline_original", + entry.as_ref().is_ok_and(|entry| !entry.is_fresh()), + ); + entry + .map(|entry| entry.into_value()) .map_err(|error| CacheLoadSnafu { error }.build()) } @@ -114,6 +127,7 @@ impl PipelineCache { generate_pipeline_cache_key(schema, name, version), ] { if let Some(content) = self.failover_cache.get(&key).await { + metrics::record_cache_lookup("pipeline_failover", true); return Ok(Some(content)); } } @@ -126,7 +140,7 @@ impl PipelineCache { .filter(|(k, _)| k.ends_with(&suffix)) .collect::>(); - match found.len() { + let result = match found.len() { 0 => Ok(None), 1 => Ok(Some(found.remove(0).1)), _ => MultiPipelineWithDiffSchemaSnafu { @@ -139,7 +153,12 @@ impl PipelineCache { .join(","), } .fail(), - } + }; + metrics::record_cache_lookup( + "pipeline_failover", + result.as_ref().is_ok_and(|content| content.is_some()), + ); + result } pub(crate) async fn insert_failover_cache(&self, content: PipelineContent, with_latest: bool) { @@ -191,7 +210,8 @@ impl PipelineCache { mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; - use tokio::sync::Barrier; + use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; + use tokio::sync::Notify; use super::*; @@ -205,36 +225,149 @@ mod tests { } } - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + #[tokio::test] async fn test_concurrent_misses_run_one_loader() { const CONCURRENCY: usize = 8; - let cache = Arc::new(PipelineCache::new(Duration::from_secs(60))); - let loads = Arc::new(AtomicUsize::new(0)); - let barrier = Arc::new(Barrier::new(CONCURRENCY)); - - let handles = (0..CONCURRENCY) + let hits = CACHE_HIT.with_label_values(&["pipeline_original"]); + let misses = CACHE_MISS.with_label_values(&["pipeline_original"]); + let before = (hits.get(), misses.get()); + let cache = PipelineCache::new(Duration::from_secs(60)); + let loads = AtomicUsize::new(0); + let release = Notify::new(); + let mut requests = (0..CONCURRENCY) .map(|_| { - let (cache, loads, barrier) = (cache.clone(), loads.clone(), barrier.clone()); - tokio::spawn(async move { - barrier.wait().await; + Box::pin(cache.get_pipeline_str_with("db", "p", None, async { + loads.fetch_add(1, Ordering::SeqCst); + release.notified().await; + Ok(content_at(1)) + })) + }) + .collect::>(); + for request in &mut requests { + assert!(futures::poll!(request.as_mut()).is_pending()); + } + assert_eq!(loads.load(Ordering::SeqCst), 1); + release.notify_waiters(); + for request in requests { + assert_eq!(request.await.unwrap(), content_at(1)); + } + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!( + (hits.get(), misses.get()), + (before.0 + CONCURRENCY as u64 - 1, before.1 + 1) + ); + } + + #[tokio::test] + async fn test_pipeline_lookup_metrics() { + let cache = PipelineCache::new(Duration::from_secs(60)); + for compiled in [false, true] { + let cache_type = if compiled { + "pipeline" + } else { + "pipeline_original" + }; + let hits = CACHE_HIT.with_label_values(&[cache_type]); + let misses = CACHE_MISS.with_label_values(&[cache_type]); + let before = (hits.get(), misses.get()); + let load = async |fail| { + if compiled { + cache + .get_pipeline_with("db", "p", None, async { + if fail { + return Err(crate::error::PipelineNotFoundSnafu { + name: "p", + version: None, + } + .build()); + } + Ok(Arc::new( + crate::manager::table::PipelineTable::compile_pipeline( + "transform:\n - fields: [message]\n type: string", + )?, + )) + }) + .await + .map(|_| ()) + } else { cache .get_pipeline_str_with("db", "p", None, async { - loads.fetch_add(1, Ordering::SeqCst); - // Hold the loader open so every caller is waiting on it. - tokio::time::sleep(Duration::from_millis(100)).await; + if fail { + return Err(crate::error::PipelineNotFoundSnafu { + name: "p", + version: None, + } + .build()); + } Ok(content_at(1)) }) .await - .unwrap() - }) - }) - .collect::>(); - - for handle in handles { - assert_eq!(handle.await.unwrap(), content_at(1)); + .map(|_| ()) + } + }; + assert!(load(true).await.is_err()); + load(false).await.unwrap(); + // A warm lookup must not evaluate the failing initializer. + load(true).await.unwrap(); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 2)); } - assert_eq!(loads.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_failover_lookup_metrics() { + let hits = CACHE_HIT.with_label_values(&["pipeline_failover"]); + let misses = CACHE_MISS.with_label_values(&["pipeline_failover"]); + let before = (hits.get(), misses.get()); + let cache = PipelineCache::new(Duration::from_secs(60)); + assert!( + cache + .get_failover_cache("db", "p", None) + .await + .unwrap() + .is_none() + ); + let local = PipelineContent { + schema: "a".into(), + ..content_at(1) + }; + cache.insert_failover_cache(local.clone(), true).await; + // Direct schema lookup and cross-schema fallback each count one hit. + assert_eq!( + cache.get_failover_cache("a", "p", None).await.unwrap(), + Some(local.clone()) + ); + assert_eq!( + cache.get_failover_cache("db", "p", None).await.unwrap(), + Some(local) + ); + cache + .insert_failover_cache( + PipelineContent { + schema: "b".into(), + ..content_at(2) + }, + true, + ) + .await; + assert!(cache.get_failover_cache("db", "p", None).await.is_err()); + // Global pipelines win over local entries inserted before or after them. + for (schema, version) in [(EMPTY_SCHEMA_NAME, 3), ("c", 4)] { + cache + .insert_failover_cache( + PipelineContent { + schema: schema.into(), + ..content_at(version) + }, + true, + ) + .await; + assert_eq!( + cache.get_failover_cache("db", "p", None).await.unwrap(), + Some(content_at(3)) + ); + } + assert_eq!((hits.get(), misses.get()), (before.0 + 4, before.1 + 2)); } #[tokio::test] @@ -288,25 +421,4 @@ mod tests { let failover = cache.get_failover_cache("b", "p", None).await.unwrap(); assert_eq!(failover.map(|c| c.version), Some(v2.version)); } - - #[tokio::test] - async fn test_failover_serves_global_pipeline_to_unwarmed_schema() { - let cache = PipelineCache::new(Duration::from_secs(60)); - let content = content_at(1); - - cache.insert_failover_cache(content.clone(), true).await; - - let found = cache.get_failover_cache("b", "p", None).await.unwrap(); - assert_eq!(found, Some(content.clone())); - - // A same-named pipeline under another schema must not shadow the global one. - let schema_local = PipelineContent { - schema: "x".to_string(), - ..content_at(2) - }; - cache.insert_failover_cache(schema_local, true).await; - - let found = cache.get_failover_cache("b", "p", None).await.unwrap(); - assert_eq!(found, Some(content)); - } } diff --git a/src/pipeline/src/manager/pipeline_operator.rs b/src/pipeline/src/manager/pipeline_operator.rs index c9731b36f20..cace74171c2 100644 --- a/src/pipeline/src/manager/pipeline_operator.rs +++ b/src/pipeline/src/manager/pipeline_operator.rs @@ -19,6 +19,7 @@ use std::time::{Duration, Instant}; use api::v1::CreateTableExpr; use catalog::{CatalogManagerRef, RegisterSystemTableRequest}; use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, default_engine}; +use common_frontend::metrics; use common_meta::rpc::ddl::TriggerReason; use common_telemetry::info; use common_time::FOREVER; @@ -164,7 +165,9 @@ impl PipelineOperator { /// Get a pipeline table from the cache. pub fn get_pipeline_table_from_cache(&self, catalog: &str) -> Option { - self.tables.read().unwrap().get(catalog).cloned() + let table = self.tables.read().unwrap().get(catalog).cloned(); + metrics::record_cache_lookup("pipeline_table", table.is_some()); + table } } diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md index 006ac59771c..d20704cf616 100644 --- a/src/servers/AGENTS.md +++ b/src/servers/AGENTS.md @@ -41,3 +41,5 @@ cargo nextest run -p servers ``` Keep protocol translation here and permissions/database behavior in frontend. +Async table batching returns storage-completion futures with its early +acknowledgements. Keep completion reporting independent of the caller's lifetime. diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 8a2b6ab2149..0be0e47eab7 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -40,6 +40,7 @@ common-batcher.workspace = true common-catalog.workspace = true common-decimal.workspace = true common-error.workspace = true +common-frontend.workspace = true common-grpc.workspace = true common-macro.workspace = true common-mem-prof = { workspace = true, optional = true } diff --git a/src/servers/src/batcher/table.rs b/src/servers/src/batcher/table.rs index 67b3431a810..daae05251f8 100644 --- a/src/servers/src/batcher/table.rs +++ b/src/servers/src/batcher/table.rs @@ -30,6 +30,8 @@ use common_batcher::flush_limiter::FlushLimiter; use common_batcher::flush_policy::timing::TimingFlushPolicy; use common_batcher::request_limiter::RequestLimiter; use common_batcher::worker_registry::WorkerRegistry; +use common_error::ext::BoxedError; +use common_query::{Output, WriteCompletion}; use operator::batcher::PendingRowsBatcher; use operator::error::{BatchFlushSnafu, Result, UnexpectedSnafu}; use operator::insert::Inserter; @@ -171,10 +173,10 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { batch: RecordBatch, ctx: QueryContextRef, permit: Arc, - ) -> Result { + ) -> Result { let total_rows = batch.num_rows(); if total_rows == 0 { - return Ok(0); + return Ok(Output::new_with_affected_rows(0)); } if table_info.catalog_name != ctx.current_catalog() || table_info.schema_name != ctx.current_schema() @@ -213,19 +215,30 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { } .fail(); } - if !self.pending_rows_batch_sync { - return Ok(total_rows); + let completion = async move { + response_rx + .await + .map_err(|_| { + UnexpectedSnafu { + violated: "batch worker stopped before reporting its write result" + .to_string(), + } + .build() + })? + .context(BatchFlushSnafu) + }; + let mut output = Output::new_with_affected_rows(total_rows); + if self.pending_rows_batch_sync { + completion.await?; + } else { + output + .meta + .write_completions + .push(WriteCompletion::new(async move { + completion.await.map_err(BoxedError::new) + })); } - response_rx - .await - .map_err(|_| { - UnexpectedSnafu { - violated: "batch worker stopped before reporting its write result".to_string(), - } - .build() - })? - .context(BatchFlushSnafu)?; - Ok(total_rows) + Ok(output) } } @@ -386,17 +399,20 @@ mod tests { .unwrap() .unwrap(); let mut submitted = Some(submitted); + let mut completion = None; if sync { assert!(!submitted.as_ref().unwrap().is_finished()); } else { - assert_eq!( - timeout(Duration::from_secs(5), submitted.take().unwrap()) - .await - .unwrap() - .unwrap() - .unwrap(), - 1 - ); + let mut output = timeout(Duration::from_secs(5), submitted.take().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(output.extract_rows_and_cost().0, 1); + assert_eq!(output.meta.write_completions.len(), 1); + let mut pending = Box::pin(output.meta.write_completions.pop().unwrap().wait()); + assert!(futures::poll!(pending.as_mut()).is_pending()); + completion = Some(pending); } // Early acknowledgement must not release capacity before completion. let acquire = batcher.acquire(); @@ -420,9 +436,15 @@ mod tests { .unwrap(); assert_eq!(result.is_err(), fail); if !fail { - assert_eq!(result.unwrap(), 1); + let output = result.unwrap(); + assert_eq!(output.extract_rows_and_cost().0, 1); + assert!(output.meta.write_completions.is_empty()); } } + if let Some(completion) = completion { + let result = timeout(Duration::from_secs(5), completion).await.unwrap(); + assert_eq!(result.is_err(), fail); + } timeout(Duration::from_secs(5), acquire) .await .unwrap() @@ -533,9 +555,13 @@ mod tests { let (first_result, second_result) = timeout(Duration::from_secs(5), async { if shared_request { let permit = batcher.acquire().await.unwrap(); - tokio::join!( + let (first, second) = tokio::join!( batcher.submit(table.clone(), first, influx_ctx, permit.clone()), batcher.submit(table.clone(), second, opentsdb_ctx, permit), + ); + ( + first.map(|output| output.extract_rows_and_cost().0), + second.map(|output| output.extract_rows_and_cost().0), ) } else { let inserter = Inserter::new( diff --git a/src/servers/src/mysql/handler.rs b/src/servers/src/mysql/handler.rs index d0dd1f25279..bdfc7a72057 100644 --- a/src/servers/src/mysql/handler.rs +++ b/src/servers/src/mysql/handler.rs @@ -23,6 +23,7 @@ use async_trait::async_trait; use chrono::{NaiveDate, NaiveDateTime}; use common_catalog::parse_optional_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; +use common_frontend::metrics; use common_query::Output; use common_telemetry::{debug, error, tracing, warn}; use common_time::Timezone; @@ -182,8 +183,9 @@ impl MysqlInstanceShim { /// Retrieve the query and logical plan by a given statement key fn plan(&self, stmt_key: &str) -> Option { - let guard = self.prepared_stmts.read(); - guard.get(stmt_key).cloned() + let plan = self.prepared_stmts.read().get(stmt_key).cloned(); + metrics::record_cache_lookup("mysql_prepared_stmt", plan.is_some()); + plan } /// Save the prepared statement and return the parameters and result columns @@ -1053,6 +1055,23 @@ mod tests { ) } + #[test] + fn test_prepared_statement_cache_metrics() { + use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; + + let hits = CACHE_HIT.with_label_values(&["mysql_prepared_stmt"]); + let misses = CACHE_MISS.with_label_values(&["mysql_prepared_stmt"]); + let before = (hits.get(), misses.get()); + let mut shim = create_shim(); + assert!(shim.plan("stmt").is_none()); + shim.save_plan(SqlPlan::Shortcut("SELECT 1".into()), "stmt".into()) + .unwrap(); + assert!(shim.plan("stmt").is_some()); + shim.do_close("stmt".into()); + assert!(shim.plan("stmt").is_none()); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 2)); + } + #[test] fn test_batching_context() { for enabled in [false, true] { diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 2c90f4ee2ae..79df85475d9 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -123,6 +123,7 @@ sqlx = { workspace = true, features = [ "chrono", ] } store-api.workspace = true +temp-env.workspace = true time = "0.3" tokio-postgres = { workspace = true } tokio-util.workspace = true diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 29c8891a695..3e447bb0194 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -14,17 +14,28 @@ #[cfg(test)] mod test { - use std::sync::Arc; + use std::sync::{Arc, Mutex}; + use std::time::Duration; use TimeUnit as ArrowTimeUnit; + use api::region::RegionResponse; + use api::v1::region::{RegionRequest, region_request}; use client::{DEFAULT_CATALOG_NAME, OutputData}; - use common_recordbatch::RecordBatches; + use common_meta::cache::LayeredCacheRegistryBuilder; + use common_meta::node_manager::{ + Datanode, DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef, NodeManagerRef, + }; + use common_meta::peer::Peer; + use common_query::request::QueryRequest; + use common_recordbatch::{RecordBatches, SendableRecordBatchStream}; use datatypes::arrow::array::AsArray; use datatypes::arrow::datatypes::{ DataType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, }; use frontend::instance::Instance; + use frontend::instance::builder::FrontendBuilder; + use frontend::service_config::BatcherOptions; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::common::v1::any_value::Value as Val; @@ -39,14 +50,574 @@ mod test { use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; use pipeline::{GreptimePipelineParams, PipelineWay}; use serde_json::json; + use servers::batcher::BatchingProtocol; + use servers::otlp::trace::{TraceAuxData, span}; use servers::query_handler::OpenTelemetryProtocolHandler; use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; + use store_api::region_engine::RegionEngine; + use store_api::storage::RegionId; + use tokio::sync::oneshot; - use crate::standalone::GreptimeDbStandaloneBuilder; + use crate::standalone::{GreptimeDbStandalone, GreptimeDbStandaloneBuilder}; use crate::tests; + #[tokio::test(flavor = "multi_thread")] + async fn test_trace_aux_cache() { + check_trace_aux_cache(false).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_trace_aux_cache_with_batching_v0_v1() { + check_trace_aux_cache(true).await; + } + + #[test] + fn test_trace_aux_async_completion() { + temp_env::with_var("PENDING_ROWS_BATCH_SYNC", Some("false"), || { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + .block_on(check_trace_aux_async_completion()); + }); + } + + #[test] + fn test_trace_aux_admission() { + temp_env::with_var("PENDING_ROWS_BATCH_SYNC", Some("false"), || { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + tokio::time::timeout(Duration::from_secs(60), check_trace_aux_admission()) + .await + .expect("auxiliary admission must make progress"); + }); + }); + } + + struct TraceAuxWriteGate { + region_id: RegionId, + started: oneshot::Sender<()>, + release: oneshot::Receiver, + } + + struct GatedTraceNodeManager { + inner: NodeManagerRef, + gate: Arc>>, + } + + #[async_trait::async_trait] + impl DatanodeManager for GatedTraceNodeManager { + async fn datanode(&self, peer: &Peer) -> DatanodeRef { + Arc::new(GatedTraceDatanode { + inner: self.inner.datanode(peer).await, + gate: self.gate.clone(), + }) + } + } + + #[async_trait::async_trait] + impl FlownodeManager for GatedTraceNodeManager { + async fn flownode(&self, peer: &Peer) -> FlownodeRef { + self.inner.flownode(peer).await + } + } + + struct GatedTraceDatanode { + inner: DatanodeRef, + gate: Arc>>, + } + + #[async_trait::async_trait] + impl Datanode for GatedTraceDatanode { + async fn handle( + &self, + request: RegionRequest, + ) -> common_meta::error::Result { + let gate = { + let mut gate = self.gate.lock().unwrap(); + if gate.as_ref().is_some_and(|gate| { + matches!(&request.body, Some(region_request::Body::Inserts(inserts)) + if inserts.requests.iter().any(|insert| insert.region_id == gate.region_id.as_u64())) + }) { + gate.take() + } else { + None + } + }; + if let Some(gate) = gate { + gate.started.send(()).unwrap(); + if !gate.release.await.unwrap() { + return common_meta::error::UnexpectedSnafu { + err_msg: "injected auxiliary write failure", + } + .fail(); + } + } + self.inner.handle(request).await + } + + async fn handle_query( + &self, + request: QueryRequest, + ) -> common_meta::error::Result { + self.inner.handle_query(request).await + } + } + + async fn check_trace_aux_admission() { + let standalone = GreptimeDbStandaloneBuilder::new("trace_aux_admission") + .with_table_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_secs(3600), + max_batch_rows: 1, + max_inflight_requests: 1, + ..Default::default() + }) + .build() + .await; + let source = standalone.fe_instance(); + let gate = Arc::new(Mutex::new(None)); + let cache_registry = Arc::new( + cache::with_default_composite_cache_registry( + LayeredCacheRegistryBuilder::default().add_cache_registry( + cache::build_fundamental_cache_registry(standalone.kv_backend.clone()), + ), + ) + .unwrap() + .build(), + ); + let mut options = standalone.opts.frontend_options(); + options.otlp.trace_ingest_chunk_size = 1; + let instance = Arc::new( + FrontendBuilder::new( + options, + standalone.kv_backend.clone(), + cache_registry, + source.catalog_manager().clone(), + Arc::new(GatedTraceNodeManager { + inner: source.node_manager().clone(), + gate: gate.clone(), + }), + source.procedure_executor().clone(), + Arc::new(catalog::process_manager::ProcessManager::new( + "trace_aux_admission".into(), + None, + )), + ) + .try_build() + .await + .unwrap(), + ); + + for version in 0..=1 { + let table_name = format!("trace_aux_admission_v{version}"); + let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + context.set_extension( + common_catalog::consts::TRACE_TABLE_NAME_SESSION_KEY, + &table_name, + ); + context.set_extension(table::requests::TRACE_TABLE_PARTITIONS_HINT_KEY, "1"); + context.set_batching_enabled(true); + let ctx = Arc::new(context); + let write = async |request| { + instance + .traces( + instance.clone(), + request, + if version == 0 { + PipelineWay::OtlpTraceDirectV0 + } else { + PipelineWay::OtlpTraceDirectV1 + }, + GreptimePipelineParams::default(), + table_name.clone(), + ctx.clone(), + ) + .await + .unwrap() + }; + let cached = |request: ExportTraceServiceRequest| { + let mut data = TraceAuxData::default(); + for span in span::parse(request).iter().flat_map(|group| &group.spans) { + data.observe_span(span); + } + instance.trace_aux_data_cached_for_test(&table_name, data, &ctx) + }; + let warm = trace_aux_request(); + assert_eq!(1, write(warm.clone()).await.accepted_spans); + while !cached(warm.clone()) { + tokio::task::yield_now().await; + } + let table = instance + .catalog_manager() + .table( + DEFAULT_CATALOG_NAME, + "public", + &format!("{table_name}_operations"), + None, + ) + .await + .unwrap() + .unwrap(); + let route = instance + .partition_manager() + .find_physical_table_route(table.table_info().table_id()) + .await + .unwrap(); + let region_id = route.region_routes[0].region.id; + + for success in [true, false] { + let mut cold = warm.clone(); + let spans = &mut cold.resource_spans[0].scope_spans[0].spans; + spans[0].name = format!("cold_{success}"); + spans.push(spans[0].clone()); + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + *gate.lock().unwrap() = Some(TraceAuxWriteGate { + region_id, + started: started_tx, + release: release_rx, + }); + + // Two chunks must progress with just one main and one auxiliary slot. + assert_eq!(2, write(cold.clone()).await.accepted_spans); + started_rx.await.unwrap(); + assert!( + !cached(cold.clone()), + "pending writes must not warm the cache" + ); + let before = trace_aux_sequences(&standalone, &table_name).await; + let mut waiting = Box::pin(write(cold.clone())); + assert!(futures::poll!(waiting.as_mut()).is_pending()); + + // Main admission is free while auxiliary storage is blocked. Warm + // requests must bypass auxiliary admission and keep making progress. + assert_eq!(1, write(warm.clone()).await.accepted_spans); + assert!(futures::poll!(waiting.as_mut()).is_pending()); + release_tx.send(success).unwrap(); + assert_eq!(2, waiting.await.accepted_spans); + if !success { + // The failed write cannot satisfy the waiting request's miss. + while !cached(cold.clone()) { + tokio::task::yield_now().await; + } + } + let after = trace_aux_sequences(&standalone, &table_name).await; + assert_eq!(before[1], after[1], "cached service must not be rewritten"); + assert_eq!( + before[2] + 1, + after[2], + "exactly one successful operation write" + ); + } + } + } + + fn trace_aux_request() -> ExportTraceServiceRequest { + serde_json::from_value(json!({ + "resourceSpans": [{ + "resource": {"attributes": [ + {"key": "service.name", "value": {"stringValue": "svc"}} + ]}, + "scopeSpans": [{"spans": [{ + "traceId": "c05d7a4ec8e1f231f02ed6e8da8655b4", + "spanId": "9630f2916e2f7909", + "name": "operation", + "kind": 2, + "startTimeUnixNano": "1736480942444376000", + "endTimeUnixNano": "1736480942444499000" + }]}] + }] + })) + .unwrap() + } + + async fn trace_aux_sequences(standalone: &GreptimeDbStandalone, table_name: &str) -> Vec { + let instance = standalone.fe_instance(); + let mut sequences = Vec::new(); + for suffix in ["", "_services", "_operations"] { + let table = instance + .catalog_manager() + .table( + DEFAULT_CATALOG_NAME, + "public", + &format!("{table_name}{suffix}"), + None, + ) + .await + .unwrap() + .unwrap(); + let route = instance + .partition_manager() + .find_physical_table_route(table.table_info().table_id()) + .await + .unwrap(); + assert_eq!(route.region_routes.len(), 1); + sequences.push( + standalone + .mito_engine + .get_committed_sequence(route.region_routes[0].region.id) + .await + .unwrap(), + ); + } + sequences + } + + async fn check_trace_aux_async_completion() { + let standalone = GreptimeDbStandaloneBuilder::new("trace_aux_async_completion") + .with_table_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_secs(3600), + max_batch_rows: 2, + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + // V2 bulk ingestion has a separate Variant/Struct conversion failure. + for version in 0..=1 { + let table_name = format!("trace_aux_async_v{version}"); + let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + context.set_extension( + common_catalog::consts::TRACE_TABLE_NAME_SESSION_KEY, + &table_name, + ); + context.set_extension(table::requests::TRACE_TABLE_PARTITIONS_HINT_KEY, "1"); + context.set_batching_enabled(true); + let ctx = Arc::new(context); + let request = trace_aux_request(); + let write = async |request, ctx| { + let outcome = tokio::time::timeout( + Duration::from_secs(30), + instance.traces( + instance.clone(), + request, + if version == 0 { + PipelineWay::OtlpTraceDirectV0 + } else { + PipelineWay::OtlpTraceDirectV1 + }, + GreptimePipelineParams::default(), + table_name.clone(), + ctx, + ), + ) + .await + .expect("async trace ingestion must return before the batch flush") + .unwrap(); + assert_eq!((outcome.accepted_spans, outcome.rejected_spans), (1, 0)); + assert!(outcome.error_message.is_none()); + }; + write(request.clone(), ctx.clone()).await; + assert_trace_v2_query( + instance, + &format!("SELECT COUNT(*) = 0 FROM {table_name}"), + ctx.clone(), + ) + .await + .unwrap(); + for suffix in ["_services", "_operations"] { + assert!( + instance + .catalog_manager() + .table( + DEFAULT_CATALOG_NAME, + "public", + &format!("{table_name}{suffix}"), + None + ) + .await + .unwrap() + .is_none() + ); + } + + // The second submission deterministically triggers the main flush. + // Omitting the service avoids a second concurrent auxiliary write. + let mut flush_request = request.clone(); + flush_request.resource_spans[0] + .resource + .as_mut() + .unwrap() + .attributes + .clear(); + write(flush_request, ctx.clone()).await; + tokio::time::timeout(Duration::from_secs(30), async { + loop { + let mut aux_data = TraceAuxData::default(); + for span in span::parse(request.clone()) + .into_iter() + .flat_map(|group| group.spans) + { + aux_data.observe_span(&span); + } + if instance.trace_aux_data_cached_for_test(&table_name, aux_data, &ctx) { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("deferred auxiliary writes must populate the cache"); + for suffix in ["_services", "_operations"] { + assert_trace_v2_query( + instance, + &format!("SELECT COUNT(*) = 1 FROM {table_name}{suffix}"), + ctx.clone(), + ) + .await + .unwrap(); + } + assert_trace_v2_query( + instance, + &format!("SELECT COUNT(*) > 0 FROM {table_name}"), + ctx.clone(), + ) + .await + .unwrap(); + + let before = trace_aux_sequences(&standalone, &table_name).await; + // A direct warm request finishes before we compare storage sequences. + let mut warm_ctx = ctx.fork(); + warm_ctx.set_batching_enabled(false); + write(request, Arc::new(warm_ctx)).await; + let after = trace_aux_sequences(&standalone, &table_name).await; + assert!(after[0] > before[0]); + assert_eq!(after[1..], before[1..], "v{version}"); + } + } + + async fn check_trace_aux_cache(batching: bool) { + let standalone = GreptimeDbStandaloneBuilder::new("trace_aux_cache") + .with_table_batcher(if batching { + BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + // Two main spans trigger a flush, while each auxiliary table + // has only one row and would wait an hour if batched. + pending_rows_flush_interval: Duration::from_secs(3600), + max_batch_rows: 2, + ..Default::default() + } + } else { + BatcherOptions::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + let cloned_instance = Arc::new(instance.as_ref().clone()); + + // V2 bulk writes currently fail on Variant/Struct schema conversion + // before reaching auxiliary ingestion; cover V2 with direct main writes. + let max_version = if batching { 1 } else { 2 }; + for version in 0..=max_version { + let table_name = format!("trace_aux_v{version}"); + let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + context.set_extension( + common_catalog::consts::TRACE_TABLE_NAME_SESSION_KEY, + &table_name, + ); + context.set_extension(table::requests::TRACE_TABLE_PARTITIONS_HINT_KEY, "1"); + context.set_batching_enabled(batching); + let ctx = Arc::new(context); + let mut request = trace_aux_request(); + if batching { + let spans = &mut request.resource_spans[0].scope_spans[0].spans; + spans.push(spans[0].clone()); + } + let accepted_spans = if batching { 2 } else { 1 }; + let write = async |frontend: &Arc, request| { + tokio::time::timeout( + Duration::from_secs(30), + frontend.traces( + frontend.clone(), + request, + match version { + 0 => PipelineWay::OtlpTraceDirectV0, + 1 => PipelineWay::OtlpTraceDirectV1, + _ => PipelineWay::OtlpTraceDirectV2, + }, + GreptimePipelineParams::default(), + table_name.clone(), + ctx.clone(), + ), + ) + .await + .expect("auxiliary writes must not wait for the batch flush interval") + .unwrap() + }; + + // An incompatible auxiliary column fails after the main span is + // accepted. No candidate key may be cached from the failed batch. + instance + .do_query( + &format!( + "CREATE TABLE {table_name}_operations (\ + \"timestamp\" TIMESTAMP(9) TIME INDEX, service_name STRING, \ + span_name STRING, span_kind BIGINT)" + ), + ctx.clone(), + ) + .await + .remove(0) + .unwrap(); + let failed_aux = write(instance, request.clone()).await; + assert_eq!( + (failed_aux.accepted_spans, failed_aux.rejected_spans), + (accepted_spans, 0) + ); + assert!( + failed_aux + .error_message + .unwrap() + .contains("Auxiliary trace tables") + ); + + instance + .do_query(&format!("DROP TABLE {table_name}_operations"), ctx.clone()) + .await + .remove(0) + .unwrap(); + let retry = write(instance, request.clone()).await; + assert_eq!( + (retry.accepted_spans, retry.rejected_spans), + (accepted_spans, 0) + ); + assert!(retry.error_message.is_none(), "{retry:?}"); + let before = trace_aux_sequences(&standalone, &table_name).await; + assert!(before.iter().all(|sequence| *sequence > 0)); + + // Storage sequences, unlike row counts in an upsert table, prove + // the writes were skipped. Cloned Instances must share the cache. + let warm = write(&cloned_instance, request.clone()).await; + assert_eq!( + (warm.accepted_spans, warm.rejected_spans), + (accepted_spans, 0) + ); + assert!(warm.error_message.is_none(), "{warm:?}"); + let after = trace_aux_sequences(&standalone, &table_name).await; + assert!(after[0] > before[0]); + assert_eq!(after[1..], before[1..], "v{version}"); + + request.resource_spans[0].scope_spans[0].spans[0].name = "new_operation".into(); + let new_operation = write(instance, request).await; + assert_eq!( + (new_operation.accepted_spans, new_operation.rejected_spans), + (accepted_spans, 0) + ); + assert!(new_operation.error_message.is_none(), "{new_operation:?}"); + let new_sequences = trace_aux_sequences(&standalone, &table_name).await; + assert!(new_sequences[0] > after[0]); + assert_eq!(new_sequences[1], after[1]); + assert!(new_sequences[2] > after[2]); + } + } + #[tokio::test(flavor = "multi_thread")] async fn test_trace_v2_on_standalone() -> Result<(), Box> { let standalone = GreptimeDbStandaloneBuilder::new("trace_v2_standalone") diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 54a4e2d67b8..b3293bf6421 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -3523,6 +3523,7 @@ fn drop_lines_with_inconsistent_results(input: String) -> String { "disable_ec2_metadata =", "cache_path =", "cache_capacity =", + "trace_aux_cache_size =", "memory_pool_size =", "scan_memory_limit =", "sas_token =",