From f94cb4eb394c6411dbe17332341338430a7f4c94 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/4] 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 5c54b39afb1..3a95ed57779 100644 --- a/src/rust/jsg-test/tests/mod.rs +++ b/src/rust/jsg-test/tests/mod.rs @@ -17,6 +17,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 fb40118334c..a156550917e 100644 --- a/src/rust/jsg/lib.rs +++ b/src/rust/jsg/lib.rs @@ -453,13 +453,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 05d31f89d97fbfc1603cb1d24546dbc7436bf29f 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/4] 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 | 58 +++++- src/workerd/io/io-context/io-context-test.c++ | 170 ++++++++++++++++ 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, 683 insertions(+), 78 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 136e61c6aac..a15076d6e10 100644 --- a/src/workerd/io/io-context.h +++ b/src/workerd/io/io-context.h @@ -668,6 +668,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. @@ -1445,21 +1466,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 { @@ -1706,7 +1723,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 @@ -1715,29 +1732,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 @@ -1753,35 +1761,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. @@ -1791,38 +1850,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 @@ -1830,6 +1860,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 7804781537b..36e93e40904 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 @@ -76,4 +78,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 8f42e2b1ca0..986e7e2da53 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 9e540e7041f..8dc4bff325d 100644 --- a/src/workerd/io/io-context/ffi.rs +++ b/src/workerd/io/io-context/ffi.rs @@ -6,12 +6,16 @@ #![allow( unsafe_code, - reason = "the cxx bridge expands to unsafe FFI glue, and current_id() reads through the \ - IoContext reference it returns" + reason = "the cxx bridge expands to unsafe FFI glue, current_id() reads through the \ + IoContext reference it returns, 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 { #[namespace = "workerd"] @@ -24,11 +28,29 @@ pub mod bridge { unsafe extern "C++" { include!("workerd/io/io-context/ffi.h"); + type IoOwnedObject; + /// The caller chooses the lifetime of the result; see `current_id()`. unsafe fn io_context_try_current<'a>() -> KjMaybe>; fn io_context_id(context: &IoContext) -> u64; 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); } } @@ -40,3 +62,35 @@ pub fn current_id() -> Option { unsafe { bridge::io_context_try_current() }.into(); context.map(|context| bridge::io_context_id(&context)) } + +/// 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 c09a1b51683..f1f8f2421fa 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 @@ -16,6 +17,9 @@ #include #include +#include +#include + namespace workerd::io_context::test { namespace { @@ -179,6 +183,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 aef86d47fca..83ddcc0c5a6 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 message of the error C++ `IoContext::current()` throws outside a request. // A copy of `kAsyncIoErrorMessage` in src/workerd/io/io-context.c++. When changing this message, @@ -88,6 +95,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 f4f02e1e4d1..33fa88aeedf 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 current_error(isolate: usize) -> String; @@ -120,3 +154,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)) +} From 59c270f4d999b2d22a520649311590499be4dc21 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bartek=20Iwa=C5=84czuk?= Date: Tue, 27 Oct 2026 19:43:26 +0100 Subject: [PATCH 3/4] Add trace spans to the Rust io-context crate IoContext::make_trace_span() and make_user_trace_span() start spans as C++ makeTraceSpan() and makeUserTraceSpan() do, returning a SpanBuilder or TraceContext that holds the C++ object and ends the span when dropped. SpanBuilder::set_tag(), SpanBuilder::new_child() and TraceContext::set_tag() mirror their C++ counterparts. Span names and tag keys are &'static CStr, written as c"..." literals. C++ wraps them without copying, as it does its own string literals. A tag value can be a bool, an integer, a float, a literal, or a borrowed string, which is copied only if the span is observed. A user trace span started before await_io() is adopted by it for async tracing, as for awaitIo(), and a test checks this. --- src/workerd/io/io-context/BUILD.bazel | 1 + src/workerd/io/io-context/ffi.c++ | 73 ++++++++ src/workerd/io/io-context/ffi.h | 26 +++ src/workerd/io/io-context/ffi.rs | 52 +++++- src/workerd/io/io-context/io-context-test.c++ | 167 +++++++++++++++++- src/workerd/io/io-context/lib.rs | 26 +++ src/workerd/io/io-context/test_bridge.rs | 64 +++++++ src/workerd/io/io-context/trace.rs | 152 ++++++++++++++++ 8 files changed, 559 insertions(+), 2 deletions(-) create mode 100644 src/workerd/io/io-context/trace.rs diff --git a/src/workerd/io/io-context/BUILD.bazel b/src/workerd/io/io-context/BUILD.bazel index f1d0a74b8bc..0fe05e677f0 100644 --- a/src/workerd/io/io-context/BUILD.bazel +++ b/src/workerd/io/io-context/BUILD.bazel @@ -10,6 +10,7 @@ wd_rust_crate( "ffi.rs", "lib.rs", "lib-test.rs", + "trace.rs", ], cxx_bridge_deps = ["//src/workerd/io"], cxx_bridge_hdrs = ["ffi.h"], diff --git a/src/workerd/io/io-context/ffi.c++ b/src/workerd/io/io-context/ffi.c++ index 36e93e40904..e0421037917 100644 --- a/src/workerd/io/io-context/ffi.c++ +++ b/src/workerd/io/io-context/ffi.c++ @@ -27,6 +27,27 @@ IoContext& currentWithId(uint64_t id) { return context; } +// Wraps a `Literal` (see ffi.h) without copying it. +kj::ConstString wrapLiteral(Literal bytes) { + KJ_ASSERT(bytes.size() > 0 && bytes.back() == '\0', "the literal is not NUL-terminated"); + return kj::ConstString(reinterpret_cast(bytes.data()), bytes.size() - 1, + kj::NullArrayDisposer::instance); +} + +// Sets a tag on a SpanBuilder or TraceContext, if the span is observed; see ffi.h for which +// strings are copied. +template +void setTag(Span& span, Literal key, Value value) { + if (!span.isObserved()) return; + if constexpr (kj::isSameType()) { + span.setTag(wrapLiteral(key), kj::str(kj::arrayPtr(value.data(), value.size()))); + } else if constexpr (kj::isSameType()) { + span.setTag(wrapLiteral(key), wrapLiteral(value)); + } else { + span.setTag(wrapLiteral(key), value); + } +} + } // namespace kj::Maybe io_context_try_current() { @@ -98,4 +119,56 @@ kj::Own io_context_add_object(uint64_t id, ::rust::Box(kj::heap(kj::mv(object))); } +kj::Own io_context_make_trace_span(uint64_t id, Literal name) { + return kj::heap(currentWithId(id).makeTraceSpan(wrapLiteral(name))); +} + +kj::Own io_context_make_user_trace_span(uint64_t id, Literal name) { + return kj::heap(currentWithId(id).makeUserTraceSpan(wrapLiteral(name))); +} + +kj::Own span_builder_new_child(SpanBuilder& span, Literal name) { + return kj::heap(span.newChild(wrapLiteral(name))); +} + +void span_builder_set_tag_bool(SpanBuilder& span, Literal key, bool value) { + setTag(span, key, value); +} + +void span_builder_set_tag_int(SpanBuilder& span, Literal key, int64_t value) { + setTag(span, key, value); +} + +void span_builder_set_tag_float(SpanBuilder& span, Literal key, double value) { + setTag(span, key, value); +} + +void span_builder_set_tag_str(SpanBuilder& span, Literal key, ::rust::Str value) { + setTag(span, key, value); +} + +void span_builder_set_tag_literal(SpanBuilder& span, Literal key, Literal value) { + setTag(span, key, value); +} + +void trace_context_set_tag_bool(TraceContext& context, Literal key, bool value) { + setTag(context, key, value); +} + +void trace_context_set_tag_int(TraceContext& context, Literal key, int64_t value) { + setTag(context, key, value); +} + +void trace_context_set_tag_float(TraceContext& context, Literal key, double value) { + setTag(context, key, value); +} + +void trace_context_set_tag_str(TraceContext& context, Literal key, ::rust::Str value) { + setTag(context, key, value); +} + +void trace_context_set_tag_literal(TraceContext& context, Literal key, Literal value) { + setTag(context, key, value); +} + } // namespace workerd::io_context diff --git a/src/workerd/io/io-context/ffi.h b/src/workerd/io/io-context/ffi.h index 986e7e2da53..6c3f9752776 100644 --- a/src/workerd/io/io-context/ffi.h +++ b/src/workerd/io/io-context/ffi.h @@ -54,4 +54,30 @@ class IoOwnedObject { // IoContext::addObject() for `slot`. kj::Own io_context_add_object(uint64_t id, ::rust::Box slot); +// Span names, tag keys and literal tag values below are passed as `literal` bytes: a +// NUL-terminated string including its NUL, which must outlive the span, as a string literal +// does. They are not copied. +using Literal = ::rust::Slice; + +// IoContext::makeTraceSpan() and makeUserTraceSpan(), in the current async context frame. +kj::Own io_context_make_trace_span(uint64_t id, Literal name); +kj::Own io_context_make_user_trace_span(uint64_t id, Literal name); + +// SpanBuilder::newChild(). +kj::Own span_builder_new_child(SpanBuilder& span, Literal name); + +// SpanBuilder::setTag() and TraceContext::setTag(). On a span that is not observed, these do +// nothing. On one that is, a `::rust::Str` value is copied, since it is not NUL-terminated and +// may not outlive the call; the key and a `Literal` value are not. +void span_builder_set_tag_bool(SpanBuilder& span, Literal key, bool value); +void span_builder_set_tag_int(SpanBuilder& span, Literal key, int64_t value); +void span_builder_set_tag_float(SpanBuilder& span, Literal key, double value); +void span_builder_set_tag_str(SpanBuilder& span, Literal key, ::rust::Str value); +void span_builder_set_tag_literal(SpanBuilder& span, Literal key, Literal value); +void trace_context_set_tag_bool(TraceContext& context, Literal key, bool value); +void trace_context_set_tag_int(TraceContext& context, Literal key, int64_t value); +void trace_context_set_tag_float(TraceContext& context, Literal key, double value); +void trace_context_set_tag_str(TraceContext& context, Literal key, ::rust::Str value); +void trace_context_set_tag_literal(TraceContext& context, Literal key, Literal value); + } // namespace workerd::io_context diff --git a/src/workerd/io/io-context/ffi.rs b/src/workerd/io/io-context/ffi.rs index 8dc4bff325d..8a496e7909d 100644 --- a/src/workerd/io/io-context/ffi.rs +++ b/src/workerd/io/io-context/ffi.rs @@ -7,7 +7,8 @@ #![allow( unsafe_code, reason = "the cxx bridge expands to unsafe FFI glue, current_id() reads through the \ - IoContext reference it returns, and io_callback_run() opens a scope C++ has set up" + IoContext reference it returns, io_callback_run() opens a scope C++ has set up, and \ + the trace span calls take literals that C++ does not copy" )] use std::any::Any; @@ -23,6 +24,8 @@ pub mod bridge { include!("workerd/io/io-context.h"); type IoContext; + type SpanBuilder; + type TraceContext; } unsafe extern "C++" { @@ -41,6 +44,53 @@ pub mod bridge { callback: Box, ) -> Result<()>; fn io_context_add_object(id: u64, slot: Box) -> Result>; + + // The `&[u8]` names, keys and literal values below are NUL-terminated strings with the + // NUL, which C++ requires and does not copy, so they must outlive the spans: the + // `&'static CStr`s that `trace.rs` passes. + unsafe fn io_context_make_trace_span(id: u64, name: &[u8]) -> Result>; + unsafe fn io_context_make_user_trace_span( + id: u64, + name: &[u8], + ) -> Result>; + unsafe fn span_builder_new_child( + span: Pin<&mut SpanBuilder>, + name: &[u8], + ) -> KjOwn; + unsafe fn span_builder_set_tag_bool(span: Pin<&mut SpanBuilder>, key: &[u8], value: bool); + unsafe fn span_builder_set_tag_int(span: Pin<&mut SpanBuilder>, key: &[u8], value: i64); + unsafe fn span_builder_set_tag_float(span: Pin<&mut SpanBuilder>, key: &[u8], value: f64); + unsafe fn span_builder_set_tag_str(span: Pin<&mut SpanBuilder>, key: &[u8], value: &str); + unsafe fn span_builder_set_tag_literal( + span: Pin<&mut SpanBuilder>, + key: &[u8], + value: &[u8], + ); + unsafe fn trace_context_set_tag_bool( + context: Pin<&mut TraceContext>, + key: &[u8], + value: bool, + ); + unsafe fn trace_context_set_tag_int( + context: Pin<&mut TraceContext>, + key: &[u8], + value: i64, + ); + unsafe fn trace_context_set_tag_float( + context: Pin<&mut TraceContext>, + key: &[u8], + value: f64, + ); + unsafe fn trace_context_set_tag_str( + context: Pin<&mut TraceContext>, + key: &[u8], + value: &str, + ); + unsafe fn trace_context_set_tag_literal( + context: Pin<&mut TraceContext>, + key: &[u8], + value: &[u8], + ); } extern "Rust" { diff --git a/src/workerd/io/io-context/io-context-test.c++ b/src/workerd/io/io-context/io-context-test.c++ index f1f8f2421fa..39a65a6624c 100644 --- a/src/workerd/io/io-context/io-context-test.c++ +++ b/src/workerd/io/io-context/io-context-test.c++ @@ -17,6 +17,7 @@ #include #include +#include #include #include @@ -233,7 +234,7 @@ KJ_TEST("await_io() runs its callback in the async context frame that was curren fixture.drainAndDestroy(kj::mv(request)); } -// Records the KJ-to-JS bridges an async tracker creates, and the causes of turns. +// Records the KJ-to-JS bridges and operations an async tracker creates, and the causes of turns. class BridgeListener final: public AsyncTraceListener { public: explicit BridgeListener(kj::Vector& events): events(events) {} @@ -241,6 +242,7 @@ class BridgeListener final: public AsyncTraceListener { 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)); + if (e.kind == AsyncKind::OPERATION) events.add(kj::str("operation ", e.id, " ", e.name)); } void onSettle(uint64_t, AsyncId, AsyncOutcome, uint64_t) override {} void onBefore(uint64_t, AsyncId, uint64_t) override {} @@ -298,6 +300,169 @@ KJ_TEST("await_io() is recorded for async tracing, as awaitIo() is") { KJ_EXPECT(causedTurn, kj::strArray(events, "\n")); } +KJ_TEST("await_io() reports its wait under the operation of a user trace span started before") { + kj::Vector events; + TestFixture fixture({ + .useRealTimers = false, + .isolateObserver = kj::atomicRefcounted(events), + }); + auto request = fixture.newIncomingRequest(); + fixture.enterContext(*request, + [&](const TestFixture::Environment& env) { await_in_user_trace_span(isolateOf(env)); }); + fixture.pollEventLoop(); + fixture.drainAndDestroy(kj::mv(request)); + + // No bridge of its own: the operation causes the turn that runs the callback. + kj::Maybe operationId; + for (auto& event: events) { + KJ_EXPECT(!event.endsWith(" awaitIo"), kj::strArray(events, "\n")); + if (event.startsWith("operation ") && event.endsWith(" rust_operation")) { + operationId = kj::str(event.first(event.size() - 15).slice(10)); + } + } + auto& id = KJ_ASSERT_NONNULL(operationId, 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")); +} + +// A span a SpanRecorder records: its name, the index of its parent, its tags as "key=type:value" +// in key order, and whether it has ended. +struct RecordedSpan { + kj::String name; + uint parent; + kj::String tags; + bool ended = false; +}; + +struct SpanRecords: public kj::Refcounted { + kj::Vector spans; + + kj::Maybe find(kj::StringPtr name) { + for (auto& span: spans) { + if (span.name == name) return span; + } + return kj::none; + } +}; + +class SpanRecorder final: public SpanObserver { + public: + SpanRecorder(kj::Rc records, uint parent) + : records(kj::mv(records)), + index(this->records->spans.size()) { + this->records->spans.add(RecordedSpan{.name = kj::str(), .parent = parent, .tags = kj::str()}); + } + + kj::Rc newChild() override { + return kj::rc(records.addRef(), index); + } + void onOpen(kj::ConstString operationName, kj::Date) override { + records->spans[index].name = kj::str(operationName); + } + void onClose(kj::Date, Span::TagMap&& tags, kj::Vector&&) override { + auto entries = KJ_MAP(entry, tags) { + kj::String value; + KJ_SWITCH_ONEOF(entry.value) { + KJ_CASE_ONEOF(str, kj::ConstString) { + value = kj::str("str:", str); + } + KJ_CASE_ONEOF(b, bool) { + value = kj::str("bool:", b); + } + KJ_CASE_ONEOF(d, double) { + value = kj::str("float:", d); + } + KJ_CASE_ONEOF(i, int64_t) { + value = kj::str("int:", i); + } + } + return kj::str(entry.key, "=", value); + }; + std::sort(entries.begin(), entries.end()); + records->spans[index].tags = kj::strArray(entries, " "); + records->spans[index].ended = true; + } + + private: + kj::Rc records; + uint index; +}; + +class SpanRequestObserver final: public RequestObserver { + public: + explicit SpanRequestObserver(kj::Rc records) + : span(kj::rc(kj::mv(records), kj::maxValue), "request"_kjc) {} + + SpanParent getSpan() override { + return SpanParent(span); + } + + private: + SpanBuilder span; +}; + +TestFixture::SetupParams recordingSpans(kj::Rc& records) { + return { + .requestObserverFactory = + kj::Function()>([&records]() -> kj::Own { + return kj::refcounted(records.addRef()); + }), + }; +} + +KJ_TEST("trace spans are recorded with their tags, under the request's span") { + auto records = kj::rc(); + TestFixture fixture(recordingSpans(records)); + auto request = fixture.newIncomingRequest(); + fixture.enterContext( + *request, [&](const TestFixture::Environment& env) { record_trace_spans(isolateOf(env)); }); + fixture.drainAndDestroy(kj::mv(request)); + + auto& root = KJ_ASSERT_NONNULL(records->find("request")); + auto rootIndex = &root - records->spans.begin(); + auto& span = KJ_ASSERT_NONNULL(records->find("rust_span")); + KJ_EXPECT(span.parent == rootIndex); + KJ_EXPECT(span.ended); + KJ_EXPECT( + span.tags == "bool=bool:true float=float:1.5 int=int:42 literal=str:literal str=str:dynamic", + span.tags); + + auto& child = KJ_ASSERT_NONNULL(records->find("rust_child")); + KJ_EXPECT(child.parent == &span - records->spans.begin()); + KJ_EXPECT(child.tags == "replaced=int:2", child.tags); + + // Its internal half; the request has no user tracing. + auto& user = KJ_ASSERT_NONNULL(records->find("rust_user_span")); + KJ_EXPECT(user.parent == rootIndex); + KJ_EXPECT(user.ended); + KJ_EXPECT(user.tags == + "bool=bool:false float=float:0.5 int=int:7 literal=str:user_literal str=str:user", + user.tags); +} + +KJ_TEST("a trace span ends when its SpanBuilder is dropped") { + auto records = kj::rc(); + TestFixture fixture(recordingSpans(records)); + auto request = fixture.newIncomingRequest(); + fixture.enterContext( + *request, [&](const TestFixture::Environment& env) { hold_trace_span(isolateOf(env)); }); + KJ_EXPECT(!KJ_ASSERT_NONNULL(records->find("held_span")).ended); + drop_held_trace_span(); + KJ_EXPECT(KJ_ASSERT_NONNULL(records->find("held_span")).ended); + fixture.drainAndDestroy(kj::mv(request)); +} + +KJ_TEST("trace spans record nothing, and tags are ignored, when the request is not traced") { + TestFixture fixture; + auto request = fixture.newIncomingRequest(); + fixture.enterContext( + *request, [&](const TestFixture::Environment& env) { record_trace_spans(isolateOf(env)); }); + fixture.drainAndDestroy(kj::mv(request)); +} + KJ_TEST("await_io() drops its future and callback when the request ends first") { TestFixture fixture; auto request = fixture.newIncomingRequest(); diff --git a/src/workerd/io/io-context/lib.rs b/src/workerd/io/io-context/lib.rs index 83ddcc0c5a6..3abfa8f412e 100644 --- a/src/workerd/io/io-context/lib.rs +++ b/src/workerd/io/io-context/lib.rs @@ -9,8 +9,10 @@ //! `jsg`. mod ffi; +mod trace; use std::any::Any; +use std::ffi::CStr; use std::future::Future; use std::marker::PhantomData; use std::sync::Arc; @@ -23,6 +25,9 @@ use std::time::UNIX_EPOCH; use ffi::bridge; use kj_rs::KjOwn; +pub use trace::SpanBuilder; +pub use trace::TagValue; +pub use trace::TraceContext; /// The message of the error C++ `IoContext::current()` throws outside a request. // A copy of `kAsyncIoErrorMessage` in src/workerd/io/io-context.c++. When changing this message, @@ -154,6 +159,27 @@ impl<'s> IoContext<'s> { _object: object, }) } + + /// Starts a span of internal tracing, as C++ `IoContext::makeTraceSpan()` does. It is a child + /// of the span current in the async context frame, or of the request's span, and ends when + /// the returned [`SpanBuilder`] is dropped. If the request is not traced, it records nothing. + /// + /// # Errors + /// Fails with an internal error if this `IoContext` is no longer the current one. + pub fn make_trace_span(&self, name: &'static CStr) -> jsg::Result { + SpanBuilder::start(self.id, name) + } + + /// Starts a span of both internal and user tracing, as C++ `IoContext::makeUserTraceSpan()` + /// does, which ends when the returned [`TraceContext`] is dropped. If the request has async + /// tracing, the span is an operation, which the next [`IoContext::await_io`] in the same turn + /// reports its wait under. + /// + /// # Errors + /// Fails with an internal error if this `IoContext` is no longer the current one. + pub fn make_user_trace_span(&self, name: &'static CStr) -> jsg::Result { + TraceContext::start(self.id, name) + } } /// A value an `IoContext` owns, from [`IoContext::add_object`]. diff --git a/src/workerd/io/io-context/test_bridge.rs b/src/workerd/io/io-context/test_bridge.rs index 33fa88aeedf..bd4ef4bb911 100644 --- a/src/workerd/io/io-context/test_bridge.rs +++ b/src/workerd/io/io-context/test_bridge.rs @@ -18,6 +18,7 @@ use std::time::UNIX_EPOCH; use io_context::IoContext; use io_context::IoOwned; +use io_context::SpanBuilder; use jsg::ToJS; use jsg::v8; @@ -78,6 +79,10 @@ mod ffi { fn owned_value_dropped() -> bool; fn take_owned_value() -> bool; fn drop_owned_handle(); + unsafe fn record_trace_spans(isolate: usize) -> Result<()>; + unsafe fn hold_trace_span(isolate: usize) -> Result<()>; + fn drop_held_trace_span(); + unsafe fn await_in_user_trace_span(isolate: usize) -> Result<()>; unsafe fn has_current(isolate: usize) -> bool; unsafe fn current_error(isolate: usize) -> String; @@ -160,6 +165,7 @@ thread_local! { const { Cell::new(ffi::AwaitObservation::Pending) }; static AWAITED_VALUE: Cell = const { Cell::new(0) }; static OWNED: RefCell>> = const { RefCell::new(None) }; + static HELD_SPAN: RefCell> = const { RefCell::new(None) }; } /// Sets `AWAIT_OBSERVATION` to `Dropped` if dropped while it is still `Pending`, that is, if the @@ -271,6 +277,64 @@ fn drop_owned_handle() { OWNED.with(|cell| cell.borrow_mut().take()); } +/// In a request, makes the spans `io-context-test.c++` expects, with a tag of each kind, and ends +/// them. +unsafe fn record_trace_spans(isolate: usize) -> Result<()> { + // SAFETY: see the module documentation. + jsg::callback_scope!(unsafe let lock, isolate as *mut v8::ffi::Isolate); + let context = IoContext::current(lock).map_err(|err| to_io(&err))?; + let mut span = context + .make_trace_span(c"rust_span") + .map_err(|err| to_io(&err))?; + span.set_tag(c"bool", true); + span.set_tag(c"int", 42_i64); + span.set_tag(c"float", 1.5); + let dynamic = String::from("dynamic"); + span.set_tag(c"str", &dynamic); + span.set_tag(c"literal", c"literal"); + let mut child = span.new_child(c"rust_child"); + child.set_tag(c"replaced", 1_i64); + child.set_tag(c"replaced", 2_i64); + let mut user = context + .make_user_trace_span(c"rust_user_span") + .map_err(|err| to_io(&err))?; + user.set_tag(c"bool", false); + user.set_tag(c"int", 7_i64); + user.set_tag(c"float", 0.5); + user.set_tag(c"str", "user"); + user.set_tag(c"literal", c"user_literal"); + Ok(()) +} + +/// In a request, starts a span that this thread keeps until `drop_held_trace_span()`. +unsafe fn hold_trace_span(isolate: usize) -> Result<()> { + // SAFETY: see the module documentation. + jsg::callback_scope!(unsafe let lock, isolate as *mut v8::ffi::Isolate); + let span = IoContext::current(lock) + .and_then(|context| context.make_trace_span(c"held_span")) + .map_err(|err| to_io(&err))?; + HELD_SPAN.with(|cell| *cell.borrow_mut() = Some(span)); + Ok(()) +} + +/// Drops the span `hold_trace_span()` started. +fn drop_held_trace_span() { + HELD_SPAN.with(|cell| cell.borrow_mut().take()); +} + +/// In a request, starts a user trace span, then awaits a future that is already complete, ending +/// the span in the callback. +unsafe fn await_in_user_trace_span(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| { + let span = context.make_user_trace_span(c"rust_operation")?; + context.await_io(std::future::ready(()), move |_lock, ()| drop(span)) + }) + .map_err(|err| to_io(&err)) +} + /// 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<()> { diff --git a/src/workerd/io/io-context/trace.rs b/src/workerd/io/io-context/trace.rs new file mode 100644 index 00000000000..8d2a6e7a63e --- /dev/null +++ b/src/workerd/io/io-context/trace.rs @@ -0,0 +1,152 @@ +// 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 + +//! Trace spans of the current request, from +//! [`IoContext::make_trace_span`](crate::IoContext::make_trace_span) and +//! [`IoContext::make_user_trace_span`](crate::IoContext::make_user_trace_span). +//! +//! Span names and tag keys are `&'static CStr`s, written as `c"..."` literals, which C++ uses +//! without copying, as it does its own string literals. + +#![allow( + unsafe_code, + reason = "the bridge takes names, keys and literal values that C++ does not copy, which \ + must therefore outlive the span" +)] + +use std::ffi::CStr; + +use kj_rs::KjOwn; + +use crate::bridge; + +/// The value of a span tag. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum TagValue<'a> { + Bool(bool), + Int(i64), + Float(f64), + /// A string, which is copied if the span is recorded. + Str(&'a str), + /// A string literal, which is not copied. + Literal(&'static CStr), +} + +impl From for TagValue<'_> { + fn from(value: bool) -> Self { + Self::Bool(value) + } +} + +impl From for TagValue<'_> { + fn from(value: i64) -> Self { + Self::Int(value) + } +} + +impl From for TagValue<'_> { + fn from(value: f64) -> Self { + Self::Float(value) + } +} + +impl<'a> From<&'a str> for TagValue<'a> { + fn from(value: &'a str) -> Self { + Self::Str(value) + } +} + +impl<'a> From<&'a String> for TagValue<'a> { + fn from(value: &'a String) -> Self { + Self::Str(value) + } +} + +impl From<&'static CStr> for TagValue<'_> { + fn from(value: &'static CStr) -> Self { + Self::Literal(value) + } +} + +/// A span being recorded, like C++ `SpanBuilder`, which it holds. The span ends when this is +/// dropped. It is neither `Send` nor `Sync`. +pub struct SpanBuilder(KjOwn); + +impl SpanBuilder { + /// [`IoContext::make_trace_span`](crate::IoContext::make_trace_span), for the `IoContext` + /// with id `context_id`. + pub(crate) fn start(context_id: u64, name: &'static CStr) -> jsg::Result { + // SAFETY: `name` is `'static`. + let span = unsafe { bridge::io_context_make_trace_span(context_id, literal(name)) }?; + Ok(Self(span)) + } + + /// Sets tag `key` to `value`, replacing any value it had, as C++ `SpanBuilder::setTag()` + /// does. Does nothing if the span is not recorded. + pub fn set_tag<'v>(&mut self, key: &'static CStr, value: impl Into>) { + let (span, key) = (self.0.pin_mut(), literal(key)); + // SAFETY: `key`, and a `Literal` value, are `'static`. + unsafe { + match value.into() { + TagValue::Bool(value) => bridge::span_builder_set_tag_bool(span, key, value), + TagValue::Int(value) => bridge::span_builder_set_tag_int(span, key, value), + TagValue::Float(value) => bridge::span_builder_set_tag_float(span, key, value), + TagValue::Str(value) => bridge::span_builder_set_tag_str(span, key, value), + TagValue::Literal(value) => { + bridge::span_builder_set_tag_literal(span, key, literal(value)); + } + } + } + } + + /// Starts a child span, as C++ `SpanBuilder::newChild()` does. + #[must_use] + pub fn new_child(&mut self, name: &'static CStr) -> Self { + // SAFETY: `name` is `'static`. + Self(unsafe { bridge::span_builder_new_child(self.0.pin_mut(), literal(name)) }) + } +} + +/// A span of both internal and user tracing, like C++ `TraceContext`, which it holds. The spans +/// end when this is dropped. It is neither `Send` nor `Sync`. +pub struct TraceContext(KjOwn); + +impl TraceContext { + /// [`IoContext::make_user_trace_span`](crate::IoContext::make_user_trace_span), for the + /// `IoContext` with id `context_id`. + pub(crate) fn start(context_id: u64, name: &'static CStr) -> jsg::Result { + // SAFETY: `name` is `'static`. + let context = + unsafe { bridge::io_context_make_user_trace_span(context_id, literal(name)) }?; + Ok(Self(context)) + } + + /// Sets tag `key` to `value` on both spans, as C++ `TraceContext::setTag()` does. + pub fn set_tag<'v>(&mut self, key: &'static CStr, value: impl Into>) { + let (context, key) = (self.0.pin_mut(), literal(key)); + // SAFETY: `key`, and a `Literal` value, are `'static`. + unsafe { + match value.into() { + TagValue::Bool(value) => bridge::trace_context_set_tag_bool(context, key, value), + TagValue::Int(value) => bridge::trace_context_set_tag_int(context, key, value), + TagValue::Float(value) => bridge::trace_context_set_tag_float(context, key, value), + TagValue::Str(value) => bridge::trace_context_set_tag_str(context, key, value), + TagValue::Literal(value) => { + bridge::trace_context_set_tag_literal(context, key, literal(value)); + } + } + } + } +} + +/// Holds no JavaScript handles. +impl jsg::Traced for SpanBuilder {} + +/// Holds no JavaScript handles. +impl jsg::Traced for TraceContext {} + +/// A literal as the bridge passes it: the bytes of `str` with its NUL. +fn literal(str: &'static CStr) -> &'static [u8] { + str.to_bytes_with_nul() +} From 73517e9d7f9473624c995faeed6a1119603dc1e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bartek=20Iwa=C5=84czuk?= Date: Wed, 28 Oct 2026 11:02:06 +0100 Subject: [PATCH 4/4] Tag a MemoryCache fallback as successful only once it is The memory_cache_fallback span was tagged fallback_success=true before the fallback's result was validated. If validation then threw, as for a NaN expiration or a value that cannot be serialized, the catch handler overwrote the tag with false. The final tags were right, but overwriting a tag is treated as a programming error, logged at ERROR in predictable mode. The tag is now set once the result has been accepted. A test records the spans of two rejected results and checks their tags, failing on the overwrite's log without the fix. --- src/workerd/api/BUILD.bazel | 1 + src/workerd/api/memory-cache-test.c++ | 107 ++++++++++++++++++++++++++ src/workerd/api/memory-cache.c++ | 5 +- 3 files changed, 111 insertions(+), 2 deletions(-) diff --git a/src/workerd/api/BUILD.bazel b/src/workerd/api/BUILD.bazel index 7a930881110..84205ee0fa5 100644 --- a/src/workerd/api/BUILD.bazel +++ b/src/workerd/api/BUILD.bazel @@ -671,6 +671,7 @@ kj_test( ":memory-cache", "//src/workerd/io:trace", "//src/workerd/tests:test-fixture", + "//src/workerd/util:thread-scopes", ], ) diff --git a/src/workerd/api/memory-cache-test.c++ b/src/workerd/api/memory-cache-test.c++ index 86228b19756..4c160f167b8 100644 --- a/src/workerd/api/memory-cache-test.c++ +++ b/src/workerd/api/memory-cache-test.c++ @@ -6,10 +6,13 @@ #include #include +#include #include #include +#include + namespace workerd::api { namespace { @@ -209,5 +212,109 @@ KJ_TEST("fallback callback stores the value and survives binding destruction") { }); } +// Records the spans of a request as they end, each as "name <- parent: key=type:value ...", with +// its tags in key order. +class SpanRecorder final: public SpanObserver { + public: + SpanRecorder(kj::Vector& spans, kj::StringPtr parent) + : spans(spans), + parent(kj::str(parent)) {} + + kj::Rc newChild() override { + return kj::rc(spans, name); + } + void onOpen(kj::ConstString operationName, kj::Date) override { + name = kj::str(operationName); + } + void onClose(kj::Date, Span::TagMap&& tags, kj::Vector&&) override { + auto entries = KJ_MAP(entry, tags) { + KJ_SWITCH_ONEOF(entry.value) { + KJ_CASE_ONEOF(str, kj::ConstString) { + return kj::str(entry.key, "=str:", str); + } + KJ_CASE_ONEOF(b, bool) { + return kj::str(entry.key, "=bool:", b); + } + KJ_CASE_ONEOF(d, double) { + return kj::str(entry.key, "=float:", d); + } + KJ_CASE_ONEOF(i, int64_t) { + return kj::str(entry.key, "=int:", i); + } + } + KJ_UNREACHABLE; + }; + std::sort(entries.begin(), entries.end()); + spans.add(kj::str(name, " <- ", parent, ": ", kj::strArray(entries, " "))); + } + + private: + kj::Vector& spans; + kj::String parent; + kj::String name; +}; + +class SpanRequestObserver final: public RequestObserver { + public: + explicit SpanRequestObserver(kj::Vector& spans) + : span(kj::rc(spans, "none"_kj), "request"_kjc) {} + + SpanParent getSpan() override { + return SpanParent(span); + } + + private: + SpanBuilder span; +}; + +KJ_TEST("a fallback whose result is rejected is tagged as failed, without overwriting a tag") { + // In predictable mode, overwriting a span's tag logs an error, which fails this test. It logs + // only the first overwrite in the process, so no other test here may overwrite one. + setPredictableModeForTest(); + MemoryCacheProvider provider; + kj::Vector spans; + TestFixture fixture({ + .mainModuleSource = R"SCRIPT( + export default { + async fetch() { + const settle = (promise) => promise.then(() => {}, () => {}); + await settle(globalThis.CACHE.read("nan", () => ({ value: 1, expiration: NaN }))); + await settle(globalThis.CACHE.read("function", () => ({ value: () => {} }))); + return new Response("ok"); + } + }; + )SCRIPT"_kj, + .requestObserverFactory = + kj::Function()>([&spans]() -> kj::Own { + return kj::refcounted(spans); + }), + }); + fixture.runInIoContext([&](const TestFixture::Environment& env) { + auto& handler = KJ_REQUIRE_NONNULL(env.js.tryGetTypeHandler>()); + auto cache = + handler.wrap(env.js, env.js.alloc(provider.getUse(kj::none, testLimits()))); + jsg::check( + env.js.v8Context()->Global()->Set(env.js.v8Context(), env.js.str("CACHE"_kj), cache)); + }); + auto response = fixture.runRequest(kj::HttpMethod::GET, "http://worker/"_kj, ""_kj); + KJ_EXPECT(response.statusCode == 200, response.body); + + kj::Vector fallbacks; + for (auto& span: spans) { + if (span.startsWith("memory_cache_fallback ")) fallbacks.add(span); + } + KJ_ASSERT(fallbacks.size() == 2, kj::strArray(spans, "\n")); + KJ_EXPECT(fallbacks[0] == + "memory_cache_fallback <- memory_cache_read: " + "fallback_error=str:TypeError: Expiration time must not be NaN. " + "fallback_result_size=float:4 fallback_success=bool:false key=str:nan", + fallbacks[0]); + // kj::str() of the DOMException that serialization throws is empty. + KJ_EXPECT(fallbacks[1] == + "memory_cache_fallback <- memory_cache_read: fallback_error=str: " + "fallback_success=bool:false key=str:function", + fallbacks[1]); +} + } // namespace } // namespace workerd::api diff --git a/src/workerd/api/memory-cache.c++ b/src/workerd/api/memory-cache.c++ index a3e577d2cbf..c46ad433157 100644 --- a/src/workerd/api/memory-cache.c++ +++ b/src/workerd/api/memory-cache.c++ @@ -366,8 +366,6 @@ jsg::Promise> MemoryCache::read(jsg::Lock& js, CacheValueProduceResult result) mutable -> jsg::JsRef { // NOTE: `callback` is IoPtr, not IoOwn. The catch block gets the IoOwn, which // ensures the object still exists at this point. - fallbackSpan->setTag("fallback_success"_kjc, true); - auto serialized = hackySerialize(js, result.value); fallbackSpan->setTag( "fallback_result_size"_kjc, static_cast(serialized.size())); @@ -379,6 +377,9 @@ jsg::Promise> MemoryCache::read(jsg::Lock& js, } else { fallbackSpan->setTag("has_expiration"_kjc, false); } + // Only now, so that the catch block, which handles what this throws, does not + // overwrite it. + fallbackSpan->setTag("fallback_success"_kjc, true); (*callback)(MemoryCacheUse::FallbackResult{kj::mv(serialized), result.expiration}, *fallbackSpan); return kj::mv(result.value);