fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
// Safe because we're !Unpin + !Drop mapping to a ?Unpin value
let gen = unsafe { Pin::map_unchecked_mut(self, |s| &mut s.0) };
- set_task_context(cx, || match gen.resume() {
+ let _guard = unsafe { set_task_context(cx) };
+ match gen.resume() {
GeneratorState::Yielded(()) => Poll::Pending,
GeneratorState::Complete(x) => Poll::Ready(x),
- })
+ }
}
}
}
}
-#[doc(hidden)]
-#[unstable(feature = "gen_future", issue = "50547")]
-/// Sets the thread-local task context used by async/await futures.
-pub fn set_task_context<F, R>(cx: &mut Context<'_>, f: F) -> R
-where
- F: FnOnce() -> R
-{
+// Safety: the returned guard must drop before `cx` is dropped and before
+// any previous guard is dropped.
+unsafe fn set_task_context(cx: &mut Context<'_>) -> SetOnDrop {
// transmute the context's lifetime to 'static so we can store it.
- let cx = unsafe {
- core::mem::transmute::<&mut Context<'_>, &mut Context<'static>>(cx)
- };
+ let cx = core::mem::transmute::<&mut Context<'_>, &mut Context<'static>>(cx);
let old_cx = TLS_CX.with(|tls_cx| {
tls_cx.replace(Some(NonNull::from(cx)))
});
- let _reset = SetOnDrop(old_cx);
- f()
+ SetOnDrop(old_cx)
}
#[doc(hidden)]
#[unstable(feature = "gen_future", issue = "50547")]
-/// Retrieves the thread-local task context used by async/await futures.
-///
-/// This function acquires exclusive access to the task context.
-///
-/// Panics if no context has been set or if the context has already been
-/// retrieved by a surrounding call to get_task_context.
-pub fn get_task_context<F, R>(f: F) -> R
+/// Polls a future in the current thread-local task waker.
+pub fn poll_with_tls_context<F>(f: Pin<&mut F>) -> Poll<F::Output>
where
- F: FnOnce(&mut Context<'_>) -> R
+ F: Future
{
let cx_ptr = TLS_CX.with(|tls_cx| {
// Clear the entry so that nested `get_task_waker` calls
//
// The pointer that was inserted came from an `&mut Context<'_>`,
// so it is safe to treat as mutable.
- unsafe { f(cx_ptr.as_mut()) }
-}
-
-#[doc(hidden)]
-#[unstable(feature = "gen_future", issue = "50547")]
-/// Polls a future in the current thread-local task waker.
-pub fn poll_with_tls_context<F>(f: Pin<&mut F>) -> Poll<F::Output>
-where
- F: Future
-{
- get_task_context(|cx| F::poll(f, cx))
+ unsafe { F::poll(f, cx_ptr.as_mut()) }
}