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(), 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); 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..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"], @@ -24,7 +25,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 +44,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..e0421037917 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 @@ -25,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() { @@ -76,4 +99,76 @@ 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))); +} + +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 8f42e2b1ca0..6c3f9752776 100644 --- a/src/workerd/io/io-context/ffi.h +++ b/src/workerd/io/io-context/ffi.h @@ -29,4 +29,55 @@ 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); + +// 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 9e540e7041f..8a496e7909d 100644 --- a/src/workerd/io/io-context/ffi.rs +++ b/src/workerd/io/io-context/ffi.rs @@ -6,12 +6,17 @@ #![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, 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; +use std::future::Future; use std::pin::Pin; +use jsg::v8; + #[cxx::bridge(namespace = "workerd::io_context")] pub mod bridge { #[namespace = "workerd"] @@ -19,16 +24,83 @@ pub mod bridge { include!("workerd/io/io-context.h"); type IoContext; + type SpanBuilder; + type TraceContext; } 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>; + + // 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" { + 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 +112,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..39a65a6624c 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,10 @@ #include #include +#include +#include +#include + namespace workerd::io_context::test { namespace { @@ -179,6 +184,336 @@ 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 and operations 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)); + 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 {} + 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() 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(); + 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..3abfa8f412e 100644 --- a/src/workerd/io/io-context/lib.rs +++ b/src/workerd/io/io-context/lib.rs @@ -9,13 +9,25 @@ //! `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; +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; +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, @@ -88,6 +100,135 @@ 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, + }) + } + + /// 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`]. +/// +/// 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..bd4ef4bb911 100644 --- a/src/workerd/io/io-context/test_bridge.rs +++ b/src/workerd/io/io-context/test_bridge.rs @@ -9,10 +9,16 @@ //! 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 io_context::SpanBuilder; use jsg::ToJS; use jsg::v8; @@ -41,9 +47,42 @@ 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 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; @@ -120,3 +159,192 @@ 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) }; + static HELD_SPAN: 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, 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<()> { + // 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)) +} 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() +}