Added resource check interval option and updated the resource monitoring interval to be taken from the cli via --resource-check-interval flag. Also wrapped the middleware around cpu intensive query endpoints available in query.rs.

Author: vkhinvasara

src/cli.rs

@@ -318,6 +318,15 @@ pub struct Options {
     pub parquet_compression: Compression,
 
     // Resource monitoring
+    #[arg(
+        long,
+        env = "P_RESOURCE_CHECK_INTERVAL",
+        default_value = "30",
+        value_parser = validation::validate_seconds,
+        help = "Resource monitoring check interval in seconds"
+    )]
+    pub resource_check_interval: u64,
+
     #[arg(
         long,
         env = "P_CPU_THRESHOLD",

src/handlers/http/modal/query_server.rs

@@ -22,13 +22,14 @@ use std::thread;
 use crate::handlers::airplane;
 use crate::handlers::http::cluster::{self, init_cluster_metrics_schedular};
 use crate::handlers::http::middleware::{DisAllowRootUser, RouteExt};
-use crate::handlers::http::{base_path, prism_base_path};
+use crate::handlers::http::{base_path, prism_base_path, resource_check};
 use crate::handlers::http::{logstream, MAX_EVENT_PAYLOAD_SIZE};
 use crate::handlers::http::{rbac, role};
 use crate::hottier::HotTierManager;
 use crate::rbac::role::Action;
 use crate::sync::sync_start;
 use crate::{analytics, migration, storage, sync};
+use actix_web::middleware::from_fn;
 use actix_web::web::{resource, ServiceConfig};
 use actix_web::{web, Scope};
 use actix_web_prometheus::PrometheusMetrics;
@@ -53,7 +54,10 @@ impl ParseableServer for QueryServer {
             .service(
                 web::scope(&base_path())
                     .service(Server::get_correlation_webscope())
-                    .service(Server::get_query_factory())
+                    .service(
+                        Server::get_query_factory()
+                            .wrap(from_fn(resource_check::check_resource_utilization_middleware))
+                    )
                     .service(Server::get_liveness_factory())
                     .service(Server::get_readiness_factory())
                     .service(Server::get_about_factory())
@@ -66,7 +70,10 @@ impl ParseableServer for QueryServer {
                     .service(Server::get_oauth_webscope(oidc_client))
                     .service(Self::get_user_role_webscope())
                     .service(Server::get_roles_webscope())
-                    .service(Server::get_counts_webscope())
+                    .service(
+                        Server::get_counts_webscope()
+                            .wrap(from_fn(resource_check::check_resource_utilization_middleware))
+                    )
                     .service(Server::get_metrics_webscope())
                     .service(Server::get_alerts_webscope())
                     .service(Self::get_cluster_web_scope()),
@@ -143,13 +150,19 @@ impl ParseableServer for QueryServer {
         let (cancel_tx, cancel_rx) = oneshot::channel();
         thread::spawn(|| sync::handler(cancel_rx));
 
+        // Start resource monitor
+        let (resource_shutdown_tx, resource_shutdown_rx) = oneshot::channel();
+        resource_check::spawn_resource_monitor(resource_shutdown_rx);
+
         tokio::spawn(airplane::server());
 
         let result = self
             .start(shutdown_rx, prometheus.clone(), PARSEABLE.options.openid())
             .await?;
         // Cancel sync jobs
         cancel_tx.send(()).expect("Cancellation should not fail");
+        // Shutdown resource monitor
+        let _ = resource_shutdown_tx.send(());
         if let Err(join_err) = startup_sync_handle.await {
             tracing::warn!("startup sync task panicked: {join_err}");
         }

src/handlers/http/modal/server.rs

@@ -73,7 +73,10 @@ impl ParseableServer for Server {
             .service(
                 web::scope(&base_path())
                     .service(Self::get_correlation_webscope())
-                    .service(Self::get_query_factory())
+                    .service(
+                        Self::get_query_factory()
+                            .wrap(from_fn(resource_check::check_resource_utilization_middleware))
+                    )
                     .service(
                         Self::get_ingest_factory()
                             .wrap(from_fn(resource_check::check_resource_utilization_middleware))
@@ -90,7 +93,10 @@ impl ParseableServer for Server {
                     .service(Self::get_oauth_webscope(oidc_client))
                     .service(Self::get_user_role_webscope())
                     .service(Self::get_roles_webscope())
-                    .service(Self::get_counts_webscope())
+                    .service(
+                        Self::get_counts_webscope()
+                            .wrap(from_fn(resource_check::check_resource_utilization_middleware))
+                    )
                     .service(Self::get_alerts_webscope())
                     .service(Self::get_metrics_webscope()),
             )

src/handlers/http/resource_check.rs

@@ -36,7 +36,8 @@ static RESOURCE_CHECK_ENABLED:LazyLock<Arc<AtomicBool>> = LazyLock::new(|| Arc::
 /// Spawn a background task to monitor system resources
 pub fn spawn_resource_monitor(shutdown_rx: tokio::sync::oneshot::Receiver<()>) {
     tokio::spawn(async move {
-        let mut check_interval = interval(Duration::from_secs(30));
+        let resource_check_interval = PARSEABLE.options.resource_check_interval;
+        let mut check_interval = interval(Duration::from_secs(resource_check_interval));
         let mut shutdown_rx = shutdown_rx;
 
         let cpu_threshold = PARSEABLE.options.cpu_utilization_threshold;

src/option.rs

@@ -187,6 +187,13 @@ pub mod validation {
         }
     }
 
+    pub fn validate_seconds(s: &str) -> Result<u64, String> {
+        if let Ok(seconds) = s.parse::<u64>() {
+            Ok(seconds)
+        } else {
+            Err("Invalid value for seconds. It should be a positive integer".to_string())
+        }
+    }
     pub fn validate_dataset_fields_allowed_limit(s: &str) -> Result<usize, String> {
         if let Ok(size) = s.parse::<usize>() {
             if (1..=DATASET_FIELD_COUNT_LIMIT).contains(&size) {