diff --git a/Cargo.lock b/Cargo.lock index 6997f0d782b..834b0b541bb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -13589,6 +13589,7 @@ dependencies = [ "hex", "hostname 0.3.1", "http 1.5.0", + "http-body 1.0.1", "humantime", "humantime-serde", "hyper 1.6.0", diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index c1af56a6281..f9c3b45832f 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -70,6 +70,7 @@ headers = "0.4" hex.workspace = true hostname = "0.3" http.workspace = true +http-body = "1" humantime.workspace = true humantime-serde.workspace = true hyper = { workspace = true, features = ["full"] } diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index 47f9953a187..af6b2499e7f 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -319,6 +319,18 @@ pub enum Error { error: std::io::Error, }, + #[snafu(display( + "Decompressed request body is too large: {} bytes exceeds the limit {} bytes", + size, + limit + ))] + DecompressedBodyTooLarge { + size: u64, + limit: u64, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to compress prometheus remote request"))] CompressPromRemoteRequest { #[snafu(implicit)] @@ -785,6 +797,7 @@ impl ErrorExt for Error { | DecompressSnappyPromRemoteRequest { .. } | DecompressSnappyLokiRequest { .. } | DecompressZstdPromRemoteRequest { .. } + | DecompressedBodyTooLarge { .. } | InvalidPromRemoteRequest { .. } | InvalidFlightTicket { .. } | InvalidPrepareStatement { .. } diff --git a/src/servers/src/grpc.rs b/src/servers/src/grpc.rs index 50dd0b69c4a..1be1f8215b5 100644 --- a/src/servers/src/grpc.rs +++ b/src/servers/src/grpc.rs @@ -50,6 +50,7 @@ use tonic::{Request, Response, Status}; use tonic_reflection::server::v1::{ServerReflection, ServerReflectionServer}; use crate::error::{AlreadyStartedSnafu, InternalSnafu, Result, StartGrpcSnafu, TcpBindSnafu}; +use crate::grpc::memory_limit::MemoryLimiterExtensionService; use crate::install_default_crypto_provider; use crate::metrics::MetricsMiddlewareLayer; use crate::otel_arrow::{HeaderInterceptor, OtelArrowServiceHandler}; @@ -213,6 +214,14 @@ impl FlightCompression { } } +/// The wrapped OTLP Arrow service type used by [`GrpcServer`]. +type OtelArrowService = MemoryLimiterExtensionService< + InterceptedService< + ArrowMetricsServiceServer>, + HeaderInterceptor, + >, +>; + pub struct GrpcServer { // states shutdown_tx: Mutex>>, @@ -224,14 +233,7 @@ pub struct GrpcServer { // tls config tls_config: Option, // Otel arrow service - otel_arrow_service: Mutex< - Option< - InterceptedService< - ArrowMetricsServiceServer>, - HeaderInterceptor, - >, - >, - >, + otel_arrow_service: Mutex>, bind_addr: Option, name: Option, config: GrpcServerConfig, diff --git a/src/servers/src/grpc/builder.rs b/src/servers/src/grpc/builder.rs index d240cde3f1b..99513dcceb3 100644 --- a/src/servers/src/grpc/builder.rs +++ b/src/servers/src/grpc/builder.rs @@ -34,12 +34,13 @@ use tonic::codegen::Service; use tonic::service::RoutesBuilder; use tonic::service::interceptor::InterceptedService; use tonic::transport::{Identity, ServerTlsConfig}; -use tower::Layer; +use tower::{Layer, ServiceBuilder}; use crate::grpc::database::DatabaseService; use crate::grpc::flight::{FlightCraftRef, FlightCraftWrapper}; use crate::grpc::frontend_grpc_handler::FrontendGrpcHandler; use crate::grpc::greptime_handler::GreptimeRequestHandler; +use crate::grpc::memory_limit::{MemoryLimiterExtensionLayer, MemoryLimiterExtensionService}; use crate::grpc::prom_query_gateway::PrometheusGatewayService; use crate::grpc::region_server::{RegionServerHandlerRef, RegionServerRequestHandler}; use crate::grpc::{GrpcServer, GrpcServerConfig}; @@ -71,6 +72,7 @@ macro_rules! add_service { let service_with_limiter = $crate::tower::ServiceBuilder::new() .layer(MemoryLimiterExtensionLayer::new( $builder.memory_limiter().clone(), + max_recv_message_size, )) .service(service_builder); @@ -87,9 +89,11 @@ pub struct GrpcServerBuilder { routes_builder: RoutesBuilder, tls_config: Option, otel_arrow_service: Option< - InterceptedService< - ArrowMetricsServiceServer>, - HeaderInterceptor, + MemoryLimiterExtensionService< + InterceptedService< + ArrowMetricsServiceServer>, + HeaderInterceptor, + >, >, >, memory_limiter: ServerMemoryLimiter, @@ -185,6 +189,13 @@ impl GrpcServerBuilder { .accept_compressed(CompressionEncoding::Zstd) .send_compressed(CompressionEncoding::Zstd); let svc = InterceptedService::new(server, HeaderInterceptor {}); + // Same pre-decode memory admission as `add_service!`. + let svc = ServiceBuilder::new() + .layer(MemoryLimiterExtensionLayer::new( + self.memory_limiter.clone(), + self.config.max_recv_message_size, + )) + .service(svc); self.otel_arrow_service = Some(svc); self } diff --git a/src/servers/src/grpc/database.rs b/src/servers/src/grpc/database.rs index 2a63a982bef..237d760269f 100644 --- a/src/servers/src/grpc/database.rs +++ b/src/servers/src/grpc/database.rs @@ -25,6 +25,7 @@ use session::context::Channel; use tonic::{Request, Response, Status, Streaming}; use crate::grpc::greptime_handler::GreptimeRequestHandler; +use crate::grpc::memory_limit::PreDecodeMemoryReservation; use crate::grpc::{TonicResult, cancellation}; use crate::hint_headers; use crate::request_memory_limiter::ServerMemoryLimiter; @@ -57,7 +58,20 @@ impl GreptimeDatabase for DatabaseService { remote_addr, hints ); - let _guard = if let Some(limiter) = request.extensions().get::() { + // Retain the pre-decode reservation for the whole request: the + // extension holding the guard would be dropped when the request is + // consumed below, but the post-decode charge is skipped while it is + // active. + let _pre_reservation = request + .extensions() + .get::() + .cloned(); + let _guard = if _pre_reservation.is_some() { + // Compressed requests already reserved the worst-case decoded + // size before tonic decompressed the message; skip the exact + // post-decode charge to avoid double accounting. + None + } else if let Some(limiter) = request.extensions().get::() { let message_size = request.get_ref().encoded_len() as u64; Some(limiter.acquire(message_size).await?) } else { @@ -117,16 +131,31 @@ impl GreptimeDatabase for DatabaseService { ); let limiter = request.extensions().get::().cloned(); + // For compressed streams the whole stream's decoding memory was + // reserved before tonic started decompressing; messages are decoded + // one at a time, so the reservation covers each of them. The + // reservation is retained below for the whole stream: the extension + // holding the guard would otherwise be dropped when the request is + // consumed, while per-message charges stay skipped. + let reservation = request + .extensions() + .get::() + .cloned(); + let pre_reserved = reservation.is_some(); let handler = self.handler.clone(); let request_future = async move { let mut affected_rows = 0; + // Hold the pre-decode reservation until the stream is exhausted. + let _reservation = reservation; let mut stream = request.into_inner(); while let Some(request) = stream.next().await { let request = request?; - let _guard = if let Some(limiter_ref) = &limiter { + let _guard = if pre_reserved { + None + } else if let Some(limiter_ref) = &limiter { let message_size = request.encoded_len() as u64; Some(limiter_ref.acquire(message_size).await?) } else { diff --git a/src/servers/src/grpc/memory_limit.rs b/src/servers/src/grpc/memory_limit.rs index 722a2c68c08..c9c62e88dd1 100644 --- a/src/servers/src/grpc/memory_limit.rs +++ b/src/servers/src/grpc/memory_limit.rs @@ -12,22 +12,69 @@ // See the License for the specific language governing permissions and // limitations under the License. +//! Aggregate memory admission for gRPC services. +//! +//! Tonic decompresses and decodes a complete request message (up to +//! `max_recv_message_size`, 512 MiB by default) *before* the typed handler is +//! invoked, so a handler that charges the [`ServerMemoryLimiter`] only sees +//! the request after the memory is already allocated. For transport-compressed +//! messages (`grpc-encoding: gzip/zstd`) a few KiB on the wire can expand to +//! hundreds of MiB during that window, which would bypass any aggregate quota. +//! +//! [`MemoryLimiterExtensionLayer`] therefore reserves the configured maximum +//! decoded message size for every compressed request *before* the inner +//! (tonic) service runs, and holds the reservation for the whole request. +//! Handlers can detect the reservation via [`PreDecodeMemoryReservation`] in +//! the request extensions and skip their own post-decode charge to avoid +//! double accounting. +//! +//! The reservation is intentionally conservative (it upper-bounds the decoded +//! size before it is known); with the default unlimited limiter it is a no-op. + +use std::convert::Infallible; +use std::sync::Arc; use std::task::{Context, Poll}; +use axum::response::IntoResponse; +use common_memory_manager::MemoryGuard; use futures::future::BoxFuture; +use http::Request; +use tonic::Status; +use tonic::codegen::Service; use tonic::server::NamedService; -use tower::{Layer, Service}; +use tower::Layer; use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; + +/// The gRPC request compression selector header. +const GRPC_ENCODING_HEADER: &str = "grpc-encoding"; +/// The "no compression" encoding value. +const IDENTITY_ENCODING: &str = "identity"; + +/// Present in the request extensions when memory for the (compressed) +/// request's decoded message was reserved before tonic decompression. +/// +/// Handlers that charge the [`ServerMemoryLimiter`] after decoding should skip +/// that charge when this marker is present: the reservation already covers the +/// peak decoding memory and stays alive for the whole request. +#[derive(Clone)] +pub(crate) struct PreDecodeMemoryReservation { + _guard: Arc>, +} #[derive(Clone)] pub struct MemoryLimiterExtensionLayer { limiter: ServerMemoryLimiter, + max_decoding_message_size: usize, } impl MemoryLimiterExtensionLayer { - pub fn new(limiter: ServerMemoryLimiter) -> Self { - Self { limiter } + pub fn new(limiter: ServerMemoryLimiter, max_decoding_message_size: usize) -> Self { + Self { + limiter, + max_decoding_message_size, + } } } @@ -38,6 +85,7 @@ impl Layer for MemoryLimiterExtensionLayer { MemoryLimiterExtensionService { inner: service, limiter: self.limiter.clone(), + max_decoding_message_size: self.max_decoding_message_size, } } } @@ -46,27 +94,199 @@ impl Layer for MemoryLimiterExtensionLayer { pub struct MemoryLimiterExtensionService { inner: S, limiter: ServerMemoryLimiter, + max_decoding_message_size: usize, } impl NamedService for MemoryLimiterExtensionService { const NAME: &'static str = S::NAME; } -impl Service> for MemoryLimiterExtensionService +impl Service> for MemoryLimiterExtensionService where - S: Service>, + S: Service, Error = Infallible> + Clone + Send + 'static, + S::Response: axum::response::IntoResponse, S::Future: Send + 'static, + ReqBody: Send + 'static, { - type Response = S::Response; - type Error = S::Error; + type Response = axum::response::Response; + type Error = Infallible; type Future = BoxFuture<'static, Result>; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_ready(cx) } - fn call(&mut self, mut req: http::Request) -> Self::Future { - req.extensions_mut().insert(self.limiter.clone()); - Box::pin(self.inner.call(req)) + fn call(&mut self, mut req: Request) -> Self::Future { + // Own a clone of the service so the returned future does not borrow + // `self` (tower's load-shedding pattern). + let mut this = self.clone(); + Box::pin(async move { + req.extensions_mut().insert(this.limiter.clone()); + + let compressed = req + .headers() + .get(GRPC_ENCODING_HEADER) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| !value.eq_ignore_ascii_case(IDENTITY_ENCODING)); + + if compressed { + // A compressed message can expand up to the decoding limit + // inside tonic, before any handler runs. Reserve that worst + // case against the aggregate quota so the decoding phase is + // admitted as well. No-op for an unlimited limiter. + let reservation = this.max_decoding_message_size as u64; + match this.limiter.acquire(reservation).await { + Ok(guard) => { + req.extensions_mut().insert(PreDecodeMemoryReservation { + _guard: Arc::new(guard), + }); + } + Err(e) => { + return Ok(Status::resource_exhausted(format!( + "request memory limit exceeded: {e}" + )) + .into_http::() + .into_response()); + } + } + } + + match this.inner.call(req).await { + Ok(response) => Ok(response.into_response()), + Err(e) => match e {}, + } + }) + } +} + +#[cfg(test)] +mod tests { + use std::convert::Infallible; + + use axum::body::Body; + use axum::response::IntoResponse; + use common_memory_manager::OnExhaustedPolicy; + use futures_util::future::BoxFuture; + use http::{HeaderValue, StatusCode}; + use tower::ServiceExt; + + use super::*; + use crate::grpc::memory_limit::MemoryLimiterExtensionLayer; + + #[derive(Clone)] + struct EchoService; + + impl Service> for EchoService { + type Response = axum::response::Response; + type Error = Infallible; + type Future = BoxFuture<'static, Result>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, req: Request) -> Self::Future { + let saw_reservation = req + .extensions() + .get::() + .is_some(); + let saw_limiter = req.extensions().get::().is_some(); + Box::pin(async move { + Ok(( + StatusCode::OK, + format!("reservation={saw_reservation} limiter={saw_limiter}"), + ) + .into_response()) + }) + } + } + + impl NamedService for EchoService { + const NAME: &'static str = "test.Echo"; + } + + #[tokio::test] + async fn test_inserts_limiter_for_uncompressed_requests() { + let limiter = ServerMemoryLimiter::new(1024 * 1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService); + + let req = Request::builder().body(Body::empty()).unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!( + &body[..], + &b"reservation=false limiter=true"[..], + "uncompressed requests must not be pre-reserved" + ); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_reserves_max_message_size_for_compressed_requests() { + let max_size = 512 * 1024 * 1024; + let limiter = ServerMemoryLimiter::new(1024 * 1024 * 1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), max_size).layer(EchoService); + + // The inner service observes the reservation; use a oneshot call and + // check the limiter was charged while the call is in flight via the + // response extensions observed by the echo service. + let req = Request::builder() + .header(GRPC_ENCODING_HEADER, HeaderValue::from_static("gzip")) + .body(Body::empty()) + .unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!( + &body[..], + &b"reservation=true limiter=true"[..], + "compressed requests must carry a pre-decode reservation" + ); + // The reservation is released once the inner call (request) finishes. + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_rejects_compressed_request_when_quota_exhausted() { + // Quota smaller than the decoding limit: the reservation cannot fit. + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService); + + let req = Request::builder() + .header(GRPC_ENCODING_HEADER, HeaderValue::from_static("zstd")) + .body(Body::empty()) + .unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + // gRPC errors are HTTP 200 with a grpc-status header. + assert_eq!( + res.headers() + .get("grpc-status") + .and_then(|v| v.to_str().ok()), + Some("8"), + "expected RESOURCE_EXHAUSTED grpc status" + ); + } + + #[tokio::test] + async fn test_identity_encoding_is_not_reserved() { + let limiter = ServerMemoryLimiter::new(1024 * 1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService); + + let req = Request::builder() + .header( + GRPC_ENCODING_HEADER, + HeaderValue::from_static(IDENTITY_ENCODING), + ) + .body(Body::empty()) + .unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!(&body[..], &b"reservation=false limiter=true"[..]); } } diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index b9fcd0a9420..a33e52a1d40 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -716,7 +716,7 @@ impl HttpServerBuilder { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/influxdb"), - HttpServer::route_influxdb(handler), + HttpServer::route_influxdb(handler, self.memory_limiter.clone()), ), ..self } @@ -736,6 +736,7 @@ impl HttpServerBuilder { prom_store_with_metric_engine, prom_validation_mode, pending_rows_batcher, + memory_limiter: self.memory_limiter.clone(), }; Self { @@ -765,7 +766,7 @@ impl HttpServerBuilder { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/otlp"), - HttpServer::route_otlp(handler, with_metric_engine), + HttpServer::route_otlp(handler, with_metric_engine, self.memory_limiter.clone()), ), ..self } @@ -799,23 +800,23 @@ impl HttpServerBuilder { let router = self.router.nest( &format!("/{HTTP_API_VERSION}"), - HttpServer::route_pipelines(log_state.clone()), + HttpServer::route_pipelines(log_state.clone(), self.memory_limiter.clone()), ); // deprecated since v0.11.0. Use `/logs` and `/pipelines` instead. let router = router.nest( &format!("/{HTTP_API_VERSION}/events"), #[allow(deprecated)] - HttpServer::route_log_deprecated(log_state.clone()), + HttpServer::route_log_deprecated(log_state.clone(), self.memory_limiter.clone()), ); let router = router.nest( &format!("/{HTTP_API_VERSION}/loki"), - HttpServer::route_loki(log_state.clone()), + HttpServer::route_loki(log_state.clone(), self.memory_limiter.clone()), ); let router = router.nest( &format!("/{HTTP_API_VERSION}/elasticsearch"), - HttpServer::route_elasticsearch(log_state.clone()), + HttpServer::route_elasticsearch(log_state.clone(), self.memory_limiter.clone()), ); let router = router.nest( @@ -827,7 +828,7 @@ impl HttpServerBuilder { let router = router.nest( &format!("/{HTTP_API_VERSION}/splunk"), - HttpServer::route_splunk(log_state), + HttpServer::route_splunk(log_state, self.memory_limiter.clone()), ); Self { router, ..self } @@ -858,7 +859,7 @@ impl HttpServerBuilder { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/dashboards"), - HttpServer::route_dashboard(handler), + HttpServer::route_dashboard(handler, self.memory_limiter.clone()), ), ..self } @@ -1184,9 +1185,13 @@ impl HttpServer { .with_state(metrics_handler) } - fn route_loki(log_state: LogState) -> Router { + fn route_loki(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router { Router::new() .route("/api/v1/push", routing::post(loki::loki_ingest)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1198,7 +1203,7 @@ impl HttpServer { .with_state(log_state) } - fn route_splunk(log_state: LogState) -> Router { + fn route_splunk(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router { Router::new() .route( "/services/collector/health", @@ -1226,6 +1231,10 @@ impl HttpServer { "/services/collector/raw/1.0", routing::post(splunk::handle_raw), ) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1237,7 +1246,10 @@ impl HttpServer { .with_state(log_state) } - fn route_elasticsearch(log_state: LogState) -> Router { + fn route_elasticsearch( + log_state: LogState, + memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new() // Return fake responsefor HEAD '/' request. .route( @@ -1309,6 +1321,10 @@ impl HttpServer { axum::Json(serde_json::json!({})), )), ) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer(ServiceBuilder::new().layer(RequestDecompressionLayer::new())) .layer(middleware::from_fn_with_state( BatchingProtocol::Elasticsearch, @@ -1318,7 +1334,10 @@ impl HttpServer { } #[deprecated(since = "0.11.0", note = "Use `route_pipelines()` instead.")] - fn route_log_deprecated(log_state: LogState) -> Router { + fn route_log_deprecated( + log_state: LogState, + memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new() .route("/logs", routing::post(event::log_ingester)) .route( @@ -1334,6 +1353,10 @@ impl HttpServer { routing::delete(event::delete_pipeline), ) .route("/pipelines/dryrun", routing::post(event::pipeline_dryrun)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1345,7 +1368,7 @@ impl HttpServer { .with_state(log_state) } - fn route_pipelines(log_state: LogState) -> Router { + fn route_pipelines(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router { Router::new() .route("/ingest", routing::post(event::log_ingester)) .route( @@ -1365,6 +1388,10 @@ impl HttpServer { routing::delete(event::delete_pipeline), ) .route("/pipelines/_dryrun", routing::post(event::pipeline_dryrun)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1465,10 +1492,17 @@ impl HttpServer { .with_state(state) } - fn route_influxdb(influxdb_handler: InfluxdbLineProtocolHandlerRef) -> Router { + fn route_influxdb( + influxdb_handler: InfluxdbLineProtocolHandlerRef, + memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new() .route("/write", routing::post(influxdb_write_v1)) .route("/api/v2/write", routing::post(influxdb_write_v2)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1495,11 +1529,16 @@ impl HttpServer { fn route_otlp( otlp_handler: OpenTelemetryProtocolHandlerRef, with_metric_engine: bool, + memory_limiter: ServerMemoryLimiter, ) -> Router { Router::new() .route("/v1/metrics", routing::post(otlp::metrics)) .route("/v1/traces", routing::post(otlp::traces)) .route("/v1/logs", routing::post(otlp::logs)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1540,13 +1579,20 @@ impl HttpServer { } #[cfg(feature = "dashboard")] - fn route_dashboard(handler: DashboardHandlerRef) -> Router { + fn route_dashboard( + handler: DashboardHandlerRef, + memory_limiter: ServerMemoryLimiter, + ) -> Router { use crate::http::dashboard::{add_dashboard, delete_dashboard, list_dashboards}; Router::new() .route("/", routing::get(list_dashboards)) .route("/{dashboard_name}", routing::post(add_dashboard)) .route("/{dashboard_name}", routing::delete(delete_dashboard)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1555,7 +1601,10 @@ impl HttpServer { } #[cfg(not(feature = "dashboard"))] - fn route_dashboard(handler: DashboardHandlerRef) -> Router { + fn route_dashboard( + handler: DashboardHandlerRef, + _memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new().with_state(DashboardState { handler }) } } diff --git a/src/servers/src/http/loki.rs b/src/servers/src/http/loki.rs index 4e12eeb931e..161664a2ce1 100644 --- a/src/servers/src/http/loki.rs +++ b/src/servers/src/http/loki.rs @@ -108,6 +108,7 @@ pub async fn loki_ingest( TypedHeader(content_type): TypedHeader, LogTableName(table_name): LogTableName, pipeline_info: PipelineInfo, + Extension(memory_limiter): Extension, bytes: Bytes, ) -> Result { ctx.set_channel(Channel::Loki); @@ -127,6 +128,7 @@ pub async fn loki_ingest( pipeline_info, bytes, &ctx, + &memory_limiter, ) .await?; @@ -182,6 +184,7 @@ async fn build_loki_context_req( pipeline_info: PipelineInfo, bytes: Bytes, ctx: &QueryContextRef, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, ) -> Result { // A pipeline header switches Loki into the generic pipeline path; without // it, Loki writes directly to the target log table. @@ -195,10 +198,12 @@ async fn build_loki_context_req( pipeline_params: pipeline_info.pipeline_params, bytes, }; - build_loki_pipeline_context_req(handler, pipeline_req, ctx).await + build_loki_pipeline_context_req(handler, pipeline_req, ctx, memory_limiter).await } None => { - let req = build_loki_raw_insert_request(content_type, table_name, bytes)?; + let req = + build_loki_raw_insert_request(content_type, table_name, bytes, memory_limiter) + .await?; Ok(ContextReq::default_opt_with_reqs(vec![req])) } } @@ -208,6 +213,7 @@ async fn build_loki_pipeline_context_req( handler: &PipelineHandlerRef, pipeline_req: LokiPipelineContextReq, ctx: &QueryContextRef, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, ) -> Result { let LokiPipelineContextReq { content_type, @@ -223,7 +229,8 @@ async fn build_loki_pipeline_context_req( PipelineDefinition::from_name(&pipeline_name, version, None).context(PipelineSnafu)?; let pipeline_ctx = PipelineContext::new(&def, &pipeline_params, Channel::Loki); - let values = extract_item::(content_type, bytes)? + let values = extract_item::(content_type, bytes, memory_limiter) + .await? .map(|item| item.map) .collect::>(); @@ -235,14 +242,16 @@ async fn build_loki_pipeline_context_req( run_pipeline(handler, &pipeline_ctx, req, ctx, true).await } -fn build_loki_raw_insert_request( +async fn build_loki_raw_insert_request( content_type: ContentType, table_name: String, bytes: Bytes, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, ) -> Result { let mut schema_info = SchemaInfo::from_schema_list(LOKI_INIT_SCHEMAS.clone()); let mut rows = Vec::with_capacity(256); - for loki_row in extract_item::(content_type, bytes)? { + let loki_rows = extract_item::(content_type, bytes, memory_limiter).await?; + for loki_row in loki_rows { let mut row = init_row( schema_info.schema.len(), loki_row.ts, @@ -275,7 +284,11 @@ fn build_loki_raw_insert_request( /// JSON push bodies become `LokiMiddleItem`, protobuf push bodies /// become `LokiMiddleItem>`, and the generic `Into` /// conversion selects either direct-write `LokiRawItem` or pipeline `LokiPipeline`. -fn extract_item(content_type: ContentType, bytes: Bytes) -> Result>> +async fn extract_item( + content_type: ContentType, + bytes: Bytes, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, +) -> Result>> where LokiMiddleItem: Into, LokiMiddleItem>: Into, @@ -285,7 +298,9 @@ where LokiJsonParser::from_bytes(bytes)?.flat_map(|item| item.into_iter().map(|i| i.into())), )), x if x == *PB_CONTENT_TYPE => Ok(Box::new( - LokiPbParser::from_bytes(bytes)?.flat_map(|item| item.into_iter().map(|i| i.into())), + LokiPbParser::from_bytes(bytes, memory_limiter) + .await? + .flat_map(|item| item.into_iter().map(|i| i.into())), )), _ => UnsupportedContentTypeSnafu { content_type }.fail(), } @@ -554,8 +569,13 @@ pub struct LokiPbParser { } impl LokiPbParser { - pub fn from_bytes(bytes: Bytes) -> Result { - let decompressed = snappy_decompress_loki_request(&bytes)?; + pub async fn from_bytes( + bytes: Bytes, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, + ) -> Result { + // `decompressed` carries the memory permits for the decoded bytes and + // keeps them alive until the protobuf decoding below is finished. + let decompressed = snappy_decompress_loki_request(&bytes, memory_limiter).await?; let req = loki_proto::logproto::PushRequest::decode(&decompressed[..]) .context(DecodeLokiRequestSnafu)?; @@ -565,13 +585,34 @@ impl LokiPbParser { } } -fn snappy_decompress_loki_request(buf: &[u8]) -> Result> { +async fn snappy_decompress_loki_request( + buf: &[u8], + limiter: &crate::request_memory_limiter::ServerMemoryLimiter, +) -> Result { // Loki's protobuf push body is Snappy-compressed independent of HTTP // content-encoding, so keep this decode step explicit. + // + // The decoded size is validated before the output buffer is allocated: + // the raw Snappy header carries the decoded length, and without this cap + // a tiny body could declare multiple GiB of output. The declared size is + // also charged to the aggregate request-memory limiter before the + // allocation, so concurrent Loki pushes cannot bypass the quota. The + // returned buffer carries the permits, keeping the reservation alive + // until the caller is done with the decoded bytes. + let decoded_len = snap::raw::decompress_len(buf).context(DecompressSnappyLokiRequestSnafu)?; + ensure!( + decoded_len <= crate::prom_store::MAX_DECOMPRESSED_REQUEST_SIZE, + crate::error::DecompressedBodyTooLargeSnafu { + size: decoded_len as u64, + limit: crate::prom_store::MAX_DECOMPRESSED_REQUEST_SIZE as u64, + } + ); + let guard = limiter.acquire(decoded_len as u64).await?; let mut decoder = Decoder::new(); - decoder + let data = decoder .decompress_vec(buf) - .context(DecompressSnappyLokiRequestSnafu) + .context(DecompressSnappyLokiRequestSnafu)?; + Ok(crate::prom_store::ChargedBuffer::new(data, vec![guard])) } impl Iterator for LokiPbParser { @@ -819,6 +860,7 @@ mod tests { use super::*; use crate::error::Error::{DecompressSnappyLokiRequest, InvalidLokiLabels}; use crate::prom_store::snappy_compress; + use crate::request_memory_limiter::ServerMemoryLimiter; const JSON_PAYLOAD: &[u8] = br#"{ "streams": [ @@ -869,13 +911,15 @@ mod tests { assert_eq!(prost_ts_to_nano(&ts), 1731748568804293888); } - #[test] - fn test_json_direct_ingest_builds_schema_and_pads_rows() { + #[tokio::test] + async fn test_json_direct_ingest_builds_schema_and_pads_rows() { let request = build_loki_raw_insert_request( JSON_CONTENT_TYPE.clone(), "custom_loki".to_string(), Bytes::from_static(JSON_PAYLOAD), + &ServerMemoryLimiter::default(), ) + .await .unwrap(); assert_eq!(request.table_name, "custom_loki"); @@ -916,12 +960,14 @@ mod tests { assert_eq!(row_string_value(second, 5), Some("worker-0")); } - #[test] - fn test_json_pipeline_conversion_names_loki_fields() { + #[tokio::test] + async fn test_json_pipeline_conversion_names_loki_fields() { let items = extract_item::( JSON_CONTENT_TYPE.clone(), Bytes::from_static(JSON_PAYLOAD), + &ServerMemoryLimiter::default(), ) + .await .unwrap() .collect::>(); @@ -951,8 +997,8 @@ mod tests { ); } - #[test] - fn test_protobuf_parser_decodes_snappy_push_request() { + #[tokio::test] + async fn test_protobuf_parser_decodes_snappy_push_request() { let request = PushRequest { streams: vec![StreamAdapter { labels: r#"{job="api"}"#.to_string(), @@ -973,9 +1019,14 @@ mod tests { }; let bytes = snappy_compress(&request.encode_to_vec()).unwrap(); - let items = extract_item::(PB_CONTENT_TYPE.clone(), Bytes::from(bytes)) - .unwrap() - .collect::>(); + let items = extract_item::( + PB_CONTENT_TYPE.clone(), + Bytes::from(bytes), + &ServerMemoryLimiter::default(), + ) + .await + .unwrap() + .collect::>(); assert_eq!(items.len(), 1); assert_eq!(items[0].ts, 1731748568804293888); @@ -987,9 +1038,14 @@ mod tests { assert!(!items[0].structured_metadata.is_empty()); } - #[test] - fn test_protobuf_parser_rejects_invalid_snappy_payload() { - let err = match LokiPbParser::from_bytes(Bytes::from_static(b"not-snappy")) { + #[tokio::test] + async fn test_protobuf_parser_rejects_invalid_snappy_payload() { + let err = match LokiPbParser::from_bytes( + Bytes::from_static(b"not-snappy"), + &ServerMemoryLimiter::default(), + ) + .await + { Ok(_) => panic!("expected invalid snappy payload to fail"), Err(err) => err, }; @@ -997,6 +1053,51 @@ mod tests { assert!(matches!(err, DecompressSnappyLokiRequest { .. })); } + #[tokio::test] + async fn test_protobuf_push_charges_decoded_size() { + use common_memory_manager::OnExhaustedPolicy; + + // A valid snappy payload whose decoded size exceeds the 1 KiB quota. + let request = loki_proto::logproto::PushRequest { + streams: vec![loki_proto::logproto::StreamAdapter { + labels: "{job=\"quota\"}".to_string(), + entries: std::iter::repeat_with(|| loki_proto::logproto::EntryAdapter { + timestamp: Some(Timestamp { + seconds: 1, + nanos: 0, + }), + line: "x".repeat(256), + ..Default::default() + }) + .take(8) + .collect(), + ..Default::default() + }], + }; + let bytes = Bytes::from(snappy_compress(&request.encode_to_vec()).unwrap()); + + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + let err = match LokiPbParser::from_bytes(bytes.clone(), &limiter).await { + Ok(_) => panic!("expected quota exhaustion to fail"), + Err(err) => err, + }; + assert!(matches!( + err, + crate::error::Error::MemoryLimitExceeded { .. } + )); + assert_eq!(0, limiter.used_bytes(), "failed charge must be released"); + + // The same payload fits an adequately sized quota. + let limiter = ServerMemoryLimiter::new(64 * 1024, OnExhaustedPolicy::Fail); + let parser = LokiPbParser::from_bytes(bytes, &limiter).await.unwrap(); + assert_eq!(1, parser.streams.len()); + assert_eq!( + 0, + limiter.used_bytes(), + "guards must be released after decode" + ); + } + #[test] fn test_parse_loki_labels() { let mut expected = BTreeMap::new(); diff --git a/src/servers/src/http/memory_limit.rs b/src/servers/src/http/memory_limit.rs index 50b4c68f03e..8d7b99c15fa 100644 --- a/src/servers/src/http/memory_limit.rs +++ b/src/servers/src/http/memory_limit.rs @@ -13,13 +13,47 @@ // limitations under the License. //! Middleware for limiting total memory usage of concurrent HTTP request bodies. +//! +//! Admission works in two stages: +//! +//! 1. Upfront, the request's `Content-Length` (if present) is charged to the +//! shared [`ServerMemoryLimiter`], preserving the configured wait/fail +//! policy for regular requests. +//! 2. While the body is streamed, [`AccountedBody`] charges every byte beyond +//! the upfront reservation. This closes two gaps of a header-only +//! reservation: HTTP/1.1 chunked requests carry no `Content-Length` (and +//! would otherwise be admitted for free), and a client can understate the +//! header while sending a much larger body. +//! +//! Permits acquired while streaming are held by [`BodyMemoryAccounting`], +//! which is also inserted into the request extensions so they stay alive until +//! the whole request is finished (extractors collect the body into owned +//! buffers that outlive the body stream itself). +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; + +use axum::body::Body; use axum::extract::{Request, State}; use axum::middleware::Next; use axum::response::{IntoResponse, Response}; +use bytes::Bytes; +use common_memory_manager::MemoryGuard; use http::StatusCode; +use http_body::{Body as HttpBody, Frame}; +use crate::error::Result; use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; + +/// Marker inserted by [`memory_limit_middleware`] when the request arrives +/// with a non-identity `Content-Encoding` header. Route-local +/// [`decoded_body_accounting_middleware`] uses it to account the decompressed +/// body: only such requests expand during request decompression. +#[derive(Clone, Copy)] +pub(crate) struct ContentEncoded; pub async fn memory_limit_middleware( State(limiter): State, @@ -44,5 +78,569 @@ pub async fn memory_limit_middleware( } }; - next.run(req).await + let content_encoded = req + .headers() + .get(http::header::CONTENT_ENCODING) + .and_then(|v| v.to_str().ok()) + .is_some_and(|v| !v.eq_ignore_ascii_case("identity")); + + // Account whatever is actually streamed beyond the upfront reservation + // (chunked transfer has none). + let accounting = BodyMemoryAccounting::default(); + // The middleware keeps its own handle so the acquired permits stay held + // for the whole request: the body wrapper and the extensions are dropped + // once the extractors finish collecting the body, before the handler is + // done with the decoded data. + let retained_accounting = accounting.clone(); + let (mut parts, body) = req.into_parts(); + // Expose the limiter to handlers that decompress request bodies on their + // own (e.g. the Loki protobuf push) so they can charge the decoded size. + parts.extensions.insert(limiter.clone()); + parts.extensions.insert(accounting.clone()); + if content_encoded { + parts.extensions.insert(ContentEncoded); + } + let accounted = AccountedBody::new(body, limiter, content_length, accounting); + let req = Request::from_parts(parts, Body::new(accounted)); + + let response = next.run(req).await; + // `retained_accounting` stays alive until this function returns, holding + // the permits across the handler's use of the collected body. + quota_exceeded_response(&retained_accounting, response) +} + +/// Rewrites the response of a request whose body accounting hit the quota: +/// the failure aborts the body mid-stream, so the extractors reject the +/// request with a generic body error (mapped to 400). Quota exhaustion must +/// surface as 429, matching the upfront admission path, so clients can tell +/// backpressure apart from malformed input. +fn quota_exceeded_response(accounting: &BodyMemoryAccounting, response: Response) -> Response { + // The flag can only be set when the body failed before the handler ran, + // so the response is the extractor rejection and can be replaced. + if accounting.take_quota_exceeded() { + return ( + StatusCode::TOO_MANY_REQUESTS, + "Request body memory limit exceeded", + ) + .into_response(); + } + response +} + +/// Route-local counterpart of [`memory_limit_middleware`] for routes that +/// install `RequestDecompressionLayer`: it must be layered *inside* the +/// decompression layer, so the body it wraps is the **decompressed** stream. +/// +/// Only requests marked [`ContentEncoded`] by the global middleware are +/// accounted; plain (uncompressed) requests are already covered by the global +/// wire-byte accounting, and charging them here as well would double-count. +/// For decompressed requests the `Content-Length` header has been stripped by +/// the decompression layer, so the full decoded size is charged as it streams. +pub(crate) async fn decoded_body_accounting_middleware( + State(limiter): State, + req: Request, + next: Next, +) -> Response { + if req.extensions().get::().is_none() { + return next.run(req).await; + } + + let content_length = req + .headers() + .get(http::header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + .unwrap_or(0); + + let accounting = BodyMemoryAccounting::default(); + // Retain the permits for the whole request, like the global middleware. + let retained_accounting = accounting.clone(); + let (mut parts, body) = req.into_parts(); + parts.extensions.insert(accounting.clone()); + let accounted = AccountedBody::new(body, limiter, content_length, accounting); + let req = Request::from_parts(parts, Body::new(accounted)); + + let response = next.run(req).await; + // `retained_accounting` stays alive until this function returns, holding + // the permits across the handler's use of the collected body. + quota_exceeded_response(&retained_accounting, response) +} + +/// Holds the memory guards acquired while a request body is streamed. Shared +/// between the [`AccountedBody`] wrapper, the request extensions and the +/// middleware itself, so the permits are only released when the request +/// (including any collected body buffers and the handler still using them) +/// is finished. +#[derive(Clone, Default)] +struct BodyMemoryAccounting { + guards: Arc>>>, + /// Set when an incremental charge hit the quota, so the rejection can be + /// rewritten into a 429 after the middleware regains control. + quota_exceeded: Arc, +} + +impl BodyMemoryAccounting { + fn hold(&self, guard: MemoryGuard) { + self.guards.lock().unwrap().push(guard); + } + + fn mark_quota_exceeded(&self) { + self.quota_exceeded.store(true, Ordering::Release); + } + + fn take_quota_exceeded(&self) -> bool { + self.quota_exceeded.swap(false, Ordering::AcqRel) + } +} + +type AcquireFuture = + Pin>> + Send>>; + +enum ChargeOutcome { + Charged, + Pending, + Failed, +} + +/// A request body wrapper that charges the shared limiter for every byte +/// streamed beyond `pre_charged` (the upfront `Content-Length` reservation). +struct AccountedBody { + inner: Body, + limiter: ServerMemoryLimiter, + accounting: BodyMemoryAccounting, + /// Bytes of body data streamed so far. + streamed: u64, + /// Bytes of the streamed body already covered by incremental permits. + incrementally_charged: u64, + /// Bytes covered by the upfront `Content-Length` reservation. + pre_charged: u64, + /// A data frame held back while its permit is being acquired, together + /// with the number of extra bytes to charge for it. + pending: Option<(Frame, u64)>, + pending_acquire: Option, +} + +impl AccountedBody { + fn new( + inner: Body, + limiter: ServerMemoryLimiter, + pre_charged: u64, + accounting: BodyMemoryAccounting, + ) -> Self { + Self { + inner, + limiter, + accounting, + streamed: 0, + incrementally_charged: 0, + pre_charged, + pending: None, + pending_acquire: None, + } + } + + /// Bytes that still need an incremental permit for the body streamed so far. + fn uncharged(&self) -> u64 { + self.streamed + .saturating_sub(self.pre_charged) + .saturating_sub(self.incrementally_charged) + } + + /// Acquires a permit for `bytes`; only one acquisition may be in flight. + /// + /// A parked acquisition always belongs to the currently buffered frame, so + /// it must be resolved *before* any new acquisition is tried: letting the + /// fast path succeed while a waiter is parked would leave the waiter alive, + /// and its late completion would then be credited with another frame's + /// byte count (under-reserving memory relative to what is marked charged). + fn charge(&mut self, bytes: u64, cx: &mut Context<'_>) -> ChargeOutcome { + debug_assert!(bytes > 0); + if let Some(fut) = self.pending_acquire.as_mut() { + return match fut.as_mut().poll(cx) { + Poll::Ready(Ok(guard)) => { + self.accounting.hold(guard); + self.pending_acquire = None; + self.incrementally_charged += bytes; + ChargeOutcome::Charged + } + Poll::Ready(Err(_)) => { + self.pending_acquire = None; + ChargeOutcome::Failed + } + Poll::Pending => ChargeOutcome::Pending, + }; + } + // Fast path: the permit is immediately available. + if let Some(guard) = self.limiter.try_acquire(bytes) { + self.accounting.hold(guard); + self.incrementally_charged += bytes; + return ChargeOutcome::Charged; + } + let limiter = self.limiter.clone(); + let mut fut = Box::pin(async move { limiter.acquire(bytes).await }); + // Poll immediately so a ready result is not missed (returning + // `Pending` without registering progress would park the task forever). + match fut.as_mut().poll(cx) { + Poll::Ready(Ok(guard)) => { + self.accounting.hold(guard); + self.incrementally_charged += bytes; + ChargeOutcome::Charged + } + Poll::Ready(Err(_)) => ChargeOutcome::Failed, + Poll::Pending => { + self.pending_acquire = Some(fut); + ChargeOutcome::Pending + } + } + } + + /// Hands out `frame` once `bytes` of it are charged, buffering the frame + /// while the acquisition is pending. + fn hand_out( + &mut self, + frame: Frame, + bytes: u64, + cx: &mut Context<'_>, + ) -> Poll, axum::Error>>> { + if bytes == 0 { + return Poll::Ready(Some(Ok(frame))); + } + match self.charge(bytes, cx) { + ChargeOutcome::Charged => Poll::Ready(Some(Ok(frame))), + ChargeOutcome::Pending => { + self.pending = Some((frame, bytes)); + Poll::Pending + } + ChargeOutcome::Failed => { + self.accounting.mark_quota_exceeded(); + Poll::Ready(Some(Err(limit_exceeded_error()))) + } + } + } +} + +impl HttpBody for AccountedBody { + type Data = Bytes; + type Error = axum::Error; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let this = self.get_mut(); + + // Finish a pending acquisition for a buffered frame first. + if let Some((frame, bytes)) = this.pending.take() { + return this.hand_out(frame, bytes, cx); + } + + match Pin::new(&mut this.inner).poll_frame(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(None) => Poll::Ready(None), + Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))), + Poll::Ready(Some(Ok(frame))) => { + if let Some(data) = frame.data_ref() { + this.streamed += data.len() as u64; + let uncharged = this.uncharged(); + this.hand_out(frame, uncharged, cx) + } else { + // Trailers pass through uncharged. + Poll::Ready(Some(Ok(frame))) + } + } + } + } + + fn is_end_stream(&self) -> bool { + self.pending.is_none() && self.inner.is_end_stream() + } + + fn size_hint(&self) -> http_body::SizeHint { + self.inner.size_hint() + } +} + +fn limit_exceeded_error() -> axum::Error { + axum::Error::new(std::io::Error::other( + "request body exceeded the aggregate memory limit", + )) +} + +#[cfg(test)] +mod tests { + use axum::body::{Body, Bytes}; + use axum::http::{Request, StatusCode, header}; + use axum::routing::post; + use axum::{Router, middleware}; + use common_memory_manager::OnExhaustedPolicy; + use tower::ServiceExt; + + use super::memory_limit_middleware; + use crate::request_memory_limiter::ServerMemoryLimiter; + + /// Streams `body` in `chunk_size` frames without a Content-Length header + /// (like HTTP/1.1 chunked transfer). + fn chunked_request(uri: &str, body: Vec, chunk_size: usize) -> Request { + let stream = futures_util::stream::iter( + body.chunks(chunk_size) + .map(|c| Ok::<_, std::io::Error>(Bytes::copy_from_slice(c))) + .collect::>(), + ); + Request::builder() + .uri(uri) + .method("POST") + .body(Body::from_stream(stream)) + .unwrap() + } + + fn counted_body_app(limiter: ServerMemoryLimiter) -> Router { + Router::new() + .route( + "/echo", + post(|body: Bytes| async move { (StatusCode::OK, body.len().to_string()) }), + ) + .layer(middleware::from_fn_with_state( + limiter, + memory_limit_middleware, + )) + } + + #[tokio::test] + async fn test_chunked_body_larger_than_quota_is_rejected() { + // The quota cannot cover a fully streamed 64 KiB chunked body. + let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let res = app + .oneshot(chunked_request("/echo", vec![b'x'; 64 * 1024], 1024)) + .await + .unwrap(); + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "a chunked body larger than the quota must be rejected with 429, not 400" + ); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_small_chunked_body_fits() { + let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let res = app + .oneshot(chunked_request("/echo", vec![b'x'; 1024], 512)) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!(&body[..], b"1024"); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_content_length_still_admitted_upfront() { + // A Content-Length larger than the quota is rejected before the body + // is read (429), unchanged from the previous behavior. + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header(header::CONTENT_LENGTH, "4096") + .body(Body::from(vec![b'x'; 4096])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS); + } + + #[tokio::test] + async fn test_understated_content_length_is_charged_the_difference() { + // Claim 1 KiB in Content-Length but stream 16 KiB: the extra 15 KiB + // must be charged while streaming and rejected by the small quota. + let limiter = ServerMemoryLimiter::new(2 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let stream = futures_util::stream::iter(vec![Ok::<_, std::io::Error>( + Bytes::copy_from_slice(&vec![b'x'; 16 * 1024]), + )]); + let req = Request::builder() + .uri("/echo") + .method("POST") + .header(header::CONTENT_LENGTH, "1024") + .body(Body::from_stream(stream)) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_honest_content_length_is_not_double_charged() { + // With a quota large enough for the whole body, an accurate + // Content-Length must admit the request without extra charges. + let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header(header::CONTENT_LENGTH, "4096") + .body(Body::from(vec![b'x'; 4096])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + assert_eq!(0, limiter.used_bytes()); + } + + /// Mirrors the route-local stack: decoded-body accounting *inside* a + /// request-decompression layer, like the ingestion routes in `http.rs`. + fn decompressing_app(limiter: ServerMemoryLimiter) -> Router { + use tower_http::decompression::RequestDecompressionLayer; + + Router::new() + .route( + "/echo", + post(|body: Bytes| async move { (StatusCode::OK, body.len().to_string()) }), + ) + .layer(middleware::from_fn_with_state( + limiter, + super::decoded_body_accounting_middleware, + )) + .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)) + } + + fn zstd_body(decoded: &[u8]) -> Body { + let compressed = zstd::stream::encode_all(decoded, 3).unwrap(); + Body::from(compressed) + } + + #[tokio::test] + async fn test_decompressed_body_is_charged_for_content_encoded_requests() { + // The decoded size (64 KiB) exceeds the 16 KiB quota: the accounting + // layer must see the decompressed bytes and reject mid-stream. + let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail); + let app = decompressing_app(limiter.clone()); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header("content-encoding", "zstd") + .extension(super::ContentEncoded) + .body(zstd_body(&vec![b'x'; 64 * 1024])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "decoded body larger than the quota must be rejected with 429" + ); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_small_decompressed_body_fits() { + let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail); + let app = decompressing_app(limiter.clone()); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header("content-encoding", "zstd") + .extension(super::ContentEncoded) + .body(zstd_body(&vec![b'x'; 1024])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!(&body[..], b"1024", "the handler must see the decoded body"); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_pending_acquisition_is_attributed_to_its_own_bytes() { + use std::time::Duration; + + use tokio_stream::wrappers::ReceiverStream; + + // Regression test: a parked acquisition must be resolved before any + // new acquisition is tried; otherwise a late waiter completion can be + // credited with a later frame's byte count. + // + // Scenario: quota 8 KiB fully held externally; frame A (4 KiB) parks + // its acquisition; the external guard is then released so A drains; + // frame B (8 KiB) arrives. Correct behavior: B waits for its own 8 KiB + // acquisition and is rejected when it times out (only 4 KiB are free). + // Buggy behavior: the stale waiter (4 KiB) completes and is credited + // with B's 8 KiB, so B is delivered although the quota cannot cover it. + let limiter = ServerMemoryLimiter::new( + 8 * 1024, + OnExhaustedPolicy::Wait { + timeout: Duration::from_millis(500), + }, + ); + let app = counted_body_app(limiter.clone()); + + let (tx, rx) = tokio::sync::mpsc::channel::>(2); + let driver_limiter = limiter.clone(); + let driver = tokio::spawn(async move { + // Hold the whole quota so frame A's acquisition parks. + let external = driver_limiter.acquire(8 * 1024).await.unwrap(); + tx.send(Ok(Bytes::from(vec![b'a'; 4096]))).await.unwrap(); + tokio::time::sleep(Duration::from_millis(300)).await; + // Free the quota: frame A's waiter drains and A is delivered. + drop(external); + tokio::time::sleep(Duration::from_millis(300)).await; + tx.send(Ok(Bytes::from(vec![b'b'; 8192]))).await.unwrap(); + // Keep the channel open; frame B's own acquisition must time out. + tokio::time::sleep(Duration::from_millis(1500)).await; + }); + + let res = app + .oneshot( + Request::builder() + .uri("/echo") + .method("POST") + .body(Body::from_stream(ReceiverStream::new(rx))) + .unwrap(), + ) + .await + .unwrap(); + driver.await.unwrap(); + + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "frame B must wait for its own 8 KiB and time out, not borrow A's reservation" + ); + } + + #[tokio::test] + async fn test_plain_requests_are_not_charged_by_the_decoded_layer() { + // Without the ContentEncoded marker the decoded layer must be a + // no-op: plain bodies are accounted by the global wire accounting. + let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail); + let app = decompressing_app(limiter.clone()); + + // Plain body, no marker: passes through. + let req = Request::builder() + .uri("/echo") + .method("POST") + .body(Body::from(vec![b'x'; 1024])) + .unwrap(); + let res = app.clone().oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + + // Compressed body WITHOUT the marker (e.g. a route that has the + // decompression layer but no accounting): still passes through. + let req = Request::builder() + .uri("/echo") + .method("POST") + .header("content-encoding", "zstd") + .body(zstd_body(&vec![b'x'; 1024])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + } } diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 68542a6f642..05d9725f947 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -52,8 +52,9 @@ use crate::prom_remote_write::validation::PromValidationMode; use crate::prom_remote_write::{ REMOTE_WRITE_V1_VERSION, REMOTE_WRITE_V2_VERSION, decode_remote_write_request, }; -use crate::prom_store::snappy_decompress; +use crate::prom_store::{MAX_DECOMPRESSED_REQUEST_SIZE, snappy_decompress_limited}; use crate::query_handler::{PipelineHandlerRef, PromStoreProtocolHandlerRef, PromStoreResponse}; +use crate::request_memory_limiter::ServerMemoryLimiter; pub const PHYSICAL_TABLE_PARAM: &str = "physical_table"; pub const DEFAULT_ENCODING: &str = "snappy"; @@ -75,6 +76,9 @@ pub struct PromStoreState { pub prom_store_with_metric_engine: bool, pub prom_validation_mode: PromValidationMode, pub pending_rows_batcher: Option>, + /// Shared request-memory limiter used to charge decompressed remote + /// read/write bodies against the aggregate quota. + pub memory_limiter: ServerMemoryLimiter, } #[derive(Debug, Serialize, Deserialize)] @@ -146,6 +150,7 @@ async fn remote_write_v1( prom_store_with_metric_engine, prom_validation_mode, pending_rows_batcher, + memory_limiter, } = state; if let Some(response) = vm_proto_version_response(¶ms) { @@ -173,13 +178,22 @@ async fn remote_write_v1( processor.set_pipeline(pipeline_handler, query_ctx.clone(), pipeline_def); } - let mut decoded = - decode_remote_write_request(is_zstd, body, prom_validation_mode, &mut processor)?; + // Keep the decode guards alive until the batches are built: the + // `TablesBuilder` retains the decompressed buffer as its raw data. + let (mut decoded, decode_guards) = decode_remote_write_request( + is_zstd, + body, + prom_validation_mode, + &mut processor, + &memory_limiter, + ) + .await?; // Parsing borrows the decode buffer, but row building copies out of it: tag // values through `decode_string`, column names through `to_owned`, and the // borrowing `col_indexes` dies inside `as_insert_requests`. Nothing below - // references the buffer, so it need not span the write. + // references the buffer, so it need not span the write — the same goes for + // the decode guards charged against the aggregate memory quota. let req = if processor.use_pipeline { drop(decoded); processor.exec_pipeline().await? @@ -188,6 +202,7 @@ async fn remote_write_v1( drop(decoded); req }; + drop(decode_guards); let batches = into_prom_write_batches(req, query_ctx); let outcome = match write_prometheus_rows_with_progress( @@ -227,6 +242,7 @@ async fn remote_write_v2( prom_store_with_metric_engine, prom_validation_mode: _, pending_rows_batcher, + memory_limiter, } = state; if let Some(response) = vm_proto_version_response(¶ms) { @@ -240,7 +256,7 @@ async fn remote_write_v2( let (db, mut query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V2_VERSION); - let req = match decode_remote_write_v2(is_zstd, body) { + let req = match decode_remote_write_v2(is_zstd, body, &memory_limiter).await { Ok(req) => req, Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), }; @@ -840,7 +856,7 @@ pub async fn remote_read( let db = params.db.clone().unwrap_or_default(); query_ctx.set_channel(Channel::Prometheus); - let request = decode_remote_read_request(body).await?; + let request = decode_remote_read_request(body, &state.memory_limiter).await?; let query_ctx = Arc::new(query_ctx); let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_READ_ELAPSED @@ -850,8 +866,13 @@ pub async fn remote_read( state.prom_store_handler.read(request, query_ctx).await } -async fn decode_remote_read_request(body: Bytes) -> Result { - let buf = snappy_decompress(&body[..])?; +async fn decode_remote_read_request( + body: Bytes, + limiter: &ServerMemoryLimiter, +) -> Result { + // Holds the memory permits for the decompressed bytes until the protobuf + // decoding is finished. + let buf = snappy_decompress_limited(&body[..], MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await?; ReadRequest::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu) } @@ -1109,6 +1130,7 @@ mod tests { prom_store_with_metric_engine: false, prom_validation_mode: PromValidationMode::Strict, pending_rows_batcher: None, + memory_limiter: ServerMemoryLimiter::default(), } } diff --git a/src/servers/src/prom_remote_write/mod.rs b/src/servers/src/prom_remote_write/mod.rs index c9aabfb2745..06729795b5b 100644 --- a/src/servers/src/prom_remote_write/mod.rs +++ b/src/servers/src/prom_remote_write/mod.rs @@ -27,6 +27,7 @@ pub(crate) mod v2; pub mod validation; use bytes::Bytes; +use common_memory_manager::MemoryGuard; use lazy_static::lazy_static; use object_pool::Pool; use snafu::ResultExt; @@ -35,7 +36,12 @@ use crate::error; use crate::prom_remote_write::decode::{PromSeriesProcessor, PromWriteRequest}; use crate::prom_remote_write::row_builder::TablesBuilder; use crate::prom_remote_write::validation::PromValidationMode; -use crate::prom_store::{snappy_decompress, zstd_decompress}; +use crate::prom_store::{ + ChargedBuffer, MAX_DECOMPRESSED_REQUEST_SIZE, snappy_decompress_limited, + zstd_decompress_limited, +}; +use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; /// Prometheus remote write protocol versions, also used as the `version` label /// of the remote write metrics. @@ -47,43 +53,74 @@ lazy_static! { Pool::new(256, PromWriteRequest::default); } -pub fn try_decompress(is_zstd: bool, body: &[u8]) -> crate::error::Result> { - if is_zstd { - zstd_decompress(body) - } else { - snappy_decompress(body) +/// Decompresses a remote-write body with the codec selected by `is_zstd`, +/// bounded by the hard decoded-size cap and charged to the aggregate +/// request-memory limiter. +/// +/// Due to vmagent's limitation there is a chance that vmagent sends the +/// content-encoding header wrong, so decoding falls back to the other codec +/// when the first attempt fails on malformed input. Size-limit and quota +/// errors are final: the body either expanded beyond the cap under the +/// declared codec, or it could not be admitted, and the other codec cannot +/// produce a different outcome. +/// +/// see +/// see +pub(crate) async fn decompress_remote_write_body( + is_zstd: bool, + body: &[u8], + limiter: &ServerMemoryLimiter, +) -> crate::error::Result { + let decompress = |is_zstd: bool| async move { + if is_zstd { + zstd_decompress_limited(body, MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await + } else { + snappy_decompress_limited(body, MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await + } + }; + + match decompress(is_zstd).await { + Ok(buf) => Ok(buf), + Err(e) => { + if matches!( + e, + crate::error::Error::DecompressSnappyPromRemoteRequest { .. } + | crate::error::Error::DecompressZstdPromRemoteRequest { .. } + ) { + decompress(!is_zstd).await + } else { + Err(e) + } + } } } -pub fn decode_remote_write_request( +pub async fn decode_remote_write_request( is_zstd: bool, body: Bytes, prom_validation_mode: PromValidationMode, processor: &mut PromSeriesProcessor, -) -> crate::error::Result> { + limiter: &ServerMemoryLimiter, +) -> crate::error::Result<( + TablesBuilder<'static>, + Vec>, +)> { let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED .with_label_values(&["decode", REMOTE_WRITE_V1_VERSION]) .start_timer(); - // due to vmagent's limitation, there is a chance that vmagent is - // sending content type wrong so we have to apply a fallback with decoding - // the content in another method. - // - // see https://github.com/VictoriaMetrics/VictoriaMetrics/issues/5301 - // see https://github.com/GreptimeTeam/greptimedb/issues/3929 - let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { - buf - } else { - // fallback to the other compression method - try_decompress(!is_zstd, &body[..])? - }; + // Holds the memory permits for the decompressed bytes; the caller must + // keep them alive as long as the returned `TablesBuilder`, which retains + // the decompressed buffer as its raw data. + let buf = decompress_remote_write_body(is_zstd, &body[..], limiter).await?; // Decompression copied the payload out, so the compressed body is no longer needed. drop(body); - + let (data, guards) = buf.into_parts(); let mut request = PROM_WRITE_REQUEST_POOL.pull(PromWriteRequest::default); request - .decode(buf, prom_validation_mode, processor) + .decode(data, prom_validation_mode, processor) .context(error::DecodePromRemoteRequestSnafu)?; - Ok(std::mem::take(&mut request.table_data)) + let tables = std::mem::take(&mut request.table_data); + Ok((tables, guards)) } diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 77ebe5a7299..b9e34dc7f24 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -50,12 +50,13 @@ use table::requests::{ use crate::error::{self, Result}; use crate::prom_remote_write::row_builder::PromCtx; use crate::prom_remote_write::validation::validate_label_name; -use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, try_decompress}; +use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, decompress_remote_write_body}; #[allow(deprecated)] use crate::prom_store::{ DATABASE_LABEL, DATABASE_LABEL_ALT, METRIC_NAME_LABEL, PHYSICAL_TABLE_LABEL, PHYSICAL_TABLE_LABEL_ALT, SCHEMA_LABEL, }; +use crate::request_memory_limiter::ServerMemoryLimiter; use crate::row_writer::{self, TableData}; use crate::semantic::{ METRIC_TYPE_COUNTER, METRIC_TYPE_GAUGE, METRIC_TYPE_GAUGE_HISTOGRAM, METRIC_TYPE_HISTOGRAM, @@ -127,21 +128,18 @@ pub(crate) struct RemoteWriteV2WriteRequests { pub semantic_index: SemanticIndexes, } -pub(crate) fn decode_remote_write_v2( +pub(crate) async fn decode_remote_write_v2( is_zstd: bool, body: Bytes, + limiter: &ServerMemoryLimiter, ) -> Result { let decode_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED .with_label_values(&["decode", REMOTE_WRITE_V2_VERSION]) .start_timer(); - // Match the v1 decoder's VictoriaMetrics fallback: some clients may send a - // mismatched content-encoding header, so try the other compression on failure. - let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { - buf - } else { - try_decompress(!is_zstd, &body[..])? - }; + // Holds the memory permits for the decompressed bytes until the protobuf + // decoding and conversion below are finished. + let buf = decompress_remote_write_body(is_zstd, &body[..], limiter).await?; // Decompression copied the payload out, so the compressed body is no longer needed. drop(body); let request = BorrowedRequest::decode(&buf).context(error::DecodePromRemoteRequestSnafu)?; @@ -789,8 +787,9 @@ pub mod test_util { use snafu::ResultExt; use crate::error::{self, Result}; - use crate::prom_remote_write::try_decompress; + use crate::prom_remote_write::decompress_remote_write_body; use crate::prom_store::snappy_compress; + use crate::request_memory_limiter::ServerMemoryLimiter; pub fn request_with_labels_and_samples( labels: Vec<(&str, &str)>, @@ -807,14 +806,24 @@ pub mod test_util { } pub fn decode_request(is_zstd: bool, body: Bytes) -> Result { - let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { - buf - } else { - try_decompress(!is_zstd, &body[..])? - }; + let buf = block_on_decode(is_zstd, body)?; Request::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu) } + /// Runs the charged decompression on a throwaway runtime so plain `#[test]` + /// callers can stay synchronous. + fn block_on_decode(is_zstd: bool, body: Bytes) -> Result> { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async move { + decompress_remote_write_body(is_zstd, &body[..], &ServerMemoryLimiter::default()) + .await + .map(|buf| buf.to_vec()) + }) + } + pub fn write_requests( request: Request, ) -> Result<(Vec, Vec, u64, u64)> { @@ -826,7 +835,15 @@ pub mod test_util { is_zstd: bool, body: Bytes, ) -> Result<(Vec, Vec, u64, u64)> { - let requests = super::decode_remote_write_v2(is_zstd, body)?; + let requests = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(super::decode_remote_write_v2( + is_zstd, + body, + &ServerMemoryLimiter::default(), + ))?; Ok(( requests.samples.all_req().collect(), requests.histograms.all_req().collect(), @@ -939,7 +956,10 @@ mod tests { assert_eq!(decoded.timeseries[0].samples.len(), 1); assert_eq!(decoded.timeseries[0].samples[0].value, 42.0); assert_eq!(decoded.timeseries[0].metadata.as_ref().unwrap().r#type, 1); - assert_eq!(decode_remote_write_v2(true, body).unwrap().sample_count, 1); + assert_eq!( + decode_v2_on_test_runtime(true, body).unwrap().sample_count, + 1 + ); } #[test] @@ -2021,7 +2041,21 @@ mod tests { fn decode_wire(wire: &[u8]) -> Result { let body = Bytes::from(crate::prom_store::snappy_compress(wire).unwrap()); - decode_remote_write_v2(false, body) + decode_v2_on_test_runtime(false, body) + } + + /// Runs the async (charged) v2 decoder on a throwaway runtime so plain + /// `#[test]` callers can stay synchronous. + fn decode_v2_on_test_runtime(is_zstd: bool, body: Bytes) -> Result { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(decode_remote_write_v2( + is_zstd, + body, + &ServerMemoryLimiter::default(), + )) } fn decode_wire_error(wire: &[u8], name: &str) -> error::Error { @@ -2096,7 +2130,7 @@ mod tests { fn decode_test_request(request: Request) -> Result { let body = Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); - decode_remote_write_v2(false, body) + decode_v2_on_test_runtime(false, body) } fn assert_invalid(name: &str, request: Request, expected: &str) { diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index e202070d98d..ffd45b8401a 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -32,6 +32,7 @@ use arrow::datatypes::{ TimestampNanosecondType, UInt32Type, }; use common_grpc::precision::Precision; +use common_memory_manager::MemoryGuard; use common_query::prelude::{greptime_timestamp, greptime_value}; use common_recordbatch::{RecordBatch, RecordBatches}; use common_telemetry::{tracing, warn}; @@ -44,6 +45,8 @@ use snap::raw::{Decoder, Encoder}; use crate::error::{self, Result}; use crate::prom_remote_write::REMOTE_WRITE_V1_VERSION; +use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; use crate::row_writer::{self, MultiTableData}; pub const METRIC_NAME_LABEL: &str = "__name__"; @@ -602,8 +605,147 @@ pub fn to_grpc_row_insert_requests(request: &WriteRequest) -> Result<(RowInsertR Ok(multi_table_data.into_row_insert_requests()) } +/// Hard cap on the decompressed size of a handler-decompressed request body +/// (Prometheus remote read/write, Loki push). +/// +/// The HTTP body limit only bounds the *encoded* bytes on the wire. Without +/// this cap a tiny compressed payload (e.g. a five-byte raw-Snappy varint or a +/// few KiB of zstd RLE blocks) could declare/expand to multiple GiB of heap +/// before any protobuf validation runs, so the decompressed size must be +/// bounded on its own. +pub const MAX_DECOMPRESSED_REQUEST_SIZE: usize = 512 * 1024 * 1024; + +/// Size of the chunks a zstd body is decompressed and charged in. +const ZSTD_DECOMPRESS_CHUNK: usize = 256 * 1024; + +/// Decompressed bytes together with the memory permits charged for them. +/// The permits are released when the buffer is dropped, so the aggregate +/// memory accounting covers the whole lifetime of the decompressed data. +pub(crate) struct ChargedBuffer { + pub data: Vec, + guards: Vec>, +} + +impl std::ops::Deref for ChargedBuffer { + type Target = [u8]; + + fn deref(&self) -> &Self::Target { + &self.data + } +} + +impl std::fmt::Debug for ChargedBuffer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChargedBuffer") + .field("data_len", &self.data.len()) + .field("guards", &self.guards.len()) + .finish() + } +} + +impl ChargedBuffer { + /// Creates a buffer with the permits charged for it. + pub(crate) fn new(data: Vec, guards: Vec>) -> Self { + Self { data, guards } + } + + /// Splits the buffer into the raw bytes and the memory permits charged for + /// them. Hold the permits as long as the bytes are kept. + pub(crate) fn into_parts(self) -> (Vec, Vec>) { + (self.data, self.guards) + } +} + +/// Decompresses a raw Snappy body after enforcing a hard cap on the decoded +/// size and charging the aggregate request-memory limiter. +/// +/// The raw Snappy format prefixes the stream with a varint of the decoded +/// length, so the output size is known (and validated) before the output +/// buffer is allocated by `decompress_vec`. +pub(crate) async fn snappy_decompress_limited( + buf: &[u8], + max_decompressed: usize, + limiter: &ServerMemoryLimiter, +) -> Result { + let decoded_len = + snap::raw::decompress_len(buf).context(error::DecompressSnappyPromRemoteRequestSnafu)?; + ensure!( + decoded_len <= max_decompressed, + error::DecompressedBodyTooLargeSnafu { + size: decoded_len as u64, + limit: max_decompressed as u64, + } + ); + // Reserve the declared size before the allocation happens so concurrent + // decompressions cannot exceed the aggregate quota. + let guard = limiter.acquire(decoded_len as u64).await?; + let mut decoder = Decoder::new(); + let data = decoder + .decompress_vec(buf) + .context(error::DecompressSnappyPromRemoteRequestSnafu)?; + Ok(ChargedBuffer { + data, + guards: vec![guard], + }) +} + +/// Decompresses a zstd body with a hard cap on the decoded size, charging the +/// aggregate request-memory limiter incrementally while the output is +/// materialized (zstd streams don't declare their decoded size up front). +pub(crate) async fn zstd_decompress_limited( + buf: &[u8], + max_decompressed: usize, + limiter: &ServerMemoryLimiter, +) -> Result { + use std::io::Read; + + let decoder = zstd::stream::read::Decoder::new(buf) + .context(error::DecompressZstdPromRemoteRequestSnafu)?; + // Allow reading one byte past the limit so oversized streams are detected + // instead of being silently truncated. + let mut limited = decoder.take(max_decompressed as u64 + 1); + let mut data = Vec::new(); + let mut guards = Vec::new(); + let mut chunk = vec![0u8; ZSTD_DECOMPRESS_CHUNK]; + loop { + let n = limited + .read(&mut chunk) + .context(error::DecompressZstdPromRemoteRequestSnafu)?; + if n == 0 { + break; + } + ensure!( + data.len() + n <= max_decompressed, + error::DecompressedBodyTooLargeSnafu { + size: (data.len() + n) as u64, + limit: max_decompressed as u64, + } + ); + // Charge each chunk before it is appended so the aggregate quota + // bounds the memory while it is being materialized. + guards.push(limiter.acquire(n as u64).await?); + data.extend_from_slice(&chunk[..n]); + } + Ok(ChargedBuffer { data, guards }) +} + +/// Decompresses a raw Snappy body with the hard decoded-size cap enforced, +/// without charging the aggregate request-memory limiter. +/// +/// Used for response bodies (remote read replies) and tests; request bodies +/// must use [`snappy_decompress_limited`] instead so their decoded size is +/// also charged to the aggregate quota. #[inline] pub fn snappy_decompress(buf: &[u8]) -> Result> { + let decoded_len = + snap::raw::decompress_len(buf).context(error::DecompressSnappyPromRemoteRequestSnafu)?; + ensure!( + decoded_len <= MAX_DECOMPRESSED_REQUEST_SIZE, + error::DecompressedBodyTooLargeSnafu { + size: decoded_len as u64, + limit: MAX_DECOMPRESSED_REQUEST_SIZE as u64, + } + ); let mut decoder = Decoder::new(); decoder .decompress_vec(buf) @@ -618,11 +760,6 @@ pub fn snappy_compress(buf: &[u8]) -> Result> { .context(error::CompressPromRemoteRequestSnafu) } -#[inline] -pub fn zstd_decompress(buf: &[u8]) -> Result> { - zstd::stream::decode_all(buf).context(error::DecompressZstdPromRemoteRequestSnafu) -} - /// Mock timeseries for test, it is both used in servers and frontend crate /// So we present it here pub fn mock_timeseries() -> Vec { @@ -762,6 +899,7 @@ mod tests { DictionaryArray, Float64Array, StringArray, TimestampMillisecondArray, UInt32Array, }; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema, UInt32Type}; + use common_memory_manager::OnExhaustedPolicy; use common_recordbatch::DfRecordBatch; use datafusion::prelude::SessionContext; use datatypes::data_type::ConcreteDataType; @@ -1506,4 +1644,126 @@ mod tests { timeseries[1].samples ); } + + #[tokio::test] + async fn test_snappy_varint_bomb_is_rejected_before_allocation() { + // A five-byte raw Snappy header declaring a ~4 GiB decoded length. + // Regression test for the compression-bomb finding: the declared size + // must be rejected by the cap check before any buffer is allocated. + let bomb = [0xffu8, 0xff, 0xff, 0xff, 0x0f]; + let limiter = ServerMemoryLimiter::default(); + + let err = snappy_decompress_limited(&bomb, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap_err(); + assert!(matches!( + err, + crate::error::Error::DecompressedBodyTooLarge { .. } + )); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_snappy_decompress_limited_roundtrip() { + let raw = vec![b'a'; 1024 * 1024]; + let mut encoder = snap::raw::Encoder::new(); + let compressed = encoder.compress_vec(&raw).unwrap(); + let limiter = ServerMemoryLimiter::default(); + + let decoded = + snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap(); + assert_eq!(&decoded[..], &raw[..]); + } + + #[tokio::test] + async fn test_snappy_charges_and_releases_quota() { + let raw = vec![b'a'; 64 * 1024]; + let mut encoder = snap::raw::Encoder::new(); + let compressed = encoder.compress_vec(&raw).unwrap(); + let limiter = ServerMemoryLimiter::new(128 * 1024, OnExhaustedPolicy::Fail); + + let decoded = + snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap(); + // Rounding to the 1 KiB permit granularity. + assert!(limiter.used_bytes() >= 64 * 1024); + assert!(limiter.used_bytes() <= 65 * 1024); + drop(decoded); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_snappy_rejects_when_quota_exhausted() { + let raw = vec![b'a'; 64 * 1024]; + let mut encoder = snap::raw::Encoder::new(); + let compressed = encoder.compress_vec(&raw).unwrap(); + // Quota smaller than the decoded size. + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + + let err = snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap_err(); + assert!(matches!( + err, + crate::error::Error::MemoryLimitExceeded { .. } + )); + assert_eq!(0, limiter.used_bytes(), "failed charge must be released"); + } + + #[tokio::test] + async fn test_zstd_at_and_over_the_cap() { + let limiter = ServerMemoryLimiter::default(); + let small_cap = 8 * 1024; + + // Exactly at the cap: accepted. + let raw = vec![0u8; small_cap]; + let compressed = zstd::stream::encode_all(&raw[..], 3).unwrap(); + let decoded = zstd_decompress_limited(&compressed, small_cap, &limiter) + .await + .unwrap(); + assert_eq!(decoded.len(), small_cap); + + // One byte past the cap: rejected. + let raw = vec![0u8; small_cap + 1]; + let compressed = zstd::stream::encode_all(&raw[..], 3).unwrap(); + let err = zstd_decompress_limited(&compressed, small_cap, &limiter) + .await + .unwrap_err(); + assert!(matches!( + err, + crate::error::Error::DecompressedBodyTooLarge { .. } + )); + } + + #[tokio::test] + async fn test_zstd_charges_incrementally() { + // 64 KiB wire that expands to 2 MiB must be charged chunk by chunk + // against a quota that only fits the expanded size. + let raw = vec![0u8; 2 * 1024 * 1024]; + let compressed = zstd::stream::encode_all(&raw[..], 19).unwrap(); + assert!(compressed.len() < 16 * 1024); + let limiter = ServerMemoryLimiter::new(4 * 1024 * 1024, OnExhaustedPolicy::Fail); + + let decoded = zstd_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap(); + assert_eq!(decoded.len(), raw.len()); + assert!(limiter.used_bytes() >= raw.len() as u64); + drop(decoded); + assert_eq!(0, limiter.used_bytes()); + } + + #[test] + fn test_snappy_response_decompress_is_capped() { + // The response-side (uncharged) helper must enforce the cap too. + let bomb = [0xffu8, 0xff, 0xff, 0xff, 0x0f]; + let err = snappy_decompress(&bomb).unwrap_err(); + assert!(matches!( + err, + crate::error::Error::DecompressedBodyTooLarge { .. } + )); + } } diff --git a/src/servers/src/request_memory_limiter.rs b/src/servers/src/request_memory_limiter.rs index 8c3d106d9c0..ac072f51b2e 100644 --- a/src/servers/src/request_memory_limiter.rs +++ b/src/servers/src/request_memory_limiter.rs @@ -73,4 +73,9 @@ impl ServerMemoryLimiter { pub fn limit_bytes(&self) -> u64 { self.manager.limit_bytes() } + + /// Returns currently reserved (in-flight) bytes. + pub fn used_bytes(&self) -> u64 { + self.manager.used_bytes() + } } diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 62c32091d76..64f9f4c6af1 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -747,7 +747,13 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( ..Default::default() }); - let mut http_server = HttpServerBuilder::new(http_opts) + // The memory limiter must be wired before the handlers: routes capture + // the builder's limiter at wiring time (mirrors `frontend::server`). + let mut http_server = HttpServerBuilder::new(http_opts); + if let Some(limiter) = memory_limiter { + http_server = http_server.with_memory_limiter(limiter); + } + let mut http_server = http_server .with_sql_handler(instance.fe_instance().clone()) .with_log_ingest_handler(instance.fe_instance().clone(), None, None) .with_logs_handler(instance.fe_instance().clone()) @@ -762,10 +768,6 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( http_server = http_server.with_user_provider(user_provider); } - if let Some(limiter) = memory_limiter { - http_server = http_server.with_memory_limiter(limiter); - } - let http_server = http_server.build(); let app = http_server.build(http_server.make_app()).unwrap(); @@ -784,7 +786,7 @@ pub async fn setup_test_prom_app_with_frontend( store_type: StorageType, name: &str, ) -> (Router, TestGuard) { - setup_test_prom_app_with_frontend_inner(store_type, name, false).await + setup_test_prom_app_with_frontend_inner(store_type, name, false, None).await } /// Like [`setup_test_prom_app_with_frontend`] but enables the pending-rows batcher, @@ -794,13 +796,25 @@ pub async fn setup_test_prom_app_with_frontend_batched( store_type: StorageType, name: &str, ) -> (Router, TestGuard) { - setup_test_prom_app_with_frontend_inner(store_type, name, true).await + setup_test_prom_app_with_frontend_inner(store_type, name, true, None).await +} + +/// Like [`setup_test_prom_app_with_frontend`] but wires a shared request-memory +/// limiter into the server, mirroring production deployments that configure +/// `max_in_flight_write_bytes`. +pub async fn setup_test_prom_app_with_frontend_and_memory_limiter( + store_type: StorageType, + name: &str, + memory_limiter: Option, +) -> (Router, TestGuard) { + setup_test_prom_app_with_frontend_inner(store_type, name, false, memory_limiter).await } async fn setup_test_prom_app_with_frontend_inner( store_type: StorageType, name: &str, enable_batcher: bool, + memory_limiter: Option, ) -> (Router, TestGuard) { unsafe { std::env::set_var("TZ", "UTC"); @@ -849,17 +863,24 @@ async fn setup_test_prom_app_with_frontend_inner( let sql = "INSERT INTO mito(host, val, ts) VALUES (1, 1.1, 0)"; run_sql(sql, &instance).await; - let http_server = build_test_prom_server(instance.fe_instance().clone(), enable_batcher) - .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) - .build(); + let http_server = build_test_prom_server( + instance.fe_instance().clone(), + enable_batcher, + memory_limiter, + ) + .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) + .build(); let app = http_server.build(http_server.make_app()).unwrap(); (app, instance.guard) } /// Builds Prometheus HTTP routes for either a standalone or distributed frontend. +/// `memory_limiter` must be set before the handlers are wired (handlers capture +/// the limiter at wiring time, mirroring production in `frontend::server`). pub fn build_test_prom_server( frontend_ref: Arc, enable_batcher: bool, + memory_limiter: Option, ) -> HttpServerBuilder { let http_opts = HttpOptions { addr: format!("127.0.0.1:{}", ports::get_port()), @@ -886,9 +907,13 @@ pub fn build_test_prom_server( None }; assert_eq!(pending_rows_batcher.is_some(), enable_batcher); - HttpServerBuilder::new(http_opts) + let mut builder = HttpServerBuilder::new(http_opts) .with_sql_handler(frontend_ref.clone()) - .with_logs_handler(frontend_ref.clone()) + .with_logs_handler(frontend_ref.clone()); + if let Some(limiter) = memory_limiter { + builder = builder.with_memory_limiter(limiter); + } + builder .with_prom_handler( frontend_ref.clone(), Some(frontend_ref.clone()), @@ -1254,7 +1279,7 @@ pub async fn setup_pg_server_with_prom_native_histogram( let instance = setup_standalone_instance(name, store_type).await; // Prometheus remote-write HTTP app with native histograms enabled. - let http_server = build_test_prom_server(instance.fe_instance().clone(), false) + let http_server = build_test_prom_server(instance.fe_instance().clone(), false, None) .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) .build(); let app = http_server.build(http_server.make_app()).unwrap(); diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 4b78e5fe276..7cc02cfc456 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -3342,7 +3342,7 @@ async fn test_histogram_ingestion_storage_lifecycle(#[case] metric_engine: bool) let mut context = TestContext::new(builder).await; let make_router = |frontend: Arc| { let builder = if metric_engine { - build_test_prom_server(frontend.clone(), true) + build_test_prom_server(frontend.clone(), true, None) } else { HttpServerBuilder::new(HttpOptions::default()) .with_sql_handler(frontend.clone()) diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index c73f80da78e..e9cfee7e9d4 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -136,6 +136,7 @@ macro_rules! grpc_tests { test_grpc_timezone, test_grpc_tls_config, test_grpc_memory_limit, + test_grpc_compressed_memory_reservation, ); )* }; @@ -358,6 +359,104 @@ pub async fn test_grpc_message_size_limit_send(store_type: StorageType) { let _ = fe_grpc_server.shutdown().await; } +/// Transport-compressed gRPC requests reserve the worst-case decoded message +/// size (`max_recv_message_size`) against the aggregate quota *before* tonic +/// decompresses them. With a quota smaller than that bound, compressed +/// requests must fail fast with `RESOURCE_EXHAUSTED` (before decoding), while +/// uncompressed requests are still admitted through the exact post-decode +/// charge. With the default unlimited limiter, compression keeps working. +pub async fn test_grpc_compressed_memory_reservation(store_type: StorageType) { + use api::v1::query_request::Query; + use api::v1::{GreptimeRequest, QueryRequest}; + use tonic::codec::CompressionEncoding; + + let config = GrpcServerConfig { + max_recv_message_size: 1024 * 1024, + ..Default::default() + }; + // Quota smaller than the decoding bound: compressed requests cannot be + // pre-admitted. + let memory_limiter = ServerMemoryLimiter::new(512 * 1024, OnExhaustedPolicy::Fail); + let (_db, fe_grpc_server) = setup_grpc_server_with( + store_type, + "test_grpc_compressed_memory_reservation", + None, + Some(config), + Some(memory_limiter), + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + + let db = Database::new_with_dbname( + format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME), + Client::with_urls(vec![addr.clone()]), + ); + db.sql("CREATE TABLE grpc_quota(ts TIMESTAMP TIME INDEX, v DOUBLE)") + .await + .unwrap(); + + let insert_request = || { + tonic::Request::new(GreptimeRequest { + header: Some(RequestHeader { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + ..Default::default() + }), + request: Some(RequestBody::Query(QueryRequest { + query: Some(Query::Sql( + "INSERT INTO grpc_quota VALUES (1000, 1.0)".to_string(), + )), + })), + }) + }; + + // Compressed request: the pre-decode reservation (1 MiB) exceeds the + // 512 KiB quota and must be rejected before any decoding happens. + let mut compressed = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap() + .send_compressed(CompressionEncoding::Zstd); + let status = compressed.handle(insert_request()).await.unwrap_err(); + assert_eq!( + status.code(), + tonic::Code::ResourceExhausted, + "compressed request must be pre-admitted against the quota, got: {status}" + ); + + // Uncompressed request on the same server: still admitted (its exact + // serialized size fits the quota). + let mut plain = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap(); + plain.handle(insert_request()).await.unwrap(); + + let _ = fe_grpc_server.shutdown().await; + + // Default (unlimited) limiter: compressed requests keep working. + let (_db, fe_grpc_server) = setup_grpc_server_with( + store_type, + "test_grpc_compressed_memory_reservation_unlimited", + None, + Some(GrpcServerConfig::default()), + None, + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + let db = Database::new_with_dbname( + format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME), + Client::with_urls(vec![addr.clone()]), + ); + db.sql("CREATE TABLE grpc_quota(ts TIMESTAMP TIME INDEX, v DOUBLE)") + .await + .unwrap(); + let mut compressed = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap() + .send_compressed(CompressionEncoding::Zstd); + compressed.handle(insert_request()).await.unwrap(); + let _ = fe_grpc_server.shutdown().await; +} + pub async fn test_grpc_message_size_limit_recv(store_type: StorageType) { let config = GrpcServerConfig { max_recv_message_size: 10, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 780140e3500..a05c58444e3 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -89,8 +89,10 @@ use table::table_name::TableName; use tests_integration::test_util::{ MockInstanceImpl, StorageType, TestGuard, assert_wal_delta, build_test_prom_server, setup_test_http_app, setup_test_http_app_with_frontend, + setup_test_http_app_with_frontend_and_custom_options, setup_test_http_app_with_frontend_and_slow_query_threshold, setup_test_http_app_with_frontend_and_user_provider, setup_test_prom_app_with_frontend, + setup_test_prom_app_with_frontend_and_memory_limiter, setup_test_prom_app_with_frontend_batched, }; use urlencoding::encode; @@ -166,6 +168,9 @@ macro_rules! http_tests { test_promql_over_non_millisecond_physical_tables, test_prometheus_label_replace_response, test_prom_http_api, + test_prom_remote_compression_bomb_rejected, + test_prom_remote_write_decoded_body_charged, + test_influxdb_write_decoded_body_charged, test_config_api, test_dashboard_api, test_prometheus_remote_write, @@ -1725,6 +1730,185 @@ pub async fn test_prom_http_api(store_type: StorageType) { guard.remove_all().await; } +/// A tiny compressed body must not be able to demand an unbounded decoded +/// allocation on the Prometheus remote-storage endpoints. The five-byte raw +/// Snappy varint (`ff ff ff ff 0f`) declares a ~4 GiB decoded length and must +/// be rejected by the decoded-size cap. +pub async fn test_prom_remote_compression_bomb_rejected(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = setup_test_prom_app_with_frontend(store_type, "prom_remote_bomb").await; + let client = TestClient::new(app).await; + + // Remote read with the five-byte varint bomb. + let res = client + .post("/v1/prometheus/read") + .body(Vec::from([0xffu8, 0xff, 0xff, 0xff, 0x0f])) + .send() + .await; + assert_eq!(res.status(), StatusCode::BAD_REQUEST); + + // Remote write v1 with the same bomb, via both codecs (exercising the + // fallback: the zstd attempt fails on malformed input and the snappy + // fallback must reject the declared size). + for encoding in ["snappy", "zstd"] { + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", encoding) + .body(Vec::from([0xffu8, 0xff, 0xff, 0xff, 0x0f])) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::BAD_REQUEST, + "bomb with Content-Encoding {encoding} must be rejected" + ); + } + + // The server must still serve valid requests afterwards. + let write_request = WriteRequest { + timeseries: prom_store::mock_timeseries(), + ..Default::default() + }; + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + guard.remove_all().await; +} + +/// The decoded (not encoded) size of a Prometheus remote write body must be +/// charged against the aggregate request-memory quota. +pub async fn test_prom_remote_write_decoded_body_charged(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + // 4 KiB aggregate quota with fail-fast policy. + let memory_limiter = ServerMemoryLimiter::new(4096, OnExhaustedPolicy::Fail); + let (app, mut guard) = setup_test_prom_app_with_frontend_and_memory_limiter( + store_type, + "prom_remote_write_quota", + Some(memory_limiter), + ) + .await; + let client = TestClient::new(app).await; + + // A valid write request whose decoded protobuf clearly exceeds the quota + // while its Snappy wire body stays far below it. + let write_request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: "__name__".to_string(), + value: "quota_metric".to_string(), + }, + Label { + name: "payload".to_string(), + // ~8 KiB decoded; compresses to a few hundred bytes. + value: "p".repeat(8 * 1024), + }, + ], + samples: vec![Sample { + value: 1.0, + timestamp: 1000, + }], + ..Default::default() + }], + ..Default::default() + }; + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + assert!( + compressed.len() < 2048, + "wire body must stay under the quota" + ); + + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + let status = res.status(); + let body = res.text().await; + assert_eq!( + status, + StatusCode::TOO_MANY_REQUESTS, + "decoded body larger than the aggregate quota must be rejected, got body: {body}" + ); + + guard.remove_all().await; +} + +/// Decompressed ingestion bodies (routes using `RequestDecompressionLayer`) +/// must be charged against the aggregate request-memory quota as they are +/// decoded, not only their compressed wire size. +pub async fn test_influxdb_write_decoded_body_charged(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + // 4 KiB aggregate quota with fail-fast policy. + let memory_limiter = ServerMemoryLimiter::new(4096, OnExhaustedPolicy::Fail); + let http_opts = servers::http::HttpOptions { + addr: format!("127.0.0.1:{}", common_test_util::ports::get_port()), + ..Default::default() + }; + let (app, mut guard) = setup_test_http_app_with_frontend_and_custom_options( + store_type, + "test_influxdb_decoded_quota", + None, + Some(http_opts), + Some(memory_limiter), + ) + .await; + let client = TestClient::new(app).await; + + // A gzip member expanding to ~8 KiB of line protocol from a tiny wire body. + let lines: Vec = (0..200) + .map(|i| format!("decoded_quota,host=h{} v=1.0,ts={}", i % 20, 1700000000 + i)) + .collect(); + let decoded = lines.join("\n"); + assert!(decoded.len() > 4096); + let compressed = zstd::stream::encode_all(decoded.as_bytes(), 3).unwrap(); + assert!( + compressed.len() < 2048, + "wire body must stay under the quota" + ); + + // Compressed write whose decoded body exceeds the quota: rejected while + // streaming (the extractor surfaces the aborted body as 400). + let res = client + .post("/v1/influxdb/write?db=public") + .header("Content-Type", "text/plain") + .header("Content-Encoding", "zstd") + .body(compressed) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "decoded body larger than the aggregate quota must be rejected with 429" + ); + + // The server still serves a small compressed write within the quota. + let small = "decoded_quota,host=ok v=1.0,ts=1700000000"; + let compressed_small = zstd::stream::encode_all(small.as_bytes(), 3).unwrap(); + let res = client + .post("/v1/influxdb/write?db=public") + .header("Content-Type", "text/plain") + .header("Content-Encoding", "zstd") + .body(compressed_small) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::NO_CONTENT, + "small compressed write must succeed, got: {}", + res.text().await + ); + + guard.remove_all().await; +} + /// PromQL must behave identically regardless of the physical metric table's /// time index unit: the same samples are remote-written into physical tables /// pre-created with second/micro/nano time indexes (plus a millisecond @@ -3722,7 +3906,7 @@ async fn check_prometheus_remote_write_batched_skip_wal(distributed: bool, v2: b common_telemetry::init_default_ut_logging(); let mut instance = MockInstanceImpl::new(&format!("prom_bulk_skip_wal_v2_{v2}"), distributed).await; - let server = build_test_prom_server(instance.frontend(), true).build(); + let server = build_test_prom_server(instance.frontend(), true, None).build(); let client = TestClient::new(server.build(server.make_app()).unwrap()).await; write_prometheus_skip_wal_sample(&client, v2, 1000, None).await; @@ -3902,7 +4086,7 @@ async fn setup_redirecting_batched_prom_app( .with_plugin(plugins) .build() .await; - let server = build_test_prom_server(standalone.fe_instance().clone(), true) + let server = build_test_prom_server(standalone.fe_instance().clone(), true, None) .with_greptime_config_options(standalone.opts.datanode_options().to_toml().unwrap()) .build(); let client = TestClient::new(server.build(server.make_app()).unwrap()).await;