diff --git a/src/libstd/future.rs b/src/libstd/future.rs
index c65f71fb1a4e2..6de3f1d545b57 100644
--- a/src/libstd/future.rs
+++ b/src/libstd/future.rs
@@ -40,10 +40,11 @@ impl<T: Generator<Yield = ()>> Future for GenFuture<T> {
     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),
-        })
+        }
     }
 }
 
@@ -61,35 +62,23 @@ impl Drop for SetOnDrop {
     }
 }
 
-#[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
@@ -108,15 +97,5 @@ where
     //
     // 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()) }
 }