|
1 | | -use std::cell::UnsafeCell; |
| 1 | +use std::cell::{Cell,UnsafeCell}; |
2 | 2 | use std::mem::{self, ManuallyDrop}; |
3 | 3 | use std::panic::{self, AssertUnwindSafe, UnwindSafe}; |
4 | 4 | use std::pin::Pin; |
5 | 5 | use std::sync::Arc; |
6 | 6 | use std::task::{RawWaker, RawWakerVTable}; |
7 | | -use std::thread::{self, Thread}; |
| 7 | +use std::thread; |
| 8 | + |
| 9 | +use crossbeam_utils::sync::Parker; |
8 | 10 |
|
9 | 11 | use super::task; |
10 | 12 | use super::task_local; |
@@ -119,46 +121,56 @@ where |
119 | 121 | F: Future<Output = T>, |
120 | 122 | { |
121 | 123 | thread_local! { |
122 | | - static ARC_THREAD: Arc<Thread> = Arc::new(thread::current()); |
| 124 | + // May hold a pre-allocated parker that can be reused for efficiency. |
| 125 | + // |
| 126 | + // Note that each invocation of `block` needs its own parker. In particular, if `block` |
| 127 | + // recursively calls itself, we must make sure that each recursive call uses a distinct |
| 128 | + // parker instance. |
| 129 | + static CACHE: Cell<Option<Arc<Parker>>> = Cell::new(None); |
123 | 130 | } |
124 | 131 |
|
125 | 132 | pin_utils::pin_mut!(f); |
126 | 133 |
|
127 | | - ARC_THREAD.with(|arc_thread: &Arc<Thread>| { |
128 | | - let ptr = (&**arc_thread as *const Thread) as *const (); |
| 134 | + CACHE.with(|cache| { |
| 135 | + // Reuse a cached parker or create a new one for this invocation of `block`. |
| 136 | + let arc_parker: Arc<Parker> = cache.take().unwrap_or_else(|| Arc::new(Parker::new())); |
| 137 | + |
| 138 | + let ptr = (&*arc_parker as *const Parker) as *const (); |
129 | 139 | let vt = vtable(); |
130 | 140 |
|
131 | 141 | let waker = unsafe { ManuallyDrop::new(Waker::from_raw(RawWaker::new(ptr, vt))) }; |
132 | 142 | let cx = &mut Context::from_waker(&waker); |
133 | 143 |
|
134 | 144 | loop { |
135 | 145 | if let Poll::Ready(t) = f.as_mut().poll(cx) { |
| 146 | + // Save the parker for the next invocation of `block`. |
| 147 | + cache.set(Some(arc_parker)); |
136 | 148 | return t; |
137 | 149 | } |
138 | | - thread::park(); |
| 150 | + arc_parker.park(); |
139 | 151 | } |
140 | 152 | }) |
141 | 153 | } |
142 | 154 |
|
143 | 155 | fn vtable() -> &'static RawWakerVTable { |
144 | 156 | unsafe fn clone_raw(ptr: *const ()) -> RawWaker { |
145 | | - let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const Thread)); |
| 157 | + let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const Parker)); |
146 | 158 | mem::forget(arc.clone()); |
147 | 159 | RawWaker::new(ptr, vtable()) |
148 | 160 | } |
149 | 161 |
|
150 | 162 | unsafe fn wake_raw(ptr: *const ()) { |
151 | | - let arc = Arc::from_raw(ptr as *const Thread); |
152 | | - arc.unpark(); |
| 163 | + let arc = Arc::from_raw(ptr as *const Parker); |
| 164 | + arc.unparker().unpark(); |
153 | 165 | } |
154 | 166 |
|
155 | 167 | unsafe fn wake_by_ref_raw(ptr: *const ()) { |
156 | | - let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const Thread)); |
157 | | - arc.unpark(); |
| 168 | + let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const Parker)); |
| 169 | + arc.unparker().unpark(); |
158 | 170 | } |
159 | 171 |
|
160 | 172 | unsafe fn drop_raw(ptr: *const ()) { |
161 | | - drop(Arc::from_raw(ptr as *const Thread)) |
| 173 | + drop(Arc::from_raw(ptr as *const Parker)) |
162 | 174 | } |
163 | 175 |
|
164 | 176 | &RawWakerVTable::new(clone_raw, wake_raw, wake_by_ref_raw, drop_raw) |
|
0 commit comments