Auto merge of #59119 - cramertj:cx-back, r=withoutboats

Future-proof the Futures API

cc https://github.com/rust-lang/rust/issues/59113, @carllerche, @rust-lang/libs

r? @withoutboats

3 files changed, 40 insertions, 29 deletions

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<T: Generator<Yield = ()>> !Unpin for GenFuture<T> {}
 #[unstable(feature = "gen_future", issue = "50547")]
 impl<T: Generator<Yield = ()>> Future for GenFuture<T> {
     type Output = T::Return;
-    fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll<Self::Output> {
+    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_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<T: Generator<Yield = ()>> Future for GenFuture<T> {
 }
 
 thread_local! {
-    static TLS_WAKER: Cell<Option<NonNull<Waker>>> = Cell::new(None);
+    static TLS_CX: Cell<Option<NonNull<Context<'static>>>> = Cell::new(None);
 }
 
-struct SetOnDrop(Option<NonNull<Waker>>);
+struct SetOnDrop(Option<NonNull<Context<'static>>>);
 
 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<F, R>(waker: &Waker, f: F) -> R
+pub fn set_task_context<F, R>(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, R>(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, R>(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>(f: Pin<&mut F>) -> Poll<F::Output>
+pub fn poll_with_tls_context<F>(f: Pin<&mut F>) -> Poll<F::Output>
 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<T: fmt::Debug> fmt::Debug for AssertUnwindSafe<T> {
 impl<F: Future> Future for AssertUnwindSafe<F> {
     type Output = F::Output;
 
-    fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll<Self::Output> {
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let pinned_field = unsafe { Pin::map_unchecked_mut(self, |x| &mut x.0) };
-        F::poll(pinned_field, waker)
+        F::poll(pinned_field, cx)
     }
 }