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 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326
#![no_std] //! cooked_waker provides safe traits for working with //! [`std::task::Waker`][Waker] and, more importantly, a set of derives for //! safely converting normal, safe rust types into `Waker` instances. It cooks //! `RawWaker` and `RawWakerVTable`, making them safe for consumption. //! //! It provides the [`Wake`] and [`WakeRef`] traits, which correspond to the //! [`wake`][Waker::wake] and [`wake_by_ref`][Waker::wake_by_ref] methods //! on [`std::task::Waker`][Waker], and it provides implenetations of these //! types for the common reference & pointer types (`Arc`, `Rc`, `&'static`, //! etc). //! //! Additionally, it provides [`IntoWaker`], which allows converting any //! `Wake + Clone` type into a [`Waker`]. This trait is automatically derived //! for any `Wake + Clone + Send + Sync + 'static` type. //! //! # Basic example //! //! ``` //! use cooked_waker::{Wake, WakeRef, IntoWaker}; //! use std::sync::atomic::{AtomicUsize, Ordering}; //! use std::task::Waker; //! //! static wake_ref_count: AtomicUsize = AtomicUsize::new(0); //! static wake_value_count: AtomicUsize = AtomicUsize::new(0); //! static drop_count: AtomicUsize = AtomicUsize::new(0); //! //! // A simple Waker struct that atomically increments the relevant static //! // counters. //! #[derive(Debug, Clone)] //! struct StaticWaker; //! //! impl WakeRef for StaticWaker { //! fn wake_by_ref(&self) { //! wake_ref_count.fetch_add(1, Ordering::SeqCst); //! } //! } //! //! impl Wake for StaticWaker { //! fn wake(self) { //! wake_value_count.fetch_add(1, Ordering::SeqCst); //! } //! } //! //! impl Drop for StaticWaker { //! fn drop(&mut self) { //! drop_count.fetch_add(1, Ordering::SeqCst); //! } //! } //! //! assert_eq!(drop_count.load(Ordering::SeqCst), 0); //! //! let waker = StaticWaker; //! { //! let waker1: Waker = waker.into_waker(); //! //! waker1.wake_by_ref(); //! assert_eq!(wake_ref_count.load(Ordering::SeqCst), 1); //! //! let waker2: Waker = waker1.clone(); //! waker2.wake_by_ref(); //! assert_eq!(wake_ref_count.load(Ordering::SeqCst), 2); //! //! waker1.wake(); //! assert_eq!(wake_value_count.load(Ordering::SeqCst), 1); //! assert_eq!(drop_count.load(Ordering::SeqCst), 1); //! } //! assert_eq!(drop_count.load(Ordering::SeqCst), 2); //! ``` //! //! # Arc example //! //! ``` //! use cooked_waker::{Wake, WakeRef, IntoWaker}; //! use std::sync::atomic::{AtomicUsize, Ordering}; //! use std::sync::Arc; //! use std::task::Waker; //! //! // A simple struct that counts the number of times it is awoken. Can't //! // be awoken by value (because that would discard the counter), so we //! // must instead wrap it in an Arc. //! #[derive(Debug, Default)] //! struct Counter { //! // We use atomic usize because we need Send + Sync and also interior //! // mutability //! count: AtomicUsize, //! } //! //! impl Counter { //! fn get(&self) -> usize { //! self.count.load(Ordering::SeqCst) //! } //! } //! //! impl WakeRef for Counter { //! fn wake_by_ref(&self) { //! let _prev = self.count.fetch_add(1, Ordering::SeqCst); //! } //! } //! //! let counter_handle = Arc::new(Counter::default()); //! //! // Create an std::task::Waker //! let waker: Waker = counter_handle.clone().into_waker(); //! //! waker.wake_by_ref(); //! waker.wake_by_ref(); //! //! let waker2 = waker.clone(); //! waker2.wake_by_ref(); //! //! // This calls Counter::wake_by_ref because the Arc doesn't have exclusive //! // ownership of the underlying Counter //! waker2.wake(); //! //! assert_eq!(counter_handle.get(), 4); //! ``` extern crate alloc; use alloc::boxed::Box; use alloc::rc; use alloc::sync as arc; use core::task::{RawWaker, RawWakerVTable, Waker}; use stowaway::{self, Stowaway}; /// Wakers that can wake by reference. This trait is used to enable a [`Wake`] /// implementation for types that don't own an underlying handle, like `Arc<T>` /// and `&T`. /// /// This trait is implemented for most container and reference types, like /// `&T where T: WakeRef`, `Box<T: WakeRef>`, and `Arc<T: WakeRef>`. pub trait WakeRef { /// Wake up the task by reference. In general [`Wake::wake`] should be /// preferred, if available, as it's probably more efficient. /// /// A [`Waker`] created by [`IntoWaker`] will call this method through /// [`Waker::wake_by_ref`]. fn wake_by_ref(&self); } /// Wakers that can wake by value. This is the primary means of waking a task. /// /// This trait is implemented for most container types, like `Box<T: Wake>` /// and `Option<T: Wake>`. It is also implemented for shared pointer types like /// `Arc<T>` and `&T`, but those implementations call `T::wake_by_ref`, because /// they don't have ownership of the underlying `T`. pub trait Wake: WakeRef + Sized { /// Wake up the task by value. By default, this simply calls /// [`WakeRef::wake_by_ref`]. /// /// A [`Waker`] created by [`IntoWaker`] will call this method through /// [`Waker::wake`]. #[inline] fn wake(self) { self.wake_by_ref() } } /// Objects that can be converted into an [`Waker`]. This trait is /// automatically implemented for any `Wake + Clone + Send + Sync + 'static` /// type. /// /// The implementation of this trait sets up a [`RawWakerVTable`] for the type, /// and arranges a conversion into a [`Waker`] through the [`stowaway`] crate, /// which allows packing the bytes of any sized type into a pointer (boxing it /// if it's too large to fit). This means that "large" waker structs will /// simply be boxed, but wakers that contain a single `Box` or `Arc` field /// (or any data smaller or the same size as a pointer) will simply move their /// pointer directly. This `Waker` will then call the relevant `Wake`, /// `RefWake`, or `Clone` methods throughout its lifecycle. /// /// It should never be necessary to implement this trait manually. /// /// [`RawWakerVTable`]: core::task::RawWakerVTable /// [`Waker`]: core::task::Waker /// [`stowaway`]: https://docs.rs/stowaway pub trait IntoWaker: Wake + Clone + Send + Sync + 'static { /// The RawWakerVTable for this type. This should never be used directly; /// it is entirely handled by `into_waker`. It is present as an associated /// const because that's the only way for it to work in generic contexts. #[doc(hidden)] const VTABLE: &'static RawWakerVTable; /// Convert this object into a `Waker`. #[must_use] fn into_waker(self) -> Waker; } impl<T: Wake + Clone + Send + Sync + 'static> IntoWaker for T { const VTABLE: &'static RawWakerVTable = &RawWakerVTable::new( // clone |raw| { let raw = raw as *mut (); let waker: &T = unsafe { stowaway::ref_from_stowed(&raw) }; let cloned = waker.clone(); let stowed = Stowaway::new(cloned); RawWaker::new(Stowaway::into_raw(stowed), T::VTABLE) }, // wake by value |raw| { let waker: T = unsafe { stowaway::unstow(raw as *mut ()) }; Wake::wake(waker); }, // wake by ref |raw| { let raw = raw as *mut (); let waker: &T = unsafe { stowaway::ref_from_stowed(&raw) }; WakeRef::wake_by_ref(waker) }, // Drop |raw| { let _waker: Stowaway<T> = unsafe { Stowaway::from_raw(raw as *mut ()) }; }, ); fn into_waker(self) -> Waker { let stowed = Stowaway::new(self); let raw_waker = RawWaker::new(Stowaway::into_raw(stowed), T::VTABLE); unsafe { Waker::from_raw(raw_waker) } } } // Waker implementations for std types. Feel free to open PRs for additional // stdlib types here. // We'd prefer to implement WakeRef for T: Deref<Target=WakeRef>, but that // results in type coherence issues with non-deref stdlib types. impl<T: WakeRef> WakeRef for &T { #[inline] fn wake_by_ref(&self) { T::wake_by_ref(*self) } } impl<T: WakeRef> Wake for &T {} impl<T: WakeRef + ?Sized> WakeRef for Box<T> { #[inline] fn wake_by_ref(&self) { T::wake_by_ref(self.as_ref()) } } impl<T: Wake> Wake for Box<T> { #[inline] fn wake(self) { T::wake(*self) } } impl<T: WakeRef + ?Sized> WakeRef for arc::Arc<T> { #[inline] fn wake_by_ref(&self) { T::wake_by_ref(self.as_ref()) } } impl<T: WakeRef + ?Sized> Wake for arc::Arc<T> {} impl<T: WakeRef + ?Sized> WakeRef for arc::Weak<T> { #[inline] fn wake_by_ref(&self) { self.upgrade().wake() } } impl<T: WakeRef + ?Sized> Wake for arc::Weak<T> {} impl<T: WakeRef + ?Sized> WakeRef for rc::Rc<T> { #[inline] fn wake_by_ref(&self) { T::wake_by_ref(self.as_ref()) } } impl<T: WakeRef + ?Sized> Wake for rc::Rc<T> { #[inline] fn wake(self) { T::wake_by_ref(self.as_ref()) } } impl<T: WakeRef + ?Sized> WakeRef for rc::Weak<T> { #[inline] fn wake_by_ref(&self) { self.upgrade().wake() } } impl<T: WakeRef + ?Sized> Wake for rc::Weak<T> {} impl<T: WakeRef> WakeRef for Option<T> { #[inline] fn wake_by_ref(&self) { if let Some(waker) = self { waker.wake_by_ref() } } } impl<T: Wake> Wake for Option<T> { #[inline] fn wake(self) { if let Some(waker) = self { waker.wake() } } } impl WakeRef for Waker { #[inline] fn wake_by_ref(&self) { Waker::wake_by_ref(self) } } impl Wake for Waker { #[inline] fn wake(self) { Waker::wake(self) } }