From bc84f3849352f4949d4a72a3c80f27afae374ebd Mon Sep 17 00:00:00 2001 From: MohamedBassem Date: Fri, 26 Jun 2026 16:32:47 +0100 Subject: [PATCH 1/2] [admin] Introduce ratelimits knobs for datafusion queries This PR introduces ratelimiting knobs for datafusion queries as an emergency mechanism to control the rate of DF queries if they end up putting the system under pressure. By default, this rate limiting is disabled. Along the way, the PR also translates some client side DF errors to be 400s (bad request) instead of 500s. --- Cargo.lock | 4 +- crates/admin/Cargo.toml | 1 + .../cluster_controller/grpc_svc_handler.rs | 11 ++++- crates/admin/src/rest_api/query.rs | 41 +++++++++++++++-- crates/storage-query-datafusion/Cargo.toml | 1 + .../storage-query-datafusion/src/context.rs | 22 ++++++++- crates/storage-query-datafusion/src/mocks.rs | 2 +- crates/types/src/config/mod.rs | 2 + crates/types/src/config/query_engine.rs | 12 +++++ crates/types/src/config/throttling.rs | 46 +++++++++++++++++++ crates/types/src/config/worker.rs | 43 +---------------- .../unreleased/query-rate-limiting.md | 46 +++++++++++++++++++ .../src/commands/snapshot/server.rs | 23 ++++++---- 13 files changed, 196 insertions(+), 58 deletions(-) create mode 100644 crates/types/src/config/throttling.rs create mode 100644 release-notes/unreleased/query-rate-limiting.md diff --git a/Cargo.lock b/Cargo.lock index 655b0424f3..60e68790bc 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,6 +8622,7 @@ dependencies = [ "derive_more", "enumset", "futures", + "gardal", "googletest", "itertools 0.14.0", "parking_lot", 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/rest_api/query.rs b/crates/admin/src/rest_api/query.rs index ad226ad299..f283b696d6 100644 --- a/crates/admin/src/rest_api/query.rs +++ b/crates/admin/src/rest_api/query.rs @@ -12,8 +12,6 @@ use std::io::Write; use std::pin::Pin; use std::sync::Arc; -use crate::query_utils::{RecordBatchWriter, WriteRecordBatchStream}; -use crate::state::AdminServiceState; use axum::extract::State; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; @@ -29,11 +27,17 @@ use http::{HeaderMap, HeaderValue}; use http_body::Frame; use http_body_util::StreamBody; 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::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 +52,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(), }), diff --git a/crates/storage-query-datafusion/Cargo.toml b/crates/storage-query-datafusion/Cargo.toml index 42758fe1f2..56f9be9f72 100644 --- a/crates/storage-query-datafusion/Cargo.toml +++ b/crates/storage-query-datafusion/Cargo.toml @@ -40,6 +40,7 @@ datafusion-proto = { workspace = true } derive_more = { workspace = true, features = ["debug"] } enumset = { workspace = true } futures = { workspace = true } +gardal = { workspace = true } itertools = { workspace = true } parking_lot = { workspace = true } paste = { 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/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/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") => ( From 0804709e18f26237fa769ffd26bf832ddb740b15 Mon Sep 17 00:00:00 2001 From: MohamedBassem Date: Tue, 30 Jun 2026 12:04:43 +0100 Subject: [PATCH 2/2] [df] Export more metrics from query engine This PR introduces two main metrics for datafusion. 1. Latency tracking for the `/query` endpoint. Although, the exported histograms will be influnced by the query complexity submitted by the user, I found that if I'm adding num /query calls and a metric for slow durations, I might as well just track the whole endpoint latency which will give us both and a couple more. We also typically have baseline query load (for exporting metrics) which will at least help set the baseline for such histograms. 2. I've also added tracking for the number of active local scanner that we have (guage and counters) that might help track if we're for example leaking scanner (due to aborted original request). It might be useful to export the number of scanner per table, put I would like more opinions on whether this is worth the extra metrics dimensions or not. Now, for the things that I didn't add: 1. I was planning to report the active time that the scanner spent actively reading. However, this turned out to be tricky because scanners are long lived throughout the query's lifetime. So capturing their lifetime will also capture the time they're not polled (e.g. between batches RPCs from the remote scanner). I considered instrumenting the handling for the NextBatch RPC from the remote scanner manager, but that wouldn't cover local scans. I think our rocksb metrics in background iterator might be good enough to capture the cost of the scans themseleves? 2. I was planning to log some metrics from the query coordinator about the number of batches/rows and latency for remote scanner, etc. But it turned out that you can get those on demand from DF with `explin analyze $query` so ditched those as well. --- Cargo.lock | 1 + crates/admin/src/metric_definitions.rs | 11 +++ crates/admin/src/rest_api/query.rs | 43 +++++++++- crates/admin/src/service.rs | 2 + crates/admin/src/storage_accounting.rs | 4 +- crates/node/src/lib.rs | 2 + crates/storage-query-datafusion/Cargo.toml | 1 + crates/storage-query-datafusion/src/lib.rs | 2 + .../src/metric_definitions.rs | 28 +++++++ .../src/remote_query_scanner_manager.rs | 82 +++++++++++++++++-- 10 files changed, 167 insertions(+), 9 deletions(-) create mode 100644 crates/storage-query-datafusion/src/metric_definitions.rs diff --git a/Cargo.lock b/Cargo.lock index 60e68790bc..ca86fb82cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8625,6 +8625,7 @@ dependencies = [ "gardal", "googletest", "itertools 0.14.0", + "metrics", "parking_lot", "paste", "prost", 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 f283b696d6..6d889235c5 100644 --- a/crates/admin/src/rest_api/query.rs +++ b/crates/admin/src/rest_api/query.rs @@ -11,6 +11,8 @@ use std::io::Write; use std::pin::Pin; use std::sync::Arc; +use std::task::{Context, Poll}; +use std::time::Instant; use axum::extract::State; use axum::http::StatusCode; @@ -22,10 +24,11 @@ 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; @@ -34,6 +37,7 @@ use restate_core::network::TransportConnect; use restate_types::invocation::client::InvocationClient; use restate_types::schema::registry::{DiscoveryClient, MetadataService, TelemetryClient}; +use crate::metric_definitions::QUERY_ENGINE_QUERY_DURATION_SECONDS; use crate::query_utils::{RecordBatchWriter, WriteRecordBatchStream}; use crate::state::AdminServiceState; @@ -129,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); }; @@ -161,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)) @@ -168,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 56f9be9f72..282e5c840b 100644 --- a/crates/storage-query-datafusion/Cargo.toml +++ b/crates/storage-query-datafusion/Cargo.toml @@ -42,6 +42,7 @@ 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/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/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)]