feat: introduce the Limiter in frontend to limit the requests by in-f…

…light write bytes size. (#5231)

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

config/config.md

@@ -18,6 +18,7 @@
 | `init_regions_parallelism` | Integer | `16` | Parallelism of initializing regions. |
 | `max_concurrent_queries` | Integer | `0` | The maximum current queries allowed to be executed. Zero means unlimited. |
 | `enable_telemetry` | Bool | `true` | Enable telemetry to collect anonymous usage data. Enabled by default. |
+| `max_in_flight_write_bytes` | String | Unset | The maximum in-flight write bytes. |
 | `runtime` | -- | -- | The runtime options. |
 | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. |
 | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute the runtime for global write operations. |
@@ -195,6 +196,7 @@
 | Key | Type | Default | Descriptions |
 | --- | -----| ------- | ----------- |
 | `default_timezone` | String | Unset | The default timezone of the server. |
+| `max_in_flight_write_bytes` | String | Unset | The maximum in-flight write bytes. |
 | `runtime` | -- | -- | The runtime options. |
 | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. |
 | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute the runtime for global write operations. |

config/frontend.example.toml

@@ -2,6 +2,10 @@
 ## @toml2docs:none-default
 default_timezone = "UTC"
 
+## The maximum in-flight write bytes.
+## @toml2docs:none-default
+#+ max_in_flight_write_bytes = "500MB"
+
 ## The runtime options.
 #+ [runtime]
 ## The number of threads to execute the runtime for global read operations.

config/standalone.example.toml

@@ -18,6 +18,10 @@ max_concurrent_queries = 0
 ## Enable telemetry to collect anonymous usage data. Enabled by default.
 #+ enable_telemetry = true
 
+## The maximum in-flight write bytes.
+## @toml2docs:none-default
+#+ max_in_flight_write_bytes = "500MB"
+
 ## The runtime options.
 #+ [runtime]
 ## The number of threads to execute the runtime for global read operations.

src/cmd/src/standalone.rs

@@ -22,6 +22,7 @@ use catalog::information_schema::InformationExtension;
 use catalog::kvbackend::KvBackendCatalogManager;
 use clap::Parser;
 use client::api::v1::meta::RegionRole;
+use common_base::readable_size::ReadableSize;
 use common_base::Plugins;
 use common_catalog::consts::{MIN_USER_FLOW_ID, MIN_USER_TABLE_ID};
 use common_config::{metadata_store_dir, Configurable, KvBackendConfig};
@@ -152,6 +153,7 @@ pub struct StandaloneOptions {
     pub tracing: TracingOptions,
     pub init_regions_in_background: bool,
     pub init_regions_parallelism: usize,
+    pub max_in_flight_write_bytes: Option<ReadableSize>,
 }
 
 impl Default for StandaloneOptions {
@@ -181,6 +183,7 @@ impl Default for StandaloneOptions {
             tracing: TracingOptions::default(),
             init_regions_in_background: false,
             init_regions_parallelism: 16,
+            max_in_flight_write_bytes: None,
         }
     }
 }
@@ -218,6 +221,7 @@ impl StandaloneOptions {
             user_provider: cloned_opts.user_provider,
             // Handle the export metrics task run by standalone to frontend for execution
             export_metrics: cloned_opts.export_metrics,
+            max_in_flight_write_bytes: cloned_opts.max_in_flight_write_bytes,
             ..Default::default()
         }
     }

src/frontend/src/error.rs

@@ -321,6 +321,12 @@ pub enum Error {
         location: Location,
         source: BoxedError,
     },
+
+    #[snafu(display("In-flight write bytes exceeded the maximum limit"))]
+    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

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use common_base::readable_size::ReadableSize;
 use common_config::config::Configurable;
 use common_options::datanode::DatanodeClientOptions;
 use common_telemetry::logging::{LoggingOptions, TracingOptions};
@@ -46,6 +47,7 @@ pub struct FrontendOptions {
     pub user_provider: Option<String>,
     pub export_metrics: ExportMetricsOption,
     pub tracing: TracingOptions,
+    pub max_in_flight_write_bytes: Option<ReadableSize>,
 }
 
 impl Default for FrontendOptions {
@@ -68,6 +70,7 @@ impl Default for FrontendOptions {
             user_provider: None,
             export_metrics: ExportMetricsOption::default(),
             tracing: TracingOptions::default(),
+            max_in_flight_write_bytes: None,
         }
     }
 }

src/frontend/src/instance.rs

@@ -87,6 +87,7 @@ use crate::error::{
 };
 use crate::frontend::FrontendOptions;
 use crate::heartbeat::HeartbeatTask;
+use crate::limiter::LimiterRef;
 use crate::script::ScriptExecutor;
 
 #[async_trait]
@@ -126,6 +127,7 @@ pub struct Instance {
     export_metrics_task: Option<ExportMetricsTask>,
     table_metadata_manager: TableMetadataManagerRef,
     stats: StatementStatistics,
+    limiter: Option<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,14 @@ impl FrontendBuilder {
 
         plugins.insert::<StatementExecutorRef>(statement_executor.clone());
 
+        // Create the limiter if the max_in_flight_write_bytes is set.
+        let limiter = self
+            .options
+            .max_in_flight_write_bytes
+            .map(|max_in_flight_write_bytes| {
+                Arc::new(Limiter::new(max_in_flight_write_bytes.as_bytes()))
+            });
+
         Ok(Instance {
             options: self.options,
             catalog_manager: self.catalog_manager,
@@ -211,6 +220,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,16 @@ impl GrpcQueryHandler for Instance {
             .check_permission(ctx.current_user(), PermissionReq::GrpcRequest(&request))
             .context(PermissionSnafu)?;
 
+        let _guard = if let Some(limiter) = &self.limiter {
+            let result = limiter.limit_request(&request);
+            if result.is_none() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            result
+        } else {
+            None
+        };
+
         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,16 @@ impl InfluxdbLineProtocolHandler for Instance {
             .post_lines_conversion(requests, ctx.clone())
             .await?;
 
+        let _guard = if let Some(limiter) = &self.limiter {
+            let result = limiter.limit_row_inserts(&requests);
+            if result.is_none() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            result
+        } else {
+            None
+        };
+
         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,16 @@ impl Instance {
         log: RowInsertRequests,
         ctx: QueryContextRef,
     ) -> ServerResult<Output> {
+        let _guard = if let Some(limiter) = &self.limiter {
+            let result = limiter.limit_row_inserts(&log);
+            if result.is_none() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            result
+        } else {
+            None
+        };
+
         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,17 @@ impl OpentsdbProtocolHandler for Instance {
             .context(AuthSnafu)?;
 
         let (requests, _) = data_point_to_grpc_row_insert_requests(data_points)?;
+
+        let _guard = if let Some(limiter) = &self.limiter {
+            let result = limiter.limit_row_inserts(&requests);
+            if result.is_none() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            result
+        } else {
+            None
+        };
+
         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,16 @@ impl OpenTelemetryProtocolHandler for Instance {
         let (requests, rows) = otlp::metrics::to_grpc_insert_requests(request)?;
         OTLP_METRICS_ROWS.inc_by(rows as u64);
 
+        let _guard = if let Some(limiter) = &self.limiter {
+            let result = limiter.limit_row_inserts(&requests);
+            if result.is_none() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            result
+        } else {
+            None
+        };
+
         self.handle_row_inserts(requests, ctx)
             .await
             .map_err(BoxedError::new)
@@ -83,6 +93,16 @@ impl OpenTelemetryProtocolHandler for Instance {
 
         OTLP_TRACES_ROWS.inc_by(rows as u64);
 
+        let _guard = if let Some(limiter) = &self.limiter {
+            let result = limiter.limit_row_inserts(&requests);
+            if result.is_none() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            result
+        } else {
+            None
+        };
+
         self.handle_log_inserts(requests, ctx)
             .await
             .map_err(BoxedError::new)
@@ -109,6 +129,17 @@ impl OpenTelemetryProtocolHandler for Instance {
         interceptor_ref.pre_execute(ctx.clone())?;
 
         let (requests, rows) = otlp::logs::to_grpc_insert_requests(request, pipeline, table_name)?;
+
+        let _guard = if let Some(limiter) = &self.limiter {
+            let result = limiter.limit_row_inserts(&requests);
+            if result.is_none() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            result
+        } else {
+            None
+        };
+
         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,16 @@ impl PromStoreProtocolHandler for Instance {
             .get::<PromStoreProtocolInterceptorRef<servers::error::Error>>();
         interceptor_ref.pre_write(&request, ctx.clone())?;
 
+        let _guard = if let Some(limiter) = &self.limiter {
+            let result = limiter.limit_row_inserts(&request);
+            if result.is_none() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            result
+        } else {
+            None
+        };
+
         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;