diff --git a/docs/jsg.md b/docs/jsg.md index c6742a0b5eb..fc1dddae8af 100644 --- a/docs/jsg.md +++ b/docs/jsg.md @@ -1950,6 +1950,10 @@ JSG provides rudimentary async context tracking via `AsyncContextFrame` in frame. Each frame has a storage context (a map of storage cells keyed by opaque keys). When a new frame is created, the current frame's storage context is propagated. +Small frames keep their storage cells inline. Larger frames use a hash table. Dead keys are +removed before storage is inherited, so a child with few live cells can use inline storage even +when its parent used a hash table. + All JavaScript promises, timers, and microtasks propagate the async context. ```js @@ -1988,12 +1992,17 @@ kj::Arc key = KJ_DEFER(key->reset()); // Clear when done { - jsg::AsyncContextFrame::StorageScope(js, key.addRef(), value); + jsg::AsyncContextFrame::StorageScope scope(js, key.addRef(), kj::mv(value)); // code runs with this storage context } // Automatically reset to previous context ``` +To install several storage cells in one frame, pass a mutable array of +`AsyncContextFrame::StorageEntry` to `StorageScope`. The entries are moved into the frame, +and inherited storage is copied once. All of the cells remain available to asynchronous work +that captures the frame, and leaving the scope restores the previous frame. + --- ## 15. Memory Tracking (`jsg::MemoryTracker`) diff --git a/src/workerd/api/tracing.c++ b/src/workerd/api/tracing.c++ index 22c7ee1235d..2d58570cbd1 100644 --- a/src/workerd/api/tracing.c++ +++ b/src/workerd/api/tracing.c++ @@ -719,7 +719,7 @@ jsg::Optional> Tracing::getActiveSpan( auto key = ioContext.getCurrentLock().getUserTraceAsyncContextKey(); KJ_IF_SOME(value, frame.get(*key)) { auto holder = value.getHandle(js).As(); - auto& asyncContext = jsg::unwrapOpaqueRef>(js.v8Isolate, holder); + auto& scopeState = jsg::unwrapOpaqueRef(js.v8Isolate, holder); auto cacheKey = v8::Private::ForApi(js.v8Isolate, js.strIntern("workerd.activeSpan"_kjc)); if (jsg::check(holder->HasPrivate(js.v8Context(), cacheKey))) { auto cached = jsg::check(holder->GetPrivate(js.v8Context(), cacheKey)); @@ -728,14 +728,23 @@ jsg::Optional> Tracing::getActiveSpan( } } - kj::Maybe> tracer; - KJ_IF_SOME(value, asyncContext->getTracer()) { - tracer = value.addRef(); + kj::Own state; + KJ_SWITCH_ONEOF(scopeState) { + KJ_CASE_ONEOF(id, IoContext::Id) { + if (!id.isCurrent()) IoContext::throwNotCurrentJsError(typeid(UserTraceAsyncContext)); + state = kj::refcounted(); + } + KJ_CASE_ONEOF(asyncContext, IoOwn) { + kj::Maybe> tracer; + KJ_IF_SOME(value, asyncContext->getTracer()) { + tracer = value.addRef(); + } + state = kj::refcounted(asyncContext->getSpan(), + kj::mv(tracer), asyncContext->getInvocationSpanContext().map([](auto& context) { + return context.clone(); + })); + } } - kj::Own state = - kj::refcounted(asyncContext->getSpan(), kj::mv(tracer), - asyncContext->getInvocationSpanContext().map( - [](auto& context) { return context.clone(); })); auto span = js.alloc(ioContext.addObject(kj::mv(state))); auto wrapped = spanHandler.wrap(js, span.addRef()); jsg::check(holder->SetPrivate(js.v8Context(), cacheKey, wrapped)); diff --git a/src/workerd/io/BUILD.bazel b/src/workerd/io/BUILD.bazel index d48ff78a2c6..880591fc491 100644 --- a/src/workerd/io/BUILD.bazel +++ b/src/workerd/io/BUILD.bazel @@ -659,6 +659,14 @@ kj_test( ], ) +kj_test( + src = "trace-scope-test.c++", + deps = [ + ":io", + "//src/workerd/tests:test-fixture", + ], +) + kj_test( src = "io-own-test.c++", deps = [ diff --git a/src/workerd/io/incoming-request-test.c++ b/src/workerd/io/incoming-request-test.c++ index c8c83ceb04f..10302c1fa4c 100644 --- a/src/workerd/io/incoming-request-test.c++ +++ b/src/workerd/io/incoming-request-test.c++ @@ -54,6 +54,79 @@ class FrozenTimerChannel final: public TimerChannel { kj::Date currentTime = kj::UNIX_EPOCH; }; +SpanParent traceScopeParent(uint64_t id) { + return SpanParent::fromSpanContext( + tracing::SpanContext(tracing::TraceId(1, 2), tracing::SpanId(id))); +} + +uint64_t traceScopeSpanId(SpanParent parent) { + auto context = parent.toSpanContext(); + return KJ_ASSERT_NONNULL(KJ_ASSERT_NONNULL(context).getSpanId()).getId(); +} + +class TraceScopeRequestObserver final: public RequestObserver { + public: + SpanParent getSpan() override { + return traceScopeParent(11); + } +}; + +KJ_TEST("request trace scopes capture both roots and restore enclosing tracing and storage") { + for (bool observeUser: {false, true}) { + TestFixture fixture({ + .useRealTimers = false, + .requestObserverFactory = kj::Function()>( + []() -> kj::Own { return kj::refcounted(); }), + }); + auto context = fixture.newIoContext(); + kj::Maybe> tracer; + if (observeUser) { + auto trace = kj::refcounted(kj::none, kj::none, kj::none, kj::none, kj::none, + kj::Array(), kj::none, ExecutionModel::STATELESS); + auto workerTracer = + kj::rc(kj::none, kj::mv(trace), PipelineLogLevel::NONE, kj::none, kj::none); + workerTracer->setMakeUserRequestSpanFunc( + [](tracing::TraceId, kj::Maybe) { return traceScopeParent(22); }); + tracer = kj::Rc(kj::mv(workerTracer)); + } + auto request = fixture.newUndeliveredIncomingRequest(*context, kj::mv(tracer)); + request->delivered(); + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + auto key = kj::arc(); + jsg::AsyncContextFrame::StorageScope storageScope( + env.js, key.addRef(), env.js.v8Ref(v8::Integer::New(env.isolate, 55).As())); + auto internalScope = context->makeAsyncTraceScope(env.lock, traceScopeParent(33)); + auto userScope = context->makeUserAsyncTraceScope(env.lock, traceScopeParent(44)); + auto& prior = KJ_ASSERT_NONNULL(jsg::AsyncContextFrame::current(env.js)); + kj::Maybe> captured; + auto checkRequestScope = [&]() { + KJ_EXPECT(traceScopeSpanId(context->getCurrentTraceSpan()) == 11); + auto userParent = context->getCurrentUserTraceSpan(); + KJ_EXPECT(userParent.isObserved() == observeUser); + if (observeUser) { + KJ_EXPECT(traceScopeSpanId(kj::mv(userParent)) == 22); + } + auto& current = KJ_ASSERT_NONNULL(jsg::AsyncContextFrame::current(env.js)); + auto value = KJ_ASSERT_NONNULL(current.get(*key)).getHandle(env.js); + KJ_EXPECT(value.As()->Value() == 55); + }; + { + auto scope = context->makeRequestTraceScope(env.lock); + checkRequestScope(); + captured = jsg::AsyncContextFrame::currentRef(env.js); + } + KJ_EXPECT(&KJ_ASSERT_NONNULL(jsg::AsyncContextFrame::current(env.js)) == &prior); + KJ_EXPECT(traceScopeSpanId(context->getCurrentTraceSpan()) == 33); + KJ_EXPECT(traceScopeSpanId(context->getCurrentUserTraceSpan()) == 44); + { + jsg::AsyncContextFrame::Scope scope(env.js, captured); + checkRequestScope(); + } + }); + fixture.drainAndDestroy(kj::mv(request)); + } +} + KJ_TEST("trace onset synchronizes an idle actor's clock before reading it") { FrozenTimerChannel timer; kj::Function(TimerChannel&)> makeChannelFactory = diff --git a/src/workerd/io/io-context.c++ b/src/workerd/io/io-context.c++ index a0291dec9ce..acea2dc4b81 100644 --- a/src/workerd/io/io-context.c++ +++ b/src/workerd/io/io-context.c++ @@ -1239,31 +1239,38 @@ kj::Own IoContext::getCacheClient() { jsg::AsyncContextFrame::StorageScope IoContext::makeAsyncTraceScope( Worker::Lock& lock, kj::Maybe spanParentOverride) { - static const SpanParent dummySpanParent = nullptr; + return jsg::AsyncContextFrame::StorageScope( + lock, lock.getTraceAsyncContextKey(), makeTraceScopeValue(lock, kj::mv(spanParentOverride))); +} +jsg::AsyncContextFrame::StorageScope IoContext::makeUserAsyncTraceScope( + Worker::Lock& lock, kj::Maybe userSpanOverride) { + return jsg::AsyncContextFrame::StorageScope(lock, lock.getUserTraceAsyncContextKey(), + makeUserTraceScopeValue(lock, kj::mv(userSpanOverride))); +} + +jsg::AsyncContextFrame::StorageScope IoContext::makeRequestTraceScope(Worker::Lock& lock) { + jsg::AsyncContextFrame::StorageEntry entries[] = { + {lock.getTraceAsyncContextKey(), makeTraceScopeValue(lock, kj::none)}, + {lock.getUserTraceAsyncContextKey(), makeUserTraceScopeValue(lock, kj::none)}, + }; + return jsg::AsyncContextFrame::StorageScope(lock, entries); +} + +jsg::Value IoContext::makeTraceScopeValue( + Worker::Lock& lock, kj::Maybe spanParentOverride) { + auto& ioContext = IoContext::current(); jsg::Lock& js = lock; - kj::Own spanParent; - KJ_IF_SOME(spo, kj::mv(spanParentOverride)) { - spanParent = kj::heap(kj::mv(spo)); - } else { - // TODO(cleanup): Can we also elide the other memory allocations for the (unused) storage - // scope if tracing is disabled? - SpanParent metricsSpan = getMetrics().getSpan(); - if (!metricsSpan.isObserved()) { - // const_cast is ok: There's no state that could be changed in a non-observed span parent. - spanParent = kj::Own( - &const_cast(dummySpanParent), kj::NullDisposer::instance); - } else { - spanParent = kj::heap(kj::mv(metricsSpan)); - } + auto spanParent = kj::mv(spanParentOverride).orDefault([&]() { return getMetrics().getSpan(); }); + if (!spanParent.isObserved()) { + return js.v8Ref(jsg::wrapOpaque(js.v8Context(), TraceScopeState(ioContext.getId()))); } - auto ioOwnSpanParent = IoContext::current().addObject(kj::mv(spanParent)); - auto spanHandle = jsg::wrapOpaque(js.v8Context(), kj::mv(ioOwnSpanParent)); - return jsg::AsyncContextFrame::StorageScope( - js, lock.getTraceAsyncContextKey(), js.v8Ref(spanHandle)); + auto ioOwnSpanParent = ioContext.addObject(kj::heap(kj::mv(spanParent))); + auto spanHandle = jsg::wrapOpaque(js.v8Context(), TraceScopeState(kj::mv(ioOwnSpanParent))); + return js.v8Ref(spanHandle); } -jsg::AsyncContextFrame::StorageScope IoContext::makeUserAsyncTraceScope( +jsg::Value IoContext::makeUserTraceScopeValue( Worker::Lock& lock, kj::Maybe userSpanOverride) { auto& ioContext = IoContext::current(); jsg::Lock& js = lock; @@ -1274,16 +1281,18 @@ jsg::AsyncContextFrame::StorageScope IoContext::makeUserAsyncTraceScope( userSpan = getRootUserTraceSpan(); } + if (!userSpan.isObserved()) { + return js.v8Ref(jsg::wrapOpaque(js.v8Context(), UserTraceScopeState(ioContext.getId()))); + } + kj::Maybe invocationSpanContext; - if (userSpan.isObserved()) { - auto& baseContext = getCurrentIncomingRequest().getInvocationSpanContext(); - auto spanId = userSpan.getSpanId(); - if (spanId != tracing::SpanId::nullId) { - invocationSpanContext = tracing::InvocationSpanContext(baseContext.getTraceId(), - baseContext.getInvocationId(), spanId, baseContext.getTraceFlags()); - } else { - invocationSpanContext = baseContext.clone(); - } + auto& baseContext = getCurrentIncomingRequest().getInvocationSpanContext(); + auto spanId = userSpan.getSpanId(); + if (spanId != tracing::SpanId::nullId) { + invocationSpanContext = tracing::InvocationSpanContext(baseContext.getTraceId(), + baseContext.getInvocationId(), spanId, baseContext.getTraceFlags()); + } else { + invocationSpanContext = baseContext.clone(); } kj::Maybe>> tracer; @@ -1294,9 +1303,9 @@ jsg::AsyncContextFrame::StorageScope IoContext::makeUserAsyncTraceScope( auto asyncContext = kj::heap( kj::mv(userSpan), kj::mv(tracer), kj::mv(invocationSpanContext)); auto ioOwnAsyncContext = ioContext.addObject(kj::mv(asyncContext)); - auto contextHandle = jsg::wrapOpaque(js.v8Context(), kj::mv(ioOwnAsyncContext)); - return jsg::AsyncContextFrame::StorageScope( - js, lock.getUserTraceAsyncContextKey(), js.v8Ref(contextHandle)); + auto contextHandle = + jsg::wrapOpaque(js.v8Context(), UserTraceScopeState(kj::mv(ioOwnAsyncContext))); + return js.v8Ref(contextHandle); } SpanParent IoContext::getCurrentTraceSpan() { @@ -1306,11 +1315,18 @@ SpanParent IoContext::getCurrentTraceSpan() { KJ_IF_SOME(value, frame.get(*lock.getTraceAsyncContextKey())) { auto handle = value.getHandle(lock); jsg::Lock& js = lock; - auto& spanParent = jsg::unwrapOpaqueRef>(js.v8Isolate, handle); + auto& state = jsg::unwrapOpaqueRef(js.v8Isolate, handle); // The frame may have been captured by a different request (e.g. via // AsyncLocalStorage.snapshot()), whose span this request can't use. - KJ_IF_SOME(span, spanParent.tryGet()) { - return span.addRef(); + KJ_SWITCH_ONEOF(state) { + KJ_CASE_ONEOF(id, Id) { + if (id.isCurrent()) return SpanParent(nullptr); + } + KJ_CASE_ONEOF(spanParent, IoOwn) { + KJ_IF_SOME(span, spanParent.tryGet()) { + return span.addRef(); + } + } } } } @@ -1332,12 +1348,18 @@ SpanParent IoContext::getCurrentUserTraceSpan() { KJ_IF_SOME(value, frame.get(*lock.getUserTraceAsyncContextKey())) { auto handle = value.getHandle(lock); jsg::Lock& js = lock; - auto& asyncContext = - jsg::unwrapOpaqueRef>(js.v8Isolate, handle); + auto& state = jsg::unwrapOpaqueRef(js.v8Isolate, handle); // As in getCurrentTraceSpan(), the frame may belong to a different request, in which case // fall back to this request's root span. - KJ_IF_SOME(context, asyncContext.tryGet()) { - return context.getSpan(); + KJ_SWITCH_ONEOF(state) { + KJ_CASE_ONEOF(id, Id) { + if (id.isCurrent()) return SpanParent(nullptr); + } + KJ_CASE_ONEOF(asyncContext, IoOwn) { + KJ_IF_SOME(context, asyncContext.tryGet()) { + return context.getSpan(); + } + } } } } diff --git a/src/workerd/io/io-context.h b/src/workerd/io/io-context.h index c7df6065ca4..e1d23e8a121 100644 --- a/src/workerd/io/io-context.h +++ b/src/workerd/io/io-context.h @@ -539,6 +539,11 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler return id; } + // Unobserved scopes retain only their context identity. This masks inherited spans while + // allowing accessors to fall back to the current request when a captured frame is foreign. + using TraceScopeState = kj::OneOf>; + using UserTraceScopeState = kj::OneOf>; + void setEntrypointHandler(jsg::Lock& js, jsg::JsObject handler); jsg::JsObject getEntrypointHandler(jsg::Lock& js); @@ -1120,6 +1125,10 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler jsg::AsyncContextFrame::StorageScope makeUserAsyncTraceScope( Worker::Lock& lock, kj::Maybe userSpan = kj::none) KJ_WARN_UNUSED_RESULT; + // Captures the request's internal and user tracing contexts in a single async frame. + jsg::AsyncContextFrame::StorageScope makeRequestTraceScope( + Worker::Lock& lock) KJ_WARN_UNUSED_RESULT; + // Returns the current span being recorded. If called while the JS lock is held, uses the trace // information from the current async context, if available. SpanParent getCurrentTraceSpan(); @@ -1170,6 +1179,9 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler void pumpMessageLoop(); private: + jsg::Value makeTraceScopeValue(Worker::Lock& lock, kj::Maybe spanParentOverride); + jsg::Value makeUserTraceScopeValue(Worker::Lock& lock, kj::Maybe userSpanOverride); + ThreadContext& thread; kj::Own selfRef = kj::refcounted(kj::Badge(), *this); diff --git a/src/workerd/io/trace-scope-test.c++ b/src/workerd/io/trace-scope-test.c++ new file mode 100644 index 00000000000..b2ded25f8e8 --- /dev/null +++ b/src/workerd/io/trace-scope-test.c++ @@ -0,0 +1,201 @@ +// 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 + +#include +#include +#include +#include + +namespace workerd { +namespace { + +SpanParent observedParent(uint64_t id) { + return SpanParent::fromSpanContext( + tracing::SpanContext(tracing::TraceId(1, 2), tracing::SpanId(id))); +} + +kj::Rc makeTracer() { + auto trace = kj::refcounted(kj::none, kj::none, kj::none, kj::none, kj::none, + kj::Array(), kj::none, ExecutionModel::STATELESS); + return kj::rc(kj::none, kj::mv(trace), PipelineLogLevel::NONE, kj::none, kj::none); +} + +kj::Rc makeObservedTracer() { + auto tracer = makeTracer(); + tracer->setMakeUserRequestSpanFunc( + [](tracing::TraceId, kj::Maybe) { return observedParent(2); }); + return kj::mv(tracer); +} + +class ObservedRequest final: public RequestObserver { + public: + SpanParent getSpan() override { + return observedParent(1); + } +}; + +uint64_t spanId(SpanParent parent) { + auto context = parent.toSpanContext(); + return KJ_ASSERT_NONNULL(KJ_ASSERT_NONNULL(context).getSpanId()).getId(); +} + +jsg::Ref activeSpan(const TestFixture::Environment& env) { + auto tracing = env.js.alloc(); + auto& handler = KJ_ASSERT_NONNULL(env.js.tryGetTypeHandler>()); + auto span = tracing->getActiveSpan(env.js, handler); + return kj::mv(KJ_ASSERT_NONNULL(span)); +} + +KJ_TEST("unobserved scopes retain identity and materialize their no-op span lazily") { + for (bool hasTracer: {false, true}) { + TestFixture fixture; + auto context = fixture.newIoContext(); + kj::Maybe> tracer; + if (hasTracer) tracer = kj::Rc(makeTracer()); + auto request = fixture.newUndeliveredIncomingRequest(*context, kj::mv(tracer)); + request->delivered(); + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + auto scope = context->makeRequestTraceScope(env.lock); + auto& frame = KJ_ASSERT_NONNULL(jsg::AsyncContextFrame::current(env.js)); + auto internalValue = + KJ_ASSERT_NONNULL(frame.get(*env.lock.getTraceAsyncContextKey())).getHandle(env.js); + auto& internalState = + jsg::unwrapOpaqueRef(env.isolate, internalValue); + KJ_EXPECT(internalState.is()); + KJ_EXPECT(internalState.get() == context->getId()); + auto userValue = + KJ_ASSERT_NONNULL(frame.get(*env.lock.getUserTraceAsyncContextKey())).getHandle(env.js); + auto& userState = + jsg::unwrapOpaqueRef(env.isolate, userValue); + KJ_EXPECT(userState.is()); + KJ_EXPECT(userState.get() == context->getId()); + KJ_EXPECT(!context->getCurrentTraceSpan().isObserved()); + KJ_EXPECT(!context->getCurrentUserTraceSpan().isObserved()); + + auto holder = userValue.As(); + auto cacheKey = v8::Private::ForApi(env.isolate, env.js.strIntern("workerd.activeSpan"_kjc)); + KJ_EXPECT(!jsg::check(holder->HasPrivate(env.js.v8Context(), cacheKey))); + auto first = activeSpan(env); + KJ_EXPECT(!first->getIsTraced()); + first->updateName(env.js, kj::str("unused")); + first->end(); + KJ_EXPECT(!first->getIsTraced()); + KJ_EXPECT(jsg::check(holder->HasPrivate(env.js.v8Context(), cacheKey))); + auto second = activeSpan(env); + KJ_EXPECT(first.get() == second.get()); + { + auto nested = context->makeUserAsyncTraceScope(env.lock, SpanParent(nullptr)); + auto other = activeSpan(env); + KJ_EXPECT(other.get() != first.get()); + KJ_EXPECT(!other->getIsTraced()); + } + KJ_EXPECT(activeSpan(env).get() == first.get()); + }); + fixture.drainAndDestroy(kj::mv(request)); + } +} + +KJ_TEST("observed and unobserved user scopes restore their cached invocation spans") { + TestFixture fixture; + auto context = fixture.newIoContext(); + auto request = fixture.newUndeliveredIncomingRequest(*context, makeObservedTracer()); + request->delivered(); + fixture.enterContext(*request, [&](const TestFixture::Environment& env) { + auto scope = context->makeRequestTraceScope(env.lock); + auto root = activeSpan(env); + KJ_EXPECT(root->getIsTraced()); + { + auto nested = context->makeUserAsyncTraceScope(env.lock, SpanParent(nullptr)); + auto unobserved = activeSpan(env); + KJ_EXPECT(!unobserved->getIsTraced()); + KJ_EXPECT(unobserved.get() != root.get()); + } + KJ_EXPECT(root->getIsTraced()); + KJ_EXPECT(activeSpan(env).get() == root.get()); + }); + fixture.drainAndDestroy(kj::mv(request)); +} + +KJ_TEST( + "foreign unobserved frames fall back to the current request and reject active-span access") { + for (bool destroySource: {false, true}) { + for (bool cacheSpan: {false, true}) { + TestFixture fixture({ + .useRealTimers = false, + .requestObserverFactory = kj::Function()>( + []() -> kj::Own { return kj::refcounted(); }), + }); + auto sourceContext = fixture.newIoContext(); + auto source = fixture.newIncomingRequest(*sourceContext); + kj::Maybe> captured; + kj::Maybe> cached; + fixture.enterContext(*source, [&](const TestFixture::Environment& env) { + auto internalScope = sourceContext->makeAsyncTraceScope(env.lock, SpanParent(nullptr)); + auto userScope = sourceContext->makeUserAsyncTraceScope(env.lock, SpanParent(nullptr)); + KJ_EXPECT(!sourceContext->getCurrentTraceSpan().isObserved()); + KJ_EXPECT(!sourceContext->getCurrentUserTraceSpan().isObserved()); + captured = jsg::AsyncContextFrame::currentRef(env.js); + if (cacheSpan) cached = activeSpan(env); + }); + if (destroySource) { + fixture.drainAndDestroy(kj::mv(source)); + sourceContext = nullptr; + } + + auto targetContext = fixture.newIoContext(); + auto target = fixture.newUndeliveredIncomingRequest(*targetContext, makeObservedTracer()); + target->delivered(); + fixture.enterContext(*target, [&](const TestFixture::Environment& env) { + jsg::AsyncContextFrame::Scope scope(env.js, captured); + KJ_EXPECT(spanId(targetContext->getCurrentTraceSpan()) == 1); + KJ_EXPECT(spanId(targetContext->getCurrentUserTraceSpan()) == 2); + if (cacheSpan) { + auto span = activeSpan(env); + KJ_EXPECT(span.get() == KJ_ASSERT_NONNULL(cached).get()); + KJ_EXPECT_THROW_MESSAGE( + "Cannot perform I/O on behalf of a different request", span->getIsTraced()); + } else { + KJ_EXPECT_THROW_MESSAGE( + "Cannot perform I/O on behalf of a different request", activeSpan(env)); + } + }); + fixture.drainAndDestroy(kj::mv(target)); + if (source.get() != nullptr) fixture.drainAndDestroy(kj::mv(source)); + } + } +} + +KJ_TEST("an unobserved actor continuation masks a newer request's observed spans") { + TestFixture fixture({ + .actorId = Worker::Actor::Id(kj::str("trace-scopes")), + .useRealTimers = false, + .requestObserverFactory = + kj::Function()>([count = 0]() mutable -> kj::Own { + if (count++ == 0) return kj::refcounted(); + return kj::refcounted(); + }), + }); + auto context = fixture.newIoContext(); + auto first = fixture.newIncomingRequest(*context); + kj::Maybe> captured; + fixture.enterContext(*first, [&](const TestFixture::Environment& env) { + auto scope = context->makeRequestTraceScope(env.lock); + captured = jsg::AsyncContextFrame::currentRef(env.js); + }); + auto second = fixture.newUndeliveredIncomingRequest(*context, makeObservedTracer()); + second->delivered(); + fixture.drainAndDestroy(kj::mv(first)); + fixture.enterContext(*second, [&](const TestFixture::Environment& env) { + KJ_EXPECT(spanId(context->getCurrentTraceSpan()) == 1); + KJ_EXPECT(spanId(context->getCurrentUserTraceSpan()) == 2); + jsg::AsyncContextFrame::Scope scope(env.js, captured); + KJ_EXPECT(!context->getCurrentTraceSpan().isObserved()); + KJ_EXPECT(!context->getCurrentUserTraceSpan().isObserved()); + KJ_EXPECT(!activeSpan(env)->getIsTraced()); + }); + fixture.drainAndDestroy(kj::mv(second)); +} + +} // namespace +} // namespace workerd diff --git a/src/workerd/io/worker-entrypoint.c++ b/src/workerd/io/worker-entrypoint.c++ index f3b6398ab65..77a8677b5e4 100644 --- a/src/workerd/io/worker-entrypoint.c++ +++ b/src/workerd/io/worker-entrypoint.c++ @@ -500,9 +500,7 @@ kj::Promise WorkerEntrypoint::requestImpl(kj::HttpMethod method, TRACE_EVENT_END("workerd", PERFETTO_TRACK_FROM_POINTER(&context)); TRACE_EVENT( "workerd", "WorkerEntrypoint::request() run", PERFETTO_FLOW_FROM_POINTER(this)); - jsg::AsyncContextFrame::StorageScope traceScope = context.makeAsyncTraceScope(lock); - jsg::AsyncContextFrame::StorageScope userTraceScope = - context.makeUserAsyncTraceScope(lock); + auto traceScope = context.makeRequestTraceScope(lock); kj::Maybe> signal; auto featureFlags = FeatureFlags::get(lock); @@ -789,8 +787,7 @@ kj::Promise WorkerEntrypoint::connect(kj::StringPtr host, entrypointName = entrypointName.clone(), versionInfo = kj::mv(versionInfo), host = kj::str(host), clientAddress = kj::mv(clientAddress)]( Worker::Lock& lock, IoContext& context) mutable { - jsg::AsyncContextFrame::StorageScope traceScope = context.makeAsyncTraceScope(lock); - jsg::AsyncContextFrame::StorageScope userTraceScope = context.makeUserAsyncTraceScope(lock); + auto traceScope = context.makeRequestTraceScope(lock); return lock.getGlobalScope().connect(kj::mv(host), kj::mv(clientAddress), headers, connection, response, lock, @@ -933,8 +930,7 @@ kj::Promise WorkerEntrypoint::runScheduled( &metrics = incomingRequest->getMetrics()]( Worker::Lock& lock, IoContext& context) mutable { TRACE_EVENT("workerd", "WorkerEntrypoint::runScheduled() run"); - jsg::AsyncContextFrame::StorageScope traceScope = context.makeAsyncTraceScope(lock); - jsg::AsyncContextFrame::StorageScope userTraceScope = context.makeUserAsyncTraceScope(lock); + auto traceScope = context.makeRequestTraceScope(lock); lock.getGlobalScope().startScheduled(scheduledTime, cron, lock, lock.getExportedHandler( @@ -1012,9 +1008,7 @@ kj::Promise WorkerEntrypoint::runAlarmImpl( [scheduledTime, retryCount, entrypointName = entrypointName.clone(), versionInfo = kj::mv(versionInfo), props = kj::mv(props)](Worker::Lock& lock, IoContext& context) mutable { - jsg::AsyncContextFrame::StorageScope traceScope = context.makeAsyncTraceScope(lock); - jsg::AsyncContextFrame::StorageScope userTraceScope = - context.makeUserAsyncTraceScope(lock); + auto traceScope = context.makeRequestTraceScope(lock); // If we have an invalid timeout, set it to the default value of 15 minutes. auto timeout = context.getLimitEnforcer().getAlarmLimit(); @@ -1112,8 +1106,7 @@ kj::Promise WorkerEntrypoint::test() { props = kj::mv(props), &metrics = incomingRequest->getMetrics()]( Worker::Lock& lock, IoContext& context) mutable -> kj::Promise { TRACE_EVENT("workerd", "WorkerEntrypoint::test() run"); - jsg::AsyncContextFrame::StorageScope traceScope = context.makeAsyncTraceScope(lock); - jsg::AsyncContextFrame::StorageScope userTraceScope = context.makeUserAsyncTraceScope(lock); + auto traceScope = context.makeRequestTraceScope(lock); return context.awaitJs(lock, lock.getGlobalScope().test(lock, diff --git a/src/workerd/jsg/BUILD.bazel b/src/workerd/jsg/BUILD.bazel index 9c6f4461356..36bc0eddfd4 100644 --- a/src/workerd/jsg/BUILD.bazel +++ b/src/workerd/jsg/BUILD.bazel @@ -37,6 +37,13 @@ wd_cc_benchmark( deps = [":jsg"], ) +wd_cc_benchmark( + name = "bench-async-context", + srcs = ["bench-async-context.c++"], + local_defines = ["JSG_IMPLEMENTATION"], + deps = [":jsg"], +) + wd_cc_library( name = "iterator", srcs = [ diff --git a/src/workerd/jsg/async-context-test.c++ b/src/workerd/jsg/async-context-test.c++ index d1e51bed832..97cc7f2d112 100644 --- a/src/workerd/jsg/async-context-test.c++ +++ b/src/workerd/jsg/async-context-test.c++ @@ -20,6 +20,215 @@ void collectGarbage(Lock& js) { check(script->Run(js.v8Context())); } +AsyncContextFrame::StorageEntry numberEntry( + Lock& js, const kj::Arc& key, int value) { + return {key.addRef(), js.v8Ref(v8::Integer::New(js.v8Isolate, value).As())}; +} + +int storedNumber(Lock& js, AsyncContextFrame& frame, const AsyncContextFrame::StorageKey& key) { + return KJ_ASSERT_NONNULL(frame.get(key)).getHandle(js).As()->Value(); +} + +KJ_TEST("multi-entry async scopes inherit storage and restore their prior frame") { + Evaluator evaluator(v8System); + evaluator.run([&](auto& js) { + auto inheritedKey = kj::arc(); + auto internalKey = kj::arc(); + auto userKey = kj::arc(); + AsyncContextFrame::StorageEntry priorEntries[] = { + numberEntry(js, inheritedKey, 1), + numberEntry(js, internalKey, 2), + }; + AsyncContextFrame::StorageScope priorScope(js, priorEntries); + auto& prior = KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)); + kj::Maybe> captured; + { + AsyncContextFrame::StorageEntry entries[] = { + numberEntry(js, internalKey, 3), + numberEntry(js, userKey, 4), + }; + AsyncContextFrame::StorageScope scope(js, entries); + auto& current = KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)); + KJ_EXPECT(storedNumber(js, current, *inheritedKey) == 1); + KJ_EXPECT(storedNumber(js, current, *internalKey) == 3); + KJ_EXPECT(storedNumber(js, current, *userKey) == 4); + KJ_EXPECT(storedNumber(js, prior, *internalKey) == 2); + KJ_EXPECT(prior.get(*userKey) == kj::none); + captured = current.addRef(); + } + KJ_EXPECT(&KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)) == &prior); + + collectGarbage(js); + { + AsyncContextFrame::Scope scope(js, captured); + auto& current = KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)); + KJ_EXPECT(storedNumber(js, current, *inheritedKey) == 1); + KJ_EXPECT(storedNumber(js, current, *internalKey) == 3); + KJ_EXPECT(storedNumber(js, current, *userKey) == 4); + } + KJ_EXPECT(&KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)) == &prior); + }); +} + +KJ_TEST("multi-entry async scopes ignore dead keys and keep the last value for duplicate keys") { + Evaluator evaluator(v8System); + evaluator.run([&](auto& js) { + auto deadKey = kj::arc(); + auto liveKey = kj::arc(); + deadKey->reset(); + v8::Global deadValue; + kj::Maybe scope; + { + v8::HandleScope handleScope(js.v8Isolate); + auto value = v8::Object::New(js.v8Isolate); + deadValue.Reset(js.v8Isolate, value); + deadValue.SetWeak(); + AsyncContextFrame::StorageEntry entries[] = { + numberEntry(js, liveKey, 1), + {deadKey.addRef(), js.v8Ref(value.template As())}, + numberEntry(js, liveKey, 3), + }; + scope.emplace(js, entries); + } + collectGarbage(js); + KJ_EXPECT(deadValue.IsEmpty()); + auto& current = KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)); + KJ_EXPECT(storedNumber(js, current, *liveKey) == 3); + scope = kj::none; + KJ_EXPECT(AsyncContextFrame::current(js) == kj::none); + }); +} + +KJ_TEST("async frame growth preserves inherited values and replaces duplicate keys") { + Evaluator evaluator(v8System); + for (int count: {0, 1, 2, 3, 4, 8, 64}) { + evaluator.run([&](auto& js) { + kj::Vector> keys; + kj::Vector entries; + for (auto i: kj::zeroTo(count)) { + keys.add(kj::arc()); + entries.add(numberEntry(js, keys.back(), i)); + } + AsyncContextFrame::StorageScope parentScope(js, entries.asPtr()); + auto& parent = KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)); + auto newKey = kj::arc(); + KJ_EXPECT(parent.get(*newKey) == kj::none); + for (auto i: kj::zeroTo(count)) { + KJ_EXPECT(storedNumber(js, parent, *keys[i]) == i); + } + + kj::Vector updates; + if (count > 0) { + updates.add(numberEntry(js, keys.front(), -1)); + updates.add(numberEntry(js, keys.back(), -2)); + } + updates.add(numberEntry(js, newKey, -3)); + { + AsyncContextFrame::StorageScope childScope(js, updates.asPtr()); + auto& child = KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)); + for (auto i: kj::zeroTo(count)) { + int expected = i; + if (i == 0) expected = -1; + if (i == count - 1) expected = -2; + KJ_EXPECT(storedNumber(js, child, *keys[i]) == expected); + KJ_EXPECT(storedNumber(js, parent, *keys[i]) == i); + } + KJ_EXPECT(storedNumber(js, child, *newKey) == -3); + KJ_EXPECT(parent.get(*newKey) == kj::none); + } + KJ_EXPECT(&KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)) == &parent); + }); + } +} + +KJ_TEST("async frame lookup and inheritance release values belonging to dead keys") { + Evaluator evaluator(v8System); + for (int count: {2, 3, 8}) { + for (bool cleanOnLookup: {false, true}) { + evaluator.run([&](auto& js) { + kj::Vector> keys; + v8::Global deadValue; + kj::Maybe parentScope; + { + v8::HandleScope handleScope(js.v8Isolate); + kj::Vector entries; + for (auto i: kj::zeroTo(count)) { + keys.add(kj::arc()); + if (i == count - 1) { + auto value = v8::Object::New(js.v8Isolate); + deadValue.Reset(js.v8Isolate, value); + deadValue.SetWeak(); + entries.add(AsyncContextFrame::StorageEntry( + keys.back().addRef(), js.v8Ref(value.template As()))); + } else { + entries.add(numberEntry(js, keys.back(), i)); + } + } + parentScope.emplace(js, entries.asPtr()); + } + for (auto i: kj::range(1, count)) { + keys[i]->reset(); + } + auto& parent = KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)); + if (cleanOnLookup) { + KJ_EXPECT(storedNumber(js, parent, *keys.front()) == 0); + collectGarbage(js); + KJ_EXPECT(deadValue.IsEmpty()); + } + auto newKey = kj::arc(); + { + AsyncContextFrame::StorageScope childScope(js, newKey.addRef(), + js.v8Ref(v8::Integer::New(js.v8Isolate, -1).template As())); + collectGarbage(js); + KJ_EXPECT(deadValue.IsEmpty()); + auto& child = KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)); + KJ_EXPECT(storedNumber(js, child, *keys.front()) == 0); + KJ_EXPECT(storedNumber(js, child, *newKey) == -1); + } + KJ_EXPECT(storedNumber(js, parent, *keys.front()) == 0); + KJ_EXPECT(parent.get(*newKey) == kj::none); + }); + } + } +} + +KJ_TEST("captured async frames retain storage values until the frame is released") { + Evaluator evaluator(v8System); + for (int count: {1, 2, 3, 8}) { + evaluator.run([&](auto& js) { + auto key = kj::arc(); + v8::Global weakValue; + kj::Maybe> captured; + { + v8::HandleScope handleScope(js.v8Isolate); + kj::Vector entries; + auto value = v8::Object::New(js.v8Isolate); + weakValue.Reset(js.v8Isolate, value); + weakValue.SetWeak(); + entries.add(AsyncContextFrame::StorageEntry( + key.addRef(), js.v8Ref(value.template As()))); + for (int i = 1; i < count; ++i) { + entries.add(numberEntry(js, kj::arc(), i)); + } + AsyncContextFrame::StorageScope scope(js, entries.asPtr()); + captured = AsyncContextFrame::currentRef(js); + } + collectGarbage(js); + KJ_EXPECT(!weakValue.IsEmpty()); + { + v8::HandleScope handleScope(js.v8Isolate); + AsyncContextFrame::Scope scope(js, captured); + auto& frame = KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)); + auto value = KJ_ASSERT_NONNULL(frame.get(*key)).getHandle(js); + KJ_EXPECT(value == weakValue.Get(js.v8Isolate)); + } + captured = kj::none; + collectGarbage(js); + KJ_EXPECT(weakValue.IsEmpty()); + }); + } +} + KJ_TEST("AsyncContextFrame::Scope retains its prior frame") { Evaluator evaluator(v8System); evaluator.run([&](auto& js) { diff --git a/src/workerd/jsg/async-context.c++ b/src/workerd/jsg/async-context.c++ index 0f5f76350f0..75f1303d5f8 100644 --- a/src/workerd/jsg/async-context.c++ +++ b/src/workerd/jsg/async-context.c++ @@ -24,24 +24,34 @@ inline void maybeSetV8ContinuationContext( } } // namespace -AsyncContextFrame::AsyncContextFrame(Lock& js, StorageEntry storageEntry) { +AsyncContextFrame::AsyncContextFrame(Lock& js, StorageEntry storageEntry) + : AsyncContextFrame(js, kj::arrayPtr(&storageEntry, 1)) {} + +AsyncContextFrame::AsyncContextFrame(Lock& js, kj::ArrayPtr storageEntries) { KJ_IF_SOME(frame, current(js)) { // Propagate the storage context of the current frame (if any). // If current(js) returns nullptr, we assume we're in the root // frame and there is no storage to propagate. - frame.storage.eraseAll([](const auto& entry) { return entry.key->isDead(); }); - for (auto& entry: frame.storage) { - storage.insert(entry.clone(js)); + frame.eraseDeadStorage(); + KJ_SWITCH_ONEOF(frame.storage) { + KJ_CASE_ONEOF(entries, InlineStorage) { + for (auto& slot: entries) { + KJ_IF_SOME(entry, slot) { + upsertStorageEntry(entry.clone(js)); + } + } + } + KJ_CASE_ONEOF(table, TableStorage) { + for (auto& entry: table) { + upsertStorageEntry(entry.clone(js)); + } + } } } - // This case is extremely unlikely to happen but let's handle it anyway - // just out of an excess of caution. - if (storageEntry.key->isDead()) return; - - storage.upsert(kj::mv(storageEntry), [](StorageEntry& existing, StorageEntry&& row) mutable { - existing.value = kj::mv(row.value); - }); + for (auto& entry: storageEntries) { + upsertStorageEntry(kj::mv(entry)); + } } AsyncContextFrame::StorageEntry::StorageEntry(kj::Arc key, Value value) @@ -78,6 +88,11 @@ Ref AsyncContextFrame::create(Lock& js, StorageEntry storageE return js.alloc(js, kj::mv(storageEntry)); } +Ref AsyncContextFrame::create( + Lock& js, kj::ArrayPtr storageEntries) { + return js.alloc(js, storageEntries); +} + v8::Local AsyncContextFrame::wrap(Lock& js, V8Ref& fn, jsg::Function validate, @@ -152,8 +167,71 @@ v8::Local AsyncContextFrame::wrapRoot( kj::Maybe AsyncContextFrame::get(const StorageKey& key) { KJ_ASSERT(!key.isDead()); - storage.eraseAll([](const auto& entry) { return entry.key->isDead(); }); - return storage.find(key).map([](auto& entry) -> Value& { return entry.value; }); + eraseDeadStorage(); + KJ_SWITCH_ONEOF(storage) { + KJ_CASE_ONEOF(entries, InlineStorage) { + for (auto& slot: entries) { + KJ_IF_SOME(entry, slot) { + if (entry.key.get() == &key) return entry.value; + } + } + return kj::none; + } + KJ_CASE_ONEOF(table, TableStorage) { + return table.find(key).map([](auto& entry) -> Value& { return entry.value; }); + } + } + KJ_UNREACHABLE; +} + +void AsyncContextFrame::eraseDeadStorage() { + KJ_SWITCH_ONEOF(storage) { + KJ_CASE_ONEOF(entries, InlineStorage) { + for (auto& slot: entries) { + KJ_IF_SOME(entry, slot) { + if (entry.key->isDead()) slot = kj::none; + } + } + } + KJ_CASE_ONEOF(table, TableStorage) { + table.eraseAll([](const auto& entry) { return entry.key->isDead(); }); + } + } +} + +void AsyncContextFrame::upsertStorageEntry(StorageEntry entry) { + if (entry.key->isDead()) return; + KJ_SWITCH_ONEOF(storage) { + KJ_CASE_ONEOF(entries, InlineStorage) { + for (auto& slot: entries) { + KJ_IF_SOME(existing, slot) { + if (existing.key.get() == entry.key.get()) { + existing.value = kj::mv(entry.value); + return; + } + } + } + for (auto& slot: entries) { + if (slot == kj::none) { + slot = kj::mv(entry); + return; + } + } + TableStorage table; + table.reserve(entries.size() + 1); + for (auto& slot: entries) { + table.insert(kj::mv(KJ_ASSERT_NONNULL(slot))); + } + table.insert(kj::mv(entry)); + storage = kj::mv(table); + return; + } + KJ_CASE_ONEOF(table, TableStorage) { + table.upsert(kj::mv(entry), [](StorageEntry& existing, StorageEntry&& row) mutable { + existing.value = kj::mv(row.value); + }); + } + } } AsyncContextFrame::Scope::Scope(Lock& js, kj::Maybe resource) @@ -179,6 +257,10 @@ AsyncContextFrame::StorageScope::StorageScope(Lock& js, kj::Arc key, : frame(AsyncContextFrame::create(js, StorageEntry(kj::mv(key), kj::mv(store)))), scope(js, *frame) {} +AsyncContextFrame::StorageScope::StorageScope(Lock& js, kj::ArrayPtr storageEntries) + : frame(AsyncContextFrame::create(js, storageEntries)), + scope(js, *frame) {} + v8::Local AsyncContextFrame::getJSWrapper(v8::Isolate* isolate) { KJ_IF_SOME(handle, tryGetHandle(isolate)) { return handle; diff --git a/src/workerd/jsg/async-context.h b/src/workerd/jsg/async-context.h index 40205c7eb37..231792556f4 100644 --- a/src/workerd/jsg/async-context.h +++ b/src/workerd/jsg/async-context.h @@ -121,6 +121,7 @@ class AsyncContextFrame final: public Wrappable { }; AsyncContextFrame(Lock& js, StorageEntry storageEntry); + AsyncContextFrame(Lock& js, kj::ArrayPtr storageEntries); inline Ref addRef() { return JSG_THIS; @@ -141,6 +142,9 @@ class AsyncContextFrame final: public Wrappable { // frame (if any) and the given StorageEntry is added. static Ref create(Lock& js, StorageEntry storageEntry); + // Moves the entries into a single frame, copying inherited storage once. + static Ref create(Lock& js, kj::ArrayPtr storageEntries); + // Wraps the given JavaScript function such that whenever the wrapper function is called, // the root AsyncContextFrame will be entered. static v8::Local wrapRoot( @@ -205,6 +209,7 @@ class AsyncContextFrame final: public Wrappable { Scope scope; StorageScope(Lock& js, kj::Arc key, Value store); + StorageScope(Lock& js, kj::ArrayPtr storageEntries); KJ_DISALLOW_COPY(StorageScope); }; @@ -216,7 +221,19 @@ class AsyncContextFrame final: public Wrappable { } void jsgGetMemoryInfo(MemoryTracker& tracker) const override { Wrappable::jsgGetMemoryInfo(tracker); - tracker.trackField("storage", storage); + KJ_SWITCH_ONEOF(storage) { + KJ_CASE_ONEOF(entries, InlineStorage) { + for (const auto& slot: entries) { + KJ_IF_SOME(entry, slot) { + // The entry's bytes are already included in this frame's self size. + entry.jsgGetMemoryInfo(tracker); + } + } + } + KJ_CASE_ONEOF(table, TableStorage) { + tracker.trackField("storage", table); + } + } } private: @@ -234,8 +251,13 @@ class AsyncContextFrame final: public Wrappable { } }; - using Storage = kj::Table>; - Storage storage; + // Keep the internal and user tracing keys inline. Larger contexts retain indexed lookup. + using InlineStorage = kj::FixedArray, 2>; + using TableStorage = kj::Table>; + kj::OneOf storage = InlineStorage(); + + void eraseDeadStorage(); + void upsertStorageEntry(StorageEntry entry); void jsgVisitForGc(GcVisitor& visitor) override; diff --git a/src/workerd/jsg/bench-async-context.c++ b/src/workerd/jsg/bench-async-context.c++ new file mode 100644 index 00000000000..43a5f89c95f --- /dev/null +++ b/src/workerd/jsg/bench-async-context.c++ @@ -0,0 +1,102 @@ +// 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 + +#include "async-context.h" +#include "jsg-test.h" + +#include + +namespace workerd::jsg::test { +namespace { + +V8System v8System; + +struct AsyncContextBenchmarkContext final: public Object, public ContextGlobal { + JSG_RESOURCE_TYPE(AsyncContextBenchmarkContext) {} +}; +JSG_DECLARE_ISOLATE_TYPE(AsyncContextBenchmarkIsolate, AsyncContextBenchmarkContext); + +enum class ScopeLayout { SEPARATE, COMBINED }; + +template +void withTraceScopes(Lock& js, + const kj::Arc& internalKey, + const kj::Arc& userKey, + Callback&& callback) { + auto value = v8::Undefined(js.v8Isolate).As(); + if constexpr (layout == ScopeLayout::SEPARATE) { + AsyncContextFrame::StorageScope internalScope(js, internalKey.addRef(), js.v8Ref(value)); + AsyncContextFrame::StorageScope userScope(js, userKey.addRef(), js.v8Ref(value)); + callback(); + } else { + AsyncContextFrame::StorageEntry entries[] = { + {internalKey.addRef(), js.v8Ref(value)}, + {userKey.addRef(), js.v8Ref(value)}, + }; + AsyncContextFrame::StorageScope scope(js, entries); + callback(); + } +} + +template +void traceScopes(benchmark::State& state) { + Evaluator evaluator(v8System); + evaluator.run([&](auto& js) { + auto internalKey = kj::arc(); + auto userKey = kj::arc(); + kj::Vector inheritedEntries; + for (auto i: kj::zeroTo(static_cast(state.range(0)))) { + inheritedEntries.add(AsyncContextFrame::StorageEntry(kj::arc(), + js.v8Ref(v8::Integer::New(js.v8Isolate, i).template As()))); + } + kj::Maybe inheritedScope; + if (!inheritedEntries.empty()) { + inheritedScope.emplace(js, inheritedEntries.asPtr()); + } + auto run = [&]() { + v8::HandleScope handleScope(js.v8Isolate); + withTraceScopes(js, internalKey, userKey, + [&]() { benchmark::DoNotOptimize(&KJ_ASSERT_NONNULL(AsyncContextFrame::current(js))); }); + }; + for (auto _: state) { + run(); + } + bench::recordAllocationStats(state, run); + }); + state.SetItemsProcessed(state.iterations()); +} + +void SeparateTraceScopes(benchmark::State& state) { + traceScopes(state); +} + +void CombinedTraceScopes(benchmark::State& state) { + traceScopes(state); +} + +void FrameLookup(benchmark::State& state) { + Evaluator evaluator(v8System); + evaluator.run([&](auto& js) { + kj::Vector> keys; + kj::Vector entries; + for (auto i: kj::zeroTo(static_cast(state.range(0)))) { + keys.add(kj::arc()); + entries.add(AsyncContextFrame::StorageEntry(keys.back().addRef(), + js.v8Ref(v8::Integer::New(js.v8Isolate, i).template As()))); + } + AsyncContextFrame::StorageScope scope(js, entries.asPtr()); + auto& frame = KJ_ASSERT_NONNULL(AsyncContextFrame::current(js)); + for (auto _: state) { + benchmark::DoNotOptimize(&KJ_ASSERT_NONNULL(frame.get(*keys.back()))); + } + }); + state.SetItemsProcessed(state.iterations()); +} + +WD_BENCHMARK(SeparateTraceScopes)->Arg(0)->Arg(1)->Arg(2)->Arg(4)->Arg(16)->Arg(64); +WD_BENCHMARK(CombinedTraceScopes)->Arg(0)->Arg(1)->Arg(2)->Arg(4)->Arg(16)->Arg(64); +WD_BENCHMARK(FrameLookup)->Arg(1)->Arg(2)->Arg(3)->Arg(4)->Arg(16)->Arg(64); + +} // namespace +} // namespace workerd::jsg::test diff --git a/src/workerd/tests/BUILD.bazel b/src/workerd/tests/BUILD.bazel index 45f303c2f93..b46523ca4c0 100644 --- a/src/workerd/tests/BUILD.bazel +++ b/src/workerd/tests/BUILD.bazel @@ -23,6 +23,11 @@ wd_cc_library( }), ) +kj_test( + src = "bench-tools-test.c++", + deps = [":bench-tools"], +) + # Calls kj::setupAsyncIo(). The I/O backend is meant to be the BINARY's choice (kj_test and # wd_cc_benchmark link //src/workerd/util:setup-async-io into every test binary; libraries do not # carry it), but this one still does, so that binaries outside this repo which link the fixture @@ -103,6 +108,12 @@ wd_cc_benchmark( deps = [":test-fixture"], ) +wd_cc_benchmark( + name = "bench-trace-scopes", + srcs = ["bench-trace-scopes.c++"], + deps = [":test-fixture"], +) + wd_cc_benchmark( name = "bench-regex", srcs = ["bench-regex.c++"], diff --git a/src/workerd/tests/bench-tools-test.c++ b/src/workerd/tests/bench-tools-test.c++ new file mode 100644 index 00000000000..ccd662bf8ca --- /dev/null +++ b/src/workerd/tests/bench-tools-test.c++ @@ -0,0 +1,38 @@ +// 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 + +#include "bench-tools.h" + +namespace workerd::bench { +namespace { + +#ifdef WD_USE_TCMALLOC +KJ_TEST("allocation sampling resumes after being disabled") { + auto interval = tcmalloc::MallocExtension::GetProfileSamplingInterval(); + KJ_DEFER(tcmalloc::MallocExtension::SetProfileSamplingInterval(interval)); + + for (uint i = 0; i < 2; ++i) { + tcmalloc::MallocExtension::SetProfileSamplingInterval(0); + auto unsampled = kj::heapArray(4093); + benchmark::DoNotOptimize(unsampled.begin()); + + tcmalloc::MallocExtension::SetProfileSamplingInterval(1); + warmUpAllocationSampling(); + auto token = tcmalloc::MallocExtension::StartAllocationProfiling(); + auto first = kj::heapArray(4093); + auto second = kj::heapArray(4093); + benchmark::DoNotOptimize(first.begin()); + benchmark::DoNotOptimize(second.begin()); + auto profile = kj::mv(token).Stop(); + int64_t allocations = 0; + profile.Iterate([&](const tcmalloc::Profile::Sample& sample) { + if (sample.requested_size == 4093) allocations += sample.count; + }); + KJ_EXPECT(allocations == 2, allocations); + } +} +#endif + +} // namespace +} // namespace workerd::bench diff --git a/src/workerd/tests/bench-tools.h b/src/workerd/tests/bench-tools.h index 33fc97ef2e3..c537b33219a 100644 --- a/src/workerd/tests/bench-tools.h +++ b/src/workerd/tests/bench-tools.h @@ -20,9 +20,9 @@ namespace workerd::bench { struct TcmallocBenchmarkConfig { TcmallocBenchmarkConfig() { - // Disable heap profiling sampling by setting interval to max value. + // Disable heap profiling sampling. // Default is ~512KB which causes probabilistic sampling of allocations. - tcmalloc::MallocExtension::SetProfileSamplingInterval(std::numeric_limits::max()); + tcmalloc::MallocExtension::SetProfileSamplingInterval(0); // Disable GWP-ASan guarded sampling. A negative value disables it. tcmalloc::MallocExtension::SetGuardedSamplingInterval(-1); @@ -35,9 +35,60 @@ struct TcmallocBenchmarkConfig { // Global instance ensures configuration runs before main(). inline TcmallocBenchmarkConfig tcmallocBenchmarkConfig; +inline void warmUpAllocationSampling() { + // Each thread keeps its sampling countdown until it reaches the next sample. Advance it with + // probes until this thread has picked up the new interval, before measuring the callback. + constexpr size_t PROBE_SIZE = (1 << 20) + 17; + for (uint i = 0; i < 256; ++i) { + auto token = tcmalloc::MallocExtension::StartAllocationProfiling(); + auto probe = kj::heapArray(PROBE_SIZE); + benchmark::DoNotOptimize(probe.begin()); + auto profile = kj::mv(token).Stop(); + bool sampled = false; + profile.Iterate([&](const tcmalloc::Profile::Sample& sample) { + if (sample.requested_size == PROBE_SIZE) sampled = true; + }); + if (sampled) return; + } + KJ_FAIL_REQUIRE("Allocation profiling did not sample the warm-up probes"); +} + } // namespace workerd::bench #endif // defined(WD_USE_TCMALLOC) +namespace workerd::bench { + +// Collect allocation samples separately so profiling does not affect the timing measurements. +template +void recordAllocationStats(benchmark::State& state, Callback&& callback) { +#ifdef WD_USE_TCMALLOC + auto samplingInterval = tcmalloc::MallocExtension::GetProfileSamplingInterval(); + tcmalloc::MallocExtension::SetProfileSamplingInterval(1); + KJ_DEFER(tcmalloc::MallocExtension::SetProfileSamplingInterval(samplingInterval)); + warmUpAllocationSampling(); + callback(); + auto token = tcmalloc::MallocExtension::StartAllocationProfiling(); + constexpr uint SAMPLES = 128; + for (uint i = 0; i < SAMPLES; ++i) { + callback(); + } + auto profile = kj::mv(token).Stop(); + int64_t bytes = 0; + int64_t allocations = 0; + profile.Iterate([&](const tcmalloc::Profile::Sample& sample) { + bytes += sample.sum; + allocations += sample.count; + }); + state.counters["native_bytes/op"] = static_cast(bytes) / SAMPLES; + state.counters["native_allocs/op"] = static_cast(allocations) / SAMPLES; +#else + (void)state; + (void)callback; +#endif +} + +} // namespace workerd::bench + // Define a benchmark. Use microseconds instead of nanoseconds by default, most tests run long // enough to not need ns precision. #define WD_BENCHMARK(X) BENCHMARK(X)->Unit(benchmark::kMicrosecond) diff --git a/src/workerd/tests/bench-trace-scopes.c++ b/src/workerd/tests/bench-trace-scopes.c++ new file mode 100644 index 00000000000..071debb1705 --- /dev/null +++ b/src/workerd/tests/bench-trace-scopes.c++ @@ -0,0 +1,70 @@ +// 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 + +#include "bench-tools.h" +#include "test-fixture.h" + +namespace workerd { +namespace { + +enum class ScopeState { LEGACY, CURRENT }; + +// The legacy baseline uses a combined frame with the original unobserved tracing payloads. +jsg::AsyncContextFrame::StorageScope makeLegacyTraceScope(const TestFixture::Environment& env) { + static SpanParent unobservedSpan(nullptr); + auto span = kj::Own(&unobservedSpan, kj::NullDisposer::instance); + auto internalValue = jsg::wrapOpaque(env.js.v8Context(), env.context.addObject(kj::mv(span))); + auto userContext = kj::heap(SpanParent(nullptr), kj::none, kj::none); + auto userValue = jsg::wrapOpaque(env.js.v8Context(), env.context.addObject(kj::mv(userContext))); + jsg::AsyncContextFrame::StorageEntry entries[] = { + {env.lock.getTraceAsyncContextKey(), env.js.v8Ref(internalValue)}, + {env.lock.getUserTraceAsyncContextKey(), env.js.v8Ref(userValue)}, + }; + return jsg::AsyncContextFrame::StorageScope(env.js, entries); +} + +template +void unobservedTraceScopes(benchmark::State& state) { + TestFixture fixture; + fixture.runInIoContext([&](const TestFixture::Environment& env) { + kj::Vector inheritedEntries; + for (auto i: kj::zeroTo(static_cast(state.range(0)))) { + inheritedEntries.add( + jsg::AsyncContextFrame::StorageEntry(kj::arc(), + env.js.v8Ref(v8::Integer::New(env.isolate, i).As()))); + } + kj::Maybe inheritedScope; + if (!inheritedEntries.empty()) inheritedScope.emplace(env.js, inheritedEntries.asPtr()); + auto run = [&]() { + v8::HandleScope handleScope(env.isolate); + auto scope = [&]() { + if constexpr (scopeState == ScopeState::LEGACY) { + return makeLegacyTraceScope(env); + } else { + return env.context.makeRequestTraceScope(env.lock); + } + }(); + benchmark::DoNotOptimize(&KJ_ASSERT_NONNULL(jsg::AsyncContextFrame::current(env.js))); + }; + for (auto _: state) { + run(); + } + bench::recordAllocationStats(state, run); + }); + state.SetItemsProcessed(state.iterations()); +} + +void LegacyUnobservedTraceScopes(benchmark::State& state) { + unobservedTraceScopes(state); +} + +void UnobservedTraceScopes(benchmark::State& state) { + unobservedTraceScopes(state); +} + +WD_BENCHMARK(LegacyUnobservedTraceScopes)->Arg(0)->Arg(4)->Arg(16); +WD_BENCHMARK(UnobservedTraceScopes)->Arg(0)->Arg(4)->Arg(16); + +} // namespace +} // namespace workerd