diff --git a/src/cli.rs b/src/cli.rs index 1e44d36bd..f420030ab 100644 --- a/src/cli.rs +++ b/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, + #[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", diff --git a/src/handlers/http/query.rs b/src/handlers/http/query.rs index c7baf8cba..34c8a7327 100644 --- a/src/handlers/http/query.rs +++ b/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, } diff --git a/src/handlers/http/resource_check.rs b/src/handlers/http/resource_check.rs index 28e05c1db..77dc1f87e 100644 --- a/src/handlers/http/resource_check.rs +++ b/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(); diff --git a/src/query/mod.rs b/src/query/mod.rs index 728e3cdb6..36a94c37d 100644 --- a/src/query/mod.rs +++ b/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, BoxedBatchStream>, Vec), ExecuteError>; const DEFAULT_COUNTS_TOP_K: usize = 10; -// pub static QUERY_SESSION: Lazy = -// Lazy::new(|| Query::create_session_context(PARSEABLE.storage())); - pub static QUERY_SESSION_STATE: Lazy = 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) -> 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) -> Vec { 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::HttpResponse::build(self.status_code()) + .insert_header(ContentType::plaintext()) + .body(self.to_string()) + } } }