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

turbo_tasks/
capture_future.rs
