1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
use core::fmt;

/// An unsafe trait for implementing custom memory management for a
/// [`Waker`](::task::Waker).
///
/// A [`Waker`](::task::Waker) is a cloneable trait object for `Wake`, and is
/// most often essentially just `Arc<Wake>`. However, in some contexts
/// (particularly `no_std`), it's desirable to avoid `Arc` in favor of some
/// custom memory management strategy. This trait is designed to allow for such
/// customization.
///
/// A default implementation of the `UnsafeWake` trait is provided for the
/// `Arc` type in the standard library. If the `std` feature of this crate
/// is not available however, you'll be required to implement your own
/// instance of this trait to pass it into `Waker::new`.
///
/// # Unsafety
///
/// This trait manually encodes the memory management of the underlying trait
/// object. Implementors of this trait must guarantee:
///
/// * Calls to `clone_raw` produce uniquely owned `Waker` handles. These handles
/// should be independently usable and droppable.
///
/// * Calls to `drop_raw` work with `self` as a raw pointer, deallocating
///   resources associated with it. This is a pretty unsafe operation as it's
///   invalidating the `self` pointer, so extreme care needs to be taken.
///
/// In general it's recommended to review the trait documentation as well as the
/// implementation for `Arc` in this crate before attempting a custom
/// implementation.
pub unsafe trait UnsafeWake {
    /// Creates a new `Waker` from this instance of `UnsafeWake`.
    ///
    /// This function will create a new uniquely owned handle that under the
    /// hood references the same notification instance. In other words calls
    /// to `wake` on the returned handle should be equivalent to calls to
    /// `wake` on this handle.
    ///
    /// # Unsafety
    ///
    /// This is also unsafe to call because it's asserting the `UnsafeWake`
    /// value is in a consistent state, i.e. hasn't been dropped.
    unsafe fn clone_raw(&self) -> Waker;

    /// Drops this instance of `UnsafeWake`, deallocating resources
    /// associated with it.
    ///
    /// This method is intended to have a signature such as:
    ///
    /// ```ignore
    /// fn drop_raw(self: *mut Self);
    /// ```
    ///
    /// Unfortunately in Rust today that signature is not object safe.
    /// Nevertheless it's recommended to implement this function *as if* that
    /// were its signature. As such it is not safe to call on an invalid
    /// pointer, nor is the validity of the pointer guaranteed after this
    /// function returns.
    ///
    /// # Unsafety
    ///
    /// This is also unsafe to call because it's asserting the `UnsafeWake`
    /// value is in a consistent state, i.e. hasn't been dropped
    unsafe fn drop_raw(&self);

    /// Indicates that the associated task is ready to make progress and should
    /// be `poll`ed.
    ///
    /// Executors generally maintain a queue of "ready" tasks; `wake` should place
    /// the associated task onto this queue.
    ///
    /// # Panics
    ///
    /// Implementations should avoid panicking, but clients should also be prepared
    /// for panics.
    ///
    /// # Unsafety
    ///
    /// This is also unsafe to call because it's asserting the `UnsafeWake`
    /// value is in a consistent state, i.e. hasn't been dropped
    unsafe fn wake(&self);
}

/// A `Waker` is a handle for waking up a task by notifying its executor that it
/// is ready to be run.
///
/// This handle contains a trait object pointing to an instance of the `Wake`
/// trait, allowing notifications to get routed through it. Usually `Waker`
/// instances are provided by an executor.
///
/// If you're implementing an executor, the recommended way to create a `Waker`
/// is via `Waker::from` applied to an `Arc<T>` value where `T: Wake`. The
/// unsafe `new` constructor should be used only in niche, `no_std` settings.
pub struct Waker {
    inner: *const UnsafeWake,
}

unsafe impl Send for Waker {}
unsafe impl Sync for Waker {}

impl Waker {
    /// Constructs a new `Waker` directly.
    ///
    /// Note that most code will not need to call this. Implementers of the
    /// `UnsafeWake` trait will typically provide a wrapper that calls this
    /// but you otherwise shouldn't call it directly.
    ///
    /// If you're working with the standard library then it's recommended to
    /// use the `Waker::from` function instead which works with the safe
    /// `Arc` type and the safe `Wake` trait.
    #[inline]
    pub unsafe fn new(inner: *const UnsafeWake) -> Waker {
        Waker { inner: inner }
    }

    /// Wake up the task associated with this `Waker`.
    pub fn wake(&self) {
        unsafe { (*self.inner).wake() }
    }

    /// Returns whether or not this `Waker` and `other` awaken the same task.
    ///
    /// This function works on a best-effort basis, and may return false even
    /// when the `Waker`s would awaken the same task. However, if this function
    /// returns true, it is guaranteed that the `Waker`s will awaken the same
    /// task.
    ///
    /// This function is primarily used for optimization purposes.
    pub fn will_wake(&self, other: &Waker) -> bool {
        self.inner == other.inner
    }
}

impl Clone for Waker {
    #[inline]
    fn clone(&self) -> Self {
        unsafe {
            (*self.inner).clone_raw()
        }
    }
}

impl fmt::Debug for Waker {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Waker")
            .finish()
    }
}

impl Drop for Waker {
    fn drop(&mut self) {
        unsafe {
            (*self.inner).drop_raw()
        }
    }
}

if_std! {
    use std::mem;
    use std::ptr;
    use std::sync::Arc;
    use core::marker::PhantomData;

    /// A way of waking up a specific task.
    ///
    /// Any task executor must provide a way of signaling that a task it owns
    /// is ready to be `poll`ed again. Executors do so by implementing this trait.
    ///
    /// Note that, rather than working directly with `Wake` trait objects, this
    /// library instead uses a custom [`Waker`](::task::Waker) to allow for
    /// customization of memory management.
    pub trait Wake: Send + Sync {
        /// Indicates that the associated task is ready to make progress and should
        /// be `poll`ed.
        ///
        /// Executors generally maintain a queue of "ready" tasks; `wake` should place
        /// the associated task onto this queue.
        ///
        /// # Panics
        ///
        /// Implementations should avoid panicking, but clients should also be prepared
        /// for panics.
        fn wake(arc_self: &Arc<Self>);
    }

    // Safe implementation of `UnsafeWake` for `Arc` in the standard library.
    //
    // Note that this is a very unsafe implementation! The crucial pieces is that
    // these two values are considered equivalent:
    //
    // * Arc<T>
    // * *const ArcWrapped<T>
    //
    // We don't actually know the layout of `ArcWrapped<T>` as it's an
    // implementation detail in the standard library. We can work, though, by
    // casting it through and back an `Arc<T>`.
    //
    // This also means that you won't actually find `UnsafeWake for Arc<T>`
    // because it's the wrong level of indirection. These methods are sort of
    // receiving Arc<T>, but not an owned version. It's... complicated. We may be
    // one of the first users of unsafe trait objects!

    struct ArcWrapped<T>(PhantomData<T>);

    unsafe impl<T: Wake + 'static> UnsafeWake for ArcWrapped<T> {
        unsafe fn clone_raw(&self) -> Waker {
            let me: *const ArcWrapped<T> = self;
            let arc = (*(&me as *const *const ArcWrapped<T> as *const Arc<T>)).clone();
            Waker::from(arc)
        }

        unsafe fn drop_raw(&self) {
            let mut me: *const ArcWrapped<T> = self;
            let me = &mut me as *mut *const ArcWrapped<T> as *mut Arc<T>;
            ptr::drop_in_place(me);
        }

        unsafe fn wake(&self) {
            let me: *const ArcWrapped<T> = self;
            T::wake(&*(&me as *const *const ArcWrapped<T> as *const Arc<T>))
        }
    }

    impl<T> From<Arc<T>> for Waker
        where T: Wake + 'static,
    {
        fn from(rc: Arc<T>) -> Waker {
            unsafe {
                let ptr = mem::transmute::<Arc<T>, *const ArcWrapped<T>>(rc);
                Waker::new(ptr)
            }
        }
    }
}