11use std:: {
22 borrow:: Cow ,
3- cell:: RefCell ,
43 fmt:: Display ,
54 future:: Future ,
65 panic,
76 pin:: Pin ,
87 task:: { Context , Poll } ,
9- time:: { Duration , Instant } ,
108} ;
119
1210use anyhow:: Result ;
1311use pin_project_lite:: pin_project;
1412use serde:: { Deserialize , Serialize } ;
15- use turbo_tasks_malloc:: { AllocationInfo , TurboMalloc } ;
1613
1714use crate :: { backend:: TurboTasksExecutionErrorMessage , panic_hooks:: LAST_ERROR_LOCATION } ;
1815
19- struct ThreadLocalData {
20- duration : Duration ,
21- allocations : usize ,
22- deallocations : usize ,
23- }
24-
25- thread_local ! {
26- static EXTRA : RefCell <Option <* mut ThreadLocalData >> = const { RefCell :: new( None ) } ;
27- }
28-
2916pin_project ! {
3017 pub struct CaptureFuture <T , F : Future <Output = T >> {
3118 #[ pin]
3219 future: F ,
33- duration: Duration ,
34- allocations: AllocationInfo ,
3520 }
3621}
3722
3823impl < T , F : Future < Output = T > > CaptureFuture < T , F > {
3924 pub fn new ( future : F ) -> Self {
40- Self {
41- future,
42- duration : Duration :: ZERO ,
43- allocations : AllocationInfo :: ZERO ,
44- }
25+ Self { future }
4526 }
4627}
4728
48- fn 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-
59- pub fn add_duration ( duration : Duration ) {
60- try_with_thread_local_data ( |data| {
61- data. duration += duration;
62- } ) ;
63- }
64-
65- pub 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-
7229#[ derive( Debug , Clone , Serialize , Deserialize , PartialEq , Eq ) ]
7330pub struct TurboTasksPanic {
7431 pub message : TurboTasksExecutionErrorMessage ,
@@ -93,21 +50,10 @@ impl Display for TurboTasksPanic {
9350}
9451
9552impl < T , F : Future < Output = T > > Future for CaptureFuture < T , F > {
96- type Output = ( Result < T , TurboTasksPanic > , Duration , AllocationInfo ) ;
53+ type Output = Result < T , TurboTasksPanic > ;
9754
9855 fn poll ( self : Pin < & mut Self > , cx : & mut Context < ' _ > ) -> Poll < Self :: Output > {
9956 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- } ) ;
11157
11258 let result =
11359 panic:: catch_unwind ( panic:: AssertUnwindSafe ( || this. future . poll ( cx) ) ) . map_err ( |err| {
@@ -132,25 +78,10 @@ impl<T, F: Future<Output = T>> Future for CaptureFuture<T, F> {
13278 } )
13379 } ) ;
13480
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;
14081 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 ( ) ) ) ,
82+ Err ( err) => Poll :: Ready ( Err ( err) ) ,
83+ Ok ( Poll :: Ready ( r) ) => Poll :: Ready ( Ok ( r) ) ,
14384 Ok ( Poll :: Pending ) => Poll :: Pending ,
14485 }
14586 }
14687}
147-
148- struct ThreadLocalDataDropGuard ;
149-
150- impl Drop for ThreadLocalDataDropGuard {
151- fn drop ( & mut self ) {
152- EXTRA . with_borrow_mut ( |cell| {
153- * cell = None ;
154- } ) ;
155- }
156- }
0 commit comments