coderabbit comments

Author: parmesant

src/handlers/http/modal/ingest_server.rs

@@ -237,7 +237,10 @@ impl IngestServer {
                         .route(
                             web::post()
                                 .to(ingest::post_event)
-                                .authorize_for_resource(Action::Ingest),
+                                .authorize_for_resource(Action::Ingest)
+                                .wrap(from_fn(
+                                    resource_check::check_resource_utilization_middleware,
+                                )),
                         ),
                 )
                 .service(

src/handlers/http/resource_check.rs

@@ -25,23 +25,20 @@ use actix_web::{
     error::ErrorServiceUnavailable,
     middleware::Next,
 };
+use sysinfo::{MemoryRefreshKind, RefreshKind, System};
 use tokio::{
     select,
     time::{Duration, interval},
 };
 use tracing::{info, trace, warn};
 
+use crate::analytics::{SYS_INFO, refresh_sys_info};
 use crate::metrics::record_process_metrics_sample;
 use crate::parseable::PARSEABLE;
-use crate::{
-    analytics::{SYS_INFO, refresh_sys_info},
-    metrics::PROCESS_METRICS_ACCUMULATOR,
-};
 
 const PROCESS_METRICS_SAMPLE_INTERVAL: Duration = Duration::from_secs(5);
 
-static RESOURCE_CHECK_ENABLED: LazyLock<Arc<AtomicBool>> =
-    LazyLock::new(|| Arc::new(AtomicBool::new(false)));
+static SERVER_OK: LazyLock<Arc<AtomicBool>> = LazyLock::new(|| Arc::new(AtomicBool::new(true)));
 
 /// Spawn a background task to monitor system resources
 pub fn spawn_resource_monitor(shutdown_rx: tokio::sync::oneshot::Receiver<()>) {
@@ -51,7 +48,7 @@ pub fn spawn_resource_monitor(shutdown_rx: tokio::sync::oneshot::Receiver<()>) {
         let mut process_metrics_interval = interval(PROCESS_METRICS_SAMPLE_INTERVAL);
         let mut shutdown_rx = shutdown_rx;
 
-        let memory_threshold = PARSEABLE.options.memory_utilization_threshold;
+        let memory_threshold = (PARSEABLE.options.memory_utilization_threshold / 100.0) as f64;
 
         loop {
             select! {
@@ -62,19 +59,22 @@ pub fn spawn_resource_monitor(shutdown_rx: tokio::sync::oneshot::Receiver<()>) {
                     refresh_sys_info();
 
                     let mut resource_ok = true;
-                    let process_mem = PROCESS_METRICS_ACCUMULATOR.get_mem();
-                    let total_mem = PROCESS_METRICS_ACCUMULATOR.get_total_mem();
-                    let memory_usage = process_mem / total_mem;
 
-                    // Check memory utilization
-                    if memory_usage > memory_threshold as f64 {
-                        warn!("High memory usage detected: {:.1}% (threshold: {:.1}%)",
-                              memory_usage, memory_threshold);
-                        resource_ok = false;
+                    let mut s = System::new_with_specifics(
+                        RefreshKind::nothing().with_memory(MemoryRefreshKind::everything()),
+                    );
+                    if let Some(cgroup) = s.cgroup_limits() {
+                        if (cgroup.rss as f64) > memory_threshold * (cgroup.total_memory as f64) {
+                            resource_ok = false;
+                        }
+                    } else {
+                        s.refresh_memory();
+                        if (s.used_memory() as f64) > memory_threshold * (s.total_memory() as f64) {
+                            resource_ok = false;
+                        }
                     }
-
-                    let previous_state = RESOURCE_CHECK_ENABLED.load(std::sync::atomic::Ordering::SeqCst);
-                    RESOURCE_CHECK_ENABLED.store(resource_ok, std::sync::atomic::Ordering::SeqCst);
+                    let previous_state = SERVER_OK.load(std::sync::atomic::Ordering::SeqCst);
+                    SERVER_OK.store(resource_ok, std::sync::atomic::Ordering::SeqCst);
 
                     // Log state changes
                     if previous_state != resource_ok {
@@ -118,7 +118,7 @@ pub async fn check_resource_utilization_middleware(
     req: ServiceRequest,
     next: Next<impl MessageBody>,
 ) -> Result<ServiceResponse<impl MessageBody>, Error> {
-    let resource_ok = RESOURCE_CHECK_ENABLED.load(std::sync::atomic::Ordering::SeqCst);
+    let resource_ok = SERVER_OK.load(std::sync::atomic::Ordering::SeqCst);
 
     if !resource_ok {
         let error_msg = "Server resources over-utilized";

src/lib.rs

@@ -74,10 +74,10 @@ pub use openid;
 use parseable::PARSEABLE;
 use reqwest::{Client, ClientBuilder};
 pub use rustls;
+pub use sysinfo;
 pub use utils as parseable_utils;
 pub use {clap, tracing_actix_web, tracing_opentelemetry, tracing_subscriber};
 pub use {opentelemetry, opentelemetry_otlp, opentelemetry_proto, opentelemetry_sdk};
-
 // It is very unlikely that panic will occur when dealing with locks.
 pub const LOCK_EXPECT: &str = "Thread shouldn't panic while holding a lock";
 

src/main.rs

@@ -21,7 +21,7 @@ use std::process::exit;
 use parseable::connectors;
 use parseable::{
     IngestServer, ParseableServer, QueryServer, Server, banner, metrics, option::Mode,
-    parseable::PARSEABLE, rbac, storage,
+    parseable::PARSEABLE, rbac, storage,analytics::{SYS_INFO, refresh_sys_info}
 };
 use tokio::signal::ctrl_c;
 use tokio::sync::oneshot;
@@ -31,6 +31,8 @@ use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
 use tracing_subscriber::{EnvFilter, Registry, fmt};
 
+use crate::metrics::record_process_metrics_sample;
+
 #[actix_web::main]
 #[cfg_attr(feature = "hotpath", hotpath::main)]
 async fn main() -> anyhow::Result<()> {
@@ -89,6 +91,27 @@ async fn main() -> anyhow::Result<()> {
     });
 
     let prometheus = metrics::build_metrics_handler();
+    // init process metrics
+    refresh_sys_info();
+    let process_metrics = tokio::task::spawn_blocking(|| {
+        let sys = SYS_INFO.lock().unwrap();
+        let total_mem = if let Some(cgroup) = sys.cgroup_limits() {
+            cgroup.total_memory
+        } else {
+            sys.total_memory()
+        };
+        sysinfo::get_current_pid()
+            .ok()
+            .and_then(|pid| sys.process(pid))
+            .map(|process| (process.cpu_usage() as f64, process.memory(), total_mem))
+    })
+    .await
+    .unwrap();
+    // first measurement
+    if let Some((cpu_usage, memory_bytes, total_mem)) = process_metrics {
+        record_process_metrics_sample(cpu_usage, memory_bytes, total_mem);
+    }
+
     // Start servers
     #[cfg(feature = "kafka")]
     {

src/metrics/mod.rs

@@ -242,15 +242,15 @@ impl ProcessMetricsAccumulator {
     }
 
     pub fn get_cpu(&self) -> f64 {
-        self.cpu_usage_avg.get() as f64
+        self.cpu_usage_avg.get()
     }
 
     pub fn get_mem(&self) -> f64 {
-        self.memory_bytes_avg.get() as f64
+        self.memory_bytes_avg.get()
     }
 
     pub fn get_total_mem(&self) -> f64 {
-        self.total_memory.get() as f64
+        self.total_memory.get()
     }
 }