fix: bound decompressed request size and close memory-admission gaps for compressed requests (#9264)

* fix: bound decompressed request size and close memory-admission gaps for compressed requests

The request-memory accounting only bounds and charges the encoded bytes on
the wire, but a compressed body can expand far beyond that during
decompression, so tiny requests could allocate disproportionate frontend
memory before any protobuf validation or quota charge.

- Handler-level decompression (Prometheus remote read/write v1+v2, Loki)
  now enforces a hard 512 MiB decoded-size cap, checked before any output
  buffer is allocated, and charges the decoded bytes to the shared
  ServerMemoryLimiter, holding the permits for the lifetime of the
  decompressed buffer.
- gRPC requests with transport compression reserve the configured
  max_recv_message_size before tonic decompresses, so the decoding phase
  is admitted against max_in_flight_write_bytes; the later per-message
  charge is skipped to avoid double accounting.
- The HTTP memory-limit middleware keeps its upfront Content-Length charge
  but now also accounts the bytes actually streamed beyond it, so chunked
  requests and understated Content-Length headers no longer bypass the
  aggregate quota.
- Routes that decompress request bodies via RequestDecompressionLayer
  (InfluxDB, OTLP, Loki, Splunk, Elasticsearch, pipelines, dashboards)
  now charge the decompressed bytes as handlers consume them: the global
  middleware marks Content-Encoding requests, and a route-local
  accounting layer inside the decompression layer charges the decoded
  stream. Plain requests are skipped to avoid double-counting.

Signed-off-by: Ning Sun <sunning@greptime.com>

* fix: address review comments on memory admission guard lifetimes and 429 mapping

- Retain the gRPC pre-decode reservation for the whole request: the
  extensions holding the guard were dropped when the request was consumed
  (into_inner), releasing the reservation while per-message charges
  stayed skipped. Both the unary and streaming handlers now clone and
  hold the reservation marker for the duration of request handling.
- Retain the HTTP body permits across the handler: the AccountedBody
  wrapper and the request extensions are dropped once the extractors
  finish collecting the body, before the handler is done with the decoded
  data. Both accounting middlewares now keep their own accounting handle
  alive across next.run(req).await.
- Return a ChargedBuffer from the Loki snappy decompressor so the
  reservation outlives the decompressed bytes through the caller's
  protobuf decoding, matching the Prometheus path.
- Map mid-stream quota exhaustion to 429 instead of the generic body
  error (400): the accounting flags exhaustion and the middlewares
  rewrite the extractor rejection, so clients can distinguish
  backpressure from malformed input.

Signed-off-by: Ning Sun <sunning@greptime.com>

* fix: resolve the in-flight body acquisition before trying a new one

A parked acquisition in AccountedBody::charge always belongs to the
currently buffered frame, so it must be resolved before any new
acquisition is tried. Letting the fast-path try_acquire succeed while a
waiter is parked left the waiter alive, and its late completion would
then be credited with a later frame's byte count, under-reserving memory
relative to what was marked charged.

Adds a regression test that reproduces the misattribution: with the
buggy ordering the test delivers a body the quota cannot cover; with the
fix the over-quota frame waits for its own acquisition and times out.

Signed-off-by: Ning Sun <sunning@greptime.com>

---------

Signed-off-by: Ning Sun <sunning@greptime.com>
(cherry picked from commit 3131cdbcb6)

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",

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

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 { .. }

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<OtelArrowServiceHandler<OpenTelemetryProtocolHandlerRef>>,
+        HeaderInterceptor,
+    >,
+>;
+
 pub struct GrpcServer {
     // states
     shutdown_tx: Mutex<Option<Sender<()>>>,
@@ -224,14 +233,7 @@ pub struct GrpcServer {
     // tls config
     tls_config: Option<ServerTlsConfig>,
     // Otel arrow service
-    otel_arrow_service: Mutex<
-        Option<
-            InterceptedService<
-                ArrowMetricsServiceServer<OtelArrowServiceHandler<OpenTelemetryProtocolHandlerRef>>,
-                HeaderInterceptor,
-            >,
-        >,
-    >,
+    otel_arrow_service: Mutex<Option<OtelArrowService>>,
     bind_addr: Option<SocketAddr>,
     name: Option<String>,
     config: GrpcServerConfig,

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<ServerTlsConfig>,
     otel_arrow_service: Option<
-        InterceptedService<
-            ArrowMetricsServiceServer<OtelArrowServiceHandler<OpenTelemetryProtocolHandlerRef>>,
-            HeaderInterceptor,
+        MemoryLimiterExtensionService<
+            InterceptedService<
+                ArrowMetricsServiceServer<OtelArrowServiceHandler<OpenTelemetryProtocolHandlerRef>>,
+                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
     }

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::<ServerMemoryLimiter>() {
+        // 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::<PreDecodeMemoryReservation>()
+            .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::<ServerMemoryLimiter>() {
             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::<ServerMemoryLimiter>().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::<PreDecodeMemoryReservation>()
+            .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 {

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<MemoryGuard<RequestMemoryMetrics>>,
+}
 
 #[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<S> Layer<S> for MemoryLimiterExtensionLayer {
         MemoryLimiterExtensionService {
             inner: service,
             limiter: self.limiter.clone(),
+            max_decoding_message_size: self.max_decoding_message_size,
         }
     }
 }
@@ -46,27 +94,199 @@ impl<S> Layer<S> for MemoryLimiterExtensionLayer {
 pub struct MemoryLimiterExtensionService<S> {
     inner: S,
     limiter: ServerMemoryLimiter,
+    max_decoding_message_size: usize,
 }
 
 impl<S: NamedService> NamedService for MemoryLimiterExtensionService<S> {
     const NAME: &'static str = S::NAME;
 }
 
-impl<S, ReqBody> Service<http::Request<ReqBody>> for MemoryLimiterExtensionService<S>
+impl<S, ReqBody> Service<Request<ReqBody>> for MemoryLimiterExtensionService<S>
 where
-    S: Service<http::Request<ReqBody>>,
+    S: Service<Request<ReqBody>, 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<Self::Response, Self::Error>>;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         self.inner.poll_ready(cx)
     }
 
-    fn call(&mut self, mut req: http::Request<ReqBody>) -> Self::Future {
-        req.extensions_mut().insert(self.limiter.clone());
-        Box::pin(self.inner.call(req))
+    fn call(&mut self, mut req: Request<ReqBody>) -> 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::<axum::body::Body>()
+                        .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<Request<Body>> for EchoService {
+        type Response = axum::response::Response;
+        type Error = Infallible;
+        type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+        fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn call(&mut self, req: Request<Body>) -> Self::Future {
+            let saw_reservation = req
+                .extensions()
+                .get::<PreDecodeMemoryReservation>()
+                .is_some();
+            let saw_limiter = req.extensions().get::<ServerMemoryLimiter>().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"[..]);
     }
 }

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<S>(log_state: LogState) -> Router<S> {
+    fn route_loki<S>(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router<S> {
         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<S>(log_state: LogState) -> Router<S> {
+    fn route_splunk<S>(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router<S> {
         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<S>(log_state: LogState) -> Router<S> {
+    fn route_elasticsearch<S>(
+        log_state: LogState,
+        memory_limiter: ServerMemoryLimiter,
+    ) -> Router<S> {
         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<S>(log_state: LogState) -> Router<S> {
+    fn route_log_deprecated<S>(
+        log_state: LogState,
+        memory_limiter: ServerMemoryLimiter,
+    ) -> Router<S> {
         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<S>(log_state: LogState) -> Router<S> {
+    fn route_pipelines<S>(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router<S> {
         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<S>(influxdb_handler: InfluxdbLineProtocolHandlerRef) -> Router<S> {
+    fn route_influxdb<S>(
+        influxdb_handler: InfluxdbLineProtocolHandlerRef,
+        memory_limiter: ServerMemoryLimiter,
+    ) -> Router<S> {
         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<S>(
         otlp_handler: OpenTelemetryProtocolHandlerRef,
         with_metric_engine: bool,
+        memory_limiter: ServerMemoryLimiter,
     ) -> Router<S> {
         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<S>(handler: DashboardHandlerRef) -> Router<S> {
+    fn route_dashboard<S>(
+        handler: DashboardHandlerRef,
+        memory_limiter: ServerMemoryLimiter,
+    ) -> Router<S> {
         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<S>(handler: DashboardHandlerRef) -> Router<S> {
+    fn route_dashboard<S>(
+        handler: DashboardHandlerRef,
+        _memory_limiter: ServerMemoryLimiter,
+    ) -> Router<S> {
         Router::new().with_state(DashboardState { handler })
     }
 }

src/servers/src/http/loki.rs

@@ -108,6 +108,7 @@ pub async fn loki_ingest(
     TypedHeader(content_type): TypedHeader<ContentType>,
     LogTableName(table_name): LogTableName,
     pipeline_info: PipelineInfo,
+    Extension(memory_limiter): Extension<crate::request_memory_limiter::ServerMemoryLimiter>,
     bytes: Bytes,
 ) -> Result<HttpResponse> {
     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<ContextReq> {
     // 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<ContextReq> {
     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::<LokiPipeline>(content_type, bytes)?
+    let values = extract_item::<LokiPipeline>(content_type, bytes, memory_limiter)
+        .await?
         .map(|item| item.map)
         .collect::<Vec<_>>();
 
@@ -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<RowInsertRequest> {
     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::<LokiRawItem>(content_type, bytes)? {
+    let loki_rows = extract_item::<LokiRawItem>(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<VrlValue>`, protobuf push bodies
 /// become `LokiMiddleItem<Vec<LabelPairAdapter>>`, and the generic `Into<T>`
 /// conversion selects either direct-write `LokiRawItem` or pipeline `LokiPipeline`.
-fn extract_item<T>(content_type: ContentType, bytes: Bytes) -> Result<Box<dyn Iterator<Item = T>>>
+async fn extract_item<T>(
+    content_type: ContentType,
+    bytes: Bytes,
+    memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter,
+) -> Result<Box<dyn Iterator<Item = T>>>
 where
     LokiMiddleItem<VrlValue>: Into<T>,
     LokiMiddleItem<Vec<LabelPairAdapter>>: Into<T>,
@@ -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<Self> {
-        let decompressed = snappy_decompress_loki_request(&bytes)?;
+    pub async fn from_bytes(
+        bytes: Bytes,
+        memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter,
+    ) -> Result<Self> {
+        // `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<Vec<u8>> {
+async fn snappy_decompress_loki_request(
+    buf: &[u8],
+    limiter: &crate::request_memory_limiter::ServerMemoryLimiter,
+) -> Result<crate::prom_store::ChargedBuffer> {
     // 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::<LokiPipeline>(
             JSON_CONTENT_TYPE.clone(),
             Bytes::from_static(JSON_PAYLOAD),
+            &ServerMemoryLimiter::default(),
         )
+        .await
         .unwrap()
         .collect::<Vec<_>>();
 
@@ -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::<LokiRawItem>(PB_CONTENT_TYPE.clone(), Bytes::from(bytes))
-            .unwrap()
-            .collect::<Vec<_>>();
+        let items = extract_item::<LokiRawItem>(
+            PB_CONTENT_TYPE.clone(),
+            Bytes::from(bytes),
+            &ServerMemoryLimiter::default(),
+        )
+        .await
+        .unwrap()
+        .collect::<Vec<_>>();
 
         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();

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<ServerMemoryLimiter>,
@@ -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<ServerMemoryLimiter>,
+    req: Request,
+    next: Next,
+) -> Response {
+    if req.extensions().get::<ContentEncoded>().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::<u64>().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<Mutex<Vec<MemoryGuard<RequestMemoryMetrics>>>>,
+    /// 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<AtomicBool>,
+}
+
+impl BodyMemoryAccounting {
+    fn hold(&self, guard: MemoryGuard<RequestMemoryMetrics>) {
+        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<Box<dyn std::future::Future<Output = Result<MemoryGuard<RequestMemoryMetrics>>> + 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<Bytes>, u64)>,
+    pending_acquire: Option<AcquireFuture>,
+}
+
+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>,
+        bytes: u64,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Bytes>, 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<Option<Result<Frame<Self::Data>, 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<u8>, chunk_size: usize) -> Request<Body> {
+        let stream = futures_util::stream::iter(
+            body.chunks(chunk_size)
+                .map(|c| Ok::<_, std::io::Error>(Bytes::copy_from_slice(c)))
+                .collect::<Vec<_>>(),
+        );
+        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::<Result<Bytes, std::io::Error>>(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);
+    }
 }

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<Arc<LogicalTablePendingRowsBatcher>>,
+    /// 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(&params) {
@@ -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(&params) {
@@ -240,7 +256,7 @@ async fn remote_write_v2(
     let (db, mut query_ctx, _timer) =
         prepare_remote_write_context(&params, 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<ReadRequest> {
-    let buf = snappy_decompress(&body[..])?;
+async fn decode_remote_read_request(
+    body: Bytes,
+    limiter: &ServerMemoryLimiter,
+) -> Result<ReadRequest> {
+    // 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(),
         }
     }
 

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<Vec<u8>> {
-    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 <https://github.com/VictoriaMetrics/VictoriaMetrics/issues/5301>
+/// see <https://github.com/GreptimeTeam/greptimedb/issues/3929>
+pub(crate) async fn decompress_remote_write_body(
+    is_zstd: bool,
+    body: &[u8],
+    limiter: &ServerMemoryLimiter,
+) -> crate::error::Result<ChargedBuffer> {
+    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<TablesBuilder<'static>> {
+    limiter: &ServerMemoryLimiter,
+) -> crate::error::Result<(
+    TablesBuilder<'static>,
+    Vec<MemoryGuard<RequestMemoryMetrics>>,
+)> {
     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))
 }

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<RemoteWriteV2WriteRequests> {
     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<Request> {
-        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<Vec<u8>> {
+        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<RowInsertRequest>, Vec<RowInsertRequest>, u64, u64)> {
@@ -826,7 +835,15 @@ pub mod test_util {
         is_zstd: bool,
         body: Bytes,
     ) -> Result<(Vec<RowInsertRequest>, Vec<RowInsertRequest>, 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<RemoteWriteV2WriteRequests> {
         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<RemoteWriteV2WriteRequests> {
+        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<RemoteWriteV2WriteRequests> {
         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) {

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<u8>,
+    guards: Vec<MemoryGuard<RequestMemoryMetrics>>,
+}
+
+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<u8>, guards: Vec<MemoryGuard<RequestMemoryMetrics>>) -> 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<u8>, Vec<MemoryGuard<RequestMemoryMetrics>>) {
+        (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<ChargedBuffer> {
+    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<ChargedBuffer> {
+    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<Vec<u8>> {
+    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<Vec<u8>> {
         .context(error::CompressPromRemoteRequestSnafu)
 }
 
-#[inline]
-pub fn zstd_decompress(buf: &[u8]) -> Result<Vec<u8>> {
-    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<TimeSeries> {
@@ -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 { .. }
+        ));
+    }
 }

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()
+    }
 }

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<ServerMemoryLimiter>,
+) -> (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<ServerMemoryLimiter>,
 ) -> (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<Instance>,
     enable_batcher: bool,
+    memory_limiter: Option<ServerMemoryLimiter>,
 ) -> 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();

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<Instance>| {
         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())

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,

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<String> = (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;