1+ use std:: mem:: MaybeUninit ;
12use std:: pin:: Pin ;
2-
33use std:: task:: Context ;
44
5- use cxx:: ExternType ;
6-
5+ use crate :: ffi:: { GuardedRustPromiseAwaiter , GuardedRustPromiseAwaiterRepr } ;
76use crate :: waker:: try_into_kj_waker_ptr;
87
9- use crate :: lazy_pin_init:: LazyPinInit ;
10-
118// =======================================================================================
129// GuardedRustPromiseAwaiter
1310
14- #[ path = "awaiter.h.rs" ]
15- mod awaiter_h;
16- pub use awaiter_h:: GuardedRustPromiseAwaiter ;
17-
1811// Safety: KJ Promises are not associated with threads, but with event loops at construction time.
1912// Therefore, they can be polled from any thread, as long as that thread has the correct event loop
2013// active at the time of the call to `poll()`. If the correct event loop is not active,
2114// GuardedRustPromiseAwaiter's API will panic. (The Guarded- prefix refers to the C++ class template
2215// ExecutorGuarded, which enforces the correct event loop requirement.)
2316unsafe impl Send for GuardedRustPromiseAwaiter { }
2417
25- impl Drop for GuardedRustPromiseAwaiter {
26- fn drop ( & mut self ) {
27- // Pin safety:
28- // The pin crate suggests implementing drop traits for address-sensitive types with an inner
29- // function which accepts a `Pin<&mut Type>` parameter, to help uphold pinning guarantees.
30- // However, since our drop function is actually a C++ destructor to which we must pass a raw
31- // pointer, there is no benefit in creating a Pin from `self`.
32- //
33- // https://doc.rust-lang.org/std/pin/index.html#implementing-drop-for-types-with-address-sensitive-states
34- //
35- // Pointer safety:
36- // 1. Pointer to self is non-null, and obviously points to valid memory.
37- // 2. We do not read or write to the OwnPromiseNode's memory, so there are no atomicity nor
38- // interleaved pointer/reference access concerns.
39- //
40- // https://doc.rust-lang.org/std/ptr/index.html#safety
41- unsafe {
42- crate :: ffi:: guarded_rust_promise_awaiter_drop_in_place ( PtrGuardedRustPromiseAwaiter (
43- self ,
44- ) ) ;
45- }
46- }
47- }
48-
49- // Safety: We have a static_assert in awaiter.c++ which breaks if you change the size or alignment
50- // of the C++ definition of GuardedRustPromiseAwaiter, with instructions to regenerate the bindgen-
51- // generated type. I couldn't figure out how to static_assert on the actual generated Rust struct,
52- // though, so it's not perfect. Ideally we'd run bindgen in the build system.
53- //
54- // https://docs.rs/cxx/latest/cxx/trait.ExternType.html#integrating-with-bindgen-generated-types
55- unsafe impl ExternType for GuardedRustPromiseAwaiter {
56- type Id = cxx:: type_id!( "::kj_rs::GuardedRustPromiseAwaiter" ) ;
57- type Kind = cxx:: kind:: Opaque ;
58- }
59-
60- // This Ptr- type is copied from dtolnay's `Box<dyn Trait>` example:
61- // https://github.com/dtolnay/cxx/pull/672/files#diff-0ca4eb15f64c7f6bd562b9efb702a1f6020ee83a6f1c686054cfd3fb314edfe1R16
62- //
63- // I don't entirely understand why it is necessary. When I tried to replace it with raw pointers, I
64- // encountered segfaults, so I put it back.
65- //
66- // TODO(someday): Examine the generated code with and without this wrapper type and figure out why
67- // it's necessary. Do we need this sort of wrapper for our macro-generated Future and Promise
68- // types, too?
69- #[ repr( transparent) ]
70- pub struct PtrGuardedRustPromiseAwaiter ( * mut GuardedRustPromiseAwaiter ) ;
71-
72- // Safety: Raw pointers are the same size in both languages.
73- unsafe impl ExternType for PtrGuardedRustPromiseAwaiter {
74- type Id = cxx:: type_id!( "::kj_rs::PtrGuardedRustPromiseAwaiter" ) ;
75- type Kind = cxx:: kind:: Trivial ;
76- }
77-
7818// =======================================================================================
7919// Await syntax for OwnPromiseNode
8020
@@ -83,7 +23,8 @@ use crate::OwnPromiseNode;
8323pub struct PromiseAwaiter < Data : std:: marker:: Unpin > {
8424 node : Option < OwnPromiseNode > ,
8525 pub ( crate ) data : Data ,
86- awaiter : LazyPinInit < GuardedRustPromiseAwaiter > ,
26+ awaiter : MaybeUninit < GuardedRustPromiseAwaiterRepr > ,
27+ awaiter_initialized : bool ,
8728 // Safety: `option_waker` must be declared after `awaiter`, because `awaiter` contains a reference
8829 // to `option_waker`. This ensures `option_waker` will be dropped after `awaiter`.
8930 option_waker : OptionWaker ,
@@ -94,7 +35,8 @@ impl<Data: std::marker::Unpin> PromiseAwaiter<Data> {
9435 PromiseAwaiter {
9536 node : Some ( node) ,
9637 data,
97- awaiter : LazyPinInit :: uninit ( ) ,
38+ awaiter : MaybeUninit :: uninit ( ) ,
39+ awaiter_initialized : false ,
9840 option_waker : OptionWaker :: empty ( ) ,
9941 }
10042 }
@@ -104,45 +46,43 @@ impl<Data: std::marker::Unpin> PromiseAwaiter<Data> {
10446/// Panics if `node` is None.
10547#[ must_use]
10648 pub fn get_awaiter ( mut self : Pin < & mut Self > ) -> Pin < & mut GuardedRustPromiseAwaiter > {
107- // On our first invocation, `node` will be Some, and `get_awaiter` will forward its
108- // contents into GuardedRustPromiseAwaiter's constructor. On all subsequent invocations, `node`
109- // will be None and the constructor will not run.
110- let node = self . as_mut ( ) . node . take ( ) ;
111-
112- // Safety: `awaiter` stores `rust_waker_ptr` and uses it to call `wake()`. Note that
113- // `awaiter` is `self.awaiter`, which lives before `self.option_waker`. Since struct members
114- // are dropped in declaration order, the `rust_waker_ptr` that `awaiter` stores will always
115- // be valid during its lifetime.
116- //
117- // We pass a mutable pointer to C++. This is safe, because our use of the OptionWaker inside
118- // of `std::task::Waker` is synchronized by ensuring we only allow calls to `poll()` on the
119- // thread with the Promise's event loop active.
120- let rust_waker_ptr = & raw mut self . as_mut ( ) . option_waker ;
121-
122- // Safety:
123- // 1. We do not implement Unpin for PromiseAwaiter.
124- // 2. Our Drop trait implementation does not move the awaiter value, nor do we use
125- // `repr(packed)` anywhere.
126- // 3. The backing memory is inside our pinned Future, so we can be assured our Drop trait
127- // implementation will run before Rust re-uses the memory.
128- //
129- // https://doc.rust-lang.org/std/pin/index.html#choosing-pinning-to-be-structural-for-field
130- let awaiter = unsafe { self . map_unchecked_mut ( |s| & mut s. awaiter ) } ;
131-
132- // Safety:
133- // 1. We trust that LazyPinInit's implementation passed us a valid pointer to an
134- // uninitialized GuardedRustPromiseAwaiter.
135- // 2. We do not read or write to the GuardedRustPromiseAwaiter's memory, so there are no atomicity
136- // nor interleaved pointer reference access concerns.
137- //
138- // https://doc.rust-lang.org/std/ptr/index.html#safety
139- awaiter. get_or_init ( move |ptr : * mut GuardedRustPromiseAwaiter | unsafe {
140- crate :: ffi:: guarded_rust_promise_awaiter_new_in_place (
141- PtrGuardedRustPromiseAwaiter ( ptr) ,
142- rust_waker_ptr,
143- node. expect ( "node should be Some in call to init()" ) ,
144- ) ;
145- } )
49+ // Safety: We never move out of `this`.
50+ let this = unsafe { Pin :: into_inner_unchecked ( self . as_mut ( ) ) } ;
51+
52+ // Initialize the awaiter if not already done
53+ if !this. awaiter_initialized {
54+ // On our first invocation, `node` will be Some, and `get_awaiter` will forward its
55+ // contents into GuardedRustPromiseAwaiter's constructor. On all subsequent invocations, `node`
56+ // will be None and the constructor will not run.
57+ let node = this. node . take ( ) ;
58+
59+ // Safety: `awaiter` stores `rust_waker_ptr` and uses it to call `wake()`. Note that
60+ // `awaiter` is `this.awaiter`, which lives before `this.option_waker`.
61+ // Since we drop awaiter manually, the `rust_waker_ptr` that `awaiter` stores will always
62+ // be valid during its lifetime.
63+ //
64+ // We pass a mutable pointer to C++. This is safe, because our use of the OptionWaker inside
65+ // of `std::task::Waker` is synchronized by ensuring we only allow calls to `poll()` on the
66+ // thread with the Promise's event loop active.
67+ let rust_waker_ptr = & raw mut this. option_waker ;
68+
69+ // Safety: The memory slot is valid and this type ensures that it will stay pinned.
70+ unsafe {
71+ crate :: ffi:: guarded_rust_promise_awaiter_new_in_place (
72+ this. awaiter . as_mut_ptr ( ) . cast :: < GuardedRustPromiseAwaiter > ( ) ,
73+ rust_waker_ptr,
74+ node. expect ( "node should be Some in call to init()" ) ,
75+ ) ;
76+ }
77+ this. awaiter_initialized = true ;
78+ }
79+
80+ // Safety: `this.awaiter` is pinned since `self` is pinned.
81+ unsafe {
82+ let raw = this. awaiter . assume_init_mut ( ) as * mut GuardedRustPromiseAwaiterRepr ;
83+ let raw = raw. cast :: < GuardedRustPromiseAwaiter > ( ) ;
84+ Pin :: new_unchecked ( & mut * raw)
85+ }
14686 }
14787
14888 pub fn poll ( mut self : Pin < & mut Self > , cx : & mut Context ) -> bool {
@@ -153,6 +93,18 @@ impl<Data: std::marker::Unpin> PromiseAwaiter<Data> {
15393 }
15494}
15595
96+ impl < Data : std:: marker:: Unpin > Drop for PromiseAwaiter < Data > {
97+ fn drop ( & mut self ) {
98+ if self . awaiter_initialized {
99+ unsafe {
100+ crate :: ffi:: guarded_rust_promise_awaiter_drop_in_place (
101+ self . awaiter . as_mut_ptr ( ) . cast :: < GuardedRustPromiseAwaiter > ( )
102+ ) ;
103+ }
104+ }
105+ }
106+ }
107+
156108// =======================================================================================
157109// OptionWaker and WakerRef
158110
0 commit comments