resource_tracker/
main.rs

1#![warn(clippy::pedantic)]
2#![doc = include_str!("../README.md")]
3
4#[cfg(not(target_os = "linux"))]
5compile_error!(
6    "resource-tracker only supports Linux; /proc and cgroup interfaces are Linux-specific."
7);
8
9mod collector;
10mod config;
11mod metrics;
12mod output;
13mod sentinel;
14mod thread_util;
15
16extern crate libc;
17
18use collector::{
19    CpuCollector, DiskCollector, GpuCollector, MemoryCollector, NetworkCollector,
20    collect_host_info, spawn_cloud_discovery,
21};
22use config::{Config, OutputFormat};
23use metrics::CloudInfo;
24use metrics::Sample;
25use rune_redact;
26use sentinel::{BatchUploader, RunContext, SentinelClient, close_run, samples_to_csv, start_run};
27use std::fs::File;
28use std::io::{BufWriter, Write};
29use std::sync::atomic::{AtomicBool, Ordering};
30use std::sync::{Arc, Mutex};
31use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
32
33// ---------------------------------------------------------------------------
34// SIGTERM handler
35// ---------------------------------------------------------------------------
36//
37static SIGTERM_RECEIVED: AtomicBool = AtomicBool::new(false);
38
39extern "C" fn handle_sigterm(_: libc::c_int) {
40    SIGTERM_RECEIVED.store(true, Ordering::Relaxed);
41}
42
43// Install SIGTERM and SIGINT handlers so the binary can flush before exiting.
44// Both signals set the same flag and trigger the same graceful shutdown path.
45//
46fn setup_signal_handlers() {
47    unsafe {
48        libc::signal(
49            libc::SIGTERM,
50            handle_sigterm as *const () as libc::sighandler_t,
51        );
52        libc::signal(
53            libc::SIGINT,
54            handle_sigterm as *const () as libc::sighandler_t,
55        );
56    }
57}
58
59struct ResourceTracker {
60    config: Config,
61    out_file: Option<std::io::BufWriter<std::fs::File>>,
62    interval: Duration,
63
64    // Collectors
65    cpu: CpuCollector,
66    memory: MemoryCollector,
67    network: NetworkCollector,
68    disk: DiskCollector,
69    gpu: GpuCollector,
70
71    // Cloud and host info
72    host_info: metrics::HostInfo,
73    cloud_info: Option<CloudInfo>,
74    cloud_rx: Option<std::sync::mpsc::Receiver<CloudInfo>>,
75
76    // Child process
77    child: Option<std::process::Child>,
78
79    // Sentinel state
80    sentinel: Option<SentinelClient>,
81    run_ctx_arc: Option<Arc<Mutex<RunContext>>>,
82    sample_buffer: Option<Arc<Mutex<Vec<Sample>>>>,
83    upload_shutdown_flag: Option<Arc<AtomicBool>>,
84    upload_handle: Option<std::thread::JoinHandle<Vec<String>>>,
85
86    // Sample tracking
87    unflushed: Vec<Sample>,
88    prev_loop_start: Option<Instant>,
89}
90
91impl ResourceTracker {
92    fn new() -> Self {
93        let config = Config::load();
94        let out_file = Self::create_sink(&config);
95        let interval = Duration::from_secs(config.interval_secs);
96
97        // The CSV schema only has a column for aggregate steal time
98        let aggregate_cpu_steal = config.aggregate_cpu_steal || config.format == OutputFormat::Csv;
99
100        let cpu = CpuCollector::new(config.pid, aggregate_cpu_steal);
101        let memory = MemoryCollector::new();
102        let network = NetworkCollector::new();
103        let disk = DiskCollector::new(interval);
104        let gpu = GpuCollector::new();
105
106        // Collect static GPU info now so host discovery can derive GPU host fields.
107        let initial_gpus = gpu.collect().unwrap_or_default();
108
109        // Host discovery: fast, local, no I/O.
110        let host_info = collect_host_info(&initial_gpus);
111
112        // Warm-up: prime delta state in stateful collectors while cloud probes run
113        let cloud_rx = spawn_cloud_discovery();
114        let cloud_info = None;
115
116        Self {
117            config,
118            out_file,
119            interval,
120            cpu,
121            memory,
122            network,
123            disk,
124            gpu,
125            host_info,
126            cloud_info,
127            cloud_rx,
128            child: None,
129            sentinel: None,
130            run_ctx_arc: None,
131            sample_buffer: None,
132            upload_shutdown_flag: None,
133            upload_handle: None,
134            unflushed: Vec::new(),
135            prev_loop_start: None,
136        }
137    }
138
139    fn warmup_collectors(&mut self) {
140        let _ = self.cpu.collect();
141        let _ = self.network.collect();
142        let _ = self.disk.collect();
143    }
144
145    fn spawn_tracked_command(&mut self) {
146        let Some((program, args)) = self.config.command.split_first() else {
147            return;
148        };
149
150        match std::process::Command::new(program).args(args).spawn() {
151            Ok(c) => {
152                self.config.pid = Some(i32::try_from(c.id()).unwrap_or(i32::MAX));
153                self.cpu.set_tracked_pid(self.config.pid);
154                self.child = Some(c);
155            }
156
157            Err(e) => {
158                eprintln!("error: failed to spawn {:?}: {e}", program);
159                std::process::exit(1);
160            }
161        }
162    }
163
164    fn mask_sensitive_data_in_command(&mut self) {
165        for item in &mut self.config.metadata.command {
166            if let Some(redacted) = Self::try_redact(item) {
167                *item = redacted;
168            }
169        }
170    }
171
172    fn try_redact(raw: &str) -> Option<String> {
173        // such keys used to be stored in files, but we're watching
174        if raw.starts_with("-----BEGIN") {
175            return Some("[KEY]".to_string());
176        }
177
178        // missing rune_redact feature: check for ftp scheme
179        if raw.starts_with("ftp://") {
180            return Self::mask_first_word(raw, "[URL]").into();
181        }
182
183        // missing rune_redact feature: check for URL variables
184        if raw.starts_with("http://") || raw.starts_with("https://") {
185            if !raw.contains(".") || raw.contains("?") || raw.contains("&") || raw.contains("=") {
186                return Self::mask_first_word(raw, "[URL]");
187            }
188        }
189
190        let redacted = rune_redact::redact(raw);
191        (redacted != raw).then_some(redacted)
192    }
193
194    fn mask_first_word(raw: &str, mask: &str) -> Option<String> {
195        let end_pos = raw.find(' ').unwrap_or(raw.len());
196        Some(format!("{}{}", mask, &raw[end_pos..]).to_owned())
197    }
198
199    fn setup_sentinel(&mut self) {
200        self.sentinel = SentinelClient::from_env();
201
202        let Some(client) = &self.sentinel else {
203            return;
204        };
205
206        // Bounded wait: give cloud discovery a chance to complete
207        if self.cloud_info.is_none() {
208            if let Some(ref rx) = self.cloud_rx {
209                self.cloud_info = rx.recv_timeout(Duration::from_secs(3)).ok();
210            }
211        }
212
213        let default_cloud = CloudInfo::default();
214        let ctx = match start_run(
215            &client.agent,
216            &client.api_base,
217            &client.token,
218            &self.config.metadata,
219            self.config.pid,
220            &self.host_info,
221            self.cloud_info.as_ref().unwrap_or(&default_cloud),
222        ) {
223            Err(e) => {
224                eprintln!("warn: sentinel start_run failed: {e}; streaming disabled");
225                return;
226            }
227            Ok(ctx) => ctx,
228        };
229
230        let ctx_arc = Arc::new(Mutex::new(ctx));
231        let upload_interval = std::env::var("TRACKER_UPLOAD_INTERVAL")
232            .ok()
233            .and_then(|v| v.parse().ok())
234            .unwrap_or(60u64);
235        let (uploader, buf) = BatchUploader::new(upload_interval, self.config.interval_secs);
236        let flag = uploader.shutdown_flag();
237        let upload_handle = uploader.spawn(
238            Arc::clone(&ctx_arc),
239            SentinelClient::new_upload_agent(),
240            client.api_base.clone(),
241            client.token.clone(),
242        );
243        if upload_handle.is_none() {
244            eprintln!(
245                "warn: sentinel background upload disabled; samples will be flushed inline on exit"
246            );
247        }
248
249        self.run_ctx_arc = Some(ctx_arc);
250        self.sample_buffer = Some(buf);
251        self.upload_shutdown_flag = Some(flag);
252        self.upload_handle = upload_handle;
253    }
254
255    fn emit_csv_header(&mut self) {
256        if self.config.format == OutputFormat::Csv {
257            Self::emit_metric_line(&self.config, &mut self.out_file, output::csv::csv_header());
258        }
259    }
260
261    fn renice_tracker(&self) {
262        let Some(renice) = self.config.renice else {
263            return;
264        };
265
266        let result = unsafe { libc::setpriority(libc::PRIO_PROCESS, 0, renice) };
267        if result == -1 {
268            eprintln!("warn: failed to renice process, ignored");
269        }
270    }
271
272    fn poll_cloud_info(&mut self) {
273        if self.cloud_info.is_none()
274            && let Some(ref rx) = self.cloud_rx
275            && let Ok(info) = rx.try_recv()
276        {
277            self.cloud_info = Some(info);
278        }
279    }
280
281    fn collect_sample(&mut self) -> Sample {
282        let loop_start = Instant::now();
283
284        let actual_interval_ms: Option<u64> = self
285            .prev_loop_start
286            .map(|p| u64::try_from((loop_start - p).as_millis()).unwrap_or(u64::MAX));
287
288        let timestamp_secs = SystemTime::now()
289            .duration_since(UNIX_EPOCH)
290            .unwrap_or_default()
291            .as_secs();
292
293        let mut sample = Sample {
294            timestamp_secs,
295            actual_interval_ms,
296            job_name: self.config.metadata.job_name.clone(),
297            tracked_pid: self.config.pid,
298            cpu: self.cpu.collect().unwrap_or_default(),
299            memory: self.memory.collect().unwrap_or_default(),
300            network: self.network.collect().unwrap_or_default(),
301            disk: self.disk.collect().unwrap_or_default(),
302            gpu: self.gpu.collect().unwrap_or_default(),
303        };
304
305        // Augment with per-process GPU stats.
306        let (vram_mib, gpu_usage, gpu_utilized) =
307            if self.config.pid.is_some() && !sample.cpu.process_tree_pids.is_empty() {
308                let pids_u32: Vec<u32> = sample
309                    .cpu
310                    .process_tree_pids
311                    .iter()
312                    .filter_map(|&p| u32::try_from(p).ok())
313                    .collect();
314                self.gpu.process_gpu_info(&pids_u32, self.interval)
315            } else {
316                self.gpu.all_gpu_process_info(self.interval)
317            };
318        sample.cpu.process_gpu_vram_mib = vram_mib;
319        sample.cpu.process_gpu_usage = gpu_usage;
320        sample.cpu.process_gpu_utilized = gpu_utilized;
321
322        self.prev_loop_start = Some(loop_start);
323
324        sample
325    }
326
327    fn emit_sample(&mut self, sample: &Sample) {
328        match self.config.format {
329            OutputFormat::Json => match serde_json::to_value(sample) {
330                Ok(mut v) => {
331                    v[format!("{}-version", env!("CARGO_PKG_NAME"))] =
332                        serde_json::Value::String(env!("CARGO_PKG_VERSION").to_string());
333                    Self::emit_metric_line(&self.config, &mut self.out_file, &v.to_string());
334                }
335                Err(e) => eprintln!("warn: json serialize error: {e}"),
336            },
337
338            OutputFormat::Csv => {
339                Self::emit_metric_line(
340                    &self.config,
341                    &mut self.out_file,
342                    &output::csv::sample_to_csv_row(sample, self.config.interval_secs),
343                );
344            }
345        }
346    }
347
348    fn buffer_sample(&mut self, sample: Sample) {
349        // Push to sentinel buffer (if streaming is active).
350        if let Some(ref buf) = self.sample_buffer {
351            buf.lock()
352                .unwrap_or_else(|e| e.into_inner())
353                .push(sample.clone());
354        }
355        self.unflushed.push(sample);
356    }
357
358    fn check_child_exit(&mut self) -> Option<i32> {
359        let child = self.child.as_mut()?;
360
361        match child.try_wait() {
362            Ok(Some(status)) => Some(status.code().unwrap_or(1)),
363            Ok(None) => None,
364            Err(e) => {
365                eprintln!("warn: error checking child status: {e}");
366                None
367            }
368        }
369    }
370
371    fn check_signal(&self) -> bool {
372        SIGTERM_RECEIVED.load(Ordering::Relaxed)
373    }
374
375    fn sleep_until_next_interval(&self, loop_start: Instant) {
376        let elapsed = loop_start.elapsed();
377        if let Some(remaining) = self.interval.checked_sub(elapsed) {
378            std::thread::sleep(remaining);
379        }
380    }
381
382    fn shutdown(&mut self, exit_code: i32) -> ! {
383        // Take ownership of fields that need to be moved
384        let sentinel = self.sentinel.take();
385        let run_ctx = self.run_ctx_arc.take();
386        let shutdown_flag = self.upload_shutdown_flag.take();
387        let upload_handle = self.upload_handle.take();
388        let remaining = std::mem::take(&mut self.unflushed);
389
390        Self::graceful_shutdown(
391            exit_code,
392            sentinel.as_ref(),
393            run_ctx,
394            shutdown_flag,
395            upload_handle,
396            remaining,
397            self.config.interval_secs,
398        );
399    }
400
401    fn run(mut self) -> ! {
402        self.warmup_collectors();
403        std::thread::sleep(self.interval);
404
405        self.spawn_tracked_command();
406        self.mask_sensitive_data_in_command();
407        self.setup_sentinel();
408        self.emit_csv_header();
409        self.renice_tracker();
410
411        // Main sampling loop
412        loop {
413            self.poll_cloud_info();
414            let loop_start = Instant::now();
415
416            let sample = self.collect_sample();
417            self.emit_sample(&sample);
418            self.buffer_sample(sample);
419
420            if let Some(code) = self.check_child_exit() {
421                self.shutdown(code);
422            }
423            if self.check_signal() {
424                self.shutdown(0);
425            }
426
427            self.sleep_until_next_interval(loop_start);
428        }
429    }
430
431    // -----------------------------------------------------------------------
432    // Output sink: stdout (default), file (--output), or suppressed (--quiet).
433    // Warnings and errors always go to stderr via eprintln! regardless.
434    // -----------------------------------------------------------------------
435    //
436    fn create_sink(config: &Config) -> Option<BufWriter<File>> {
437        if config.quiet {
438            return None;
439        }
440
441        match config.output_file.as_deref() {
442            Some(path) => File::create(path).map(BufWriter::new).ok(),
443            None => None,
444        }
445    }
446
447    // ---------------------------------------------------------------------------
448    // Graceful shutdown
449    // ---------------------------------------------------------------------------
450    //
451    // Flush remaining samples, close the Sentinel run, then exit.
452    //
453    // Called on both shell-wrapper child exit and SIGTERM.  Replaces the former
454    // bare `std::process::exit()` calls so the upload thread always gets a chance
455    // to flush.
456    //
457    fn graceful_shutdown(
458        exit_code: i32,
459        sentinel: Option<&SentinelClient>,
460        run_ctx: Option<Arc<Mutex<RunContext>>>,
461        shutdown_flag: Option<Arc<AtomicBool>>,
462        upload_handle: Option<std::thread::JoinHandle<Vec<String>>>,
463        remaining: Vec<Sample>,
464        interval_secs: u64,
465    ) -> ! {
466        if let (Some(client), Some(ctx_arc), Some(flag), Some(handle)) =
467            (sentinel, run_ctx, shutdown_flag, upload_handle)
468        {
469            // Signal the upload thread to flush its buffer to S3, then wait for it.
470            // The thread performs one final S3 upload of any remaining buffered samples
471            // before it exits, and returns the list of all successfully uploaded URIs.
472            flag.store(true, Ordering::Relaxed);
473            let uploaded_uris = handle.join().unwrap_or_default();
474
475            // Route selection:
476            //   S3 route   -- at least one batch was uploaded; uploaded_uris is non-empty.
477            //                 The final flush is already included in uploaded_uris.
478            //   Inline route -- no S3 uploads (short run or all S3 failures); send all
479            //                   collected samples as a raw CSV string.
480            let remaining_csv = if uploaded_uris.is_empty() && !remaining.is_empty() {
481                Some(samples_to_csv(&remaining, interval_secs))
482            } else {
483                None
484            };
485
486            let ctx = ctx_arc.lock().unwrap_or_else(|e| e.into_inner());
487            if let Err(e) = close_run(
488                &client.agent,
489                &client.api_base,
490                &client.token,
491                &ctx,
492                Some(exit_code),
493                remaining_csv,
494                &uploaded_uris,
495            ) {
496                eprintln!("warn: sentinel close_run failed: {e}");
497            }
498        }
499
500        std::process::exit(exit_code);
501    }
502
503    // Emit one line of metric output to the selected sink.
504    // quiet=true  -> no-op
505    // output_file -> write to file and flush (so `tail -f` works)
506    // default     -> eprintln! to stderr (keeps stdout clean for the tracked app)
507    //
508    fn emit_metric_line(config: &Config, out_file: &mut Option<BufWriter<File>>, msg: &str) {
509        if config.quiet {
510            return;
511        }
512
513        match out_file {
514            Some(writer) => {
515                let _ = writeln!(writer, "{msg}");
516                let _ = writer.flush();
517            }
518            None => eprintln!("{msg}"),
519        }
520    }
521}
522
523// ---------------------------------------------------------------------------
524// main
525// ---------------------------------------------------------------------------
526//
527fn main() {
528    setup_signal_handlers();
529    let tracker = ResourceTracker::new();
530    tracker.run();
531}
532
533// ---------------------------------------------------------------------------
534// Tests
535// ---------------------------------------------------------------------------
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540
541    /// Verify that SIGINT sets SIGTERM_RECEIVED, triggering the same graceful
542    /// shutdown path as SIGTERM.  The test installs the handler, resets the
543    /// flag, raises SIGINT, then asserts the flag is true.
544    #[test]
545    fn test_sigint_sets_shutdown_flag() {
546        // Reset in case a previous test left the flag set.
547        SIGTERM_RECEIVED.store(false, Ordering::SeqCst);
548
549        // Install the handler for SIGINT (mirrors what main() does).
550        unsafe {
551            libc::signal(
552                libc::SIGINT,
553                handle_sigterm as *const () as libc::sighandler_t,
554            );
555        }
556
557        // Raise SIGINT on the current process.
558        unsafe {
559            libc::raise(libc::SIGINT);
560        }
561
562        assert!(
563            SIGTERM_RECEIVED.load(Ordering::SeqCst),
564            "SIGTERM_RECEIVED flag must be true after SIGINT"
565        );
566
567        // Clean up: reset the flag and restore the default SIGINT disposition
568        // so this does not interfere with other tests.
569        SIGTERM_RECEIVED.store(false, Ordering::SeqCst);
570        unsafe {
571            libc::signal(libc::SIGINT, libc::SIG_DFL);
572        }
573    }
574
575    fn test_redact(data: &str, contains: Option<&str>) {
576        let result = ResourceTracker::try_redact(data);
577
578        match (contains, result) {
579            (Some(expected), Some(redacted)) => {
580                assert!(
581                    redacted.contains(expected),
582                    "expected to be redacted, got: {redacted}"
583                );
584            }
585            (Some(_), None) => {
586                panic!("expected to be redacted, but not detected");
587            }
588            (None, Some(redacted)) => {
589                panic!("expected to be unchanged, got {redacted}");
590            }
591            (None, None) => (),
592        }
593    }
594
595    // redact: email
596
597    #[test]
598    fn test_redact_email_plain_good() {
599        test_redact("sample@example.com", Some("[EMAIL]"));
600    }
601
602    #[test]
603    fn test_redact_email_dot_in_username() {
604        test_redact("good.sample@example.com", Some("[EMAIL]"));
605    }
606
607    #[test]
608    fn test_redact_email_twitter_style() {
609        test_redact("@twitternick", None);
610    }
611
612    #[test]
613    fn test_redact_email_invalid_host() {
614        test_redact("nick@invalid_host.com", None);
615    }
616
617    // redact: URL
618
619    #[test]
620    fn test_redact_url_no_tld() {
621        test_redact("http://server04", Some("[URL]")); // reveals local machine name
622    }
623
624    #[test]
625    fn test_redact_url_http() {
626        test_redact("http://example.com", None); // not leaking any information
627    }
628
629    #[test]
630    fn test_redact_url_https() {
631        test_redact("https://example.com/path", None); // innocent
632    }
633
634    #[test]
635    fn test_redact_url_https_with_account() {
636        test_redact("https://nick@example.com/path", Some("[")); // both [URL] and [EMAIL] is okay
637    }
638
639    #[test]
640    fn test_redact_url_with_query_params() {
641        test_redact("https://example.com/page?q=search&lang=en", Some("[URL]"));
642    }
643
644    #[test]
645    fn test_redact_url_with_fragment() {
646        test_redact("https://example.com#section", None); // innocent
647    }
648
649    #[test]
650    fn test_redact_url_with_subdomain() {
651        test_redact("https://api.example.com/report/from/otherworld", None); // innocent
652    }
653
654    #[test]
655    fn test_redact_url_with_port() {
656        test_redact("https://localhost:8080/admin", Some("[URL]"));
657    }
658
659    #[test]
660    fn test_redact_url_ftp() {
661        test_redact("ftp://ftp.example.com/files", Some("[URL]"));
662    }
663
664    #[test]
665    fn test_redact_url_invalid_no_protocol() {
666        test_redact("example.com", None); // not a real URL
667    }
668
669    #[test]
670    fn test_redact_connection_string() {
671        let raw = "app.py --connection-string 'postgresql://username:ASDAD_32ejae32DWQdw2d2@foobar.db.provider.com:12345/db?sslmode=require'";
672        assert_eq!(
673            ResourceTracker::try_redact(raw).as_deref(),
674            Some("app.py --connection-string [SECRET]")
675        );
676    }
677
678    // redact: IP address
679
680    #[test]
681    fn test_redact_ipv4_standard() {
682        test_redact("192.168.1.1", Some("[IP]"));
683    }
684
685    #[test]
686    fn test_redact_ipv4_with_port() {
687        test_redact("192.168.1.1:8080", Some("[IP]"));
688    }
689
690    #[test]
691    fn test_redact_ipv4_all_zeros() {
692        test_redact("0.0.0.0", Some("[IP]"));
693    }
694
695    #[test]
696    fn test_redact_ipv4_loopback() {
697        test_redact("127.0.0.1", Some("[IP]"));
698    }
699
700    #[test]
701    fn test_redact_ipv4_broadcast() {
702        test_redact("255.255.255.255", Some("[IP]"));
703    }
704
705    #[test]
706    fn test_redact_ip_invalid_octet_overflow() {
707        test_redact("256.168.1.1", None);
708    }
709
710    #[test]
711    fn test_redact_ip_invalid_partial() {
712        test_redact("192.168.1", None);
713    }
714
715    // redact keys
716
717    #[test]
718    fn test_redact_ssl_private_key_rsa() {
719        test_redact(
720            "-----BEGIN RSA PRIVATE KEY-----\nMIIEowIBAAKCAQEA...\n-----END RSA PRIVATE KEY-----",
721            Some("[KEY]"),
722        );
723    }
724
725    #[test]
726    fn test_redact_ssl_cert() {
727        test_redact(
728            "-----BEGIN CERTIFICATE-----\nMIIEowIBAAKCAQEA...\n-----END CERTIFICATE-----",
729            Some("[KEY]"),
730        );
731    }
732
733    // token
734
735    #[test]
736    fn test_redact_possible_token() {
737        test_redact("ar4mNbYrVwZuAtJhCf7DgLeW2oI5qR8eMvXn", Some("[SECRET]"));
738    }
739}
resource_tracker/main.rs

resource_tracker/
main.rs
