From 839ce65cb891a5dd2221976e34af3603cef2a389 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bartek=20Iwa=C5=84czuk?= Date: Tue, 27 Oct 2026 18:40:29 +0100 Subject: [PATCH 1/2] Make jsg::Error::new_internal_error() public Rust APIs need it for failures that only a bug causes, such as an IoContext that is no longer current. Throwing the error gives JavaScript a generic "internal error; reference = " Error and logs the message with that ID. A new test checks the redaction. --- src/rust/jsg-test/tests/internal_error.rs | 27 +++++++++++++++++++++++ src/rust/jsg-test/tests/mod.rs | 1 + src/rust/jsg/lib.rs | 15 +++++++------ 3 files changed, 36 insertions(+), 7 deletions(-) create mode 100644 src/rust/jsg-test/tests/internal_error.rs diff --git a/src/rust/jsg-test/tests/internal_error.rs b/src/rust/jsg-test/tests/internal_error.rs new file mode 100644 index 00000000000..2c24eeb9801 --- /dev/null +++ b/src/rust/jsg-test/tests/internal_error.rs @@ -0,0 +1,27 @@ +// Copyright (c) 2026 Cloudflare, Inc. +// Licensed under the Apache 2.0 license found in the LICENSE file or at: +// https://opensource.org/licenses/Apache-2.0 + +//! An internal error reaches JavaScript without its message. + +use jsg::FromJS; + +#[test] +fn internal_error_is_redacted_when_thrown() { + crate::Harness::new().run_in_context(|lock, _ctx| { + let thrown = + lock.try_catch(|_lock| Err::<(), _>(jsg::Error::new_internal_error("secret detail")))?; + let Err(exception) = thrown else { + panic!("the internal error was not thrown"); + }; + let err = jsg::Error::from_js(lock, exception)?; + assert_eq!(err.name, jsg::ExceptionType::Error); + assert!( + err.message.starts_with("internal error; reference = "), + "{}", + err.message + ); + assert!(!err.message.contains("secret detail"), "{}", err.message); + Ok(()) + }); +} diff --git a/src/rust/jsg-test/tests/mod.rs b/src/rust/jsg-test/tests/mod.rs index d1bbdbd0ca4..79b2cbfb543 100644 --- a/src/rust/jsg-test/tests/mod.rs +++ b/src/rust/jsg-test/tests/mod.rs @@ -19,6 +19,7 @@ mod gc; mod global; mod harness; mod integer; +mod internal_error; mod isolates; mod jsg_oneof; mod jsg_struct; diff --git a/src/rust/jsg/lib.rs b/src/rust/jsg/lib.rs index 7432af5cec6..05af7cd5298 100644 --- a/src/rust/jsg/lib.rs +++ b/src/rust/jsg/lib.rs @@ -445,13 +445,14 @@ impl Error { }) } - /// Like `new_type_error()`, but marks the resulting `Error` `is_internal`, so - /// `Lock::throw_exception()` redacts `message` instead of exposing it to guest JS. - /// The JS-visible error type ends up as plain `Error` regardless of the type set - /// here (`Lock::throw_internal_error()` -> C++ `makeInternalError()` always throws - /// a plain `Error`), so `name` only matters if something inspects the `Error` - /// value directly without throwing it (e.g. in tests). - fn new_internal_error(message: impl Into) -> Self { + /// An error that only a bug causes, like a failed C++ `KJ_ASSERT`. Throwing it + /// (`Lock::throw_exception()`) redacts `message`: JavaScript gets a generic + /// `"internal error; reference = "` `Error`, and `message` is logged with that ID. + /// The JS-visible error type is plain `Error` regardless of the type set here + /// (`Lock::throw_internal_error()` -> C++ `makeInternalError()` always throws a plain + /// `Error`), so `name` only matters if something inspects the `Error` value directly + /// without throwing it (e.g. in tests). + pub fn new_internal_error(message: impl Into) -> Self { Self { name: ExceptionType::TypeError, message: message.into(), From 42efd1c78b7dd28afb4d87f283dcad69a418d03c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bartek=20Iwa=C5=84czuk?= Date: Tue, 27 Oct 2026 18:40:29 +0100 Subject: [PATCH 2/2] Add await_io() and add_object() to the Rust io-context crate IoContext::await_io() waits for a Rust future as part of the request, then runs a callback with its output in the IoContext. It goes through IoContext::awaitIoAndRun(), a new public variant of awaitIo() for callers that settle no jsg::Promise. Both now share the I/O half of awaitIoImpl(), split out as resumeAfterIo(), so the callback runs like awaitIo()'s does: holding a pending event, in the critical section and async context frame current when the wait began, and recorded for async tracing. If the request ends first, the future and callback are dropped without the callback running. IoContext::add_object() gives the IoContext ownership of a value, as C++ addObject() does. It returns IoOwned, from which the value can only be taken out, rather than used in place as IoOwn is. --- src/workerd/io/io-context.h | 184 +++++++++++------- src/workerd/io/io-context/BUILD.bazel | 8 +- src/workerd/io/io-context/ffi.c++ | 22 +++ src/workerd/io/io-context/ffi.h | 25 +++ src/workerd/io/io-context/ffi.rs | 61 +++++- src/workerd/io/io-context/io-context-test.c++ | 169 ++++++++++++++++ src/workerd/io/io-context/lib.rs | 115 +++++++++++ src/workerd/io/io-context/test-bridge.h | 15 ++ src/workerd/io/io-context/test_bridge.rs | 164 ++++++++++++++++ 9 files changed, 686 insertions(+), 77 deletions(-) create mode 100644 src/workerd/io/io-context/test-bridge.h diff --git a/src/workerd/io/io-context.h b/src/workerd/io/io-context.h index 0262b40b10c..df28f852e35 100644 --- a/src/workerd/io/io-context.h +++ b/src/workerd/io/io-context.h @@ -676,6 +676,27 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler template jsg::PromiseForResult awaitIo(jsg::Lock& js, kj::Promise promise, Func&& func); + // What a KJ promise for T produced: its value, or the exception it failed with. + template + struct ExceptionOr_ { + using Type = kj::OneOf; + }; + template <> + struct ExceptionOr_ { + using Type = kj::Maybe; + }; + template + using ExceptionOr = ExceptionOr_::Type; + + // Like awaitIo(), for callers that settle no jsg::Promise, such as APIs implemented in Rust: + // waits for `promise` as a task of this IoContext, holding a pending event, then calls + // `func(Worker::Lock&, ExceptionOr&&)` in run(), in the critical section and async context + // frame that are current when this is called. If the IoContext is destroyed first, `func` is + // dropped without being called. An exception `func` throws fails the task, as it would for + // addTask(). + template + void awaitIoAndRun(kj::Promise promise, Func&& func); + // Attach the objects to the promise by creating a continuation that holds them. // This ensures the attachments stay alive until the promise resolves. // This should ONLY be used with TraceContext or SpanBuilder objects. @@ -1453,21 +1474,17 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler inline void operator()(jsg::Lock&) const {} }; - template - struct ExceptionOr_ { - using Type = kj::OneOf; - }; - template <> - struct ExceptionOr_ { - using Type = kj::Maybe; - }; - template - using ExceptionOr = ExceptionOr_::Type; - template jsg::PromiseForResult awaitIoImpl( jsg::Lock& js, kj::Promise promise, InputLockOrMaybeCriticalSection ilOrCs, Func&& func); + // The I/O half of awaitIoImpl() and awaitIoAndRun(): waits for `promise` as a task of this + // IoContext, then calls `func(Worker::Lock&, ExceptionOr&&)` in run(), in the async context + // frame that is current when this is called, with an async trace resource as the turn's cause. + template + void resumeAfterIo( + jsg::Lock& js, kj::Promise promise, InputLockOrMaybeCriticalSection ilOrCs, Func&& func); + // Async trace resources for the awaitIo() and awaitJs() bridges. Inert if the context is not // traced. AsyncResource makeAwaitIoResource() const { @@ -1714,7 +1731,7 @@ Result throwOrReturnResult(jsg::Lock& js, IoContext::ExceptionOr&& excep }; template -jsg::PromiseForResult IoContext::awaitIoImpl( +void IoContext::resumeAfterIo( jsg::Lock& js, kj::Promise promise, InputLockOrMaybeCriticalSection ilOrCs, Func&& func) { // WARNING: The fact that `promise` has been passed by value whereas `func` is by reference is // actually important, because this means that if we throw an exception here in the function @@ -1723,29 +1740,20 @@ jsg::PromiseForResult IoContext::awaitIoImpl( requireCurrent(); - // `T` is the type produced by the input promise. `Result` is the type of the final output - // promise. `Func` transforms from `T` to `Result`. - using Result = jsg::ReturnType; - // It is necessary for us to grab a reference to the jsg::AsyncContextFrame here // and pass it into the then(). If the promise is rejected, and there is no rejection // handler attached to it, an unhandledrejection event will be scheduled, and scheduling // that event needs to be done within the appropriate frame to propagate the correct context. - // We need to catch exceptions from KJ and merge them into the result, so that they can propagate - // to JavaScript. + // We need to catch exceptions from KJ and pass them to `func`, so that they can propagate to + // JavaScript. kj::Promise> promiseExceptionOrT = promiseForExceptionOrT(kj::mv(promise)); - // Reminder: This can throw JsExceptionThrown if the execution context has been terminated. - // it's important in that case that `promiseExceptionOrT` will be destroyed before `func`. - auto [jsPromise, resolver] = js.newPromiseAndResolver>(); - // Settled when the KJ promise completes; causes the turn that resumes JavaScript. AsyncResource asyncTraceResource = makeAwaitIoResource(); addTask(promiseExceptionOrT.then( - [this, resolver = kj::mv(resolver), ilOrCs = kj::mv(ilOrCs), - maybeAsyncContext = jsg::AsyncContextFrame::currentRef(js), + [this, ilOrCs = kj::mv(ilOrCs), maybeAsyncContext = jsg::AsyncContextFrame::currentRef(js), asyncTraceResource = kj::mv(asyncTraceResource), // Reminder: It's important that `func` gets attached to the promise before the whole // thing is passed to `addTask()`, so that it's impossible for `func` to be destroyed @@ -1761,35 +1769,86 @@ jsg::PromiseForResult IoContext::awaitIoImpl( }; return run( - [resolver = kj::mv(resolver), - funcResultPair = FuncResultPair{kj::fwd(func), kj::mv(exceptionOrT)}, + [funcResultPair = FuncResultPair{kj::fwd(func), kj::mv(exceptionOrT)}, maybeAsyncContext = kj::mv(maybeAsyncContext), asyncTraceResource = kj::mv(asyncTraceResource)](Worker::Lock& lock) mutable { asyncTraceResource.enterAsTurnCause(); jsg::AsyncContextFrame::Scope asyncScope(lock, maybeAsyncContext); - jsg::Lock& js = lock; - - if constexpr (jsg::isVoid()) { - KJ_IF_SOME(e, funcResultPair.exceptionOrT) { - // We don't use `resolver.reject()` here because if we convert the kj::Exception into - // a JS Error here, it won't have a useful stack trace. V8 can generate a good stack - // trace as long as we construct the Error inside of a promise continuation, so we use - // a `.then()` below that actually extracts the kj::Exception and turn it into a JS - // Error. - resolver.resolve(js, kj::mv(e)); - } else { + funcResultPair.func(lock, kj::mv(funcResultPair.exceptionOrT)); + }, + kj::mv(ilOrCs)); + })); +} + +template +jsg::PromiseForResult IoContext::awaitIoImpl( + jsg::Lock& js, kj::Promise promise, InputLockOrMaybeCriticalSection ilOrCs, Func&& func) { + // WARNING: As in resumeAfterIo(), `promise` is passed by value and `func` by reference, so that + // if we throw an exception here, `promise` is destroyed before `func`. + + requireCurrent(); + + // `T` is the type produced by the input promise. `Result` is the type of the final output + // promise. `Func` transforms from `T` to `Result`. + using Result = jsg::ReturnType; + + // Reminder: This can throw JsExceptionThrown if the execution context has been terminated. + // it's important in that case that `promise` will be destroyed before `func`. + auto [jsPromise, resolver] = js.newPromiseAndResolver>(); + + resumeAfterIo(js, kj::mv(promise), kj::mv(ilOrCs), + [resolver = kj::mv(resolver), func = kj::fwd(func)]( + Worker::Lock& lock, ExceptionOr&& exceptionOrT) mutable { + jsg::Lock& js = lock; + + if constexpr (jsg::isVoid()) { + KJ_IF_SOME(e, exceptionOrT) { + // We don't use `resolver.reject()` here because if we convert the kj::Exception into + // a JS Error here, it won't have a useful stack trace. V8 can generate a good stack + // trace as long as we construct the Error inside of a promise continuation, so we use + // a `.then()` below that actually extracts the kj::Exception and turn it into a JS + // Error. + resolver.resolve(js, kj::mv(e)); + } else { + try { + js.tryCatch([&]() { + if constexpr (jsg::isVoid()) { + func(js); + resolver.resolve(js, kj::none); + } else { + resolver.resolve(js, func(js)); + } + }, [&](jsg::Value error) { + // Here we can just `resolver.reject` because we already have a JS exception. + resolver.reject(js, error.getHandle(js)); + }); + } catch (jsg::JsExceptionThrown&) { + // An uncatchable JS exception -- presumably, the isolate has been terminated. We + // can't convert this into a promise rejection, we need to just propagate it up. + throw; + } catch (...) { + // Again, pass along the KJ exception so we can convert it later in the right context. + resolver.resolve(js, kj::getCaughtExceptionAsKj()); + } + } + } else { + // T is not void. + KJ_SWITCH_ONEOF(exceptionOrT) { + KJ_CASE_ONEOF(exception, kj::Exception) { + // Again, pass along the KJ exception so we can convert it later in the right context. + resolver.resolve(js, kj::mv(exception)); + } + KJ_CASE_ONEOF(result, T) { try { js.tryCatch([&]() { if constexpr (jsg::isVoid()) { - funcResultPair.func(js); + func(js, kj::mv(result)); resolver.resolve(js, kj::none); } else { - resolver.resolve(js, funcResultPair.func(js)); + // Here we can just `resolver.reject` because we already have a JS exception. + resolver.resolve(js, func(js, kj::mv(result))); } - }, [&](jsg::Value error) { - // Here we can just `resolver.reject` because we already have a JS exception. - resolver.reject(js, error.getHandle(js)); - }); + }, [&](jsg::Value error) { resolver.reject(js, error.getHandle(js)); }); } catch (jsg::JsExceptionThrown&) { // An uncatchable JS exception -- presumably, the isolate has been terminated. We // can't convert this into a promise rejection, we need to just propagate it up. @@ -1799,38 +1858,9 @@ jsg::PromiseForResult IoContext::awaitIoImpl( resolver.resolve(js, kj::getCaughtExceptionAsKj()); } } - } else { - // T is not void. - KJ_SWITCH_ONEOF(funcResultPair.exceptionOrT) { - KJ_CASE_ONEOF(exception, kj::Exception) { - // Again, pass along the KJ exception so we can convert it later in the right context. - resolver.resolve(js, kj::mv(exception)); - } - KJ_CASE_ONEOF(result, T) { - try { - js.tryCatch([&]() { - if constexpr (jsg::isVoid()) { - funcResultPair.func(js, kj::mv(result)); - resolver.resolve(js, kj::none); - } else { - // Here we can just `resolver.reject` because we already have a JS exception. - resolver.resolve(js, funcResultPair.func(js, kj::mv(result))); - } - }, [&](jsg::Value error) { resolver.reject(js, error.getHandle(js)); }); - } catch (jsg::JsExceptionThrown&) { - // An uncatchable JS exception -- presumably, the isolate has been terminated. We - // can't convert this into a promise rejection, we need to just propagate it up. - throw; - } catch (...) { - // Again, pass along the KJ exception so we can convert it later in the right context. - resolver.resolve(js, kj::getCaughtExceptionAsKj()); - } - } - } } - }, - kj::mv(ilOrCs)); - })); + } + }); // Reminder: This can throw JsExceptionThrown if the execution context has been terminated. We // have already disowned `promise` and `func` by this point, though, so teardown order is no @@ -1838,6 +1868,12 @@ jsg::PromiseForResult IoContext::awaitIoImpl( return jsPromise.then(js, throwOrReturnResult); } +template +void IoContext::awaitIoAndRun(kj::Promise promise, Func&& func) { + resumeAfterIo(getCurrentLock(), promise.attach(registerPendingEvent()), getCriticalSection(), + kj::fwd(func)); +} + template kj::_::ReducePromises> IoContext::awaitJs(jsg::Lock& js, jsg::Promise jsPromise) { auto paf = kj::newPromiseAndFulfiller>(); diff --git a/src/workerd/io/io-context/BUILD.bazel b/src/workerd/io/io-context/BUILD.bazel index 5658d9fdd1d..f1d0a74b8bc 100644 --- a/src/workerd/io/io-context/BUILD.bazel +++ b/src/workerd/io/io-context/BUILD.bazel @@ -24,7 +24,10 @@ wd_cc_library( name = "ffi", srcs = ["ffi.c++"], hdrs = ["ffi.h"], - deps = ["//src/workerd/io"], + deps = [ + ":ffi.rs@cxx", + "//src/workerd/io", + ], ) # Checked by //src/rust/jsg:cpp-jsg-allowlist-test. @@ -40,11 +43,12 @@ wd_rust_crate( srcs = ["test_bridge.rs"], crate_root = "test_bridge.rs", cxx_bridge_deps = ["//src/rust/jsg"], - cxx_bridge_hdrs = [], + cxx_bridge_hdrs = ["test-bridge.h"], cxx_bridge_src = "test_bridge.rs", deps = [ ":io-context", "//src/rust/jsg", + "@crates_vendor//:futures", ], ) diff --git a/src/workerd/io/io-context/ffi.c++ b/src/workerd/io/io-context/ffi.c++ index efbbe71e23e..b2e34557d83 100644 --- a/src/workerd/io/io-context/ffi.c++ +++ b/src/workerd/io/io-context/ffi.c++ @@ -4,6 +4,8 @@ #include "ffi.h" +#include + #include #include @@ -78,4 +80,24 @@ void io_context_write_logfwdr(uint64_t id, uint32_t channel, ::rust::Slice future, ::rust::Box callback) { + kj::Promise<::rust::Box> output = io_future_wait(kj::mv(future)); + currentWithId(id).awaitIoAndRun(kj::mv(output), + [callback = kj::mv(callback)]( + Worker::Lock& lock, IoContext::ExceptionOr<::rust::Box>&& result) mutable { + // io_future_wait() produces a value and does not fail. + auto output = kj::mv(KJ_REQUIRE_NONNULL(result.tryGet<::rust::Box>(), + "the future that IoContext::await_io() waited for failed")); + io_callback_run( + reinterpret_cast(lock.getIsolate()), kj::mv(callback), kj::mv(output)); + }); +} + +kj::Own io_context_add_object(uint64_t id, ::rust::Box slot) { + auto& context = currentWithId(id); + auto object = context.addObject(kj::heap<::rust::Box>(kj::mv(slot))); + return kj::heap(kj::heap(kj::mv(object))); +} + } // namespace workerd::io_context diff --git a/src/workerd/io/io-context/ffi.h b/src/workerd/io/io-context/ffi.h index 67240317000..aac068b9d27 100644 --- a/src/workerd/io/io-context/ffi.h +++ b/src/workerd/io/io-context/ffi.h @@ -29,4 +29,29 @@ int64_t io_context_now(uint64_t id); // write waits for the output gate, as AnalyticsEngine::writeDataPoint() does. void io_context_write_logfwdr(uint64_t id, uint32_t channel, ::rust::Slice message); +// Rust types, defined in ffi.rs. +struct IoFuture; +struct IoCallback; +struct IoOwnedSlot; + +// IoContext::awaitIoAndRun(): waits for `future` as a task of the IoContext, holding a pending +// event, then runs `callback` with its output in IoContext::run(), in the critical section and +// async context frame that are current when this is called. Both are dropped if the IoContext is +// destroyed first. +void io_context_await_io( + uint64_t id, ::rust::Box future, ::rust::Box callback); + +// Holds what IoContext::addObject() returns for a Rust value, which the IoContext then destroys +// when it is destroyed, unless this is destroyed first. +class IoOwnedObject { + public: + explicit IoOwnedObject(kj::Own object): object(kj::mv(object)) {} + + private: + kj::Own object; +}; + +// IoContext::addObject() for `slot`. +kj::Own io_context_add_object(uint64_t id, ::rust::Box slot); + } // namespace workerd::io_context diff --git a/src/workerd/io/io-context/ffi.rs b/src/workerd/io/io-context/ffi.rs index 7c15117dc42..76374e430ed 100644 --- a/src/workerd/io/io-context/ffi.rs +++ b/src/workerd/io/io-context/ffi.rs @@ -4,16 +4,75 @@ //! The bridge to the C++ `IoContext` calls in `ffi.h`. -#![expect(unsafe_code, reason = "the cxx bridge expands to unsafe FFI glue")] +#![allow( + unsafe_code, + reason = "the cxx bridge expands to unsafe FFI glue, and io_callback_run() opens a scope C++ \ + has set up" +)] + +use std::any::Any; +use std::future::Future; +use std::pin::Pin; + +use jsg::v8; #[cxx::bridge(namespace = "workerd::io_context")] pub mod bridge { unsafe extern "C++" { include!("workerd/io/io-context/ffi.h"); + type IoOwnedObject; + fn io_context_no_request_message() -> &'static str; fn io_context_try_current_id() -> KjMaybe; fn io_context_now(id: u64) -> Result; fn io_context_write_logfwdr(id: u64, channel: u32, message: &[u8]) -> Result<()>; + fn io_context_await_io( + id: u64, + future: Box, + callback: Box, + ) -> Result<()>; + fn io_context_add_object(id: u64, slot: Box) -> Result>; } + + extern "Rust" { + type IoFuture; + type IoOutput; + type IoCallback; + type IoOwnedSlot; + + async fn io_future_wait(future: Box) -> Box; + unsafe fn io_callback_run(isolate: usize, callback: Box, output: Box); + } +} +/// The future that `IoContext::await_io()` waits for, with its output type erased. +pub struct IoFuture(pub Pin> + Send>>); + +/// The output of an [`IoFuture`]. +pub struct IoOutput(Box); + +/// What [`IoCallback`] calls. +pub type IoCallbackFn = Box FnOnce(&mut jsg::Lock<'s>, Box)>; + +/// The callback that `IoContext::await_io()` runs with the output of its future. +pub struct IoCallback(pub IoCallbackFn); + +/// A value that `IoContext::add_object()` hands to C++, which drops it with the `IoContext`. +/// Dropping it drops the value it holds, unless Rust took the value first. +pub struct IoOwnedSlot(#[expect(dead_code, reason = "only dropped")] pub Box); + +async fn io_future_wait(future: Box) -> Box { + Box::new(IoOutput(future.0.await)) +} + +/// Runs `callback` with `output`, from `io_context_await_io()` in the `IoContext`, with the +/// isolate locked, a context entered and a `HandleScope` open. +unsafe fn io_callback_run(isolate: usize, callback: Box, output: Box) { + let scope = std::pin::pin!( + // SAFETY: C++ calls this with the isolate locked, a context entered and a HandleScope + // open, which outlive the scope. + unsafe { jsg::CallbackScope::new(isolate as *mut v8::ffi::Isolate) } + ); + // A panic throws an internal error and terminates execution, as in a `#[jsg_method]`. + let _ = jsg::catch_panic(scope, |lock| (callback.0)(lock, output.0)); } diff --git a/src/workerd/io/io-context/io-context-test.c++ b/src/workerd/io/io-context/io-context-test.c++ index 73f68d822b3..9495217a7bc 100644 --- a/src/workerd/io/io-context/io-context-test.c++ +++ b/src/workerd/io/io-context/io-context-test.c++ @@ -5,6 +5,7 @@ // Tests the io-context crate inside and outside a request of a real worker. The Rust half is // test_bridge.rs. +#include #include #include #include @@ -17,6 +18,8 @@ #include #include +#include +#include namespace workerd::io_context::test { namespace { @@ -167,6 +170,172 @@ KJ_TEST("an IoContext id works only while its IoContext is current") { fixture.drainAndDestroy(kj::mv(first)); } +KJ_TEST("await_io() resumes the request when its future completes on another thread") { + TestFixture fixture; + auto request = fixture.newIncomingRequest(); + fixture.enterContext(*request, + [&](const TestFixture::Environment& env) { await_from_thread(isolateOf(env), 42); }); + KJ_EXPECT(await_observation() == AwaitObservation::Pending); + + // The other thread wakes this one's event loop. + for (auto i KJ_UNUSED: kj::zeroTo(2000)) { + if (await_observation() != AwaitObservation::Pending) break; + fixture.pollEventLoop(); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + KJ_EXPECT(await_observation() == AwaitObservation::RanInRequest); + KJ_EXPECT(awaited_value() == 42); + fixture.drainAndDestroy(kj::mv(request)); +} + +// The value of TEST_KEY in the frame record_async_context() last saw, or "none". +kj::String recordedAsyncContext = kj::str("not recorded"); +kj::Arc TEST_KEY = + kj::arc(); + +} // namespace + +void record_async_context(uintptr_t isolate) { + auto& js = jsg::Lock::from(reinterpret_cast(isolate)); + recordedAsyncContext = kj::str("none"); + KJ_IF_SOME(frame, jsg::AsyncContextFrame::current(js)) { + KJ_IF_SOME(value, frame.get(*TEST_KEY)) { + recordedAsyncContext = jsg::JsValue(value.getHandle(js)).toString(js); + } + } +} + +namespace { + +KJ_TEST("await_io() runs its callback in the async context frame that was current") { + TestFixture fixture; + auto request = fixture.newIncomingRequest(); + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + jsg::AsyncContextFrame::StorageScope scope(env.js, TEST_KEY.addRef(), + env.js.v8Ref(v8::Local(env.js.str("of the caller"_kj)))); + await_and_record_async_context(isolateOf(env)); + }); + fixture.pollEventLoop(); + KJ_EXPECT(recordedAsyncContext == "of the caller", recordedAsyncContext); + fixture.drainAndDestroy(kj::mv(request)); +} + +// Records the KJ-to-JS bridges an async tracker creates, and the causes of turns. +class BridgeListener final: public AsyncTraceListener { + public: + explicit BridgeListener(kj::Vector& events): events(events) {} + + void onContextBegin(uint64_t, uint64_t) override {} + void onInit(uint64_t, const AsyncInitEvent& e) override { + if (e.kind == AsyncKind::KJ_TO_JS) events.add(kj::str("bridge ", e.id, " ", e.name)); + } + void onSettle(uint64_t, AsyncId, AsyncOutcome, uint64_t) override {} + void onBefore(uint64_t, AsyncId, uint64_t) override {} + void onAfter(uint64_t, AsyncId, uint64_t) override {} + void onDestroy(uint64_t, AsyncId, uint64_t) override {} + void onTurn(uint64_t, const AsyncTurn& turn) override { + events.add(kj::str("turn cause=", turn.cause)); + } + void onContextEnd(uint64_t, uint64_t, const AsyncContextStats&) override {} + + private: + kj::Vector& events; +}; + +class BridgeObserver final: public IsolateObserver { + public: + explicit BridgeObserver(kj::Vector& events): events(events) {} + + kj::Maybe getAsyncTraceConfig() const override { + return AsyncTraceConfig{}; + } + void addAsyncTraceSinks( + AsyncTraceSinks& sinks, kj::StringPtr, kj::Maybe) const override { + sinks.add(kj::heap(events)); + } + + private: + kj::Vector& events; +}; + +KJ_TEST("await_io() is recorded for async tracing, as awaitIo() is") { + kj::Vector events; + TestFixture fixture({ + .useRealTimers = false, + .isolateObserver = kj::atomicRefcounted(events), + }); + auto request = fixture.newIncomingRequest(); + fixture.enterContext(*request, + [&](const TestFixture::Environment& env) { await_and_record_async_context(isolateOf(env)); }); + fixture.pollEventLoop(); + fixture.drainAndDestroy(kj::mv(request)); + + // The wait is a bridge, and causes the turn that runs the callback. + kj::Maybe bridgeId; + for (auto& event: events) { + if (event.startsWith("bridge ") && event.endsWith(" awaitIo")) { + bridgeId = kj::str(event.first(event.size() - 8).slice(7)); + } + } + auto& id = KJ_ASSERT_NONNULL(bridgeId, kj::strArray(events, "\n")); + bool causedTurn = false; + for (auto& event: events) { + if (event == kj::str("turn cause=", id)) causedTurn = true; + } + KJ_EXPECT(causedTurn, kj::strArray(events, "\n")); +} + +KJ_TEST("await_io() drops its future and callback when the request ends first") { + TestFixture fixture; + auto request = fixture.newIncomingRequest(); + fixture.enterContext( + *request, [&](const TestFixture::Environment& env) { await_forever(isolateOf(env)); }); + fixture.pollEventLoop(); + KJ_EXPECT(await_observation() == AwaitObservation::Pending); + fixture.drainAndDestroy(kj::mv(request)); + KJ_EXPECT(await_observation() == AwaitObservation::Dropped); +} + +KJ_TEST("add_object() drops the value with the request, unless it was taken") { + TestFixture fixture; + + auto request = fixture.newIncomingRequest(); + fixture.enterContext( + *request, [&](const TestFixture::Environment& env) { add_owned_value(isolateOf(env)); }); + fixture.drainAndDestroy(kj::mv(request)); + KJ_EXPECT(owned_value_dropped()); + KJ_EXPECT(!take_owned_value()); + + request = fixture.newIncomingRequest(); + fixture.enterContext( + *request, [&](const TestFixture::Environment& env) { add_owned_value(isolateOf(env)); }); + KJ_EXPECT(!owned_value_dropped()); + KJ_EXPECT(take_owned_value()); + KJ_EXPECT(owned_value_dropped()); + fixture.drainAndDestroy(kj::mv(request)); +} + +KJ_TEST("dropping an add_object() handle drops the value, at once only if its request is current") { + TestFixture fixture; + auto request = fixture.newIncomingRequest(); + + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + add_owned_value(isolateOf(env)); + drop_owned_handle(); + KJ_EXPECT(owned_value_dropped()); + }); + + fixture.enterContext( + *request, [&](const TestFixture::Environment& env) { add_owned_value(isolateOf(env)); }); + drop_owned_handle(); + KJ_EXPECT(!owned_value_dropped()); + // The request drops the value the next time it runs. + fixture.enterContext( + *request, [&](const TestFixture::Environment&) { KJ_EXPECT(owned_value_dropped()); }); + + fixture.drainAndDestroy(kj::mv(request)); +} + constexpr kj::StringPtr PROBE_SPECIFIER = "test:io-context-probe"_kj; KJ_TEST("the legacy registry creates a Rust module's exports outside the request resolving it") { diff --git a/src/workerd/io/io-context/lib.rs b/src/workerd/io/io-context/lib.rs index 92c6d058c9d..a40b2ff7c03 100644 --- a/src/workerd/io/io-context/lib.rs +++ b/src/workerd/io/io-context/lib.rs @@ -10,12 +10,19 @@ mod ffi; +use std::any::Any; +use std::future::Future; use std::marker::PhantomData; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::sync::PoisonError; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; use ffi::bridge; +use kj_rs::KjOwn; /// The `IoContext` of the current request. /// @@ -81,6 +88,114 @@ impl<'s> IoContext<'s> { pub fn write_logfwdr(&self, channel: u32, message: &[u8]) -> jsg::Result<()> { Ok(bridge::io_context_write_logfwdr(self.id, channel, message)?) } + + /// Waits for `future` as part of this request, then calls `callback` with its output, in + /// this `IoContext`, with a new lock. `future` may complete on any thread; `callback` runs on + /// this `IoContext`'s thread. As for C++ `IoContext::awaitIo()`, which shares its + /// implementation, `callback` runs in the async context frame that was current when + /// `await_io` was called, the one the I/O was started in, rather than whichever is current + /// when `future` completes. In a Durable Object it likewise runs in the critical section, if + /// any, of that call. The wait is recorded for async tracing. + /// + /// If the request is canceled first, both are dropped without `callback` running. `callback` + /// returns nothing: what it settles, such as a promise, it settles itself, and an exception + /// it leaves pending is a termination. + /// + /// # Errors + /// Fails with an internal error if this `IoContext` is no longer the current one. + pub fn await_io(&self, future: F, callback: C) -> jsg::Result<()> + where + T: Send + 'static, + F: Future + Send + 'static, + C: for<'t> FnOnce(&mut jsg::Lock<'t>, T) + 'static, + { + let future = ffi::IoFuture(Box::pin(async move { + let output: Box = Box::new(future.await); + output + })); + let callback = ffi::IoCallback(Box::new(move |lock, output| { + // `io_future_wait()` produced `output` from this `F`. + let output = output + .downcast::() + .unwrap_or_else(|_| unreachable!("the output of await_io()'s future")); + callback(lock, *output); + })); + Ok(bridge::io_context_await_io( + self.id, + Box::new(future), + Box::new(callback), + )?) + } + + /// Gives this `IoContext` ownership of `value`, as C++ `IoContext::addObject()` does: it is + /// dropped when the `IoContext` is destroyed, or when the returned handle is dropped, unless + /// [`IoOwned::take`] has taken it. How the value is reached differs from C++; see + /// [`IoOwned`]. + /// + /// `value` should hold no JavaScript handles: nothing traces it. + /// + /// # Errors + /// Fails with an internal error if this `IoContext` is no longer the current one. + pub fn add_object(&self, value: T) -> jsg::Result> { + let slot = Arc::new(Mutex::new(Some(value))); + let object = bridge::io_context_add_object( + self.id, + Box::new(ffi::IoOwnedSlot(Box::new(SlotRef(Arc::clone(&slot))))), + )?; + Ok(IoOwned { + slot, + _object: object, + }) + } +} + +/// A value an `IoContext` owns, from [`IoContext::add_object`]. +/// +/// Like C++ `IoOwn`, which `IoContext::addObject()` returns, it ties the value to its +/// `IoContext`, but the value can only be taken out, not used in place. Compared with `IoOwn`: +/// +/// - **When the value is dropped:** the same, when the handle is dropped or the `IoContext` is +/// destroyed, unless [`take`](Self::take) took it first. As with `IoOwn`, dropping the handle +/// while its `IoContext` is not current leaves the value to be dropped later, by the +/// `IoContext`. +/// - **Access:** `IoOwn` is dereferenced in place, and only while its `IoContext` is current. +/// Here the value is moved out with [`take`](Self::take), whichever `IoContext`, if any, is +/// current. +/// - **After the `IoContext` is destroyed:** dereferencing an `IoOwn` throws; here +/// [`take`](Self::take) returns `None`. +/// - **Threads:** an `IoOwn` is used on its `IoContext`'s thread, but may be dropped on any; the +/// value is then dropped on the `IoContext`'s thread. `IoOwned` is neither `Send` nor `Sync`, +/// so it stays on the thread that created it, which is the `IoContext`'s. +pub struct IoOwned { + slot: Arc>>, + // Dropping it drops the value, if the IoContext has not already. + _object: KjOwn, +} + +impl IoOwned { + /// Takes the value, or returns `None` if it was taken, or its `IoContext` has been + /// destroyed. + pub fn take(&self) -> Option { + lock_slot(&self.slot).take() + } +} + +/// Holds no JavaScript handles; see [`IoContext::add_object`]. +impl jsg::Traced for IoOwned {} + +/// The reference to an [`IoOwned`]'s value that its `IoContext` holds. +struct SlotRef(Arc>>); + +impl Drop for SlotRef { + fn drop(&mut self) { + let value = lock_slot(&self.0).take(); + drop(value); + } +} + +fn lock_slot(slot: &Mutex>) -> MutexGuard<'_, Option> { + // Nothing panics while holding the lock: the value is taken out of it and dropped after. + slot.lock().unwrap_or_else(PoisonError::into_inner) } #[cfg(test)] diff --git a/src/workerd/io/io-context/test-bridge.h b/src/workerd/io/io-context/test-bridge.h new file mode 100644 index 00000000000..3760746703d --- /dev/null +++ b/src/workerd/io/io-context/test-bridge.h @@ -0,0 +1,15 @@ +// Copyright (c) 2026 Cloudflare, Inc. +// Licensed under the Apache 2.0 license found in the LICENSE file or at: +// https://opensource.org/licenses/Apache-2.0 + +#pragma once + +#include + +// What test_bridge.rs calls back in io-context-test.c++, which defines it. +namespace workerd::io_context::test { + +// Records the async context frame that is current, in the isolate `isolate`. +void record_async_context(uintptr_t isolate); + +} // namespace workerd::io_context::test diff --git a/src/workerd/io/io-context/test_bridge.rs b/src/workerd/io/io-context/test_bridge.rs index 0b2b0609b5d..c7888d1ddc4 100644 --- a/src/workerd/io/io-context/test_bridge.rs +++ b/src/workerd/io/io-context/test_bridge.rs @@ -9,10 +9,15 @@ //! thread, with a context entered and a `v8::HandleScope` open. use std::cell::Cell; +use std::cell::RefCell; use std::pin::Pin; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; use std::time::UNIX_EPOCH; use io_context::IoContext; +use io_context::IoOwned; use jsg::ToJS; use jsg::v8; @@ -41,9 +46,38 @@ mod ffi { type ModuleRegistry = jsg::v8::ffi::ModuleRegistry; } + /// What the callback of the last `await_io()` started by `await_from_thread()` or + /// `await_forever()` did. + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum AwaitObservation { + /// It has not run, and still exists. + Pending, + /// It ran in a request, with the value the future produced. + RanInRequest, + /// It ran outside any request. + RanOutsideRequest, + /// It was dropped without running. + Dropped, + } + + unsafe extern "C++" { + include!("workerd/io/io-context/test-bridge.h"); + + fn record_async_context(isolate: usize); + } + extern "Rust" { + unsafe fn await_and_record_async_context(isolate: usize) -> Result<()>; fn register_probe_module(registry: Pin<&mut ModuleRegistry>); fn take_probe_observation() -> ProbeObservation; + unsafe fn await_from_thread(isolate: usize, value: i32) -> Result<()>; + unsafe fn await_forever(isolate: usize) -> Result<()>; + fn await_observation() -> AwaitObservation; + fn awaited_value() -> i32; + unsafe fn add_owned_value(isolate: usize) -> Result<()>; + fn owned_value_dropped() -> bool; + fn take_owned_value() -> bool; + fn drop_owned_handle(); unsafe fn has_current(isolate: usize) -> bool; unsafe fn now_nanos(isolate: usize) -> Result; @@ -109,3 +143,133 @@ fn register_probe_module(registry: Pin<&mut ffi::ModuleRegistry>) { fn take_probe_observation() -> ffi::ProbeObservation { PROBE_OBSERVATION.with(|cell| cell.replace(ffi::ProbeObservation::NotRun)) } + +thread_local! { + static AWAIT_OBSERVATION: Cell = + const { Cell::new(ffi::AwaitObservation::Pending) }; + static AWAITED_VALUE: Cell = const { Cell::new(0) }; + static OWNED: RefCell>> = const { RefCell::new(None) }; +} + +/// Sets `AWAIT_OBSERVATION` to `Dropped` if dropped while it is still `Pending`, that is, if the +/// callback holding it never ran. +struct DropObserver; + +impl Drop for DropObserver { + fn drop(&mut self) { + AWAIT_OBSERVATION.with(|cell| { + if cell.get() == ffi::AwaitObservation::Pending { + cell.set(ffi::AwaitObservation::Dropped); + } + }); + } +} + +fn record_await(lock: &mut jsg::Lock<'_>, value: i32) { + let observation = if IoContext::try_current(lock).is_some() { + ffi::AwaitObservation::RanInRequest + } else { + ffi::AwaitObservation::RanOutsideRequest + }; + AWAITED_VALUE.with(|cell| cell.set(value)); + AWAIT_OBSERVATION.with(|cell| cell.set(observation)); +} + +/// In a request, awaits a value that another thread sends after a delay. +unsafe fn await_from_thread(isolate: usize, value: i32) -> Result<()> { + AWAIT_OBSERVATION.with(|cell| cell.set(ffi::AwaitObservation::Pending)); + // SAFETY: see the module documentation. + jsg::callback_scope!(unsafe let lock, isolate as *mut v8::ffi::Isolate); + let (sender, receiver) = futures::channel::oneshot::channel(); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(20)); + let _ = sender.send(value); + }); + let observer = DropObserver; + IoContext::current(lock) + .and_then(|context| { + context.await_io(receiver, move |lock, received| { + let _observer = observer; + record_await(lock, received.unwrap_or(-1)); + }) + }) + .map_err(|err| to_io(&err)) +} + +/// In a request, awaits a future that never completes. +unsafe fn await_forever(isolate: usize) -> Result<()> { + AWAIT_OBSERVATION.with(|cell| cell.set(ffi::AwaitObservation::Pending)); + // SAFETY: see the module documentation. + jsg::callback_scope!(unsafe let lock, isolate as *mut v8::ffi::Isolate); + let observer = DropObserver; + IoContext::current(lock) + .and_then(|context| { + context.await_io(std::future::pending::(), move |lock, value| { + let _observer = observer; + record_await(lock, value); + }) + }) + .map_err(|err| to_io(&err)) +} + +fn await_observation() -> ffi::AwaitObservation { + AWAIT_OBSERVATION.with(Cell::get) +} + +fn awaited_value() -> i32 { + AWAITED_VALUE.with(Cell::get) +} + +static OWNED_DROPPED: AtomicBool = AtomicBool::new(false); + +/// Sets `OWNED_DROPPED` when dropped. +struct DropFlag; + +impl Drop for DropFlag { + fn drop(&mut self) { + OWNED_DROPPED.store(true, Ordering::SeqCst); + } +} + +/// Gives the current request a value, whose handle this thread keeps. +unsafe fn add_owned_value(isolate: usize) -> Result<()> { + OWNED_DROPPED.store(false, Ordering::SeqCst); + // SAFETY: see the module documentation. + jsg::callback_scope!(unsafe let lock, isolate as *mut v8::ffi::Isolate); + let owned = IoContext::current(lock) + .and_then(|context| context.add_object(DropFlag)) + .map_err(|err| to_io(&err))?; + OWNED.with(|cell| *cell.borrow_mut() = Some(owned)); + Ok(()) +} + +/// Whether the value `add_owned_value()` gave its request has been dropped. +fn owned_value_dropped() -> bool { + OWNED_DROPPED.load(Ordering::SeqCst) +} + +/// Takes the value `add_owned_value()` gave its request and drops it, then drops the handle; +/// false if there was no value. +fn take_owned_value() -> bool { + let owned = OWNED.with(|cell| cell.borrow_mut().take()); + owned.as_ref().and_then(IoOwned::take).is_some() +} + +/// Drops the handle `add_owned_value()` kept, without taking its value. +fn drop_owned_handle() { + OWNED.with(|cell| cell.borrow_mut().take()); +} + +/// In a request, awaits a future that is already complete, and from its callback has C++ record +/// the async context frame that is current. +unsafe fn await_and_record_async_context(isolate: usize) -> Result<()> { + // SAFETY: see the module documentation. + jsg::callback_scope!(unsafe let lock, isolate as *mut v8::ffi::Isolate); + IoContext::current(lock) + .and_then(|context| { + context.await_io(std::future::ready(()), move |_lock, ()| { + ffi::record_async_context(isolate); + }) + }) + .map_err(|err| to_io(&err)) +}