Skip to content

Commit 7aeafbd

Browse files
parmesantparmesant
committed
coderabbit comments
1 parent 492b5c6 commit 7aeafbd

8 files changed

Lines changed: 83 additions & 45 deletions

File tree

‎src/cli.rs‎

Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -467,7 +467,7 @@ pub struct Options {
467467
long,
468468
long = "query-mem-threshold",
469469
value_parser = validation::validate_percentage,
470-
default_value = "80.0",
470+
default_value = "100.0",
471471
env = "P_QUERY_MEMORY_THRESHOLD",
472472
help = "Set a threshold (percentage) for memory beyond which query will get queued for 10s to prevent OOM"
473473
)]
@@ -518,15 +518,6 @@ pub struct Options {
518518
)]
519519
pub resource_check_interval: u64,
520520

521-
#[arg(
522-
long,
523-
env = "P_CPU_THRESHOLD",
524-
default_value = "100.0",
525-
value_parser = validation::validate_percentage,
526-
help = "CPU utilization threshold percentage (0.0-100.0) for resource monitoring"
527-
)]
528-
pub cpu_utilization_threshold: f32,
529-
530521
#[arg(
531522
long,
532523
env = "P_MEMORY_THRESHOLD",

‎src/handlers/http/modal/ingest_server.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -244,7 +244,10 @@ impl IngestServer {
244244
.route(
245245
web::post()
246246
.to(ingest::post_event)
247-
.authorize_for_resource(Action::Ingest),
247+
.authorize_for_resource(Action::Ingest)
248+
.wrap(from_fn(
249+
resource_check::check_resource_utilization_middleware,
250+
)),
248251
),
249252
)
250253
.service(

‎src/handlers/http/modal/server.rs‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ use crate::handlers::http::otel_generator::{
3434
use crate::handlers::http::prism_base_path;
3535
use crate::handlers::http::query;
3636
use crate::handlers::http::query_context;
37+
use crate::handlers::http::resource_check;
3738
use crate::handlers::http::targets;
3839
use crate::handlers::http::users::dashboards;
3940
use crate::handlers::http::users::filters;
@@ -50,6 +51,7 @@ use crate::sync;
5051
use crate::handlers::http::alert_target_policy;
5152
use actix_web::Resource;
5253
use actix_web::Scope;
54+
use actix_web::middleware::from_fn;
5355
use actix_web::web;
5456
use actix_web::web::resource;
5557
use actix_web_prometheus::PrometheusMetrics;
@@ -85,7 +87,9 @@ impl ParseableServer for Server {
8587
web::scope(&base_path())
8688
.service(Self::get_query_factory())
8789
.service(Self::get_query_context_factory())
88-
.service(Self::get_ingest_factory())
90+
.service(Self::get_ingest_factory().wrap(from_fn(
91+
resource_check::check_resource_utilization_middleware,
92+
)))
8993
.service(Self::get_liveness_factory())
9094
.service(Self::get_readiness_factory())
9195
.service(Self::get_about_factory())
@@ -114,7 +118,9 @@ impl ParseableServer for Server {
114118
.service(Self::get_traces_webscope())
115119
.service(Self::get_dataset_stats_webscope()),
116120
)
117-
.service(Self::get_ingest_otel_factory())
121+
.service(Self::get_ingest_otel_factory().wrap(from_fn(
122+
resource_check::check_resource_utilization_middleware,
123+
)))
118124
.service(Self::get_generated());
119125
}
120126

@@ -537,7 +543,10 @@ impl Server {
537543
.route(
538544
web::post()
539545
.to(ingest::post_event)
540-
.authorize_for_resource(Action::Ingest),
546+
.authorize_for_resource(Action::Ingest)
547+
.wrap(from_fn(
548+
resource_check::check_resource_utilization_middleware,
549+
)),
541550
)
542551
// DELETE "/logstream/{logstream}" ==> Delete log stream
543552
.route(

‎src/handlers/http/resource_check.rs‎

Lines changed: 19 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -25,23 +25,20 @@ use actix_web::{
2525
error::ErrorServiceUnavailable,
2626
middleware::Next,
2727
};
28+
use sysinfo::{MemoryRefreshKind, RefreshKind, System};
2829
use tokio::{
2930
select,
3031
time::{Duration, interval},
3132
};
3233
use tracing::{info, trace, warn};
3334

35+
use crate::analytics::{SYS_INFO, refresh_sys_info};
3436
use crate::metrics::record_process_metrics_sample;
3537
use crate::parseable::PARSEABLE;
36-
use crate::{
37-
analytics::{SYS_INFO, refresh_sys_info},
38-
metrics::PROCESS_METRICS_ACCUMULATOR,
39-
};
4038

4139
const PROCESS_METRICS_SAMPLE_INTERVAL: Duration = Duration::from_secs(5);
4240

43-
static RESOURCE_CHECK_ENABLED: LazyLock<Arc<AtomicBool>> =
44-
LazyLock::new(|| Arc::new(AtomicBool::new(false)));
41+
static SERVER_OK: LazyLock<Arc<AtomicBool>> = LazyLock::new(|| Arc::new(AtomicBool::new(true)));
4542

4643
/// Spawn a background task to monitor system resources
4744
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<()>) {
5148
let mut process_metrics_interval = interval(PROCESS_METRICS_SAMPLE_INTERVAL);
5249
let mut shutdown_rx = shutdown_rx;
5350

54-
let memory_threshold = PARSEABLE.options.memory_utilization_threshold;
51+
let memory_threshold = (PARSEABLE.options.memory_utilization_threshold / 100.0) as f64;
5552

5653
loop {
5754
select! {
@@ -62,19 +59,22 @@ pub fn spawn_resource_monitor(shutdown_rx: tokio::sync::oneshot::Receiver<()>) {
6259
refresh_sys_info();
6360

6461
let mut resource_ok = true;
65-
let process_mem = PROCESS_METRICS_ACCUMULATOR.get_mem();
66-
let total_mem = PROCESS_METRICS_ACCUMULATOR.get_total_mem();
67-
let memory_usage = process_mem / total_mem;
6862

69-
// Check memory utilization
70-
if memory_usage > memory_threshold as f64 {
71-
warn!("High memory usage detected: {:.1}% (threshold: {:.1}%)",
72-
memory_usage, memory_threshold);
73-
resource_ok = false;
63+
let mut s = System::new_with_specifics(
64+
RefreshKind::nothing().with_memory(MemoryRefreshKind::everything()),
65+
);
66+
if let Some(cgroup) = s.cgroup_limits() {
67+
if (cgroup.rss as f64) > memory_threshold * (cgroup.total_memory as f64) {
68+
resource_ok = false;
69+
}
70+
} else {
71+
s.refresh_memory();
72+
if (s.used_memory() as f64) > memory_threshold * (s.total_memory() as f64) {
73+
resource_ok = false;
74+
}
7475
}
75-
76-
let previous_state = RESOURCE_CHECK_ENABLED.load(std::sync::atomic::Ordering::SeqCst);
77-
RESOURCE_CHECK_ENABLED.store(resource_ok, std::sync::atomic::Ordering::SeqCst);
76+
let previous_state = SERVER_OK.load(std::sync::atomic::Ordering::SeqCst);
77+
SERVER_OK.store(resource_ok, std::sync::atomic::Ordering::SeqCst);
7878

7979
// Log state changes
8080
if previous_state != resource_ok {
@@ -118,7 +118,7 @@ pub async fn check_resource_utilization_middleware(
118118
req: ServiceRequest,
119119
next: Next<impl MessageBody>,
120120
) -> Result<ServiceResponse<impl MessageBody>, Error> {
121-
let resource_ok = RESOURCE_CHECK_ENABLED.load(std::sync::atomic::Ordering::SeqCst);
121+
let resource_ok = SERVER_OK.load(std::sync::atomic::Ordering::SeqCst);
122122

123123
if !resource_ok {
124124
let error_msg = "Server resources over-utilized";

‎src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,10 +74,10 @@ pub use openid;
7474
use parseable::PARSEABLE;
7575
use reqwest::{Client, ClientBuilder};
7676
pub use rustls;
77+
pub use sysinfo;
7778
pub use utils as parseable_utils;
7879
pub use {clap, tracing_actix_web, tracing_opentelemetry, tracing_subscriber};
7980
pub use {opentelemetry, opentelemetry_otlp, opentelemetry_proto, opentelemetry_sdk};
80-
8181
// It is very unlikely that panic will occur when dealing with locks.
8282
pub const LOCK_EXPECT: &str = "Thread shouldn't panic while holding a lock";
8383

‎src/main.rs‎

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,12 @@ use std::process::exit;
2020
#[cfg(feature = "kafka")]
2121
use parseable::connectors;
2222
use parseable::{
23-
IngestServer, ParseableServer, QueryServer, Server, banner, metrics, option::Mode,
24-
parseable::PARSEABLE, rbac, storage,
23+
IngestServer, ParseableServer, QueryServer, Server,
24+
analytics::{SYS_INFO, refresh_sys_info},
25+
banner, metrics,
26+
option::Mode,
27+
parseable::PARSEABLE,
28+
rbac, storage,
2529
};
2630
use tokio::signal::ctrl_c;
2731
use tokio::sync::oneshot;
@@ -31,6 +35,8 @@ use tracing_subscriber::layer::SubscriberExt;
3135
use tracing_subscriber::util::SubscriberInitExt;
3236
use tracing_subscriber::{EnvFilter, Registry, fmt};
3337

38+
use crate::metrics::record_process_metrics_sample;
39+
3440
#[actix_web::main]
3541
#[cfg_attr(feature = "hotpath", hotpath::main)]
3642
async fn main() -> anyhow::Result<()> {
@@ -89,6 +95,27 @@ async fn main() -> anyhow::Result<()> {
8995
});
9096

9197
let prometheus = metrics::build_metrics_handler();
98+
// init process metrics
99+
refresh_sys_info();
100+
let process_metrics = tokio::task::spawn_blocking(|| {
101+
let sys = SYS_INFO.lock().unwrap();
102+
let total_mem = if let Some(cgroup) = sys.cgroup_limits() {
103+
cgroup.total_memory
104+
} else {
105+
sys.total_memory()
106+
};
107+
sysinfo::get_current_pid()
108+
.ok()
109+
.and_then(|pid| sys.process(pid))
110+
.map(|process| (process.cpu_usage() as f64, process.memory(), total_mem))
111+
})
112+
.await
113+
.unwrap();
114+
// first measurement
115+
if let Some((cpu_usage, memory_bytes, total_mem)) = process_metrics {
116+
record_process_metrics_sample(cpu_usage, memory_bytes, total_mem);
117+
}
118+
92119
// Start servers
93120
#[cfg(feature = "kafka")]
94121
{

‎src/metrics/mod.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -242,15 +242,15 @@ impl ProcessMetricsAccumulator {
242242
}
243243

244244
pub fn get_cpu(&self) -> f64 {
245-
self.cpu_usage_avg.get() as f64
245+
self.cpu_usage_avg.get()
246246
}
247247

248248
pub fn get_mem(&self) -> f64 {
249-
self.memory_bytes_avg.get() as f64
249+
self.memory_bytes_avg.get()
250250
}
251251

252252
pub fn get_total_mem(&self) -> f64 {
253-
self.total_memory.get() as f64
253+
self.total_memory.get()
254254
}
255255
}
256256

‎src/query/mod.rs‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ use std::pin::Pin;
5454
use std::sync::atomic::{AtomicUsize, Ordering};
5555
use std::sync::{Arc, RwLock};
5656
use std::task::{Context, Poll};
57-
use sysinfo::System;
57+
use sysinfo::{MemoryRefreshKind, RefreshKind, System};
5858
use tokio::runtime::Runtime;
5959
use tokio_stream::wrappers::UnboundedReceiverStream;
6060
use tracing::Instrument;
@@ -70,7 +70,7 @@ use crate::catalog::manifest::Manifest;
7070
use crate::catalog::snapshot::Snapshot;
7171
use crate::event::DEFAULT_TIMESTAMP_KEY;
7272
use crate::handlers::http::query::QueryError;
73-
use crate::metrics::{PROCESS_METRICS_ACCUMULATOR, increment_bytes_scanned_in_query_by_date};
73+
use crate::metrics::increment_bytes_scanned_in_query_by_date;
7474
use crate::option::Mode;
7575
use crate::parseable::{DEFAULT_TENANT, PARSEABLE};
7676
use crate::storage::{ObjectStorage, ObjectStorageProvider, ObjectStoreFormat};
@@ -169,13 +169,21 @@ impl InMemorySessionContext {
169169
}
170170

171171
async fn enough_available_memory() -> Result<(), ExecuteError> {
172+
let mut s = System::new_with_specifics(
173+
RefreshKind::nothing().with_memory(MemoryRefreshKind::everything()),
174+
);
172175
let threshold = (PARSEABLE.options.query_mem_threshold / 100.0) as f64;
173176
let f = async {
174177
loop {
175-
let process_mem = PROCESS_METRICS_ACCUMULATOR.get_mem();
176-
let total_mem = PROCESS_METRICS_ACCUMULATOR.get_total_mem();
177-
if process_mem < threshold * total_mem {
178-
return;
178+
if let Some(cgroup) = s.cgroup_limits() {
179+
if (cgroup.rss as f64) > threshold * (cgroup.total_memory as f64) {
180+
return;
181+
}
182+
} else {
183+
s.refresh_memory();
184+
if (s.used_memory() as f64) > threshold * (s.total_memory() as f64) {
185+
return;
186+
}
179187
}
180188
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
181189
}

0 commit comments

Comments
 (0)