feat: throttling for query (#1766)

* feat: throttling for query

Incoming query will stay in queue till the server's memory is under a set threshold.
The timeout is set to 10s, failing faster is better than staying in limbo.

* coderabbit suggestions

Author: parmesant

src/cli.rs

@@ -451,10 +451,20 @@ pub struct Options {
         long,
         long = "query-mempool-size",
         env = "P_QUERY_MEMORY_LIMIT",
-        help = "Set a fixed memory limit for query in GiB"
+        help = "Set a fixed memory limit for query in Bytes"
     )]
     pub query_memory_pool_size: Option<usize>,
 
+    #[arg(
+        long,
+        long = "query-mem-threshold",
+        value_parser = validation::validate_percentage,
+        default_value = "80.0",
+        env = "P_QUERY_MEMORY_THRESHOLD",
+        help = "Set a threshold (percentage) for memory beyond which query will get queued for 10s to prevent OOM"
+    )]
+    pub query_mem_threshold: f32,
+
     #[arg(
         long,
         env = "P_PARQUET_METADATA_CACHE_SIZE",

src/handlers/http/query.rs

@@ -730,7 +730,8 @@ Description: {0}"#
 impl actix_web::ResponseError for QueryError {
     fn status_code(&self) -> StatusCode {
         match self {
-            QueryError::Execute(_) | QueryError::JsonParse(_) => StatusCode::INTERNAL_SERVER_ERROR,
+            QueryError::JsonParse(_) => StatusCode::INTERNAL_SERVER_ERROR,
+            QueryError::Execute(e) => e.status_code(),
             QueryError::MetastoreError(e) => e.status_code(),
             _ => StatusCode::BAD_REQUEST,
         }

src/handlers/http/resource_check.rs

@@ -63,8 +63,11 @@ pub fn spawn_resource_monitor(shutdown_rx: tokio::sync::oneshot::Receiver<()>) {
                     refresh_sys_info();
                     let (used_memory, total_memory, cpu_usage) = tokio::task::spawn_blocking(|| {
                         let sys = SYS_INFO.lock().unwrap();
-                        let used_memory = sys.used_memory() as f32;
-                        let total_memory = sys.total_memory() as f32;
+                        let (used_memory, total_memory) = if let Some(cgroup) = sys.cgroup_limits() {
+                            (cgroup.rss as f32,cgroup.total_memory as f32)
+                        } else {
+                            (sys.used_memory() as f32,sys.total_memory() as f32)
+                        };
                         let cpu_usage = sys.global_cpu_usage();
                         (used_memory, total_memory, cpu_usage)
                     }).await.unwrap();

src/query/mod.rs

@@ -53,7 +53,7 @@ use std::pin::Pin;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, RwLock};
 use std::task::{Context, Poll};
-use sysinfo::System;
+use sysinfo::{MemoryRefreshKind, RefreshKind, System};
 use tokio::runtime::Runtime;
 use tokio_stream::wrappers::UnboundedReceiverStream;
 use tracing::Instrument;
@@ -82,9 +82,6 @@ type BoxedBatchStream = SendableRecordBatchStream;
 type QueryResult = Result<(Either<Vec<RecordBatch>, BoxedBatchStream>, Vec<String>), ExecuteError>;
 const DEFAULT_COUNTS_TOP_K: usize = 10;
 
-// pub static QUERY_SESSION: Lazy<SessionContext> =
-//     Lazy::new(|| Query::create_session_context(PARSEABLE.storage()));
-
 pub static QUERY_SESSION_STATE: Lazy<SessionState> =
     Lazy::new(|| Query::create_session_state(PARSEABLE.storage()));
 
@@ -147,10 +144,42 @@ impl InMemorySessionContext {
     }
 }
 
+async fn enough_available_memory() -> Result<(), ExecuteError> {
+    let mut s = System::new_with_specifics(
+        RefreshKind::nothing().with_memory(MemoryRefreshKind::everything()),
+    );
+    s.refresh_all();
+    let threshold = (PARSEABLE.options.query_mem_threshold / 100.0) as f64;
+
+    let f = async {
+        loop {
+            if let Some(cgroup) = s.cgroup_limits() {
+                if (cgroup.rss as f64) < threshold * (cgroup.total_memory as f64) {
+                    return;
+                }
+            } else {
+                s.refresh_memory();
+                if (s.used_memory() as f64) < threshold * (s.total_memory() as f64) {
+                    return;
+                }
+            }
+            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+        }
+    };
+    if let Err(e) = tokio::time::timeout(tokio::time::Duration::from_secs(10), f).await {
+        Err(ExecuteError::ServerBusy(e))
+    } else {
+        Ok(())
+    }
+}
+
 /// This function executes a query on the dedicated runtime, ensuring that the query is not isolated to a single thread/CPU
 /// at a time and has access to the entire thread pool, enabling better concurrent processing, and thus quicker results.
 pub async fn execute(query: Query, is_streaming: bool, tenant_id: &Option<String>) -> QueryResult {
     let id = tenant_id.clone();
+
+    // before executing query, check whether enough memory is available or not
+    enough_available_memory().await?;
     QUERY_RUNTIME
         .spawn(async move {
             tokio::time::timeout(
@@ -220,7 +249,11 @@ impl Query {
             None => {
                 let mut system = System::new();
                 system.refresh_memory();
-                let available_mem = system.available_memory();
+                let available_mem = if let Some(cgroup) = system.cgroup_limits() {
+                    cgroup.total_memory
+                } else {
+                    system.total_memory()
+                };
                 (available_mem as usize, 0.85)
             }
         };
@@ -270,6 +303,8 @@ impl Query {
             .parquet
             .schema_force_view_types = true;
 
+        config.options_mut().explain.show_statistics = true;
+
         SessionStateBuilder::new()
             .with_default_features()
             .with_config(config)
@@ -958,6 +993,7 @@ pub fn flatten_objects_for_count(objects: Vec<Value>) -> Vec<Value> {
 
 pub mod error {
     use crate::{parseable::StreamNotFound, storage::ObjectStorageError};
+    use actix_web::http::{StatusCode, header::ContentType};
     use datafusion::error::DataFusionError;
     use tokio::time::error::Elapsed;
 
@@ -971,6 +1007,26 @@ pub mod error {
         StreamNotFound(#[from] StreamNotFound),
         #[error("{0}: {1}s")]
         Timeout(Elapsed, u64),
+        #[error("{0}: Not enough memory available to serve the request")]
+        ServerBusy(Elapsed),
+    }
+
+    impl actix_web::ResponseError for ExecuteError {
+        fn status_code(&self) -> StatusCode {
+            match self {
+                ExecuteError::ObjectStorage(_) => StatusCode::INTERNAL_SERVER_ERROR,
+                ExecuteError::Datafusion(_) => StatusCode::INTERNAL_SERVER_ERROR,
+                ExecuteError::StreamNotFound(_) => StatusCode::NOT_FOUND,
+                ExecuteError::Timeout(_, _) => StatusCode::REQUEST_TIMEOUT,
+                ExecuteError::ServerBusy(_) => StatusCode::SERVICE_UNAVAILABLE,
+            }
+        }
+
+        fn error_response(&self) -> actix_web::HttpResponse<actix_web::body::BoxBody> {
+            actix_web::HttpResponse::build(self.status_code())
+                .insert_header(ContentType::plaintext())
+                .body(self.to_string())
+        }
     }
 }