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

GreptimeTeam · Dec 24, 2024 · fb3a702 · fb3a702
1 parent d51b65a
commit fb3a702
Show file tree

Hide file tree

Showing 14 changed files with 406 additions and 7 deletions.
diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs
@@ -152,6 +152,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<u64>,
 }
 
 impl Default for StandaloneOptions {
@@ -181,6 +182,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 +220,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()
         }
     }

diff --git a/src/frontend/src/error.rs b/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,
         }
     }
 

diff --git a/src/frontend/src/frontend.rs b/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: Option<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: None,
         }
     }
 }

diff --git a/src/frontend/src/instance.rs b/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: Option<LimiterRef>,
 }
 
 impl Instance {

diff --git a/src/frontend/src/instance/builder.rs b/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,13 @@ impl FrontendBuilder {
 
         plugins.insert::<StatementExecutorRef>(statement_executor.clone());
 
+        let limiter =
+            if let Some(max_in_flight_write_bytes) = self.options.max_in_flight_write_bytes {
+                Some(Arc::new(Limiter::new(max_in_flight_write_bytes)))
+            } else {
+                None
+            };
+
         Ok(Instance {
             options: self.options,
             catalog_manager: self.catalog_manager,
@@ -211,6 +219,7 @@ impl FrontendBuilder {
             export_metrics_task: None,
             table_metadata_manager: Arc::new(TableMetadataManager::new(kv_backend)),
             stats: self.stats,
+            limiter,
         })
     }
 }
diff --git a/src/frontend/src/instance/grpc.rs b/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)?;
 
+        if let Some(limiter) = &self.limiter {
+            if limiter.should_be_limited() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+
+            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?,

diff --git a/src/frontend/src/instance/influxdb.rs b/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?;
 
+        if let Some(limiter) = &self.limiter {
+            if limiter.should_be_limited() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+
+            limiter
+                .in_flight_write_bytes_counter()
+                .process_with_row_inserts(&requests);
+        }
+
         self.handle_influx_row_inserts(requests, ctx)
             .await
             .map_err(BoxedError::new)

diff --git a/src/frontend/src/instance/log_handler.rs b/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,15 @@ impl Instance {
         log: RowInsertRequests,
         ctx: QueryContextRef,
     ) -> ServerResult<Output> {
+        if let Some(limiter) = &self.limiter {
+            if limiter.should_be_limited() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            limiter
+                .in_flight_write_bytes_counter()
+                .process_with_row_inserts(&log);
+        }
+
         self.inserter
             .handle_log_inserts(log, ctx, self.statement_executor.as_ref())
             .await

diff --git a/src/frontend/src/instance/opentsdb.rs b/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,16 @@ impl OpentsdbProtocolHandler for Instance {
             .context(AuthSnafu)?;
 
         let (requests, _) = data_point_to_grpc_row_insert_requests(data_points)?;
+
+        if let Some(limiter) = &self.limiter {
+            if limiter.should_be_limited() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            limiter
+                .in_flight_write_bytes_counter()
+                .process_with_row_inserts(&requests);
+        }
+
         let output = self
             .handle_row_inserts(requests, ctx)
             .await

diff --git a/src/frontend/src/instance/otlp.rs b/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,15 @@ impl OpenTelemetryProtocolHandler for Instance {
         let (requests, rows) = otlp::metrics::to_grpc_insert_requests(request)?;
         OTLP_METRICS_ROWS.inc_by(rows as u64);
 
+        if let Some(limiter) = &self.limiter {
+            if limiter.should_be_limited() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            limiter
+                .in_flight_write_bytes_counter()
+                .process_with_row_inserts(&requests);
+        }
+
         self.handle_row_inserts(requests, ctx)
             .await
             .map_err(BoxedError::new)
@@ -83,6 +92,15 @@ impl OpenTelemetryProtocolHandler for Instance {
 
         OTLP_TRACES_ROWS.inc_by(rows as u64);
 
+        if let Some(limiter) = &self.limiter {
+            if limiter.should_be_limited() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            limiter
+                .in_flight_write_bytes_counter()
+                .process_with_row_inserts(&requests);
+        }
+
         self.handle_log_inserts(requests, ctx)
             .await
             .map_err(BoxedError::new)
@@ -109,6 +127,16 @@ impl OpenTelemetryProtocolHandler for Instance {
         interceptor_ref.pre_execute(ctx.clone())?;
 
         let (requests, rows) = otlp::logs::to_grpc_insert_requests(request, pipeline, table_name)?;
+
+        if let Some(limiter) = &self.limiter {
+            if limiter.should_be_limited() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            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))

diff --git a/src/frontend/src/instance/prom_store.rs b/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,15 @@ impl PromStoreProtocolHandler for Instance {
             .get::<PromStoreProtocolInterceptorRef<servers::error::Error>>();
         interceptor_ref.pre_write(&request, ctx.clone())?;
 
+        if let Some(limiter) = &self.limiter {
+            if limiter.should_be_limited() {
+                return InFlightWriteBytesExceededSnafu.fail();
+            }
+            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)

diff --git a/src/frontend/src/lib.rs b/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;