Skip to content

Commit 9e6bc3c

Browse files
committed
Apply review suggestions and fix tests
1 parent 01a704c commit 9e6bc3c

File tree

4 files changed

+163
-78
lines changed

4 files changed

+163
-78
lines changed

src/libcore/task/wake.rs

Lines changed: 31 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,14 @@
55
use fmt;
66
use marker::Unpin;
77

8-
/// A `RawWaker` allows the implementor of a task executor to create a `Waker`
8+
/// A `RawWaker` allows the implementor of a task executor to create a [`Waker`]
99
/// which provides customized wakeup behavior.
1010
///
1111
/// [vtable]: https://en.wikipedia.org/wiki/Virtual_method_table
1212
///
1313
/// It consists of a data pointer and a [virtual function pointer table (vtable)][vtable] that
1414
/// customizes the behavior of the `RawWaker`.
15-
#[derive(PartialEq)]
15+
#[derive(PartialEq, Debug)]
1616
pub struct RawWaker {
1717
/// A data pointer, which can be used to store arbitrary data as required
1818
/// by the executor. This could be e.g. a type-erased pointer to an `Arc`
@@ -24,55 +24,41 @@ pub struct RawWaker {
2424
pub vtable: &'static RawWakerVTable,
2525
}
2626

27-
impl fmt::Debug for RawWaker {
28-
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
29-
f.debug_struct("RawWaker")
30-
.finish()
31-
}
32-
}
33-
3427
/// A virtual function pointer table (vtable) that specifies the behavior
3528
/// of a [`RawWaker`].
3629
///
3730
/// The pointer passed to all functions inside the vtable is the `data` pointer
38-
/// from the enclosing `RawWaker` object.
39-
#[derive(PartialEq, Copy, Clone)]
31+
/// from the enclosing [`RawWaker`] object.
32+
#[derive(PartialEq, Copy, Clone, Debug)]
4033
pub struct RawWakerVTable {
41-
/// This function will be called when the `RawWaker` gets cloned, e.g. when
42-
/// the `Waker` in which the `RawWaker` is stored gets cloned.
34+
/// This function will be called when the [`RawWaker`] gets cloned, e.g. when
35+
/// the [`Waker`] in which the [`RawWaker`] is stored gets cloned.
4336
///
4437
/// The implementation of this function must retain all resources that are
45-
/// required for this additional instance of a `RawWaker` and associated
46-
/// task. Calling `wake` on the resulting `RawWaker` should result in a wakeup
47-
/// of the same task that would have been awoken by the original `RawWaker`.
38+
/// required for this additional instance of a [`RawWaker`] and associated
39+
/// task. Calling `wake` on the resulting [`RawWaker`] should result in a wakeup
40+
/// of the same task that would have been awoken by the original [`RawWaker`].
4841
pub clone: unsafe fn(*const ()) -> RawWaker,
4942

50-
/// This function will be called when `wake` is called on the `Waker`.
51-
/// It must wake up the task associated with this `RawWaker`.
43+
/// This function will be called when `wake` is called on the [`Waker`].
44+
/// It must wake up the task associated with this [`RawWaker`].
5245
pub wake: unsafe fn(*const ()),
5346

54-
/// This function gets called when a `RawWaker` gets dropped.
47+
/// This function gets called when a [`RawWaker`] gets dropped.
5548
///
5649
/// The implementation of this function must make sure to release any
57-
/// resources that are associated with this instance of a `RawWaker` and
50+
/// resources that are associated with this instance of a [`RawWaker`] and
5851
/// associated task.
59-
pub drop_fn: unsafe fn(*const ()),
60-
}
61-
62-
impl fmt::Debug for RawWakerVTable {
63-
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
64-
f.debug_struct("RawWakerVTable")
65-
.finish()
66-
}
52+
pub drop: unsafe fn(*const ()),
6753
}
6854

6955
/// A `Waker` is a handle for waking up a task by notifying its executor that it
7056
/// is ready to be run.
7157
///
72-
/// This handle encapsulates a `RawWaker` instance, which defines the
58+
/// This handle encapsulates a [`RawWaker`] instance, which defines the
7359
/// executor-specific wakeup behavior.
7460
///
75-
/// Implements `Clone`, `Send`, and `Sync`.
61+
/// Implements [`Clone`], [`Send`], and [`Sync`].
7662
#[repr(transparent)]
7763
pub struct Waker {
7864
waker: RawWaker,
@@ -87,6 +73,10 @@ impl Waker {
8773
pub fn wake(&self) {
8874
// The actual wakeup call is delegated through a virtual function call
8975
// to the implementation which is defined by the executor.
76+
77+
// SAFETY: This is safe because `Waker::new_unchecked` is the only way
78+
// to initialize `wake` and `data` requiring the user to acknowledge
79+
// that the contract of `RawWaker` is upheld.
9080
unsafe { (self.waker.vtable.wake)(self.waker.data) }
9181
}
9282

@@ -101,10 +91,11 @@ impl Waker {
10191
self.waker == other.waker
10292
}
10393

104-
/// Creates a new `Waker` from `RawWaker`.
94+
/// Creates a new `Waker` from [`RawWaker`].
10595
///
106-
/// The method cannot check whether `RawWaker` fulfills the required API
107-
/// contract to make it usable for `Waker` and is therefore unsafe.
96+
/// The behavior of the returned `Waker` is undefined if the contract defined
97+
/// in [RawWaker]'s documentation is not upheld. Therefore this method is
98+
/// unsafe.
10899
pub unsafe fn new_unchecked(waker: RawWaker) -> Waker {
109100
Waker {
110101
waker,
@@ -115,14 +106,20 @@ impl Waker {
115106
impl Clone for Waker {
116107
fn clone(&self) -> Self {
117108
Waker {
109+
// SAFETY: This is safe because `Waker::new_unchecked` is the only way
110+
// to initialize `clone` and `data` requiring the user to acknowledge
111+
// that the contract of [`RawWaker`] is upheld.
118112
waker: unsafe { (self.waker.vtable.clone)(self.waker.data) },
119113
}
120114
}
121115
}
122116

123117
impl Drop for Waker {
124118
fn drop(&mut self) {
125-
unsafe { (self.waker.vtable.drop_fn)(self.waker.data) }
119+
// SAFETY: This is safe because `Waker::new_unchecked` is the only way
120+
// to initialize `drop` and `data` requiring the user to acknowledge
121+
// that the contract of `RawWaker` is upheld.
122+
unsafe { (self.waker.vtable.drop)(self.waker.data) }
126123
}
127124
}
128125

src/test/compile-fail/must_use-in-stdlib-traits.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
use std::iter::Iterator;
55
use std::future::Future;
66

7-
use std::task::{Poll, LocalWaker};
7+
use std::task::{Poll, Waker};
88
use std::pin::Pin;
99
use std::unimplemented;
1010

@@ -13,7 +13,7 @@ struct MyFuture;
1313
impl Future for MyFuture {
1414
type Output = u32;
1515

16-
fn poll(self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<u32> {
16+
fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll<u32> {
1717
Poll::Pending
1818
}
1919
}

src/test/run-pass/async-await.rs

Lines changed: 61 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -9,17 +9,70 @@ use std::sync::{
99
atomic::{self, AtomicUsize},
1010
};
1111
use std::task::{
12-
LocalWaker, Poll, Wake,
13-
local_waker_from_nonlocal,
12+
Poll, Waker, RawWaker, RawWakerVTable,
1413
};
1514

15+
macro_rules! waker_vtable {
16+
($ty:ident) => {
17+
&RawWakerVTable {
18+
clone: clone_arc_raw::<$ty>,
19+
drop: drop_arc_raw::<$ty>,
20+
wake: wake_arc_raw::<$ty>,
21+
}
22+
};
23+
}
24+
25+
pub trait ArcWake {
26+
fn wake(arc_self: &Arc<Self>);
27+
28+
fn into_waker(wake: Arc<Self>) -> Waker where Self: Sized
29+
{
30+
let ptr = Arc::into_raw(wake) as *const();
31+
32+
unsafe {
33+
Waker::new_unchecked(RawWaker{
34+
data: ptr,
35+
vtable: waker_vtable!(Self),
36+
})
37+
}
38+
}
39+
}
40+
41+
unsafe fn increase_refcount<T: ArcWake>(data: *const()) {
42+
// Retain Arc by creating a copy
43+
let arc: Arc<T> = Arc::from_raw(data as *const T);
44+
let arc_clone = arc.clone();
45+
// Forget the Arcs again, so that the refcount isn't decrased
46+
let _ = Arc::into_raw(arc);
47+
let _ = Arc::into_raw(arc_clone);
48+
}
49+
50+
unsafe fn clone_arc_raw<T: ArcWake>(data: *const()) -> RawWaker {
51+
increase_refcount::<T>(data);
52+
RawWaker {
53+
data: data,
54+
vtable: waker_vtable!(T),
55+
}
56+
}
57+
58+
unsafe fn drop_arc_raw<T: ArcWake>(data: *const()) {
59+
// Drop Arc
60+
let _: Arc<T> = Arc::from_raw(data as *const T);
61+
}
62+
63+
unsafe fn wake_arc_raw<T: ArcWake>(data: *const()) {
64+
let arc: Arc<T> = Arc::from_raw(data as *const T);
65+
ArcWake::wake(&arc);
66+
let _ = Arc::into_raw(arc);
67+
}
68+
1669
struct Counter {
1770
wakes: AtomicUsize,
1871
}
1972

20-
impl Wake for Counter {
21-
fn wake(this: &Arc<Self>) {
22-
this.wakes.fetch_add(1, atomic::Ordering::SeqCst);
73+
impl ArcWake for Counter {
74+
fn wake(arc_self: &Arc<Self>) {
75+
arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
2376
}
2477
}
2578

@@ -29,11 +82,11 @@ fn wake_and_yield_once() -> WakeOnceThenComplete { WakeOnceThenComplete(false) }
2982

3083
impl Future for WakeOnceThenComplete {
3184
type Output = ();
32-
fn poll(mut self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<()> {
85+
fn poll(mut self: Pin<&mut Self>, waker: &Waker) -> Poll<()> {
3386
if self.0 {
3487
Poll::Ready(())
3588
} else {
36-
lw.wake();
89+
waker.wake();
3790
self.0 = true;
3891
Poll::Pending
3992
}
@@ -130,7 +183,7 @@ where
130183
{
131184
let mut fut = Box::pin(f(9));
132185
let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
133-
let waker = local_waker_from_nonlocal(counter.clone());
186+
let waker = ArcWake::into_waker(counter.clone());
134187
assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
135188
assert_eq!(Poll::Pending, fut.as_mut().poll(&waker));
136189
assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));

src/test/run-pass/futures-api.rs

Lines changed: 69 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -3,69 +3,104 @@
33

44
use std::future::Future;
55
use std::pin::Pin;
6-
use std::rc::Rc;
76
use std::sync::{
87
Arc,
98
atomic::{self, AtomicUsize},
109
};
1110
use std::task::{
12-
Poll, Wake, Waker, LocalWaker,
13-
local_waker, local_waker_from_nonlocal,
11+
Poll, Waker, RawWaker, RawWakerVTable,
1412
};
1513

16-
struct Counter {
17-
local_wakes: AtomicUsize,
18-
nonlocal_wakes: AtomicUsize,
14+
macro_rules! waker_vtable {
15+
($ty:ident) => {
16+
&RawWakerVTable {
17+
clone: clone_arc_raw::<$ty>,
18+
drop: drop_arc_raw::<$ty>,
19+
wake: wake_arc_raw::<$ty>,
20+
}
21+
};
1922
}
2023

21-
impl Wake for Counter {
22-
fn wake(this: &Arc<Self>) {
23-
this.nonlocal_wakes.fetch_add(1, atomic::Ordering::SeqCst);
24+
pub trait ArcWake {
25+
fn wake(arc_self: &Arc<Self>);
26+
27+
fn into_waker(wake: Arc<Self>) -> Waker where Self: Sized
28+
{
29+
let ptr = Arc::into_raw(wake) as *const();
30+
31+
unsafe {
32+
Waker::new_unchecked(RawWaker{
33+
data: ptr,
34+
vtable: waker_vtable!(Self),
35+
})
36+
}
2437
}
38+
}
39+
40+
unsafe fn increase_refcount<T: ArcWake>(data: *const()) {
41+
// Retain Arc by creating a copy
42+
let arc: Arc<T> = Arc::from_raw(data as *const T);
43+
let arc_clone = arc.clone();
44+
// Forget the Arcs again, so that the refcount isn't decrased
45+
let _ = Arc::into_raw(arc);
46+
let _ = Arc::into_raw(arc_clone);
47+
}
2548

26-
unsafe fn wake_local(this: &Arc<Self>) {
27-
this.local_wakes.fetch_add(1, atomic::Ordering::SeqCst);
49+
unsafe fn clone_arc_raw<T: ArcWake>(data: *const()) -> RawWaker {
50+
increase_refcount::<T>(data);
51+
RawWaker {
52+
data: data,
53+
vtable: waker_vtable!(T),
54+
}
55+
}
56+
57+
unsafe fn drop_arc_raw<T: ArcWake>(data: *const()) {
58+
// Drop Arc
59+
let _: Arc<T> = Arc::from_raw(data as *const T);
60+
}
61+
62+
unsafe fn wake_arc_raw<T: ArcWake>(data: *const()) {
63+
let arc: Arc<T> = Arc::from_raw(data as *const T);
64+
ArcWake::wake(&arc);
65+
let _ = Arc::into_raw(arc);
66+
}
67+
68+
struct Counter {
69+
wakes: AtomicUsize,
70+
}
71+
72+
impl Wake for Counter {
73+
fn wake(arc_self: &Arc<Self>) {
74+
arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
2875
}
2976
}
3077

3178
struct MyFuture;
3279

3380
impl Future for MyFuture {
3481
type Output = ();
35-
fn poll(self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<Self::Output> {
36-
// Wake once locally
37-
lw.wake();
38-
// Wake twice non-locally
39-
let waker = lw.clone().into_waker();
82+
fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll<Self::Output> {
83+
// Wake twice
4084
waker.wake();
4185
waker.wake();
4286
Poll::Ready(())
4387
}
4488
}
4589

46-
fn test_local_waker() {
90+
fn test_waker() {
4791
let counter = Arc::new(Counter {
48-
local_wakes: AtomicUsize::new(0),
49-
nonlocal_wakes: AtomicUsize::new(0),
92+
wakes: AtomicUsize::new(0),
5093
});
51-
let waker = unsafe { local_waker(counter.clone()) };
52-
assert_eq!(Poll::Ready(()), Pin::new(&mut MyFuture).poll(&waker));
53-
assert_eq!(1, counter.local_wakes.load(atomic::Ordering::SeqCst));
54-
assert_eq!(2, counter.nonlocal_wakes.load(atomic::Ordering::SeqCst));
55-
}
94+
let waker = ArcWake::into_waker(counter.clone());
95+
assert_eq!(2, Arc::strong_count(&counter));
5696

57-
fn test_local_as_nonlocal_waker() {
58-
let counter = Arc::new(Counter {
59-
local_wakes: AtomicUsize::new(0),
60-
nonlocal_wakes: AtomicUsize::new(0),
61-
});
62-
let waker: LocalWaker = local_waker_from_nonlocal(counter.clone());
6397
assert_eq!(Poll::Ready(()), Pin::new(&mut MyFuture).poll(&waker));
64-
assert_eq!(0, counter.local_wakes.load(atomic::Ordering::SeqCst));
65-
assert_eq!(3, counter.nonlocal_wakes.load(atomic::Ordering::SeqCst));
98+
assert_eq!(2, counter.wakes.load(atomic::Ordering::SeqCst));
99+
100+
drop(waker);
101+
assert_eq!(1, Arc::strong_count(&counter));
66102
}
67103

68104
fn main() {
69-
test_local_waker();
70-
test_local_as_nonlocal_waker();
105+
test_waker();
71106
}

0 commit comments

Comments
 (0)