@@ -23,51 +23,102 @@ use actix_web::{
23
23
error:: ErrorServiceUnavailable ,
24
24
middleware:: Next ,
25
25
} ;
26
- use sysinfo:: System ;
27
- use tracing:: warn;
26
+ use tokio:: { select, time:: { interval, Duration } } ;
27
+ use tokio:: sync:: RwLock ;
28
+ use tracing:: { warn, trace, info} ;
28
29
29
- const CPU_UTILIZATION_THRESHOLD : f32 = 90.0 ;
30
- const MEMORY_UTILIZATION_THRESHOLD : f32 = 90.0 ;
30
+ use crate :: analytics:: { SYS_INFO , refresh_sys_info} ;
31
+ use crate :: parseable:: PARSEABLE ;
32
+
33
+ static RESOURCE_CHECK_ENABLED : RwLock < bool > = RwLock :: const_new ( true ) ;
34
+
35
+ /// Spawn a background task to monitor system resources
36
+ pub fn spawn_resource_monitor ( shutdown_rx : tokio:: sync:: oneshot:: Receiver < ( ) > ) {
37
+ tokio:: spawn ( async move {
38
+ let mut check_interval = interval ( Duration :: from_secs ( 30 ) ) ;
39
+ let mut shutdown_rx = shutdown_rx;
40
+
41
+ let cpu_threshold = PARSEABLE . options . cpu_utilization_threshold ;
42
+ let memory_threshold = PARSEABLE . options . memory_utilization_threshold ;
43
+
44
+ info ! ( "Resource monitor started with thresholds - CPU: {:.1}%, Memory: {:.1}%" ,
45
+ cpu_threshold, memory_threshold) ;
46
+ loop {
47
+ select ! {
48
+ _ = check_interval. tick( ) => {
49
+ trace!( "Checking system resource utilization..." ) ;
50
+
51
+ refresh_sys_info( ) ;
52
+ let ( used_memory, total_memory, cpu_usage) = {
53
+ let sys = SYS_INFO . lock( ) . unwrap( ) ;
54
+ let used_memory = sys. used_memory( ) as f32 ;
55
+ let total_memory = sys. total_memory( ) as f32 ;
56
+ let cpu_usage = sys. global_cpu_usage( ) ;
57
+ ( used_memory, total_memory, cpu_usage)
58
+ } ;
59
+
60
+ let mut resource_ok = true ;
61
+
62
+ // Calculate memory usage percentage
63
+ let memory_usage = if total_memory > 0.0 {
64
+ ( used_memory / total_memory) * 100.0
65
+ } else {
66
+ 0.0
67
+ } ;
68
+
69
+ // Log current resource usage every few checks for debugging
70
+ info!( "Current resource usage - CPU: {:.1}%, Memory: {:.1}% ({:.1}GB/{:.1}GB)" ,
71
+ cpu_usage, memory_usage,
72
+ used_memory / 1024.0 / 1024.0 / 1024.0 ,
73
+ total_memory / 1024.0 / 1024.0 / 1024.0 ) ;
74
+
75
+ // Check memory utilization
76
+ if memory_usage > memory_threshold {
77
+ warn!( "High memory usage detected: {:.1}% (threshold: {:.1}%)" ,
78
+ memory_usage, memory_threshold) ;
79
+ resource_ok = false ;
80
+ }
81
+
82
+ // Check CPU utilization
83
+ if cpu_usage > cpu_threshold {
84
+ warn!( "High CPU usage detected: {:.1}% (threshold: {:.1}%)" ,
85
+ cpu_usage, cpu_threshold) ;
86
+ resource_ok = false ;
87
+ }
88
+
89
+ let previous_state = * RESOURCE_CHECK_ENABLED . read( ) . await ;
90
+ * RESOURCE_CHECK_ENABLED . write( ) . await = resource_ok;
91
+
92
+ // Log state changes
93
+ if previous_state != resource_ok {
94
+ if resource_ok {
95
+ info!( "Resource utilization back to normal - requests will be accepted" ) ;
96
+ } else {
97
+ warn!( "Resource utilization too high - requests will be rejected" ) ;
98
+ }
99
+ }
100
+ } ,
101
+ _ = & mut shutdown_rx => {
102
+ trace!( "Resource monitor shutting down" ) ;
103
+ break ;
104
+ }
105
+ }
106
+ }
107
+ } ) ;
108
+ }
31
109
32
110
/// Middleware to check system resource utilization before processing requests
33
- /// Returns 503 Service Unavailable if CPU or memory usage exceeds thresholds
111
+ /// Returns 503 Service Unavailable if resources are over-utilized
34
112
pub async fn check_resource_utilization_middleware (
35
113
req : ServiceRequest ,
36
114
next : Next < impl MessageBody > ,
37
115
) -> Result < ServiceResponse < impl MessageBody > , Error > {
38
116
39
- let mut sys = System :: new_all ( ) ;
40
- sys. refresh_cpu_usage ( ) ;
41
- sys. refresh_memory ( ) ;
117
+ let resource_ok = * RESOURCE_CHECK_ENABLED . read ( ) . await ;
42
118
43
- let used_memory = sys. used_memory ( ) as f32 ;
44
- let total_memory = sys. total_memory ( ) as f32 ;
45
-
46
- // Check memory utilization
47
- if total_memory > 0.0 {
48
- let memory_usage = ( used_memory / total_memory) * 100.0 ;
49
- if memory_usage > MEMORY_UTILIZATION_THRESHOLD {
50
- let error_msg = format ! ( "Memory is over-utilized: {:.1}%" , memory_usage) ;
51
- warn ! (
52
- "Rejecting request to {} due to high memory usage: {:.1}% (threshold: {:.1}%)" ,
53
- req. path( ) ,
54
- memory_usage,
55
- MEMORY_UTILIZATION_THRESHOLD
56
- ) ;
57
- return Err ( ErrorServiceUnavailable ( error_msg) ) ;
58
- }
59
- }
60
-
61
- // Check CPU utilization
62
- let cpu_usage = sys. global_cpu_usage ( ) ;
63
- if cpu_usage > CPU_UTILIZATION_THRESHOLD {
64
- let error_msg = format ! ( "CPU is over-utilized: {:.1}%" , cpu_usage) ;
65
- warn ! (
66
- "Rejecting request to {} due to high CPU usage: {:.1}% (threshold: {:.1}%)" ,
67
- req. path( ) ,
68
- cpu_usage,
69
- CPU_UTILIZATION_THRESHOLD
70
- ) ;
119
+ if !resource_ok {
120
+ let error_msg = "Server resources over-utilized" ;
121
+ warn ! ( "Rejecting request to {} due to resource constraints" , req. path( ) ) ;
71
122
return Err ( ErrorServiceUnavailable ( error_msg) ) ;
72
123
}
73
124
0 commit comments