rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-25 06:35:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: batch ordinary table writes across HTTP protocols (#9115)

Browse Source 
* feat: integrate table batching across HTTP protocols

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix: skip empty prepared writes before batch admission

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix: load batching protocols from environment and document frontend wiring

Signed-off-by: WenyXu <wenymedia@gmail.com>

* refactor: remove experimental prefix from pending rows batcher config

Signed-off-by: WenyXu <wenymedia@gmail.com>

* style: sort frontend test dependencies

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix: count batched ingestion once and update config snapshot

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-09-17 09:32:12 +00:00
committed by GitHub
parent d45f5d6eaf
commit be15c88e92
29 changed files with 1455 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+2
View File
@@ -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",
Cargo.toml
+1
View File
@@ -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"
config/config.md
+16 -2
View File
@@ -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.<br/>When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the<br/>`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.<br/>When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is<br/>raised to at least the largest active flush interval plus 1 second. The intervals come from<br/>`prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. |
| `http.body_limit` | String | `64MB` | HTTP request body limit.<br/>The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.<br/>Set to 0 to disable limit. |
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default<br/>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.<br/>Available values: "last_non_null", "last_row". |
| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.<br/>Legacy Prometheus batching settings under prom_store remain supported.<br/>HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.<br/>Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.<br/>Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.<br/>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.<br/>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.<br/>When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the<br/>`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.<br/>When synchronous Prometheus or shared table batching is enabled, a nonzero timeout is<br/>raised to at least the largest active flush interval plus 1 second. The intervals come from<br/>`prom_store.pending_rows_flush_interval` and `pending_rows_batcher.pending_rows_flush_interval`. |
| `http.body_limit` | String | `64MB` | HTTP request body limit.<br/>The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.<br/>Set to 0 to disable limit. |
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default<br/>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.<br/>Available values: "last_non_null", "last_row". |
| `pending_rows_batcher` | -- | -- | Shared experimental ordinary-table batching for opted-in ingestion protocols.<br/>Legacy Prometheus batching settings under prom_store remain supported.<br/>HTTP write protocols sharing this batcher. Omitted or empty disables all entrances.<br/>Supported: influxdb, opentsdb, otlp, logs, loki, splunk, elasticsearch, http_sql, prom.<br/>Prom uses ordinary-table batching without metric engine, otherwise its dedicated batcher.<br/>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. |
config/frontend.example.toml
+34 -2
View File
@@ -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.
config/standalone.example.toml
+34 -2
View File
@@ -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.
src/frontend/AGENTS.md
+6
View File
@@ -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.
src/frontend/Cargo.toml
+1
View File
@@ -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
src/frontend/src/frontend.rs
+57 -3
View File
@@ -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<T> =
Pin<Box<dyn Stream<Item = std::result::Result<T, Status>> + 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]
src/frontend/src/instance.rs
+1 -1
View File
@@ -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<Statement> {
parse_stmt(sql, &GreptimeDbDialect {}).unwrap()
src/frontend/src/instance/builder.rs
+48 -14
View File
@@ -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<Arc<dyn PendingRowsBatcher>> {
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<dyn PendingRowsBatcher>)
};
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(),
src/frontend/src/instance/influxdb.rs
+1 -1
View File
@@ -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]
src/frontend/src/instance/opentsdb.rs
+24 -9
View File
@@ -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<DataPoint>, ctx: QueryContextRef) -> ServerResult<usize> {
// 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<DataPoint>,
ctx: QueryContextRef,
) -> server_error::Result<usize> {
) -> ServerResult<usize> {
self.execute_opentsdb_write(data_points, ctx).await
}
}
impl Instance {
async fn execute_opentsdb_write(
&self,
data_points: Vec<DataPoint>,
ctx: QueryContextRef,
) -> ServerResult<usize> {
self.plugins
.get::<PermissionCheckerRef>()
.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() {
src/frontend/src/server.rs
+172 -20
View File
@@ -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<T>
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<SocketAddr> {
#[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`
src/frontend/src/service_config.rs
+2
View File
@@ -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;
src/frontend/src/service_config/influxdb.rs
+1 -1
View File
@@ -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() {
src/frontend/src/service_config/pending_rows_batcher.rs
+150
View File
@@ -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<BatchingProtocol>,
/// 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::<PendingRowsBatcherOptions>(&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::<PendingRowsBatcherOptions>("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::<PendingRowsBatcherOptions>(&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}");
}
}
}
src/operator/src/insert.rs
+248 -11
View File
@@ -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<Arc<dyn PendingRowsBatcher>>,
}
pub type InserterRef = Arc<Inserter>;
@@ -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<Arc<dyn PendingRowsBatcher>>,
) -> Self {
self.pending_rows_batcher = batcher;
self
}
pub async fn handle_column_inserts(
&self,
requests: InsertRequests,
@@ -267,6 +280,11 @@ impl Inserter {
) -> Result<Output> {
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<TableId, Arc<TableInfo>>,
ctx: QueryContextRef,
batcher: &Arc<dyn PendingRowsBatcher>,
) -> Result<Output> {
// 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::<HashMap<_, _>>();
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::<Result<usize>>()?;
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<dyn PendingRowsBatcher>> {
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<dyn PendingRowsBatcher>,
) -> Result<Output> {
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<Output> {
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<Arc<tokio::sync::OwnedSemaphorePermit>> {
panic!("empty writes must not acquire batch admission")
}
async fn submit(
&self,
_table_info: TableInfoRef,
_batch: arrow::record_batch::RecordBatch,
_ctx: QueryContextRef,
_permit: Arc<tokio::sync::OwnedSemaphorePermit>,
) -> Result<usize> {
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<dyn PendingRowsBatcher> = 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<dyn PendingRowsBatcher> = 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)));
}
}
src/operator/src/req_convert/insert.rs
+1
View File
@@ -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;
src/operator/src/req_convert/insert/stmt_to_region.rs
+36 -21
View File
@@ -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<InstantAndNormalInsertRequests> {
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<TableRef> {
src/operator/src/req_convert/insert/table_to_region.rs
+35 -7
View File
@@ -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<Rows> {
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<InstantAndNormalInsertRequests> {
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<InstantAndNormalInsertRequests> {
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;
src/servers/src/batcher/table.rs
+68 -25
View File
@@ -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,
src/servers/src/http.rs
+219 -5
View File
@@ -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<Router>,
@@ -197,6 +214,7 @@ pub struct HttpServer {
// server configs
options: HttpOptions,
batching_protocols: Vec<BatchingProtocol>,
bind_addr: Option<SocketAddr>,
/// 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<BatchingProtocol>,
mut req: Request,
next: Next,
) -> Response {
let enabled = req
.extensions()
.get::<Arc<Vec<BatchingProtocol>>>()
.is_some_and(|protocols| protocols.contains(&protocol));
if let Some(ctx) = req.extensions_mut().get_mut::<QueryContext>() {
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<BatchingProtocol>,
user_provider: Option<UserProviderRef>,
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<BatchingProtocol>) -> 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<S>(api_state: ApiState) -> Router<S> {
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<S>(state: PromStoreState) -> Router<S> {
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<S>(opentsdb_handler: OpentsdbProtocolHandlerRef) -> Router<S> {
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::<BatchingProtocol>("\"prom\"").unwrap(),
BatchingProtocol::Prom
);
for name in ["sql", "unknown"] {
assert!(serde_json::from_str::<BatchingProtocol>(&format!("\"{name}\"")).is_err());
}
assert_eq!(
serde_json::from_str::<BatchingProtocol>("\"http_sql\"").unwrap(),
BatchingProtocol::HttpSql
);
}
#[derive(Default)]
struct RecordingWriteHandler {
selections: Mutex<Vec<bool>>,
}
#[async_trait::async_trait]
impl OpentsdbProtocolHandler for RecordingWriteHandler {
async fn preflight(&self, _: &[DataPoint], _: QueryContextRef) -> ServerResult<()> {
Ok(())
}
async fn exec(&self, points: Vec<DataPoint>, ctx: QueryContextRef) -> ServerResult<usize> {
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<Output> {
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
]
);
}
}
}
src/servers/src/http/opentsdb.rs
+82 -2
View File
@@ -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<DataPoint>, _: QueryContextRef) -> Result<usize> {
self.direct.fetch_add(1, Ordering::Relaxed);
Ok(points.len())
}
async fn exec_batch(&self, points: Vec<DataPoint>, _: QueryContextRef) -> Result<usize> {
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() {
src/servers/src/query_handler.rs
+7
View File
@@ -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<DataPoint>, ctx: QueryContextRef) -> Result<usize>;
/// 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<DataPoint>, ctx: QueryContextRef) -> Result<usize> {
self.exec(data_points, ctx).await
}
}
pub struct PromStoreResponse {
src/servers/tests/http/prom_store_test.rs
+101 -3
View File
@@ -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<String>,
with_metric_engine: bool,
@@ -80,6 +81,12 @@ impl PromStoreProtocolHandler for DummyInstance {
ctx: QueryContextRef,
with_metric_engine: bool,
) -> Result<Output> {
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<Result<Output>> {
unimplemented!()
async fn do_query(&self, _: &str, ctx: QueryContextRef) -> Vec<Result<Output>> {
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<Output> {
@@ -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);
}
}
}
src/session/src/context.rs
+29 -2
View File
@@ -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<ConfigurationVariables>,
/// 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());
}
}
src/standalone/Cargo.toml
+1
View File
@@ -46,4 +46,5 @@ tokio.workspace = true
url.workspace = true
[dev-dependencies]
temp-env.workspace = true
toml.workspace = true
src/standalone/src/options.rs
+69 -3
View File
@@ -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() {
tests-integration/tests/http.rs
+9
View File
@@ -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