Skip to content

Commit dc3cdfe

Browse files
parmesantparmesant
committed
coderabbit comments
1 parent 7a47329 commit dc3cdfe

8 files changed

Lines changed: 82 additions & 44 deletions

File tree

‎src/cli.rs‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -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
@@ -237,7 +237,10 @@ impl IngestServer {
237237
.route(
238238
web::post()
239239
.to(ingest::post_event)
240-
.authorize_for_resource(Action::Ingest),
240+
.authorize_for_resource(Action::Ingest)
241+
.wrap(from_fn(
242+
resource_check::check_resource_utilization_middleware,
243+
)),
241244
),
242245
)
243246
.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;
@@ -51,6 +52,7 @@ use crate::sync::sync_start;
5152
use crate::handlers::http::alert_target_policy;
5253
use actix_web::Resource;
5354
use actix_web::Scope;
55+
use actix_web::middleware::from_fn;
5456
use actix_web::web;
5557
use actix_web::web::resource;
5658
use actix_web_prometheus::PrometheusMetrics;
@@ -86,7 +88,9 @@ impl ParseableServer for Server {
8688
web::scope(&base_path())
8789
.service(Self::get_query_factory())
8890
.service(Self::get_query_context_factory())
89-
.service(Self::get_ingest_factory())
91+
.service(Self::get_ingest_factory().wrap(from_fn(
92+
resource_check::check_resource_utilization_middleware,
93+
)))
9094
.service(Self::get_liveness_factory())
9195
.service(Self::get_readiness_factory())
9296
.service(Self::get_about_factory())
@@ -115,7 +119,9 @@ impl ParseableServer for Server {
115119
.service(Self::get_traces_webscope())
116120
.service(Self::get_dataset_stats_webscope()),
117121
)
118-
.service(Self::get_ingest_otel_factory())
122+
.service(Self::get_ingest_otel_factory().wrap(from_fn(
123+
resource_check::check_resource_utilization_middleware,
124+
)))
119125
.service(Self::get_generated());
120126
}
121127

@@ -529,7 +535,10 @@ impl Server {
529535
.route(
530536
web::post()
531537
.to(ingest::post_event)
532-
.authorize_for_resource(Action::Ingest),
538+
.authorize_for_resource(Action::Ingest)
539+
.wrap(from_fn(
540+
resource_check::check_resource_utilization_middleware,
541+
)),
533542
)
534543
// DELETE "/logstream/{logstream}" ==> Delete log stream
535544
.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)