336 lines · 9.8 KB · Rust
RawDownload
use scheduler::Instant;use std::{    cell::LazyCell,    collections::HashMap,    hash::Hasher,    hash::{DefaultHasher, Hash},    sync::Arc,    thread::ThreadId,}; use serde::{Deserialize, Serialize}; use crate::SharedString; #[doc(hidden)]#[derive(Debug, Copy, Clone)]pub struct TaskTiming {    pub location: &'static core::panic::Location<'static>,    pub start: Instant,    pub end: Option<Instant>,} #[doc(hidden)]#[derive(Debug, Clone)]pub struct ThreadTaskTimings {    pub thread_name: Option<String>,    pub thread_id: ThreadId,    pub timings: Vec<TaskTiming>,    pub total_pushed: u64,} impl ThreadTaskTimings {    /// Convert global thread timings into their structured format.    pub fn convert(timings: &[GlobalThreadTimings]) -> Vec<Self> {        timings            .iter()            .filter_map(|t| match t.timings.upgrade() {                Some(timings) => Some((t.thread_id, timings)),                _ => None,            })            .map(|(thread_id, timings)| {                let timings = timings.lock();                let thread_name = timings.thread_name.clone();                let total_pushed = timings.total_pushed;                let timings = &timings.timings;                 let mut vec = Vec::with_capacity(timings.len());                 let (s1, s2) = timings.as_slices();                vec.extend_from_slice(s1);                vec.extend_from_slice(s2);                 ThreadTaskTimings {                    thread_name,                    thread_id,                    timings: vec,                    total_pushed,                }            })            .collect()    }} /// Serializable variant of [`core::panic::Location`]#[derive(Debug, Clone, Serialize, Deserialize)]pub struct SerializedLocation {    /// Name of the source file    pub file: SharedString,    /// Line in the source file    pub line: u32,    /// Column in the source file    pub column: u32,} impl From<&core::panic::Location<'static>> for SerializedLocation {    fn from(value: &core::panic::Location<'static>) -> Self {        SerializedLocation {            file: value.file().into(),            line: value.line(),            column: value.column(),        }    }} /// Serializable variant of [`TaskTiming`]#[derive(Debug, Clone, Serialize, Deserialize)]pub struct SerializedTaskTiming {    /// Location of the timing    pub location: SerializedLocation,    /// Time at which the measurement was reported in nanoseconds    pub start: u128,    /// Duration of the measurement in nanoseconds    pub duration: u128,} impl SerializedTaskTiming {    /// Convert an array of [`TaskTiming`] into their serializable format    ///    /// # Params    ///    /// `anchor` - [`Instant`] that should be earlier than all timings to use as base anchor    pub fn convert(anchor: Instant, timings: &[TaskTiming]) -> Vec<SerializedTaskTiming> {        let serialized = timings            .iter()            .map(|timing| {                let start = timing.start.duration_since(anchor).as_nanos();                let duration = timing                    .end                    .unwrap_or_else(|| Instant::now())                    .duration_since(timing.start)                    .as_nanos();                SerializedTaskTiming {                    location: timing.location.into(),                    start,                    duration,                }            })            .collect::<Vec<_>>();         serialized    }} /// Serializable variant of [`ThreadTaskTimings`]#[derive(Debug, Clone, Serialize, Deserialize)]pub struct SerializedThreadTaskTimings {    /// Thread name    pub thread_name: Option<String>,    /// Hash of the thread id    pub thread_id: u64,    /// Timing records for this thread    pub timings: Vec<SerializedTaskTiming>,} impl SerializedThreadTaskTimings {    /// Convert [`ThreadTaskTimings`] into their serializable format    ///    /// # Params    ///    /// `anchor` - [`Instant`] that should be earlier than all timings to use as base anchor    pub fn convert(anchor: Instant, timings: ThreadTaskTimings) -> SerializedThreadTaskTimings {        let serialized_timings = SerializedTaskTiming::convert(anchor, &timings.timings);         let mut hasher = DefaultHasher::new();        timings.thread_id.hash(&mut hasher);        let thread_id = hasher.finish();         SerializedThreadTaskTimings {            thread_name: timings.thread_name,            thread_id,            timings: serialized_timings,        }    }} #[doc(hidden)]#[derive(Debug, Clone)]pub struct ThreadTimingsDelta {    /// Hashed thread id    pub thread_id: u64,    /// Thread name, if known    pub thread_name: Option<String>,    /// New timings since the last call. If the circular buffer wrapped around    /// since the previous poll, some entries may have been lost.    pub new_timings: Vec<SerializedTaskTiming>,} /// Tracks which timing events have already been seen so that callers can request only unseen events.#[doc(hidden)]pub struct ProfilingCollector {    startup_time: Instant,    cursors: HashMap<ThreadId, u64>,} impl ProfilingCollector {    pub fn new(startup_time: Instant) -> Self {        Self {            startup_time,            cursors: HashMap::default(),        }    }     pub fn startup_time(&self) -> Instant {        self.startup_time    }     pub fn collect_unseen(        &mut self,        all_timings: Vec<ThreadTaskTimings>,    ) -> Vec<ThreadTimingsDelta> {        let mut deltas = Vec::with_capacity(all_timings.len());         for thread in all_timings {            let mut hasher = DefaultHasher::new();            thread.thread_id.hash(&mut hasher);            let hashed_id = hasher.finish();             let prev_cursor = self.cursors.get(&thread.thread_id).copied().unwrap_or(0);            let buffer_len = thread.timings.len() as u64;            let buffer_start = thread.total_pushed.saturating_sub(buffer_len);             let mut slice = if prev_cursor < buffer_start {                // Cursor fell behind the buffer — some entries were evicted.                // Return everything still in the buffer.                thread.timings.as_slice()            } else {                let skip = (prev_cursor - buffer_start) as usize;                &thread.timings[skip.min(thread.timings.len())..]            };             // Don't emit the last entry if it's still in-progress (end: None).            let incomplete_at_end = slice.last().is_some_and(|t| t.end.is_none());            if incomplete_at_end {                slice = &slice[..slice.len() - 1];            }             let cursor_advance = if incomplete_at_end {                thread.total_pushed.saturating_sub(1)            } else {                thread.total_pushed            };             self.cursors.insert(thread.thread_id, cursor_advance);             if slice.is_empty() {                continue;            }             let new_timings = SerializedTaskTiming::convert(self.startup_time, slice);             deltas.push(ThreadTimingsDelta {                thread_id: hashed_id,                thread_name: thread.thread_name,                new_timings,            });        }         deltas    }     pub fn reset(&mut self) {        self.cursors.clear();    }} // Allow 20mb of task timing entriesconst MAX_TASK_TIMINGS: usize = (20 * 1024 * 1024) / core::mem::size_of::<TaskTiming>(); #[doc(hidden)]pub type TaskTimings = circular_buffer::CircularBuffer<MAX_TASK_TIMINGS, TaskTiming>;#[doc(hidden)]pub type GuardedTaskTimings = spin::Mutex<ThreadTimings>; #[doc(hidden)]pub struct GlobalThreadTimings {    pub thread_id: ThreadId,    pub timings: std::sync::Weak<GuardedTaskTimings>,} #[doc(hidden)]pub static GLOBAL_THREAD_TIMINGS: spin::Mutex<Vec<GlobalThreadTimings>> =    spin::Mutex::new(Vec::new()); thread_local! {    #[doc(hidden)]    pub static THREAD_TIMINGS: LazyCell<Arc<GuardedTaskTimings>> = LazyCell::new(|| {        let current_thread = std::thread::current();        let thread_name = current_thread.name();        let thread_id = current_thread.id();        let timings = ThreadTimings::new(thread_name.map(|e| e.to_string()), thread_id);        let timings = Arc::new(spin::Mutex::new(timings));         {            let timings = Arc::downgrade(&timings);            let global_timings = GlobalThreadTimings {                thread_id: std::thread::current().id(),                timings,            };            GLOBAL_THREAD_TIMINGS.lock().push(global_timings);        }         timings    });} #[doc(hidden)]pub struct ThreadTimings {    pub thread_name: Option<String>,    pub thread_id: ThreadId,    pub timings: Box<TaskTimings>,    pub total_pushed: u64,} impl ThreadTimings {    pub fn new(thread_name: Option<String>, thread_id: ThreadId) -> Self {        ThreadTimings {            thread_name,            thread_id,            timings: TaskTimings::boxed(),            total_pushed: 0,        }    }} impl Drop for ThreadTimings {    fn drop(&mut self) {        let mut thread_timings = GLOBAL_THREAD_TIMINGS.lock();         let Some((index, _)) = thread_timings            .iter()            .enumerate()            .find(|(_, t)| t.thread_id == self.thread_id)        else {            return;        };        thread_timings.swap_remove(index);    }} #[doc(hidden)]#[allow(dead_code)] // Used by Linux and Windows dispatchers, not macOSpub fn add_task_timing(timing: TaskTiming) {    THREAD_TIMINGS.with(|timings| {        let mut timings = timings.lock();         if let Some(last_timing) = timings.timings.back_mut() {            if last_timing.location == timing.location && last_timing.start == timing.start {                last_timing.end = timing.end;                return;            }        }         timings.timings.push_back(timing);        timings.total_pushed += 1;    });}