turbo_tasks/
capture_future.rs

1use std::{
2    cell::RefCell,
3    fmt::Display,
4    future::Future,
5    panic,
6    pin::Pin,
7    task::{Context, Poll},
8    time::{Duration, Instant},
9};
10
11use anyhow::Result;
12use pin_project_lite::pin_project;
13use turbo_tasks_malloc::{AllocationInfo, TurboMalloc};
14
15use crate::{LAST_ERROR_LOCATION, backend::TurboTasksExecutionErrorMessage};
16
17struct ThreadLocalData {
18    duration: Duration,
19    allocations: usize,
20    deallocations: usize,
21}
22
23thread_local! {
24    static EXTRA: RefCell<Option<*mut ThreadLocalData>> = const { RefCell::new(None) };
25}
26
27pin_project! {
28    pub struct CaptureFuture<T, F: Future<Output = T>> {
29        #[pin]
30        future: F,
31        duration: Duration,
32        allocations: usize,
33        deallocations: usize,
34    }
35}
36
37impl<T, F: Future<Output = T>> CaptureFuture<T, F> {
38    pub fn new(future: F) -> Self {
39        Self {
40            future,
41            duration: Duration::ZERO,
42            allocations: 0,
43            deallocations: 0,
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)]
73pub struct TurboTasksPanic {
74    pub message: TurboTasksExecutionErrorMessage,
75    pub location: Option<String>,
76}
77
78impl Display for TurboTasksPanic {
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        write!(f, "{}", self.message)
81    }
82}
83
84impl<T, F: Future<Output = T>> Future for CaptureFuture<T, F> {
85    type Output = (Result<T, TurboTasksPanic>, Duration, usize);
86
87    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
88        let this = self.project();
89        let start = Instant::now();
90        let start_allocations = TurboMalloc::allocation_counters();
91        let guard = ThreadLocalDataDropGuard;
92        let mut data = ThreadLocalData {
93            duration: Duration::ZERO,
94            allocations: 0,
95            deallocations: 0,
96        };
97        EXTRA.with_borrow_mut(|cell| {
98            *cell = Some(&mut data as *mut ThreadLocalData);
99        });
100
101        let result =
102            panic::catch_unwind(panic::AssertUnwindSafe(|| this.future.poll(cx))).map_err(|err| {
103                let message = match err.downcast_ref::<&'static str>() {
104                    Some(s) => TurboTasksExecutionErrorMessage::PIISafe(s),
105                    None => match err.downcast_ref::<String>() {
106                        Some(s) => TurboTasksExecutionErrorMessage::NonPIISafe(s.clone()),
107                        None => {
108                            let error_message = err
109                                .downcast_ref::<Box<dyn Display>>()
110                                .map(|e| e.to_string())
111                                .unwrap_or_else(|| String::from("<unknown panic>"));
112
113                            TurboTasksExecutionErrorMessage::NonPIISafe(error_message)
114                        }
115                    },
116                };
117
118                LAST_ERROR_LOCATION.with_borrow(|loc| TurboTasksPanic {
119                    message,
120                    location: loc.clone(),
121                })
122            });
123
124        drop(guard);
125        let elapsed = start.elapsed();
126        let allocations = start_allocations.until_now();
127        *this.duration += elapsed + data.duration;
128        *this.allocations += allocations.allocations + data.allocations;
129        *this.deallocations += allocations.deallocations + data.deallocations;
130        match result {
131            Err(err) => {
132                let memory_usage = this.allocations.saturating_sub(*this.deallocations);
133                Poll::Ready((Err(err), *this.duration, memory_usage))
134            }
135            Ok(Poll::Ready(r)) => {
136                let memory_usage = this.allocations.saturating_sub(*this.deallocations);
137                Poll::Ready((Ok(r), *this.duration, memory_usage))
138            }
139            Ok(Poll::Pending) => Poll::Pending,
140        }
141    }
142}
143
144struct ThreadLocalDataDropGuard;
145
146impl Drop for ThreadLocalDataDropGuard {
147    fn drop(&mut self) {
148        EXTRA.with_borrow_mut(|cell| {
149            *cell = None;
150        });
151    }
152}
turbo_tasks/capture_future.rs

turbo_tasks/
capture_future.rs
