diff --git a/Cargo.lock b/Cargo.lock index dba349c74e5..2badf7c6605 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5684,6 +5684,7 @@ dependencies = [ "store-api", "strfmt", "table", + "temp-env", "tokio", "tokio-util", "toml 0.8.23", @@ -14361,6 +14362,7 @@ dependencies = [ "snafu 0.8.6", "store-api", "table", + "temp-env", "tokio", "toml 0.8.23", "url", diff --git a/Cargo.toml b/Cargo.toml index 886c1bb9104..96ce1229621 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -253,6 +253,7 @@ sqlx = { version = "0.8", default-features = false, features = [ stringprep = "0.1" strum = { version = "0.27", features = ["derive"] } sysinfo = "0.33" +temp-env = "0.3" tempfile = "3" tokio = { version = "1.47", features = ["full"] } tokio-postgres = "0.7.18" diff --git a/config/config.md b/config/config.md index dd5b980f1ba..4dcffe40654 100644 --- a/config/config.md +++ b/config/config.md @@ -33,7 +33,7 @@ | `runtime.experimental_workload_scheduler.sample_every_polls` | Integer | `16` | Number of polls between scheduler fairness samples. Must be greater than zero. | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | -| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the
`prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. | +| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. | | `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. | | `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions | | `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. | @@ -74,6 +74,13 @@ | `influxdb` | -- | -- | InfluxDB protocol options. | | `influxdb.enable` | Bool | `true` | Whether to enable InfluxDB protocol in HTTP API. | | `influxdb.default_merge_mode` | String | `last_non_null` | Default merge mode for tables automatically created by InfluxDB protocol.
Available values: "last_non_null", "last_row". | +| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.
Legacy Prometheus batching settings under prom_store remain supported.
HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.
Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.
Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.
Effective shared Prom settings take precedence; existing prom_store settings remain compatible. | +| `pending_rows_batcher.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | +| `pending_rows_batcher.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | +| `pending_rows_batcher.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by the frontend batcher. | +| `pending_rows_batcher.worker_channel_capacity` | Integer | `65526` | Maximum queued submissions per table worker. | +| `pending_rows_batcher.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | +| `pending_rows_batcher.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued table Flow notifications. | | `jaeger` | -- | -- | Jaeger protocol options. | | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | @@ -279,7 +286,7 @@ | `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | -| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the
`prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. | +| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is
raised to at least the largest active flush interval plus 1 second. The intervals come from
`prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. | | `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. | | `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions | | `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. | @@ -332,6 +339,13 @@ | `influxdb` | -- | -- | InfluxDB protocol options. | | `influxdb.enable` | Bool | `true` | Whether to enable InfluxDB protocol in HTTP API. | | `influxdb.default_merge_mode` | String | `last_non_null` | Default merge mode for tables automatically created by InfluxDB protocol.
Available values: "last_non_null", "last_row". | +| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.
Legacy Prometheus batching settings under prom_store remain supported.
HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.
Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.
Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.
Effective shared Prom settings take precedence; existing prom_store settings remain compatible. | +| `pending_rows_batcher.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | +| `pending_rows_batcher.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | +| `pending_rows_batcher.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by the frontend batcher. | +| `pending_rows_batcher.worker_channel_capacity` | Integer | `65526` | Maximum queued submissions per table worker. | +| `pending_rows_batcher.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | +| `pending_rows_batcher.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued table Flow notifications. | | `jaeger` | -- | -- | Jaeger protocol options. | | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 7c6862916e4..95550931d46 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -52,8 +52,9 @@ default_column_prefix = "greptime" ## The address to bind the HTTP server. addr = "127.0.0.1:4000" ## HTTP request timeout. Set to 0 to disable timeout. -## When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the -## `prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. +## When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is +## raised to at least the largest active flush interval plus 1 second. The intervals come from +## `prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. timeout = "0s" ## HTTP request body limit. ## The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`. @@ -226,6 +227,37 @@ enable = true ## Available values: "last_non_null", "last_row". default_merge_mode = "last_non_null" +## Shared experimental ordinary-table batching for opted-in ingestion protocols. +## Legacy Prometheus batching settings under prom_store remain supported. +## HTTP write protocols sharing this batcher. Omitted or empty disables all entrances. +## Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom. +## Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher. +## Effective shared Prom settings take precedence; existing prom_store settings remain compatible. +[pending_rows_batcher] +# protocols = [ +# "influxdb", +# "opentsdb", +# "otlp", +# "logs", +# "loki", +# "splunk", +# "elasticsearch", +# "http_sql", +# "prom", +# ] +## Flush interval measured from the first pending submission. Zero disables batching. +pending_rows_flush_interval = "0s" +## Flush after a complete submission reaches this row threshold. +max_batch_rows = 100000 +## Maximum concurrent flushes shared by the frontend batcher. +max_concurrent_flushes = 256 +## Maximum queued submissions per table worker. +worker_channel_capacity = 65526 +## Maximum admitted original requests awaiting completion. +max_inflight_requests = 3000 +## Maximum number of queued table Flow notifications. +flow_notification_queue_capacity = 1024 + ## Jaeger protocol options. [jaeger] ## Whether to enable Jaeger protocol in HTTP API. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index e0a5f4bcd11..54a1555232e 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -77,8 +77,9 @@ max_concurrent_queries = 0 ## The address to bind the HTTP server. addr = "127.0.0.1:4000" ## HTTP request timeout. Set to 0 to disable timeout. -## When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the -## `prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. +## When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is +## raised to at least the largest active flush interval plus 1 second. The intervals come from +## `prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. timeout = "0s" ## HTTP request body limit. ## The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`. @@ -205,6 +206,37 @@ enable = true ## Available values: "last_non_null", "last_row". default_merge_mode = "last_non_null" +## Shared experimental ordinary-table batching for opted-in ingestion protocols. +## Legacy Prometheus batching settings under prom_store remain supported. +## HTTP write protocols sharing this batcher. Omitted or empty disables all entrances. +## Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom. +## Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher. +## Effective shared Prom settings take precedence; existing prom_store settings remain compatible. +[pending_rows_batcher] +# protocols = [ +# "influxdb", +# "opentsdb", +# "otlp", +# "logs", +# "loki", +# "splunk", +# "elasticsearch", +# "http_sql", +# "prom", +# ] +## Flush interval measured from the first pending submission. Zero disables batching. +pending_rows_flush_interval = "0s" +## Flush after a complete submission reaches this row threshold. +max_batch_rows = 100000 +## Maximum concurrent flushes shared by the frontend batcher. +max_concurrent_flushes = 256 +## Maximum queued submissions per table worker. +worker_channel_capacity = 65526 +## Maximum admitted original requests awaiting completion. +max_inflight_requests = 3000 +## Maximum number of queued table Flow notifications. +flow_notification_queue_capacity = 1024 + ## Jaeger protocol options. [jaeger] ## Whether to enable Jaeger protocol in HTTP API. diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index b5c0871fc11..9740d160164 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -44,6 +44,12 @@ remote datanodes via `operator`/`client`. auto-create, partition routing) → local `RegionServer` (standalone) or RPC to datanodes (distributed). +- **Table batching** (`instance/builder.rs`): protocol entry points opt in through + `QueryContext`. The primary inserter prepares eligible ordinary-table writes + for `servers::batcher::table::TablePendingRowsBatcher`. A separate execution-only + inserter, with no batcher attached, sends the prepared bulk writes to datanodes, + avoiding recursive batching. The batcher handles successful-write Flow notifications. + ## Public surface - `Instance` (`instance.rs`) — the business-logic container. diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index 677dbbbe55f..b59e1e83a4a 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -98,5 +98,6 @@ hyper-util = { workspace = true, features = ["tokio"] } reqwest.workspace = true serde_json.workspace = true strfmt = "0.2" +temp-env.workspace = true tower.workspace = true uuid.workspace = true diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index bb15abd7bbf..793ab94205e 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -35,8 +35,8 @@ use crate::error::Result; use crate::heartbeat::HeartbeatTask; use crate::instance::Instance; use crate::service_config::{ - InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions, PostgresOptions, - PromStoreOptions, + InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions, + PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions, }; #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] @@ -65,6 +65,8 @@ pub struct FrontendOptions { pub postgres: PostgresOptions, pub opentsdb: OpentsdbOptions, pub influxdb: InfluxdbOptions, + /// Shared experimental ordinary-table batching; independent of Prom batching. + pub pending_rows_batcher: PendingRowsBatcherOptions, pub prom_store: PromStoreOptions, pub jaeger: JaegerOptions, pub otlp: OtlpOptions, @@ -100,6 +102,7 @@ impl Default for FrontendOptions { postgres: PostgresOptions::default(), opentsdb: OpentsdbOptions::default(), influxdb: InfluxdbOptions::default(), + pending_rows_batcher: PendingRowsBatcherOptions::default(), jaeger: JaegerOptions::default(), prom_store: PromStoreOptions::default(), otlp: OtlpOptions::default(), @@ -124,6 +127,7 @@ impl Configurable for FrontendOptions { "heartbeat_env_vars", "meta_client.metasrv_addrs", "event_recorder.event_types", + "pending_rows_batcher.protocols", ]) } } @@ -208,7 +212,7 @@ mod tests { use tonic::codegen::tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status, Streaming}; - use super::*; + use crate::frontend::*; use crate::heartbeat::{ FrontendHeartbeatExtension, FrontendHeartbeatExtensionResult, FrontendHeartbeatExtensions, }; @@ -218,6 +222,54 @@ mod tests { type GrpcStream = Pin> + Send + Sync + 'static>>; + #[test] + fn test_batcher_protocols_from_env() { + temp_env::with_vars( + [( + "FRONTEND_BATCHER_TEST__PENDING_ROWS_BATCHER__PROTOCOLS", + Some("influxdb,http_sql"), + )], + || { + let options = + FrontendOptions::load_layered_options(None, "FRONTEND_BATCHER_TEST").unwrap(); + assert_eq!( + options.pending_rows_batcher.protocols, + vec![ + servers::http::BatchingProtocol::Influxdb, + servers::http::BatchingProtocol::HttpSql + ] + ); + }, + ); + } + + #[test] + fn test_protocol_pending_rows_batcher_config() { + let defaults: FrontendOptions = toml::from_str("").unwrap(); + assert!( + !defaults + .pending_rows_batcher + .pending_rows_batching_enabled() + ); + let options: FrontendOptions = toml::from_str( + r#" +[pending_rows_batcher] +protocols = ["influxdb", "http_sql"] +pending_rows_flush_interval = "5ms" +max_batch_rows = 25 +"#, + ) + .unwrap(); + assert_eq!(options.pending_rows_batcher.max_batch_rows, 25); + assert_eq!(options.pending_rows_batcher.protocols.len(), 2); + assert!(options.pending_rows_batcher.pending_rows_batching_enabled()); + let serialized = toml::to_string(&options).unwrap(); + let parsed: FrontendOptions = toml::from_str(&serialized).unwrap(); + assert_eq!(options.influxdb, parsed.influxdb); + assert_eq!(options.opentsdb, parsed.opentsdb); + assert_eq!(options.pending_rows_batcher, parsed.pending_rows_batcher); + } + #[test] fn test_toml() { let opts = FrontendOptions::default(); @@ -225,6 +277,8 @@ mod tests { assert!(toml_string.contains("experimental_enable_exponential_histogram = false")); let parsed: FrontendOptions = toml::from_str(&toml_string).unwrap(); assert_eq!(parsed.otlp, opts.otlp); + assert_eq!(parsed.influxdb, opts.influxdb); + assert_eq!(parsed.opentsdb, opts.opentsdb); } #[test] diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index f834722cca5..7225987d752 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1854,9 +1854,9 @@ mod tests { use tokio::sync::{mpsc, oneshot}; use tower::ServiceExt; - use super::*; use crate::frontend::FrontendOptions; use crate::instance::builder::FrontendBuilder; + use crate::instance::*; fn parse_test_sql(sql: &str) -> Vec { parse_stmt(sql, &GreptimeDbDialect {}).unwrap() diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index e49b75bbf89..7a6ba474168 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -24,7 +24,7 @@ use catalog::system_schema::semantic_graph::EntityGraphProviderRef; use common_base::Plugins; use common_datasource::object_store::LocalFileAccess; use common_event_recorder::{EventRecorderImpl, EventRecorderRef}; -use common_meta::cache::{LayeredCacheRegistryRef, TableRouteCacheRef}; +use common_meta::cache::{LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableRouteCacheRef}; use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator}; use common_meta::key::TableMetadataManager; use common_meta::key::flow::FlowMetadataManager; @@ -32,6 +32,7 @@ use common_meta::kv_backend::KvBackendRef; use common_meta::node_manager::NodeManagerRef; use common_meta::procedure_executor::ProcedureExecutorRef; use dashmap::DashMap; +use operator::batcher::PendingRowsBatcher; use operator::delete::Deleter; use operator::flow::FlowServiceOperator; use operator::insert::Inserter; @@ -49,6 +50,8 @@ use partition::manager::PartitionRuleManager; use pipeline::pipeline_operator::PipelineOperator; use query::QueryEngineFactory; use query::region_query::RegionQueryHandlerFactoryRef; +use servers::batcher::table::TablePendingRowsBatcher; +use servers::http::BatchingProtocol; use snafu::{OptionExt, ResultExt}; use crate::error::{self, DataFusionSnafu, ExternalSnafu, Result}; @@ -58,6 +61,7 @@ use crate::heartbeat::frontend_peer_addr; use crate::instance::Instance; use crate::instance::entity_graph::EntityGraphProviderImpl; use crate::instance::region_query::FrontendRegionQueryHandler; +use crate::service_config::PendingRowsBatcherOptions; /// The frontend [`Instance`] builder. pub struct FrontendBuilder { @@ -215,20 +219,50 @@ impl FrontendBuilder { FrontendRegionQueryHandler::arc(partition_manager.clone(), node_manager.clone()) }; - let table_flownode_cache = - self.layered_cache_registry - .get() - .context(error::CacheRequiredSnafu { - name: TABLE_FLOWNODE_SET_CACHE_NAME, - })?; + let table_flownode_cache: TableFlownodeSetCacheRef = self + .layered_cache_registry + .get() + .context(error::CacheRequiredSnafu { + name: TABLE_FLOWNODE_SET_CACHE_NAME, + })?; - let inserter = Arc::new(Inserter::new( - self.catalog_manager.clone(), - partition_manager.clone(), - node_manager.clone(), - table_flownode_cache, - self.options.auto_create_table, - )); + let create_inserter = || { + Inserter::new( + self.catalog_manager.clone(), + partition_manager.clone(), + node_manager.clone(), + table_flownode_cache.clone(), + self.options.auto_create_table, + ) + }; + // The execution-only inserter owns no batchers, avoiding an Arc cycle. + let bulk_inserter = Arc::new(create_inserter()); + let build_batcher = + |options: &PendingRowsBatcherOptions| -> Option> { + if !options.pending_rows_batching_enabled() + || (self.options.prom_store.with_metric_engine + && options + .protocols + .iter() + .all(|protocol| *protocol == BatchingProtocol::Prom)) + { + return None; + } + TablePendingRowsBatcher::try_new( + options.pending_rows_flush_interval, + options.max_batch_rows, + options.max_concurrent_flushes, + options.worker_channel_capacity, + options.max_inflight_requests, + options.flow_notification_queue_capacity, + bulk_inserter.clone(), + ) + .map(|batcher| batcher as Arc) + }; + let inserter = Arc::new( + create_inserter() + .with_pending_rows_batcher(build_batcher(&self.options.pending_rows_batcher)), + ); let deleter = Arc::new(Deleter::new( self.catalog_manager.clone(), partition_manager.clone(), diff --git a/src/frontend/src/instance/influxdb.rs b/src/frontend/src/instance/influxdb.rs index c23065394f1..8e4d0af87d8 100644 --- a/src/frontend/src/instance/influxdb.rs +++ b/src/frontend/src/instance/influxdb.rs @@ -208,7 +208,7 @@ mod tests { use session::context::QueryContext; use store_api::mito_engine_options::MERGE_MODE_KEY; - use super::*; + use crate::instance::influxdb::*; use crate::service_config::influxdb::InfluxdbMergeMode; #[test] diff --git a/src/frontend/src/instance/opentsdb.rs b/src/frontend/src/instance/opentsdb.rs index e5fda529594..17908519ad3 100644 --- a/src/frontend/src/instance/opentsdb.rs +++ b/src/frontend/src/instance/opentsdb.rs @@ -21,7 +21,7 @@ use auth::{ }; use common_error::ext::BoxedError; use common_telemetry::tracing; -use servers::error::{self as server_error, AuthSnafu, ExecuteGrpcQuerySnafu}; +use servers::error::{AuthSnafu, ExecuteGrpcQuerySnafu, Result as ServerResult}; use servers::opentsdb::codec::DataPoint; use servers::opentsdb::data_point_to_grpc_row_insert_requests; use servers::query_handler::OpentsdbProtocolHandler; @@ -44,11 +44,7 @@ fn permission_targets(data_points: &[DataPoint], ctx: &QueryContextRef) -> Permi #[async_trait] impl OpentsdbProtocolHandler for Instance { - async fn preflight( - &self, - data_points: &[DataPoint], - ctx: QueryContextRef, - ) -> server_error::Result<()> { + async fn preflight(&self, data_points: &[DataPoint], ctx: QueryContextRef) -> ServerResult<()> { self.check_table_permission( &ctx, PermissionReq::Action(OPENTSDB_WRITE), @@ -59,11 +55,30 @@ impl OpentsdbProtocolHandler for Instance { } #[tracing::instrument(skip_all, fields(protocol = "opentsdb"))] - async fn exec( + async fn exec(&self, data_points: Vec, ctx: QueryContextRef) -> ServerResult { + // Keep diagnostic per-point errors independent of other batched writes. + let mut ctx = ctx.fork(); + ctx.set_batching_enabled(false); + self.execute_opentsdb_write(data_points, Arc::new(ctx)) + .await + } + + #[tracing::instrument(skip_all, fields(protocol = "opentsdb"))] + async fn exec_batch( &self, data_points: Vec, ctx: QueryContextRef, - ) -> server_error::Result { + ) -> ServerResult { + self.execute_opentsdb_write(data_points, ctx).await + } +} + +impl Instance { + async fn execute_opentsdb_write( + &self, + data_points: Vec, + ctx: QueryContextRef, + ) -> ServerResult { self.plugins .get::() .as_ref() @@ -99,7 +114,7 @@ impl OpentsdbProtocolHandler for Instance { mod tests { use session::context::QueryContext; - use super::*; + use crate::instance::opentsdb::*; #[test] fn test_permission_targets_do_not_require_row_conversion() { diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 1812dc71a51..257e5243e89 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -33,7 +33,7 @@ use servers::grpc::{GrpcOptions, GrpcServer}; use servers::http::event::LogValidatorRef; use servers::http::result::error_result::ErrorResponse; use servers::http::utils::router::RouterConfigurator; -use servers::http::{HttpOptions, HttpServer, HttpServerBuilder}; +use servers::http::{BatchingProtocol, HttpOptions, HttpServer, HttpServerBuilder}; use servers::interceptor::LogIngestInterceptorRef; use servers::metrics_handler::MetricsHandler; use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef}; @@ -49,6 +49,7 @@ use tonic::Status; use crate::error::{self, Result, StartServerSnafu, TomlFormatSnafu}; use crate::frontend::FrontendOptions; use crate::instance::Instance; +use crate::service_config::PromStoreOptions; pub struct Services where @@ -106,6 +107,7 @@ where request_memory_limiter: ServerMemoryLimiter, ) -> HttpServerBuilder { let mut builder = HttpServerBuilder::new(effective_http_options(opts)) + .with_batching_protocols(opts.pending_rows_batcher.protocols.clone()) .with_memory_limiter(request_memory_limiter) .with_sql_handler(self.instance.clone()); @@ -127,21 +129,22 @@ where builder = builder.with_influxdb_handler(self.instance.clone()); } - if opts.prom_store.enable { - let pending_rows_batcher = if opts.prom_store.with_metric_engine { + let prom_store = effective_prom_store_options(opts); + if prom_store.enable { + let pending_rows_batcher = if prom_store.with_metric_engine { PendingRowsBatcher::try_new( self.instance.partition_manager().clone(), self.instance.node_manager().clone(), self.instance.catalog_manager().clone(), self.instance.table_flownode_set_cache().clone(), - opts.prom_store.with_metric_engine, + prom_store.with_metric_engine, self.instance.clone(), - opts.prom_store.pending_rows_flush_interval, - opts.prom_store.max_batch_rows, - opts.prom_store.max_concurrent_flushes, - opts.prom_store.worker_channel_capacity, - opts.prom_store.max_inflight_requests, - opts.prom_store.flow_notification_queue_capacity, + prom_store.pending_rows_flush_interval, + prom_store.max_batch_rows, + prom_store.max_concurrent_flushes, + prom_store.worker_channel_capacity, + prom_store.max_inflight_requests, + prom_store.flow_notification_queue_capacity, ) } else { None @@ -426,22 +429,45 @@ where } } +/// Selected shared controls override legacy Prom batching knobs, not protocol behavior. +fn effective_prom_store_options(opts: &FrontendOptions) -> PromStoreOptions { + let mut prom_store = opts.prom_store.clone(); + let shared = &opts.pending_rows_batcher; + if shared.protocols.contains(&BatchingProtocol::Prom) && shared.pending_rows_batching_enabled() + { + prom_store.pending_rows_flush_interval = shared.pending_rows_flush_interval; + prom_store.max_batch_rows = shared.max_batch_rows; + prom_store.max_concurrent_flushes = shared.max_concurrent_flushes; + prom_store.worker_channel_capacity = shared.worker_channel_capacity; + prom_store.max_inflight_requests = shared.max_inflight_requests; + prom_store.flow_notification_queue_capacity = shared.flow_notification_queue_capacity; + } + prom_store +} + fn effective_http_options(opts: &FrontendOptions) -> HttpOptions { effective_http_options_with_sync(opts, pending_rows_batch_sync_enabled()) } fn effective_http_options_with_sync(opts: &FrontendOptions, batch_sync: bool) -> HttpOptions { let mut http = opts.http.clone(); - let flush_interval = opts.prom_store.pending_rows_flush_interval; + let prom_store = effective_prom_store_options(opts); + let shared = &opts.pending_rows_batcher; + // Ordinary-table batching always waits for its flush, independently of the + // dedicated Prom batcher's asynchronous acknowledgement mode. + let common_enabled = shared.pending_rows_batching_enabled() + && shared.protocols.iter().any(|protocol| { + *protocol != BatchingProtocol::Prom + || (prom_store.enable && !prom_store.with_metric_engine) + }); + let common_interval = common_enabled.then_some(shared.pending_rows_flush_interval); + let prom_interval = (prom_store.pending_rows_batching_enabled() && batch_sync) + .then_some(prom_store.pending_rows_flush_interval); + let Some(flush_interval) = common_interval.into_iter().chain(prom_interval).max() else { + return http; + }; let fallback_timeout = flush_interval.saturating_add(Duration::from_secs(1)); - // In asynchronous batch mode submissions return right after enqueue and - // no request waits for a pending-row flush, so the timeout must not be - // raised either. - if !opts.prom_store.pending_rows_batching_enabled() - || !batch_sync - || http.timeout.is_zero() - || http.timeout > fallback_timeout - { + if http.timeout.is_zero() || http.timeout > fallback_timeout { return http; } @@ -462,6 +488,7 @@ fn parse_addr(addr: &str) -> Result { #[cfg(test)] mod tests { + use std::num::NonZeroUsize; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; @@ -478,8 +505,98 @@ mod tests { use servers::grpc::flight::{FlightCraft, FlightCraftRef, TonicStream}; use tonic::{Code, Request, Response, Status, Streaming}; - use super::*; use crate::instance::builder::FrontendBuilder; + use crate::server::*; + + #[test] + fn test_effective_prom_batching_controls() { + // Only an enabled shared Prom selection replaces the legacy controls. + for (selected, shared_enabled, metric_engine, prom_enabled) in [ + (true, true, true, true), + (false, true, true, true), + (true, false, true, true), + (true, true, false, true), + (true, true, true, false), + ] { + let mut opts = FrontendOptions::default(); + opts.http.timeout = Duration::from_millis(1); + opts.prom_store.pending_rows_flush_interval = Duration::from_secs(2); + opts.prom_store.with_metric_engine = metric_engine; + opts.prom_store.enable = prom_enabled; + opts.prom_store + .experimental_enable_prometheus_native_histogram = true; + let shared = &mut opts.pending_rows_batcher; + shared.protocols = vec![if selected { + BatchingProtocol::Prom + } else { + BatchingProtocol::Influxdb + }]; + shared.pending_rows_flush_interval = if shared_enabled { + Duration::from_secs(5) + } else { + Duration::ZERO + }; + shared.max_batch_rows = 7; + shared.max_concurrent_flushes = 3; + shared.worker_channel_capacity = 11; + shared.max_inflight_requests = 13; + shared.flow_notification_queue_capacity = NonZeroUsize::new(17).unwrap(); + + let mut expected = opts.prom_store.clone(); + if selected && shared_enabled { + expected.pending_rows_flush_interval = shared.pending_rows_flush_interval; + expected.max_batch_rows = shared.max_batch_rows; + expected.max_concurrent_flushes = shared.max_concurrent_flushes; + expected.worker_channel_capacity = shared.worker_channel_capacity; + expected.max_inflight_requests = shared.max_inflight_requests; + expected.flow_notification_queue_capacity = shared.flow_notification_queue_capacity; + } + let actual = effective_prom_store_options(&opts); + assert_eq!(actual, expected); + assert_eq!( + actual.pending_rows_batching_enabled(), + metric_engine && prom_enabled + ); + } + } + + #[test] + fn test_http_timeout_covers_synchronous_batchers() { + // Shared ordinary writes remain synchronous even when Prom is asynchronous. + for ( + protocols, + metric_engine, + batch_sync, + shared_secs, + legacy_secs, + timeout_secs, + expected_secs, + ) in [ + (vec![BatchingProtocol::Prom], false, false, 5, 2, 1, 6), + (vec![BatchingProtocol::Prom], true, false, 5, 2, 1, 1), + (vec![BatchingProtocol::Prom], true, true, 5, 2, 1, 6), + (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 1, 6), + (vec![BatchingProtocol::Influxdb], true, true, 5, 8, 1, 9), + (vec![BatchingProtocol::Influxdb], true, true, 8, 5, 1, 9), + (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 0, 0), + (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 10, 10), + (vec![BatchingProtocol::Prom], false, false, 0, 2, 1, 1), + (vec![], true, false, 5, 2, 1, 1), + (vec![], true, true, 5, 2, 1, 3), + ] { + let mut opts = FrontendOptions::default(); + opts.http.timeout = Duration::from_secs(timeout_secs); + opts.prom_store.with_metric_engine = metric_engine; + opts.prom_store.pending_rows_flush_interval = Duration::from_secs(legacy_secs); + opts.pending_rows_batcher.protocols = protocols; + opts.pending_rows_batcher.pending_rows_flush_interval = + Duration::from_secs(shared_secs); + assert_eq!( + effective_http_options_with_sync(&opts, batch_sync).timeout, + Duration::from_secs(expected_secs) + ); + } + } struct CountingFlightCraft { inner: FlightCraftRef, @@ -561,6 +678,41 @@ mod tests { ); } + #[test] + fn test_invalid_shared_batching_preserves_prom_store_options() { + type KnobMutator = fn(&mut FrontendOptions); + let cases: [KnobMutator; 5] = [ + |opts| opts.pending_rows_batcher.max_concurrent_flushes = usize::MAX, + |opts| opts.pending_rows_batcher.worker_channel_capacity = usize::MAX, + |opts| opts.pending_rows_batcher.max_inflight_requests = usize::MAX, + |opts| { + opts.pending_rows_batcher.flow_notification_queue_capacity = + NonZeroUsize::new(usize::MAX).unwrap() + }, + |opts| opts.pending_rows_batcher.pending_rows_flush_interval = Duration::MAX, + ]; + for invalidate in cases { + let mut opts = FrontendOptions::default(); + opts.http.timeout = Duration::from_secs(1); + opts.prom_store.pending_rows_flush_interval = Duration::from_secs(5); + opts.pending_rows_batcher.protocols = + vec![BatchingProtocol::Prom, BatchingProtocol::Influxdb]; + opts.pending_rows_batcher.pending_rows_flush_interval = Duration::from_secs(10); + invalidate(&mut opts); + assert!(!opts.pending_rows_batcher.pending_rows_batching_enabled()); + assert_eq!(opts.prom_store, effective_prom_store_options(&opts)); + assert_eq!( + Duration::from_secs(6), + effective_http_options_with_sync(&opts, true).timeout + ); + opts.prom_store.pending_rows_flush_interval = Duration::ZERO; + assert_eq!( + Duration::from_secs(1), + effective_http_options_with_sync(&opts, true).timeout + ); + } + } + #[test] fn test_effective_http_timeout_skips_fallback_when_batcher_disabled() { // Mirrors the conditions under which `PendingRowsBatcher::try_new` diff --git a/src/frontend/src/service_config.rs b/src/frontend/src/service_config.rs index 03ec12e2c9f..bcd1097cfcf 100644 --- a/src/frontend/src/service_config.rs +++ b/src/frontend/src/service_config.rs @@ -17,6 +17,7 @@ pub mod jaeger; pub mod mysql; pub mod opentsdb; pub mod otlp; +pub mod pending_rows_batcher; pub mod postgres; pub mod prom_store; @@ -25,5 +26,6 @@ pub use jaeger::JaegerOptions; pub use mysql::MysqlOptions; pub use opentsdb::OpentsdbOptions; pub use otlp::OtlpOptions; +pub use pending_rows_batcher::PendingRowsBatcherOptions; pub use postgres::PostgresOptions; pub use prom_store::PromStoreOptions; diff --git a/src/frontend/src/service_config/influxdb.rs b/src/frontend/src/service_config/influxdb.rs index 912e342471d..af28dee935f 100644 --- a/src/frontend/src/service_config/influxdb.rs +++ b/src/frontend/src/service_config/influxdb.rs @@ -49,7 +49,7 @@ impl Default for InfluxdbOptions { #[cfg(test)] mod tests { - use super::InfluxdbOptions; + use crate::service_config::influxdb::InfluxdbOptions; #[test] fn test_influxdb_options() { diff --git a/src/frontend/src/service_config/pending_rows_batcher.rs b/src/frontend/src/service_config/pending_rows_batcher.rs new file mode 100644 index 00000000000..ef7ce14ef02 --- /dev/null +++ b/src/frontend/src/service_config/pending_rows_batcher.rs @@ -0,0 +1,150 @@ +// 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. + +use std::num::NonZeroUsize; +use std::time::Duration; + +use common_batcher::flush_policy::timing::TimingFlushPolicy; +use serde::{Deserialize, Serialize}; +use servers::http::BatchingProtocol; +use tokio::sync::Semaphore; + +/// Experimental table write batching options shared by HTTP ingestion protocols. +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default)] +pub struct PendingRowsBatcherOptions { + /// HTTP write protocols sharing this batcher; empty disables all entrances. + pub protocols: Vec, + /// Time from the first pending submission to a timed flush. Zero disables batching. + #[serde(with = "humantime_serde")] + pub pending_rows_flush_interval: Duration, + /// Row threshold checked after appending a complete submission. + pub max_batch_rows: usize, + /// Maximum concurrent flushes shared by the frontend batcher. + pub max_concurrent_flushes: usize, + /// Maximum number of queued submissions per table worker. + pub worker_channel_capacity: usize, + /// Maximum number of original requests awaiting completion. + pub max_inflight_requests: usize, + /// Maximum number of table Flow notifications waiting in the shared queue. + pub flow_notification_queue_capacity: NonZeroUsize, +} + +impl PendingRowsBatcherOptions { + /// Returns whether a protocol opts in and its controls pass construction validation. + pub fn pending_rows_batching_enabled(&self) -> bool { + !self.protocols.is_empty() + && TimingFlushPolicy::validate(self.pending_rows_flush_interval) + && self.max_batch_rows > 0 + && [ + self.max_concurrent_flushes, + self.worker_channel_capacity, + self.max_inflight_requests, + self.flow_notification_queue_capacity.get(), + ] + .into_iter() + .all(|capacity| (1..=Semaphore::MAX_PERMITS).contains(&capacity)) + } +} + +impl Default for PendingRowsBatcherOptions { + fn default() -> Self { + Self { + protocols: Vec::new(), + pending_rows_flush_interval: Duration::ZERO, + max_batch_rows: 100_000, + max_concurrent_flushes: 256, + worker_channel_capacity: 65526, + max_inflight_requests: 3000, + flow_notification_queue_capacity: NonZeroUsize::new(1024).unwrap_or(NonZeroUsize::MIN), + } + } +} + +#[cfg(test)] +mod tests { + use crate::service_config::pending_rows_batcher::*; + + #[test] + fn test_protocols() { + let options: PendingRowsBatcherOptions = toml::from_str( + "protocols = ['influxdb', 'opentsdb', 'otlp', 'logs', 'loki', 'splunk', 'elasticsearch', 'http_sql', 'prom']", + ).unwrap(); + assert_eq!(options.protocols.len(), 9); + assert!(options.protocols.contains(&BatchingProtocol::HttpSql)); + assert!(PendingRowsBatcherOptions::default().protocols.is_empty()); + for invalid in ["sql", "jaeger", "unknown"] { + assert!( + toml::from_str::(&format!("protocols = ['{invalid}']")) + .is_err() + ); + } + } + + #[test] + fn test_notification_capacity() { + let default = PendingRowsBatcherOptions::default(); + assert_eq!(default.flow_notification_queue_capacity.get(), 1024); + let configured: PendingRowsBatcherOptions = + toml::from_str("flow_notification_queue_capacity = 8").unwrap(); + assert_eq!(configured.flow_notification_queue_capacity.get(), 8); + assert!( + toml::from_str::("flow_notification_queue_capacity = 0") + .is_err() + ); + } + + #[test] + fn test_defaults_and_roundtrip() { + let options: PendingRowsBatcherOptions = toml::from_str("").unwrap(); + assert_eq!(options, PendingRowsBatcherOptions::default()); + assert!(!options.pending_rows_batching_enabled()); + assert_eq!(options.max_batch_rows, 100_000); + assert_eq!(options.max_concurrent_flushes, 256); + assert_eq!(options.worker_channel_capacity, 65526); + assert_eq!(options.max_inflight_requests, 3000); + let serialized = toml::to_string(&options).unwrap(); + assert_eq!( + options, + toml::from_str::(&serialized).unwrap() + ); + } + + #[test] + fn test_partial_options_and_zero_controls() { + let options: PendingRowsBatcherOptions = + toml::from_str("pending_rows_flush_interval = '5ms'").unwrap(); + assert_eq!( + options.pending_rows_flush_interval, + Duration::from_millis(5) + ); + assert!(!options.pending_rows_batching_enabled()); + let enabled: PendingRowsBatcherOptions = + toml::from_str("protocols = ['http_sql']\npending_rows_flush_interval = '5ms'") + .unwrap(); + assert!(enabled.pending_rows_batching_enabled()); + for field in [ + "max_batch_rows", + "max_concurrent_flushes", + "worker_channel_capacity", + "max_inflight_requests", + ] { + let options: PendingRowsBatcherOptions = toml::from_str(&format!( + "protocols = ['influxdb']\npending_rows_flush_interval = '5ms'\n{field} = 0" + )) + .unwrap(); + assert!(!options.pending_rows_batching_enabled(), "{field}"); + } + } +} diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index d88006b18ad..7fb538dc775 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -24,7 +24,7 @@ use api::v1::region::{ }; use api::v1::{ AlterTableExpr, ColumnDataType, ColumnSchema, CreateTableExpr, InsertRequests, - RowInsertRequest, RowInsertRequests, SemanticType, + RowInsertRequest, RowInsertRequests, Rows, SemanticType, }; use catalog::CatalogManagerRef; use client::{OutputData, OutputMeta}; @@ -65,7 +65,7 @@ use store_api::mito_engine_options::{ }; use store_api::storage::{RegionId, TableId}; use table::TableRef; -use table::metadata::TableInfo; +use table::metadata::{TableInfo, TableInfoRef}; use table::requests::{ AUTO_CREATE_TABLE_KEY, InsertRequest as TableInsertRequest, SEMANTIC_PER_TABLE_INDEX_KEY, TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1, TRACE_TABLE_PARTITIONS_HINT_KEY, @@ -73,6 +73,7 @@ use table::requests::{ }; use table::table_reference::TableReference; +use crate::batcher::PendingRowsBatcher; use crate::error::{ CatalogSnafu, ColumnOptionsSnafu, CreatePartitionRulesSnafu, FindRegionLeaderSnafu, InvalidInsertRequestSnafu, JoinTaskSnafu, RequestInsertsSnafu, Result, TableNotFoundSnafu, @@ -81,7 +82,8 @@ use crate::expr_helper; use crate::region_req_factory::RegionRequestFactory; use crate::req_convert::common::preprocess_row_insert_requests; use crate::req_convert::insert::{ - ColumnToRow, RowToRegion, StatementToRegion, TableToRegion, fill_reqs_with_impure_default, + ColumnToRow, ImpureDefaultFiller, RowToRegion, StatementToRegion, TableToRegion, + fill_reqs_with_impure_default, rows_to_record_batch, }; use crate::statement::StatementExecutor; @@ -94,6 +96,7 @@ pub struct Inserter { /// When `false`, missing tables are never auto-created regardless of the /// per-request `auto_create_table` hint. When `true`, the hint still applies. auto_create_table: bool, + pending_rows_batcher: Option>, } pub type InserterRef = Arc; @@ -163,9 +166,19 @@ impl Inserter { node_manager, table_flownode_set_cache, auto_create_table, + pending_rows_batcher: None, } } + /// Installs the shared batcher; callers explicitly select its ingestion entry point. + pub fn with_pending_rows_batcher( + mut self, + batcher: Option>, + ) -> Self { + self.pending_rows_batcher = batcher; + self + } + pub async fn handle_column_inserts( &self, requests: InsertRequests, @@ -267,6 +280,11 @@ impl Inserter { ) -> Result { let skip_wal = ctx.skip_wal(); + let batcher = self + .pending_rows_batcher + .as_ref() + .filter(|_| ctx.batching_enabled()); + // remove empty requests requests.inserts.retain(|req| { req.rows @@ -290,6 +308,19 @@ impl Inserter { ) .await?; + // Instant tables have no persisted data for dirty-window Flow to read. + // Metric tables keep their existing dedicated ingestion path. + if let Some(batcher) = batcher + && instant_table_ids.is_empty() + && table_infos + .values() + .all(|info| info.meta.engine == default_engine()) + { + return self + .submit_pending_rows(requests, table_infos, ctx, batcher) + .await; + } + let name_to_info = table_infos .values() .map(|info| (info.name.clone(), info.clone())) @@ -305,6 +336,85 @@ impl Inserter { self.do_request(inserts, &table_infos, &ctx).await } + async fn submit_pending_rows( + &self, + mut requests: RowInsertRequests, + table_infos: HashMap>, + ctx: QueryContextRef, + batcher: &Arc, + ) -> Result { + // All entry points, including single-table and SQL writes, skip empty input + // before evaluating defaults or converting prepared rows. + requests.inserts.retain(|request| { + request + .rows + .as_ref() + .is_some_and(|rows| !rows.rows.is_empty()) + }); + let by_name = table_infos + .values() + .map(|info| (info.name.as_str(), info)) + .collect::>(); + let mut prepared = Vec::with_capacity(requests.inserts.len()); + for request in &mut requests.inserts { + let table_info = + by_name + .get(request.table_name.as_str()) + .context(TableNotFoundSnafu { + table_name: &request.table_name, + })?; + let Some(rows) = &mut request.rows else { + continue; + }; + ImpureDefaultFiller::new((*table_info).clone())?.fill_rows(rows); + let batch = rows_to_record_batch(rows, table_info)?; + prepared.push(((*table_info).clone(), batch)); + } + // Preserve the existing meter input and original request boundary. These + // envelopes are only for accounting; routing happens after batching. + let metered = InstantAndNormalInsertRequests { + normal_requests: RegionInsertRequests { + requests: requests + .inserts + .into_iter() + .map(|request| RegionInsertRequest { + rows: request.rows, + ..Default::default() + }) + .collect(), + }, + instant_requests: RegionInsertRequests::default(), + }; + let write_cost = write_meter!( + ctx.current_catalog(), + ctx.current_schema(), + metered, + ctx.channel() as u8 + ); + prepared.retain(|(_, batch)| batch.num_rows() != 0); + let results = if prepared.is_empty() { + Vec::new() + } else { + // One original request shares admission across all table submissions. + let permit = batcher.acquire().await?; + let submissions = prepared.into_iter().map(|(info, batch)| { + // Routing uses the target database; metering above retains the + // original request context, including fully qualified SQL writes. + let mut target_ctx = ctx.fork(); + target_ctx.set_current_catalog(&info.catalog_name); + target_ctx.set_current_schema(&info.schema_name); + batcher.submit(info, batch, Arc::new(target_ctx), permit.clone()) + }); + // 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 _), + )) + } + /// Handles row inserts request with metric engine. pub async fn handle_metric_row_inserts( &self, @@ -353,6 +463,36 @@ impl Inserter { self.do_request(inserts, &table_infos, &ctx).await } + fn table_batcher( + &self, + table_info: &TableInfoRef, + ctx: &QueryContextRef, + ) -> Option<&Arc> { + self.pending_rows_batcher.as_ref().filter(|_| { + ctx.batching_enabled() + && !table_info.is_ttl_instant_table() + && table_info.meta.engine == default_engine() + }) + } + + async fn submit_table_rows( + &self, + rows: Rows, + table_info: TableInfoRef, + ctx: QueryContextRef, + batcher: &Arc, + ) -> Result { + let requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: table_info.name.clone(), + rows: Some(rows), + }], + }; + let table_infos = HashMap::from_iter([(table_info.table_id(), table_info)]); + self.submit_pending_rows(requests, table_infos, ctx, batcher) + .await + } + pub async fn handle_table_insert( &self, request: TableInsertRequest, @@ -367,9 +507,13 @@ impl Inserter { })?; let table_info = table.table_info(); - let inserts = TableToRegion::new(&table_info, &self.partition_manager) - .convert(request) - .await?; + let converter = TableToRegion::new(&table_info, &self.partition_manager); + let skip_wal = request.skip_wal; + let rows = converter.prepare(request)?; + if let Some(batcher) = self.table_batcher(&table_info, &ctx) { + return self.submit_table_rows(rows, table_info, ctx, batcher).await; + } + let inserts = converter.partition(rows, skip_wal).await?; let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]); @@ -381,10 +525,15 @@ impl Inserter { insert: &Insert, ctx: &QueryContextRef, ) -> Result { - let (inserts, table_info) = - StatementToRegion::new(self.catalog_manager.as_ref(), &self.partition_manager, ctx) - .convert(insert, ctx) - .await?; + let converter = + StatementToRegion::new(self.catalog_manager.as_ref(), &self.partition_manager, ctx); + let (rows, table_info) = converter.prepare(insert, ctx).await?; + if let Some(batcher) = self.table_batcher(&table_info, ctx) { + return self + .submit_table_rows(rows, table_info, ctx.clone(), batcher) + .await; + } + let inserts = converter.partition(rows, table_info.clone(), ctx).await?; let table_infos = HashMap::from_iter([(table_info.table_id(), table_info.clone())]); @@ -1537,7 +1686,7 @@ mod tests { use table::dist_table::DummyDataSource; use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType}; - use super::*; + use crate::insert::*; use crate::test_util::{create_partition_rule_manager, prepare_mocked_backend}; fn make_table_ref_with_schema( @@ -1933,4 +2082,92 @@ mod tests { table_options.get(MERGE_MODE_KEY).map(String::as_str) ); } + + struct UnexpectedBatcher; + + #[async_trait::async_trait] + impl PendingRowsBatcher for UnexpectedBatcher { + async fn acquire(&self) -> Result> { + panic!("empty writes must not acquire batch admission") + } + + async fn submit( + &self, + _table_info: TableInfoRef, + _batch: arrow::record_batch::RecordBatch, + _ctx: QueryContextRef, + _permit: Arc, + ) -> Result { + panic!("empty writes must not submit a batch") + } + } + + async fn batcher_test_inserter() -> Inserter { + let kv_backend = prepare_mocked_backend().await; + Inserter::new( + catalog::memory::MemoryCatalogManager::new(), + create_partition_rule_manager(kv_backend.clone()).await, + Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)), + Arc::new(new_table_flownode_set_cache( + String::new(), + Cache::new(100), + kv_backend, + )), + true, + ) + } + + #[tokio::test] + async fn test_instant_table_bypasses_batcher() { + let batcher: Arc = Arc::new(UnexpectedBatcher); + let inserter = batcher_test_inserter() + .await + .with_pending_rows_batcher(Some(batcher)); + let mut ctx = session::context::QueryContextBuilder::default().build(); + ctx.set_batching_enabled(true); + let ctx = Arc::new(ctx); + let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype()) + .table_info(); + assert!(inserter.table_batcher(&table, &ctx).is_some()); + let mut instant = (*table).clone(); + instant.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant); + assert!(inserter.table_batcher(&Arc::new(instant), &ctx).is_none()); + } + + #[tokio::test] + async fn test_empty_prepared_rows_skip_batcher() { + let inserter = batcher_test_inserter().await; + let table = make_table_ref_with_schema("ts", "value", ConcreteDataType::float64_datatype()) + .table_info(); + let batcher: Arc = Arc::new(UnexpectedBatcher); + let ctx = QueryContext::arc(); + let output = inserter + .submit_table_rows( + Rows { + schema: vec![], + rows: vec![], + }, + table.clone(), + ctx.clone(), + &batcher, + ) + .await + .unwrap(); + assert!(matches!(output.data, OutputData::AffectedRows(0))); + let output = inserter + .submit_pending_rows( + RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: table.name.clone(), + rows: None, + }], + }, + HashMap::from_iter([(table.table_id(), table)]), + ctx, + &batcher, + ) + .await + .unwrap(); + assert!(matches!(output.data, OutputData::AffectedRows(0))); + } } diff --git a/src/operator/src/req_convert/insert.rs b/src/operator/src/req_convert/insert.rs index e36caf2f8c3..079e98d4f9c 100644 --- a/src/operator/src/req_convert/insert.rs +++ b/src/operator/src/req_convert/insert.rs @@ -22,6 +22,7 @@ mod timestamps; use api::v1::SemanticType; pub use column_to_row::ColumnToRow; +pub(crate) use fill_impure_default::ImpureDefaultFiller; pub use fill_impure_default::fill_reqs_with_impure_default; pub use row_to_batch::rows_to_record_batch; pub use row_to_region::RowToRegion; diff --git a/src/operator/src/req_convert/insert/stmt_to_region.rs b/src/operator/src/req_convert/insert/stmt_to_region.rs index e811facf74d..50ae90d587f 100644 --- a/src/operator/src/req_convert/insert/stmt_to_region.rs +++ b/src/operator/src/req_convert/insert/stmt_to_region.rs @@ -64,6 +64,41 @@ impl<'a> StatementToRegion<'a> { stmt: &Insert, query_ctx: &QueryContextRef, ) -> Result<(InstantAndNormalInsertRequests, TableInfoRef)> { + let (rows, table_info) = self.prepare(stmt, query_ctx).await?; + let requests = self.partition(rows, table_info.clone(), query_ctx).await?; + Ok((requests, table_info)) + } + + /// Routes already prepared rows while retaining TTL and WAL behavior. + pub async fn partition( + &self, + rows: Rows, + table_info: TableInfoRef, + query_ctx: &QueryContextRef, + ) -> Result { + let requests = Partitioner::new(self.partition_manager) + .partition_insert_requests(&table_info, rows, query_ctx.skip_wal()) + .await?; + let requests = RegionInsertRequests { requests }; + if table_info.is_ttl_instant_table() { + Ok(InstantAndNormalInsertRequests { + normal_requests: Default::default(), + instant_requests: requests, + }) + } else { + Ok(InstantAndNormalInsertRequests { + normal_requests: requests, + instant_requests: Default::default(), + }) + } + } + + /// Resolves SQL values and their table schema without partition routing. + pub async fn prepare( + &self, + stmt: &Insert, + query_ctx: &QueryContextRef, + ) -> Result<(Rows, TableInfoRef)> { let name = stmt.table_name().context(ParseSqlSnafu)?; let (catalog, schema, table_name) = self.get_full_name(name)?; let table = self.get_table(&catalog, &schema, &table_name).await?; @@ -153,27 +188,7 @@ impl<'a> StatementToRegion<'a> { schema.push(grpc_column_schema); } - let requests = Partitioner::new(self.partition_manager) - .partition_insert_requests(&table_info, Rows { schema, rows }, query_ctx.skip_wal()) - .await?; - let requests = RegionInsertRequests { requests }; - if table_info.is_ttl_instant_table() { - Ok(( - InstantAndNormalInsertRequests { - normal_requests: Default::default(), - instant_requests: requests, - }, - table_info, - )) - } else { - Ok(( - InstantAndNormalInsertRequests { - normal_requests: requests, - instant_requests: Default::default(), - }, - table_info, - )) - } + Ok((Rows { schema, rows }, table_info)) } async fn get_table(&self, catalog: &str, schema: &str, table: &str) -> Result { diff --git a/src/operator/src/req_convert/insert/table_to_region.rs b/src/operator/src/req_convert/insert/table_to_region.rs index adc40118ec3..4ee7f0201e4 100644 --- a/src/operator/src/req_convert/insert/table_to_region.rs +++ b/src/operator/src/req_convert/insert/table_to_region.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use api::helper::vectors_to_rows; use api::v1::Rows; use api::v1::region::InsertRequests as RegionInsertRequests; use partition::manager::PartitionRuleManager; @@ -36,17 +37,31 @@ impl<'a> TableToRegion<'a> { } } + /// Converts column vectors into rows without partition routing. + pub fn prepare(&self, request: TableInsertRequest) -> Result { + let row_count = row_count(&request.columns_values)?; + let schema = column_schema(self.table_info, &request.columns_values)?; + let rows = vectors_to_rows(request.columns_values.values(), row_count); + Ok(Rows { schema, rows }) + } + pub async fn convert( &self, request: TableInsertRequest, ) -> Result { - let row_count = row_count(&request.columns_values)?; - let schema = column_schema(self.table_info, &request.columns_values)?; - let rows = api::helper::vectors_to_rows(request.columns_values.values(), row_count); + let skip_wal = request.skip_wal; + let rows = self.prepare(request)?; + self.partition(rows, skip_wal).await + } - let rows = Rows { schema, rows }; + /// Routes already prepared rows while retaining TTL and WAL behavior. + pub async fn partition( + &self, + rows: Rows, + skip_wal: bool, + ) -> Result { let requests = Partitioner::new(self.partition_manager) - .partition_insert_requests(self.table_info, rows, request.skip_wal) + .partition_insert_requests(self.table_info, rows, skip_wal) .await?; let requests = RegionInsertRequests { requests }; @@ -72,16 +87,29 @@ mod tests { use api::v1::helper::tag_column_schema; use api::v1::region::InsertRequest as RegionInsertRequest; use api::v1::value::ValueData; - use api::v1::{ColumnDataType, PartitionExprVersion, Row, Value}; + use api::v1::{ColumnDataType, PartitionExprVersion, Row, Rows, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use datatypes::vectors::{Int32Vector, VectorRef}; use store_api::storage::RegionId; + use table::requests::InsertRequest as TableInsertRequest; - use super::*; + use crate::req_convert::insert::table_to_region::TableToRegion; use crate::test_util::{ create_partition_rule_manager, new_test_table_info, prepare_mocked_backend, }; + #[tokio::test] + async fn test_prepare_preserves_rows_before_routing() { + let backend = prepare_mocked_backend().await; + let partition_manager = create_partition_rule_manager(backend).await; + let table_info = new_test_table_info(1, "table_1", vec![0u32, 1, 2].into_iter()); + let converter = TableToRegion::new(&table_info, &partition_manager); + let values = vec![Some(1), None, Some(11), Some(101)]; + let request = build_table_request(Arc::new(Int32Vector::from(values.clone()))); + let expected = build_region_request(values, 0, None, false).rows.unwrap(); + assert_eq!(converter.prepare(request).unwrap(), expected); + } + #[tokio::test] async fn test_insert_request_table_to_region() { check_insert_request_table_to_region(false).await; diff --git a/src/servers/src/batcher/table.rs b/src/servers/src/batcher/table.rs index 3b88a4ca827..1a46a7f10f3 100644 --- a/src/servers/src/batcher/table.rs +++ b/src/servers/src/batcher/table.rs @@ -225,6 +225,7 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::num::NonZeroUsize; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; @@ -238,19 +239,23 @@ mod tests { use arrow::array::{Int32Array, TimestampMillisecondArray}; use arrow::datatypes::Schema as ArrowSchema; use arrow::record_batch::RecordBatch; + use catalog::memory::MemoryCatalogManager; 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_catalog::consts::default_engine; use common_grpc::flight::FlightDecoder; use common_meta::error::Result as MetaResult; use common_meta::peer::Peer; use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager}; + use common_query::OutputData; use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; use common_telemetry::info; use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder}; use datatypes::value::Value as DtValue; + use datatypes::vectors::{Int32Vector, TimestampMillisecondVector, VectorRef}; use operator::batcher::PendingRowsBatcher; use operator::error::Error; use operator::insert::Inserter; @@ -261,7 +266,9 @@ mod tests { }; use session::context::{Channel, QueryContext}; use store_api::storage::RegionId; + use table::dist_table::DistTable; use table::metadata::TableInfoRef; + use table::requests::InsertRequest; use tokio::sync::{broadcast, mpsc, oneshot}; use tokio::time::timeout; @@ -364,14 +371,15 @@ mod tests { expected_schema: schema_name.clone(), })); let inserter = Arc::new(Inserter::new( - catalog::memory::MemoryCatalogManager::new(), - partitions, - nodes, + MemoryCatalogManager::new(), + partitions.clone(), + nodes.clone(), mock_table_flownode_cache(1, vec![]).await, true, )); let mut table = new_test_table_info(1, "table_1", [1, 2, 3].into_iter()); table.schema_name = schema_name; + table.meta.engine = default_engine().to_string(); let mut columns = table.meta.schema.column_schemas().to_vec(); columns[2] = columns[2] .clone() @@ -399,32 +407,67 @@ mod tests { inserter, ) .unwrap(); - let influx_ctx = Arc::new(QueryContext::with_channel( - "greptime", - &table.schema_name, - Channel::Influx, - )); - let opentsdb_ctx = Arc::new(QueryContext::with_channel( - "greptime", - &table.schema_name, - Channel::Opentsdb, - )); + let context = |channel| { + let mut ctx = QueryContext::with_channel("greptime", &table.schema_name, channel); + ctx.set_batching_enabled(true); + Arc::new(ctx) + }; + let influx_ctx = context(Channel::Influx); + let opentsdb_ctx = context(Channel::Opentsdb); let ingest_count = DIST_INGEST_ROW_COUNT.with_label_values(&[influx_ctx.get_db_string().as_str()]); assert_eq!(0, ingest_count.get()); influx_ctx.set_skip_wal(skip_wal); opentsdb_ctx.set_skip_wal(skip_wal); - let first_permit = batcher.acquire().await.unwrap(); - let second_permit = if shared_request { - first_permit.clone() - } else { - batcher.acquire().await.unwrap() - }; let (first_result, second_result) = timeout(Duration::from_secs(5), async { - tokio::join!( - batcher.submit(table.clone(), first, influx_ctx, first_permit), - batcher.submit(table.clone(), second, opentsdb_ctx, second_permit), - ) + if shared_request { + let permit = batcher.acquire().await.unwrap(); + tokio::join!( + batcher.submit(table.clone(), first, influx_ctx, permit.clone()), + batcher.submit(table.clone(), second, opentsdb_ctx, permit), + ) + } else { + let inserter = Inserter::new( + MemoryCatalogManager::new_with_table(DistTable::table(table.clone())), + partitions, + nodes, + mock_table_flownode_cache(1, vec![]).await, + true, + ) + .with_pending_rows_batcher(Some(batcher)); + let request = |value: i32| InsertRequest { + catalog_name: table.catalog_name.clone(), + schema_name: table.schema_name.clone(), + table_name: table.name.clone(), + columns_values: HashMap::from([ + ( + "a".to_string(), + Arc::new(Int32Vector::from_slice([value])) as VectorRef, + ), + ( + "ts".to_string(), + Arc::new(TimestampMillisecondVector::from_slice([ + 1000 + i64::from(value) + ])) as VectorRef, + ), + ]), + skip_wal, + }; + let (first, second) = tokio::join!( + inserter.handle_table_insert(request(1), influx_ctx), + inserter.handle_table_insert(request(2), opentsdb_ctx), + ); + ( + first.map(|output| match output.data { + OutputData::AffectedRows(rows) => rows, + _ => panic!("expected affected rows"), + }), + second.map(|output| match output.data { + OutputData::AffectedRows(rows) => rows, + _ => panic!("expected affected rows"), + }), + ) + } }) .await .expect("the row threshold did not dispatch the combined bulk insert"); @@ -474,7 +517,7 @@ mod tests { } #[tokio::test] - async fn test_submit_combines_then_routes_bulk_with_defaults() { + async fn test_inserter_counts_batched_rows_once() { for report_missing_row in [false, true] { assert_eq!(1, run_bulk_case(2, report_missing_row, false).await); } @@ -583,7 +626,7 @@ mod tests { FlowNotifier::new(cache.clone(), nodes.clone(), NonZeroUsize::new(16).unwrap()) .unwrap(); let inserter = Arc::new(Inserter::new( - catalog::memory::MemoryCatalogManager::new(), + MemoryCatalogManager::new(), partitions, nodes, cache, diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 40b7bb515b8..27411a5f8a6 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -21,13 +21,13 @@ use std::time::Duration; use async_trait::async_trait; use auth::UserProviderRef; -use axum::extract::{DefaultBodyLimit, Request}; +use axum::extract::{DefaultBodyLimit, Request, State}; use axum::http::StatusCode as HttpStatusCode; use axum::middleware::Next; use axum::response::{IntoResponse, Response}; use axum::routing::Route; use axum::serve::ListenerExt; -use axum::{Router, middleware, routing}; +use axum::{Extension, Router, middleware, routing}; use common_base::readable_size::ReadableSize; use common_recordbatch::RecordBatch; use common_telemetry::{error, info}; @@ -40,6 +40,7 @@ use futures::FutureExt; use http::{HeaderValue, Method}; use serde::{Deserialize, Serialize}; use serde_json::Value; +use session::context::QueryContext; use snafu::{ResultExt, ensure}; use tokio::sync::Mutex; use tokio::sync::oneshot::{self, Sender}; @@ -188,6 +189,22 @@ pub(crate) enum HttpServerKind { Api, } +/// HTTP write protocols eligible for the shared pending-row batcher. +/// Prometheus uses this selector only when metric-engine storage is disabled. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BatchingProtocol { + Prom, + Influxdb, + Opentsdb, + Otlp, + Logs, + Loki, + Splunk, + Elasticsearch, + HttpSql, +} + #[derive(Default)] pub struct HttpServer { router: StdMutex, @@ -197,6 +214,7 @@ pub struct HttpServer { // server configs options: HttpOptions, + batching_protocols: Vec, bind_addr: Option, /// What this server instance exposes. See [`HttpServerKind`]. kind: HttpServerKind, @@ -217,6 +235,22 @@ pub fn is_api_listener_path(path: &str) -> bool { is_namespace(path, HTTP_API_PREFIX_WITHOUT_TRAILING_SLASH) || is_namespace(path, "/dashboard") } +/// Sets a local-only write selector after authentication creates the context. +async fn set_http_write_batching( + State(protocol): State, + mut req: Request, + next: Next, +) -> Response { + let enabled = req + .extensions() + .get::>>() + .is_some_and(|protocols| protocols.contains(&protocol)); + if let Some(ctx) = req.extensions_mut().get_mut::() { + ctx.set_batching_enabled(enabled); + } + next.run(req).await +} + /// Outer guard for the API listener. Rejects any path outside the API surface /// with `404 Not Found` before it can reach a handler (and before auth side /// effects). Reachability is kept separate from authentication. @@ -242,6 +276,8 @@ impl HttpServer { pub struct HttpOptions { pub addr: String, + /// Request timeout; zero disables it. Frontend raises a nonzero timeout to at least + /// the largest active synchronous Prom or shared table batch flush interval plus one second. #[serde(with = "humantime_serde")] pub timeout: Duration, @@ -613,6 +649,7 @@ pub struct DashboardState { pub struct HttpServerBuilder { options: HttpOptions, + batching_protocols: Vec, user_provider: Option, router: Router, memory_limiter: ServerMemoryLimiter, @@ -622,12 +659,19 @@ impl HttpServerBuilder { pub fn new(options: HttpOptions) -> Self { Self { options, + batching_protocols: Vec::new(), user_provider: None, router: Router::new(), memory_limiter: ServerMemoryLimiter::default(), } } + /// Selects HTTP write protocols allowed to use the shared batcher. + pub fn with_batching_protocols(mut self, protocols: Vec) -> Self { + self.batching_protocols = protocols; + self + } + /// Set a global memory limiter for all server protocols. pub fn with_memory_limiter(mut self, limiter: ServerMemoryLimiter) -> Self { self.memory_limiter = limiter; @@ -849,6 +893,7 @@ impl HttpServerBuilder { pub fn build(self) -> HttpServer { HttpServer { options: self.options, + batching_protocols: self.batching_protocols.clone(), user_provider: self.user_provider, shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router), @@ -882,6 +927,7 @@ impl HttpServerBuilder { let internal = HttpServer { options: self.options, + batching_protocols: self.batching_protocols.clone(), user_provider: self.user_provider.clone(), shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router.clone()), @@ -899,6 +945,7 @@ impl HttpServerBuilder { }; Some(HttpServer { options: api_options, + batching_protocols: self.batching_protocols, user_provider: self.user_provider.clone(), shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router), @@ -1056,6 +1103,7 @@ impl HttpServer { AuthState::new(self.user_provider.clone()), authorize::check_http_auth, )) + .layer(Extension(Arc::new(self.batching_protocols.clone()))) .layer(middleware::from_fn(hints::extract_hints)) .layer(middleware::from_fn(client_ip::log_error_with_client_ip)) .layer(middleware::from_fn( @@ -1142,6 +1190,10 @@ impl HttpServer { ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), ) + .layer(middleware::from_fn_with_state( + BatchingProtocol::Loki, + set_http_write_batching, + )) .with_state(log_state) } @@ -1177,6 +1229,10 @@ impl HttpServer { ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), ) + .layer(middleware::from_fn_with_state( + BatchingProtocol::Splunk, + set_http_write_batching, + )) .with_state(log_state) } @@ -1253,6 +1309,10 @@ impl HttpServer { )), ) .layer(ServiceBuilder::new().layer(RequestDecompressionLayer::new())) + .layer(middleware::from_fn_with_state( + BatchingProtocol::Elasticsearch, + set_http_write_batching, + )) .with_state(log_state) } @@ -1277,6 +1337,10 @@ impl HttpServer { ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), ) + .layer(middleware::from_fn_with_state( + BatchingProtocol::Logs, + set_http_write_batching, + )) .with_state(log_state) } @@ -1304,12 +1368,24 @@ impl HttpServer { ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), ) + .layer(middleware::from_fn_with_state( + BatchingProtocol::Logs, + set_http_write_batching, + )) .with_state(log_state) } fn route_sql(api_state: ApiState) -> Router { Router::new() - .route("/sql", routing::get(handler::sql).post(handler::sql)) + .route( + "/sql", + routing::get(handler::sql).post(handler::sql).layer( + middleware::from_fn_with_state( + BatchingProtocol::HttpSql, + set_http_write_batching, + ), + ), + ) .route( "/sql/parse", routing::get(handler::sql_parse).post(handler::sql_parse), @@ -1365,9 +1441,18 @@ impl HttpServer { /// [read]: https://prometheus.io/docs/prometheus/latest/querying/remote_read_api/ /// [write]: https://prometheus.io/docs/concepts/remote_write_spec/ fn route_prom(state: PromStoreState) -> Router { + let write = routing::post(prom_store::remote_write); + let write = if state.prom_store_with_metric_engine { + write + } else { + write.layer(middleware::from_fn_with_state( + BatchingProtocol::Prom, + set_http_write_batching, + )) + }; Router::new() .route("/read", routing::post(prom_store::remote_read)) - .route("/write", routing::post(prom_store::remote_write)) + .route("/write", write) .with_state(state) } @@ -1381,12 +1466,20 @@ impl HttpServer { ) .route("/ping", routing::get(influxdb_ping)) .route("/health", routing::get(influxdb_health)) + .layer(middleware::from_fn_with_state( + BatchingProtocol::Influxdb, + set_http_write_batching, + )) .with_state(influxdb_handler) } fn route_opentsdb(opentsdb_handler: OpentsdbProtocolHandlerRef) -> Router { Router::new() .route("/api/put", routing::post(opentsdb::put)) + .layer(middleware::from_fn_with_state( + BatchingProtocol::Opentsdb, + set_http_write_batching, + )) .with_state(opentsdb_handler) } @@ -1403,6 +1496,10 @@ impl HttpServer { ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), ) + .layer(middleware::from_fn_with_state( + BatchingProtocol::Otlp, + set_http_write_batching, + )) .with_state(OtlpState { with_metric_engine, experimental_enable_exponential_histogram, @@ -1575,8 +1672,8 @@ mod test { use tokio::sync::mpsc; use tokio::time::Instant; - use super::*; use crate::http::test_helpers::TestClient; + use crate::http::*; use crate::prom_remote_write::validation::validate_label_name; use crate::query_handler::sql::SqlQueryHandler; @@ -2311,3 +2408,120 @@ mod test { assert!(!validate_label_name(&[0xff, 0xfe])); } } + +#[cfg(test)] +mod batching_tests { + use std::sync::{Arc, Mutex}; + + use axum::http::StatusCode; + use common_query::Output; + use session::context::QueryContextRef; + + use crate::error::Result as ServerResult; + use crate::http::test_helpers::TestClient; + use crate::http::{BatchingProtocol, HttpOptions, HttpServerBuilder}; + use crate::influxdb::InfluxdbRequest; + use crate::opentsdb::codec::DataPoint; + use crate::query_handler::{InfluxdbLineProtocolHandler, OpentsdbProtocolHandler}; + + #[test] + fn test_protocol_names_reject_unknown_values() { + assert_eq!( + serde_json::from_str::("\"prom\"").unwrap(), + BatchingProtocol::Prom + ); + for name in ["sql", "unknown"] { + assert!(serde_json::from_str::(&format!("\"{name}\"")).is_err()); + } + assert_eq!( + serde_json::from_str::("\"http_sql\"").unwrap(), + BatchingProtocol::HttpSql + ); + } + + #[derive(Default)] + struct RecordingWriteHandler { + selections: Mutex>, + } + + #[async_trait::async_trait] + impl OpentsdbProtocolHandler for RecordingWriteHandler { + async fn preflight(&self, _: &[DataPoint], _: QueryContextRef) -> ServerResult<()> { + Ok(()) + } + + async fn exec(&self, points: Vec, ctx: QueryContextRef) -> ServerResult { + self.selections.lock().unwrap().push(ctx.batching_enabled()); + Ok(points.len()) + } + } + + #[async_trait::async_trait] + impl InfluxdbLineProtocolHandler for RecordingWriteHandler { + async fn exec(&self, _: InfluxdbRequest, ctx: QueryContextRef) -> ServerResult { + self.selections.lock().unwrap().push(ctx.batching_enabled()); + Ok(Output::new_with_affected_rows(1)) + } + } + + #[tokio::test] + async fn test_real_influx_and_opentsdb_routes_keep_selection() { + for protocols in [ + vec![], + vec![BatchingProtocol::Influxdb], + vec![BatchingProtocol::Opentsdb], + vec![BatchingProtocol::Influxdb, BatchingProtocol::Opentsdb], + ] { + let handler = Arc::new(RecordingWriteHandler::default()); + let influx_enabled = protocols.contains(&BatchingProtocol::Influxdb); + let opentsdb_enabled = protocols.contains(&BatchingProtocol::Opentsdb); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_batching_protocols(protocols) + .with_influxdb_handler(handler.clone()) + .with_opentsdb_handler(handler.clone()) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + for path in [ + "/v1/influxdb/write", + "/v1/influxdb/api/v2/write?bucket=public", + ] { + assert_eq!( + client + .post(path) + .body("cpu value=1 42") + .send() + .await + .status(), + StatusCode::NO_CONTENT + ); + } + let point = + serde_json::json!({"metric":"cpu", "timestamp":42, "value":1, "tags":{"host":"a"}}); + for query in ["", "?summary", "?details"] { + let response = client + .post(&format!("/v1/opentsdb/api/put{query}")) + .json(&point) + .send() + .await; + assert_eq!( + response.status(), + if query.is_empty() { + StatusCode::NO_CONTENT + } else { + StatusCode::OK + } + ); + } + assert_eq!( + *handler.selections.lock().unwrap(), + vec![ + influx_enabled, + influx_enabled, + opentsdb_enabled, + opentsdb_enabled, + opentsdb_enabled + ] + ); + } + } +} diff --git a/src/servers/src/http/opentsdb.rs b/src/servers/src/http/opentsdb.rs index 565fa82ae8e..51667db2fe7 100644 --- a/src/servers/src/http/opentsdb.rs +++ b/src/servers/src/http/opentsdb.rs @@ -97,7 +97,7 @@ pub async fn put( } let response = if !summary && !details { - if let Err(e) = opentsdb_handler.exec(data_points, ctx.clone()).await { + if let Err(e) = opentsdb_handler.exec_batch(data_points, ctx.clone()).await { // Not debugging purpose, failed fast. return error::InternalSnafu { err_msg: e.to_string(), @@ -172,7 +172,87 @@ impl OpentsdbDebuggingResponse { #[cfg(test)] mod test { - use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use session::context::QueryContextRef; + + use crate::http::opentsdb::*; + use crate::query_handler::OpentsdbProtocolHandler; + + #[derive(Default)] + struct CountingHandler { + direct: AtomicUsize, + batched: AtomicUsize, + preflight: AtomicUsize, + } + + #[async_trait::async_trait] + impl OpentsdbProtocolHandler for CountingHandler { + async fn preflight(&self, _: &[DataPoint], _: QueryContextRef) -> Result<()> { + self.preflight.fetch_add(1, Ordering::Relaxed); + Ok(()) + } + + async fn exec(&self, points: Vec, _: QueryContextRef) -> Result { + self.direct.fetch_add(1, Ordering::Relaxed); + Ok(points.len()) + } + + async fn exec_batch(&self, points: Vec, _: QueryContextRef) -> Result { + self.batched.fetch_add(1, Ordering::Relaxed); + Ok(points.len()) + } + } + + #[tokio::test] + async fn test_put_batches_only_non_debug_requests() { + for params in [ + HashMap::new(), + HashMap::from([("summary".to_string(), String::new())]), + HashMap::from([("details".to_string(), String::new())]), + HashMap::from([ + ("summary".to_string(), String::new()), + ("details".to_string(), String::new()), + ]), + ] { + let debug = !params.is_empty(); + let handler = Arc::new(CountingHandler::default()); + let body = Bytes::from_static( + br#"[ + {"metric":"cpu","timestamp":1000,"value":1,"tags":{"host":"a"}}, + {"metric":"cpu","timestamp":1001,"value":2,"tags":{"host":"b"}} + ]"#, + ); + let (status, Json(response)) = put( + State(handler.clone() as OpentsdbProtocolHandlerRef), + Query(params), + Extension(QueryContext::with("greptime", "public")), + body, + ) + .await + .unwrap(); + assert_eq!( + usize::from(debug) * 2, + handler.direct.load(Ordering::Relaxed) + ); + assert_eq!(usize::from(!debug), handler.batched.load(Ordering::Relaxed)); + assert_eq!( + usize::from(debug), + handler.preflight.load(Ordering::Relaxed) + ); + if debug { + assert_eq!(HttpStatusCode::OK, status); + let OpentsdbPutResponse::Debug(response) = response else { + panic!("expected debug response") + }; + assert_eq!(2, response.success); + assert_eq!(0, response.failed); + } else { + assert_eq!(HttpStatusCode::NO_CONTENT, status); + assert!(matches!(response, OpentsdbPutResponse::Empty)); + } + } + } #[test] fn test_into_opentsdb_data_point() { diff --git a/src/servers/src/query_handler.rs b/src/servers/src/query_handler.rs index 35ec1cc2ff5..e74088a2a5d 100644 --- a/src/servers/src/query_handler.rs +++ b/src/servers/src/query_handler.rs @@ -94,6 +94,13 @@ pub trait OpentsdbProtocolHandler { /// A successful request will not return a response. /// Only on error will the socket return a line of data. async fn exec(&self, data_points: Vec, ctx: QueryContextRef) -> Result; + + /// Executes an ordinary HTTP put with optional batching. Debug and socket + /// callers retain [`Self::exec`]; the frontend clears HTTP batching selection + /// there so diagnostic requests preserve direct, per-point error attribution. + async fn exec_batch(&self, data_points: Vec, ctx: QueryContextRef) -> Result { + self.exec(data_points, ctx).await + } } pub struct PromStoreResponse { diff --git a/src/servers/tests/http/prom_store_test.rs b/src/servers/tests/http/prom_store_test.rs index 8f3e2b0e770..7c07e3b0344 100644 --- a/src/servers/tests/http/prom_store_test.rs +++ b/src/servers/tests/http/prom_store_test.rs @@ -40,7 +40,7 @@ use servers::error::{self, Result}; use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF}; use servers::http::prom_store::PHYSICAL_TABLE_PARAM; use servers::http::test_helpers::{TestClient, TestResponse}; -use servers::http::{HttpOptions, HttpServerBuilder}; +use servers::http::{BatchingProtocol, HttpOptions, HttpServerBuilder}; use servers::prom_remote_write::v2::test_util as remote_write_v2; use servers::prom_remote_write::validation::PromValidationMode; use servers::prom_store; @@ -62,6 +62,7 @@ struct DummyInstance { } struct RemoteWriteCapture { + batching_enabled: bool, schema: String, physical_table: Option, with_metric_engine: bool, @@ -80,6 +81,12 @@ impl PromStoreProtocolHandler for DummyInstance { ctx: QueryContextRef, with_metric_engine: bool, ) -> Result { + if with_metric_engine { + assert!( + !ctx.batching_enabled(), + "metric-engine Prom must retain its dedicated batcher" + ); + } let write_call = self.write_calls.fetch_add(1, Ordering::SeqCst) + 1; if self.fail_write_call == Some(write_call) { return error::InvalidPromRemoteRequestSnafu { @@ -91,6 +98,7 @@ impl PromStoreProtocolHandler for DummyInstance { let _ = self .write_tx .send(RemoteWriteCapture { + batching_enabled: ctx.batching_enabled(), schema: ctx.current_schema(), physical_table: ctx.extension(PHYSICAL_TABLE_PARAM).map(ToString::to_string), with_metric_engine, @@ -150,8 +158,12 @@ impl PromStoreProtocolHandler for DummyInstance { #[async_trait] impl SqlQueryHandler for DummyInstance { - async fn do_query(&self, _: &str, _: QueryContextRef) -> Vec> { - unimplemented!() + async fn do_query(&self, _: &str, ctx: QueryContextRef) -> Vec> { + self.read_tx + .send((ctx.current_schema(), vec![u8::from(ctx.batching_enabled())])) + .await + .unwrap(); + vec![Ok(Output::new_with_affected_rows(1))] } async fn do_analyze_stream_query(&self, _: &str, _: QueryContextRef) -> Result { @@ -237,6 +249,17 @@ fn make_test_app_with_write_failure_inner( fail_write_call, }); let server = HttpServerBuilder::new(http_opts) + .with_batching_protocols(vec![ + BatchingProtocol::Prom, + BatchingProtocol::Influxdb, + BatchingProtocol::Opentsdb, + BatchingProtocol::Otlp, + BatchingProtocol::Logs, + BatchingProtocol::Loki, + BatchingProtocol::Splunk, + BatchingProtocol::Elasticsearch, + BatchingProtocol::HttpSql, + ]) .with_sql_handler(instance.clone()) .with_prom_handler( instance, @@ -795,3 +818,78 @@ fn assert_remote_write_v2_written_headers_with_histograms( .map(|x| x.to_str().unwrap()) ); } + +#[tokio::test] +async fn test_http_sql_protocol_selection() { + for protocols in [ + vec![], + vec![BatchingProtocol::Influxdb], + vec![BatchingProtocol::HttpSql], + ] { + let expected = protocols.contains(&BatchingProtocol::HttpSql); + let (read_tx, mut read_rx) = mpsc::channel(1); + let (write_tx, _write_rx) = mpsc::channel(1); + let instance = Arc::new(DummyInstance { + read_tx, + write_tx, + write_calls: Arc::new(AtomicUsize::new(0)), + fail_write_call: None, + }); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_batching_protocols(protocols) + .with_sql_handler(instance) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + let response = client.get("/v1/sql?sql=SELECT%201").send().await; + assert!(response.status().is_success()); + assert_eq!( + read_rx.recv().await.unwrap(), + ("public".to_string(), vec![u8::from(expected)]) + ); + } +} + +#[tokio::test] +async fn test_prom_batching_depends_on_protocol_and_metric_engine() { + for with_metric_engine in [false, true] { + for enabled in [false, true] { + let (read_tx, _read_rx) = mpsc::channel(1); + let (write_tx, mut write_rx) = mpsc::channel(16); + let instance = Arc::new(DummyInstance { + read_tx, + write_tx, + write_calls: Arc::new(AtomicUsize::new(0)), + fail_write_call: None, + }); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_batching_protocols(if enabled { + vec![BatchingProtocol::Prom] + } else { + vec![] + }) + .with_prom_handler( + instance, + None, + with_metric_engine, + PromValidationMode::Unchecked, + false, + None, + ) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + let request = WriteRequest { + timeseries: prom_store::mock_timeseries(), + ..Default::default() + }; + let response = client + .post("/v1/prometheus/write") + .body(snappy_compress(&request.encode_to_vec()).unwrap()) + .send() + .await; + assert_eq!(response.status(), 204); + let capture = write_rx.recv().await.unwrap(); + assert_eq!(capture.batching_enabled, enabled && !with_metric_engine); + assert_eq!(capture.with_metric_engine, with_metric_engine); + } + } +} diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 2e8b49dae37..2833abd2578 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -79,6 +79,9 @@ pub struct QueryContext { /// The configuration parameter are used to store the parameters that are set by the user #[builder(default)] configuration_parameter: Arc, + /// Local-only write batching selection; never transported in protobuf extensions. + #[builder(default)] + batching_enabled: bool, /// Track which protocol the query comes from. #[builder(default)] channel: Channel, @@ -430,6 +433,16 @@ impl QueryContext { &self.configuration_parameter } + /// Whether the local HTTP entry point selected write batching. + pub fn batching_enabled(&self) -> bool { + self.batching_enabled + } + + /// Sets local write batching selection without adding a wire-visible extension. + pub fn set_batching_enabled(&mut self, enabled: bool) { + self.batching_enabled = enabled; + } + pub fn channel(&self) -> Channel { self.channel } @@ -594,6 +607,7 @@ impl QueryContextBuilder { .configuration_parameter .unwrap_or_else(|| Arc::new(ConfigurationVariables::default())), channel, + batching_enabled: self.batching_enabled.unwrap_or_default(), process_id: self.process_id.unwrap_or_default(), conn_info: self.conn_info.unwrap_or_default(), protocol_ctx: self.protocol_ctx.unwrap_or_default(), @@ -709,9 +723,8 @@ mod test { use common_catalog::consts::DEFAULT_CATALOG_NAME; - use super::*; use crate::Session; - use crate::context::Channel; + use crate::context::{Channel, *}; #[test] fn test_session() { @@ -860,4 +873,18 @@ mod test { ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "another-query-id"); assert!(!ctx.live_analyze_metrics_enabled()); } + #[test] + fn test_batching_selection_is_local_only() { + let mut ctx = QueryContextBuilder::default().build(); + assert!(!ctx.batching_enabled()); + ctx.set_batching_enabled(true); + assert!(ctx.clone().batching_enabled()); + assert!(ctx.fork().batching_enabled()); + let wire: api::v1::QueryContext = ctx.into(); + assert!(!QueryContext::from(wire).batching_enabled()); + let ctx = QueryContextBuilder::default() + .set_extension("batching_enabled".to_string(), "true".to_string()) + .build(); + assert!(!ctx.batching_enabled()); + } } diff --git a/src/standalone/Cargo.toml b/src/standalone/Cargo.toml index 3dcb1d5f108..a36983203ee 100644 --- a/src/standalone/Cargo.toml +++ b/src/standalone/Cargo.toml @@ -46,4 +46,5 @@ tokio.workspace = true url.workspace = true [dev-dependencies] +temp-env.workspace = true toml.workspace = true diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 9f20d73952a..4884b5ee091 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -24,8 +24,8 @@ use file_engine::config::EngineConfig as FileEngineConfig; use flow::FlowConfig; use frontend::frontend::FrontendOptions; use frontend::service_config::{ - InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions, PostgresOptions, - PromStoreOptions, + InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions, + PendingRowsBatcherOptions, PostgresOptions, PromStoreOptions, }; use mito2::config::MitoConfig; use pipeline::PipelineOptions; @@ -56,6 +56,8 @@ pub struct StandaloneOptions { pub postgres: PostgresOptions, pub opentsdb: OpentsdbOptions, pub influxdb: InfluxdbOptions, + /// Shared experimental ordinary-table batching; independent of Prom batching. + pub pending_rows_batcher: PendingRowsBatcherOptions, pub jaeger: JaegerOptions, pub otlp: OtlpOptions, pub prom_store: PromStoreOptions, @@ -97,6 +99,7 @@ impl Default for StandaloneOptions { postgres: PostgresOptions::default(), opentsdb: OpentsdbOptions::default(), influxdb: InfluxdbOptions::default(), + pending_rows_batcher: PendingRowsBatcherOptions::default(), jaeger: JaegerOptions::default(), otlp: OtlpOptions::default(), prom_store: PromStoreOptions::default(), @@ -130,6 +133,7 @@ impl Configurable for StandaloneOptions { "heartbeat_env_vars", "wal.broker_endpoints", "event_recorder.event_types", + "pending_rows_batcher.protocols", ]) } } @@ -158,6 +162,7 @@ impl StandaloneOptions { postgres: cloned_opts.postgres, opentsdb: cloned_opts.opentsdb, influxdb: cloned_opts.influxdb, + pending_rows_batcher: cloned_opts.pending_rows_batcher, jaeger: cloned_opts.jaeger, otlp: cloned_opts.otlp, prom_store: cloned_opts.prom_store, @@ -208,7 +213,68 @@ mod tests { use common_event_recorder::EventTypeFilter; - use super::*; + use crate::options::*; + + #[test] + fn test_batcher_protocols_from_env() { + temp_env::with_vars( + [( + "STANDALONE_BATCHER_TEST__PENDING_ROWS_BATCHER__PROTOCOLS", + Some("influxdb,http_sql"), + )], + || { + let options = + StandaloneOptions::load_layered_options(None, "STANDALONE_BATCHER_TEST") + .unwrap(); + assert_eq!( + options.pending_rows_batcher.protocols, + vec![ + servers::http::BatchingProtocol::Influxdb, + servers::http::BatchingProtocol::HttpSql + ] + ); + }, + ); + } + + #[test] + fn test_protocol_pending_rows_batcher_config() { + let defaults: StandaloneOptions = toml::from_str("").unwrap(); + assert!( + !defaults + .pending_rows_batcher + .pending_rows_batching_enabled() + ); + let options: StandaloneOptions = toml::from_str( + r#" +[pending_rows_batcher] +protocols = ["influxdb", "http_sql"] +pending_rows_flush_interval = "5ms" +max_batch_rows = 25 +flow_notification_queue_capacity = 17 +"#, + ) + .unwrap(); + assert_eq!(options.pending_rows_batcher.max_batch_rows, 25); + assert_eq!(options.pending_rows_batcher.protocols.len(), 2); + assert_eq!( + options + .pending_rows_batcher + .flow_notification_queue_capacity + .get(), + 17 + ); + assert!(options.pending_rows_batcher.pending_rows_batching_enabled()); + let serialized = toml::to_string(&options).unwrap(); + let parsed: StandaloneOptions = toml::from_str(&serialized).unwrap(); + assert_eq!(options.influxdb, parsed.influxdb); + assert_eq!(options.opentsdb, parsed.opentsdb); + assert_eq!(options.pending_rows_batcher, parsed.pending_rows_batcher); + let frontend = options.frontend_options(); + assert_eq!(options.influxdb, frontend.influxdb); + assert_eq!(options.opentsdb, frontend.opentsdb); + assert_eq!(options.pending_rows_batcher, frontend.pending_rows_batcher); + } #[test] fn test_event_recorder_event_types_preserve_filter_semantics() { diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index d6cb0cf13ec..f8a4de57654 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2306,6 +2306,15 @@ enable = true enable = true default_merge_mode = "last_non_null" +[pending_rows_batcher] +protocols = [] +pending_rows_flush_interval = "0s" +max_batch_rows = 100000 +max_concurrent_flushes = 256 +worker_channel_capacity = 65526 +max_inflight_requests = 3000 +flow_notification_queue_capacity = 1024 + [jaeger] enable = true