Revert "[workerd-cxx] remove + Send" (#37)

This reverts commit fc3d2a8.

Author: mikea

kj-rs/awaiter.c++

@@ -13,10 +13,10 @@ namespace kj_rs {
 // that end, we use bindgen to generate an opaque FFI type of known size for RustPromiseAwaiter in
 // awaiter.h.rs.
 //
-static_assert(sizeof(RustPromiseAwaiter) == sizeof(RustPromiseAwaiterRepr),
-    "RustPromiseAwaiter size changed, you must update lib.rs ffi");
-static_assert(alignof(RustPromiseAwaiter) == alignof(RustPromiseAwaiterRepr),
-    "RustPromiseAwaiter alignment changed, you must update lib.rs ffi");
+static_assert(sizeof(GuardedRustPromiseAwaiter) == sizeof(GuardedRustPromiseAwaiterRepr),
+    "GuardedRustPromiseAwaiter size changed, you must update lib.rs ffi");
+static_assert(alignof(GuardedRustPromiseAwaiter) == alignof(GuardedRustPromiseAwaiterRepr),
+    "GuardedRustPromiseAwaiter alignment changed, you must update lib.rs ffi");
 
 RustPromiseAwaiter::RustPromiseAwaiter(
     OptionWaker& optionWaker, OwnPromiseNode nodeParam, kj::SourceLocation location)
@@ -132,11 +132,11 @@ OwnPromiseNode RustPromiseAwaiter::take_own_promise_node() {
   return kj::mv(node);
 }
 
-void rust_promise_awaiter_new_in_place(
-    RustPromiseAwaiter* ptr, OptionWaker* optionWaker, OwnPromiseNode node) {
+void guarded_rust_promise_awaiter_new_in_place(
+    GuardedRustPromiseAwaiter* ptr, OptionWaker* optionWaker, OwnPromiseNode node) {
   kj::ctor(*ptr, *optionWaker, kj::mv(node));
 }
-void rust_promise_awaiter_drop_in_place(RustPromiseAwaiter* ptr) {
+void guarded_rust_promise_awaiter_drop_in_place(GuardedRustPromiseAwaiter* ptr) {
   kj::dtor(*ptr);
 }
 
@@ -204,16 +204,20 @@ void FuturePollEvent::tracePromise(kj::_::TraceBuilder& builder, bool stopAtNext
   }
 }
 
-FuturePollEvent::PollScope::PollScope(FuturePollEvent& futurePollEvent): event(futurePollEvent) {
+FuturePollEvent::PollScope::PollScope(FuturePollEvent& futurePollEvent): holder(futurePollEvent) {
   futurePollEvent.enterPollScope();
 }
 
 FuturePollEvent::PollScope::~PollScope() noexcept(false) {
-  event.exitPollScope(reset());
+  holder.get().futurePollEvent.exitPollScope(reset());
 }
 
 kj::Maybe<FuturePollEvent&> FuturePollEvent::PollScope::tryGetFuturePollEvent() const {
-  return event;
+  KJ_IF_SOME(h, holder.tryGet()) {
+    return h.futurePollEvent;
+  } else {
+    return kj::none;
+  }
 }
 
 }  // namespace kj_rs

kj-rs/awaiter.h

@@ -106,8 +106,24 @@ class RustPromiseAwaiter final: public kj::_::Event,
   OwnPromiseNode node;
 };
 
-void rust_promise_awaiter_new_in_place(RustPromiseAwaiter*, OptionWaker*, OwnPromiseNode);
-void rust_promise_awaiter_drop_in_place(RustPromiseAwaiter*);
+// We force Rust to call our `poll()` overloads using this ExecutorGuarded wrapper around the actual
+// RustPromiseAwaiter class. This allows us to assume all calls that reach RustPromiseAwaiter itself
+// are on the correct thread.
+struct GuardedRustPromiseAwaiter: ExecutorGuarded<RustPromiseAwaiter> {
+  // We need to inherit constructors or else placement-new will try to aggregate-initialize us.
+  using ExecutorGuarded<RustPromiseAwaiter>::ExecutorGuarded;
+
+  bool poll(const WakerRef& waker, const KjWaker* maybeKjWaker) {
+    return get().poll(waker, maybeKjWaker);
+  }
+  OwnPromiseNode take_own_promise_node() {
+    return get().take_own_promise_node();
+  }
+};
+
+void guarded_rust_promise_awaiter_new_in_place(
+    GuardedRustPromiseAwaiter*, OptionWaker*, OwnPromiseNode);
+void guarded_rust_promise_awaiter_drop_in_place(GuardedRustPromiseAwaiter*);
 
 // =======================================================================================
 // FuturePollEvent
@@ -171,7 +187,10 @@ class FuturePollEvent::PollScope: public LazyArcWaker {
   kj::Maybe<FuturePollEvent&> tryGetFuturePollEvent() const override;
 
  private:
-  FuturePollEvent& event;
+  struct FuturePollEventHolder {
+    FuturePollEvent& futurePollEvent;
+  };
+  ExecutorGuarded<FuturePollEventHolder> holder;
 };
 
 // =======================================================================================

kj-rs/awaiter.rs

@@ -2,9 +2,19 @@ use std::mem::MaybeUninit;
 use std::pin::Pin;
 use std::task::Context;
 
-use crate::ffi::{RustPromiseAwaiter, RustPromiseAwaiterRepr};
+use crate::ffi::{GuardedRustPromiseAwaiter, GuardedRustPromiseAwaiterRepr};
 use crate::waker::try_into_kj_waker_ptr;
 
+// =======================================================================================
+// GuardedRustPromiseAwaiter
+
+// Safety: KJ Promises are not associated with threads, but with event loops at construction time.
+// Therefore, they can be polled from any thread, as long as that thread has the correct event loop
+// active at the time of the call to `poll()`. If the correct event loop is not active,
+// GuardedRustPromiseAwaiter's API will panic. (The Guarded- prefix refers to the C++ class template
+// ExecutorGuarded, which enforces the correct event loop requirement.)
+unsafe impl Send for GuardedRustPromiseAwaiter {}
+
 // =======================================================================================
 // Await syntax for OwnPromiseNode
 
@@ -13,7 +23,7 @@ use crate::OwnPromiseNode;
 pub struct PromiseAwaiter<Data: std::marker::Unpin> {
     node: Option<OwnPromiseNode>,
     pub(crate) data: Data,
-    awaiter: MaybeUninit<RustPromiseAwaiterRepr>,
+    awaiter: MaybeUninit<GuardedRustPromiseAwaiterRepr>,
     awaiter_initialized: bool,
     // Safety: `option_waker` must be declared after `awaiter`, because `awaiter` contains a reference
     // to `option_waker`. This ensures `option_waker` will be dropped after `awaiter`.
@@ -35,19 +45,19 @@ impl<Data: std::marker::Unpin> PromiseAwaiter<Data> {
     ///
     /// Panics if `node` is None.
     #[must_use]
-    pub fn get_awaiter(mut self: Pin<&mut Self>) -> Pin<&mut RustPromiseAwaiter> {
+    pub fn get_awaiter(mut self: Pin<&mut Self>) -> Pin<&mut GuardedRustPromiseAwaiter> {
         // Safety: We never move out of `this`.
         let this = unsafe { Pin::into_inner_unchecked(self.as_mut()) };
 
         // Initialize the awaiter if not already done
         if !this.awaiter_initialized {
             // On our first invocation, `node` will be Some, and `get_awaiter` will forward its
-            // contents into RustPromiseAwaiter's constructor. On all subsequent invocations, `node`
+            // contents into GuardedRustPromiseAwaiter's constructor. On all subsequent invocations, `node`
             // will be None and the constructor will not run.
             let node = this.node.take();
 
             // Safety: `awaiter` stores `rust_waker_ptr` and uses it to call `wake()`. Note that
-            // `awaiter` is `this.awaiter`, which lives before `this.option_waker`.
+            // `awaiter` is `this.awaiter`, which lives before `this.option_waker`. 
             // Since we drop awaiter manually, the `rust_waker_ptr` that `awaiter` stores will always
             // be valid during its lifetime.
             //
@@ -58,8 +68,8 @@ impl<Data: std::marker::Unpin> PromiseAwaiter<Data> {
 
             // Safety: The memory slot is valid and this type ensures that it will stay pinned.
             unsafe {
-                crate::ffi::rust_promise_awaiter_new_in_place(
-                    this.awaiter.as_mut_ptr().cast::<RustPromiseAwaiter>(),
+                crate::ffi::guarded_rust_promise_awaiter_new_in_place(
+                    this.awaiter.as_mut_ptr().cast::<GuardedRustPromiseAwaiter>(),
                     rust_waker_ptr,
                     node.expect("node should be Some in call to init()"),
                 );
@@ -69,8 +79,8 @@ impl<Data: std::marker::Unpin> PromiseAwaiter<Data> {
 
         // Safety: `this.awaiter` is pinned since `self` is pinned.
         unsafe {
-            let raw = this.awaiter.assume_init_mut() as *mut RustPromiseAwaiterRepr;
-            let raw = raw.cast::<RustPromiseAwaiter>();
+            let raw = this.awaiter.assume_init_mut() as *mut GuardedRustPromiseAwaiterRepr;
+            let raw = raw.cast::<GuardedRustPromiseAwaiter>();
             Pin::new_unchecked(&mut *raw)
         }
     }
@@ -87,8 +97,8 @@ impl<Data: std::marker::Unpin> Drop for PromiseAwaiter<Data> {
     fn drop(&mut self) {
         if self.awaiter_initialized {
             unsafe {
-                crate::ffi::rust_promise_awaiter_drop_in_place(
-                    self.awaiter.as_mut_ptr().cast::<RustPromiseAwaiter>(),
+                crate::ffi::guarded_rust_promise_awaiter_drop_in_place(
+                    self.awaiter.as_mut_ptr().cast::<GuardedRustPromiseAwaiter>()
                 );
             }
         }

kj-rs/future.rs

@@ -38,19 +38,19 @@ pub(crate) mod repr {
 
     type DropCallback = unsafe extern "C" fn(fut: *mut c_void);
 
-    type FuturePtr<'a, T> = *mut (dyn Future<Output = Result<T, String>> + 'a);
+    type FuturePtr<'a, T> = *mut (dyn Future<Output = Result<T, String>> + Send + 'a);
 
-    /// Represents a `dyn Future<Output = Result<T, String>>`.
+    /// Represents a `dyn Future<Output = Result<T, String>>` + Send.
     #[repr(C)]
     pub struct RustFuture<'a, T> {
         pub fut: FuturePtr<'a, T>,
         pub poll: PollCallback,
         pub drop: DropCallback,
     }
 
-    type InfallibleFuturePtr<'a, T> = *mut (dyn Future<Output = T> + 'a);
+    type InfallibleFuturePtr<'a, T> = *mut (dyn Future<Output = T> + Send + 'a);
 
-    /// Represents a `dyn Future<Output = T>` where T is not a Result.
+    /// Represents a `dyn Future<Output = T> + Send` where T is not a Result.
     #[repr(C)]
     pub struct RustInfallibleFuture<'a, T> {
         pub fut: InfallibleFuturePtr<'a, T>,
@@ -127,7 +127,7 @@ pub(crate) mod repr {
 
     #[must_use]
     pub fn future<'a, T: Unpin>(
-        fut: Pin<Box<dyn Future<Output = Result<T, String>> + 'a>>,
+        fut: Pin<Box<dyn Future<Output = Result<T, String>> + Send + 'a>>,
     ) -> RustFuture<'a, T> {
         let fut = Box::into_raw(unsafe { Pin::into_inner_unchecked(fut) });
         let poll = RustFuture::<T>::poll;
@@ -137,7 +137,7 @@ pub(crate) mod repr {
 
     #[must_use]
     pub fn infallible_future<'a, T: Unpin>(
-        fut: Pin<Box<dyn Future<Output = T> + 'a>>,
+        fut: Pin<Box<dyn Future<Output = T> + Send + 'a>>,
     ) -> RustInfallibleFuture<'a, T> {
         let fut = Box::into_raw(unsafe { Pin::into_inner_unchecked(fut) });
         let poll = RustInfallibleFuture::<T>::poll;

kj-rs/lib.rs

@@ -26,10 +26,10 @@ pub type Error = std::io::Error;
 #[allow(clippy::needless_lifetimes)]
 mod ffi {
 
-    /// Representation of a `RustPromiseAwaiter` in C++. The size of the blob should match.
+    /// Representation of a `GuardedRustPromiseAwaiter` in C++. The size of the blob should match.
     #[derive(Debug)]
-    pub struct RustPromiseAwaiterRepr {
-        _bindgen_opaque_blob: [u64; 14usize],
+    pub struct GuardedRustPromiseAwaiterRepr {
+        _bindgen_opaque_blob: [u64; 15usize],
     }
 
     extern "Rust" {
@@ -73,22 +73,22 @@ mod ffi {
     unsafe extern "C++" {
         include!("kj-rs/awaiter.h");
 
-        type RustPromiseAwaiter;
+        type GuardedRustPromiseAwaiter;
 
-        unsafe fn rust_promise_awaiter_new_in_place(
-            ptr: *mut RustPromiseAwaiter,
+        unsafe fn guarded_rust_promise_awaiter_new_in_place(
+            ptr: *mut GuardedRustPromiseAwaiter,
             rust_waker_ptr: *mut OptionWaker,
             node: OwnPromiseNode,
         );
-        unsafe fn rust_promise_awaiter_drop_in_place(ptr: *mut RustPromiseAwaiter);
+        unsafe fn guarded_rust_promise_awaiter_drop_in_place(ptr: *mut GuardedRustPromiseAwaiter);
 
         unsafe fn poll(
-            self: Pin<&mut RustPromiseAwaiter>,
+            self: Pin<&mut GuardedRustPromiseAwaiter>,
             waker: &WakerRef,
             maybe_kj_waker: *const KjWaker,
         ) -> bool;
 
         #[must_use]
-        fn take_own_promise_node(self: Pin<&mut RustPromiseAwaiter>) -> OwnPromiseNode;
+        fn take_own_promise_node(self: Pin<&mut GuardedRustPromiseAwaiter>) -> OwnPromiseNode;
     }
 }

kj-rs/promise.rs

@@ -19,6 +19,15 @@ type CxxResult<T> = std::result::Result<T, cxx::Exception>;
 #[repr(transparent)]
 pub struct OwnPromiseNode(*mut c_void /* kj::_::PromiseNode* */);
 
+// Safety: KJ Promises are not associated with threads, but with event loops at construction time.
+// Therefore, they can be polled from any thread, as long as that thread has the correct event loop
+// active at the time of the call to `poll()`. If the correct event loop is not active, the
+// OwnPromiseNode's API will typically panic, undefined behavior could be possible. However, Rust
+// doesn't have direct access to OwnPromiseNode's API. Instead, it can only use the Promise by
+// having GuardedRustPromiseAwaiter consume it, and GuardedRustPromiseAwaiter implements the
+// correct-executor guarantee.
+unsafe impl Send for OwnPromiseNode {}
+
 // Note: drop is not the only way for OwnPromiseNode to be destroyed.
 // It is forgotten using `MaybeUninit` and its ownership passed over to c++ in `unwrap`.
 impl Drop for OwnPromiseNode {
@@ -142,3 +151,5 @@ impl<T> KjPromise for CallbacksFuture<T> {
         Ok(unsafe { ret.assume_init() })
     }
 }
+
+unsafe impl<T: Send> Send for CallbacksFuture<T> {}

kj-rs/tests/test_futures.rs

@@ -149,7 +149,10 @@ pub async fn new_layered_ready_future_void() -> Result<()> {
 }
 
 // From example at https://doc.rust-lang.org/std/future/fn.poll_fn.html#capturing-a-pinned-state
-async fn naive_select<T>(a: impl Future<Output = T>, b: impl Future<Output = T>) -> T {
+async fn naive_select<T>(
+    a: impl Future<Output = T> + Send,
+    b: impl Future<Output = T> + Send,
+) -> T {
     let (mut a, mut b) = (pin!(a), pin!(b));
     future::poll_fn(move |cx| {
         if let Poll::Ready(r) = a.as_mut().poll(cx) {

macro/src/expand.rs

@@ -1268,9 +1268,9 @@ fn expand_rust_function_shim_super(
         };
 
         if fut.throws_tokens.is_some() {
-            quote!(-> std::pin::Pin<Box<dyn ::std::future::Future<Output = ::std::result::Result<#output, String>> #lifetimes>>)
+            quote!(-> std::pin::Pin<Box<dyn ::std::future::Future<Output = ::std::result::Result<#output, String>> + Send #lifetimes>>)
         } else {
-            quote!(-> std::pin::Pin<Box<dyn ::std::future::Future<Output = #output> #lifetimes>>)
+            quote!(-> std::pin::Pin<Box<dyn ::std::future::Future<Output = #output> + Send #lifetimes>>)
         }
     } else {
         expand_return_type(&sig.ret)