feat: introduct Limiter to limit in-flight write bytes size in frontend

src/frontend/src/error.rs

@@ -321,6 +321,12 @@ pub enum Error {
         location: Location,
         source: BoxedError,
     },
+
+    #[snafu(display("In-flight write bytes exceeded"))]
+    InFlightWriteBytesExceeded {
+        #[snafu(implicit)]
+        location: Location,
+    },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -392,6 +398,8 @@ impl ErrorExt for Error {
             Error::StartScriptManager { source, .. } => source.status_code(),
 
             Error::TableOperation { source, .. } => source.status_code(),
+
+            Error::InFlightWriteBytesExceeded { .. } => StatusCode::RateLimited,
         }
     }
 

src/frontend/src/frontend.rs

@@ -46,6 +46,7 @@ pub struct FrontendOptions {
     pub user_provider: Option<String>,
     pub export_metrics: ExportMetricsOption,
     pub tracing: TracingOptions,
+    pub max_in_flight_write_bytes: u64,
 }
 
 impl Default for FrontendOptions {
@@ -68,6 +69,7 @@ impl Default for FrontendOptions {
             user_provider: None,
             export_metrics: ExportMetricsOption::default(),
             tracing: TracingOptions::default(),
+            max_in_flight_write_bytes: u64::MAX,
         }
     }
 }

src/frontend/src/instance.rs

@@ -86,6 +86,7 @@ use crate::error::{
 };
 use crate::frontend::FrontendOptions;
 use crate::heartbeat::HeartbeatTask;
+use crate::limiter::LimiterRef;
 use crate::script::ScriptExecutor;
 
 #[async_trait]
@@ -124,6 +125,7 @@ pub struct Instance {
     export_metrics_task: Option<ExportMetricsTask>,
     table_metadata_manager: TableMetadataManagerRef,
     stats: StatementStatistics,
+    limiter: LimiterRef,
 }
 
 impl Instance {

src/frontend/src/instance/builder.rs

@@ -43,6 +43,7 @@ use crate::frontend::FrontendOptions;
 use crate::heartbeat::HeartbeatTask;
 use crate::instance::region_query::FrontendRegionQueryHandler;
 use crate::instance::Instance;
+use crate::limiter::Limiter;
 use crate::script::ScriptExecutor;
 
 /// The frontend [`Instance`] builder.
@@ -196,6 +197,7 @@ impl FrontendBuilder {
 
         plugins.insert::<StatementExecutorRef>(statement_executor.clone());
 
+        let limiter = Arc::new(Limiter::new(self.options.max_in_flight_write_bytes));
         Ok(Instance {
             options: self.options,
             catalog_manager: self.catalog_manager,
@@ -211,6 +213,7 @@ impl FrontendBuilder {
             export_metrics_task: None,
             table_metadata_manager: Arc::new(TableMetadataManager::new(kv_backend)),
             stats: self.stats,
+            limiter,
         })
     }
 }

src/frontend/src/instance/grpc.rs

@@ -29,8 +29,8 @@ use snafu::{ensure, OptionExt, ResultExt};
 use table::table_name::TableName;
 
 use crate::error::{
-    Error, IncompleteGrpcRequestSnafu, NotSupportedSnafu, PermissionSnafu, Result,
-    TableOperationSnafu,
+    Error, InFlightWriteBytesExceededSnafu, IncompleteGrpcRequestSnafu, NotSupportedSnafu,
+    PermissionSnafu, Result, TableOperationSnafu,
 };
 use crate::instance::{attach_timer, Instance};
 use crate::metrics::{GRPC_HANDLE_PROMQL_ELAPSED, GRPC_HANDLE_SQL_ELAPSED};
@@ -50,6 +50,14 @@ impl GrpcQueryHandler for Instance {
             .check_permission(ctx.current_user(), PermissionReq::GrpcRequest(&request))
             .context(PermissionSnafu)?;
 
+        if self.limiter.should_be_limited() {
+            return InFlightWriteBytesExceededSnafu.fail();
+        }
+
+        self.limiter
+            .in_flight_write_bytes_counter()
+            .process_with_request(&request);
+
         let output = match request {
             Request::Inserts(requests) => self.handle_inserts(requests, ctx.clone()).await?,
             Request::RowInserts(requests) => self.handle_row_inserts(requests, ctx.clone()).await?,

src/frontend/src/instance/influxdb.rs

@@ -16,7 +16,7 @@ use async_trait::async_trait;
 use auth::{PermissionChecker, PermissionCheckerRef, PermissionReq};
 use client::Output;
 use common_error::ext::BoxedError;
-use servers::error::{AuthSnafu, Error};
+use servers::error::{AuthSnafu, Error, InFlightWriteBytesExceededSnafu};
 use servers::influxdb::InfluxdbRequest;
 use servers::interceptor::{LineProtocolInterceptor, LineProtocolInterceptorRef};
 use servers::query_handler::InfluxdbLineProtocolHandler;
@@ -46,6 +46,14 @@ impl InfluxdbLineProtocolHandler for Instance {
             .post_lines_conversion(requests, ctx.clone())
             .await?;
 
+        if self.limiter.should_be_limited() {
+            return InFlightWriteBytesExceededSnafu.fail();
+        }
+
+        self.limiter
+            .in_flight_write_bytes_counter()
+            .process_with_row_inserts(&requests);
+
         self.handle_influx_row_inserts(requests, ctx)
             .await
             .map_err(BoxedError::new)

src/frontend/src/instance/log_handler.rs

@@ -22,7 +22,8 @@ use common_error::ext::BoxedError;
 use pipeline::pipeline_operator::PipelineOperator;
 use pipeline::{GreptimeTransformer, Pipeline, PipelineInfo, PipelineVersion};
 use servers::error::{
-    AuthSnafu, Error as ServerError, ExecuteGrpcRequestSnafu, PipelineSnafu, Result as ServerResult,
+    AuthSnafu, Error as ServerError, ExecuteGrpcRequestSnafu, InFlightWriteBytesExceededSnafu,
+    PipelineSnafu, Result as ServerResult,
 };
 use servers::interceptor::{LogIngestInterceptor, LogIngestInterceptorRef};
 use servers::query_handler::PipelineHandler;
@@ -110,6 +111,14 @@ impl Instance {
         log: RowInsertRequests,
         ctx: QueryContextRef,
     ) -> ServerResult<Output> {
+        if self.limiter.should_be_limited() {
+            return InFlightWriteBytesExceededSnafu.fail();
+        }
+
+        self.limiter
+            .in_flight_write_bytes_counter()
+            .process_with_row_inserts(&log);
+
         self.inserter
             .handle_log_inserts(log, ctx, self.statement_executor.as_ref())
             .await

src/frontend/src/instance/opentsdb.rs

@@ -17,7 +17,7 @@ use auth::{PermissionChecker, PermissionCheckerRef, PermissionReq};
 use common_error::ext::BoxedError;
 use common_telemetry::tracing;
 use servers::error as server_error;
-use servers::error::AuthSnafu;
+use servers::error::{AuthSnafu, InFlightWriteBytesExceededSnafu};
 use servers::opentsdb::codec::DataPoint;
 use servers::opentsdb::data_point_to_grpc_row_insert_requests;
 use servers::query_handler::OpentsdbProtocolHandler;
@@ -41,6 +41,13 @@ impl OpentsdbProtocolHandler for Instance {
             .context(AuthSnafu)?;
 
         let (requests, _) = data_point_to_grpc_row_insert_requests(data_points)?;
+        if self.limiter.should_be_limited() {
+            return InFlightWriteBytesExceededSnafu.fail();
+        }
+
+        self.limiter
+            .in_flight_write_bytes_counter()
+            .process_with_row_inserts(&requests);
         let output = self
             .handle_row_inserts(requests, ctx)
             .await

src/frontend/src/instance/otlp.rs

@@ -21,7 +21,7 @@ use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest;
 use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest;
 use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
 use pipeline::PipelineWay;
-use servers::error::{self, AuthSnafu, Result as ServerResult};
+use servers::error::{self, AuthSnafu, InFlightWriteBytesExceededSnafu, Result as ServerResult};
 use servers::interceptor::{OpenTelemetryProtocolInterceptor, OpenTelemetryProtocolInterceptorRef};
 use servers::otlp;
 use servers::query_handler::OpenTelemetryProtocolHandler;
@@ -53,6 +53,14 @@ impl OpenTelemetryProtocolHandler for Instance {
         let (requests, rows) = otlp::metrics::to_grpc_insert_requests(request)?;
         OTLP_METRICS_ROWS.inc_by(rows as u64);
 
+        if self.limiter.should_be_limited() {
+            return InFlightWriteBytesExceededSnafu.fail();
+        }
+
+        self.limiter
+            .in_flight_write_bytes_counter()
+            .process_with_row_inserts(&requests);
+
         self.handle_row_inserts(requests, ctx)
             .await
             .map_err(BoxedError::new)
@@ -83,6 +91,14 @@ impl OpenTelemetryProtocolHandler for Instance {
 
         OTLP_TRACES_ROWS.inc_by(rows as u64);
 
+        if self.limiter.should_be_limited() {
+            return InFlightWriteBytesExceededSnafu.fail();
+        }
+
+        self.limiter
+            .in_flight_write_bytes_counter()
+            .process_with_row_inserts(&requests);
+
         self.handle_log_inserts(requests, ctx)
             .await
             .map_err(BoxedError::new)
@@ -109,6 +125,15 @@ impl OpenTelemetryProtocolHandler for Instance {
         interceptor_ref.pre_execute(ctx.clone())?;
 
         let (requests, rows) = otlp::logs::to_grpc_insert_requests(request, pipeline, table_name)?;
+
+        if self.limiter.should_be_limited() {
+            return InFlightWriteBytesExceededSnafu.fail();
+        }
+
+        self.limiter
+            .in_flight_write_bytes_counter()
+            .process_with_row_inserts(&requests);
+
         self.handle_log_inserts(requests, ctx)
             .await
             .inspect(|_| OTLP_LOGS_ROWS.inc_by(rows as u64))

src/frontend/src/instance/prom_store.rs

@@ -30,7 +30,7 @@ use common_telemetry::{debug, tracing};
 use operator::insert::InserterRef;
 use operator::statement::StatementExecutor;
 use prost::Message;
-use servers::error::{self, AuthSnafu, Result as ServerResult};
+use servers::error::{self, AuthSnafu, InFlightWriteBytesExceededSnafu, Result as ServerResult};
 use servers::http::header::{collect_plan_metrics, CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF};
 use servers::http::prom_store::PHYSICAL_TABLE_PARAM;
 use servers::interceptor::{PromStoreProtocolInterceptor, PromStoreProtocolInterceptorRef};
@@ -175,6 +175,14 @@ impl PromStoreProtocolHandler for Instance {
             .get::<PromStoreProtocolInterceptorRef<servers::error::Error>>();
         interceptor_ref.pre_write(&request, ctx.clone())?;
 
+        if self.limiter.should_be_limited() {
+            return InFlightWriteBytesExceededSnafu.fail();
+        }
+
+        self.limiter
+            .in_flight_write_bytes_counter()
+            .process_with_row_inserts(&request);
+
         let output = if with_metric_engine {
             let physical_table = ctx
                 .extension(PHYSICAL_TABLE_PARAM)

src/frontend/src/lib.rs

@@ -18,6 +18,7 @@ pub mod error;
 pub mod frontend;
 pub mod heartbeat;
 pub mod instance;
+pub(crate) mod limiter;
 pub(crate) mod metrics;
 mod script;
 pub mod server;