timetrace 0.2.2

Lightweight Rust profiling library with RAII and Chrome/Perfetto trace support
Documentation
#[cfg(feature = "json")]
use std::fs;
#[cfg(feature = "json")]
use std::fs::File;
#[cfg(feature = "json")]
use std::io::Write;
#[cfg(feature = "json")]
use std::time::Instant;

#[cfg(feature = "json")]
use crate::backends::base::TracingBackend;
#[cfg(feature = "json")]
use crate::profile_result::ProfileResult;

#[cfg(feature = "json")]
pub struct JsonProfiler {
    current_session: Option<String>,
    file: Option<File>,
    profile_count: usize,

    max_frames: Option<usize>,
    pub frame_count: usize,

    max_duration_ms: Option<u64>,
    start_time: Option<Instant>,
}

#[cfg(feature = "json")]
impl JsonProfiler {
    pub fn new() -> Self {
        Self {
            current_session: None,
            file: None,
            profile_count: 0,

            max_frames: None,
            frame_count: 0,

            max_duration_ms: None,
            start_time: None,
        }
    }

    pub fn write_profile(&mut self, result: &ProfileResult) {
        if self.current_session.is_none() {
            return;
        }

        if let Some(file) = &mut self.file {
            if self.profile_count > 0 {
                file.write_all(b",").unwrap();
            }

            self.profile_count += 1;

            let thread_name = result
                .args
                .as_ref()
                .and_then(|a| a.get("thread_name"))
                .and_then(|v| v.as_str())
                .unwrap_or("unnamed")
                .to_string();

            let mut result = result.clone();

            result.args = Some(serde_json::json!({
                "frame": self.frame_count,
                "thread_name": thread_name,
            }));

            let json = serde_json::to_string(&result).unwrap();
            file.write_all(json.as_bytes()).unwrap();
            file.flush().unwrap();
        }
    }

    fn write_header(&mut self) {
        if let Some(file) = &mut self.file {
            file.write_all(b"{\"otherData\": {},\"traceEvents\":[")
                .unwrap();
            file.flush().unwrap();
        }
    }

    fn write_footer(&mut self) {
        if let Some(file) = &mut self.file {
            file.write_all(b"]}").unwrap();
            file.flush().unwrap();
        }
    }
}

#[cfg(feature = "json")]
impl TracingBackend for JsonProfiler {
    fn begin_session_limited(
        &mut self,
        name: &str,
        filepath: &str,
        max_frames: Option<usize>,
        max_duration_ms: Option<u64>,
    ) {
        self.begin_session(name, filepath);

        self.max_frames = max_frames;
        self.max_duration_ms = max_duration_ms;
        self.frame_count = 0;
        self.start_time = Some(Instant::now());
    }

    fn begin_session(&mut self, name: &str, filepath: &str) {
        if self.current_session.is_some() {
            self.end_session().unwrap();
        }

        if let Some(parent) = std::path::Path::new(filepath).parent() {
            fs::create_dir_all(parent).unwrap();
        }

        let file = File::create(filepath).unwrap();

        self.file = Some(file);
        self.current_session = Some(name.to_string());
        self.profile_count = 0;

        self.write_header();
    }

    fn end_session(&mut self) -> Result<(), String> {
        if self.current_session.is_none() {
            return Ok(());
        }

        self.write_footer();

        self.file = None;
        self.current_session = None;
        self.profile_count = 0;

        Ok(())
    }

    fn new_frame(&mut self) {
        self.frame_count += 1;

        if let Some(max) = self.max_frames {
            if self.frame_count >= max {
                self.end_session().unwrap();
                return;
            }
        }

        if let Some(max_ms) = self.max_duration_ms {
            if let Some(start) = self.start_time {
                let elapsed = start.elapsed().as_millis() as u64;

                if elapsed >= max_ms {
                    self.end_session().unwrap();
                }
            }
        }
    }

    fn record_span(&mut self, _name: &str) {}

    fn record_event(&mut self, event: &ProfileResult) {
        self.write_profile(event);
    }
}