turbo_tasks/
task_statistics.rs

1use std::sync::{Arc, OnceLock};
2
3use serde::{Serialize, Serializer, ser::SerializeMap};
4
5use crate::{FxDashMap, macro_helpers::NativeFunction};
6
7/// An API for optionally enabling, updating, and reading aggregated statistics.
8#[derive(Default)]
9pub struct TaskStatisticsApi {
10    inner: OnceLock<Arc<TaskStatistics>>,
11}
12
13impl TaskStatisticsApi {
14    pub fn enable(&self) -> &Arc<TaskStatistics> {
15        self.inner.get_or_init(|| {
16            Arc::new(TaskStatistics {
17                inner: FxDashMap::with_hasher(Default::default()),
18            })
19        })
20    }
21
22    // Calls `func` if statistics have been enabled (via
23    // [`TaskStatisticsApi::enable`]).
24    pub fn map<T>(&self, func: impl FnOnce(&Arc<TaskStatistics>) -> T) -> Option<T> {
25        self.get().map(func)
26    }
27
28    // Returns the statistics if they have been enabled (via
29    // [`TaskStatisticsApi::enable`]).
30    pub fn get(&self) -> Option<&Arc<TaskStatistics>> {
31        self.inner.get()
32    }
33}
34
35/// A type representing the enabled state of [`TaskStatisticsApi`]. Implements [`serde::Serialize`].
36pub struct TaskStatistics {
37    inner: FxDashMap<&'static NativeFunction, TaskFunctionStatistics>,
38}
39
40impl TaskStatistics {
41    pub fn increment_cache_hit(&self, native_fn: &'static NativeFunction) {
42        self.with_task_type_statistics(native_fn, |stats| stats.cache_hit += 1)
43    }
44
45    pub fn increment_cache_miss(&self, native_fn: &'static NativeFunction) {
46        self.with_task_type_statistics(native_fn, |stats| stats.cache_miss += 1)
47    }
48
49    fn with_task_type_statistics(
50        &self,
51        native_fn: &'static NativeFunction,
52        func: impl Fn(&mut TaskFunctionStatistics),
53    ) {
54        func(self.inner.entry(native_fn).or_default().value_mut())
55    }
56
57    pub fn get(&self, f: &'static NativeFunction) -> TaskFunctionStatistics {
58        self.inner.get(f).unwrap().value().clone()
59    }
60}
61
62/// Statistics for an individual function.
63#[derive(Default, Serialize, Clone)]
64pub struct TaskFunctionStatistics {
65    pub cache_hit: u32,
66    pub cache_miss: u32,
67}
68
69impl Serialize for TaskStatistics {
70    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
71    where
72        S: Serializer,
73    {
74        let mut map = serializer.serialize_map(Some(self.inner.len()))?;
75        for entry in &self.inner {
76            map.serialize_entry(entry.key().global_name, entry.value())?;
77        }
78        map.end()
79    }
80}