diff --git a/config/config.md b/config/config.md
index 92f192e6df6a..177642c90783 100644
--- a/config/config.md
+++ b/config/config.md
@@ -25,12 +25,14 @@
| `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. |
| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout. |
| `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. |
+| `http.max_total_body_memory` | String | Unset | Maximum total memory for all concurrent HTTP request bodies.
Set to 0 to disable the limit. Default: "0" (unlimited) |
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions |
| `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. |
| `http.prom_validation_mode` | String | `strict` | Whether to enable validation for Prometheus remote write requests.
Available options:
- strict: deny invalid UTF-8 strings (default).
- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).
- unchecked: do not valid strings. |
| `grpc` | -- | -- | The gRPC server options. |
| `grpc.bind_addr` | String | `127.0.0.1:4001` | The address to bind the gRPC server. |
| `grpc.runtime_size` | Integer | `8` | The number of server worker threads. |
+| `grpc.max_total_message_memory` | String | Unset | Maximum total memory for all concurrent gRPC request messages.
Set to 0 to disable the limit. Default: "0" (unlimited) |
| `grpc.max_connection_age` | String | Unset | The maximum connection age for gRPC connection.
The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour.
Refer to https://grpc.io/docs/guides/keepalive/ for more details. |
| `grpc.tls` | -- | -- | gRPC server TLS options, see `mysql.tls` section. |
| `grpc.tls.mode` | String | `disable` | TLS mode. |
@@ -235,6 +237,7 @@
| `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. |
| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout. |
| `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. |
+| `http.max_total_body_memory` | String | Unset | Maximum total memory for all concurrent HTTP request bodies.
Set to 0 to disable the limit. Default: "0" (unlimited) |
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions |
| `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. |
| `http.prom_validation_mode` | String | `strict` | Whether to enable validation for Prometheus remote write requests.
Available options:
- strict: deny invalid UTF-8 strings (default).
- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).
- unchecked: do not valid strings. |
@@ -242,6 +245,7 @@
| `grpc.bind_addr` | String | `127.0.0.1:4001` | The address to bind the gRPC server. |
| `grpc.server_addr` | String | `127.0.0.1:4001` | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. |
| `grpc.runtime_size` | Integer | `8` | The number of server worker threads. |
+| `grpc.max_total_message_memory` | String | Unset | Maximum total memory for all concurrent gRPC request messages.
Set to 0 to disable the limit. Default: "0" (unlimited) |
| `grpc.flight_compression` | String | `arrow_ipc` | Compression mode for frontend side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `none` |
| `grpc.max_connection_age` | String | Unset | The maximum connection age for gRPC connection.
The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour.
Refer to https://grpc.io/docs/guides/keepalive/ for more details. |
| `grpc.tls` | -- | -- | gRPC server TLS options, see `mysql.tls` section. |
diff --git a/config/frontend.example.toml b/config/frontend.example.toml
index b26d88323e49..9ffcdad54070 100644
--- a/config/frontend.example.toml
+++ b/config/frontend.example.toml
@@ -31,6 +31,10 @@ timeout = "0s"
## The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
## Set to 0 to disable limit.
body_limit = "64MB"
+## Maximum total memory for all concurrent HTTP request bodies.
+## Set to 0 to disable the limit. Default: "0" (unlimited)
+## @toml2docs:none-default
+#+ max_total_body_memory = "1GB"
## HTTP CORS support, it's turned on by default
## This allows browser to access http APIs without CORS restrictions
enable_cors = true
@@ -54,6 +58,10 @@ bind_addr = "127.0.0.1:4001"
server_addr = "127.0.0.1:4001"
## The number of server worker threads.
runtime_size = 8
+## Maximum total memory for all concurrent gRPC request messages.
+## Set to 0 to disable the limit. Default: "0" (unlimited)
+## @toml2docs:none-default
+#+ max_total_message_memory = "1GB"
## Compression mode for frontend side Arrow IPC service. Available options:
## - `none`: disable all compression
## - `transport`: only enable gRPC transport compression (zstd)
diff --git a/config/standalone.example.toml b/config/standalone.example.toml
index 5fae0f444fda..744dbbe7517d 100644
--- a/config/standalone.example.toml
+++ b/config/standalone.example.toml
@@ -36,6 +36,10 @@ timeout = "0s"
## The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
## Set to 0 to disable limit.
body_limit = "64MB"
+## Maximum total memory for all concurrent HTTP request bodies.
+## Set to 0 to disable the limit. Default: "0" (unlimited)
+## @toml2docs:none-default
+#+ max_total_body_memory = "1GB"
## HTTP CORS support, it's turned on by default
## This allows browser to access http APIs without CORS restrictions
enable_cors = true
@@ -56,6 +60,10 @@ prom_validation_mode = "strict"
bind_addr = "127.0.0.1:4001"
## The number of server worker threads.
runtime_size = 8
+## Maximum total memory for all concurrent gRPC request messages.
+## Set to 0 to disable the limit. Default: "0" (unlimited)
+## @toml2docs:none-default
+#+ max_total_message_memory = "1GB"
## The maximum connection age for gRPC connection.
## The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour.
## Refer to https://grpc.io/docs/guides/keepalive/ for more details.
diff --git a/src/flow/src/server.rs b/src/flow/src/server.rs
index 3f46203ba068..eae97756a565 100644
--- a/src/flow/src/server.rs
+++ b/src/flow/src/server.rs
@@ -490,6 +490,7 @@ impl<'a> FlownodeServiceBuilder<'a> {
let config = GrpcServerConfig {
max_recv_message_size: opts.grpc.max_recv_message_size.as_bytes() as usize,
max_send_message_size: opts.grpc.max_send_message_size.as_bytes() as usize,
+ max_total_message_memory: opts.grpc.max_total_message_memory.as_bytes() as usize,
tls: opts.grpc.tls.clone(),
max_connection_age: opts.grpc.max_connection_age,
};
diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs
index d36bdd1494f1..30ff4cec58fb 100644
--- a/src/servers/src/error.rs
+++ b/src/servers/src/error.rs
@@ -164,6 +164,18 @@ pub enum Error {
location: Location,
},
+ #[snafu(display(
+ "Too many concurrent large requests, limit: {}, request size: {} bytes",
+ limit,
+ request_size
+ ))]
+ TooManyConcurrentRequests {
+ limit: usize,
+ request_size: usize,
+ #[snafu(implicit)]
+ location: Location,
+ },
+
#[snafu(display("Invalid query: {}", reason))]
InvalidQuery {
reason: String,
@@ -729,6 +741,8 @@ impl ErrorExt for Error {
InvalidUtf8Value { .. } | InvalidHeaderValue { .. } => StatusCode::InvalidArguments,
+ TooManyConcurrentRequests { .. } => StatusCode::RuntimeResourcesExhausted,
+
ParsePromQL { source, .. } => source.status_code(),
Other { source, .. } => source.status_code(),
diff --git a/src/servers/src/grpc.rs b/src/servers/src/grpc.rs
index 2f759db2a0bb..1c479a04de36 100644
--- a/src/servers/src/grpc.rs
+++ b/src/servers/src/grpc.rs
@@ -19,6 +19,7 @@ mod database;
pub mod flight;
pub mod frontend_grpc_handler;
pub mod greptime_handler;
+pub mod memory_limit;
pub mod prom_query_gateway;
pub mod region_server;
@@ -51,6 +52,7 @@ use crate::error::{AlreadyStartedSnafu, InternalSnafu, Result, StartGrpcSnafu, T
use crate::metrics::MetricsMiddlewareLayer;
use crate::otel_arrow::{HeaderInterceptor, OtelArrowServiceHandler};
use crate::query_handler::OpenTelemetryProtocolHandlerRef;
+use crate::request_limiter::RequestMemoryLimiter;
use crate::server::Server;
use crate::tls::TlsOption;
@@ -67,6 +69,8 @@ pub struct GrpcOptions {
pub max_recv_message_size: ReadableSize,
/// Max gRPC sending(encoding) message size
pub max_send_message_size: ReadableSize,
+ /// Maximum total memory for all concurrent gRPC request messages. 0 disables the limit.
+ pub max_total_message_memory: ReadableSize,
/// Compression mode in Arrow Flight service.
pub flight_compression: FlightCompression,
pub runtime_size: usize,
@@ -116,6 +120,7 @@ impl GrpcOptions {
GrpcServerConfig {
max_recv_message_size: self.max_recv_message_size.as_bytes() as usize,
max_send_message_size: self.max_send_message_size.as_bytes() as usize,
+ max_total_message_memory: self.max_total_message_memory.as_bytes() as usize,
tls: self.tls.clone(),
max_connection_age: self.max_connection_age,
}
@@ -134,6 +139,7 @@ impl Default for GrpcOptions {
server_addr: String::new(),
max_recv_message_size: DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE,
max_send_message_size: DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE,
+ max_total_message_memory: ReadableSize(0),
flight_compression: FlightCompression::ArrowIpc,
runtime_size: 8,
tls: TlsOption::default(),
@@ -153,6 +159,7 @@ impl GrpcOptions {
server_addr: format!("127.0.0.1:{}", DEFAULT_INTERNAL_GRPC_ADDR_PORT),
max_recv_message_size: DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE,
max_send_message_size: DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE,
+ max_total_message_memory: ReadableSize(0),
flight_compression: FlightCompression::ArrowIpc,
runtime_size: 8,
tls: TlsOption::default(),
@@ -217,6 +224,7 @@ pub struct GrpcServer {
bind_addr: Option,
name: Option,
config: GrpcServerConfig,
+ memory_limiter: RequestMemoryLimiter,
}
/// Grpc Server configuration
@@ -226,6 +234,8 @@ pub struct GrpcServerConfig {
pub max_recv_message_size: usize,
// Max gRPC sending(encoding) message size
pub max_send_message_size: usize,
+ /// Maximum total memory for all concurrent gRPC request messages. 0 disables the limit.
+ pub max_total_message_memory: usize,
pub tls: TlsOption,
/// Maximum time that a channel may exist.
/// Useful when the server wants to control the reconnection of its clients.
@@ -238,6 +248,7 @@ impl Default for GrpcServerConfig {
Self {
max_recv_message_size: DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE.as_bytes() as usize,
max_send_message_size: DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE.as_bytes() as usize,
+ max_total_message_memory: 0,
tls: TlsOption::default(),
max_connection_age: None,
}
@@ -277,6 +288,11 @@ impl GrpcServer {
}
Ok(())
}
+
+ /// Get the memory limiter for monitoring current memory usage
+ pub fn memory_limiter(&self) -> &RequestMemoryLimiter {
+ &self.memory_limiter
+ }
}
pub struct HealthCheckHandler;
diff --git a/src/servers/src/grpc/builder.rs b/src/servers/src/grpc/builder.rs
index 75a0bb13c3c0..ae5c22613850 100644
--- a/src/servers/src/grpc/builder.rs
+++ b/src/servers/src/grpc/builder.rs
@@ -38,6 +38,7 @@ use crate::grpc::{GrpcServer, GrpcServerConfig};
use crate::otel_arrow::{HeaderInterceptor, OtelArrowServiceHandler};
use crate::prometheus_handler::PrometheusHandlerRef;
use crate::query_handler::OpenTelemetryProtocolHandlerRef;
+use crate::request_limiter::RequestMemoryLimiter;
use crate::tls::TlsOption;
/// Add a gRPC service (`service`) to a `builder`([RoutesBuilder]).
@@ -57,7 +58,17 @@ macro_rules! add_service {
.send_compressed(CompressionEncoding::Gzip)
.send_compressed(CompressionEncoding::Zstd);
- $builder.routes_builder_mut().add_service(service_builder);
+ // Apply memory limiter layer
+ use $crate::grpc::memory_limit::MemoryLimiterExtensionLayer;
+ let service_with_limiter = $crate::tower::ServiceBuilder::new()
+ .layer(MemoryLimiterExtensionLayer::new(
+ $builder.memory_limiter().clone(),
+ ))
+ .service(service_builder);
+
+ $builder
+ .routes_builder_mut()
+ .add_service(service_with_limiter);
};
}
@@ -73,10 +84,12 @@ pub struct GrpcServerBuilder {
HeaderInterceptor,
>,
>,
+ memory_limiter: RequestMemoryLimiter,
}
impl GrpcServerBuilder {
pub fn new(config: GrpcServerConfig, runtime: Runtime) -> Self {
+ let memory_limiter = RequestMemoryLimiter::new(config.max_total_message_memory);
Self {
name: None,
config,
@@ -84,6 +97,7 @@ impl GrpcServerBuilder {
routes_builder: RoutesBuilder::default(),
tls_config: None,
otel_arrow_service: None,
+ memory_limiter,
}
}
@@ -95,6 +109,10 @@ impl GrpcServerBuilder {
&self.runtime
}
+ pub fn memory_limiter(&self) -> &RequestMemoryLimiter {
+ &self.memory_limiter
+ }
+
pub fn name(self, name: Option) -> Self {
Self { name, ..self }
}
@@ -198,6 +216,7 @@ impl GrpcServerBuilder {
bind_addr: None,
name: self.name,
config: self.config,
+ memory_limiter: self.memory_limiter,
}
}
}
diff --git a/src/servers/src/grpc/database.rs b/src/servers/src/grpc/database.rs
index 13c328399daf..5d132c434ef4 100644
--- a/src/servers/src/grpc/database.rs
+++ b/src/servers/src/grpc/database.rs
@@ -20,11 +20,14 @@ use common_error::status_code::StatusCode;
use common_query::OutputData;
use common_telemetry::{debug, warn};
use futures::StreamExt;
+use prost::Message;
use tonic::{Request, Response, Status, Streaming};
use crate::grpc::greptime_handler::GreptimeRequestHandler;
use crate::grpc::{TonicResult, cancellation};
use crate::hint_headers;
+use crate::metrics::{METRIC_GRPC_MEMORY_USAGE_BYTES, METRIC_GRPC_REQUESTS_REJECTED_TOTAL};
+use crate::request_limiter::RequestMemoryLimiter;
pub(crate) struct DatabaseService {
handler: GreptimeRequestHandler,
@@ -48,6 +51,27 @@ impl GreptimeDatabase for DatabaseService {
"GreptimeDatabase::Handle: request from {:?} with hints: {:?}",
remote_addr, hints
);
+
+ let _guard = request
+ .extensions()
+ .get::()
+ .filter(|limiter| limiter.is_enabled())
+ .and_then(|limiter| {
+ let message_size = request.get_ref().encoded_len();
+ limiter
+ .try_acquire(message_size)
+ .map(|guard| {
+ guard.inspect(|g| {
+ METRIC_GRPC_MEMORY_USAGE_BYTES.set(g.current_usage() as i64);
+ })
+ })
+ .inspect_err(|_| {
+ METRIC_GRPC_REQUESTS_REJECTED_TOTAL.inc();
+ })
+ .transpose()
+ })
+ .transpose()?;
+
let handler = self.handler.clone();
let request_future = async move {
let request = request.into_inner();
@@ -94,6 +118,9 @@ impl GreptimeDatabase for DatabaseService {
"GreptimeDatabase::HandleRequests: request from {:?} with hints: {:?}",
remote_addr, hints
);
+
+ let limiter = request.extensions().get::().cloned();
+
let handler = self.handler.clone();
let request_future = async move {
let mut affected_rows = 0;
@@ -101,6 +128,25 @@ impl GreptimeDatabase for DatabaseService {
let mut stream = request.into_inner();
while let Some(request) = stream.next().await {
let request = request?;
+
+ let _guard = limiter
+ .as_ref()
+ .filter(|limiter| limiter.is_enabled())
+ .and_then(|limiter| {
+ let message_size = request.encoded_len();
+ limiter
+ .try_acquire(message_size)
+ .map(|guard| {
+ guard.inspect(|g| {
+ METRIC_GRPC_MEMORY_USAGE_BYTES.set(g.current_usage() as i64);
+ })
+ })
+ .inspect_err(|_| {
+ METRIC_GRPC_REQUESTS_REJECTED_TOTAL.inc();
+ })
+ .transpose()
+ })
+ .transpose()?;
let output = handler.handle_request(request, hints.clone()).await?;
match output.data {
OutputData::AffectedRows(rows) => affected_rows += rows,
diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs
index bb431bfdaeb0..44b307fe717f 100644
--- a/src/servers/src/grpc/flight.rs
+++ b/src/servers/src/grpc/flight.rs
@@ -45,6 +45,8 @@ use crate::error::{InvalidParameterSnafu, ParseJsonSnafu, Result, ToJsonSnafu};
pub use crate::grpc::flight::stream::FlightRecordBatchStream;
use crate::grpc::greptime_handler::{GreptimeRequestHandler, get_request_type};
use crate::grpc::{FlightCompression, TonicResult, context_auth};
+use crate::metrics::{METRIC_GRPC_MEMORY_USAGE_BYTES, METRIC_GRPC_REQUESTS_REJECTED_TOTAL};
+use crate::request_limiter::{RequestMemoryGuard, RequestMemoryLimiter};
use crate::{error, hint_headers};
pub type TonicStream = Pin> + Send + 'static>>;
@@ -211,7 +213,9 @@ impl FlightCraft for GreptimeRequestHandler {
&self,
request: Request>,
) -> TonicResult>> {
- let (headers, _, stream) = request.into_parts();
+ let (headers, extensions, stream) = request.into_parts();
+
+ let limiter = extensions.get::().cloned();
let query_ctx = context_auth::create_query_context_from_grpc_metadata(&headers)?;
context_auth::check_auth(self.user_provider.clone(), &headers, query_ctx.clone()).await?;
@@ -225,6 +229,7 @@ impl FlightCraft for GreptimeRequestHandler {
query_ctx.current_catalog().to_string(),
query_ctx.current_schema(),
),
+ limiter,
};
self.put_record_batches(stream, tx, query_ctx).await;
@@ -248,10 +253,15 @@ pub(crate) struct PutRecordBatchRequest {
pub(crate) table_name: TableName,
pub(crate) request_id: i64,
pub(crate) data: FlightData,
+ pub(crate) _guard: Option,
}
impl PutRecordBatchRequest {
- fn try_new(table_name: TableName, flight_data: FlightData) -> Result {
+ fn try_new(
+ table_name: TableName,
+ flight_data: FlightData,
+ limiter: Option<&RequestMemoryLimiter>,
+ ) -> Result {
let request_id = if !flight_data.app_metadata.is_empty() {
let metadata: DoPutMetadata =
serde_json::from_slice(&flight_data.app_metadata).context(ParseJsonSnafu)?;
@@ -259,10 +269,30 @@ impl PutRecordBatchRequest {
} else {
0
};
+
+ let _guard = limiter
+ .filter(|limiter| limiter.is_enabled())
+ .map(|limiter| {
+ let message_size = flight_data.encoded_len();
+ limiter
+ .try_acquire(message_size)
+ .map(|guard| {
+ guard.inspect(|g| {
+ METRIC_GRPC_MEMORY_USAGE_BYTES.set(g.current_usage() as i64);
+ })
+ })
+ .inspect_err(|_| {
+ METRIC_GRPC_REQUESTS_REJECTED_TOTAL.inc();
+ })
+ })
+ .transpose()?
+ .flatten();
+
Ok(Self {
table_name,
request_id,
data: flight_data,
+ _guard,
})
}
}
@@ -270,6 +300,7 @@ impl PutRecordBatchRequest {
pub(crate) struct PutRecordBatchRequestStream {
flight_data_stream: Streaming,
state: PutRecordBatchRequestStreamState,
+ limiter: Option,
}
enum PutRecordBatchRequestStreamState {
@@ -298,6 +329,7 @@ impl Stream for PutRecordBatchRequestStream {
}
let poll = ready!(self.flight_data_stream.poll_next_unpin(cx));
+ let limiter = self.limiter.clone();
let result = match &mut self.state {
PutRecordBatchRequestStreamState::Init(catalog, schema) => match poll {
@@ -311,8 +343,11 @@ impl Stream for PutRecordBatchRequestStream {
Err(e) => return Poll::Ready(Some(Err(e.into()))),
};
- let request =
- PutRecordBatchRequest::try_new(table_name.clone(), flight_data);
+ let request = PutRecordBatchRequest::try_new(
+ table_name.clone(),
+ flight_data,
+ limiter.as_ref(),
+ );
let request = match request {
Ok(request) => request,
Err(e) => return Poll::Ready(Some(Err(e.into()))),
@@ -333,8 +368,12 @@ impl Stream for PutRecordBatchRequestStream {
},
PutRecordBatchRequestStreamState::Started(table_name) => poll.map(|x| {
x.and_then(|flight_data| {
- PutRecordBatchRequest::try_new(table_name.clone(), flight_data)
- .map_err(Into::into)
+ PutRecordBatchRequest::try_new(
+ table_name.clone(),
+ flight_data,
+ limiter.as_ref(),
+ )
+ .map_err(Into::into)
})
}),
};
diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs
index e19fc4352b78..095c36abb197 100644
--- a/src/servers/src/grpc/greptime_handler.rs
+++ b/src/servers/src/grpc/greptime_handler.rs
@@ -160,6 +160,7 @@ impl GreptimeRequestHandler {
table_name,
request_id,
data,
+ _guard,
} = request;
let timer = metrics::GRPC_BULK_INSERT_ELAPSED.start_timer();
diff --git a/src/servers/src/grpc/memory_limit.rs b/src/servers/src/grpc/memory_limit.rs
new file mode 100644
index 000000000000..a3dee9da575a
--- /dev/null
+++ b/src/servers/src/grpc/memory_limit.rs
@@ -0,0 +1,72 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::task::{Context, Poll};
+
+use futures::future::BoxFuture;
+use tonic::server::NamedService;
+use tower::{Layer, Service};
+
+use crate::request_limiter::RequestMemoryLimiter;
+
+#[derive(Clone)]
+pub struct MemoryLimiterExtensionLayer {
+ limiter: RequestMemoryLimiter,
+}
+
+impl MemoryLimiterExtensionLayer {
+ pub fn new(limiter: RequestMemoryLimiter) -> Self {
+ Self { limiter }
+ }
+}
+
+impl Layer for MemoryLimiterExtensionLayer {
+ type Service = MemoryLimiterExtensionService;
+
+ fn layer(&self, service: S) -> Self::Service {
+ MemoryLimiterExtensionService {
+ inner: service,
+ limiter: self.limiter.clone(),
+ }
+ }
+}
+
+#[derive(Clone)]
+pub struct MemoryLimiterExtensionService {
+ inner: S,
+ limiter: RequestMemoryLimiter,
+}
+
+impl NamedService for MemoryLimiterExtensionService {
+ const NAME: &'static str = S::NAME;
+}
+
+impl Service> for MemoryLimiterExtensionService
+where
+ S: Service>,
+ S::Future: Send + 'static,
+{
+ type Response = S::Response;
+ type Error = S::Error;
+ type Future = BoxFuture<'static, Result>;
+
+ fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> {
+ self.inner.poll_ready(cx)
+ }
+
+ fn call(&mut self, mut req: http::Request) -> Self::Future {
+ req.extensions_mut().insert(self.limiter.clone());
+ Box::pin(self.inner.call(req))
+ }
+}
diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs
index e373e0a78050..68fb0f04e906 100644
--- a/src/servers/src/http.rs
+++ b/src/servers/src/http.rs
@@ -82,6 +82,7 @@ use crate::query_handler::{
OpenTelemetryProtocolHandlerRef, OpentsdbProtocolHandlerRef, PipelineHandlerRef,
PromStoreProtocolHandlerRef,
};
+use crate::request_limiter::RequestMemoryLimiter;
use crate::server::Server;
pub mod authorize;
@@ -97,6 +98,7 @@ pub mod jaeger;
pub mod logs;
pub mod loki;
pub mod mem_prof;
+mod memory_limit;
pub mod opentsdb;
pub mod otlp;
pub mod pprof;
@@ -129,6 +131,7 @@ pub struct HttpServer {
router: StdMutex,
shutdown_tx: Mutex