turbo_tasks/
capture_future.rs

1use std::{
2    borrow::Cow,
3    cell::RefCell,
4    fmt::Display,
5    future::Future,
6    panic,
7    pin::Pin,
8    task::{Context, Poll},
9    time::{Duration, Instant},
10};
11
12use anyhow::Result;
13use pin_project_lite::pin_project;
14use serde::{Deserialize, Serialize};
15use turbo_tasks_malloc::{AllocationInfo, TurboMalloc};
16
17use crate::{backend::TurboTasksExecutionErrorMessage, panic_hooks::LAST_ERROR_LOCATION};
18
19struct ThreadLocalData {
20    duration: Duration,
21    allocations: usize,
22    deallocations: usize,
23}
24
25thread_local! {
26    static EXTRA: RefCell<Option<*mut ThreadLocalData>> = const { RefCell::new(None) };
27}
28
29pin_project! {
30    pub struct CaptureFuture<T, F: Future<Output = T>> {
31        #[pin]
32        future: F,
33        duration: Duration,
34        allocations: AllocationInfo,
35    }
36}
37
38impl<T, F: Future<Output = T>> CaptureFuture<T, F> {
39    pub fn new(future: F) -> Self {
40        Self {
41            future,
42            duration: Duration::ZERO,
43            allocations: AllocationInfo::ZERO,
44        }
45    }
46}
47
48fn try_with_thread_local_data(f: impl FnOnce(&mut ThreadLocalData)) {
49    EXTRA.with_borrow(|cell| {
50        if let Some(data) = cell {
51            // Safety: This data is thread local and only accessed in this thread
52            unsafe {
53                f(&mut **data);
54            }
55        }
56    });
57}
58
59pub fn add_duration(duration: Duration) {
60    try_with_thread_local_data(|data| {
61        data.duration += duration;
62    });
63}
64
65pub fn add_allocation_info(alloc_info: AllocationInfo) {
66    try_with_thread_local_data(|data| {
67        data.allocations += alloc_info.allocations;
68        data.deallocations += alloc_info.deallocations;
69    });
70}
71
72#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
73pub struct TurboTasksPanic {
74    pub message: TurboTasksExecutionErrorMessage,
75    pub location: Option<String>,
76}
77
78impl TurboTasksPanic {
79    pub fn into_panic(self) -> Box<dyn std::any::Any + Send> {
80        Box::new(format!(
81            "{} at {}",
82            self.message,
83            self.location
84                .unwrap_or_else(|| "unknown location".to_string())
85        ))
86    }
87}
88
89impl Display for TurboTasksPanic {
90    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91        write!(f, "{}", self.message)
92    }
93}
94
95impl<T, F: Future<Output = T>> Future for CaptureFuture<T, F> {
96    type Output = (Result<T, TurboTasksPanic>, Duration, AllocationInfo);
97
98    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
99        let this = self.project();
100        let start = Instant::now();
101        let start_allocations = TurboMalloc::allocation_counters();
102        let guard = ThreadLocalDataDropGuard;
103        let mut data = ThreadLocalData {
104            duration: Duration::ZERO,
105            allocations: 0,
106            deallocations: 0,
107        };
108        EXTRA.with_borrow_mut(|cell| {
109            *cell = Some(&mut data as *mut ThreadLocalData);
110        });
111
112        let result =
113            panic::catch_unwind(panic::AssertUnwindSafe(|| this.future.poll(cx))).map_err(|err| {
114                let message = match err.downcast_ref::<&'static str>() {
115                    Some(s) => TurboTasksExecutionErrorMessage::PIISafe(Cow::Borrowed(s)),
116                    None => match err.downcast_ref::<String>() {
117                        Some(s) => TurboTasksExecutionErrorMessage::NonPIISafe(s.clone()),
118                        None => {
119                            let error_message = err
120                                .downcast_ref::<Box<dyn Display>>()
121                                .map(|e| e.to_string())
122                                .unwrap_or_else(|| String::from("<unknown panic>"));
123
124                            TurboTasksExecutionErrorMessage::NonPIISafe(error_message)
125                        }
126                    },
127                };
128
129                LAST_ERROR_LOCATION.with_borrow(|loc| TurboTasksPanic {
130                    message,
131                    location: loc.clone(),
132                })
133            });
134
135        drop(guard);
136        let elapsed = start.elapsed();
137        let allocations = start_allocations.until_now();
138        *this.duration += elapsed + data.duration;
139        *this.allocations += allocations;
140        match result {
141            Err(err) => Poll::Ready((Err(err), *this.duration, this.allocations.clone())),
142            Ok(Poll::Ready(r)) => Poll::Ready((Ok(r), *this.duration, this.allocations.clone())),
143            Ok(Poll::Pending) => Poll::Pending,
144        }
145    }
146}
147
148struct ThreadLocalDataDropGuard;
149
150impl Drop for ThreadLocalDataDropGuard {
151    fn drop(&mut self) {
152        EXTRA.with_borrow_mut(|cell| {
153            *cell = None;
154        });
155    }
156}
turbo_tasks/capture_future.rs

turbo_tasks/
capture_future.rs
