-
Notifications
You must be signed in to change notification settings - Fork 36
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
improve readability, provenance compliance, and documentation for sha…
…red arc waker
- Loading branch information
Showing
5 changed files
with
141 additions
and
88 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,92 +1,136 @@ | ||
| use core::task::{RawWaker, RawWakerVTable, Waker}; | ||
| use std::sync::Arc; | ||
| //! To save on allocations, we avoid making a separate Arc Waker for every subfuture. | ||
| //! Rather, we have all N Wakers share a single Arc, and use a "redirect" mechanism to allow different wakers to be distinguished. | ||
| //! The mechanism works as follows. | ||
| //! The Arc contains 2 things: | ||
| //! - the Readiness structure ([ReadinessArray][super::array::ReadinessArray] / [ReadinessVec][super::vec::ReadinessVec]) | ||
| //! - the redirect array. | ||
| //! The redirect array contains N repeated copies of the pointer to the Arc itself (obtained by `Arc::into_raw`). | ||
| //! The Waker for the `i`th subfuture points to the `i`th item in the redirect array. | ||
| //! (i.e. the Waker pointer is `*const *const A` where `A` is the type of the item in the Arc) | ||
| //! When the Waker is woken, we deref it twice (giving reference to the content of the Arc), | ||
| //! and compare it to the address of the redirect slice. | ||
| //! The difference tells us the index of the waker. We can then record this woken index in the Readiness. | ||
| //! | ||
| //! ```text | ||
| //! ┌───────────────────────────┬──────────────┬──────────────┐ | ||
| //! │ │ │ │ | ||
| //! │ / ┌─────────────┬──────┼───────┬──────┼───────┬──────┼───────┬─────┐ \ | ||
| //! ▼ / │ │ │ │ │ │ │ │ │ \ | ||
| //! Arc < │ Readiness │ redirect[0] │ redirect[1] │ redirect[2] │ ... │ > | ||
| //! ▲ \ │ │ │ │ │ │ / | ||
| //! │ \ └─────────────┴──────▲───────┴──────▲───────┴──────▲───────┴─────┘ / | ||
| //! │ │ │ │ | ||
| //! └─┐ ┌───────────────┘ │ │ | ||
| //! │ │ │ │ | ||
| //! │ │ ┌──────────────────┘ │ | ||
| //! │ │ │ │ | ||
| //! │ │ │ ┌─────────────────────┘ | ||
| //! │ │ │ │ | ||
| //! │ │ │ │ | ||
| //! ┌────┼────┬────┼──────┬────┼──────┬────┼──────┬─────┐ | ||
| //! │ │ │ │ │ │ │ │ │ │ | ||
| //! │ │ wakers[0] │ wakers[1] │ wakers[2] │ ... │ | ||
| //! │ │ │ │ │ │ | ||
| //! └─────────┴───────────┴───────────┴───────────┴─────┘ | ||
| //! ``` | ||
| // In the diagram below, `A` is the upper block. | ||
| // It is a struct that implements WakeDataContainer (so either WakerVecInner or WakerArrayInner). | ||
| // The lower block is either WakerVec or WakerArray. Each waker there points to a slot of wake_data in `A`. | ||
| // Every one of these slots contain a pointer to the Arc wrapping `A` itself. | ||
| // Wakers figure out their indices by comparing the address they are pointing to to `wake_data`'s start address. | ||
| // | ||
| // ┌───────────────────────────┬──────────────┬──────────────┐ | ||
| // │ │ │ │ | ||
| // │ / ┌─────────────┬──────┼───────┬──────┼───────┬──────┼───────┬─────┐ \ | ||
| // ▼ / │ │ │ │ │ │ │ │ │ \ | ||
| // Arc < │ Readiness │ wake_data[0] │ wake_data[1] │ wake_data[2] │ ... │ > | ||
| // ▲ \ │ │ │ │ │ │ / | ||
| // │ \ └─────────────┴──────▲───────┴──────▲───────┴──────▲───────┴─────┘ / | ||
| // │ │ │ │ | ||
| // └─┐ ┌───────────────┘ │ │ | ||
| // │ │ │ │ | ||
| // │ │ ┌──────────────────┘ │ | ||
| // │ │ │ │ | ||
| // │ │ │ ┌─────────────────────┘ | ||
| // │ │ │ │ | ||
| // │ │ │ │ | ||
| // ┌────┼────┬────┼──────┬────┼──────┬────┼──────┬─────┐ | ||
| // │ │ │ │ │ │ │ │ │ │ | ||
| // │ Inner │ wakers[0] │ wakers[1] │ wakers[2] │ ... │ | ||
| // │ │ │ │ │ │ | ||
| // └─────────┴───────────┴───────────┴───────────┴─────┘ | ||
| // TODO: Right now each waker gets its own redirect slot. | ||
| // We can save space by making size_of::<*const _>() wakers share the same slot. | ||
| // With such change, in 64-bit system, the redirect array/vec would only need ⌈N/8⌉ slots instead of N. | ||
|
|
||
| // TODO: Right now each waker gets its own wake_data slot. | ||
| // We can save space by making size_of::<usize>() wakers share the same slot. | ||
| // With such change, in 64-bit system, the wake_data array/vec would only need ⌈N/8⌉ slots instead of N. | ||
| use core::task::{RawWaker, RawWakerVTable, Waker}; | ||
| use std::sync::Arc; | ||
|
|
||
| pub(super) trait WakeDataContainer { | ||
| /// Get the reference of the wake_data slice. This is used to compute the index. | ||
| fn get_wake_data_slice(&self) -> &[*const Self]; | ||
| /// A trait to be implemented on [super::WakerArray] and [super::WakerVec] for polymorphism. | ||
| /// These are the type that goes in the Arc. They both contain the Readiness and the redirect array/vec. | ||
| pub(super) trait SharedArcContent { | ||
| /// Get the reference of the redirect slice. This is used to compute the index. | ||
| fn get_redirect_slice(&self) -> &[*const Self]; | ||
| /// Called when the subfuture at the specified index should be polled. | ||
| /// Should call `Readiness::set_ready`. | ||
| fn wake_index(&self, index: usize); | ||
| } | ||
| pub(super) unsafe fn waker_for_wake_data_slot<A: WakeDataContainer>( | ||
|
|
||
| /// Create one waker following the mechanism described in the [module][self] doc. | ||
| /// The following must be upheld for safety: | ||
| /// - `pointer` must points to a slot in the redirect array. | ||
| /// - that slot must contain a pointer obtained by `Arc::<A>::into_raw`. | ||
| /// - the Arc must still be alive at the time this function is called. | ||
| /// The following should be upheld for correct behavior: | ||
| /// - calling `SharedArcContent::get_redirect_slice` on the content of the Arc should give the redirect array within which `pointer` points to. | ||
| #[deny(unsafe_op_in_unsafe_fn)] | ||
| pub(super) unsafe fn waker_from_redirect_position<A: SharedArcContent>( | ||
| pointer: *const *const A, | ||
| ) -> Waker { | ||
| unsafe fn clone_waker<A: WakeDataContainer>(pointer: *const ()) -> RawWaker { | ||
| /// Create a Waker from a type-erased pointer. | ||
| /// The pointer must satisfy the safety constraints listed in the wrapping function's documentation. | ||
| unsafe fn create_waker<A: SharedArcContent>(pointer: *const ()) -> RawWaker { | ||
| // Retype the type-erased pointer. | ||
| let pointer = pointer as *const *const A; | ||
| let raw = *pointer; // This is the raw pointer of Arc<Inner>. | ||
|
|
||
| // We're creating a new Waker, so we need to increment the count. | ||
| Arc::increment_strong_count(raw); | ||
| // SAFETY: The constraints listed for the wrapping function documentation means | ||
| // - `*pointer` is an `*const A` obtained from `Arc::<A>::into_raw`. | ||
| // - the Arc is alive. | ||
| // So this operation is safe. | ||
| unsafe { Arc::increment_strong_count(*pointer) }; | ||
|
|
||
| RawWaker::new(pointer as *const (), create_vtable::<A>()) | ||
| } | ||
|
|
||
| // Convert a pointer to a wake_data slot to the Arc<Inner>. | ||
| unsafe fn to_arc<A: WakeDataContainer>(pointer: *const *const A) -> Arc<A> { | ||
| let raw = *pointer; | ||
| Arc::from_raw(raw) | ||
| } | ||
| unsafe fn wake<A: WakeDataContainer, const BY_REF: bool>(pointer: *const ()) { | ||
| /// Invoke `SharedArcContent::wake_index` with the index in the redirect slice where this pointer points to. | ||
| /// The pointer must satisfy the safety constraints listed in the wrapping function's documentation. | ||
| unsafe fn wake_by_ref<A: SharedArcContent>(pointer: *const ()) { | ||
| // Retype the type-erased pointer. | ||
| let pointer = pointer as *const *const A; | ||
| let arc = to_arc::<A>(pointer); | ||
| // Calculate the index | ||
| let index = ((pointer as usize) // This is the slot our pointer points to. | ||
| - (arc.get_wake_data_slice() as *const [*const A] as *const () as usize)) // This is the starting address of wake_data. | ||
| / std::mem::size_of::<*const A>(); | ||
|
|
||
| arc.wake_index(index); | ||
| // SAFETY: we are already requiring `pointer` to point to a slot in the redirect array. | ||
| let raw: *const A = unsafe { *pointer }; | ||
| // SAFETY: we are already requiring the pointer in the redirect array slot to be obtained from `Arc::into_raw`. | ||
| let arc_content: &A = unsafe { &*raw }; | ||
|
|
||
| // Dropping the Arc would decrement the strong count. | ||
| // We only want to do that when we're not waking by ref. | ||
| if BY_REF { | ||
| std::mem::forget(arc); | ||
| } else { | ||
| std::mem::drop(arc); | ||
| } | ||
| // Calculate the index. | ||
| // This is your familiar pointer math | ||
| // `item_address = array_address + (index * item_size)` | ||
| // rearranged to | ||
| // `index = (item_address - array_address) / item_size`. | ||
| let item_address = sptr::Strict::addr(pointer); | ||
| let redirect_slice_address = sptr::Strict::addr(arc_content.get_redirect_slice().as_ptr()); | ||
| let redirect_item_size = core::mem::size_of::<*const A>(); // the size of each item in the redirect slice | ||
| let index = (item_address - redirect_slice_address) / redirect_item_size; | ||
|
|
||
| arc_content.wake_index(index); | ||
| } | ||
| unsafe fn drop_waker<A: WakeDataContainer>(pointer: *const ()) { | ||
|
|
||
| /// The pointer must satisfy the safety constraints listed in the wrapping function's documentation. | ||
| unsafe fn drop_waker<A: SharedArcContent>(pointer: *const ()) { | ||
| // Retype the type-erased pointer. | ||
| let pointer = pointer as *const *const A; | ||
| let arc = to_arc::<A>(pointer); | ||
| // Decrement the strong count by dropping the Arc. | ||
| std::mem::drop(arc); | ||
|
|
||
| // SAFETY: we are already requiring `pointer` to point to a slot in the redirect array. | ||
| let raw = unsafe { *pointer }; | ||
| // SAFETY: we are already requiring the pointer in the redirect array slot to be obtained from `Arc::into_raw`. | ||
| unsafe { Arc::decrement_strong_count(raw) }; | ||
| } | ||
| fn create_vtable<A: WakeDataContainer>() -> &'static RawWakerVTable { | ||
|
|
||
| /// The pointer must satisfy the safety constraints listed in the wrapping function's documentation. | ||
| unsafe fn wake<A: SharedArcContent>(pointer: *const ()) { | ||
| // SAFETY: we are already requiring the constraints of `wake_by_ref` and `drop_waker`. | ||
| unsafe { | ||
| wake_by_ref::<A>(pointer); | ||
| drop_waker::<A>(pointer); | ||
| } | ||
| } | ||
|
|
||
| fn create_vtable<A: SharedArcContent>() -> &'static RawWakerVTable { | ||
| &RawWakerVTable::new( | ||
| clone_waker::<A>, | ||
| wake::<A, false>, | ||
| wake::<A, true>, | ||
| create_waker::<A>, | ||
| wake::<A>, | ||
| wake_by_ref::<A>, | ||
| drop_waker::<A>, | ||
| ) | ||
| } | ||
| Waker::from_raw(clone_waker::<A>(pointer as *const ())) | ||
| // SAFETY: All our vtable functions adhere to the RawWakerVTable contract, | ||
| // and we are already requiring that `pointer` is what our functions expect. | ||
| unsafe { Waker::from_raw(create_waker::<A>(pointer as *const ())) } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters