diff --git a/Cargo.lock b/Cargo.lock index 655b0424f3..ca86fb82cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7269,6 +7269,7 @@ dependencies = [ "derive_builder", "derive_more", "futures", + "gardal", "googletest", "http 1.4.0", "http-body 1.0.1", @@ -7658,7 +7659,7 @@ dependencies = [ "http-body 1.0.1", "http-body-util", "jiff", - "object_store 0.13.2", + "object_store", "restate-cli-util", "restate-core", "restate-limiter", @@ -8621,8 +8622,10 @@ dependencies = [ "derive_more", "enumset", "futures", + "gardal", "googletest", "itertools 0.14.0", + "metrics", "parking_lot", "paste", "prost", diff --git a/crates/admin/Cargo.toml b/crates/admin/Cargo.toml index 56c6df2d98..507e1c6cd3 100644 --- a/crates/admin/Cargo.toml +++ b/crates/admin/Cargo.toml @@ -57,6 +57,7 @@ http = { workspace = true } http-body = { workspace = true } http-body-util = { workspace = true } hyper-util = { workspace = true } +gardal = { workspace = true } itertools = { workspace = true } metrics = { workspace = true } mime_guess = { version = "2.0.5", optional = true } diff --git a/crates/admin/src/cluster_controller/grpc_svc_handler.rs b/crates/admin/src/cluster_controller/grpc_svc_handler.rs index 80c93779e5..9388d6d337 100644 --- a/crates/admin/src/cluster_controller/grpc_svc_handler.rs +++ b/crates/admin/src/cluster_controller/grpc_svc_handler.rs @@ -36,7 +36,7 @@ use restate_core::protobuf::cluster_ctrl_svc::{ }; use restate_core::{Metadata, MetadataWriter}; use restate_metadata_store::WriteError; -use restate_storage_query_datafusion::context::QueryContext; +use restate_storage_query_datafusion::context::{QueryContext, QueryError}; use restate_storage_query_datafusion::node_fan_out::NodeWarnings; use restate_types::config::{MetadataClientKind, MetadataClientOptions, NetworkingOptions}; use restate_types::identifiers::PartitionId; @@ -429,7 +429,7 @@ impl ClusterCtrlSvc for ClusterCtrlSvcHandler { .query_context .execute(&request.query) .await - .map_err(datafusion_error_to_status)?; + .map_err(datafusion_query_error_to_status)?; let node_warnings = query_result.node_warnings; @@ -651,6 +651,13 @@ fn drain_node_warnings(node_warnings: &[NodeWarnings]) -> Vec { out } +fn datafusion_query_error_to_status(err: QueryError) -> Status { + match err { + QueryError::DataFusion(e) => datafusion_error_to_status(e), + err @ QueryError::RateLimited(_) => Status::resource_exhausted(err.to_string()), + } +} + fn datafusion_error_to_status(err: DataFusionError) -> Status { match err { DataFusionError::SQL(..) diff --git a/crates/admin/src/metric_definitions.rs b/crates/admin/src/metric_definitions.rs index ceb377bf9b..ced079e6b1 100644 --- a/crates/admin/src/metric_definitions.rs +++ b/crates/admin/src/metric_definitions.rs @@ -10,12 +10,17 @@ use metrics::{Unit, describe_counter, describe_gauge, describe_histogram}; +// Storage Accounting metrics pub(crate) const USAGE_STATE_STORAGE_BYTES: &str = "restate.usage.state.storage_bytes"; pub(crate) const USAGE_STATE_STORAGE_BYTE_SECONDS: &str = "restate.usage.state.storage_byte_seconds.total"; pub(crate) const USAGE_STATE_SIZE_ACCOUNTING_QUERY_DURATION_SECONDS: &str = "restate.usage.state_size_accounting.query_duration_seconds"; +// Query Engine metrics +pub(crate) const QUERY_ENGINE_QUERY_DURATION_SECONDS: &str = + "restate.query_engine.query_duration_seconds"; + pub(crate) fn describe_metrics() { describe_gauge!( USAGE_STATE_STORAGE_BYTES, @@ -33,5 +38,11 @@ pub(crate) fn describe_metrics() { USAGE_STATE_SIZE_ACCOUNTING_QUERY_DURATION_SECONDS, Unit::Seconds, "Accounting query execution duration" + ); + + describe_histogram!( + QUERY_ENGINE_QUERY_DURATION_SECONDS, + Unit::Seconds, + "End-to-end duration of a query, including planning, execution and streaming the response" ) } diff --git a/crates/admin/src/rest_api/query.rs b/crates/admin/src/rest_api/query.rs index ad226ad299..6d889235c5 100644 --- a/crates/admin/src/rest_api/query.rs +++ b/crates/admin/src/rest_api/query.rs @@ -11,9 +11,9 @@ use std::io::Write; use std::pin::Pin; use std::sync::Arc; +use std::task::{Context, Poll}; +use std::time::Instant; -use crate::query_utils::{RecordBatchWriter, WriteRecordBatchStream}; -use crate::state::AdminServiceState; use axum::extract::State; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; @@ -24,16 +24,24 @@ use datafusion::arrow::datatypes::Schema; use datafusion::arrow::ipc::writer::StreamWriter; use datafusion::arrow::json::writer::JsonArray; use datafusion::common::DataFusionError; -use futures::{StreamExt, TryStreamExt}; +use futures::{Stream, StreamExt, TryStreamExt}; use http::{HeaderMap, HeaderValue}; use http_body::Frame; use http_body_util::StreamBody; +use metrics::histogram; use parking_lot::Mutex; +use serde::Serialize; + use restate_admin_rest_model::query::QueryRequest; use restate_core::network::TransportConnect; use restate_types::invocation::client::InvocationClient; use restate_types::schema::registry::{DiscoveryClient, MetadataService, TelemetryClient}; -use serde::Serialize; + +use crate::metric_definitions::QUERY_ENGINE_QUERY_DURATION_SECONDS; +use crate::query_utils::{RecordBatchWriter, WriteRecordBatchStream}; +use crate::state::AdminServiceState; + +const RETRY_AFTER_HEADER: &str = "Retry-After"; /// Error response for query endpoint. #[derive(Debug, Serialize, utoipa::ToSchema)] @@ -48,16 +56,47 @@ pub(crate) enum QueryError { Datafusion(#[from] datafusion::error::DataFusionError), #[error("Query service not available")] Unavailable, + #[error("Rate limited")] + RateLimited(#[from] gardal::RateLimited), +} + +impl From for QueryError { + fn from(err: restate_storage_query_datafusion::context::QueryError) -> Self { + match err { + restate_storage_query_datafusion::context::QueryError::DataFusion(e) => { + Self::Datafusion(e) + } + restate_storage_query_datafusion::context::QueryError::RateLimited(e) => { + Self::RateLimited(e) + } + } + } } impl IntoResponse for QueryError { fn into_response(self) -> Response { let status_code = match &self { + QueryError::Datafusion(datafusion::error::DataFusionError::Plan(_)) + | QueryError::Datafusion(datafusion::error::DataFusionError::SchemaError(_, _)) + | QueryError::Datafusion(datafusion::error::DataFusionError::SQL(_, _)) => { + StatusCode::BAD_REQUEST + } QueryError::Datafusion(_) => StatusCode::INTERNAL_SERVER_ERROR, QueryError::Unavailable => StatusCode::SERVICE_UNAVAILABLE, + QueryError::RateLimited(_) => StatusCode::TOO_MANY_REQUESTS, }; + + let mut headers = http::HeaderMap::new(); + if let QueryError::RateLimited(e) = &self { + headers.insert( + RETRY_AFTER_HEADER, + HeaderValue::from(e.earliest_retry_after().as_secs()), + ); + } + ( status_code, + headers, Json(QueryErrorBody { message: self.to_string(), }), @@ -94,6 +133,10 @@ where Invocations: InvocationClient + Send + Sync + Clone + 'static, Transport: TransportConnect, { + let duration_guard = QueryDurationGuard { + start: Instant::now(), + }; + let Some(query_context) = state.query_context.as_ref() else { return Err(QueryError::Unavailable); }; @@ -126,6 +169,13 @@ where return Err(err.into()); } + // Move the guard into the body so the duration is recorded once the stream is fully + // consumed or dropped (e.g. on client disconnect). + let result_stream = Timed { + inner: result_stream, + _guard: duration_guard, + }; + Ok(Response::builder() .header(http::header::CONTENT_TYPE, content_type) .body(StreamBody::new(result_stream)) @@ -133,6 +183,32 @@ where .into_response()) } +/// Records the elapsed time since `start` into the query-duration histogram when dropped. +struct QueryDurationGuard { + start: Instant, +} + +impl Drop for QueryDurationGuard { + fn drop(&mut self) { + histogram!(QUERY_ENGINE_QUERY_DURATION_SECONDS).record(self.start.elapsed().as_secs_f64()); + } +} + +/// Wraps a stream, keeping a [`QueryDurationGuard`] alive for as long as the stream is held. +/// The guard records the query duration when this wrapper (the response body) is dropped. +struct Timed { + inner: S, + _guard: QueryDurationGuard, +} + +impl Stream for Timed { + type Item = S::Item; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_next_unpin(cx) + } +} + #[derive(Clone)] // unfortunately the json writer doesnt give a way to get a mutable reference to the underlying writer, so we need another pointer in to its buffer // we use a lock here to help make the writer send/sync, despite it being totally uncontended :( diff --git a/crates/admin/src/service.rs b/crates/admin/src/service.rs index 930c41dcee..bb98b37601 100644 --- a/crates/admin/src/service.rs +++ b/crates/admin/src/service.rs @@ -37,6 +37,7 @@ use restate_types::net::listener::Listeners; use restate_types::schema::registry::SchemaRegistry; use restate_util_time::DurationExt; +use crate::metric_definitions::describe_metrics; use crate::rest_api::{MAX_ADMIN_API_VERSION, MIN_ADMIN_API_VERSION}; use crate::schema_registry_integration::{MetadataService, TelemetryClient}; use crate::{rest_api, state}; @@ -71,6 +72,7 @@ where service_discovery: ServiceDiscovery, telemetry_http_client: Option, ) -> Self { + describe_metrics(); let metadata_client = metadata_writer.raw_metadata_store_client().clone(); Self { listeners, diff --git a/crates/admin/src/storage_accounting.rs b/crates/admin/src/storage_accounting.rs index 7228d03708..af2ff38c4c 100644 --- a/crates/admin/src/storage_accounting.rs +++ b/crates/admin/src/storage_accounting.rs @@ -25,7 +25,7 @@ use restate_util_time::DurationExt; use crate::metric_definitions::{ USAGE_STATE_SIZE_ACCOUNTING_QUERY_DURATION_SECONDS, USAGE_STATE_STORAGE_BYTE_SECONDS, - USAGE_STATE_STORAGE_BYTES, describe_metrics, + USAGE_STATE_STORAGE_BYTES, }; pub(crate) const STORAGE_QUERY: &str = @@ -49,8 +49,6 @@ impl StorageAccountingTask { } pub async fn run(mut self) -> anyhow::Result<()> { - describe_metrics(); - let effective_interval = self.update_interval.add_jitter(0.1); let start_at = tokio::time::Instant::now() + effective_interval; let mut update_interval = tokio::time::interval_at(start_at, effective_interval); diff --git a/crates/node/src/lib.rs b/crates/node/src/lib.rs index 20672abc25..6b1ca81941 100644 --- a/crates/node/src/lib.rs +++ b/crates/node/src/lib.rs @@ -163,6 +163,8 @@ impl Node { mut address_book: AddressBook, ) -> Result { metric_definitions::describe_metrics(); + restate_storage_query_datafusion::describe_metrics(); + let mut server_builder = NetworkServerBuilder::new(&mut address_book); let config = updateable_config.pinned(); diff --git a/crates/storage-query-datafusion/Cargo.toml b/crates/storage-query-datafusion/Cargo.toml index 42758fe1f2..282e5c840b 100644 --- a/crates/storage-query-datafusion/Cargo.toml +++ b/crates/storage-query-datafusion/Cargo.toml @@ -40,7 +40,9 @@ datafusion-proto = { workspace = true } derive_more = { workspace = true, features = ["debug"] } enumset = { workspace = true } futures = { workspace = true } +gardal = { workspace = true } itertools = { workspace = true } +metrics = { workspace = true } parking_lot = { workspace = true } paste = { workspace = true } prost = { workspace = true } diff --git a/crates/storage-query-datafusion/src/context.rs b/crates/storage-query-datafusion/src/context.rs index 3a0ae95f3b..68be90454f 100644 --- a/crates/storage-query-datafusion/src/context.rs +++ b/crates/storage-query-datafusion/src/context.rs @@ -48,6 +48,16 @@ use crate::empty_invoker_status_handle::EmptyInvokerStatusHandle; use crate::node_fan_out::NodeWarnings; use crate::remote_query_scanner_manager::RemoteScannerManager; +type RateLimiter = gardal::SharedTokenBucket; + +#[derive(thiserror::Error, Debug)] +pub enum QueryError { + #[error("Datafusion error: {0}")] + DataFusion(#[from] datafusion::common::DataFusionError), + #[error("Rate limited")] + RateLimited(#[from] gardal::RateLimited), +} + const SYS_INVOCATION_VIEW: &str = "CREATE VIEW sys_invocation as SELECT ss.id, ss.vqueue_id, @@ -496,6 +506,7 @@ where pub struct QueryContext { sql_options: SQLOptions, datafusion_context: SessionContext, + rate_limiter: Option, } impl QueryContext { @@ -508,6 +519,9 @@ impl QueryContext { options.tmp_dir.clone(), options.query_parallelism(), &options.datafusion_options, + options.rate_limiting.as_ref().map(|limit| { + RateLimiter::new(gardal::Limit::from(limit.clone()), gardal::TokioClock) + }), )?; registerer.register(&ctx).await?; @@ -572,6 +586,7 @@ impl QueryContext { temp_folder: Option, default_parallelism: Option, datafusion_options: &HashMap, + rate_limiter: Option, ) -> Result { // // build the runtime @@ -625,10 +640,15 @@ impl QueryContext { Ok(Self { sql_options, datafusion_context: ctx, + rate_limiter, }) } - pub async fn execute(&self, sql: &str) -> datafusion::common::Result { + pub async fn execute(&self, sql: &str) -> Result { + if let Some(limiter) = self.rate_limiter.as_ref() { + limiter.try_consume_one()?; + } + let state = self.datafusion_context.state(); let statement = state.sql_to_statement(sql, &datafusion::config::Dialect::PostgreSQL)?; let plan = state.statement_to_plan(statement).await?; diff --git a/crates/storage-query-datafusion/src/lib.rs b/crates/storage-query-datafusion/src/lib.rs index 3aff9990a2..a6918f7482 100644 --- a/crates/storage-query-datafusion/src/lib.rs +++ b/crates/storage-query-datafusion/src/lib.rs @@ -23,6 +23,7 @@ mod keyed_service_status; mod locks; mod log; pub mod loglet_worker; +mod metric_definitions; mod node; pub mod node_fan_out; mod partition; @@ -57,6 +58,7 @@ use datafusion::arrow::record_batch::RecordBatch; use datafusion::error::DataFusionError; use datafusion::execution::TaskContext; use datafusion::physical_plan::PhysicalExpr; +pub use metric_definitions::describe_metrics; use prost::Message; #[cfg(test)] diff --git a/crates/storage-query-datafusion/src/metric_definitions.rs b/crates/storage-query-datafusion/src/metric_definitions.rs new file mode 100644 index 0000000000..7f224c4daf --- /dev/null +++ b/crates/storage-query-datafusion/src/metric_definitions.rs @@ -0,0 +1,28 @@ +// Copyright (c) 2023 - 2026 Restate Software, Inc., Restate GmbH. +// All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +use metrics::{Unit, describe_counter, describe_gauge}; + +pub(crate) const DATAFUSION_NUM_ACTIVE_SCANNERS: &str = "restate.datafusion.num_active_scanners"; +pub(crate) const DATAFUSION_NUM_SCANNERS: &str = "restate.datafusion.num_scanners.total"; + +pub fn describe_metrics() { + describe_gauge!( + DATAFUSION_NUM_ACTIVE_SCANNERS, + Unit::Count, + "Number of currently active DataFusion scanners" + ); + + describe_counter!( + DATAFUSION_NUM_SCANNERS, + Unit::Count, + "Total number of DataFusion scanners created" + ); +} diff --git a/crates/storage-query-datafusion/src/mocks.rs b/crates/storage-query-datafusion/src/mocks.rs index e1a4a0eedd..c203698801 100644 --- a/crates/storage-query-datafusion/src/mocks.rs +++ b/crates/storage-query-datafusion/src/mocks.rs @@ -250,7 +250,7 @@ impl MockQueryEngine { pub async fn execute( &self, sql: impl AsRef + Send, - ) -> datafusion::common::Result { + ) -> Result { self.2.execute(sql.as_ref()).await } } diff --git a/crates/storage-query-datafusion/src/remote_query_scanner_manager.rs b/crates/storage-query-datafusion/src/remote_query_scanner_manager.rs index 5e9e22ae75..07dfdbd273 100644 --- a/crates/storage-query-datafusion/src/remote_query_scanner_manager.rs +++ b/crates/storage-query-datafusion/src/remote_query_scanner_manager.rs @@ -10,16 +10,21 @@ use std::collections::BTreeMap; use std::fmt::{Debug, Formatter}; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; +use std::task::{Context, Poll}; use anyhow::{anyhow, bail}; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::DataFusionError; -use datafusion::execution::SendableRecordBatchStream; +use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream}; use datafusion::physical_plan::PhysicalExpr; use datafusion::physical_plan::metrics::Time; +use futures::Stream; +use metrics::{counter, gauge}; use parking_lot::Mutex; use restate_core::Metadata; @@ -29,6 +34,7 @@ use restate_types::identifiers::PartitionId; use restate_types::net::remote_query_scanner::{RemoteQueryScannerOpen, ScannerId}; use restate_types::sharding::KeyRange; +use crate::metric_definitions::{DATAFUSION_NUM_ACTIVE_SCANNERS, DATAFUSION_NUM_SCANNERS}; use crate::remote_query_scanner_client::{ RemoteScanner, RemoteScannerService, remote_scan_as_datafusion_stream, }; @@ -201,8 +207,10 @@ impl RemoteScannerManager { // make the local scanner available to serve a remote RPC. // see usages of [[local_partition_scanner]] // we use the table_name to associate a remote scanner with its local counterpart. - self.local_store_scanners - .register(name.clone(), local_scanner.clone()); + self.local_store_scanners.register( + name.clone(), + Arc::new(InstrumentedScanPartition(local_scanner)), + ); } RemotePartitionsScanner::new(self.clone(), name) @@ -213,8 +221,12 @@ impl RemoteScannerManager { /// as a `ScanPartition` adapter so it integrates with the existing remote /// scanner server infrastructure. pub fn register_node_scanner(&self, table_name: impl Into, scanner: Arc) { - self.local_store_scanners - .register(table_name, Arc::new(ScanToScanPartitionAdapter(scanner))); + self.local_store_scanners.register( + table_name, + Arc::new(InstrumentedScanPartition(Arc::new( + ScanToScanPartitionAdapter(scanner), + ))), + ); } pub fn local_partition_scanner(&self, table: &str) -> Option> { @@ -252,6 +264,66 @@ impl RemotePartitionsScanner { } } +/// Wraps a `SendableRecordBatchStream` to instrument metrics associated with +/// the stream lifetime. +struct InstrumentedRecordBatchStream(SendableRecordBatchStream); + +impl InstrumentedRecordBatchStream { + fn new(inner: SendableRecordBatchStream) -> Self { + gauge!(DATAFUSION_NUM_ACTIVE_SCANNERS).increment(1); + counter!(DATAFUSION_NUM_SCANNERS).increment(1); + Self(inner) + } +} + +impl Stream for InstrumentedRecordBatchStream { + type Item = datafusion::common::Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.0.as_mut().poll_next(cx) + } +} + +impl RecordBatchStream for InstrumentedRecordBatchStream { + fn schema(&self) -> SchemaRef { + self.0.schema() + } +} + +impl Drop for InstrumentedRecordBatchStream { + fn drop(&mut self) { + gauge!(DATAFUSION_NUM_ACTIVE_SCANNERS).decrement(1); + } +} + +/// Wraps `ScanPartition`s so returned record batch streams are instrumented. +#[derive(Debug)] +struct InstrumentedScanPartition(Arc); + +impl ScanPartition for InstrumentedScanPartition { + fn scan_partition( + &self, + partition_id: PartitionId, + range: KeyRange, + projection: SchemaRef, + predicate: Option>, + batch_size: usize, + limit: Option, + elapsed_compute: Time, + ) -> anyhow::Result { + let stream = self.0.scan_partition( + partition_id, + range, + projection, + predicate, + batch_size, + limit, + elapsed_compute, + )?; + Ok(Box::pin(InstrumentedRecordBatchStream::new(stream))) + } +} + /// Adapts a node-level `Scan` into a `ScanPartition` for the remote scanner /// server. The partition_id and range are ignored since this is a node-scoped table. #[derive(Debug)] diff --git a/crates/types/src/config/mod.rs b/crates/types/src/config/mod.rs index e7ae506667..80751c007c 100644 --- a/crates/types/src/config/mod.rs +++ b/crates/types/src/config/mod.rs @@ -29,6 +29,7 @@ mod networking; mod object_store; mod query_engine; mod rocksdb; +mod throttling; mod worker; pub use admin::*; @@ -49,6 +50,7 @@ pub use networking::*; pub use object_store::*; pub use query_engine::*; pub use rocksdb::*; +pub use throttling::*; pub use worker::*; use std::fs; diff --git a/crates/types/src/config/query_engine.rs b/crates/types/src/config/query_engine.rs index 0e5b22f443..06999d6c10 100644 --- a/crates/types/src/config/query_engine.rs +++ b/crates/types/src/config/query_engine.rs @@ -16,6 +16,8 @@ use serde_with::serde_as; use restate_util_bytecount::NonZeroByteCount; +use crate::config::throttling::ThrottlingOptions; + /// # Storage query engine options #[serde_as] #[derive(Debug, Clone, Serialize, Deserialize, derive_builder::Builder)] @@ -44,6 +46,15 @@ pub struct QueryEngineOptions { #[serde(default, skip_serializing_if = "HashMap::is_empty")] #[cfg_attr(feature = "schemars", schemars(skip))] pub datafusion_options: HashMap, + + /// # Query rate limiting + /// + /// Per-node rate limiting options for datafusion query execution. When exceeded, + /// the query will fail with TOO_MANY_REQUESTS. If unset, no rate limiting is applied. + /// + /// Since v1.8.0 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub rate_limiting: Option, } impl QueryEngineOptions { @@ -60,6 +71,7 @@ impl Default for QueryEngineOptions { tmp_dir: None, query_parallelism: None, datafusion_options: HashMap::new(), + rate_limiting: None, } } } diff --git a/crates/types/src/config/throttling.rs b/crates/types/src/config/throttling.rs new file mode 100644 index 0000000000..91e3139393 --- /dev/null +++ b/crates/types/src/config/throttling.rs @@ -0,0 +1,46 @@ +use std::num::NonZeroU32; + +use serde::{Deserialize, Serialize}; + +use crate::rate::Rate; + +/// # Throttling options +/// +/// Token-bucket throttling options. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +#[serde(rename_all = "kebab-case")] +pub struct ThrottlingOptions { + /// # Refill rate + /// + /// The rate at which the tokens are replenished. + /// + /// Syntax: `/` where `` is `s|sec|second`, `m|min|minute`, or `h|hr|hour`. + /// unit defaults to per second if not specified. + #[cfg_attr(feature = "schemars", schemars(with = "String"))] + pub rate: Rate, + + /// # Burst capacity + /// + /// The maximum number of tokens the bucket can hold. + /// Default to the rate value if not specified. + pub capacity: Option, +} + +impl From for gardal::Limit { + fn from(options: ThrottlingOptions) -> Self { + use gardal::Limit; + + let mut limit = match options.rate { + Rate::Second(rate) => Limit::per_second(rate), + Rate::Minute(rate) => Limit::per_minute(rate), + Rate::Hour(rate) => Limit::per_hour(rate), + }; + + if let Some(capacity) = options.capacity { + limit = limit.with_burst(capacity); + } + + limit + } +} diff --git a/crates/types/src/config/worker.rs b/crates/types/src/config/worker.rs index 5b19f2a428..ccc03926c0 100644 --- a/crates/types/src/config/worker.rs +++ b/crates/types/src/config/worker.rs @@ -24,13 +24,13 @@ use super::{ BackgroundWorkBudget, CommonOptions, DEFAULT_MESSAGE_SIZE_LIMIT, NetworkingOptions, ObjectStoreOptions, RocksDbOptions, }; +use crate::config::throttling::ThrottlingOptions; use crate::config::{ AwsLambdaOptions, DeprecatedAwsLambdaOptions, DeprecatedHttpOptions, HttpOptions, IngestionOptions, }; use crate::identifiers::PartitionId; use crate::net::connect_opts::MESSAGE_SIZE_OVERHEAD; -use crate::rate::Rate; use crate::retries::RetryPolicy; const MIN_ROCKSDB_MEMORY: NonZeroByteCount = @@ -1062,47 +1062,6 @@ impl SnapshotsOptions { } } -/// # Throttling options -/// -/// Throttling options per invoker. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] -#[serde(rename_all = "kebab-case")] -pub struct ThrottlingOptions { - /// # Refill rate - /// - /// The rate at which the tokens are replenished. - /// - /// Syntax: `/` where `` is `s|sec|second`, `m|min|minute`, or `h|hr|hour`. - /// unit defaults to per second if not specified. - #[cfg_attr(feature = "schemars", schemars(with = "String"))] - pub rate: Rate, - - /// # Burst capacity - /// - /// The maximum number of tokens the bucket can hold. - /// Default to the rate value if not specified. - pub capacity: Option, -} - -impl From for gardal::Limit { - fn from(options: ThrottlingOptions) -> Self { - use gardal::Limit; - - let mut limit = match options.rate { - Rate::Second(rate) => Limit::per_second(rate), - Rate::Minute(rate) => Limit::per_minute(rate), - Rate::Hour(rate) => Limit::per_hour(rate), - }; - - if let Some(capacity) = options.capacity { - limit = limit.with_burst(capacity); - } - - limit - } -} - mod serde_helpers { use std::num::NonZeroUsize; diff --git a/release-notes/unreleased/query-rate-limiting.md b/release-notes/unreleased/query-rate-limiting.md new file mode 100644 index 0000000000..65d9f8648a --- /dev/null +++ b/release-notes/unreleased/query-rate-limiting.md @@ -0,0 +1,46 @@ +# Release Notes: Optional Rate limiting for storage query execution + +## New Feature + +### What Changed +Datafusion query execution (the admin REST `/query` endpoint and the +cluster-controller gRPC query service) can now be rate limited via a new, +optional `rate-limiting` knob on the query engine options. The limiter is a +token bucket: each query consumes one token, and queries that exceed the limit +fail fast instead of executing. + +When the limit is exceeded, the admin REST endpoint responds with `429 Too Many Requests` and a `Retry-After` header. + +If `rate-limiting` is unset (the default), no rate limiting is applied and +behavior is unchanged. + +### Why This Matters +Expensive or high-volume ad-hoc SQL queries can compete with the node's core +workload for CPU and memory. Operators can now cap query throughput on demand to protect +the rest of the system. + +### Impact on Users +- New deployments: no rate limiting unless explicitly configured. +- Existing deployments: no change on upgrade unless explicitly configured. + +### Migration Guidance +To rate limit queries, configure the rate (and optionally a burst capacity) in +the query engine options: + +```toml +[admin.query-engine.rate-limiting] +rate = "100/s" # /, unit is s|sec|second, m|min|minute, or h|hr|hour +capacity = 200 # optional burst capacity; defaults to the rate value +``` + +## Behavioral Change + +### What Changed +The admin REST `/query` endpoint now returns `400 Bad Request` (instead of +`500 Internal Server Error`) for client-side query errors — Datafusion `Plan`, +`SchemaError`, and `SQL` errors (e.g. invalid SQL or unknown columns/tables). + +### Impact on Users +- Clients that distinguish failures by status code will now see `4xx` for + malformed queries instead of `5xx`. Internal/execution failures continue to + return `500`. diff --git a/tools/restate-doctor/src/commands/snapshot/server.rs b/tools/restate-doctor/src/commands/snapshot/server.rs index fed6d26313..093d12dd17 100644 --- a/tools/restate-doctor/src/commands/snapshot/server.rs +++ b/tools/restate-doctor/src/commands/snapshot/server.rs @@ -55,21 +55,28 @@ struct QueryErrorBody { } /// Errors that can occur when executing a query. -#[derive(Debug, thiserror::Error)] -enum QueryError { - #[error("Datafusion error: {0}")] - Datafusion(#[from] DataFusionError), +struct QueryError(restate_storage_query_datafusion::context::QueryError); + +impl From for QueryError { + fn from(err: datafusion::error::DataFusionError) -> Self { + Self(restate_storage_query_datafusion::context::QueryError::DataFusion(err)) + } } impl IntoResponse for QueryError { fn into_response(self) -> Response { - let status_code = match &self { - QueryError::Datafusion(_) => StatusCode::INTERNAL_SERVER_ERROR, + let status_code = match &self.0 { + restate_storage_query_datafusion::context::QueryError::DataFusion(_) => { + StatusCode::INTERNAL_SERVER_ERROR + } + restate_storage_query_datafusion::context::QueryError::RateLimited(_) => { + StatusCode::TOO_MANY_REQUESTS + } }; ( status_code, Json(QueryErrorBody { - message: self.to_string(), + message: self.0.to_string(), }), ) .into_response() @@ -104,7 +111,7 @@ async fn query( headers: HeaderMap, Json(payload): Json, ) -> Result { - let query_result = ctx.execute(&payload.query).await?; + let query_result = ctx.execute(&payload.query).await.map_err(QueryError)?; let (result_stream, content_type) = match headers.get(http::header::ACCEPT) { Some(v) if v == HeaderValue::from_static("application/json") => (