Adding optional load shedding and concurrency limitation to bulk ingest. (#4991)

Author: fulmicoton

quickwit/quickwit-common/src/lib.rs

@@ -93,6 +93,20 @@ pub fn get_from_env<T: FromStr + Debug>(key: &str, default_value: T) -> T {
     default_value
 }
 
+pub fn get_from_env_opt<T: FromStr + Debug>(key: &str) -> Option<T> {
+    let Some(value_str) = std::env::var(key).ok() else {
+        info!("{key} is not set");
+        return None;
+    };
+    if let Ok(value) = T::from_str(&value_str) {
+        info!(value=?value, "setting `{}` from environment", key);
+        Some(value)
+    } else {
+        error!(value_str=%value_str, "failed to parse `{}` from environment", key);
+        None
+    }
+}
+
 pub fn truncate_str(text: &str, max_len: usize) -> &str {
     if max_len > text.len() {
         return text;

quickwit/quickwit-serve/src/decompression.rs

@@ -18,6 +18,7 @@
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 use std::io::Read;
+use std::sync::OnceLock;
 
 use bytes::Bytes;
 use flate2::read::GzDecoder;
@@ -27,6 +28,13 @@ use thiserror::Error;
 use warp::reject::Reject;
 use warp::Filter;
 
+use crate::load_shield::{LoadShield, LoadShieldPermit};
+
+fn get_ingest_load_shield() -> &'static LoadShield {
+    static LOAD_SHIELD: OnceLock<LoadShield> = OnceLock::new();
+    LOAD_SHIELD.get_or_init(|| LoadShield::new("ingest"))
+}
+
 /// There are two ways to decompress the body:
 /// - Stream the body through an async decompressor
 /// - Fetch the body and then decompress the bytes
@@ -83,22 +91,27 @@ pub(crate) fn get_body_bytes() -> impl Filter<Extract = (Body,), Error = warp::R
     warp::header::optional("content-encoding")
         .and(warp::body::bytes())
         .and_then(|encoding: Option<String>, body: Bytes| async move {
-            decompress_body(encoding, body).await.map(Body::from)
+            let permit = get_ingest_load_shield().acquire_permit().await?;
+            decompress_body(encoding, body)
+                .await
+                .map(|content| Body::new(content, permit))
         })
 }
 
 pub(crate) struct Body {
     pub content: Bytes,
     _gauge_guard: GaugeGuard<'static>,
+    _permit: LoadShieldPermit,
 }
 
-impl From<Bytes> for Body {
-    fn from(content: Bytes) -> Self {
+impl Body {
+    pub fn new(content: Bytes, load_shield_permit: LoadShieldPermit) -> Body {
         let mut gauge_guard = GaugeGuard::from_gauge(&MEMORY_METRICS.in_flight.rest_server);
         gauge_guard.add(content.len() as i64);
         Body {
             content,
             _gauge_guard: gauge_guard,
+            _permit: load_shield_permit,
         }
     }
 }

quickwit/quickwit-serve/src/lib.rs

@@ -32,6 +32,7 @@ mod index_api;
 mod indexing_api;
 mod ingest_api;
 mod jaeger_api;
+mod load_shield;
 mod metrics;
 mod metrics_api;
 mod node_info_handler;

quickwit/quickwit-serve/src/load_shield.rs

@@ -0,0 +1,99 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+use std::time::Duration;
+
+use quickwit_common::metrics::{GaugeGuard, IntGauge};
+use tokio::sync::{Semaphore, SemaphorePermit};
+
+use crate::rest::TooManyRequests;
+
+pub struct LoadShield {
+    in_flight_semaphore_opt: Option<Semaphore>, // This one is doing the load shedding.
+    concurrency_semaphore_opt: Option<Semaphore>,
+    ongoing_gauge: IntGauge,
+    pending_gauge: IntGauge,
+}
+
+pub struct LoadShieldPermit {
+    _concurrency_permit_opt: Option<SemaphorePermit<'static>>,
+    _in_flight_permit_opt: Option<SemaphorePermit<'static>>,
+    _ongoing_gauge_guard: GaugeGuard<'static>,
+}
+
+impl LoadShield {
+    pub fn new(endpoint_group: &'static str) -> LoadShield {
+        let endpoint_group_uppercase = endpoint_group.to_ascii_uppercase();
+        let max_in_flight_env_key = format!("QW_{endpoint_group_uppercase}_MAX_IN_FLIGHT");
+        let max_concurrency_env_key = format!("QW_{endpoint_group_uppercase}_MAX_CONCURRENCY");
+        let max_in_flight_opt: Option<usize> =
+            quickwit_common::get_from_env_opt(&max_in_flight_env_key);
+        let max_concurrency_opt: Option<usize> =
+            quickwit_common::get_from_env_opt(&max_concurrency_env_key);
+        let in_flight_semaphore_opt = max_in_flight_opt.map(Semaphore::new);
+        let concurrency_semaphore_opt = max_concurrency_opt.map(Semaphore::new);
+        let pending_gauge = crate::metrics::SERVE_METRICS
+            .pending_requests
+            .with_label_values([endpoint_group]);
+        let ongoing_gauge = crate::metrics::SERVE_METRICS
+            .ongoing_requests
+            .with_label_values([endpoint_group]);
+        LoadShield {
+            in_flight_semaphore_opt,
+            concurrency_semaphore_opt,
+            ongoing_gauge,
+            pending_gauge,
+        }
+    }
+
+    async fn acquire_in_flight_permit(
+        &'static self,
+    ) -> Result<Option<SemaphorePermit<'static>>, warp::Rejection> {
+        let Some(in_flight_semaphore) = &self.in_flight_semaphore_opt else {
+            return Ok(None);
+        };
+        let Ok(in_flight_permit) = in_flight_semaphore.try_acquire() else {
+            // Wait a little to deal before load shedding. The point is to lower the load associated
+            // with super aggressive clients.
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            return Err(warp::reject::custom(TooManyRequests));
+        };
+        Ok(Some(in_flight_permit))
+    }
+
+    async fn acquire_concurrency_permit(&'static self) -> Option<SemaphorePermit<'static>> {
+        let concurrency_semaphore = self.concurrency_semaphore_opt.as_ref()?;
+        Some(concurrency_semaphore.acquire().await.unwrap())
+    }
+
+    pub async fn acquire_permit(&'static self) -> Result<LoadShieldPermit, warp::Rejection> {
+        let mut pending_gauge_guard = GaugeGuard::from_gauge(&self.pending_gauge);
+        pending_gauge_guard.add(1);
+        let in_flight_permit_opt = self.acquire_in_flight_permit().await?;
+        let concurrency_permit_opt = self.acquire_concurrency_permit().await;
+        drop(pending_gauge_guard);
+        let mut ongoing_gauge_guard = GaugeGuard::from_gauge(&self.ongoing_gauge);
+        ongoing_gauge_guard.add(1);
+        Ok(LoadShieldPermit {
+            _in_flight_permit_opt: in_flight_permit_opt,
+            _concurrency_permit_opt: concurrency_permit_opt,
+            _ongoing_gauge_guard: ongoing_gauge_guard,
+        })
+    }
+}

quickwit/quickwit-serve/src/metrics.rs

@@ -18,11 +18,15 @@
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 use once_cell::sync::Lazy;
-use quickwit_common::metrics::{new_counter_vec, new_histogram_vec, HistogramVec, IntCounterVec};
+use quickwit_common::metrics::{
+    new_counter_vec, new_gauge_vec, new_histogram_vec, HistogramVec, IntCounterVec, IntGaugeVec,
+};
 
 pub struct RestMetrics {
     pub http_requests_total: IntCounterVec<2>,
     pub request_duration_secs: HistogramVec<2>,
+    pub ongoing_requests: IntGaugeVec<1>,
+    pub pending_requests: IntGaugeVec<1>,
 }
 
 impl Default for RestMetrics {
@@ -43,6 +47,20 @@ impl Default for RestMetrics {
                 ["method", "status_code"],
                 quickwit_common::metrics::exponential_buckets(0.02, 2.0, 8).unwrap(),
             ),
+            ongoing_requests: new_gauge_vec(
+                "ingest_ongoing_requests",
+                "Number of ongoing ingest requests.",
+                "",
+                &[],
+                ["endpoint_group"],
+            ),
+            pending_requests: new_gauge_vec(
+                "ingest_pending_requests",
+                "Number of pending ingest requests.",
+                "",
+                &[],
+                ["endpoint_group"],
+            ),
         }
     }
 }

quickwit/quickwit-serve/src/rest.rs

@@ -17,6 +17,7 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
+use std::fmt::Formatter;
 use std::net::SocketAddr;
 use std::sync::Arc;
 
@@ -65,6 +66,28 @@ pub(crate) struct InvalidArgument(pub String);
 
 impl warp::reject::Reject for InvalidArgument {}
 
+#[derive(Debug)]
+pub struct TooManyRequests;
+
+impl warp::reject::Reject for TooManyRequests {}
+
+impl std::fmt::Display for TooManyRequests {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "too many requests")
+    }
+}
+
+#[derive(Debug)]
+pub struct InternalError(pub String);
+
+impl warp::reject::Reject for InternalError {}
+
+impl std::fmt::Display for InternalError {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "internal error: {}", self.0)
+    }
+}
+
 /// Starts REST services.
 pub(crate) async fn start_rest_server(
     rest_listen_addr: SocketAddr,
@@ -325,6 +348,11 @@ fn get_status_with_error(rejection: Rejection) -> RestApiError {
             status_code: StatusCode::PAYLOAD_TOO_LARGE,
             message: error.to_string(),
         }
+    } else if let Some(err) = rejection.find::<TooManyRequests>() {
+        RestApiError {
+            status_code: StatusCode::TOO_MANY_REQUESTS,
+            message: err.to_string(),
+        }
     } else {
         error!("REST server error: {:?}", rejection);
         RestApiError {