11use std:: { sync:: Arc , time:: Duration } ;
22
3+ #[ cfg( test) ]
4+ use mockall:: { automock, predicate:: * } ;
35use observability_deps:: tracing:: debug;
4- use sysinfo:: { ProcessRefreshKind , System } ;
6+ use sysinfo:: { Pid , ProcessRefreshKind , System } ;
57
8+ use crate :: store:: TelemetryStore ;
69use crate :: Result ;
7- use crate :: { store:: TelemetryStore , TelemetryError } ;
810
9- struct CpuAndMemorySampler {
11+ #[ cfg_attr( test, automock) ]
12+ pub trait SystemInfoProvider : Send + Sync + ' static {
13+ fn refresh_metrics ( & mut self , pid : Pid ) ;
14+
15+ fn get_pid ( & self ) -> Result < Pid , & ' static str > ;
16+
17+ fn get_process_specific_metrics ( & self , pid : Pid ) -> Option < ( f32 , u64 ) > ;
18+ }
19+
20+ struct SystemInfo {
1021 system : System ,
1122}
1223
13- impl CpuAndMemorySampler {
14- pub fn new ( system : System ) -> Self {
15- Self { system }
24+ impl SystemInfo {
25+ pub fn new ( ) -> SystemInfo {
26+ Self {
27+ system : System :: new ( ) ,
28+ }
1629 }
30+ }
1731
32+ impl SystemInfoProvider for SystemInfo {
1833 /// This method picks the memory and cpu usage for this process using the
1934/// pid.
20- pub fn get_cpu_and_mem_used ( & mut self ) -> Result < ( f32 , u64 ) > {
21- let pid = sysinfo:: get_current_pid ( ) . map_err ( TelemetryError :: CannotGetPid ) ?;
35+ fn refresh_metrics ( & mut self , pid : Pid ) {
2236 self . system . refresh_pids_specifics (
2337 & [ pid] ,
2438 ProcessRefreshKind :: new ( )
2539 . with_cpu ( )
2640 . with_memory ( )
2741 . with_disk_usage ( ) ,
2842 ) ;
43+ }
2944
30- let process = self
31- . system
32- . process ( pid)
33- . unwrap_or_else ( || panic ! ( "cannot get process with pid: {}" , pid) ) ;
45+ fn get_pid ( & self ) -> Result < Pid , & ' static str > {
46+ sysinfo:: get_current_pid ( )
47+ }
48+
49+ fn get_process_specific_metrics < ' a > ( & self , pid : Pid ) -> Option < ( f32 , u64 ) > {
50+ let process = self . system . process ( pid) ?;
3451
35- let memory_used = process. memory ( ) ;
3652 let cpu_used = process. cpu_usage ( ) ;
53+ let memory_used = process. memory ( ) ;
54+ Some ( ( cpu_used, memory_used) )
55+ }
56+ }
57+
58+ struct CpuAndMemorySampler {
59+ system : Box < dyn SystemInfoProvider > ,
60+ }
3761
62+ impl CpuAndMemorySampler {
63+ pub fn new ( system : impl SystemInfoProvider ) -> Self {
64+ Self {
65+ system : Box :: new ( system) ,
66+ }
67+ }
68+
69+ pub fn get_cpu_and_mem_used ( & mut self ) -> Option < ( f32 , u64 ) > {
70+ let pid = self . system . get_pid ( ) . ok ( ) ?;
71+ self . system . refresh_metrics ( pid) ;
72+ let ( cpu_used, memory_used) = self . system . get_process_specific_metrics ( pid) ?;
3873 debug ! (
39- mem_used = ?memory_used,
4074 cpu_used = ?cpu_used,
75+ mem_used = ?memory_used,
4176 "trying to sample data for cpu/memory" ) ;
4277
43- Ok ( ( cpu_used, memory_used) )
78+ Some ( ( cpu_used, memory_used) )
4479 }
4580}
4681
@@ -49,18 +84,81 @@ pub(crate) async fn sample_metrics(
4984 duration_secs : Duration ,
5085) -> tokio:: task:: JoinHandle < ( ) > {
5186 tokio:: spawn ( async move {
52- let mut sampler = CpuAndMemorySampler :: new ( System :: new ( ) ) ;
87+ let mut sampler = CpuAndMemorySampler :: new ( SystemInfo :: new ( ) ) ;
5388
5489 // sample every minute
5590 let mut interval = tokio:: time:: interval ( duration_secs) ;
5691 interval. set_missed_tick_behavior ( tokio:: time:: MissedTickBehavior :: Skip ) ;
5792
5893 loop {
5994 interval. tick ( ) . await ;
60- if let Ok ( ( cpu_used, memory_used) ) = sampler. get_cpu_and_mem_used ( ) {
61- store. add_cpu_and_memory ( cpu_used, memory_used) ;
62- store. rollup_events ( ) ;
63- }
95+ sample_all_metrics ( & mut sampler, & store) ;
6496 }
6597 } )
6698}
99+
100+ fn sample_all_metrics ( sampler : & mut CpuAndMemorySampler , store : & Arc < TelemetryStore > ) {
101+ if let Some ( ( cpu_used, memory_used) ) = sampler. get_cpu_and_mem_used ( ) {
102+ store. add_cpu_and_memory ( cpu_used, memory_used) ;
103+ } else {
104+ debug ! ( "Cannot get cpu/mem usage stats for this process" ) ;
105+ }
106+ store. rollup_events ( ) ;
107+ }
108+
109+ #[ cfg( test) ]
110+ mod tests {
111+
112+ use crate :: ParquetMetrics ;
113+
114+ use super :: * ;
115+
116+ #[ derive( Debug ) ]
117+ struct MockParquetMetrics ;
118+
119+ impl ParquetMetrics for MockParquetMetrics {
120+ fn get_metrics ( & self ) -> ( u64 , f64 , u64 ) {
121+ ( 10 , 20.0 , 30 )
122+ }
123+ }
124+
125+ #[ test]
126+ fn test_sample_all_metrics ( ) {
127+ let mut mock_sys_info_provider = MockSystemInfoProvider :: new ( ) ;
128+ let store = TelemetryStore :: new_without_background_runners ( Arc :: from ( MockParquetMetrics ) ) ;
129+
130+ mock_sys_info_provider
131+ . expect_get_pid ( )
132+ . return_const ( Ok ( Pid :: from ( 5 ) ) ) ;
133+ mock_sys_info_provider
134+ . expect_refresh_metrics ( )
135+ . return_const ( ( ) ) ;
136+ mock_sys_info_provider
137+ . expect_get_process_specific_metrics ( )
138+ . return_const ( Some ( ( 10.0f32 , 100u64 ) ) ) ;
139+
140+ let mut sampler = CpuAndMemorySampler :: new ( mock_sys_info_provider) ;
141+
142+ sample_all_metrics ( & mut sampler, & store) ;
143+ }
144+
145+ #[ test]
146+ fn test_sample_all_metrics_with_call_failure ( ) {
147+ let mut mock_sys_info_provider = MockSystemInfoProvider :: new ( ) ;
148+ let store = TelemetryStore :: new_without_background_runners ( Arc :: from ( MockParquetMetrics ) ) ;
149+
150+ mock_sys_info_provider
151+ . expect_get_pid ( )
152+ . return_const ( Ok ( Pid :: from ( 5 ) ) ) ;
153+ mock_sys_info_provider
154+ . expect_refresh_metrics ( )
155+ . return_const ( ( ) ) ;
156+ mock_sys_info_provider
157+ . expect_get_process_specific_metrics ( )
158+ . return_const ( None ) ;
159+
160+ let mut sampler = CpuAndMemorySampler :: new ( mock_sys_info_provider) ;
161+
162+ sample_all_metrics ( & mut sampler, & store) ;
163+ }
164+ }
0 commit comments