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    pub fn is_enabled(&self) -> bool {
23        self.inner.get().is_some()
24    }
25
26    // Calls `func` if statistics have been enabled (via
27    // [`TaskStatisticsApi::enable`]).
28    pub fn map<T>(&self, func: impl FnOnce(&Arc<TaskStatistics>) -> T) -> Option<T> {
29        self.get().map(func)
30    }
31
32    // Calls `func` if statistics have been enabled (via
33    // [`TaskStatisticsApi::enable`]).
34    pub fn get(&self) -> Option<&Arc<TaskStatistics>> {
35        self.inner.get()
36    }
37}
38
39/// A type representing the enabled state of [`TaskStatisticsApi`]. Implements [`serde::Serialize`].
40pub struct TaskStatistics {
41    inner: FxDashMap<&'static NativeFunction, TaskFunctionStatistics>,
42}
43
44impl TaskStatistics {
45    pub fn increment_cache_hit(&self, native_fn: &'static NativeFunction) {
46        self.with_task_type_statistics(native_fn, |stats| stats.cache_hit += 1)
47    }
48
49    pub fn increment_cache_miss(&self, native_fn: &'static NativeFunction) {
50        self.with_task_type_statistics(native_fn, |stats| stats.cache_miss += 1)
51    }
52
53    fn with_task_type_statistics(
54        &self,
55        native_fn: &'static NativeFunction,
56        func: impl Fn(&mut TaskFunctionStatistics),
57    ) {
58        func(self.inner.entry(native_fn).or_default().value_mut())
59    }
60}
61
62/// Statistics for an individual function.
63#[derive(Default, Serialize)]
64struct TaskFunctionStatistics {
65    cache_hit: u32,
66    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}