[df] Export more metrics from query engine

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 }

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<QueryWarning> {
     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(..)

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"
     )
 }

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<restate_storage_query_datafusion::context::QueryError> 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,13 +169,46 @@ 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))
         .expect("content-type header is correct")
         .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<S> {
+    inner: S,
+    _guard: QueryDurationGuard,
+}
+
+impl<S: Stream + Unpin> Stream for Timed<S> {
+    type Item = S::Item;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        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 :(

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<HttpClient>,
     ) -> Self {
+        describe_metrics();
         let metadata_client = metadata_writer.raw_metadata_store_client().clone();
         Self {
             listeners,

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);

crates/node/src/lib.rs

@@ -163,6 +163,8 @@ impl Node {
         mut address_book: AddressBook,
     ) -> Result<Self, BuildError> {
         metric_definitions::describe_metrics();
+        restate_storage_query_datafusion::describe_metrics();
+
         let mut server_builder = NetworkServerBuilder::new(&mut address_book);
         let config = updateable_config.pinned();
 

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 }

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<gardal::TokioClock>;
+
+#[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<RateLimiter>,
 }
 
 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<String>,
         default_parallelism: Option<usize>,
         datafusion_options: &HashMap<String, String>,
+        rate_limiter: Option<RateLimiter>,
     ) -> Result<Self, DataFusionError> {
         //
         // 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<QueryResult> {
+    pub async fn execute(&self, sql: &str) -> Result<QueryResult, QueryError> {
+        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?;

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)]