Skip to content

Commit 37fdb45

Browse files
committed
Future-proof the Futures API
1 parent f22dca0 commit 37fdb45

File tree

11 files changed

+174
-70
lines changed

11 files changed

+174
-70
lines changed

src/liballoc/boxed.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ use core::ops::{
8181
CoerceUnsized, DispatchFromDyn, Deref, DerefMut, Receiver, Generator, GeneratorState
8282
};
8383
use core::ptr::{self, NonNull, Unique};
84-
use core::task::{Waker, Poll};
84+
use core::task::{Context, Poll};
8585

8686
use crate::vec::Vec;
8787
use crate::raw_vec::RawVec;
@@ -911,7 +911,7 @@ impl<G: ?Sized + Generator> Generator for Pin<Box<G>> {
911911
impl<F: ?Sized + Future + Unpin> Future for Box<F> {
912912
type Output = F::Output;
913913

914-
fn poll(mut self: Pin<&mut Self>, waker: &Waker) -> Poll<Self::Output> {
915-
F::poll(Pin::new(&mut *self), waker)
914+
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
915+
F::poll(Pin::new(&mut *self), cx)
916916
}
917917
}

src/libcore/future/future.rs

+10-8
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
use marker::Unpin;
66
use ops;
77
use pin::Pin;
8-
use task::{Poll, Waker};
8+
use task::{Context, Poll};
99

1010
/// A future represents an asynchronous computation.
1111
///
@@ -44,8 +44,9 @@ pub trait Future {
4444
/// Once a future has finished, clients should not `poll` it again.
4545
///
4646
/// When a future is not ready yet, `poll` returns `Poll::Pending` and
47-
/// stores a clone of the [`Waker`] to be woken once the future can
48-
/// make progress. For example, a future waiting for a socket to become
47+
/// stores a clone of the [`Waker`] copied from the current [`Context`].
48+
/// This [`Waker`] is then woken once the future can make progress.
49+
/// For example, a future waiting for a socket to become
4950
/// readable would call `.clone()` on the [`Waker`] and store it.
5051
/// When a signal arrives elsewhere indicating that the socket is readable,
5152
/// `[Waker::wake]` is called and the socket future's task is awoken.
@@ -88,16 +89,17 @@ pub trait Future {
8889
///
8990
/// [`Poll::Pending`]: ../task/enum.Poll.html#variant.Pending
9091
/// [`Poll::Ready(val)`]: ../task/enum.Poll.html#variant.Ready
92+
/// [`Context`]: ../task/struct.Context.html
9193
/// [`Waker`]: ../task/struct.Waker.html
9294
/// [`Waker::wake`]: ../task/struct.Waker.html#method.wake
93-
fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll<Self::Output>;
95+
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>;
9496
}
9597

9698
impl<F: ?Sized + Future + Unpin> Future for &mut F {
9799
type Output = F::Output;
98100

99-
fn poll(mut self: Pin<&mut Self>, waker: &Waker) -> Poll<Self::Output> {
100-
F::poll(Pin::new(&mut **self), waker)
101+
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
102+
F::poll(Pin::new(&mut **self), cx)
101103
}
102104
}
103105

@@ -108,7 +110,7 @@ where
108110
{
109111
type Output = <<P as ops::Deref>::Target as Future>::Output;
110112

111-
fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll<Self::Output> {
112-
Pin::get_mut(self).as_mut().poll(waker)
113+
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
114+
Pin::get_mut(self).as_mut().poll(cx)
113115
}
114116
}

src/libcore/task/mod.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -8,4 +8,4 @@ mod poll;
88
pub use self::poll::Poll;
99

1010
mod wake;
11-
pub use self::wake::{Waker, RawWaker, RawWakerVTable};
11+
pub use self::wake::{Context, Waker, RawWaker, RawWakerVTable};

src/libcore/task/wake.rs

+95-4
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
issue = "50547")]
44

55
use fmt;
6-
use marker::Unpin;
6+
use marker::{PhantomData, Unpin};
77

88
/// A `RawWaker` allows the implementor of a task executor to create a [`Waker`]
99
/// which provides customized wakeup behavior.
@@ -36,6 +36,10 @@ impl RawWaker {
3636
/// The `vtable` customizes the behavior of a `Waker` which gets created
3737
/// from a `RawWaker`. For each operation on the `Waker`, the associated
3838
/// function in the `vtable` of the underlying `RawWaker` will be called.
39+
#[rustc_promotable]
40+
#[unstable(feature = "futures_api",
41+
reason = "futures in libcore are unstable",
42+
issue = "50547")]
3943
pub const fn new(data: *const (), vtable: &'static RawWakerVTable) -> RawWaker {
4044
RawWaker {
4145
data,
@@ -63,21 +67,103 @@ pub struct RawWakerVTable {
6367
/// required for this additional instance of a [`RawWaker`] and associated
6468
/// task. Calling `wake` on the resulting [`RawWaker`] should result in a wakeup
6569
/// of the same task that would have been awoken by the original [`RawWaker`].
66-
pub clone: unsafe fn(*const ()) -> RawWaker,
70+
clone: unsafe fn(*const ()) -> RawWaker,
6771

6872
/// This function will be called when `wake` is called on the [`Waker`].
6973
/// It must wake up the task associated with this [`RawWaker`].
7074
///
7175
/// The implemention of this function must not consume the provided data
7276
/// pointer.
73-
pub wake: unsafe fn(*const ()),
77+
wake: unsafe fn(*const ()),
78+
79+
/// This function gets called when a [`RawWaker`] gets dropped.
80+
///
81+
/// The implementation of this function must make sure to release any
82+
/// resources that are associated with this instance of a [`RawWaker`] and
83+
/// associated task.
84+
drop: unsafe fn(*const ()),
85+
}
7486

87+
impl RawWakerVTable {
88+
/// Creates a new `RawWakerVTable` from the provided `clone`, `wake`, and
89+
/// `drop` functions.
90+
///
91+
/// # `clone`
92+
///
93+
/// This function will be called when the [`RawWaker`] gets cloned, e.g. when
94+
/// the [`Waker`] in which the [`RawWaker`] is stored gets cloned.
95+
///
96+
/// The implementation of this function must retain all resources that are
97+
/// required for this additional instance of a [`RawWaker`] and associated
98+
/// task. Calling `wake` on the resulting [`RawWaker`] should result in a wakeup
99+
/// of the same task that would have been awoken by the original [`RawWaker`].
100+
///
101+
/// # `wake`
102+
///
103+
/// This function will be called when `wake` is called on the [`Waker`].
104+
/// It must wake up the task associated with this [`RawWaker`].
105+
///
106+
/// The implemention of this function must not consume the provided data
107+
/// pointer.
108+
///
109+
/// # `drop`
110+
///
75111
/// This function gets called when a [`RawWaker`] gets dropped.
76112
///
77113
/// The implementation of this function must make sure to release any
78114
/// resources that are associated with this instance of a [`RawWaker`] and
79115
/// associated task.
80-
pub drop: unsafe fn(*const ()),
116+
#[rustc_promotable]
117+
#[unstable(feature = "futures_api",
118+
reason = "futures in libcore are unstable",
119+
issue = "50547")]
120+
pub const fn new(
121+
clone: unsafe fn(*const ()) -> RawWaker,
122+
wake: unsafe fn(*const ()),
123+
drop: unsafe fn(*const ()),
124+
) -> Self {
125+
Self {
126+
clone,
127+
wake,
128+
drop,
129+
}
130+
}
131+
}
132+
133+
/// The `Context` of an asynchronous task.
134+
///
135+
/// Currently, `Context` only serves to provide access to a `&Waker`
136+
/// which can be used to wake the current task.
137+
pub struct Context<'a> {
138+
waker: &'a Waker,
139+
// Ensure we future-proof against variance changes by including
140+
// an `&mut` reference with the provided lifetime.
141+
_marker: PhantomData<&'a mut ()>,
142+
}
143+
144+
impl<'a> Context<'a> {
145+
/// Create a new `Context` from a `&Waker`.
146+
#[inline]
147+
pub fn from_waker(waker: &'a Waker) -> Self {
148+
Context {
149+
waker,
150+
_marker: PhantomData,
151+
}
152+
}
153+
154+
/// Returns a reference to the `Waker` for the current task.
155+
#[inline]
156+
pub fn waker(&self) -> &'a Waker {
157+
&self.waker
158+
}
159+
}
160+
161+
impl fmt::Debug for Context<'_> {
162+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
163+
f.debug_struct("Context")
164+
.field("waker", &self.waker)
165+
.finish()
166+
}
81167
}
82168

83169
/// A `Waker` is a handle for waking up a task by notifying its executor that it
@@ -98,6 +184,7 @@ unsafe impl Sync for Waker {}
98184

99185
impl Waker {
100186
/// Wake up the task associated with this `Waker`.
187+
#[inline]
101188
pub fn wake(&self) {
102189
// The actual wakeup call is delegated through a virtual function call
103190
// to the implementation which is defined by the executor.
@@ -115,6 +202,7 @@ impl Waker {
115202
/// returns `true`, it is guaranteed that the `Waker`s will awaken the same task.
116203
///
117204
/// This function is primarily used for optimization purposes.
205+
#[inline]
118206
pub fn will_wake(&self, other: &Waker) -> bool {
119207
self.waker == other.waker
120208
}
@@ -124,6 +212,7 @@ impl Waker {
124212
/// The behavior of the returned `Waker` is undefined if the contract defined
125213
/// in [`RawWaker`]'s and [`RawWakerVTable`]'s documentation is not upheld.
126214
/// Therefore this method is unsafe.
215+
#[inline]
127216
pub unsafe fn new_unchecked(waker: RawWaker) -> Waker {
128217
Waker {
129218
waker,
@@ -132,6 +221,7 @@ impl Waker {
132221
}
133222

134223
impl Clone for Waker {
224+
#[inline]
135225
fn clone(&self) -> Self {
136226
Waker {
137227
// SAFETY: This is safe because `Waker::new_unchecked` is the only way
@@ -143,6 +233,7 @@ impl Clone for Waker {
143233
}
144234

145235
impl Drop for Waker {
236+
#[inline]
146237
fn drop(&mut self) {
147238
// SAFETY: This is safe because `Waker::new_unchecked` is the only way
148239
// to initialize `drop` and `data` requiring the user to acknowledge

src/libstd/future.rs

+36-25
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use core::marker::Unpin;
55
use core::pin::Pin;
66
use core::option::Option;
77
use core::ptr::NonNull;
8-
use core::task::{Waker, Poll};
8+
use core::task::{Context, Poll};
99
use core::ops::{Drop, Generator, GeneratorState};
1010

1111
#[doc(inline)]
@@ -32,72 +32,83 @@ impl<T: Generator<Yield = ()>> !Unpin for GenFuture<T> {}
3232
#[unstable(feature = "gen_future", issue = "50547")]
3333
impl<T: Generator<Yield = ()>> Future for GenFuture<T> {
3434
type Output = T::Return;
35-
fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll<Self::Output> {
35+
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
3636
// Safe because we're !Unpin + !Drop mapping to a ?Unpin value
3737
let gen = unsafe { Pin::map_unchecked_mut(self, |s| &mut s.0) };
38-
set_task_waker(waker, || match gen.resume() {
38+
set_task_context(cx, || match gen.resume() {
3939
GeneratorState::Yielded(()) => Poll::Pending,
4040
GeneratorState::Complete(x) => Poll::Ready(x),
4141
})
4242
}
4343
}
4444

4545
thread_local! {
46-
static TLS_WAKER: Cell<Option<NonNull<Waker>>> = Cell::new(None);
46+
static TLS_CX: Cell<Option<NonNull<Context<'static>>>> = Cell::new(None);
4747
}
4848

49-
struct SetOnDrop(Option<NonNull<Waker>>);
49+
struct SetOnDrop(Option<NonNull<Context<'static>>>);
5050

5151
impl Drop for SetOnDrop {
5252
fn drop(&mut self) {
53-
TLS_WAKER.with(|tls_waker| {
54-
tls_waker.set(self.0.take());
53+
TLS_CX.with(|tls_cx| {
54+
tls_cx.set(self.0.take());
5555
});
5656
}
5757
}
5858

5959
#[unstable(feature = "gen_future", issue = "50547")]
6060
/// Sets the thread-local task context used by async/await futures.
61-
pub fn set_task_waker<F, R>(waker: &Waker, f: F) -> R
61+
pub fn set_task_context<F, R>(cx: &mut Context<'_>, f: F) -> R
6262
where
6363
F: FnOnce() -> R
6464
{
65-
let old_waker = TLS_WAKER.with(|tls_waker| {
66-
tls_waker.replace(Some(NonNull::from(waker)))
65+
// transmute the context's lifetime to 'static so we can store it.
66+
let cx = unsafe {
67+
core::mem::transmute::<&mut Context<'_>, &mut Context<'static>>(cx)
68+
};
69+
let old_cx = TLS_CX.with(|tls_cx| {
70+
tls_cx.replace(Some(NonNull::from(cx)))
6771
});
68-
let _reset_waker = SetOnDrop(old_waker);
72+
let _reset = SetOnDrop(old_cx);
6973
f()
7074
}
7175

7276
#[unstable(feature = "gen_future", issue = "50547")]
73-
/// Retrieves the thread-local task waker used by async/await futures.
77+
/// Retrieves the thread-local task context used by async/await futures.
7478
///
75-
/// This function acquires exclusive access to the task waker.
79+
/// This function acquires exclusive access to the task context.
7680
///
77-
/// Panics if no waker has been set or if the waker has already been
78-
/// retrieved by a surrounding call to get_task_waker.
79-
pub fn get_task_waker<F, R>(f: F) -> R
81+
/// Panics if no context has been set or if the context has already been
82+
/// retrieved by a surrounding call to get_task_context.
83+
pub fn get_task_context<F, R>(f: F) -> R
8084
where
81-
F: FnOnce(&Waker) -> R
85+
F: FnOnce(&mut Context<'_>) -> R
8286
{
83-
let waker_ptr = TLS_WAKER.with(|tls_waker| {
87+
let cx_ptr = TLS_CX.with(|tls_cx| {
8488
// Clear the entry so that nested `get_task_waker` calls
8589
// will fail or set their own value.
86-
tls_waker.replace(None)
90+
tls_cx.replace(None)
8791
});
88-
let _reset_waker = SetOnDrop(waker_ptr);
92+
let _reset = SetOnDrop(cx_ptr);
8993

90-
let waker_ptr = waker_ptr.expect(
91-
"TLS Waker not set. This is a rustc bug. \
94+
let mut cx_ptr = cx_ptr.expect(
95+
"TLS Context not set. This is a rustc bug. \
9296
Please file an issue on https://github.com/rust-lang/rust.");
93-
unsafe { f(waker_ptr.as_ref()) }
97+
98+
// Safety: we've ensured exclusive access to the context by
99+
// removing the pointer from TLS, only to be replaced once
100+
// we're done with it.
101+
//
102+
// The pointer that was inserted came from an `&mut Context<'_>`,
103+
// so it is safe to treat as mutable.
104+
unsafe { f(cx_ptr.as_mut()) }
94105
}
95106

96107
#[unstable(feature = "gen_future", issue = "50547")]
97108
/// Polls a future in the current thread-local task waker.
98-
pub fn poll_with_tls_waker<F>(f: Pin<&mut F>) -> Poll<F::Output>
109+
pub fn poll_with_tls_context<F>(f: Pin<&mut F>) -> Poll<F::Output>
99110
where
100111
F: Future
101112
{
102-
get_task_waker(|waker| F::poll(f, waker))
113+
get_task_context(|cx| F::poll(f, cx))
103114
}

src/libstd/macros.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -332,7 +332,7 @@ macro_rules! r#await {
332332
let mut pinned = $e;
333333
loop {
334334
if let $crate::task::Poll::Ready(x) =
335-
$crate::future::poll_with_tls_waker(unsafe {
335+
$crate::future::poll_with_tls_context(unsafe {
336336
$crate::pin::Pin::new_unchecked(&mut pinned)
337337
})
338338
{

src/libstd/panic.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ use crate::panicking;
1212
use crate::ptr::{Unique, NonNull};
1313
use crate::rc::Rc;
1414
use crate::sync::{Arc, Mutex, RwLock, atomic};
15-
use crate::task::{Waker, Poll};
15+
use crate::task::{Context, Poll};
1616
use crate::thread::Result;
1717

1818
#[stable(feature = "panic_hooks", since = "1.10.0")]
@@ -323,9 +323,9 @@ impl<T: fmt::Debug> fmt::Debug for AssertUnwindSafe<T> {
323323
impl<F: Future> Future for AssertUnwindSafe<F> {
324324
type Output = F::Output;
325325

326-
fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll<Self::Output> {
326+
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
327327
let pinned_field = unsafe { Pin::map_unchecked_mut(self, |x| &mut x.0) };
328-
F::poll(pinned_field, waker)
328+
F::poll(pinned_field, cx)
329329
}
330330
}
331331

0 commit comments

Comments
 (0)