From 0d01ff51eb259d7257bfabef5a0d2284cafda27a Mon Sep 17 00:00:00 2001 From: Dan Lapid Date: Fri, 14 Aug 2026 12:44:38 -0400 Subject: [PATCH] kj-rs: rework the async bridge waker/event machinery Groundwork for the tokio-backed Rust I/O backend (kj-rs-tokio / kj-rs-io, landing separately); kj-rs itself stays a pure Promise<->Future bridge. - Replace KjWaker with FutureWakerCell: every cloned waker is a same-thread cell (non-atomic kj::Refcounted; the bridge's single-thread axiom) that arms the owning FuturePollEvent directly via Event::armDepthFirst(). The cell's link to the event is weak and structurally invalidated when the event dies, so wakers Rust retains past the future's lifetime (e.g. parked in a channel's AtomicWaker) neutralize into safe no-ops instead of arming a freed event. Waker ownership round-trips through RawWaker data slots via kj::Rc::disown()/reown() -- hence the capnp-cpp pin bump to the current v2 head, which carries those (merged upstream). - Replace the LinkedGroup machinery (linked-group.h + test) with an intrusive weak link (RustPromiseAwaiter::link / FuturePollEvent::leaves). - Make bridged kj::Promises eager by default: the Rust future is polled to its first suspension point at conversion, so KJ callers no longer need a manual .eagerlyEvaluate(nullptr); RustFuture::lazily() is the escape hatch for the rare cold case. - Convert panics escaping a bridged future's poll into kj::Exceptions (a rejected promise) instead of aborting the process, mirroring the sync bridge's catch_unwind path. - Add a thread-local armed-waker hook so an integrating kj::EventPort that drives tokio tasks inside its own wait() (the upcoming kj-rs-tokio) can nudge itself out of a blocking park; null/no-op by default. - Split the cxx bridge module out of lib.rs into ffi.rs, and quarantine unsafe into named FFI islands: deny(unsafe_code) crate-wide, re-allowed per-module only where the FFI seam genuinely needs it. - Depend on @capnp-cpp//src/kj:kj-async-core instead of the :kj-async umbrella, keeping kj-rs (and everything built on it) off the concrete kj OS event loop. - Qualify bare uint as kj::uint in src/rust/kj/tests/ffi-test.c++: these dependency changes make that target newly compile on Windows CI, where no global uint exists (POSIX gets one from sys/types.h). - New tests: waker neutralization across event death (neutralize-waker-test), FuturePollEvent shared-event semantics (shared-event-test), and expanded future/awaiter coverage. --- src/rust/cxx/AGENTS.md | 15 + src/rust/cxx/kj-rs/BUILD.bazel | 13 +- src/rust/cxx/kj-rs/awaiter.c++ | 205 +++++----- src/rust/cxx/kj-rs/awaiter.h | 156 ++++---- src/rust/cxx/kj-rs/awaiter.rs | 47 ++- src/rust/cxx/kj-rs/executor-guarded.c++ | 25 +- src/rust/cxx/kj-rs/executor-guarded.h | 30 +- src/rust/cxx/kj-rs/ffi.rs | 151 ++++++++ src/rust/cxx/kj-rs/future.h | 30 +- src/rust/cxx/kj-rs/future.rs | 180 +++++++-- src/rust/cxx/kj-rs/lib.rs | 118 ++---- src/rust/cxx/kj-rs/linked-group.h | 315 --------------- src/rust/cxx/kj-rs/maybe.rs | 60 ++- src/rust/cxx/kj-rs/own.rs | 39 +- src/rust/cxx/kj-rs/promise.c++ | 4 +- src/rust/cxx/kj-rs/promise.h | 2 +- src/rust/cxx/kj-rs/promise.rs | 52 ++- src/rust/cxx/kj-rs/refcount.rs | 46 ++- src/rust/cxx/kj-rs/tests/BUILD.bazel | 26 +- .../cxx/kj-rs/tests/awaitables-cc-test.c++ | 190 ++++++++- src/rust/cxx/kj-rs/tests/lib.rs | 82 +++- .../cxx/kj-rs/tests/linked-group-test.c++ | 364 ------------------ .../cxx/kj-rs/tests/neutralize-waker-test.c++ | 144 +++++++ .../cxx/kj-rs/tests/shared-event-test.c++ | 321 +++++++++++++++ src/rust/cxx/kj-rs/tests/test_futures.rs | 268 +++++++++++-- src/rust/cxx/kj-rs/tests/test_own.rs | 63 +-- src/rust/cxx/kj-rs/waker.c++ | 163 ++------ src/rust/cxx/kj-rs/waker.h | 300 ++++++++------- src/rust/cxx/kj-rs/waker.rs | 216 +++++++---- src/rust/kj/tests/ffi-test.c++ | 5 +- 30 files changed, 2152 insertions(+), 1478 deletions(-) create mode 100644 src/rust/cxx/kj-rs/ffi.rs delete mode 100644 src/rust/cxx/kj-rs/linked-group.h delete mode 100644 src/rust/cxx/kj-rs/tests/linked-group-test.c++ create mode 100644 src/rust/cxx/kj-rs/tests/neutralize-waker-test.c++ create mode 100644 src/rust/cxx/kj-rs/tests/shared-event-test.c++ diff --git a/src/rust/cxx/AGENTS.md b/src/rust/cxx/AGENTS.md index bb8fb9d3ca2..83b9c30e5b1 100644 --- a/src/rust/cxx/AGENTS.md +++ b/src/rust/cxx/AGENTS.md @@ -36,6 +36,21 @@ Bazel module, Cargo workspace, toolchain configuration, or external `workerd-cxx - `tests/` and `kj-rs/tests/` — Rust and C++ bridge integration tests - `tools/bazel/` — Bazel bridge-generation macro used by this component's tests +## Async bridge semantics + +- Marking a fn `async` in `extern "Rust"` yields a `kj::Promise` in C++; `async` in + `extern "C++"` yields an `impl Future` in Rust. +- Bridged `kj::Promise`s are **eager by default**: the Rust future is polled to its first + suspension point at the call (KJ code assumes hot promises), so callers never need + `.eagerlyEvaluate(nullptr)`. `RustFuture::lazily()` (kj-rs/future.h) is the C++-side + escape hatch for the rare cold case. +- The waker bridge honors `Waker: Send + Sync` for real: a Rust `.await` of a KJ promise links + to the `FuturePollEvent` via an intrusive weak link (`RustPromiseAwaiter::link` / + `FuturePollEvent::leaves`), and a cloned waker is an atomically-refcounted `FutureWakerCell` + whose wake checks the owning executor — on the owning loop's thread it arms the + `FuturePollEvent` directly (the hot path), from any other thread it delivers through a + cross-thread fulfiller that the event renews each poll. + ## Conventions - Follow the parent `src/rust/AGENTS.md` and repository `AGENTS.md`. diff --git a/src/rust/cxx/kj-rs/BUILD.bazel b/src/rust/cxx/kj-rs/BUILD.bazel index 971b7e12512..a80c2a575bb 100644 --- a/src/rust/cxx/kj-rs/BUILD.bazel +++ b/src/rust/cxx/kj-rs/BUILD.bazel @@ -11,7 +11,10 @@ wd_cc_library( "@platforms//os:windows": True, "//conditions:default": False, }), - visibility = ["//src/rust/cxx/tests:__pkg__"], + visibility = [ + "//src/rust/cxx/kj-rs-tokio:__pkg__", + "//src/rust/cxx/tests:__pkg__", + ], deps = [ ":bridge", ], @@ -54,13 +57,17 @@ rust_test( rust_cxx_bridge( name = "bridge", - src = "lib.rs", + src = "ffi.rs", hdrs = glob(["*.h"]), include_prefix = "kj-rs", visibility = ["//src/rust/cxx/tests:__pkg__"], deps = [ "//src/rust/cxx:core", "@capnp-cpp//src/kj:kj", - "@capnp-cpp//src/kj:kj-async", + # kj-rs is the base cxx<->rust Promise/Future bridge: it uses only the abstract async + # core (kj::Promise / kj::EventLoop via async.h), no kj OS I/O. Depending on + # :kj-async-core (not the :kj-async umbrella) keeps the whole kj-rs stack -- and thus + # kj-rs-io / kj-rs-tokio built on it -- off the concrete kj OS event loop (:kj-async-os). + "@capnp-cpp//src/kj:kj-async-core", ], ) diff --git a/src/rust/cxx/kj-rs/awaiter.c++ b/src/rust/cxx/kj-rs/awaiter.c++ index 7bd3aa0fafb..44f1107d463 100644 --- a/src/rust/cxx/kj-rs/awaiter.c++ +++ b/src/rust/cxx/kj-rs/awaiter.c++ @@ -1,6 +1,6 @@ #include "awaiter.h" -#include +#include #include @@ -28,13 +28,29 @@ RustPromiseAwaiter::RustPromiseAwaiter( } RustPromiseAwaiter::~RustPromiseAwaiter() noexcept(false) { - // Our `tracePromise()` implementation checks for a null `node`, so we don't have to sever our - // LinkedGroup relationship before destroying `node`. If our FuturePollEvent (our LinkedGroup) - // tries to trace us between now and our destructor completing, `tracePromise()` will ignore the - // null `node`. + // Sever our weak link to any FuturePollEvent before we go away, so it can't trace into or arm a + // destroyed awaiter. Our `tracePromise()` implementation also checks for a null `node`, so even + // between clearPollEvent() and node reset we are safe to trace. + clearPollEvent(); unwindDetector.catchExceptionsIfUnwinding([this]() { node = nullptr; }); } +void RustPromiseAwaiter::setPollEvent(FuturePollEvent& futurePollEvent) { + KJ_IF_SOME(old, maybePollEvent) { + if (&old == &futurePollEvent) return; + old.leaves.remove(*this); + } + futurePollEvent.leaves.add(*this); + maybePollEvent = futurePollEvent; +} + +void RustPromiseAwaiter::clearPollEvent() { + KJ_IF_SOME(old, maybePollEvent) { + old.leaves.remove(*this); + maybePollEvent = kj::none; + } +} + void RustPromiseAwaiter::fire() { // Safety: Our Event can only fire on the event loop which was active when our Event base class // was constructed. Therefore, we don't need to check that we're on the correct event loop. @@ -42,10 +58,10 @@ void RustPromiseAwaiter::fire() { // Nullify our `maybeOptionWaker` to signal that we are done. KJ_DEFER(maybeOptionWaker = kj::none); - KJ_IF_SOME(futurePollEvent, linkedGroup().tryGet()) { + KJ_IF_SOME(futurePollEvent, maybePollEvent) { // Optimized path: we're still linked to a FuturePollEvent. Arm it directly. futurePollEvent.armDepthFirst(); - linkedGroup().set(kj::none); + clearPollEvent(); } else KJ_IF_SOME(optionWaker, maybeOptionWaker) { // We use wake_if_some() rather than an unconditional wake because the OptionWaker may be empty. This // happens when poll() took the optimized path (clearing the OptionWaker and linking to a @@ -68,7 +84,7 @@ void RustPromiseAwaiter::traceEvent(kj::_::TraceBuilder& builder) { node->tracePromise(builder, true); } // TODO(someday): Can we add an entry for the `.await` expression in Rust here? - KJ_IF_SOME(futurePollEvent, linkedGroup().tryGet()) { + KJ_IF_SOME(futurePollEvent, maybePollEvent) { futurePollEvent.traceEvent(builder); } } @@ -82,48 +98,21 @@ void RustPromiseAwaiter::tracePromise(kj::_::TraceBuilder& builder, bool stopAtN // TODO(someday): Can we add an entry for the `.await` expression in Rust here? } -bool RustPromiseAwaiter::poll(const WakerRef& waker, const KjWaker* maybeKjWaker) { +bool RustPromiseAwaiter::poll(const WakerRef& waker) { // TODO(perf): If `this->isNext()` is true, meaning our event is next in line to fire, can we // disarm it, set `done = true`, etc.? If we can only suspend if our enclosing KJ coroutine has - // suspended at least once, we may be able to check for that through LazyArcWaker, but this path + // suspended at least once, we may be able to check for that through PollWaker, but this path // doesn't have access to one. KJ_IF_SOME(optionWaker, maybeOptionWaker) { // Our Promise is not yet ready. - // Check for an optimized wake path. - KJ_IF_SOME(kjWaker, maybeKjWaker) { - KJ_IF_SOME(futurePollEvent, kjWaker.tryGetFuturePollEvent()) { - // Optimized path. The Future which is polling our Promise is in turn being polled by a - // `co_await` expression somewhere up the stack from us. We can arrange to arm the - // `co_await` expression's KJ Event directly when our Promise is ready. - - // Drop any Waker stored in OptionWaker. We'll use the LinkedGroup to wake instead. - // - // Note: this leaves OptionWaker empty while maybeOptionWaker is still Some(ref). If the - // FuturePollEvent is later destroyed (severing the LinkedGroup link) before our Promise - // fires, fire() will find no LinkedGroup AND an empty OptionWaker. fire() handles this - // via wake_if_some(), which is a no-op on an empty OptionWaker. - optionWaker.set_none(); - - // Store a reference to the current `co_await` expression's Future polling Event. The - // reference is weak, and will be cleared if the `co_await` expression happens to end before - // our Promise is ready. In the more likely case that our Promise becomes ready while the - // `co_await` expression is still active, we'll arm its Event so it can `poll()` us again. - linkedGroup().set(futurePollEvent); - - return false; - } - } - - // Unoptimized fallback path. - // Tell our OptionWaker to store a clone of whatever Waker we were given. optionWaker.set(waker); - // Clearing our reference to the FuturePollEvent (if we have one) tells our fire() + // Clearing our weak reference to the FuturePollEvent (if we have one) tells our fire() // implementation to use our OptionWaker to perform the wake. - linkedGroup().set(kj::none); + clearPollEvent(); return false; } else { @@ -132,6 +121,40 @@ bool RustPromiseAwaiter::poll(const WakerRef& waker, const KjWaker* maybeKjWaker } } +bool RustPromiseAwaiter::poll(const WakerRef& waker, const PollWaker& pollWaker) { + KJ_IF_SOME(futurePollEvent, pollWaker.tryGetFuturePollEvent()) { + KJ_IF_SOME(optionWaker, maybeOptionWaker) { + // Our Promise is not yet ready, and we have an optimized wake path. The Future which is + // polling our Promise is in turn being polled by a `co_await` expression somewhere up the + // stack from us. We can arrange to arm the `co_await` expression's KJ Event directly when + // our Promise is ready. + + // Drop any Waker stored in OptionWaker. We'll use our weak link to the FuturePollEvent to + // wake instead. + // + // Note: this leaves OptionWaker empty while maybeOptionWaker is still Some(ref). If the + // FuturePollEvent is later destroyed (severing our weak link) before our Promise fires, + // fire() will find no linked FuturePollEvent AND an empty OptionWaker. fire() handles this + // via wake_if_some(), which is a no-op on an empty OptionWaker. + optionWaker.set_none(); + + // Store a weak reference to the current `co_await` expression's Future polling Event. The + // reference is weak, and will be cleared if the `co_await` expression happens to end before + // our Promise is ready. In the more likely case that our Promise becomes ready while the + // `co_await` expression is still active, we'll arm its Event so it can `poll()` us again. + setPollEvent(futurePollEvent); + + return false; + } else { + // Our Promise is ready. + return true; + } + } + // The PollWaker exposes no FuturePollEvent (its owning thread's kj::Executor is not ours -- + // cannot normally happen in the single-thread world). Fall back to the generic path. + return poll(waker); +} + OwnPromiseNode RustPromiseAwaiter::take_own_promise_node() { KJ_ASSERT(maybeOptionWaker == kj::none, "take_own_promise_node() should only be called after poll() " @@ -151,43 +174,56 @@ void guarded_rust_promise_awaiter_drop_in_place(GuardedRustPromiseAwaiter* ptr) // ======================================================================================= // FuturePollEvent -void FuturePollEvent::exitPollScope(kj::Maybe> maybePromise) { - // Await any LazyArcWaker promise that got created during the call to `poll()`. Note that if a - // Future returns Ready _and_ synchronously wakes its Waker, the work done to await the - // LazyArcWaker promise is wasted, since we will immediately tear the entire BoxFutureAwaiter - // down. However, that's an unlikely case, and this work here isn't likely to be a significant - // source of overhead. - KJ_IF_SOME(promise, maybePromise) { - auto& node = arcWakerPromise.emplace(kj::_::PromiseNode::from(kj::mv(promise))); - node->setSelfPointer(&node); - node->onReady(this); - } +FuturePollEvent::FuturePollEvent(kj::SourceLocation location) + : Event(location), + // Created eagerly (not lazily on first clone): `&Waker` is Sync, so even the borrowed + // per-poll waker may be cloned or woken from a foreign thread during the very first poll, + // and both paths need the cell to already exist. One small allocation per awaited future, + // next to the coroutine frame and promise nodes already being allocated. + wakerCell{kj::arc(static_cast(*this))} { + ensureCrossThreadWakeArmed(); } -void FuturePollEvent::enterPollScope() noexcept { - // Clear out any previous LazyArcWaker promise the FuturePollEvent was holding onto. Note that - // since there is no code path which rejects this Promise, this is not strictly required for - // correctness, but nevertheless serves as a useful assertion. - KJ_IF_SOME(node, arcWakerPromise) { - kj::_::ExceptionOr output; +FuturePollEvent::~FuturePollEvent() noexcept(false) { + // Our FutureWakerCell is neutralized by the wakerCell guard's destructor during member + // destruction, so any waker reference Rust retained past our lifetime observes a dead weak link + // on a later same-thread wake and is a safe no-op, rather than arming this freed event. The + // cross-thread path neutralizes independently: `crossThreadWakePromise` dies with us, making + // fulfills of the cell's retained fulfiller no-ops. + + // Sever our weak links to all leaves, so a RustPromiseAwaiter that outlives us (e.g. a stashed + // PromiseFuture) never arms this freed event. + for (;;) { + auto it = leaves.begin(); + if (it == leaves.end()) break; + auto& leaf = *it; + leaves.remove(leaf); + leaf.maybePollEvent = kj::none; + } +} - node->get(output); - KJ_IF_SOME(exception, kj::runCatchingExceptions([this]() { arcWakerPromise = kj::none; })) { - output.addException(kj::mv(exception)); +void FuturePollEvent::ensureCrossThreadWakeArmed() { + if (!wakerCell.cell->needsFreshCrossThreadFulfiller()) { + return; + } + auto paf = kj::newPromiseAndCrossThreadFulfiller(); + crossThreadWakePromise = paf.promise + .then([this]() { + // A foreign thread woke the cell: arm the event from the owning thread, exactly as a + // same-thread wake would. `this` is safe: the promise chain running this continuation is + // owned by the event itself. No renewal here (see ensureCrossThreadWakeArmed's declaration); + // further cross-thread wakes before the next poll coalesce into the pending poll. + armDepthFirst(); + if (futurePollArmNudge != nullptr) { + futurePollArmNudge(); } + }).eagerlyEvaluate(nullptr); + wakerCell.cell->replaceCrossThreadFulfiller(kj::mv(paf.fulfiller)); +} - // NOTE: `node` is now dangling. - - KJ_IF_SOME(exception, output.exception) { - // We should only ever receive a WakeInstruction, never an exception. If we do receive an - // exception, it would be because our ArcWaker implementation allowed its cross-thread promise - // fulfiller to be destroyed without being fulfilled, or because we foolishly added an - // explicit call to the fulfiller's reject() function. Either way, it is a programming error, - // so we abort the process here by re-throwing across a noexcept boundary. This avoids having - // implement the ability to "reject" the Future poll() Event. - kj::throwFatalException(kj::mv(exception)); - } - } +kj::Arc FuturePollEvent::cloneWakerCell() { + // Hand out a new strong reference for Rust to retain. + return wakerCell.cell.addRef(); } void FuturePollEvent::tracePromise(kj::_::TraceBuilder& builder, bool stopAtNextEvent) { @@ -199,32 +235,9 @@ void FuturePollEvent::tracePromise(kj::_::TraceBuilder& builder, bool stopAtNext // When tracing, we can only pick one branch to follow. Arbitrarily, I'm following the first // RustPromiseAwaiter branch, similar to how ExclusiveJoinPromiseNode chooses its left branch. In // the common case, this will be whatever OwnPromiseNode our Rust Future is currently `.await`ing. - auto rustPromiseAwaiters = linkedObjects(); - if (rustPromiseAwaiters.begin() != rustPromiseAwaiters.end()) { + if (!leaves.empty()) { // Our Rust Future is awaiting an OwnPromiseNode. We'll pick the first one in our list. - rustPromiseAwaiters.front().tracePromise(builder, false); - } else KJ_IF_SOME(node, arcWakerPromise) { - // Our Rust Future is not awaiting any OwnPromiseNode, and instead cloned our Waker. We'll trace - // our ArcWaker Promise instead. - if (node.get() != nullptr) { - node->tracePromise(builder, false); - } - } -} - -FuturePollEvent::PollScope::PollScope(FuturePollEvent& futurePollEvent): holder(futurePollEvent) { - futurePollEvent.enterPollScope(); -} - -FuturePollEvent::PollScope::~PollScope() noexcept(false) { - holder.get().futurePollEvent.exitPollScope(reset()); -} - -kj::Maybe FuturePollEvent::PollScope::tryGetFuturePollEvent() const { - KJ_IF_SOME(h, holder.tryGet()) { - return h.futurePollEvent; - } else { - return kj::none; + leaves.front().tracePromise(builder, false); } } diff --git a/src/rust/cxx/kj-rs/awaiter.h b/src/rust/cxx/kj-rs/awaiter.h index ec707c2753a..18d32d60818 100644 --- a/src/rust/cxx/kj-rs/awaiter.h +++ b/src/rust/cxx/kj-rs/awaiter.h @@ -1,19 +1,20 @@ #pragma once #include "kj-rs/executor-guarded.h" -#include "kj-rs/linked-group.h" +#include "kj-rs/promise.h" #include "kj-rs/waker.h" #include +#include namespace kj_rs { // ======================================================================================= // Opaque Rust types // -// The following types are defined in lib.rs, and thus in lib.rs.h. lib.rs.h depends on our C++ -// headers, including awaiter.h (the file you're currently reading), so we forward-declare some types -// here for use in the C++ headers. +// The following types are defined in the cxx bridge (ffi.rs), and thus in ffi.rs.h. ffi.rs.h +// depends on our C++ headers, including awaiter.h (the file you're currently reading), so we +// forward-declare some types here for use in the C++ headers. // Wrapper around an `&std::task::Waker`, passed to `RustPromiseAwaiter::poll()`. This indirection // is required because cxx-rs does not permit us to expose opaque Rust types to C++ defined outside @@ -46,17 +47,17 @@ struct OptionWaker; // alignment using bindgen. See inside awaiter.c++ for a static_assert to remind us to re-run // bindgen. // -// RustPromiseAwaiter has two base classes: KJ Event, and a LinkedObject template instantiation. We -// use the Event to discover when our wrapped Promise is ready. Our Event fire() implementation -// records the fact that we are done, then wakes our Waker or arms the FuturePollEvent, if we -// have one. We access the FuturePollEvent via our LinkedObject base class mixin. It gives us the -// ability to store a weak reference to the FuturePollEvent, if we were last polled by one. +// RustPromiseAwaiter has one base class: KJ Event. We use the Event, via the native +// `node->onReady(this)` mechanism, to discover when our wrapped Promise is ready (exactly as a KJ +// coroutine registers its own Event on the promise it `co_await`s). Our Event fire() implementation +// records the fact that we are done, then wakes our Waker or arms the FuturePollEvent, if we have +// one. We hold a weak reference to the FuturePollEvent that last polled us (see maybePollEvent +// below), so a fired Promise can arm that poll event directly. // // Cancellation: Dropping the RustPromiseAwaiter destroys its OwnPromiseNode, cancelling the // wrapped KJ promise. If the RustPromiseAwaiter was never constructed, Rust's OwnPromiseNode::drop() // cancels the promise directly. -class RustPromiseAwaiter final: public kj::_::Event, - public LinkedObject { +class RustPromiseAwaiter final: public kj::_::Event { public: // The Rust code which constructs RustPromiseAwaiter passes us a pointer to a OptionWaker, which can // be thought of as a Rust-native component RustPromiseAwaiter. Its job is to hold a clone of @@ -84,18 +85,29 @@ class RustPromiseAwaiter final: public kj::_::Event, // Poll this Promise for readiness. // - // If the Waker is a KjWaker, you may pass the KjWaker pointer as a second parameter. This may - // allow the implementation of `poll()` to optimize the wake by arming a KJ Event directly when - // the wrapped Promise becomes ready. - // - // If the Waker is not a KjWaker, the `maybeKjWaker` pointer argument must be nullptr. - bool poll(const WakerRef& waker, const KjWaker* maybeKjWaker); + // The two-argument overload is for polls driven by a PollWaker (i.e. a `co_await`ed Future's + // poll): it may optimize the wake by arming a KJ Event directly when the wrapped Promise + // becomes ready. Polls driven by any other Waker use the one-argument overload. + bool poll(const WakerRef& waker); + bool poll(const WakerRef& waker, const PollWaker& pollWaker); // Release ownership of the OwnPromiseNode. Asserts if called before the Promise is ready; that // is, `poll()` must have returned true prior to calling `take_own_promise_node()`. OwnPromiseNode take_own_promise_node(); private: + // Purpose-built one-to-many weak link to the FuturePollEvent that last polled us. The link is + // weak in both directions: destroying either side severs it (see + // clearPollEvent() and ~FuturePollEvent()), so a fired Promise never arms, and tracing never + // touches, a destroyed FuturePollEvent. A RustPromiseAwaiter may outlive the FuturePollEvent that + // first polled it (e.g. a stashed PromiseFuture) and later re-link to a different one. + friend class FuturePollEvent; + void setPollEvent(FuturePollEvent& futurePollEvent); + void clearPollEvent(); + + kj::Maybe maybePollEvent; + kj::ListLink link; + // The Rust code which instantiates RustPromiseAwaiter does so with a OptionWaker object right // next to the RustPromiseAwaiter, such that it is dropped after RustPromiseAwaiter. Thus, our // reference to our OptionWaker is stable. We use the OptionWaker to (optionally) store a clone of @@ -117,8 +129,11 @@ struct GuardedRustPromiseAwaiter: ExecutorGuarded { // We need to inherit constructors or else placement-new will try to aggregate-initialize us. using ExecutorGuarded::ExecutorGuarded; - bool poll(const WakerRef& waker, const KjWaker* maybeKjWaker) { - return get().poll(waker, maybeKjWaker); + bool poll(const WakerRef& waker) { + return get().poll(waker); + } + bool pollWithPollWaker(const WakerRef& waker, const PollWaker& pollWaker) { + return get().poll(waker, pollWaker); } OwnPromiseNode take_own_promise_node() { return get().take_own_promise_node(); @@ -136,20 +151,20 @@ void guarded_rust_promise_awaiter_drop_in_place(GuardedRustPromiseAwaiter*); // `Event::fire()` override which actually polls the Future; this class implements all other base // class virtual functions. // -// A FuturePollEvent contains an optional ArcWakerPromiseAwaiter and a list of zero or more -// RustPromiseAwaiters. These "sub-Promise awaiters" all wrap a KJ Promise of some sort, and arrange -// to arm the FuturePollEvent when their Promises become ready. +// A FuturePollEvent owns an optional FutureWakerCell (handed out to Rust by PollWaker::cloneCell()) +// and a list of zero or more RustPromiseAwaiters. These "sub-Promise awaiters" all wrap a KJ +// Promise of some sort, and arrange to arm the FuturePollEvent when their Promises become ready; a +// woken FutureWakerCell arms it the same way. // // The PromiseNode base class is a hack to implement async tracing. That is, we only implement the // `tracePromise()` function, and decide which Promise to trace into if/when the coroutine calls our // `tracePromise()` implementation. This primarily makes the lifetimes easier to manage: our -// RustPromiseAwaiter LinkedObjects have independent lifetimes from the FuturePollEvent, so we -// mustn't leave references to them, or their members, lying around in the Coroutine class. -class FuturePollEvent: public kj::_::PromiseNode, - public kj::_::Event, - public LinkedGroup { +// weakly-linked RustPromiseAwaiter leaves have independent lifetimes from the FuturePollEvent, so +// we mustn't leave references to them, or their members, lying around in the Coroutine class. +class FuturePollEvent: public kj::_::PromiseNode, public kj::_::Event { public: - FuturePollEvent(kj::SourceLocation location = {}): Event(location) {} + FuturePollEvent(kj::SourceLocation location = {}); + ~FuturePollEvent() noexcept(false); // ------------------------------------------------------- // PromiseNode API @@ -159,42 +174,50 @@ class FuturePollEvent: public kj::_::PromiseNode, void tracePromise(kj::_::TraceBuilder& builder, bool stopAtNextEvent) override; - protected: - // PollScope is a LazyArcWaker which is associated with a specific FuturePollEvent, allowing - // optimized Promise `.await`s. Additionally, PollScope's destructor arranges to await any - // ArcWaker promise which was lazily created. - // - // Used by FutureAwaiter, our derived class. - class PollScope; - private: - // Private API for PollScope. - void enterPollScope() noexcept; - void exitPollScope(kj::Maybe> maybeLazyArcWakerPromise); - - kj::Maybe arcWakerPromise; -}; - -class FuturePollEvent::PollScope: public LazyArcWaker { - public: - // `futurePollEvent` is the FuturePollEvent responsible for calling `Future::poll()`, and must - // outlive this PollScope. - PollScope(FuturePollEvent& futurePollEvent); - ~PollScope() noexcept(false); - KJ_DISALLOW_COPY_AND_MOVE(PollScope); - - // The Event which is using this PollScope to poll() a Future. Waking this FuturePollEvent's - // PollScope arms this Event (possibly via a cross-thread promise fulfiller). We also arm the - // Event directly in the RustPromiseAwaiter class, to more optimally `.await` KJ Promises from - // within Rust. If the current thread's kj::Executor is not the same as the one which owns the - // FuturePollEvent, this function returns kj::none. - kj::Maybe tryGetFuturePollEvent() const override; + // Hand out a new strong reference to this event's FutureWakerCell. Used by + // PollWaker::cloneCell(). The cell is created eagerly with this event (so even the borrowed + // per-poll waker can be cloned or woken from any thread) and lives until both this event and + // every Rust reference are gone; ~FuturePollEvent neutralizes it so late wakes no-op. + friend class PollWaker; + kj::Arc cloneWakerCell(); + const FutureWakerCell& wakerCellRef() { + return *wakerCell.cell; + } - private: - struct FuturePollEventHolder { - FuturePollEvent& futurePollEvent; + // (Re-)install a fresh cross-thread fulfiller in the cell if the previous one was consumed by + // a foreign-thread wake (or never existed). The promise side is chained to arm this event. + // Owning thread only; called at construction and at the top of each poll (PollWaker's + // constructor) — never from within the consumed promise's own continuation, which would + // destroy the very chain the continuation is running from. + void ensureCrossThreadWakeArmed(); + + // Weakly-linked list of the RustPromiseAwaiters ("leaves") this Future is currently `.await`ing + // and which may arm this poll event when their Promises become ready. Severed on destruction so + // a leaf that outlives us never arms a freed event. + friend class RustPromiseAwaiter; + kj::List leaves; + + // The FutureWakerCell handed out by cloneWakerCell(). We hold one strong reference through + // this guard, whose destructor neutralizes the cell (nulling its weak Event link so retained + // Rust references become safe no-ops) before releasing it — the invalidation is tied to this + // event's destruction structurally, not by a destructor body remembering to call it. + struct NeutralizeGuard { + kj::Arc cell; + ~NeutralizeGuard() noexcept(false) { + if (cell.get() != nullptr) { + cell->neutralize(); + } + } }; - ExecutorGuarded holder; + NeutralizeGuard wakerCell; + + // The promise side of the cell's cross-thread fulfiller: resolves when a foreign-thread wake + // fulfills it, and its continuation arms this event from the owning thread. Renewed by + // ensureCrossThreadWakeArmed(). Declared after `wakerCell` so it is destroyed first (its + // continuation captures `this`); once it dies with this event, fulfilling the cell's retained + // fulfiller becomes a safe no-op — the cross-thread flavor of neutralization. + kj::Promise crossThreadWakePromise = nullptr; }; // ======================================================================================= @@ -204,7 +227,7 @@ template concept Future = requires(F f) { typename F::Output; { - f.poll(kj::instance(), + f.poll(kj::instance(), kj::instance&>()) } -> std::same_as; }; @@ -256,15 +279,10 @@ class FutureAwaiter final: public FuturePollEvent { void poll() { if (isDone()) return; - // TODO(perf): Check if we already have an ArcWaker from a previous suspension and give it to - // LazyArcWaker for cloning if we have the last reference to it at this point. This could save - // memory allocations, but would depend on making XThreadFulfiller and XThreadPaf resettable - // to really benefit. - { - PollScope pollScope(*this); + PollWaker pollWaker(*this); - future.poll(pollScope, result); + future.poll(pollWaker, result); if (isDone()) { onReadyEvent.arm(); } diff --git a/src/rust/cxx/kj-rs/awaiter.rs b/src/rust/cxx/kj-rs/awaiter.rs index 9c32cc6faf3..706ea604099 100644 --- a/src/rust/cxx/kj-rs/awaiter.rs +++ b/src/rust/cxx/kj-rs/awaiter.rs @@ -1,3 +1,8 @@ +//! FFI island (see crate-root `#![deny(unsafe_code)]`): placement bridge for the C++ +//! `GuardedRustPromiseAwaiter` and the `.await` poll glue — Pin projection and placement +//! new/drop over Rust-owned memory. A genuine unsafe seam. +#![allow(unsafe_code)] + use std::mem::MaybeUninit; use std::pin::Pin; use std::task::Context; @@ -7,7 +12,7 @@ use crate::OwnPromiseNode; // Await syntax for OwnPromiseNode use crate::ffi::GuardedRustPromiseAwaiter; use crate::ffi::GuardedRustPromiseAwaiterRepr; -use crate::waker::try_into_kj_waker_ptr; +use crate::waker::try_poll_waker; pub struct PromiseAwaiter { node: Option, @@ -17,6 +22,16 @@ pub struct PromiseAwaiter { // 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`. option_waker: OptionWaker, + // Suppresses the auto `Unpin` impl (for this type and any wrapper like `PromiseFuture`). + // After the first poll, `awaiter` holds an in-place-constructed C++ object that is (a) a + // `kj::_::Event` registered with the event loop, (b) the target of the promise node's + // self-pointer (`setSelfPointer` points INTO this memory), (c) the holder of a reference to + // the sibling `option_waker` field, and (d) possibly threaded into a `FuturePollEvent`'s + // intrusive `leaves` list. Moving `self` after that (e.g. `let g = f;` after a `&mut f` + // partial await, which `Unpin` would permit in safe code) leaves all four pointers dangling + // -- use-after-free when the promise fires. `PhantomPinned` turns that into a compile error; + // ordinary `.await` pins structurally and is unaffected. + _pinned: std::marker::PhantomPinned, } impl PromiseAwaiter { @@ -27,6 +42,7 @@ impl PromiseAwaiter { awaiter: MaybeUninit::uninit(), awaiter_initialized: false, option_waker: OptionWaker::empty(), + _pinned: std::marker::PhantomPinned, } } @@ -45,6 +61,12 @@ impl PromiseAwaiter { // contents into GuardedRustPromiseAwaiter's constructor. On all subsequent invocations, `node` // will be None and the constructor will not run. let node = this.node.take(); + // `node` is `Some` on this first (initializing) invocation; see the comment above. + #[expect( + clippy::expect_used, + reason = "get_awaiter initializes exactly once while node is Some (awaiter_initialized guards re-entry); None here is an unreachable internal-invariant violation" + )] + let node = node.expect("node should be Some in call to init()"); // Safety: `awaiter` stores `rust_waker_ptr` and uses it to call `wake()`. Note that // `awaiter` is `this.awaiter`, which lives before `this.option_waker`. @@ -64,7 +86,7 @@ impl PromiseAwaiter { .as_mut_ptr() .cast::(), rust_waker_ptr, - node.expect("node should be Some in call to init()"), + node, ); } this.awaiter_initialized = true; @@ -81,20 +103,25 @@ impl PromiseAwaiter { } pub fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> bool { - let maybe_kj_waker = try_into_kj_waker_ptr(cx.waker()); - let awaiter = self.as_mut().get_awaiter(); - // Safety: The awaiter is initialized by `get_awaiter()` above. `WakerRef` borrows the - // context's waker, which is alive for the duration of the call. `maybe_kj_waker` is null - // or points to the KjWaker inside the waker (validated by `try_into_kj_waker_ptr`). - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe { awaiter.poll(&WakerRef(cx.waker()), maybe_kj_waker) } + // If the Waker driving this poll lends out a C++ PollWaker, take the optimized entry + // point, which may arm the enclosing `co_await`'s KJ Event directly. Both borrows live + // off `cx` for the duration of the call. + match try_poll_waker(cx.waker()) { + Some(poll_waker) => self + .as_mut() + .get_awaiter() + .poll_with_poll_waker(&WakerRef(cx.waker()), poll_waker), + None => self.as_mut().get_awaiter().poll(&WakerRef(cx.waker())), + } } } impl Drop for PromiseAwaiter { fn drop(&mut self) { if self.awaiter_initialized { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `awaiter_initialized` is true, so `self.awaiter` holds a + // `GuardedRustPromiseAwaiter` constructed in place by `get_awaiter`; drop it in + // place exactly once here (this is the only drop path, and `self` is being dropped). unsafe { crate::ffi::guarded_rust_promise_awaiter_drop_in_place( self.awaiter diff --git a/src/rust/cxx/kj-rs/executor-guarded.c++ b/src/rust/cxx/kj-rs/executor-guarded.c++ index a026dfef756..71b5cc93494 100644 --- a/src/rust/cxx/kj-rs/executor-guarded.c++ +++ b/src/rust/cxx/kj-rs/executor-guarded.c++ @@ -4,12 +4,35 @@ namespace kj_rs { +namespace { + +const kj::Executor* tryGetCurrentThreadExecutor() { + // kj/async.h exposes no "maybe" variant of getCurrentThreadExecutor() (verified: only the + // KJ_REQUIRE-ing accessor exists, and the thread-local EventLoop pointer it checks is private + // to async.c++), so probe by catching its recoverable "No event loop is running on this + // thread" requirement failure and treating it as "no executor". This runs on teardown paths + // only (see requireCurrentOrTearingDown below), so the exception cost is irrelevant. + const kj::Executor* current = nullptr; + auto maybeException = + kj::runCatchingExceptions([&]() { current = &kj::getCurrentThreadExecutor(); }); + // (maybeException != kj::none) <=> no live EventLoop on this thread; leave `current` null. + (void)maybeException; + return current; +} + +} // namespace + bool isCurrent(const kj::Executor& executor) { - return &executor == &kj::getCurrentThreadExecutor(); + return tryGetCurrentThreadExecutor() == &executor; } void requireCurrent(const kj::Executor& executor, kj::LiteralStringConst message) { KJ_REQUIRE(isCurrent(executor), message); } +void requireCurrentOrTearingDown(const kj::Executor& executor, kj::LiteralStringConst message) { + const kj::Executor* current = tryGetCurrentThreadExecutor(); + KJ_REQUIRE(current == &executor || current == nullptr, message); +} + } // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/executor-guarded.h b/src/rust/cxx/kj-rs/executor-guarded.h index e7f68c8d7d9..b00fb2944ca 100644 --- a/src/rust/cxx/kj-rs/executor-guarded.h +++ b/src/rust/cxx/kj-rs/executor-guarded.h @@ -4,11 +4,20 @@ namespace kj_rs { -// Return true if `executor`'s event loop is active on the current thread. +// Return true if `executor`'s event loop is active on the current thread. Never throws: a +// thread with no live event loop at all (e.g. after ~EventLoop during teardown) reports false. bool isCurrent(const kj::Executor& executor); // Assert that `executor`'s event loop is active on the current thread, or throw an exception // containing `message`. void requireCurrent(const kj::Executor& executor, kj::LiteralStringConst message); +// Like requireCurrent(), but tolerant of full loop teardown: if the current thread has NO event +// loop at all (~EventLoop already ran), return quietly instead of throwing. Used by +// ~ExecutorGuarded, which can legitimately run after loop teardown from Rust drop glue (e.g. the +// tokio runtime cancelling still-pending LocalSet tasks in TokioPort::drop, after +// TokioAsyncIoContext destroyed the WaitScope and EventLoop); a throw there would unwind through +// a cxx `prevent_unwind` boundary and abort the process. Destruction on a thread running a +// DIFFERENT live event loop is still a contract violation and throws. +void requireCurrentOrTearingDown(const kj::Executor& executor, kj::LiteralStringConst message); // ExecutorGuarded is a helper class which allows mutable access to a wrapped value to any thread // running the KJ event loop that was active at the time of construction. Any access attempts by a @@ -19,7 +28,10 @@ class ExecutorGuarded { template ExecutorGuarded(Args&&... args): value(kj::fwd(args)...) {} ~ExecutorGuarded() noexcept(false) { - requireCurrent(executor, "destruction on wrong event loop"_kjc); + // Teardown-tolerant (see requireCurrentOrTearingDown above): destruction with no event loop + // on the thread proceeds quietly (best-effort destruction of `value`); destruction on a + // thread running a different live loop still throws. + requireCurrentOrTearingDown(*executor, "destruction on wrong event loop"_kjc); } KJ_DISALLOW_COPY_AND_MOVE(ExecutorGuarded); @@ -29,7 +41,7 @@ class ExecutorGuarded { // Throws an exception with `message` if the current thread is not running the expected event // loop. T& get(kj::LiteralStringConst message = "access on wrong event loop"_kjc) const { - requireCurrent(executor, message); + requireCurrent(*executor, message); // Safety: const_cast is okay because we know that we are being accessed on a thread running our // original event loop. All successful accesses through `get()` are effectively single-threaded, @@ -38,7 +50,7 @@ class ExecutorGuarded { } kj::Maybe tryGet() const { - if (isCurrent(executor)) { + if (isCurrent(*executor)) { // Safety: const_cast is okay because we know that we are being accessed on a thread running our // original event loop. All successful accesses through `get()` are effectively single-threaded, // even though the event loop, and this object, may collectively move between threads. @@ -49,7 +61,15 @@ class ExecutorGuarded { } private: - const kj::Executor& executor = kj::getCurrentThreadExecutor(); + // Owned (via `addRef()`) rather than a bare `const kj::Executor&`, so the Executor object stays + // alive as long as this guard does. A guarded value can be destroyed by Rust *after* the KJ + // event loop has been torn down (e.g. a bridged future dropped during teardown); with a bare + // reference the destructor's executor check would take the address of — and a reused address + // could alias — a freed Executor. `addRef()` keeps the Executor at a stable, valid address; if + // the loop is gone, `isCurrent` reports false without throwing and the destructor's + // `requireCurrentOrTearingDown` lets destruction proceed quietly (`get()`/`requireCurrent` + // still throw, since post-teardown *access* remains a contract violation). + kj::Own executor = kj::getCurrentThreadExecutor().addRef(); T value; }; diff --git a/src/rust/cxx/kj-rs/ffi.rs b/src/rust/cxx/kj-rs/ffi.rs new file mode 100644 index 00000000000..cbecd8964d8 --- /dev/null +++ b/src/rust/cxx/kj-rs/ffi.rs @@ -0,0 +1,151 @@ +//! The `#[cxx::bridge]` FFI island for kj-rs. +//! +//! This is the crate's dedicated `#[cxx::bridge]` file (file-top `#![allow(unsafe_code)]`): the +//! cxx bridge DSL and its generated glue are inherently unsafe (extern "C++" vtables, placement +//! new/drop, C-ABI signatures) — a genuine seam. The bridge module is re-exported as +//! `crate::ffi::*`, the path the rest of the crate (awaiter.rs, promise.rs, waker.rs) uses. +//! +//! The crate root `lib.rs` is wholly-safe. The per-type vocabulary islands +//! (`future.rs`, `awaiter.rs`, `waker.rs`, `promise.rs`, `own.rs`, `refcount.rs`, `maybe.rs`) +//! carry their own file-top `#![allow(unsafe_code)]`: they implement individual unsafe primitive +//! types, distinct from this — the crate's cxx bridge. +#![allow(unsafe_code)] + +pub use bridge::*; + +use crate::awaiter::OptionWaker; +use crate::awaiter::WakerRef; + +// SAFETY: `FutureWakerCell` is C++'s thread-safe waker cell (waker.h). Everything Rust can reach +// through a shared reference or a `KjArc` handle is safe from any thread: `addRef`/`reown` and +// handle drops are atomic refcount operations (kj::AtomicRefcounted), `wakeByRef` routes wakes +// through an owning-executor check (same-thread direct arm, cross-thread fulfiller otherwise), +// and the last handle may destroy the cell on any thread (its members — an Executor own, a +// mutexed fulfiller, and an owning-thread-only weak link that destruction does not dereference — +// all tolerate that). These impls are what let `KjArc` (and thus the +// `std::task::Waker` built over it in waker.rs) be `Send + Sync`, as the `Waker` contract +// requires. +unsafe impl Send for bridge::FutureWakerCell {} +// SAFETY: see the `Send` impl above. +unsafe impl Sync for bridge::FutureWakerCell {} + +#[cxx::bridge(namespace = "kj_rs")] +// The cxx bridge DSL and its generated glue are inherently unsafe (extern "C++" vtables, placement +// new/drop, C-ABI signatures). This is a genuine seam. +// missing_safety_doc: the `# Safety` docs on `reown` below are for human readers only — the cxx +// macro does not forward doc comments to the generated unsafe shim, so the lint cannot be +// satisfied by documentation here. +#[expect(clippy::missing_safety_doc)] +mod bridge { + + /// Representation of a `GuardedRustPromiseAwaiter` in C++. The size of the blob should match. + #[derive(Debug)] + pub struct GuardedRustPromiseAwaiterRepr { + _bindgen_opaque_blob: [u64; 14usize], + } + + extern "Rust" { + type WakerRef<'a>; + } + + extern "Rust" { + // We expose the Rust Waker type to C++ through this OptionWaker reference wrapper. cxx-rs + // does not allow us to export types defined outside this crate, such as Waker, directly. + // + // `LazyRustPromiseAwaiter` (the implementation of `.await` syntax/the IntoFuture trait), + // stores a OptionWaker immediately after `GuardedRustPromiseAwaiter` in declaration order. + // pass the Waker to the `RustPromiseAwaiter` class, which is implemented in C++ + type OptionWaker; + fn set(&mut self, waker: &WakerRef); + fn set_none(&mut self); + fn wake_if_some(&mut self); + } + + unsafe extern "C++" { + include!("kj-rs/waker.h"); + + /// The stack-owned waker C++ passes to `Future::poll()`. Rust only ever borrows it; the + /// `Waker` built from it (waker.rs) has a no-op drop and clones by taking a real strong + /// reference to the event's `FutureWakerCell` via `clone_cell()`. Both operations are + /// safe from any thread (`&Waker` is `Sync`). + type PollWaker; + #[cxx_name = "wakeByRef"] + fn wake_by_ref(self: &PollWaker); + #[cxx_name = "cloneCell"] + fn clone_cell(self: &PollWaker) -> KjArc; + + /// The atomically-refcounted, thread-safe cell behind every retained waker; waking + /// it arms the owning FuturePollEvent — directly on the owning loop's thread, through a + /// cross-thread fulfiller from any other — and is a safe no-op after that event is + /// destroyed. + type FutureWakerCell; + #[cxx_name = "wakeByRef"] + fn wake_by_ref(self: &FutureWakerCell); + #[cxx_name = "addRef"] + fn add_ref(self: &FutureWakerCell) -> KjArc; + /// Re-own a strong reference previously disowned into a `RawWaker` data slot (waker.rs's + /// owned-cell vtable). + /// + /// # Safety + /// + /// `self` must carry exactly such a surrendered reference — this mints an owned handle + /// without incrementing the count. Dropping the returned handle releases the reference. + unsafe fn reown(self: &FutureWakerCell) -> KjArc; + } + + unsafe extern "C++" { + include!("kj-rs/promise.h"); + + type OwnPromiseNode = crate::OwnPromiseNode; + + // Takes `&mut` (not a raw pointer): this is a placement-destruct of a live + // `OwnPromiseNode` whose backing memory is owned by Rust and only reached through the + // `&mut self` in `OwnPromiseNode`'s `Drop`. The reference is valid for the call; the + // value is logically dead only after, inside `drop`, so no use-after-free is possible. + // Expressing it as a borrow lets cxx generate a safe-to-call binding. + fn own_promise_node_drop_in_place(node: &mut OwnPromiseNode); + } + + unsafe extern "C++" { + include!("kj-rs/awaiter.h"); + + type GuardedRustPromiseAwaiter; + + /// Placement-new of the C++ awaiter into Rust-owned storage. + /// + /// # Safety + /// + /// - `ptr` must point to uninitialized storage of (at least) the size and alignment of + /// `GuardedRustPromiseAwaiterRepr`, valid for writes, and must stay pinned for the + /// awaiter's lifetime. + /// - `rust_waker_ptr` must point to a valid `OptionWaker` that outlives the awaiter (the + /// C++ side stores and dereferences this pointer on later `poll()`s). + /// - The awaiter must eventually be destroyed exactly once via + /// `guarded_rust_promise_awaiter_drop_in_place`. + unsafe fn guarded_rust_promise_awaiter_new_in_place( + ptr: *mut GuardedRustPromiseAwaiter, + rust_waker_ptr: *mut OptionWaker, + node: OwnPromiseNode, + ); + /// Placement-destruct of the awaiter constructed by + /// `guarded_rust_promise_awaiter_new_in_place`. + /// + /// # Safety + /// + /// `ptr` must point to a live awaiter previously constructed in that storage by + /// `guarded_rust_promise_awaiter_new_in_place`, and the awaiter must not be used again + /// afterwards (at most one drop per construction). + unsafe fn guarded_rust_promise_awaiter_drop_in_place(ptr: *mut GuardedRustPromiseAwaiter); + + fn poll(self: Pin<&mut GuardedRustPromiseAwaiter>, waker: &WakerRef) -> bool; + #[cxx_name = "pollWithPollWaker"] + fn poll_with_poll_waker( + self: Pin<&mut GuardedRustPromiseAwaiter>, + waker: &WakerRef, + poll_waker: &PollWaker, + ) -> bool; + + #[must_use] + fn take_own_promise_node(self: Pin<&mut GuardedRustPromiseAwaiter>) -> OwnPromiseNode; + } +} diff --git a/src/rust/cxx/kj-rs/future.h b/src/rust/cxx/kj-rs/future.h index 39aae452a09..78b2abbb5b5 100644 --- a/src/rust/cxx/kj-rs/future.h +++ b/src/rust/cxx/kj-rs/future.h @@ -67,7 +67,7 @@ namespace repr { // ::kj_rs::repr::PollCallback using PollCallback = kj_rs::FuturePollStatus (*)( - void /* RustFuture::fut */* fut, const void* waker, void /* T */* ret); + void /* RustFuture::fut */* fut, const ::kj_rs::PollWaker& waker, void /* T */* ret); // ::kj_rs::repr::DropCallback using DropCallback = void (*)(void /* RustFuture::fut */* fut); @@ -79,8 +79,31 @@ using DropCallback = void (*)(void /* RustFuture::fut */* fut); // which drops the Rust Future and transitively cancels any KJ sub-promises it was .await'ing. struct RustFuture { + // Eager-by-default conversion: the returned promise starts running immediately, without + // being awaited — the future is polled synchronously up to its first suspension point + // (exactly like calling a KJ coroutine, which runs to its first co_await), and continues + // on the event loop from there. KJ code universally assumes promises are "hot" (a stored + // promise still makes progress), so this is the right default for every bridged + // `async fn`; before this conversion was eager, every consumer had to remember a manual + // `.eagerlyEvaluate(nullptr)`. + // + // Requires a current kj::EventLoop on this thread (same requirement as awaiting the + // promise, just enforced at creation). Cancellation is unchanged: dropping the promise + // still synchronously cancels the Rust future and everything it is awaiting. + // + // The rare consumer that genuinely wants a cold future can call `lazily()` below on the + // raw RustFuture instead of going through this conversion (the bridge's generated shims + // always convert eagerly, so that consumer must obtain the RustFuture itself). template operator kj::Promise() { + return lazily().eagerlyEvaluate(nullptr); + } + + // Lazy (cold) conversion: nothing runs until the returned promise is first awaited. + // This is the raw adapter the eager conversion above builds on, and the C++-side escape + // hatch for code that genuinely needs a cold promise. + template + kj::Promise lazily() { struct Impl { using ExceptionOrValue = ::kj::_::ExceptionOr<::kj::_::FixVoid>; using Output = ::kj::_::FixVoid; @@ -100,10 +123,9 @@ struct RustFuture { KJ_DISALLOW_COPY(Impl); - void poll(const ::kj_rs::KjWaker& waker, ExceptionOrValue& output) noexcept { + void poll(const ::kj_rs::PollWaker& waker, ExceptionOrValue& output) noexcept { ::kj_rs::FuturePoller poller; - poller.poll( - [this, &waker](void* result) { return fut.poll(&fut, &waker, result); }, output); + poller.poll([this, &waker](void* result) { return fut.poll(&fut, waker, result); }, output); } RustFuture fut; diff --git a/src/rust/cxx/kj-rs/future.rs b/src/rust/cxx/kj-rs/future.rs index aab7b9dbd32..707eb53618f 100644 --- a/src/rust/cxx/kj-rs/future.rs +++ b/src/rust/cxx/kj-rs/future.rs @@ -1,3 +1,8 @@ +//! FFI island (see crate-root `#![deny(unsafe_code)]`): the `RustFuture` C-ABI vtable that drives +//! all bridged async — `unsafe extern "C"` poll/drop callbacks, raw-pointer result writes, and +//! `Pin`/`Box` raw conversions. A genuine unsafe seam. +#![allow(unsafe_code)] + // This file contains boilerplate which must occur once per crate, rather than once per type. use std::pin::Pin; @@ -33,11 +38,59 @@ pub mod repr { use static_assertions::assert_eq_size; use super::FuturePollStatus; - use crate::KjWaker; + use crate::ffi::PollWaker; + + /// Converts a panic payload (from `std::panic::catch_unwind`) escaping a bridged future + /// into a heap-allocated `kj::Exception` written to the poll callback's output parameter, + /// so C++ observes a rejected promise instead of a process abort. + /// + /// This mirrors the sync bridge path (`cxx::private::try_unwind`/`catch_unwind` in + /// src/unwind.rs), which converts panics in `extern "Rust"` functions into + /// `kj::Exception`s. One divergence: a `cxx::CanceledException` payload (produced when an + /// infallible `extern "C++"` call throws `kj::CanceledException`) cannot be propagated as + /// a distinct "canceled" state here, because `FuturePollStatus` has no Canceled arm; it is + /// reported as a regular `kj::Exception` describing the cancellation instead. + /// + /// # Safety + /// + /// `ret` must point to storage valid for holding a `kj::Exception*` (the C++ + /// `FuturePoller` union guarantees this for the Error arm). + #[expect( + clippy::needless_pass_by_value, + reason = "takes ownership of the panic payload, mirroring std::panic::catch_unwind's Err arm" + )] + unsafe fn write_panic_as_exception( + ret: *mut c_void, + err: Box, + ) -> FuturePollStatus { + let msg = if let Some(s) = err.downcast_ref::<&'static str>() { + format!("panic in bridged future poll: {s}") + } else if let Some(s) = err.downcast_ref::() { + format!("panic in bridged future poll: {s}") + } else if err.downcast_ref::().is_some() { + "panic in bridged future poll: kj::CanceledException".to_owned() + } else { + "panic in bridged future poll".to_owned() + }; + let exception = cxx::IntoKjException::into_kj_exception( + cxx::KjError::new(cxx::KjExceptionType::Failed, msg), + file!(), + line!(), + ); + // SAFETY: `ret` points to storage valid for a `kj::Exception*` per this fn's + // `# Safety` contract (the C++ `FuturePoller` Error arm). + unsafe { + std::ptr::write( + ret.cast::<*mut c_void>(), + exception.into_raw().as_ptr().cast(), + ); + } + FuturePollStatus::ERROR + } - type PollCallback = unsafe extern "C" fn( + type PollCallback = for<'a> unsafe extern "C" fn( fut: *mut c_void, - waker: *const c_void, + waker: &'a PollWaker, ret: *mut c_void, ) -> FuturePollStatus; @@ -71,27 +124,41 @@ pub mod repr { assert_eq_size!(RustInfallibleFuture<()>, [*mut c_void; 4]); impl RustFuture<'_, T> { + /// # Safety + /// + /// C++ `RustFuture` vtable protocol (future.h): `fut` must be the `fut` field of a + /// live, not-yet-dropped `RustFuture` created by [`future`]; `ret` must point to + /// storage suitable for a `T` + /// (Complete) or a `kj::Exception*` (Error), per `FuturePoller`'s union. + /// + /// Unwind safety: any panic escaping the wrapped future's `poll` (or the waker + /// machinery) is caught here and converted into an errored completion, because + /// unwinding out of an `extern "C"` fn is instant process abort (Rust >= 1.81). + /// This makes a panicking bridged `async fn` surface to C++ as a rejected + /// `kj::Promise` carrying a `kj::Exception`, matching the sync bridge path. pub(crate) unsafe extern "C" fn poll( fut: *mut c_void, - waker: *const c_void, + waker: &PollWaker, ret: *mut c_void, ) -> FuturePollStatus { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: per this fn's `# Safety` contract, `fut` is the `fut` field of a live + // `RustFuture`, i.e. a valid `*mut FuturePtr` we may read and then pin. let fut = unsafe { *(fut.cast::>()) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the boxed future is never moved out of its heap allocation, so pinning + // the `&mut` reborrow is sound. let fut = unsafe { Pin::new_unchecked(&mut *fut) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - let waker = unsafe { &*waker.cast::() }; - let waker = Waker::from(waker); - let mut context = Context::from_waker(&waker); - match fut.poll(&mut context) { - Poll::Ready(Ok(value)) => { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let waker = Waker::from(waker); + let mut context = Context::from_waker(&waker); + fut.poll(&mut context) + })) { + Ok(Poll::Ready(Ok(value))) => { + // SAFETY: `ret` points to storage suitable for a `T` (Complete arm). unsafe { std::ptr::write(ret.cast::(), value) }; FuturePollStatus::COMPLETE } - Poll::Ready(Err(error)) => { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + Ok(Poll::Ready(Err(error))) => { + // SAFETY: `ret` points to storage for a `kj::Exception*` (Error arm). unsafe { std::ptr::write( ret.cast::<*mut c_void>(), @@ -100,53 +167,90 @@ pub mod repr { }; FuturePollStatus::ERROR } - Poll::Pending => FuturePollStatus::PENDING, + Ok(Poll::Pending) => FuturePollStatus::PENDING, + // SAFETY: `ret` is the Error-arm storage; forwarded to `write_panic_as_exception`. + Err(panic_payload) => unsafe { write_panic_as_exception(ret, panic_payload) }, } } + /// # Safety + /// + /// C++ `RustFuture` vtable protocol (future.h): `fut` must be the `fut` field of a + /// live `RustFuture` created by [`future`], and must not be used again afterwards + /// (drop-exactly-once, enforced by `Impl`'s move semantics on the C++ side). + /// + /// Unwind safety: a panic in the future's destructor has no error channel (this is + /// called from C++ destructors/cancellation paths), so it is converted into a + /// deterministic, labeled abort via `cxx::private::prevent_unwind` — the same + /// semantics the sync bridge uses for panics that cannot be reported (rather than + /// the unlabeled langdef abort of unwinding out of `extern "C"`). pub(crate) unsafe extern "C" fn drop_in_place(fut: *mut c_void) { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: per this fn's `# Safety` contract, `fut` is the `fut` field of a live + // `RustFuture` not used again, so we may read the pointer, reclaim the box, and pin. let fut = unsafe { *(fut.cast::>()) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `fut` was produced by `Box::into_raw` in [`future`]; reclaim ownership once. let fut = unsafe { Box::from_raw(fut) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the box is never moved out of its allocation, so pinning it is sound. let fut = unsafe { Pin::new_unchecked(fut) }; - drop(fut); + cxx::private::prevent_unwind("kj_rs::repr::RustFuture::drop_in_place", move || { + drop(fut); + }); } } impl RustInfallibleFuture<'_, T> { + /// # Safety + /// + /// Same contract as [`RustFuture::poll`], with `fut` created by [`infallible_future`]. + /// Although the future itself cannot return an error, a panic escaping its `poll` is + /// still converted into an errored completion (`FuturePollStatus::ERROR` writing a + /// `kj::Exception*`): the C++ `FuturePoller` handles the Error arm identically for + /// infallible futures, and `kj::Promise` can always carry an exception. pub(crate) unsafe extern "C" fn poll( fut: *mut c_void, - waker: *const c_void, + waker: &PollWaker, ret: *mut c_void, ) -> FuturePollStatus { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: per this fn's `# Safety` contract, `fut` is the `fut` field of a live + // `RustInfallibleFuture`, i.e. a valid `*mut InfallibleFuturePtr` to read+pin. let fut = unsafe { *(fut.cast::>()) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the boxed future is never moved out of its allocation, so pinning is sound. let fut = unsafe { Pin::new_unchecked(&mut *fut) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - let waker = unsafe { &*waker.cast::() }; - let waker = Waker::from(waker); - let mut context = Context::from_waker(&waker); - match fut.poll(&mut context) { - Poll::Ready(value) => { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let waker = Waker::from(waker); + let mut context = Context::from_waker(&waker); + fut.poll(&mut context) + })) { + Ok(Poll::Ready(value)) => { + // SAFETY: `ret` points to storage suitable for a `T` (Complete arm). unsafe { std::ptr::write(ret.cast::(), value) }; FuturePollStatus::COMPLETE } - Poll::Pending => FuturePollStatus::PENDING, + Ok(Poll::Pending) => FuturePollStatus::PENDING, + // SAFETY: `ret` is the Error-arm storage; forwarded to `write_panic_as_exception`. + Err(panic_payload) => unsafe { write_panic_as_exception(ret, panic_payload) }, } } + /// # Safety + /// + /// Same contract as [`RustFuture::drop_in_place`], with `fut` created by + /// [`infallible_future`]. A panic in the destructor aborts deterministically with a + /// label (see there for rationale). pub(crate) unsafe extern "C" fn drop_in_place(fut: *mut c_void) { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: per this fn's `# Safety` contract, `fut` is the `fut` field of a live + // `RustInfallibleFuture` not used again; read the pointer, reclaim the box, and pin. let fut = unsafe { *(fut.cast::>()) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `fut` came from `Box::into_raw` in [`infallible_future`]; reclaim once. let fut = unsafe { Box::from_raw(fut) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the box is never moved out of its allocation, so pinning it is sound. let fut = unsafe { Pin::new_unchecked(fut) }; - drop(fut); + cxx::private::prevent_unwind( + "kj_rs::repr::RustInfallibleFuture::drop_in_place", + move || { + drop(fut); + }, + ); } } @@ -154,7 +258,8 @@ pub mod repr { pub fn future<'a, T: Unpin>( fut: Pin> + 'a>>, ) -> RustFuture<'a, T> { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the box is immediately re-boxed via `Box::into_raw` and only ever reconstituted + // (and re-pinned) in `drop_in_place`, so the pinned future is never moved. let fut = Box::into_raw(unsafe { Pin::into_inner_unchecked(fut) }); let poll = RustFuture::::poll; let drop = RustFuture::::drop_in_place; @@ -165,7 +270,8 @@ pub mod repr { pub fn infallible_future<'a, T: Unpin>( fut: Pin + 'a>>, ) -> RustInfallibleFuture<'a, T> { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the box is immediately re-boxed via `Box::into_raw` and only ever reconstituted + // (and re-pinned) in `drop_in_place`, so the pinned future is never moved. let fut = Box::into_raw(unsafe { Pin::into_inner_unchecked(fut) }); let poll = RustInfallibleFuture::::poll; let drop = RustInfallibleFuture::::drop_in_place; diff --git a/src/rust/cxx/kj-rs/lib.rs b/src/rust/cxx/kj-rs/lib.rs index c7df9fcd7d2..e2a2d995d36 100644 --- a/src/rust/cxx/kj-rs/lib.rs +++ b/src/rust/cxx/kj-rs/lib.rs @@ -1,6 +1,39 @@ -use awaiter::OptionWaker; +// Safety & panic enforcement walls. Inherent-FFI crate: unsafe is +// concentrated at the bridge and every op must sit in an explicit, documented unsafe +// block; prod code returns Result/KjError rather than panicking (a panic on the async +// poll path is a process abort). Test code is exempted below. +#![deny(unsafe_op_in_unsafe_fn)] +// Quarantine unsafe into named FFI islands: deny unsafe crate-wide, then re-allow it only on the +// modules that genuinely need it (each carries its own `#![allow(unsafe_code)]`). Any module +// without that opt-in — and any newly-added module — is compiler-proven unsafe-free, and no future +// edit can smuggle unsafe into non-island code without tripping this deny. +#![deny(unsafe_code)] +#![deny(clippy::undocumented_unsafe_blocks)] +#![deny( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::unreachable, + clippy::todo, + clippy::unimplemented +)] +#![cfg_attr( + test, + allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::unreachable, + clippy::todo, + clippy::unimplemented + ) +)] + +// The cxx bridge expands vocabulary builtins (KjRc, KjMaybe, ...) to `::kj_rs::...` paths; make +// that path resolve inside this crate itself, since ffi.rs's bridge uses them too. +extern crate self as kj_rs; + pub use awaiter::PromiseAwaiter; -use awaiter::WakerRef; pub use date::KjDate; pub use future::FuturePollStatus; pub use future::map_err; @@ -14,10 +47,12 @@ pub use promise::new_callbacks_promise_future; pub use refcount::repr::KjArc; pub use refcount::repr::KjRc; -pub use crate::ffi::KjWaker; +pub use crate::ffi::FutureWakerCell; +pub use crate::ffi::PollWaker; mod awaiter; mod date; +mod ffi; mod future; pub mod maybe; mod own; @@ -36,80 +71,3 @@ pub type Result = std::io::Result; pub type Error = std::io::Error; pub trait JsgStruct {} - -#[cxx::bridge(namespace = "kj_rs")] -mod ffi { - - /// Representation of a `GuardedRustPromiseAwaiter` in C++. The size of the blob should match. - #[derive(Debug)] - pub struct GuardedRustPromiseAwaiterRepr { - _bindgen_opaque_blob: [u64; 13usize], - } - - extern "Rust" { - type WakerRef<'a>; - } - - extern "Rust" { - // We expose the Rust Waker type to C++ through this OptionWaker reference wrapper. cxx-rs - // does not allow us to export types defined outside this crate, such as Waker, directly. - // - // `LazyRustPromiseAwaiter` (the implementation of `.await` syntax/the IntoFuture trait), - // stores a OptionWaker immediately after `GuardedRustPromiseAwaiter` in declaration order. - // pass the Waker to the `RustPromiseAwaiter` class, which is implemented in C++ - type OptionWaker; - fn set(&mut self, waker: &WakerRef); - fn set_none(&mut self); - fn wake_if_some(&mut self); - } - - unsafe extern "C++" { - include!("kj-rs/waker.h"); - - // Match the definition of the abstract virtual class in the C++ header. - type KjWaker; - #[cxx_name = "clone"] - fn clone_kj_waker(&self) -> *const KjWaker; - fn wake(&self); - fn wake_by_ref(&self); - fn drop(&self); - } - - unsafe extern "C++" { - include!("kj-rs/promise.h"); - - type OwnPromiseNode = crate::OwnPromiseNode; - - /// # Safety - /// `node` must point to a live `OwnPromiseNode`. - unsafe fn own_promise_node_drop_in_place(node: *mut OwnPromiseNode); - } - - unsafe extern "C++" { - include!("kj-rs/awaiter.h"); - - type GuardedRustPromiseAwaiter; - - /// # Safety - /// The pointers must identify valid storage and a live waker for the awaiter's lifetime. - unsafe fn guarded_rust_promise_awaiter_new_in_place( - ptr: *mut GuardedRustPromiseAwaiter, - rust_waker_ptr: *mut OptionWaker, - node: OwnPromiseNode, - ); - /// # Safety - /// `ptr` must point to an initialized guarded awaiter. - unsafe fn guarded_rust_promise_awaiter_drop_in_place(ptr: *mut GuardedRustPromiseAwaiter); - - /// # Safety - /// `maybe_kj_waker`, when non-null, must point to a live `KjWaker`. - unsafe fn poll( - self: Pin<&mut GuardedRustPromiseAwaiter>, - waker: &WakerRef, - maybe_kj_waker: *const KjWaker, - ) -> bool; - - #[must_use] - fn take_own_promise_node(self: Pin<&mut GuardedRustPromiseAwaiter>) -> OwnPromiseNode; - } -} diff --git a/src/rust/cxx/kj-rs/linked-group.h b/src/rust/cxx/kj-rs/linked-group.h deleted file mode 100644 index 34567a71e2b..00000000000 --- a/src/rust/cxx/kj-rs/linked-group.h +++ /dev/null @@ -1,315 +0,0 @@ -#pragma once - -#include - -namespace kj_rs { - -// `LinkedGroup` and `LinkedObject` are CRTP mixins which allow derived classes G and O -// to weakly refer to each other in a one-to-many relationship. -// -// For example, say you have two classes, Group and Object. There exists a natural one-to-many -// relationship between the two. Given a Group, you would like to be able to dererefence its -// Objects, and, given an Object, you would like be able to dereference its Group. Further suppose -// the objects have independent lifetimes: Objects may be destroyed before their Groups, and Groups -// may be destroyed before their Objects. -// -// If you are operating in a single-threaded context (or can provide sufficient synchronization), -// and if Group and Object are both immobile (non-copyable, non-moveable) classes, then -// `LinkedGroup` and `LinkedObject` can be used to implement the above -// scenario safely. To do so, first: -// -// - Your Group class must publicly inherit from `LinkedGroup`. -// - Your Object class must publicly inherit from `LinkedObject`. -// -// This will add one protected member function to each of your derived classes: -// `Object::linkedGroup()`, and `Group::linkedObjects()`. They are protected so that they are not -// part of your type's public API unless you explicitly want them to be, e.g., with a public `using` -// statement like `using LinkedGroup::linkedObjects`. -// -// You can use `Object::linkedGroup()` to manage Group membership and dereference Groups from -// Objects: -// -// - `object.linkedGroup().set(group)` adds an Object to a Group. -// This function implicitly removes the Object from its current Group, if any. -// - `object.linkedGroup().set(kj::none)` removes an Object from its current Group, if any. -// - `object.linkedGroup().tryGet()` dereferences the Object's current Group, if any. -// -// You can use `Group::linkedObjects()` to iterate over the list of currently linked Objects. -// -// - `group.linkedObjects().begin()` obtains an iterator to the beginning of the list of Objects. -// - `group.linkedObjects().end()` obtains an iterator to the end of the list of Objets. -// - `group.linkedObjects().front()` dereferences the front of the list of Objects. -// Calling `front()` on an empty list (`begin() == end()`) is undefined behavior. -// - `group.linkedObjects().empty()` is true if there are no Objects in the list. -// -// Finally, destroying either the Group or its Object safely severs their relationship(s). -// -// - Destroying an Object implicitly calls `object.linkedGroup().set(kj::none)` on itself. -// - Destroying a Group implicitly calls `object.linkedGroup().set(kj::none)` on all its objects. -// -// Considerations: -// -// - Your Group object's destructor will contain a _O(n)_ algorithm inside it, with _n_ being the -// number of linked objects at destruction time. If Groups frequently outlive large sets of -// Objects, this may be an issue to consider. -// - It is valid to remove the front Object in a `Group::linkedObjects()` list while iterating -// over the list. Removing an Object in any other position in the list will invalidate all -// existing iterators. -// -// TODO(someday): Multiple inheritance if an object must join multiple groups, or a group must -// have multiple linked object types? Can we write something like `linkedGroup()` in the -// LinkedObject derived class, and `linkedObjects()` in the LinkedGroup derived class? -template -class LinkedGroup; -template -class LinkedObject; - -template -class StaticCastIterator; - -// CRTP mixin for derived class G. -template -class LinkedGroup { - public: - LinkedGroup() = default; - ~LinkedGroup() noexcept(false) { - for (auto& object: list) { - object.removeFromGroup(*this); - } - } - KJ_DISALLOW_COPY_AND_MOVE(LinkedGroup); - - private: - // We'll refer to the `LinkedObject` type quite a bit below, so we shadow the class - // template with our own convenience typedef. But, we need to give LinkedObject friend access to - // us first. - friend class LinkedObject; - using LinkedObject = LinkedObject; - - using List = kj::List; - - using ListIterator = kj::ListIterator; - using ConstListIterator = kj::ListIterator; - - using Iterator = StaticCastIterator; - using ConstIterator = StaticCastIterator; - - protected: - // A proxy class representing this LinkedGroup's list of LinkedObjects, if any. Instead of - // exposing multiple functions on LinkedGroup, we expose one: `linkedObjects()`, and that function - // returns an object of this proxy class (or the similar ConstLinkedObjectList class below). - class LinkedObjectList { - public: - LinkedObjectList(List& list): list(list) {} - Iterator begin() { - return list.begin(); - } - Iterator end() { - return list.end(); - } - decltype(*kj::instance()) front() { - return *begin(); - } - bool empty() const { - return list.empty(); - } - - private: - List& list; - }; - - class ConstLinkedObjectList { - public: - ConstLinkedObjectList(const List& list): list(list) {} - ConstIterator begin() const { - return list.begin(); - } - ConstIterator end() const { - return list.end(); - } - decltype(*kj::instance()) front() const { - return *begin(); - } - bool empty() const { - return list.empty(); - } - - private: - const List& list; - }; - - LinkedObjectList linkedObjects() { - return LinkedObjectList(list); - } - ConstLinkedObjectList linkedObjects() const { - return ConstLinkedObjectList(list); - } - - private: - kj::List list; -}; - -// CRTP mixin for derived class O. -template -class LinkedObject { - public: - LinkedObject() = default; - ~LinkedObject() noexcept(false) { - invalidateGroup(); - } - KJ_DISALLOW_COPY_AND_MOVE(LinkedObject); - - private: - // We'll refer to the `LinkedGroup` type quite a bit below, so we shadow the class template - // with our own convenience typedef. But, we need to give LinkedGroup friend access to us first. - friend class LinkedGroup; - using LinkedGroup = LinkedGroup; - - protected: - // A proxy class representing this LinkedObject's LinkedGroup, if any. Instead of exposing - // multiple functions on LinkedObject, we expose one: `linkedGroup()`, and that function returns - // an object of this proxy class (or the similar ConstLinkedGroupProxy class below). - class LinkedGroupProxy { - public: - LinkedGroupProxy(LinkedObject& self): self(self) {} - void set(LinkedGroup& newGroup) { - self.setGroup(newGroup); - } - void set(kj::None) { - self.invalidateGroup(); - } - kj::Maybe tryGet() { - return self.tryGetGroup(); - } - - private: - LinkedObject& self; - }; - - // Const version of LinkedGroupProxy, exposing only `tryGet()`. - class ConstLinkedGroupProxy { - public: - ConstLinkedGroupProxy(const LinkedObject& self): self(self) {} - kj::Maybe tryGet() const { - return self.tryGetGroup(); - } - - private: - const LinkedObject& self; - }; - - // Provide access to this Object's LinkedGroup, if any. - LinkedGroupProxy linkedGroup() { - return *this; - } - ConstLinkedGroupProxy linkedGroup() const { - return *this; - } - - private: - void setGroup(LinkedGroup& newGroup) { - // Invalidate our current group membership, if any. - KJ_IF_SOME(oldGroup, maybeGroup) { - // If we're already a member of `newGroup`, we're done. Otherwise, we must remove ourselves - // from the old group. - if (&newGroup == &oldGroup) { - return; - } else { - removeFromGroup(oldGroup); - } - } else { - KJ_IREQUIRE(!link.isLinked()); - } - - // Add ourselves to the new group. - newGroup.list.add(*this); - maybeGroup = newGroup; - } - - kj::Maybe tryGetGroup() { - KJ_IF_SOME(group, maybeGroup) { - KJ_IREQUIRE(link.isLinked()); - return static_cast(group); - } else { - KJ_IREQUIRE(!link.isLinked()); - return kj::none; - } - } - - kj::Maybe tryGetGroup() const { - KJ_IF_SOME(group, maybeGroup) { - KJ_IREQUIRE(link.isLinked()); - return static_cast(group); - } else { - KJ_IREQUIRE(!link.isLinked()); - return kj::none; - } - } - - void invalidateGroup() { - KJ_IF_SOME(group, maybeGroup) { - removeFromGroup(group); - } else { - KJ_IREQUIRE(!link.isLinked()); - } - } - - // Helper for `setGroup()`, `invalidateGroup()`, and `~LinkedGroup()`. - void removeFromGroup(LinkedGroup& group) { - KJ_IREQUIRE(link.isLinked()); - group.list.remove(*this); - maybeGroup = kj::none; - } - - kj::ListLink link; - kj::Maybe maybeGroup; -}; - -// An iterator which wraps `InnerIterator` and `static_cast`s all mutable dereferences to -// `MaybeConstT&`, and all const dereferences to `const T&`. -// -// With the Ranges TS, all of this nonsense could be boiled down to a one-liner based on -// `std::views::transform()`. I encountered too many puzzles to solve while trying to get that -// working, so here we are. -template -class StaticCastIterator { - public: - // Construct an iterator using a default-constructed InnerIterator. In practice, this constructs - // an end iterator. - StaticCastIterator() = default; - - // Construct an iterator wrapping `inner`. - StaticCastIterator(InnerIterator inner): inner(inner) {} - - MaybeConstT& operator*() { - return static_cast(*inner); - } - const T& operator*() const { - return static_cast(*inner); - } - MaybeConstT* operator->() { - return static_cast(inner.operator->()); - } - const T* operator->() const { - return static_cast(inner.operator->()); - } - - inline StaticCastIterator& operator++() { - ++inner; - return *this; - } - inline StaticCastIterator operator++(int) { - StaticCastIterator result = *this; - ++inner; - return result; - } - - inline bool operator==(const StaticCastIterator& other) const { - return inner == other.inner; - } - - private: - InnerIterator inner; -}; - -} // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/maybe.rs b/src/rust/cxx/kj-rs/maybe.rs index 2b5b015c5ee..2a528d39aff 100644 --- a/src/rust/cxx/kj-rs/maybe.rs +++ b/src/rust/cxx/kj-rs/maybe.rs @@ -1,3 +1,9 @@ +//! FFI island: the `KjMaybe` representation of `kj::Maybe`. +//! +//! (See crate-root `#![deny(unsafe_code)]`.) Carries the `unsafe trait` niche contracts +//! (`HasNiche`/`MaybeItem`) and `assume_init` on the discriminated union. A genuine unsafe seam. +#![allow(unsafe_code)] + use std::mem::MaybeUninit; use std::pin::Pin; @@ -31,11 +37,12 @@ unsafe trait HasNiche: Sized { fn is_niche(value: *const Self) -> bool; } -// In Rust, references are not allowed to be null, so a null `MaybeUninit<&T>` is a niche -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: in Rust, references are not allowed to be null, so a null `MaybeUninit<&T>` is a +// niche (see the `HasNiche` trait contract above). unsafe impl HasNiche for &T { fn is_niche(value: *const &T) -> bool { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `value` points to a valid `&T`; we read it as a `*const *const T` (never as a + // reference, which the compiler assumes non-null) to test the pointer for null. unsafe { // We must cast it as pointing to a pointer, as opposed to a reference, // because the rust compiler assumes a reference is never null, and @@ -45,10 +52,10 @@ unsafe impl HasNiche for &T { } } -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: as for `&T` — a null `&mut T` is the niche (see the `HasNiche` trait contract). unsafe impl HasNiche for &mut T { fn is_niche(value: *const &mut T) -> bool { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `value` points to a valid `&mut T`; read as `*const *mut T` to null-check. unsafe { // We must cast it as pointing to a pointer, as opposed to a reference, // because the rust compiler assumes a reference is never null, and @@ -58,10 +65,11 @@ unsafe impl HasNiche for &mut T { } } -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: as for `&mut T` — a null pointee is the niche (see the `HasNiche` trait contract). unsafe impl HasNiche for Pin<&mut T> { fn is_niche(value: *const Pin<&mut T>) -> bool { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `value` points to a valid `Pin<&mut T>` (layout-identical to `&mut T`); read + // as `*const *mut T` to null-check. unsafe { // We must cast it as pointing to a pointer, as opposed to a reference, // because the rust compiler assumes a reference is never null, and @@ -72,10 +80,10 @@ unsafe impl HasNiche for Pin<&mut T> { } // In `kj`, `kj::Own` are considered `none` in a `Maybe` if the data pointer is null -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: a `KjOwn` with a null data pointer is `kj::none` (see the `HasNiche` trait contract). unsafe impl HasNiche for crate::repr::KjOwn { fn is_niche(value: *const Self) -> bool { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `value` points to a valid `KjOwn`; querying its data pointer is sound. unsafe { (*value).as_ptr().is_null() } } } @@ -111,7 +119,8 @@ pub unsafe trait MaybeItem: Sized { } fn drop_in_place(value: &mut KjMaybe) { if ::is_some(value) { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `is_some` just confirmed the `some` union member is initialized, so + // dropping it in place is sound. `KjMaybe`'s `Drop` calls this exactly once. unsafe { value.some.assume_init_drop(); } @@ -123,7 +132,9 @@ pub unsafe trait MaybeItem: Sized { /// Avoids running into generic specialization problems. macro_rules! impl_maybe_item_for_has_niche { ($ty:ty) => { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `$ty` is only ever a `HasNiche` type (enforced at the macro's use sites), so + // it carries a `()` discriminant and detects `none` via its null niche — matching kj's + // niche-value-optimized `Maybe` layout, as the `MaybeItem` trait contract requires. unsafe impl MaybeItem for $ty { type Discriminant = (); @@ -160,7 +171,9 @@ macro_rules! impl_maybe_item_for_has_niche { /// Avoids running into generic specialization problems. macro_rules! impl_maybe_item_for_primitive { ($ty:ty) => { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: primitives have no niche, so this mirrors kj's non-niche + // `kj::_::NullableValue` layout with an explicit `bool` discriminant (`is_set`) + // followed by the value, exactly as the `MaybeItem` trait contract requires. unsafe impl MaybeItem for $ty { type Discriminant = bool; @@ -198,7 +211,8 @@ impl_maybe_item_for_primitive!( u8, u16, u32, u64, u128, usize, i8, i16, i32, i64, i128, isize, f32, f64, bool, &str, String ); -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: `&[T]` is a fat pointer with no usable niche here, so it uses the explicit +// `bool`-discriminant (non-niche) `MaybeItem` representation, matching kj's layout. unsafe impl MaybeItem for &[T] { type Discriminant = bool; @@ -233,7 +247,9 @@ unsafe impl MaybeItem for &[T] { // // We therefore mirror that layout with a `bool` discriminant here, exactly like // the primitive types above, rather than implementing [`HasNiche`]. -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// +// SAFETY: `kj::Rc` defines no `Maybe` niche members, so `kj::Maybe>` uses the +// non-niche `bool`-discriminant `NullableValue` layout mirrored here (see comment above). unsafe impl MaybeItem for crate::KjRc { type Discriminant = bool; @@ -260,7 +276,9 @@ unsafe impl MaybeItem for crate::KjRc { } } -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: like `kj::Rc`, `kj::Arc` defines no `Maybe` niche members, so +// `kj::Maybe>` uses the non-niche `bool`-discriminant `NullableValue` layout +// mirrored here. unsafe impl MaybeItem for crate::KjArc { type Discriminant = bool; @@ -387,7 +405,9 @@ pub(crate) mod repr { if value.is_some() { // We can't move out of value so we copy it and forget it in // order to perform a "manual" move out of value - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `is_some` confirmed `some` is initialized; `assume_init_read` copies + // it out, and the immediately following `mem::forget(value)` prevents the + // source from being dropped, so ownership moves out exactly once. let ret = unsafe { Some(value.some.assume_init_read()) }; std::mem::forget(value); ret @@ -408,10 +428,10 @@ pub(crate) mod repr { if self.is_none() { write!(f, "Maybe::None") } else { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - write!(f, "Maybe::Some({:?})", unsafe { - self.some.assume_init_ref() - }) + // SAFETY: the `is_none()` branch above is false here, so `some` is + // initialized and may be borrowed for formatting. + let value = unsafe { self.some.assume_init_ref() }; + write!(f, "Maybe::Some({value:?})") } } } diff --git a/src/rust/cxx/kj-rs/own.rs b/src/rust/cxx/kj-rs/own.rs index 2061fb60168..a89fc5a1f09 100644 --- a/src/rust/cxx/kj-rs/own.rs +++ b/src/rust/cxx/kj-rs/own.rs @@ -1,4 +1,8 @@ //! The `workerd-cxx` module containing the [`Own`] type, which is bindings to the `kj::Own` C++ type +//! +//! FFI island (see crate-root `#![deny(unsafe_code)]`): `KjOwn` mirrors `kj::Own` — raw-pointer +//! deref and `extern "C"` disposer/refcount calls. A genuine unsafe seam. +#![allow(unsafe_code)] use std::fmt; use std::marker::PhantomData; @@ -21,18 +25,16 @@ impl NonNullExceptMaybe { } pub unsafe fn as_ref(&self) -> &T { - // Safety: - // This value will only be null when in a [`Maybe`], which does niche value optimization - // for a null pointer, so the inner [`Own`] can never be accessed if it is null - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self.0` is null only when this `NonNullExceptMaybe` lives inside a + // `Maybe` (which niche-optimizes the null pointer and never dereferences the inner + // `Own`), so here — reached only through the non-null `Own` API — it is a valid, + // live pointer. The caller's `unsafe` obligation is that `self` outlives the borrow. unsafe { &*self.0 } } pub unsafe fn as_mut(&mut self) -> &mut T { - // Safety: - // This value will only be null when in a [`Maybe`], which does niche value optimization - // for a null pointer, so the inner [`Own`] can never be accessed if it is null - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: as in `as_ref`, `self.0` is non-null and live when reached through the + // `Own` API; `&mut self` gives exclusive access, so the mutable reborrow is unique. unsafe { &mut *self.0 } } } @@ -119,11 +121,19 @@ pub mod repr { } } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe impl Send for KjOwn where T: Send {} - - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe impl Sync for KjOwn where T: Sync {} + // NO `Send`/`Sync` impls, deliberately. + // + // A `KjOwn` carries a type-erased `kj::Disposer*` alongside the object pointer, and + // dropping the `KjOwn` runs that disposer on whichever thread the drop happens on. A + // bound on `T` alone (e.g. `T: Send`) says nothing about the disposer: `kj::Own`s minted + // from `kj::Rc::toOwn()`/`kj::refcounted` (non-atomic refcount decrement), arena-backed + // objects, or any other custom disposer are NOT safe to destroy from another thread, and + // nothing at the bridge boundary guarantees disposer thread-safety. + // + // All current consumers keep `KjOwn`s on the KJ event-loop thread that created them, so + // no impls are needed. If a genuine cross-thread use case appears, it must come with an + // explicit opt-in mechanism that asserts the *disposer* is thread-safe (not just `T`); + // do not re-add blanket impls here. impl Deref for KjOwn { type Target = T; @@ -207,7 +217,8 @@ pub mod repr { } let this = std::ptr::from_mut::(self).cast::(); - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `this` points to this live `KjOwn` being dropped exactly once; the C++ + // `own$drop` shim invokes the type-erased `kj::Disposer` stored alongside `ptr`. unsafe { __drop(this); } diff --git a/src/rust/cxx/kj-rs/promise.c++ b/src/rust/cxx/kj-rs/promise.c++ index 4a03da29f1b..8382c8175a5 100644 --- a/src/rust/cxx/kj-rs/promise.c++ +++ b/src/rust/cxx/kj-rs/promise.c++ @@ -11,8 +11,8 @@ namespace kj_rs { static_assert(sizeof(OwnPromiseNode) == sizeof(uint64_t) * 1, "OwnPromiseNode size changed"); static_assert(alignof(OwnPromiseNode) == alignof(uint64_t) * 1, "OwnPromiseNode alignment changed"); -void own_promise_node_drop_in_place(OwnPromiseNode* node) { - kj::dtor(*node); +void own_promise_node_drop_in_place(OwnPromiseNode& node) { + kj::dtor(node); } } // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/promise.h b/src/rust/cxx/kj-rs/promise.h index b4d79d07ae9..fc5a858df4a 100644 --- a/src/rust/cxx/kj-rs/promise.h +++ b/src/rust/cxx/kj-rs/promise.h @@ -10,7 +10,7 @@ namespace kj_rs { using OwnPromiseNode = kj::_::OwnPromiseNode; -void own_promise_node_drop_in_place(OwnPromiseNode*); +void own_promise_node_drop_in_place(OwnPromiseNode&); namespace repr { diff --git a/src/rust/cxx/kj-rs/promise.rs b/src/rust/cxx/kj-rs/promise.rs index 1f221b8af3b..b7d6abda451 100644 --- a/src/rust/cxx/kj-rs/promise.rs +++ b/src/rust/cxx/kj-rs/promise.rs @@ -1,3 +1,8 @@ +//! FFI island (see crate-root `#![deny(unsafe_code)]`): `OwnPromiseNode`/`PromiseFuture` bridge — +//! `unsafe impl ExternType`, `unsafe extern "C"` unwrap callbacks, and `Pin` projection. A genuine +//! unsafe seam. +#![allow(unsafe_code)] + use std::ffi::c_void; use std::future::Future; use std::marker::PhantomData; @@ -20,16 +25,11 @@ pub struct OwnPromiseNode(*mut c_void /* kj::_::PromiseNode* */); // It is forgotten using `MaybeUninit` and its ownership passed over to c++ in `unwrap`. impl Drop for OwnPromiseNode { fn drop(&mut self) { - // Safety: - // 1. Pointer to self is non-null, and obviously points to valid memory. - // 2. We do not read or write to the OwnPromiseNode's memory, so there are no atomicity nor - // interleaved pointer/reference access concerns. - // - // https://doc.rust-lang.org/std/ptr/index.html#safety - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe { - crate::ffi::own_promise_node_drop_in_place(self); - } + // `own_promise_node_drop_in_place` placement-destructs the node behind `self`. The + // borrow is valid for the call; the value is only logically dead afterwards, inside + // this `drop`, and the inner `*mut c_void` has no drop glue, so there is no + // use-after-free or double-free. Expressed as a `&mut` binding, so no `unsafe` needed. + crate::ffi::own_promise_node_drop_in_place(self); } } @@ -141,12 +141,36 @@ impl KjPromise for CallbacksFuture { // unwrap will take over node ownership let node = ManuallyDrop::new(node); - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `node.0` is a live `OwnPromiseNode` whose ownership the callback takes over + // (wrapped in `ManuallyDrop` so we don't also drop it); `ret` is valid, suitably-aligned + // uninitialized storage for `Output`, which the callback initializes on the success path. unsafe { (callbacks.unwrap)(node.0, ret.as_mut_ptr().cast::()).into_result() }?; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the `?` above propagated any error, so on this path the callback reported + // success and therefore initialized `ret`. Ok(unsafe { ret.assume_init() }) } } -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. -unsafe impl Send for CallbacksFuture {} +// No `unsafe impl Send for CallbacksFuture`, deliberately. +// +// `CallbacksFuture` is only ever wrapped in `PromiseFuture`, whose `PromiseAwaiter` holds an +// `Option` (a raw pointer, hence `!Send`), so the composed future is `!Send` +// regardless. The bridged async machinery is confined to the KJ event-loop thread and `spawn` is +// `spawn_local`-backed (no `Send` requirement), so nothing needs a `Send` impl. Asserting the +// wrapper stays `!Send` locks that in. +#[cfg(test)] +mod send_guards { + use static_assertions::assert_not_impl_any; + + use super::CallbacksFuture; + use super::PromiseFuture; + + // The raw `*mut c_void` node makes this `!Send`/`!Sync` on its own; guard against a future + // hand-written impl silently introducing cross-thread transfer of a KJ promise node. + assert_not_impl_any!(CallbacksFuture: Send, Sync); + + // After its first poll, `PromiseFuture`'s embedded awaiter memory is self-referential and + // event-loop-linked (see `PromiseAwaiter::_pinned`); it must stay `!Unpin` so safe code + // cannot move it between polls (`&mut`-based awaits require `Unpin`). + assert_not_impl_any!(PromiseFuture>: Unpin); +} diff --git a/src/rust/cxx/kj-rs/refcount.rs b/src/rust/cxx/kj-rs/refcount.rs index 0489a75884d..ebe88bcb7e1 100644 --- a/src/rust/cxx/kj-rs/refcount.rs +++ b/src/rust/cxx/kj-rs/refcount.rs @@ -1,4 +1,8 @@ //! Module for both [`KjRc`] and [`KjArc`], since they're nearly identical types +//! +//! FFI island (see crate-root `#![deny(unsafe_code)]`): `KjRc`/`KjArc` mirror `kj::Rc`/`kj::Arc` — +//! `unsafe impl Send/Sync`, `extern "C"` refcount ops, and `Pin` projection. A genuine unsafe seam. +#![allow(unsafe_code)] use static_assertions::assert_eq_align; use static_assertions::assert_eq_size; @@ -30,10 +34,26 @@ pub mod repr { ptr: NonNull, } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe impl Send for KjArc where T: Send {} - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe impl Sync for KjArc where T: Sync {} + // Safety: `KjArc` mirrors `std::sync::Arc`'s thread-safety contract, and therefore + // requires the same `T: Send + Sync` bound for both `Send` and `Sync`: + // + // - `T: Sync` is required because clones can be sent to other threads, giving multiple + // threads concurrent `&T` access to the same pointee. + // - `T: Send` is required because the last `KjArc` to drop destroys the pointee on + // whichever thread it happens to live on, effectively transferring ownership of `T` + // to that thread. (Likewise, `get_mut()` can hand out exclusive access on any thread.) + // + // The reference count itself is managed on the C++ side by `kj::AtomicRefcounted` + // (atomic increments/decrements; the bridge's clone/drop shims require the pointee to be + // atomic-refcounted), so concurrent clone/drop of separate handles is safe once `T` + // satisfies the bounds above. + // + // A weaker `Send where T: Send` bound would be unsound: with `T: Send + !Sync`, cloning and + // sending a clone yields concurrent `&T` on two threads, so both impls require `T: Send + Sync`. + unsafe impl Send for KjArc where T: Send + Sync {} + // SAFETY: see the `Send` impl above — `KjArc` mirrors `std::sync::Arc`'s `T: Send + Sync` + // contract for `Sync` for the same reasons. + unsafe impl Sync for KjArc where T: Send + Sync {} impl KjRc { #[must_use] @@ -43,7 +63,8 @@ pub mod repr { fn __is_shared(this: *const c_void) -> bool; } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjRc` (`&self`), so its `*const c_void` refcounted + // pointer is valid for the C++ `is_shared` query. unsafe { __is_shared(std::ptr::from_ref(self).cast::()) } } @@ -73,7 +94,8 @@ pub mod repr { fn __drop(this: *mut c_void); } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjRc` being dropped exactly once; the C++ `drop` shim + // releases its refcount handle. unsafe { __drop(std::ptr::from_mut(self).cast::()); } @@ -88,7 +110,8 @@ pub mod repr { fn __is_shared(this: *const c_void) -> bool; } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjArc` (`&self`), so its `*const c_void` refcounted + // pointer is valid for the C++ `is_shared` query. unsafe { __is_shared(std::ptr::from_ref(self).cast::()) } } @@ -140,7 +163,8 @@ pub mod repr { } let mut ret = std::mem::MaybeUninit::::uninit(); - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjRc`; the C++ `clone` shim bumps the refcount and + // initializes `ret` with a valid `KjRc`, so `assume_init` is sound afterwards. unsafe { __clone( std::ptr::from_ref(self).cast::(), @@ -159,7 +183,8 @@ pub mod repr { } let mut ret = std::mem::MaybeUninit::::uninit(); - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjArc`; the C++ `clone` shim bumps the atomic refcount + // and initializes `ret` with a valid `KjArc`, so `assume_init` is sound afterwards. unsafe { __clone( std::ptr::from_ref(self).cast::(), @@ -177,7 +202,8 @@ pub mod repr { fn __drop(this: *mut c_void); } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjArc` being dropped exactly once; the C++ `drop` shim + // releases its refcount handle. unsafe { __drop(std::ptr::from_mut(self).cast::()); } diff --git a/src/rust/cxx/kj-rs/tests/BUILD.bazel b/src/rust/cxx/kj-rs/tests/BUILD.bazel index 204d928f8f1..b513baa6bf7 100644 --- a/src/rust/cxx/kj-rs/tests/BUILD.bazel +++ b/src/rust/cxx/kj-rs/tests/BUILD.bazel @@ -31,6 +31,7 @@ rust_library( # TODO(cleanup): Why isn't :cxx transitive? "//src/rust/cxx", "//src/rust/cxx/kj-rs", + "@crates_vendor//:static_assertions", ], ) @@ -139,10 +140,31 @@ wd_cc_library( ) cc_test( - name = "linked-group-test", + name = "shared-event-test", size = "small", srcs = [ - "linked-group-test.c++", + "shared-event-test.c++", + ], + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + deps = [ + "//src/rust/cxx/kj-rs", + "//src/rust/cxx/third-party:runtime", + "@capnp-cpp//src/kj:kj-test", + ], +) + +cc_test( + name = "neutralize-waker-test", + size = "small", + srcs = [ + "neutralize-waker-test.c++", ], linkstatic = select({ "@platforms//os:windows": True, diff --git a/src/rust/cxx/kj-rs/tests/awaitables-cc-test.c++ b/src/rust/cxx/kj-rs/tests/awaitables-cc-test.c++ index 22e22776ad7..535ed09e277 100644 --- a/src/rust/cxx/kj-rs/tests/awaitables-cc-test.c++ +++ b/src/rust/cxx/kj-rs/tests/awaitables-cc-test.c++ @@ -8,6 +8,16 @@ #include +// Raw-RustFuture test helpers, defined in tests/lib.rs and tests/test_futures.rs. The +// bridge's generated `async fn` shims always apply RustFuture's eager-by-default +// kj::Promise conversion, so tests that need a *cold* promise receive the not-yet-converted +// RustFuture through these and call `.lazily()` (kj-rs/future.h) themselves. +extern "C" { +void kj_rs_demo_lazy_side_effect_future(::kj_rs::repr::RustFuture* out); +void kj_rs_demo_lazy_future_awaiting_cancellable_promise(::kj_rs::repr::RustFuture* out); +void kj_rs_demo_work_before_poll(uint64_t* target, ::kj_rs::repr::RustFuture* out); +} + namespace kj_rs_demo { namespace { @@ -168,6 +178,57 @@ KJ_TEST(".awaiting a Promise from Rust can produce an Err Result") { waitScope); } +KJ_TEST("a panicking bridged async fn surfaces as a catchable kj::Exception, not an abort") { + // Unwind protection in the RustFuture vtable (kj-rs/future.rs): panics escaping poll() are + // converted into errored completions, mirroring the sync bridge's panic -> kj::Exception + // conversion. Before that protection, any of these would abort the process (unwinding out + // of an extern "C" fn). + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + { + // Fallible future, panic on first poll. + auto exception = KJ_ASSERT_NONNULL( + kj::runCatchingExceptions([&]() { new_panicking_future_void().wait(waitScope); })); + KJ_EXPECT(exception.getDescription().contains("bridged future panicked on purpose"), + exception.getDescription()); + } + + { + // Infallible future: the promise can still reject (kj::Promise always carries an + // exception channel even when the Rust signature is infallible). + auto exception = KJ_ASSERT_NONNULL(kj::runCatchingExceptions( + [&]() { new_panicking_infallible_future_void().wait(waitScope); })); + KJ_EXPECT(exception.getDescription().contains("bridged infallible future panicked on purpose"), + exception.getDescription()); + } + + { + // Panic after a suspension point: exercises the event-loop-driven poll path (not the + // eager creation-time poll). + auto exception = KJ_ASSERT_NONNULL(kj::runCatchingExceptions( + [&]() { new_panicking_after_await_future_void().wait(waitScope); })); + KJ_EXPECT(exception.getDescription().contains("panicked after a suspension point"), + exception.getDescription()); + } + + // A panicking bridged future can also be caught from a KJ coroutine. + []() -> kj::Promise { + kj::Maybe maybeException; + try { + co_await new_panicking_future_void(); + } catch (...) { + maybeException = kj::getCaughtExceptionAsKj(); + } + auto& exception = KJ_ASSERT_NONNULL(maybeException, "should have thrown"); + KJ_EXPECT(exception.getDescription().contains("bridged future panicked on purpose"), + exception.getDescription()); + }().wait(waitScope); + + // The loop is still healthy after the panics: run a normal future to completion. + []() -> kj::Promise { co_await new_ready_future_void(); }().wait(waitScope); +} + KJ_TEST("Rust can await Promise") { kj::EventLoop loop; kj::WaitScope waitScope(loop); @@ -188,25 +249,133 @@ KJ_TEST("C++ can receive asynchronous wakes after poll()") { kj::WaitScope waitScope(loop); auto promise = new_threaded_delay_future_void(); - // It's not ready yet. + // It's not ready yet: the future stashed a clone of its waker and returned Pending. KJ_EXPECT(!promise.poll(waitScope)); - // But later it is. + // Wake the stashed waker on the loop thread; this arms the FuturePollEvent so the next poll + // completes. Exercises a cloned-waker wake that arrives after poll() has already returned. + wake_delayed_future(); promise.wait(waitScope); } +KJ_TEST("Waker woken from another thread") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + // The future clones its waker into a spawned std::thread (one with no KJ event loop), which + // wakes it from there ~10ms later. `std::task::Waker` is Send + Sync, so this is plain safe + // Rust; the bridge must deliver the foreign-thread wake to this loop (waker.h's cross-thread + // fulfiller path) rather than losing it -- this wait() hangs if it does. + new_cross_thread_wake_future_void().wait(waitScope); + + // Again, but with the loop parked in wait() the whole time (no poll()-first warmup), so the + // wake is guaranteed to arrive while the loop sleeps rather than racing the first poll. + new_cross_thread_wake_future_void().wait(waitScope); +} + +KJ_TEST("Waker woken from another thread across multiple polls") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + // Four rounds of pending -> foreign-thread wake -> re-poll. Each round consumes the cell's + // cross-thread fulfiller, which the next poll must renew (waker.h); if renewal ever fails, + // a later round's wake is lost and this wait() hangs. + new_multi_round_cross_thread_wake_future_void().wait(waitScope); +} + +KJ_TEST("Retained waker woken from another thread, before and after future destruction") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + { + // Alive case: the future stashes a waker clone on first poll; wake it from a joined foreign + // thread, after poll() has returned. The wake must be delivered (the wait completes). + auto promise = new_threaded_delay_future_void(); + KJ_EXPECT(!promise.poll(waitScope)); + wake_stashed_waker_from_background_thread(); + promise.wait(waitScope); + } + + { + // Dead case: destroy the future (and its FuturePollEvent) while the stashed clone lives on, + // then wake from a foreign thread. The neutralized cell must make this a safe no-op — under + // ASan this is the use-after-free regression for waking a freed event cross-thread. + auto promise = new_threaded_delay_future_void(); + KJ_EXPECT(!promise.poll(waitScope)); + { auto dropped = kj::mv(promise); } + wake_stashed_waker_from_background_thread(); + } +} + KJ_TEST("Work before poll") { kj::EventLoop loop; kj::WaitScope waitScope(loop); uint64_t val = 0; - // It should be possible for rust function to do work before returning the future - // even if we don't poll or cancel it. - auto promise = work_before_poll(val); + // It should be possible for a Rust function to do work before returning the future + // even if we don't poll or cancel it. The future panics if polled, so it is converted + // with RustFuture::lazily() (the eager-by-default conversion polls at creation); this + // also proves cold promises really are never polled unawaited. + ::kj_rs::repr::RustFuture fut; + kj_rs_demo_work_before_poll(&val, &fut); + auto promise = fut.lazily(); KJ_EXPECT(val == 42); } +// ======================================================================================= +// Eager-by-default vs RustFuture::lazily(): bridged async fns surface as *eager* +// kj::Promises (polled to their first suspension at creation, like a KJ coroutine); +// `.lazily()` is the C++-side escape hatch that restores the cold future. + +KJ_TEST("bridged async fns are eager by default: the body runs at promise creation") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + reset_side_effect_counter(); + { + auto promise = new_side_effect_future_void(); + // No suspension points, so it ran to completion synchronously at creation, before any + // await or event-loop turn. + KJ_EXPECT(get_side_effect_counter() == 1); + } + KJ_EXPECT(get_side_effect_counter() == 1); +} + +KJ_TEST("RustFuture::lazily() opts out: nothing runs until the promise is first awaited") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + reset_side_effect_counter(); + ::kj_rs::repr::RustFuture fut; + kj_rs_demo_lazy_side_effect_future(&fut); + auto promise = fut.lazily(); + KJ_EXPECT(get_side_effect_counter() == 0); + + // Turning the event loop without awaiting the promise doesn't run it either. + kj::evalLater([]() {}).wait(waitScope); + KJ_EXPECT(get_side_effect_counter() == 0); + + promise.wait(waitScope); + KJ_EXPECT(get_side_effect_counter() == 1); +} + +KJ_TEST("eager promises still cancel on drop (never explicitly polled by the caller)") { + // Cancellation semantics are unchanged by eager evaluation: dropping the promise + // synchronously cancels the Rust future and the KJ promise it is awaiting. Here the + // caller never polls or awaits — creation alone started the future (suspending it at its + // .await), and destruction alone cancels it. + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + reset_cancellation_counter(); + { + auto promise = new_future_awaiting_cancellable_promise(); + KJ_EXPECT(get_cancellation_counter() == 0); + } + KJ_EXPECT(get_cancellation_counter() == 1); +} + // TODO(someday): More test cases. -// - Standalone ArcWaker tests. Ensure Rust calls ArcWaker destructor when we expect. +// - Standalone FutureWakerCell tests. Ensure Rust drops cloned waker cells when we expect. // - Throwing an exception from PromiseNode functions, including destructor. // ======================================================================================= @@ -219,11 +388,16 @@ KJ_TEST("Work before poll") { KJ_TEST("Cancellation: drop never-polled Rust future") { // Dropping a kj::Promise wrapping a Rust future that was never polled should not crash. Since the // future was never polled, the Rust async function body was never entered, so no sub-promises - // exist to cancel. + // exist to cancel. Uses a raw RustFuture converted with `.lazily()`: eager-by-default promises + // are always polled at least once (at creation), so only a cold promise can reach this path. kj::EventLoop loop; kj::WaitScope waitScope(loop); - { auto promise = new_future_awaiting_cancellable_promise(); } + { + ::kj_rs::repr::RustFuture fut; + kj_rs_demo_lazy_future_awaiting_cancellable_promise(&fut); + auto promise = fut.lazily(); + } } KJ_TEST("Cancellation: C++ dropping promise cancels Rust future's awaited KJ promise") { diff --git a/src/rust/cxx/kj-rs/tests/lib.rs b/src/rust/cxx/kj-rs/tests/lib.rs index 9a9aa7ae39d..519d031a1ed 100644 --- a/src/rust/cxx/kj-rs/tests/lib.rs +++ b/src/rust/cxx/kj-rs/tests/lib.rs @@ -13,24 +13,34 @@ mod test_own; mod test_refcount; use kj_rs::KjOwn; +use test_futures::get_side_effect_counter; +use test_futures::new_cross_thread_wake_future_void; use test_futures::new_drop_cancellable_promise_without_polling; use test_futures::new_error_handling_future_void_infallible; use test_futures::new_errored_future_void; use test_futures::new_future_awaiting_cancellable_promise; use test_futures::new_kj_errored_future_void; use test_futures::new_layered_ready_future_void; +use test_futures::new_multi_round_cross_thread_wake_future_void; use test_futures::new_naive_select_future_void; +use test_futures::new_panicking_after_await_future_void; +use test_futures::new_panicking_future_void; +use test_futures::new_panicking_infallible_future_void; use test_futures::new_pending_future_void; use test_futures::new_promise_i32_awaiting_future_void; use test_futures::new_ready_future_i32; use test_futures::new_ready_future_void; use test_futures::new_select_with_cancellation; +use test_futures::new_side_effect_future_void; use test_futures::new_threaded_delay_future_void; use test_futures::new_two_step_cancellable_future; use test_futures::new_waking_future_void; use test_futures::new_wrapped_waker_future_void; use test_futures::poll_and_stash_promise_future; +use test_futures::reset_side_effect_counter; use test_futures::unstash_and_await_promise_future; +use test_futures::wake_delayed_future; +use test_futures::wake_stashed_waker_from_background_thread; use test_maybe::take_maybe_own; use test_maybe::take_maybe_own_ret; use test_maybe::take_maybe_ref; @@ -260,6 +270,19 @@ pub mod ffi { async fn new_ready_future_shared_type() -> Shared; async fn new_waking_future_void(cloning_action: CloningAction, waking_action: WakingAction); async fn new_threaded_delay_future_void(); + // Woken from a spawned foreign thread (no KJ event loop): exercises the cross-thread + // `Waker` contract end to end. + async fn new_cross_thread_wake_future_void(); + // Foreign-thread wakes across multiple polls: exercises the cross-thread fulfiller's + // per-poll renewal. + async fn new_multi_round_cross_thread_wake_future_void(); + // Takes the waker stashed by `new_threaded_delay_future_void` and wakes it from a joined + // foreign thread — used both while the future is alive and after it has been destroyed + // (the neutralized cross-thread late wake). + fn wake_stashed_waker_from_background_thread(); + // Wakes the waker stashed by `new_threaded_delay_future_void`'s future, on the loop thread, + // to drive an asynchronous same-thread wake after poll() has returned. + fn wake_delayed_future(); async fn new_layered_ready_future_void() -> Result<()>; async fn new_naive_select_future_void() -> Result<()>; @@ -267,6 +290,12 @@ pub mod ffi { async fn new_errored_future_void() -> Result<()>; + // Unwind protection (kj-rs/future.rs): panics escaping a bridged future's poll() + // must become rejected promises (kj::Exception), not process aborts. + async fn new_panicking_future_void() -> Result<()>; + async fn new_panicking_infallible_future_void(); + async fn new_panicking_after_await_future_void() -> Result<()>; + async fn new_kj_errored_future_void() -> Result<()>; async fn new_error_handling_future_void_infallible(); @@ -275,7 +304,17 @@ pub mod ffi { async fn new_ready_future_i32(value: i32) -> Result; async fn new_pass_through_feature_shared() -> Shared; - async unsafe fn work_before_poll<'a>(target: &'a mut u64) -> Result<()>; + // Eager-by-default test helpers. The bridge's conversion polls the future + // synchronously to its first suspension at promise creation. The cold-promise + // (`RustFuture::lazily()`) counterparts bypass the bridge: they hand C++ the raw + // `RustFuture` through plain `extern "C"` helpers (see `test_futures.rs`). + #[expect(clippy::allow_attributes)] // Only called from C++ tests; #[expect(dead_code)] fails in builds where the lint does not fire + #[allow(dead_code)] + fn reset_side_effect_counter(); + #[expect(clippy::allow_attributes)] // Only called from C++ tests; #[expect(dead_code)] fails in builds where the lint does not fire + #[allow(dead_code)] + fn get_side_effect_counter() -> u64; + async fn new_side_effect_future_void(); // Cancellation test helpers. async fn new_future_awaiting_cancellable_promise() -> Result<()>; @@ -314,6 +353,20 @@ unsafe impl Send for ffi::OpaqueAtomicRefcountedClass {} // Safety: the test type follows the thread-safety contract of its C++ implementation. unsafe impl Sync for ffi::OpaqueAtomicRefcountedClass {} +// Compile-time thread-safety contracts (kj-rs/own.rs, kj-rs/refcount.rs): +// +// KjOwn is never Send/Sync: its type-erased kj disposer (kj::Rc, arena, ...) may not be +// thread-safe, regardless of T. +static_assertions::assert_not_impl_any!(KjOwn: Send, Sync); +static_assertions::assert_not_impl_any!(KjOwn: Send, Sync); +// KjArc matches std::sync::Arc: Send/Sync require T: Send + Sync... +static_assertions::assert_impl_all!(kj_rs::KjArc: Send, Sync); +// ...so a Send + !Sync payload (Cell) must make KjArc neither Send nor Sync (clones would +// otherwise hand concurrent &T to multiple threads). +static_assertions::assert_not_impl_any!(kj_rs::KjArc>: Send, Sync); +// KjRc (non-atomic refcount) must never be Send or Sync. +static_assertions::assert_not_impl_any!(kj_rs::KjRc: Send, Sync); + pub fn modify_own_return(mut own: KjOwn) -> KjOwn { own.pin_mut().set_data(72); own @@ -357,6 +410,33 @@ fn work_before_poll(target: &mut u64) -> impl Future> { } } +/// Hands C++ the raw, not-yet-converted future from [`work_before_poll`]. +/// +/// The returned future must never be polled (its body panics), so it cannot go through a +/// bridged `async fn` shim: those always apply `RustFuture`'s eager-by-default +/// `kj::Promise` conversion, which polls at creation. The C++ test converts it with +/// `RustFuture::lazily()` instead (see `awaitables-cc-test.c++`). +/// +/// # Safety +/// +/// `target` must be a valid, exclusive `u64` pointer that outlives the future; `out` must +/// point to uninitialized storage for one `::kj_rs::repr::RustFuture` (future.h), which the +/// caller takes ownership of. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn kj_rs_demo_work_before_poll<'a>( + target: &'a mut u64, + out: *mut kj_rs::repr::RustFuture<'a, ()>, +) { + let fut = kj_rs::repr::future(Box::pin(kj_rs::map_err( + work_before_poll(target), + file!(), + line!(), + ))); + // SAFETY: `out` points to uninitialized storage for one RustFuture, per this fn's + // `# Safety` contract. + unsafe { out.write(fut) }; +} + #[cfg(test)] mod tests { use crate::ffi; diff --git a/src/rust/cxx/kj-rs/tests/linked-group-test.c++ b/src/rust/cxx/kj-rs/tests/linked-group-test.c++ deleted file mode 100644 index 77a86ba2182..00000000000 --- a/src/rust/cxx/kj-rs/tests/linked-group-test.c++ +++ /dev/null @@ -1,364 +0,0 @@ -#include "kj-rs/linked-group.h" - -#include - -namespace kj_rs { -namespace { - -// Minimal concrete types for testing. -class TestGroup; -class TestObject; - -class TestGroup: public LinkedGroup { - public: - explicit TestGroup(int id): id(id) {} - - // Expose the protected member for testing. - using LinkedGroup::linkedObjects; - - int id; -}; - -class TestObject: public LinkedObject { - public: - explicit TestObject(int id): id(id) {} - - // Expose the protected member for testing. - using LinkedObject::linkedGroup; - - int id; -}; - -// --------------------------------------------------------------------------- -// Basic membership -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: object can join a group") { - TestGroup group(1); - TestObject object(10); - - KJ_EXPECT(group.linkedObjects().empty()); - KJ_EXPECT(object.linkedGroup().tryGet() == kj::none); - - object.linkedGroup().set(group); - - KJ_EXPECT(!group.linkedObjects().empty()); - KJ_EXPECT(group.linkedObjects().front().id == 10); - KJ_IF_SOME(g, object.linkedGroup().tryGet()) { - KJ_EXPECT(g.id == 1); - } else { - KJ_FAIL_EXPECT("expected group to be set"); - } -} - -KJ_TEST("LinkedGroup: object can leave a group") { - TestGroup group(1); - TestObject object(10); - - object.linkedGroup().set(group); - KJ_EXPECT(!group.linkedObjects().empty()); - - object.linkedGroup().set(kj::none); - KJ_EXPECT(group.linkedObjects().empty()); - KJ_EXPECT(object.linkedGroup().tryGet() == kj::none); -} - -KJ_TEST("LinkedGroup: object can switch groups") { - TestGroup group1(1); - TestGroup group2(2); - TestObject object(10); - - object.linkedGroup().set(group1); - KJ_EXPECT(!group1.linkedObjects().empty()); - KJ_EXPECT(group2.linkedObjects().empty()); - - object.linkedGroup().set(group2); - KJ_EXPECT(group1.linkedObjects().empty()); - KJ_EXPECT(!group2.linkedObjects().empty()); - KJ_IF_SOME(g, object.linkedGroup().tryGet()) { - KJ_EXPECT(g.id == 2); - } else { - KJ_FAIL_EXPECT("expected group to be set"); - } -} - -// --------------------------------------------------------------------------- -// Insertion order -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: objects are iterable in insertion order") { - TestGroup group(1); - TestObject a(1), b(2), c(3); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - c.linkedGroup().set(group); - - // Verify iteration order matches insertion order. - auto it = group.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it->id == 3); - ++it; - KJ_EXPECT(it == group.linkedObjects().end()); -} - -// --------------------------------------------------------------------------- -// Redundant set() is a no-op -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: redundant set() does not change position") { - TestGroup group(1); - TestObject a(1), b(2), c(3); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - c.linkedGroup().set(group); - - // Re-set b to the same group — order should be unchanged. - b.linkedGroup().set(group); - - auto it = group.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it->id == 3); - ++it; - KJ_EXPECT(it == group.linkedObjects().end()); -} - -// --------------------------------------------------------------------------- -// Lifetimes: object destroyed before group -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: destroying an object removes it from the group") { - TestGroup group(1); - TestObject a(1); - { - TestObject b(2); - b.linkedGroup().set(group); - a.linkedGroup().set(group); - - // Both present. - auto it = group.linkedObjects().begin(); - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it == group.linkedObjects().end()); - } - // b is destroyed; only a remains. - auto it = group.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it == group.linkedObjects().end()); -} - -// --------------------------------------------------------------------------- -// Lifetimes: group destroyed before objects -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: destroying a group unlinks all objects") { - TestObject a(1), b(2); - { - TestGroup group(1); - a.linkedGroup().set(group); - b.linkedGroup().set(group); - KJ_EXPECT(a.linkedGroup().tryGet() != kj::none); - KJ_EXPECT(b.linkedGroup().tryGet() != kj::none); - } - // Group destroyed — objects should no longer reference it. - KJ_EXPECT(a.linkedGroup().tryGet() == kj::none); - KJ_EXPECT(b.linkedGroup().tryGet() == kj::none); -} - -// --------------------------------------------------------------------------- -// Iteration and removal of the front element -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: removing the front element during iteration is safe") { - TestGroup group(1); - TestObject a(1), b(2), c(3); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - c.linkedGroup().set(group); - - // The header documents that removing the *front* element during iteration is valid. - kj::Vector collected; - for (auto it = group.linkedObjects().begin(); it != group.linkedObjects().end();) { - auto& obj = *it; - ++it; // advance before removing - collected.add(obj.id); - obj.linkedGroup().set(kj::none); - } - - KJ_EXPECT(collected.size() == 3); - KJ_EXPECT(collected[0] == 1); - KJ_EXPECT(collected[1] == 2); - KJ_EXPECT(collected[2] == 3); - KJ_EXPECT(group.linkedObjects().empty()); -} - -// --------------------------------------------------------------------------- -// Const access -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: const access to group's objects") { - TestGroup group(1); - TestObject a(1), b(2); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - - const TestGroup& cgroup = group; - KJ_EXPECT(!cgroup.linkedObjects().empty()); - KJ_EXPECT(cgroup.linkedObjects().front().id == 1); - - auto it = cgroup.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it == cgroup.linkedObjects().end()); -} - -KJ_TEST("LinkedGroup: const access to object's group") { - TestGroup group(1); - TestObject object(10); - - object.linkedGroup().set(group); - - const TestObject& cobject = object; - KJ_IF_SOME(g, cobject.linkedGroup().tryGet()) { - KJ_EXPECT(g.id == 1); - } else { - KJ_FAIL_EXPECT("expected group to be set"); - } -} - -// --------------------------------------------------------------------------- -// Empty state -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: default-constructed objects have no group") { - TestObject object(1); - KJ_EXPECT(object.linkedGroup().tryGet() == kj::none); -} - -KJ_TEST("LinkedGroup: default-constructed groups have no objects") { - TestGroup group(1); - KJ_EXPECT(group.linkedObjects().empty()); - KJ_EXPECT(group.linkedObjects().begin() == group.linkedObjects().end()); -} - -// --------------------------------------------------------------------------- -// Multiple objects, various removal patterns -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: removing a middle object leaves others intact") { - TestGroup group(1); - TestObject a(1), b(2), c(3); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - c.linkedGroup().set(group); - - b.linkedGroup().set(kj::none); - - auto it = group.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it->id == 3); - ++it; - KJ_EXPECT(it == group.linkedObjects().end()); -} - -KJ_TEST("LinkedGroup: removing all objects one by one") { - TestGroup group(1); - TestObject a(1), b(2), c(3); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - c.linkedGroup().set(group); - - a.linkedGroup().set(kj::none); - KJ_EXPECT(!group.linkedObjects().empty()); - - b.linkedGroup().set(kj::none); - KJ_EXPECT(!group.linkedObjects().empty()); - - c.linkedGroup().set(kj::none); - KJ_EXPECT(group.linkedObjects().empty()); -} - -// --------------------------------------------------------------------------- -// set(kj::none) on an unlinked object is a no-op -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: set(none) on an unlinked object is safe") { - TestObject object(1); - // Should not crash or assert. - object.linkedGroup().set(kj::none); - KJ_EXPECT(object.linkedGroup().tryGet() == kj::none); -} - -// --------------------------------------------------------------------------- -// Multiple groups are independent -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: multiple groups are independent") { - TestGroup g1(1), g2(2); - TestObject a(1), b(2), c(3), d(4); - - a.linkedGroup().set(g1); - b.linkedGroup().set(g1); - c.linkedGroup().set(g2); - d.linkedGroup().set(g2); - - // Verify g1 has {a, b}. - { - auto it = g1.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it == g1.linkedObjects().end()); - } - - // Verify g2 has {c, d}. - { - auto it = g2.linkedObjects().begin(); - KJ_EXPECT(it->id == 3); - ++it; - KJ_EXPECT(it->id == 4); - ++it; - KJ_EXPECT(it == g2.linkedObjects().end()); - } - - // Move b from g1 to g2. - b.linkedGroup().set(g2); - - { - auto it = g1.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it == g1.linkedObjects().end()); - } - { - auto it = g2.linkedObjects().begin(); - KJ_EXPECT(it->id == 3); - ++it; - KJ_EXPECT(it->id == 4); - ++it; - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it == g2.linkedObjects().end()); - } -} - -} // namespace -} // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/tests/neutralize-waker-test.c++ b/src/rust/cxx/kj-rs/tests/neutralize-waker-test.c++ new file mode 100644 index 00000000000..210a3e09eeb --- /dev/null +++ b/src/rust/cxx/kj-rs/tests/neutralize-waker-test.c++ @@ -0,0 +1,144 @@ +// Regression test: FutureWakerCell "neutralize-on-drop". +// +// The bridge's same-thread waker is a FutureWakerCell whose wake() arms the owning FuturePollEvent. +// The hazard: a waker clone is retained (e.g. handed to some sub-future) and its wake() fires AFTER +// the FuturePollEvent (and the boxed future it owns) has been torn down -- arming a freed +// kj::_::Event would be a use-after-free. +// +// This guards the refcounted-cell handling: the cell holds an Event*; the FuturePollEvent holds +// the strong ref and NULLS the cell on destruction (BEFORE the boxed future / sub-wakers drop). A +// retained clone that calls wake() after teardown observes null and is a safe no-op. +// +// Run under ASAN to confirm no UAF: +// bazel test //kj-rs/tests:neutralize-waker-test --config=asan + +#include +#include +#include +#include + +namespace kj_rs { +namespace { + +using kj::uint; + +// The refcounted cell shared between the FutureEvent and any retained Waker clones. The bridge is +// single-threaded (no cross-thread wakes), so a plain (non-atomic) kj::Refcounted with a bare +// Event* is sufficient -- no mutex/atomic needed. +class WakerCell: public kj::Refcounted { + public: + kj::_::Event* event = nullptr; + + bool observedNullOnWake = false; + uint wakeArmCount = 0; + + // Called by the FutureEvent's destructor to neutralize all outstanding waker clones. + void neutralize() { + event = nullptr; + } + + // The Waker's wake(): arm the FutureEvent, or no-op if it's been neutralized. + void wake() { + if (event != nullptr) { + event->armDepthFirst(); + ++wakeArmCount; + } else { + observedNullOnWake = true; // SAFE no-op: no arm of a freed Event, no UAF. + } + } +}; + +// Models the boxed Rust future (and its sub-wakers) owned inline by the FutureEvent. Its whole job +// here is to ASSERT the ordering requirement: by the time it is destroyed, the cell must already +// have been neutralized -- i.e. nulled BEFORE the boxed future / sub-wakers drop. +class BoxedFutureStandin { + public: + explicit BoxedFutureStandin(WakerCell& cell): cell(cell) {} + ~BoxedFutureStandin() noexcept(false) { + KJ_ASSERT(cell.event == nullptr, + "ordering violation: cell must be neutralized BEFORE the boxed future drops"); + } + + private: + WakerCell& cell; +}; + +// Stand-in for the FutureEvent: a kj Event that owns the WakerCell strong ref and the boxed future. +class FutureEventStandin final: public kj::_::Event { + public: + explicit FutureEventStandin(kj::Rc cellParam, kj::SourceLocation location = {}) + : Event(location), + cell(kj::mv(cellParam)), + boxedFuture(*cell) { + cell->event = this; + } + + ~FutureEventStandin() noexcept(false) { + // ORDERING: neutralize the cell FIRST (destructor body runs before member subobjects are + // destroyed). Member destruction order is reverse-declaration: `boxedFuture` then `cell`. + // So when boxedFuture's dtor asserts, the cell is already nulled. + cell->neutralize(); + } + + uint fireCount = 0; + void traceEvent(kj::_::TraceBuilder&) override {} + + private: + void fire() override { + ++fireCount; + } + + kj::Rc cell; // declared first -> destroyed LAST + BoxedFutureStandin boxedFuture; // declared second -> destroyed FIRST +}; + +KJ_TEST("FutureWaker neutralize-on-drop: retained clone wake() is a safe no-op after teardown") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + auto cell = kj::rc(); + auto retainedClone = cell.addRef(); // a second handle that will OUTLIVE the FutureEvent. + + auto event = kj::heap(kj::mv(cell)); + + // Sanity: while the FutureEvent is alive, a wake via the retained clone arms it. + retainedClone->wake(); + waitScope.poll(); + KJ_EXPECT(event->fireCount == 1); + KJ_EXPECT(retainedClone->wakeArmCount == 1); + KJ_EXPECT(!retainedClone->observedNullOnWake); + + // TEARDOWN: destroy the FutureEvent. Its dtor neutralizes the cell (asserted to happen before + // boxedFuture drops). The retained clone keeps the cell object itself alive. + event = nullptr; + + // The retained clone's wake() now observes a null Event* -> SAFE no-op. Without neutralize-on- + // drop this would arm a freed Event (UAF -- caught by ASAN under --config=asan). + retainedClone->wake(); + waitScope.poll(); // must not fire anything, must not crash + KJ_EXPECT(retainedClone->observedNullOnWake); + KJ_EXPECT(retainedClone->wakeArmCount == 1); // unchanged: the post-teardown wake armed nothing +} + +KJ_TEST("FutureWaker neutralize-on-drop: multiple retained clones all neutralized together") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + auto cell = kj::rc(); + auto cloneA = cell.addRef(); + auto cloneB = cell.addRef(); + + auto event = kj::heap(kj::mv(cell)); + event = nullptr; // teardown + + // Both retained clones observe null; neither arms a freed Event. + cloneA->wake(); + cloneB->wake(); + waitScope.poll(); + KJ_EXPECT(cloneA->observedNullOnWake); + KJ_EXPECT(cloneB->observedNullOnWake); + KJ_EXPECT(cloneA->wakeArmCount == 0); +} + +} // namespace +} // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/tests/shared-event-test.c++ b/src/rust/cxx/kj-rs/tests/shared-event-test.c++ new file mode 100644 index 00000000000..15da02bbfbf --- /dev/null +++ b/src/rust/cxx/kj-rs/tests/shared-event-test.c++ @@ -0,0 +1,321 @@ +// Regression test: one shared kj Event as the onReady target of MANY concurrent pending nodes. +// +// The bridge registers the SAME FuturePollEvent as the onReady target of every kj::Promise a Rust +// future is `.await`ing (many nodes -> one event, kj's native mechanism). This guards the +// kj Event/onReady properties that relies on: +// (a) idempotent arming: fulfilling several nodes in the SAME turn arms the one event once +// (Event::armDepthFirst's `if (prev == nullptr)` guard, async.c++:2201), +// (b) re-poll: fire() re-polls, ready nodes are consumed, not-yet-ready nodes stay registered +// and re-arm the event when THEY later resolve, +// (c) arm-while-firing: fulfilling a node DURING the shared event's own fire() safely re-arms +// it for the next turn (no "Promise callback destroyed itself" abort async.c++:2188, no +// lost wake) -- because turn() unlinks the event before fire(), so prev==nullptr during +// fire and armDepthFirst re-inserts it. +// +// Two test groups: +// GROUP A -- "pure kj primitive": N PromiseNodes each call node->onReady(&sharedEvent) DIRECTLY. +// Proves kj's OnReadyEvent::arm() -> Event::armDepthFirst() coalescing on one shared +// target. Readiness bookkeeping is test-driven (kj exposes no per-node readiness +// query): the point of Group A is the wake/arm coalescing path. +// GROUP B -- "faithful future model": each leaf is a trivial ~6-line per-leaf arm Event whose +// fire() sets ready=true and arms ONE shared re-poll event, and readiness here is +// genuinely kj-detected. + +#include +#include +#include +#include +#include + +namespace kj_rs { +namespace { + +using kj::uint; +using kj::_::Event; +using kj::_::ExceptionOr; +using kj::_::OwnPromiseNode; +using kj::_::PromiseNode; +using kj::_::Void; + +// A stable slot holding a fulfiller/node pair so we can call setSelfPointer() and later get(). +struct Slot { + kj::Own> fulfiller; + OwnPromiseNode node; + bool fulfilled = false; // test-side bookkeeping (Group A) + bool consumed = false; + + static kj::Own make() { + auto paf = kj::newPromiseAndFulfiller(); + auto self = kj::heap(); + self->fulfiller = kj::mv(paf.fulfiller); + self->node = PromiseNode::from(kj::mv(paf.promise)); + self->node->setSelfPointer(&self->node); + return self; + } + + void consume() { + ExceptionOr output; + node->get(output); + KJ_ASSERT(output.exception == kj::none); + consumed = true; + } +}; + +// ======================================================================================= +// GROUP A -- N nodes -> ONE shared event as direct onReady target. + +class SharedRepollEvent final: public Event { + public: + SharedRepollEvent(kj::ArrayPtr> slots, kj::SourceLocation location = {}) + : Event(location), + slots(slots) {} + + uint fireCount = 0; + uint consumedCount = 0; + + // If set, invoked once during the NEXT fire() -- models a sub-future resolving mid-poll. + kj::Function* armWhileFiringHook = nullptr; + + void traceEvent(kj::_::TraceBuilder&) override {} + + private: + kj::ArrayPtr> slots; + + void fire() override { + ++fireCount; + + // Model "the rust future re-polls all its sub-futures": consume every ready (fulfilled), + // not-yet-consumed node; leave the rest registered. + for (auto& slot: slots) { + if (slot->fulfilled && !slot->consumed) { + slot->consume(); + ++consumedCount; + } + } + + // ARM-WHILE-FIRING: fulfill another node during our own fire(). Its onReady points at us; + // arming us here must be safe (we were unlinked before fire, so prev==nullptr -> re-inserts). + if (armWhileFiringHook != nullptr) { + auto* hook = armWhileFiringHook; + armWhileFiringHook = nullptr; + (*hook)(); + } + // NOTE: we intentionally do NOT self-destruct; a FutureEvent lives until its future resolves. + } +}; + +KJ_TEST("SharedEvent(A): fulfilling several nodes in one turn arms the shared event idempotently") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + auto slots = kj::heapArray>(4); + for (auto& s: slots) s = Slot::make(); + + SharedRepollEvent shared(slots); + for (auto& s: slots) s->node->onReady(&shared); + + // Fulfill 3 of the 4 in the SAME turn (no loop run in between). + for (uint i: {0u, 1u, 2u}) { + slots[i]->fulfiller->fulfill(); + slots[i]->fulfilled = true; + } + + // Exactly one fire should happen: the 3 same-turn arms coalesced into a single armed event. + // (If they had NOT coalesced, fireCount would be 3.) + waitScope.poll(); + KJ_EXPECT(shared.fireCount == 1); + KJ_EXPECT(shared.consumedCount == 3); + + // (b) The 4th node was never fulfilled: it stays registered on the shared event. Fulfilling it + // now must re-arm the (idle) shared event and fire again. + slots[3]->fulfiller->fulfill(); + slots[3]->fulfilled = true; + waitScope.poll(); + KJ_EXPECT(shared.fireCount == 2); + KJ_EXPECT(shared.consumedCount == 4); +} + +KJ_TEST("SharedEvent(A): arm-while-firing re-arms safely for the next turn (no lost wake)") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + auto slots = kj::heapArray>(3); + for (auto& s: slots) s = Slot::make(); + + SharedRepollEvent shared(slots); + for (auto& s: slots) s->node->onReady(&shared); + + // During the shared event's FIRST fire(), fulfill slot[2]. Its onReady == &shared, so this arms + // `shared` while `shared` is mid-fire. Must not abort; must re-arm for the next turn. + kj::Function hook = [&]() { + slots[2]->fulfiller->fulfill(); + slots[2]->fulfilled = true; + }; + shared.armWhileFiringHook = &hook; + + // Kick off: fulfill slots 0 and 1 in this turn. + for (uint i: {0u, 1u}) { + slots[i]->fulfiller->fulfill(); + slots[i]->fulfilled = true; + } + + // A single poll() drains: fire #1 consumes 0,1 and (via the hook) fulfills slot[2], which arms + // `shared` mid-fire; fire #2 (the re-arm) consumes slot[2]. Proves no abort + no lost wake. + waitScope.poll(); + KJ_EXPECT(shared.fireCount == 2); + KJ_EXPECT(shared.consumedCount == 3); +} + +// ======================================================================================= +// GROUP B -- faithful future model: trivial per-leaf arm events + ONE shared re-poll event. + +class FutureEvent; + +// A leaf `.await` of a kj::Promise. ~6 lines of real logic: on its promise's readiness, record +// ready and arm the shared FutureEvent. +class LeafAwaiter final: public Event { + public: + LeafAwaiter(OwnPromiseNode nodeParam, FutureEvent& futureEvent, kj::SourceLocation location = {}); + ~LeafAwaiter() noexcept(false) { + node = nullptr; + } + + bool ready = false; + bool consumed = false; + + void consume() { + KJ_ASSERT(ready && !consumed); + ExceptionOr output; + node->get(output); + KJ_ASSERT(output.exception == kj::none); + consumed = true; + } + + void traceEvent(kj::_::TraceBuilder&) override {} + + private: + OwnPromiseNode node; + FutureEvent& futureEvent; + void fire() override; // defined after FutureEvent +}; + +// The ONE event that "is the future": its fire() re-polls the whole future (== consume any ready +// leaves, leave the rest). Every leaf arms THIS single event. +class FutureEvent final: public Event { + public: + FutureEvent(kj::SourceLocation location = {}): Event(location) {} + + uint fireCount = 0; + + // If set, invoked once during the NEXT fire() -- models a sub-future resolving mid-poll. + kj::Function* armWhileFiringHook = nullptr; + + void addLeaf(kj::Own leaf) { + leaves.add(kj::mv(leaf)); + } + + bool allConsumed() const { + for (auto& l: leaves) + if (!l->consumed) return false; + return true; + } + uint consumedCount() const { + uint n = 0; + for (auto& l: leaves) + if (l->consumed) ++n; + return n; + } + + void traceEvent(kj::_::TraceBuilder&) override {} + + private: + kj::Vector> leaves; + + void fire() override { + ++fireCount; + for (auto& l: leaves) { + if (l->ready && !l->consumed) l->consume(); + } + if (armWhileFiringHook != nullptr) { + auto* hook = armWhileFiringHook; + armWhileFiringHook = nullptr; + (*hook)(); // fulfills another leaf -> its LeafAwaiter arms `this` mid-fire + } + } +}; + +LeafAwaiter::LeafAwaiter(OwnPromiseNode nodeParam, FutureEvent& fe, kj::SourceLocation location) + : Event(location), + node(kj::mv(nodeParam)), + futureEvent(fe) { + node->setSelfPointer(&node); + node->onReady(this); +} + +void LeafAwaiter::fire() { + ready = true; + futureEvent.armDepthFirst(); // arm the ONE shared future event (idempotent across leaves) +} + +KJ_TEST( + "SharedEvent(B): many leaves arm one FutureEvent; ready consumed, pending stay registered") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + FutureEvent future; + kj::Vector>> fulfillers; + + for (uint i = 0; i < 4; ++i) { + auto paf = kj::newPromiseAndFulfiller(); + fulfillers.add(kj::mv(paf.fulfiller)); + future.addLeaf(kj::heap(PromiseNode::from(kj::mv(paf.promise)), future)); + } + + // Fulfill 3 leaves in one turn -> 3 leaf Events fire (each arms `future`), then `future` fires + // ONCE (coalesced). turn() runs the leaf events + the single future event. + for (uint i: {0u, 1u, 2u}) fulfillers[i]->fulfill(); + + loop.run(64); + KJ_EXPECT(future.fireCount >= 1); + KJ_EXPECT(future.consumedCount() == 3); + KJ_EXPECT(!future.allConsumed()); + + uint fireCountAfter3 = future.fireCount; + + // 4th leaf stays registered; fulfilling it re-arms the future event. + fulfillers[3]->fulfill(); + loop.run(64); + KJ_EXPECT(future.fireCount > fireCountAfter3); + KJ_EXPECT(future.allConsumed()); +} + +KJ_TEST( + "SharedEvent(B): leaf resolving during the future's fire re-arms the future (no lost wake)") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + FutureEvent future; + kj::Vector>> fulfillers; + + for (uint i = 0; i < 3; ++i) { + auto paf = kj::newPromiseAndFulfiller(); + fulfillers.add(kj::mv(paf.fulfiller)); + future.addLeaf(kj::heap(PromiseNode::from(kj::mv(paf.promise)), future)); + } + + // During the future's FIRST fire(), fulfill leaf 2. Its LeafAwaiter event will arm `future` + // while `future` is still mid-fire -> must safely re-arm for the next turn. + kj::Function hook = [&]() { fulfillers[2]->fulfill(); }; + future.armWhileFiringHook = &hook; + + // Kick off with leaves 0 and 1. + fulfillers[0]->fulfill(); + fulfillers[1]->fulfill(); + + loop.run(64); + KJ_EXPECT(future.fireCount >= 2); // at least: fire consuming 0/1, then the re-armed fire for 2 + KJ_EXPECT(future.allConsumed()); // leaf 2 (fulfilled during a fire) was not lost +} + +} // namespace +} // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/tests/test_futures.rs b/src/rust/cxx/kj-rs/tests/test_futures.rs index cb8213e6eba..c9307ac0cb1 100644 --- a/src/rust/cxx/kj-rs/tests/test_futures.rs +++ b/src/rust/cxx/kj-rs/tests/test_futures.rs @@ -25,6 +25,54 @@ pub async fn new_ready_future_void() { std::future::ready(()).await } +// Eager-by-default vs `RustFuture::lazily()` test helpers: they record when their body ran +// so the C++ driver can observe eager promises running at creation and `.lazily()` ones only +// when awaited. The bridge's generated shims always apply the eager conversion, so the +// `.lazily()` tests receive the raw `RustFuture` through the plain `extern "C"` helpers +// below instead of bridged `async fn` declarations. + +static SIDE_EFFECT_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + +pub fn reset_side_effect_counter() { + SIDE_EFFECT_COUNTER.store(0, std::sync::atomic::Ordering::SeqCst); +} + +pub fn get_side_effect_counter() -> u64 { + SIDE_EFFECT_COUNTER.load(std::sync::atomic::Ordering::SeqCst) +} + +/// Increments the side-effect counter when its body runs (first poll), then completes. +pub async fn new_side_effect_future_void() { + SIDE_EFFECT_COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst); +} + +/// Same body as [`new_side_effect_future_void`], but handed to C++ as a raw `RustFuture` +/// (via [`kj_rs_demo_lazy_side_effect_future`]) and converted with `RustFuture::lazily()`: +/// the C++ promise is cold, so the counter only moves once the promise is first awaited. +pub async fn new_lazy_side_effect_future_void() { + SIDE_EFFECT_COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst); +} + +/// Hands C++ the raw, not-yet-converted [`new_lazy_side_effect_future_void`] future. +/// +/// `awaitables-cc-test.c++` converts it with `RustFuture::lazily()` (future.h). Plain +/// `extern "C"` because the bridge's generated `async fn` shims always apply the +/// eager-by-default `kj::Promise` conversion before C++ ever sees the future. +/// +/// # Safety +/// +/// `out` must point to uninitialized storage for one `::kj_rs::repr::RustFuture` +/// (future.h), which the caller takes ownership of. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn kj_rs_demo_lazy_side_effect_future( + out: *mut kj_rs::repr::RustInfallibleFuture<'static, ()>, +) { + let fut = kj_rs::repr::infallible_future(Box::pin(new_lazy_side_effect_future_void())); + // SAFETY: `out` points to uninitialized storage for one RustFuture, per this fn's + // `# Safety` contract. + unsafe { out.write(fut) }; +} + struct WakingFuture { done: bool, cloning_action: CloningAction, @@ -64,6 +112,15 @@ fn do_cloned_wake(waker: Waker, waking_action: WakingAction) { } } +/// Runs `f` on a freshly spawned thread (one with no KJ event loop) and returns its result, +/// propagating panics. Used to exercise the cross-thread arms of the `Waker` contract. +fn on_background_thread(f: impl FnOnce() -> T + Send) -> T { + std::thread::scope(|scope| match scope.spawn(f).join() { + Ok(value) => value, + Err(payload) => std::panic::resume_unwind(payload), + }) +} + impl Future for WakingFuture { type Output = (); fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context) -> Poll<()> { @@ -102,44 +159,154 @@ pub async fn new_waking_future_void(cloning_action: CloningAction, waking_action WakingFuture::new(cloning_action, waking_action).await } -struct ThreadedDelayFuture { - handle: Option>, +// A future that, on its first poll, stashes a clone of its waker and returns Pending WITHOUT +// waking; a later call to `wake_delayed_future()` — made by the C++ driver on the event loop's own +// thread, after poll() has already returned — wakes the stashed clone, arming the FuturePollEvent +// so the next poll completes. This exercises an *asynchronous* wake (one that arrives after poll() +// returned, via a cloned waker) on the loop thread; `wake_stashed_waker_from_background_thread()` +// below reuses the stash to exercise the same shape from a foreign thread, including after the +// future has been destroyed. + +thread_local! { + static DELAYED_WAKER: std::cell::RefCell> = const { std::cell::RefCell::new(None) }; } -impl ThreadedDelayFuture { - fn new() -> Self { - Self { handle: None } - } +struct DelayedWakeFuture { + done: bool, } -/// Run a function, `f`, on a thread in the background and return its result. -fn on_background_thread(f: impl FnOnce() -> T + Send) -> T { - std::thread::scope(|scope| match scope.spawn(f).join() { - Ok(value) => value, - Err(payload) => std::panic::resume_unwind(payload), - }) +impl DelayedWakeFuture { + fn new() -> Self { + Self { done: false } + } } -impl Future for ThreadedDelayFuture { +impl Future for DelayedWakeFuture { type Output = (); fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context) -> Poll<()> { - if let Some(handle) = self.handle.take() { - let _ = handle.join(); + if self.done { return Poll::Ready(()); } + // Stash a clone of the waker for the C++ driver to wake later, on this same thread. + DELAYED_WAKER.with(|w| *w.borrow_mut() = Some(cx.waker().clone())); + self.done = true; + Poll::Pending + } +} - let waker = cx.waker(); - let waker = on_background_thread(|| waker.clone()); - self.handle = Some(std::thread::spawn(|| { - std::thread::sleep(std::time::Duration::from_millis(100)); +pub async fn new_threaded_delay_future_void() { + DelayedWakeFuture::new().await +} + +/// Wake the waker stashed by [`DelayedWakeFuture`]. Called by the C++ test driver on the event +/// loop thread, after the future's first poll has returned Pending. +pub fn wake_delayed_future() { + DELAYED_WAKER.with(|w| { + if let Some(waker) = w.borrow_mut().take() { waker.wake(); - })); + } + }); +} + +/// Take the waker stashed by [`DelayedWakeFuture`] and wake it from a spawned foreign thread +/// (joined before returning, so the wake has fully happened when this returns). The C++ driver +/// calls this either while the future is alive (the wake must arm its event, cross-thread) or +/// after the future has been destroyed (the wake must be a safe neutralized no-op — the +/// asan-visible regression for a freed `FuturePollEvent`). +pub fn wake_stashed_waker_from_background_thread() { + let waker = DELAYED_WAKER.with(|w| w.borrow_mut().take()); + if let Some(waker) = waker { + let handle = std::thread::spawn(move || waker.wake()); + let _ = handle.join(); + } +} + +/// A future woken from a foreign thread across MULTIPLE polls: each round hands a fresh waker +/// clone to a spawned thread, which wakes it from there; the round only completes once that wake +/// has actually been delivered (spurious polls stay Pending). Exercises the cross-thread +/// fulfiller's per-poll renewal (waker.h): every round's wake must be delivered, not just the +/// first. +struct MultiRoundCrossThreadWakeFuture { + rounds_left: u32, + /// `Some(flag)` while a round's foreign wake is in flight; the thread sets the flag (then + /// wakes), and the next poll that observes it completes the round. + pending_wake: Option>, +} + +impl Future for MultiRoundCrossThreadWakeFuture { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + if let Some(flag) = &self.pending_wake { + if !flag.load(std::sync::atomic::Ordering::SeqCst) { + // Spurious poll: this round's foreign wake hasn't been delivered yet. + return Poll::Pending; + } + self.pending_wake = None; + self.rounds_left -= 1; + } + if self.rounds_left == 0 { + return Poll::Ready(()); + } + // Start the next round: this poll's waker goes to a fresh foreign thread. + let flag = Arc::new(std::sync::atomic::AtomicBool::new(false)); + self.pending_wake = Some(Arc::clone(&flag)); + let waker = cx.waker().clone(); + std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_millis(5)); + flag.store(true, std::sync::atomic::Ordering::SeqCst); + waker.wake(); + }); Poll::Pending } } -pub async fn new_threaded_delay_future_void() { - ThreadedDelayFuture::new().await +pub async fn new_multi_round_cross_thread_wake_future_void() { + MultiRoundCrossThreadWakeFuture { + rounds_left: 4, + pending_wake: None, + } + .await +} + +/// A future woken FROM ANOTHER THREAD: its first poll clones the waker and hands it to a spawned +/// `std::thread`, which stores a flag and calls `wake()` from that thread — one with no KJ event +/// loop at all. Exercises the full cross-thread `Waker` contract (`Waker: Send + Sync`): the +/// clone crosses the thread boundary, the wake runs there, and so does the drop of that handle. +struct CrossThreadWakeFuture { + woken: Arc, + spawned: bool, +} + +impl Future for CrossThreadWakeFuture { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + if self.woken.load(std::sync::atomic::Ordering::SeqCst) { + return Poll::Ready(()); + } + if !self.spawned { + self.spawned = true; + let waker = cx.waker().clone(); + let woken = Arc::clone(&self.woken); + std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_millis(10)); + woken.store(true, std::sync::atomic::Ordering::SeqCst); + // The wake may even race poll() returning Pending; both orders must work (the + // cross-thread delivery coalesces into the next poll). + waker.wake(); + }); + } + Poll::Pending + } +} + +pub async fn new_cross_thread_wake_future_void() { + CrossThreadWakeFuture { + woken: Arc::new(std::sync::atomic::AtomicBool::new(false)), + spawned: false, + } + .await } pub async fn new_layered_ready_future_void() -> Result<()> { @@ -214,6 +381,26 @@ pub async fn new_errored_future_void() -> Result<()> { Err(std::io::Error::other("test error")) } +// Unwind-protection helpers (kj-rs/future.rs vtable): a panic escaping a bridged future's +// poll() must surface to C++ as a rejected kj::Promise carrying a kj::Exception, not a +// process abort. These panic at different points to cover both the fallible and infallible +// vtables and both the first-poll and event-loop-driven poll paths. + +pub async fn new_panicking_future_void() -> Result<()> { + panic!("bridged future panicked on purpose"); +} + +pub async fn new_panicking_infallible_future_void() { + panic!("bridged infallible future panicked on purpose"); +} + +pub async fn new_panicking_after_await_future_void() -> Result<()> { + crate::ffi::new_ready_promise_void() + .await + .expect("should not throw"); + panic!("bridged future panicked after a suspension point"); +} + pub async fn new_kj_errored_future_void() -> std::result::Result<(), cxx::KjError> { Err(cxx::KjError::new( cxx::KjExceptionType::Overloaded, @@ -258,6 +445,43 @@ pub async fn new_future_awaiting_cancellable_promise() -> Result<()> { Ok(()) } +/// Like [`new_future_awaiting_cancellable_promise`], but handed to C++ as a raw +/// `RustFuture` (via [`kj_rs_demo_lazy_future_awaiting_cancellable_promise`]) and converted +/// with `RustFuture::lazily()`, so the C++ promise really is never polled unless awaited — +/// the only way to exercise the "drop a never-polled future" path now that bridged promises +/// are eager by default. +pub async fn new_lazy_future_awaiting_cancellable_promise() -> Result<()> { + crate::ffi::new_cancellation_detecting_promise_void() + .await + .map_err(Error::other)?; + Ok(()) +} + +/// Hands C++ the raw, not-yet-converted [`new_lazy_future_awaiting_cancellable_promise`] +/// future. +/// +/// `awaitables-cc-test.c++` converts it with `RustFuture::lazily()` (future.h). Plain +/// `extern "C"` because the bridge's generated `async fn` shims always apply the +/// eager-by-default `kj::Promise` conversion before C++ ever sees the future. +/// +/// # Safety +/// +/// `out` must point to uninitialized storage for one `::kj_rs::repr::RustFuture` +/// (future.h), which the caller takes ownership of. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn kj_rs_demo_lazy_future_awaiting_cancellable_promise( + out: *mut kj_rs::repr::RustFuture<'static, ()>, +) { + let fut = kj_rs::repr::future(Box::pin(kj_rs::map_err( + new_lazy_future_awaiting_cancellable_promise(), + file!(), + line!(), + ))); + // SAFETY: `out` points to uninitialized storage for one RustFuture, per this fn's + // `# Safety` contract. + unsafe { out.write(fut) }; +} + /// Two-step future: the first step completes normally, and the second step awaits a /// cancellation-detecting promise that never resolves. After one poll, the future will have /// advanced past step 1 and be suspended at step 2. diff --git a/src/rust/cxx/kj-rs/tests/test_own.rs b/src/rust/cxx/kj-rs/tests/test_own.rs index 61d0235b08c..292c3950ed7 100644 --- a/src/rust/cxx/kj-rs/tests/test_own.rs +++ b/src/rust/cxx/kj-rs/tests/test_own.rs @@ -201,20 +201,10 @@ pub mod tests { assert!(!debug_str.is_empty()); } - #[test] - fn test_own_send_between_threads() { - use std::thread; - - let own = ffi::cxx_kj_own(); - let handle = thread::spawn(move || { - // Own should be Send, so this should work - assert_eq!(own.get_data(), 42); - own - }); - - let returned_own = handle.join().unwrap(); - assert_eq!(returned_own.get_data(), 42); - } + // NOTE: `KjOwn` is deliberately neither `Send` nor `Sync`, even for `T: Send + Sync` + // (like `OpaqueCxxClass` here): the `KjOwn` carries a type-erased C++ disposer which may + // not be thread-safe (kj::Rc, arena allocations, ...). The negative is asserted at compile + // time in lib.rs (assert_not_impl_any). #[test] fn test_own_concurrent_creation() { @@ -254,49 +244,4 @@ pub mod tests { let expected: Vec = (0..num_threads).map(|i| i as u64 * 100).collect(); assert_eq!(results, expected); } - - // This is one test generated by Claude. I am unsure it sufficiently tests multithreading. - #[test] - fn test_own_stress_multithreaded() { - use std::sync::mpsc; - use std::thread; - - let (tx, rx) = mpsc::channel(); - let num_threads: u64 = 12; - let items_per_thread: u64 = 100; - - for thread_id in 0..num_threads { - let tx_clone = tx.clone(); - thread::spawn(move || { - for i in 0..items_per_thread { - let mut own = ffi::cxx_kj_own(); - let value = thread_id * items_per_thread + i; - own.pin_mut().set_data(value); - - // Send the Own across thread boundary - tx_clone.send(own).unwrap(); - } - }); - } - drop(tx); // Close the sending side - - // Collect all Owns from all threads - let mut received_owns = Vec::new(); - while let Ok(own) = rx.recv() { - received_owns.push(own); - } - - // Verify we received the expected number - assert_eq!( - received_owns.len(), - (num_threads * items_per_thread) as usize - ); - - // Verify all values are correct - let mut values: Vec = received_owns.iter().map(|own| own.get_data()).collect(); - values.sort_unstable(); - - let expected: Vec = (0..(num_threads * items_per_thread)).collect(); - assert_eq!(values, expected); - } } diff --git a/src/rust/cxx/kj-rs/waker.c++ b/src/rust/cxx/kj-rs/waker.c++ index 0c993fdaf7e..387ef8a8aa2 100644 --- a/src/rust/cxx/kj-rs/waker.c++ +++ b/src/rust/cxx/kj-rs/waker.c++ @@ -1,150 +1,51 @@ #include "waker.h" -#include +#include "awaiter.h" namespace kj_rs { -// ======================================================================================= -// ArcWakerPromiseNode - -ArcWakerPromiseNode::ArcWakerPromiseNode(kj::Promise promise) - : node(PromiseNode::from(kj::mv(promise))) { - node->setSelfPointer(&node); -} - -void ArcWakerPromiseNode::destroy() noexcept { - auto drop = kj::mv(owner); -} - -void ArcWakerPromiseNode::onReady(kj::_::Event* event) noexcept { - node->onReady(event); -} - -void ArcWakerPromiseNode::get(kj::_::ExceptionOrValue& output) noexcept { - node->get(output); - KJ_IF_SOME(exception, kj::runCatchingExceptions([this]() { node = nullptr; })) { - output.addException(kj::mv(exception)); - } -} - -void ArcWakerPromiseNode::tracePromise(kj::_::TraceBuilder& builder, bool stopAtNextEvent) { - // TODO(someday): Is it possible to get the address of the Rust code which cloned our Waker? - - if (node.get() != nullptr) { - node->tracePromise(builder, stopAtNextEvent); - } -} - -// ======================================================================================= -// ArcWaker - -PromiseArcWakerPair ArcWaker::create(const kj::Executor& executor) { - // TODO(perf): newPromiseAndCrossThreadFulfiller() makes two heap allocations, but it is probably - // optimizable to one. - // TODO(perf): This heap allocation could also probably be collapsed into the fulfiller's. - auto waker = - kj::arc(kj::Badge(), executor.newPromiseAndCrossThreadFulfiller()); - auto promise = const_cast(waker.get())->getPromise(); - return { - .promise = kj::mv(promise), - .waker = kj::mv(waker), - }; -} - -kj::Promise ArcWaker::getPromise() { - KJ_REQUIRE(node.owner == nullptr); - node.owner = addRefToThis(); - return kj::_::PromiseNode::to>(OwnPromiseNode(&node)); -} - -ArcWaker::ArcWaker(kj::Badge, kj::PromiseCrossThreadFulfillerPair paf) - : node(kj::mv(paf.promise)), - fulfiller(kj::mv(paf.fulfiller)) {} - -const KjWaker* ArcWaker::clone() const { - return addRefToThis().disown(); -} -void ArcWaker::wake() const { - wake_by_ref(); - drop(); -} -void ArcWaker::wake_by_ref() const { - fulfiller->fulfill(); -} -void ArcWaker::drop() const { - auto drop = kj::Arc::reown(this); -} +// Definition of the arm-nudge hook declared in waker.h; null until an integrating event port +// (kj-rs-tokio's TokioEventPort) installs itself. Thread-local: one loop/port per thread. +thread_local void (*futurePollArmNudge)() = nullptr; // ======================================================================================= -// LazyArcWaker - -const KjWaker* LazyArcWaker::clone() const { - // Rust code wants to suspend and wait for something. We'll start handing out ArcWakers if we - // haven't already been woken synchronously. +// PollWaker +// +// These are defined here rather than inline in waker.h because they reach into FuturePollEvent, +// which is only a complete type once awaiter.h is included. - if (wakeCount.load(std::memory_order_relaxed) > 0) { - // We were already woken synchronously, so there's no point handing out more wakers for the - // current call to `Future::poll()`. We can hand out a noop waker by returning nullptr. - return nullptr; - } - - auto lock = cloned.lockExclusive(); - - if (*lock == kj::none) { - // We haven't been cloned before, so make a new ArcWaker. - *lock = ArcWaker::create(executor); - } - - return KJ_ASSERT_NONNULL(*lock).waker->clone(); +PollWaker::PollWaker(FuturePollEvent& futurePollEvent) + : holder(FuturePollEventHolder{futurePollEvent}), + cell(futurePollEvent.wakerCellRef()) { + // Every poll starts on the owning thread; renew the cell's cross-thread fulfiller if a + // foreign-thread wake consumed it getting us here. Doing it at poll-top (rather than inside + // the consumed promise's own continuation) means we never destroy a promise chain from within + // its own continuation frame. + futurePollEvent.ensureCrossThreadWakeArmed(); } -void LazyArcWaker::wake() const { - // LazyArcWakers are only exposed to Rust by const borrow, meaning Rust can never arrange to call - // `wake()`, which drops `self`, on this object. - KJ_UNIMPLEMENTED("Rust user code should never have possess a consumable " - "reference to LazyArcWaker"); -} +PollWaker::~PollWaker() noexcept(false) {} -void LazyArcWaker::wake_by_ref() const { - // Woken synchronously during a call to `future.poll(awaitWaker)`. - wakeCount.fetch_add(1, std::memory_order_relaxed); +void PollWaker::wakeByRef() const { + // Delegate to the cell, which handles both threads: on the owning thread this is a synchronous + // same-turn wake (armDepthFirst() is idempotent and safe from within the event's own fire(), + // so it works whether we were reached from onReady() or fire(), and causes an immediate + // re-poll); from a foreign thread it goes through the cross-thread fulfiller — `&Waker` is + // Sync, so even this borrowed waker may legally be woken from another thread during the poll. + cell.wakeByRef(); } -void LazyArcWaker::drop() const { - ++dropCount; +kj::Arc PollWaker::cloneCell() const { + // Rust wants a waker it can retain and wake later: hand out a strong reference to the event's + // FutureWakerCell. Atomic refcount, safe from any thread. + return cell.addRef(); } -kj::Maybe> LazyArcWaker::reset() { - // This function is only called after `future.poll(awaitWaker)` has returned, meaning Rust has - // dropped its reference. Thus, we don't need to worry about thread-safety here, and can call - // `cloned.getWithoutLock()`, for example. - - KJ_ASSERT(dropCount == 1); - KJ_DEFER(dropCount = 0); - KJ_DEFER(wakeCount.store(0, std::memory_order_relaxed)); - - // Reset the ArcWaker on our way out. Since we only return the ArcWaker's promise to our caller, - // we ensure that Rust owns the only remaining ArcWaker clones, if any. - // - // TODO(perf): If ArcWakers were resettable, we could instead return the ArcWaker for our caller - // to cache for later use. - KJ_DEFER(cloned.getWithoutLock() = kj::none); - - if (wakeCount.load(std::memory_order_relaxed) > 0) { - // The future returned Pending, but synchronously called `wake_by_ref()` on the LazyArcWaker, - // indicating it wants to immediately be polled again. We should arm our event right now, - // which will call `await_ready()` again on the event loop. - return kj::Promise(kj::READY_NOW); - } else KJ_IF_SOME(arcWakerPair, cloned.getWithoutLock()) { - // The future returned Pending and cloned an ArcWaker to notify us later. We'll arrange for - // the ArcWaker's promise to arm our event once it's fulfilled. - return kj::mv(arcWakerPair.promise); - } else { - // The future returned Pending, did not call `wake_by_ref()` on the LazyArcWaker, and did not - // clone an ArcWaker. Rust is either awaiting a KJ promise, or the Rust equivalent of - // kj::NEVER_DONE. - return kj::none; +kj::Maybe PollWaker::tryGetFuturePollEvent() const { + KJ_IF_SOME(h, holder.tryGet()) { + return h.futurePollEvent; } + return kj::none; } } // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/waker.h b/src/rust/cxx/kj-rs/waker.h index 382572aabb4..6009f0c73e0 100644 --- a/src/rust/cxx/kj-rs/waker.h +++ b/src/rust/cxx/kj-rs/waker.h @@ -1,176 +1,194 @@ #pragma once -#include "promise.h" +#include "kj-rs/executor-guarded.h" #include #include -#include #include -#include - namespace kj_rs { -using kj::uint; +class FuturePollEvent; -// ======================================================================================= -// KjWaker +// Hook invoked when a waker arms its FuturePollEvent (below) from the owning loop's thread. An +// integrating kj::EventPort that drives tokio tasks inside its own wait() (see kj-rs-tokio) +// installs this to nudge itself out of a blocking park: a tokio task completing during the port's +// block_on() may arm a KJ event same-thread, and KJ's edge-triggered setRunnable() misses that +// arm when the loop's runnable state is already set (e.g. left true by a prior timer). Null +// (no-op) by default; thread-local, one integrating port per loop thread. +extern thread_local void (*futurePollArmNudge)(); -class FuturePollEvent; +// ======================================================================================= +// FutureWakerCell -// KjWaker is an abstract base class which defines an interface mirroring Rust's RawWakerVTable -// struct. Rust has four trampoline functions, defined in waker.rs, which translate Waker::clone(), -// Waker::wake(), etc. calls to the virtual member functions on this class. +// FutureWakerCell is the thread-safe "waker cell" behind every Rust waker that outlives a single +// `Future::poll()` call. `std::task::Waker` is `Send + Sync`, so safe Rust may clone, wake, and +// drop it from any thread; the cell upholds that contract: +// +// - clone/drop are atomic refcount operations (kj::AtomicRefcounted, handed across the FFI as +// `kj::Arc`). +// - wakeByRef() checks whether it is running on the owning event loop's thread. On the owning +// thread — the overwhelmingly common case, e.g. every tokio I/O readiness event under +// kj-rs-tokio — it arms the owning FuturePollEvent's `kj::_::Event` directly via +// `Event::armDepthFirst()`, which is idempotent across same-turn arms and safe to call from +// within the event's own `fire()` (turn() unlinks the event before firing). This is the same +// way a RustPromiseAwaiter leaf arms the FuturePollEvent when its Promise becomes ready — a +// wake is just another leaf arming the same event. From any other thread, wakeByRef() +// fulfills a `kj::CrossThreadPromiseFulfiller` instead (safe from threads with no KJ event +// loop at all, e.g. tokio's blocking pool); the promise side lives in the FuturePollEvent and +// arms it from the owning thread, where the weak event link below is safe to read. The +// consumed fulfiller is renewed at the top of the next poll (see PollWaker's constructor), so +// repeated cross-thread wakes coalesce into the pending poll — exactly the coalescing the +// Waker contract permits. +// +// Neutralize-on-drop: the cell's link to the Event is weak — a `kj::Maybe` invalidated +// (structurally, by the owning FuturePollEvent's RAII guard; see awaiter.h) when that event is +// destroyed. The link is only ever read or written on the owning thread (same-thread wakes, +// neutralize()), so it needs no synchronization; a cell reference that Rust retains past the +// Future's lifetime (e.g. parked in a channel's AtomicWaker) observes a dead link — or, cross- +// thread, a fulfiller whose promise side died with the event — and the wake is a safe no-op +// rather than arming a freed Event. // -// Rust requires Wakers to be Send and Sync, meaning all of the functions defined here may be called -// concurrently by any thread. Derived class implementations of these functions must handle this, -// which is why all of the virtual member functions are `const`-qualified. -class KjWaker { +// Ownership only ever crosses the FFI as real `kj::Arc` handles (PollWaker:: +// cloneCell(), addRef()). Rust's RawWakerVTable island (waker.rs) carries its handle in the +// RawWaker data slot, disowning/reowning it at the vtable edge — the one place `std::task::Waker` +// forces a raw pointer. +class FutureWakerCell final: public kj::AtomicRefcounted { public: - // Return a pointer to a new strong ref to a KjWaker. Note that `clone()` may return nullptr, - // in which case the Rust implementation in waker.rs will treat it as a no-op Waker. Rust - // immediately wraps this pointer in its own Waker object, which is responsible for later - // releasing the strong reference. - // - // TODO(cleanup): Build kj::Arc into cxx-rs so we can return one instead of a raw pointer. - virtual const KjWaker* clone() const = 0; - - // Wake and drop this waker. - virtual void wake() const = 0; - - // Wake this waker, but do not drop it. - virtual void wake_by_ref() const = 0; - - // Drop this waker. - virtual void drop() const = 0; - - // If this KjWaker implementation has an associated FuturePollEvent, C++ code can request access - // to it here. The RustPromiseAwaiter class (which helps Rust `.await` KJ Promises) uses this to - // optimize awaits, when possible. - virtual kj::Maybe tryGetFuturePollEvent() const { - return kj::none; + explicit FutureWakerCell(kj::_::Event& event) + : executor(kj::getCurrentThreadExecutor().addRef()), + event(event) {} + + // Called by `~FuturePollEvent` (owning thread) to neutralize this cell and every outstanding + // Rust reference to it, making any subsequent same-thread wake a safe no-op. (Cross-thread + // wakes neutralize independently: fulfilling the cell's fulfiller after the event destroyed + // the promise side is already a no-op.) Const because `kj::Arc` (like the FFI) only hands out + // const access; `event` is owning-thread-only interior state. + void neutralize() const { + event = kj::none; } -}; -// ======================================================================================= -// ArcWakerPromiseNode + // Wake from any thread: arm the owning FuturePollEvent (directly on the owning thread, via the + // cross-thread fulfiller otherwise), or no-op if it has been neutralized. Const because Rust + // reaches it through `&self`. + void wakeByRef() const { + if (isCurrent(*executor)) { + // Owning thread: arm the event directly. `event` is only touched on this thread. + KJ_IF_SOME(e, event) { + e.armDepthFirst(); + // Nudge an integrating event port out of a blocking park (see `futurePollArmNudge`). + // No-op unless a port installed the hook and is currently parked. + if (futurePollArmNudge != nullptr) { + futurePollArmNudge(); + } + } + } else { + // Foreign thread (possibly one with no KJ event loop): deliver through the cross-thread + // fulfiller. Fulfilling twice before the owning loop renews it, or after the promise side + // died with the event, is a documented no-op — wakes coalesce. + auto lock = crossThreadWake.lockShared(); + if (*lock != nullptr) { + (*lock)->fulfill(); + } + } + } -class ArcWaker; + // Install a fresh cross-thread fulfiller, replacing any consumed one. Called on the owning + // thread by FuturePollEvent (at construction and at the top of each poll); the lock is only + // ever contended by a concurrent foreign-thread wakeByRef(). + void replaceCrossThreadFulfiller( + kj::Own> fulfiller) const { + *crossThreadWake.lockExclusive() = kj::mv(fulfiller); + } -class ArcWakerPromiseNode: public kj::_::PromiseNode { - public: - ArcWakerPromiseNode(kj::Promise promise); - KJ_DISALLOW_COPY_AND_MOVE(ArcWakerPromiseNode); + // True if the current cross-thread fulfiller has been consumed (or discarded) and should be + // renewed before the next park. Owning thread only. + bool needsFreshCrossThreadFulfiller() const { + auto lock = crossThreadWake.lockShared(); + return *lock == nullptr || !(*lock)->isWaiting(); + } - void destroy() noexcept override; - void onReady(kj::_::Event* event) noexcept override; - void get(kj::_::ExceptionOrValue& output) noexcept override; - void tracePromise(kj::_::TraceBuilder& builder, bool stopAtNextEvent) override; + // Hand out a new strong reference. Const + const_cast because Rust reaches it through `&self`: + // cells are always heap-allocated non-const (kj::arc in FuturePollEvent), and the atomic + // refcount bump is safe from any thread. + kj::Arc addRef() const { + return const_cast(*this).addRefToThis(); + } - private: - kj::Arc owner = nullptr; - OwnPromiseNode node; + // Re-own a strong reference previously surrendered to a raw pointer: waker.rs disowns the + // kj::Arc it parks in a RawWaker data slot, and its vtable's drop calls this to reclaim it. + // Exposed to Rust as an `unsafe fn`: `this` must carry exactly such a surrendered reference, + // and dropping the returned handle releases it. + kj::Arc reown() const { + return kj::Arc::reown(this); + } - friend class ArcWaker; + private: + // The owning event loop's executor, used to route wakes: captured at construction (which + // happens on the owning thread), immutable afterwards, safe to read from any thread. Owned via + // addRef() so a cell retained by Rust past loop teardown still has a valid Executor to ask + // (isCurrent() then reports false and the wake takes the — dead, no-op — fulfiller path). + kj::Own executor; + + // Weak, owner-invalidated reference to the owning FuturePollEvent's Event base: non-owning (the + // event lives in the promise graph; the cell must observe its death, never extend its life) and + // nulled by `neutralize()` when that event is destroyed. Only read and written on the owning + // thread, so no synchronization is required; mutable because the cell is only ever reached + // const (kj::Arc / FFI `&self`). + mutable kj::Maybe event; + + // Cross-thread wake delivery: fulfilled by foreign-thread wakeByRef(), renewed by the owning + // thread (replaceCrossThreadFulfiller). The promise side lives in the FuturePollEvent; if the + // event dies first, fulfilling is a safe no-op. Mutex-guarded because the owning thread's + // renewal races with foreign-thread fulfills; same-thread wakes never touch it. + kj::MutexGuarded>> crossThreadWake; }; // ======================================================================================= -// ArcWaker - -class ArcWaker; +// PollWaker -struct PromiseArcWakerPair { - kj::Promise promise; - kj::Arc waker; -}; - -// ArcWaker is an atomic-refcounted wrapper around a `CrossThreadPromiseFulfiller`. -// The atomic-refcounted aspect makes it safe to call `clone()` and `drop()` concurrently, while the -// `CrossThreadPromiseFulfiller` aspect makes it safe to call `wake_by_ref()` concurrently. Finally, -// `wake()` is implemented in terms of `wake_by_ref()` and `drop()`. +// PollWaker is the waker C++ passes to `Future::poll()`. It lives on the stack / in a coroutine +// frame for the duration of a single poll, and Rust only ever borrows it (waker.rs wraps it in a +// Waker whose drop is a no-op). // -// This class is mostly an implementation detail of LazyArcWaker. -class ArcWaker: public kj::AtomicRefcounted, public KjWaker { +// - wakeByRef() delegates to the event's FutureWakerCell, which handles both the same-thread +// (synchronous same-turn re-poll) and foreign-thread cases — `&Waker` is Sync, so even the +// borrowed waker may legally be woken from another thread during the poll. +// - cloneCell() is how Rust retains a waker past the poll: it hands out a strong reference to +// the event's FutureWakerCell, so a later wake from any thread arms the same event. +// - tryGetFuturePollEvent() lets RustPromiseAwaiter (which helps Rust `.await` KJ Promises) +// arm the event directly instead of going through a waker, when possible (owning thread +// only). +class PollWaker final { public: - // Construct a new promise and ArcWaker promise pair, with the Promise to be scheduled on the - // event loop associated with `executor`. - static PromiseArcWakerPair create(const kj::Executor& executor); + // `futurePollEvent` is the FuturePollEvent responsible for calling `Future::poll()`, and must + // outlive this PollWaker. Construction happens on the owning thread at the top of each poll, + // and renews the cell's cross-thread fulfiller if a foreign-thread wake consumed it. + explicit PollWaker(FuturePollEvent& futurePollEvent); + ~PollWaker() noexcept(false); + KJ_DISALLOW_COPY_AND_MOVE(PollWaker); - ArcWaker(kj::Badge, kj::PromiseCrossThreadFulfillerPair paf); - KJ_DISALLOW_COPY_AND_MOVE(ArcWaker); + // Wake from any thread: arm the associated FuturePollEvent so it (re-)polls. + void wakeByRef() const; - const KjWaker* clone() const override; - void wake() const override; - void wake_by_ref() const override; - void drop() const override; + // Hand out a new strong reference to the event's FutureWakerCell, for Rust to retain and wake + // later. Safe from any thread (atomic refcount). + kj::Arc cloneCell() const; - private: - kj::Promise getPromise(); - - ArcWakerPromiseNode node; - kj::Own> fulfiller; -}; - -// ======================================================================================= -// LazyArcWaker - -// LazyArcWaker is intended to live locally on the stack or in a coroutine frame. Trying to -// `clone()` it will cause it to allocate an ArcWaker for the caller. -class LazyArcWaker: public KjWaker { - public: - // Create a new or clone an existing ArcWaker, leak its pointer, and return it. This may be called - // by any thread. - const KjWaker* clone() const override; - - // Unimplemented, because Rust user code cannot consume the `std::task::Waker` we create which - // wraps this LazyArcWaker. - void wake() const override; - - // Rust user code can wake us synchronously during the execution of `future.poll()` using this - // function. This may be called by any thread. - void wake_by_ref() const override; - - // Does not actually destroy this object. Instead, we increment a counter so we can assert that it - // was dropped exactly once before `future.poll()` returned. This can only be called on the thread - // which is doing the awaiting, because our implementation of `future.poll()` never transfers the - // Waker object to a different thread. - void drop() const override; - - // Used by the owner of LazyArcWaker after `future.poll()` has returned, to retrieve the - // LazyArcWaker's state for further processing. This is non-const, because by the time this is - // called, Rust has dropped all of its borrows to this class, meaning we no longer have to worry - // about thread safety. - // - // This function will assert if `drop()` has not been called since LazyArcWaker was constructed, - // or since the last call to `reset()`. - // - // Returns `kj::none` the LazyArcWaker was neither woken nor cloned before being dropped. Returns - // `kj::READY_NOW` if the LazyArcWaker was synchronously woken. Otherwise, if `clone()` was - // called, return the promise associated with the cloned ArcWaker. - kj::Maybe> reset(); + // The FuturePollEvent whose poll() this waker was created for, if the current thread's + // kj::Executor is the one which owns it. + kj::Maybe tryGetFuturePollEvent() const; private: - // We store the kj::Executor for the constructing thread so that we can lazily instantiate a - // CrossThreadPromiseFulfiller from any thread in our `clone()` implementation. - const kj::Executor& executor = kj::getCurrentThreadExecutor(); - - // Initialized by `clone()`, which may be called by any thread. This could almost be a - // `kj::Lazy`, but we need to be able to detect when we haven't been cloned. - kj::MutexGuarded> cloned; - - // Incremented by `wake_by_ref()`, which may be called by any thread. All operations use relaxed - // memory order, because this counter does not guard any memory. - mutable std::atomic wakeCount{0}; - - // Incremented by `drop()`, so we can validate that `drop()` is only called once on this object. - // - // Rust requires that Wakers be droppable by any thread. However, we own the implementation of - // `poll()` to which `LazyArcWaker&` is passed, and those implementations store the Rust - // `std::task::Waker` object on the stack,, and never move it elsewhere. Since that object is - // responsible for calling `LazyArcWaker::drop()`, we know for sure that `drop()` will only ever be - // called on the thread which constructed it. Therefore, there is no need to make `dropCount` - // thread-safe. - mutable uint dropCount = 0; + struct FuturePollEventHolder { + FuturePollEvent& futurePollEvent; + }; + ExecutorGuarded holder; + + // The event's cell, cached here so wakeByRef()/cloneCell() work from any thread without going + // through the executor-guarded holder. Valid for this PollWaker's whole life: the cell is + // created eagerly with the FuturePollEvent, which outlives the poll. + const FutureWakerCell& cell; }; } // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/waker.rs b/src/rust/cxx/kj-rs/waker.rs index 472a259aaa3..3c9ec78ea41 100644 --- a/src/rust/cxx/kj-rs/waker.rs +++ b/src/rust/cxx/kj-rs/waker.rs @@ -1,102 +1,184 @@ +//! FFI island (see crate-root `#![deny(unsafe_code)]`): the two `RawWakerVTable`s bridging the +//! C++ wakers into `std::task::Waker`. `std`'s vtable ABI is four raw-pointer functions, so this +//! file is where ownership must round-trip through a raw `RawWaker` data slot — everywhere else +//! (including the whole C++ interface) waker ownership is a real `kj::Arc` handle. A genuine +//! unsafe seam. +#![allow(unsafe_code)] + use std::task::RawWaker; use std::task::RawWakerVTable; use std::task::Waker; -use crate::ffi::KjWaker; +use crate::KjArc; +use crate::ffi::FutureWakerCell; +use crate::ffi::PollWaker; -// Safety: We use the type system to express the Sync nature of KjWaker in the cxx-rs FFI boundary. -// Specifically, we only allow invocations on const KjWakers, and in KJ C++, use of const-qualified -// functions is thread-safe by convention. Our implementations of KjWakers in C++ respect this -// convention. -// -// Note: Implementing these traits does not seem to be required for building, but the Waker -// documentation makes it clear Send and Sync are a requirement of the pointed-to type. +// Thread-safety: `std::task::Waker` documents that the vtable functions must be thread-safe, and +// `Waker: Send + Sync` means safe Rust may clone, wake, and drop these from any thread. The +// vtables uphold that for real: clone/drop are atomic refcount operations on the +// `FutureWakerCell` (`Send + Sync`; see the impls in ffi.rs), and `wakeByRef` routes wakes +// through an owning-executor check on the C++ side — same-thread wakes arm the event directly, +// foreign-thread wakes go through a cross-thread fulfiller (waker.h). + +// ======================================================================================= +// Borrowed vtable: Wakers lending out the PollWaker C++ passes to `Future::poll()` // -// https://doc.rust-lang.org/std/task/struct.RawWaker.html -// https://doc.rust-lang.org/std/task/struct.RawWakerVTable.html -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. -unsafe impl Send for KjWaker {} -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. -unsafe impl Sync for KjWaker {} - -impl From<&KjWaker> for Waker { - fn from(waker: &KjWaker) -> Self { - let waker = RawWaker::new( - std::ptr::from_ref::(waker).cast::<()>(), - &KJ_WAKER_VTABLE, +// `data` is the `&PollWaker` the Waker was built from — borrowed, never null, and alive for the +// duration of the poll (the Waker is created and dropped inside the poll bridge in future.rs). +// Dropping such a Waker frees nothing; cloning it takes a real strong reference to the event's +// FutureWakerCell and switches to the owned-cell vtable below. + +impl From<&PollWaker> for Waker { + fn from(waker: &PollWaker) -> Self { + let raw = RawWaker::new( + std::ptr::from_ref::(waker).cast::<()>(), + &POLL_WAKER_VTABLE, ); - // Safety: KjWaker's Rust-exposed interface is Send and Sync and its RawWakerVTable - // implementation functions are all thread-safe. - // - // https://doc.rust-lang.org/std/task/struct.Waker.html#safety-1 - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe { Self::from_raw(waker) } + // Safety: the vtable functions below uphold the RawWaker contract (see the thread-safety + // note above); `data` outlives the Waker because future.rs drops the Waker before poll + // returns. + unsafe { Self::from_raw(raw) } } } -// Helper function for use in KjWaker's RawWakerVTable implementation to factor out a tedious null -// pointer check. -fn deref_kj_waker<'a>(data: *const ()) -> Option<&'a KjWaker> { +/// # Safety +/// +/// `data` must be the pointer a [`From<&PollWaker>`] conversion was made with, still live per the +/// `RawWaker` contract (upheld because these Wakers only exist within a single `poll` call). +unsafe fn poll_waker_clone(data: *const ()) -> RawWaker { + // Safety: forwarded from this fn's `# Safety` contract. + let waker = unsafe { &*data.cast::() }; + RawWaker::new( + cell_into_raw(waker.clone_cell()).cast::<()>(), + &CELL_WAKER_VTABLE, + ) +} + +/// # Safety +/// +/// Same contract as [`poll_waker_clone`]. +unsafe fn poll_waker_wake_by_ref(data: *const ()) { + // Safety: forwarded from this fn's `# Safety` contract. + let waker = unsafe { &*data.cast::() }; + waker.wake_by_ref(); +} + +/// # Safety +/// +/// Same contract as [`poll_waker_clone`]. Consuming a borrowed Waker owns nothing, so `wake` is +/// just `wake_by_ref` (the paired drop is a no-op). +unsafe fn poll_waker_wake(data: *const ()) { + // Safety: forwarded from this fn's `# Safety` contract. + unsafe { poll_waker_wake_by_ref(data) } +} + +fn poll_waker_drop(_data: *const ()) { + // No-op: the PollWaker is stack-owned by the C++ poll scope; this Waker only borrowed it. +} + +static POLL_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new( + poll_waker_clone, + poll_waker_wake, + poll_waker_wake_by_ref, + poll_waker_drop, +); + +// ======================================================================================= +// Owned-cell vtable: retained Wakers holding a strong reference to a FutureWakerCell +// +// `data` carries one strong reference (a disowned `KjArc`). Clone takes another +// reference; drop re-owns and releases the carried one. All of it is safe from any thread: the +// refcount is atomic and the cell's wake is executor-routed (waker.h). + +/// Surrender the handle's strong reference into a bare pointer for a `RawWaker` data slot. +/// Reversed by `FutureWakerCell::reown` in [`cell_waker_drop`]. +fn cell_into_raw(cell: KjArc) -> *const FutureWakerCell { + let ptr = cell.get(); + std::mem::forget(cell); + ptr +} + +/// # Safety +/// +/// `data` must be either null or a pointer produced by [`cell_into_raw`] whose strong reference +/// is still carried by this `RawWaker` (upheld by the `Waker`/`RawWaker` contract: these vtable +/// entries are only installed alongside such pointers, by [`poll_waker_clone`] and the functions +/// below). +unsafe fn cell_deref<'a>(data: *const ()) -> Option<&'a FutureWakerCell> { if data.is_null() { None } else { - let p = data.cast::(); - // Safety: - // 1. p is guaranteed non-null by the check above. - // 2. This function is only used in the implementations of our RawWakerVTable for KjWaker. - // All vtable implementation functions are trivially guaranteed that their owning Waker - // object is still alive. We assume the Waker was constructed correctly to begin with, - // and that therefore the pointer still points to valid memory. - // 3. We do not read or write the KjWaker's memory, so there are no atomicity concerns nor - // interleaved pointer/reference access concerns. - // - // https://doc.rust-lang.org/std/ptr/index.html#safety - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - Some(unsafe { &*p }) + // Safety: non-null per the check; live per this fn's `# Safety` contract (the carried + // strong reference keeps the cell alive). + Some(unsafe { &*data.cast::() }) } } -pub fn kj_waker_clone(data: *const ()) -> RawWaker { - let new_data = if let Some(kj_waker) = deref_kj_waker(data) { - kj_waker.clone_kj_waker().cast::<()>() +/// # Safety +/// +/// Same contract as [`cell_deref`]. +unsafe fn cell_waker_clone(data: *const ()) -> RawWaker { + // Safety: forwarded from this fn's `# Safety` contract. + let new_data = if let Some(cell) = unsafe { cell_deref(data) } { + cell_into_raw(cell.add_ref()) } else { std::ptr::null() }; - RawWaker::new(new_data, &KJ_WAKER_VTABLE) + RawWaker::new(new_data.cast::<()>(), &CELL_WAKER_VTABLE) } -pub fn kj_waker_wake(data: *const ()) { - if let Some(kj_waker) = deref_kj_waker(data) { - kj_waker.wake(); +/// # Safety +/// +/// Same contract as [`cell_deref`]. +unsafe fn cell_waker_wake_by_ref(data: *const ()) { + // Safety: forwarded from this fn's `# Safety` contract. + if let Some(cell) = unsafe { cell_deref(data) } { + cell.wake_by_ref(); } } -pub fn kj_waker_wake_by_ref(data: *const ()) { - if let Some(kj_waker) = deref_kj_waker(data) { - kj_waker.wake_by_ref(); +/// # Safety +/// +/// Same contract as [`cell_deref`], and the carried strong reference is released (the `RawWaker` +/// must not be used again — guaranteed by the `Waker` contract for `drop`). +unsafe fn cell_waker_drop(data: *const ()) { + // Safety: forwarded from this fn's `# Safety` contract. + if let Some(cell) = unsafe { cell_deref(data) } { + // Safety: `data` carries a strong reference per this fn's `# Safety` contract; re-own it + // and let the handle fall, releasing the reference. + let _cell = unsafe { cell.reown() }; } } -pub fn kj_waker_drop(data: *const ()) { - if let Some(kj_waker) = deref_kj_waker(data) { - kj_waker.drop(); +/// # Safety +/// +/// Same contract as [`cell_waker_drop`]. +unsafe fn cell_waker_wake(data: *const ()) { + // Safety: forwarded from this fn's `# Safety` contract; wake-then-release. + unsafe { + cell_waker_wake_by_ref(data); + cell_waker_drop(data); } } -static KJ_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new( - kj_waker_clone, - kj_waker_wake, - kj_waker_wake_by_ref, - kj_waker_drop, +static CELL_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new( + cell_waker_clone, + cell_waker_wake, + cell_waker_wake_by_ref, + cell_waker_drop, ); -/// If `waker` wraps a `KjWaker`, return the `KjWaker` pointer it was originally constructed with, -/// or null if `waker` does not wrap a `KjWaker`. Note that the `KjWaker` pointer originally used -/// to construct `waker` may itself by null. -pub fn try_into_kj_waker_ptr(waker: &Waker) -> *const KjWaker { - if waker.vtable() == &KJ_WAKER_VTABLE { - waker.data().cast::() +/// If `waker` lends out a C++ `PollWaker` (borrowed vtable above), return a reference to it, +/// borrowed from `waker` itself. Owned-cell and foreign Wakers both return `None`: neither +/// exposes a `FuturePollEvent` to arm directly, so `RustPromiseAwaiter` takes its generic +/// fallback path for them. +pub fn try_poll_waker(waker: &Waker) -> Option<&PollWaker> { + if waker.vtable() == &POLL_WAKER_VTABLE { + // Safety: Wakers carrying POLL_WAKER_VTABLE are only ever built by `From<&PollWaker>` + // above, so `data` is a `&PollWaker` that outlives `waker` (the PollWaker is stack-owned + // by the C++ poll driving this call); the returned borrow is tied to `waker`'s lifetime. + Some(unsafe { &*waker.data().cast::() }) } else { - std::ptr::null() + None } } diff --git a/src/rust/kj/tests/ffi-test.c++ b/src/rust/kj/tests/ffi-test.c++ index e4218071c45..34cee77d7ee 100644 --- a/src/rust/kj/tests/ffi-test.c++ +++ b/src/rust/kj/tests/ffi-test.c++ @@ -48,10 +48,11 @@ class MockHttpService: public kj::HttpService { class TestConnectResponse: public kj::HttpService::ConnectResponse { public: - void accept(uint statusCode, kj::StringPtr statusText, const kj::HttpHeaders& headers) override { + void accept( + kj::uint statusCode, kj::StringPtr statusText, const kj::HttpHeaders& headers) override { KJ_UNIMPLEMENTED("not exercised by test"); } - kj::Own reject(uint statusCode, + kj::Own reject(kj::uint statusCode, kj::StringPtr statusText, const kj::HttpHeaders& headers, kj::Maybe expectedBodySize = kj::none) override {