Skip to main content

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        // Sort by `global_name` so the emitted JSON is deterministic — the
75        // underlying `FxDashMap` has unspecified iteration order. The map is
76        // small (~1500 entries in practice), so the sort cost is negligible
77        // and not worth optimizing.
78        let mut entries: Vec<_> = self
79            .inner
80            .iter()
81            .map(|e| (e.key().ty.global_name, e.value().clone()))
82            .collect();
83        entries.sort_unstable_by_key(|(name, _)| *name);
84        let mut map = serializer.serialize_map(Some(entries.len()))?;
85        for (name, stats) in &entries {
86            map.serialize_entry(name, stats)?;
87        }
88        map.end()
89    }
90}