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(), 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)) +}