From 1691e06db661a19a5e25276c18cd165386f027bb Mon Sep 17 00:00:00 2001 From: Taylor Cramer Date: Mon, 11 Mar 2019 16:56:00 -0700 Subject: Future-proof the Futures API --- src/libstd/future.rs | 61 +++++++++++++++++++++++++++++++--------------------- src/libstd/macros.rs | 2 +- src/libstd/panic.rs | 6 +++--- 3 files changed, 40 insertions(+), 29 deletions(-) (limited to 'src/libstd') diff --git a/src/libstd/future.rs b/src/libstd/future.rs index aa784746122..898387cb9f5 100644 --- a/src/libstd/future.rs +++ b/src/libstd/future.rs @@ -5,7 +5,7 @@ use core::marker::Unpin; use core::pin::Pin; use core::option::Option; use core::ptr::NonNull; -use core::task::{Waker, Poll}; +use core::task::{Context, Poll}; use core::ops::{Drop, Generator, GeneratorState}; #[doc(inline)] @@ -32,10 +32,10 @@ impl> !Unpin for GenFuture {} #[unstable(feature = "gen_future", issue = "50547")] impl> Future for GenFuture { type Output = T::Return; - fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll { + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { // 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_waker(waker, || match gen.resume() { + set_task_context(cx, || match gen.resume() { GeneratorState::Yielded(()) => Poll::Pending, GeneratorState::Complete(x) => Poll::Ready(x), }) @@ -43,61 +43,72 @@ impl> Future for GenFuture { } thread_local! { - static TLS_WAKER: Cell>> = Cell::new(None); + static TLS_CX: Cell>>> = Cell::new(None); } -struct SetOnDrop(Option>); +struct SetOnDrop(Option>>); impl Drop for SetOnDrop { fn drop(&mut self) { - TLS_WAKER.with(|tls_waker| { - tls_waker.set(self.0.take()); + TLS_CX.with(|tls_cx| { + tls_cx.set(self.0.take()); }); } } #[unstable(feature = "gen_future", issue = "50547")] /// Sets the thread-local task context used by async/await futures. -pub fn set_task_waker(waker: &Waker, f: F) -> R +pub fn set_task_context(cx: &mut Context<'_>, f: F) -> R where F: FnOnce() -> R { - let old_waker = TLS_WAKER.with(|tls_waker| { - tls_waker.replace(Some(NonNull::from(waker))) + // 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 old_cx = TLS_CX.with(|tls_cx| { + tls_cx.replace(Some(NonNull::from(cx))) }); - let _reset_waker = SetOnDrop(old_waker); + let _reset = SetOnDrop(old_cx); f() } #[unstable(feature = "gen_future", issue = "50547")] -/// Retrieves the thread-local task waker used by async/await futures. +/// Retrieves the thread-local task context used by async/await futures. /// -/// This function acquires exclusive access to the task waker. +/// This function acquires exclusive access to the task context. /// -/// Panics if no waker has been set or if the waker has already been -/// retrieved by a surrounding call to get_task_waker. -pub fn get_task_waker(f: F) -> R +/// 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: F) -> R where - F: FnOnce(&Waker) -> R + F: FnOnce(&mut Context<'_>) -> R { - let waker_ptr = TLS_WAKER.with(|tls_waker| { + let cx_ptr = TLS_CX.with(|tls_cx| { // Clear the entry so that nested `get_task_waker` calls // will fail or set their own value. - tls_waker.replace(None) + tls_cx.replace(None) }); - let _reset_waker = SetOnDrop(waker_ptr); + let _reset = SetOnDrop(cx_ptr); - let waker_ptr = waker_ptr.expect( - "TLS Waker not set. This is a rustc bug. \ + let mut cx_ptr = cx_ptr.expect( + "TLS Context not set. This is a rustc bug. \ Please file an issue on https://github.com/rust-lang/rust."); - unsafe { f(waker_ptr.as_ref()) } + + // Safety: we've ensured exclusive access to the context by + // removing the pointer from TLS, only to be replaced once + // we're done with it. + // + // The pointer that was inserted came from an `&mut Context<'_>`, + // so it is safe to treat as mutable. + unsafe { f(cx_ptr.as_mut()) } } #[unstable(feature = "gen_future", issue = "50547")] /// Polls a future in the current thread-local task waker. -pub fn poll_with_tls_waker(f: Pin<&mut F>) -> Poll +pub fn poll_with_tls_context(f: Pin<&mut F>) -> Poll where F: Future { - get_task_waker(|waker| F::poll(f, waker)) + get_task_context(|cx| F::poll(f, cx)) } diff --git a/src/libstd/macros.rs b/src/libstd/macros.rs index be4db1f737d..44ca56e2611 100644 --- a/src/libstd/macros.rs +++ b/src/libstd/macros.rs @@ -346,7 +346,7 @@ macro_rules! r#await { let mut pinned = $e; loop { if let $crate::task::Poll::Ready(x) = - $crate::future::poll_with_tls_waker(unsafe { + $crate::future::poll_with_tls_context(unsafe { $crate::pin::Pin::new_unchecked(&mut pinned) }) { diff --git a/src/libstd/panic.rs b/src/libstd/panic.rs index cc147d851de..5a8101e2301 100644 --- a/src/libstd/panic.rs +++ b/src/libstd/panic.rs @@ -12,7 +12,7 @@ use crate::panicking; use crate::ptr::{Unique, NonNull}; use crate::rc::Rc; use crate::sync::{Arc, Mutex, RwLock, atomic}; -use crate::task::{Waker, Poll}; +use crate::task::{Context, Poll}; use crate::thread::Result; #[stable(feature = "panic_hooks", since = "1.10.0")] @@ -323,9 +323,9 @@ impl fmt::Debug for AssertUnwindSafe { impl Future for AssertUnwindSafe { type Output = F::Output; - fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll { + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let pinned_field = unsafe { Pin::map_unchecked_mut(self, |x| &mut x.0) }; - F::poll(pinned_field, waker) + F::poll(pinned_field, cx) } } -- cgit 1.4.1-3-g733a5