Optimize tracing memory usage

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<AsyncContextFrame::StorageKey> 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`)

src/workerd/api/tracing.c++

@@ -719,7 +719,7 @@ jsg::Optional<jsg::Ref<user_tracing::Span>> Tracing::getActiveSpan(
     auto key = ioContext.getCurrentLock().getUserTraceAsyncContextKey();
     KJ_IF_SOME(value, frame.get(*key)) {
       auto holder = value.getHandle(js).As<v8::Object>();
-      auto& asyncContext = jsg::unwrapOpaqueRef<IoOwn<UserTraceAsyncContext>>(js.v8Isolate, holder);
+      auto& scopeState = jsg::unwrapOpaqueRef<IoContext::UserTraceScopeState>(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<jsg::Ref<user_tracing::Span>> Tracing::getActiveSpan(
         }
       }
 
-      kj::Maybe<kj::Own<BaseTracer::WeakRef>> tracer;
-      KJ_IF_SOME(value, asyncContext->getTracer()) {
-        tracer = value.addRef();
+      kj::Own<user_tracing::SpanState> state;
+      KJ_SWITCH_ONEOF(scopeState) {
+        KJ_CASE_ONEOF(id, IoContext::Id) {
+          if (!id.isCurrent()) IoContext::throwNotCurrentJsError(typeid(UserTraceAsyncContext));
+          state = kj::refcounted<user_tracing::NoopSpanState>();
+        }
+        KJ_CASE_ONEOF(asyncContext, IoOwn<UserTraceAsyncContext>) {
+          kj::Maybe<kj::Own<BaseTracer::WeakRef>> tracer;
+          KJ_IF_SOME(value, asyncContext->getTracer()) {
+            tracer = value.addRef();
+          }
+          state = kj::refcounted<user_tracing::InvocationSpanState>(asyncContext->getSpan(),
+              kj::mv(tracer), asyncContext->getInvocationSpanContext().map([](auto& context) {
+            return context.clone();
+          }));
+        }
       }
-      kj::Own<user_tracing::SpanState> state =
-          kj::refcounted<user_tracing::InvocationSpanState>(asyncContext->getSpan(), kj::mv(tracer),
-              asyncContext->getInvocationSpanContext().map(
-                  [](auto& context) { return context.clone(); }));
       auto span = js.alloc<user_tracing::Span>(ioContext.addObject(kj::mv(state)));
       auto wrapped = spanHandler.wrap(js, span.addRef());
       jsg::check(holder->SetPrivate(js.v8Context(), cacheKey, wrapped));

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 = [

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<RequestObserver>()>(
+          []() -> kj::Own<RequestObserver> { return kj::refcounted<TraceScopeRequestObserver>(); }),
+    });
+    auto context = fixture.newIoContext();
+    kj::Maybe<kj::Rc<BaseTracer>> tracer;
+    if (observeUser) {
+      auto trace = kj::refcounted<Trace>(kj::none, kj::none, kj::none, kj::none, kj::none,
+          kj::Array<kj::String>(), kj::none, ExecutionModel::STATELESS);
+      auto workerTracer =
+          kj::rc<WorkerTracer>(kj::none, kj::mv(trace), PipelineLogLevel::NONE, kj::none, kj::none);
+      workerTracer->setMakeUserRequestSpanFunc(
+          [](tracing::TraceId, kj::Maybe<tracing::TraceFlags>) { return traceScopeParent(22); });
+      tracer = kj::Rc<BaseTracer>(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::StorageKey>();
+      jsg::AsyncContextFrame::StorageScope storageScope(
+          env.js, key.addRef(), env.js.v8Ref(v8::Integer::New(env.isolate, 55).As<v8::Value>()));
+      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<jsg::Ref<jsg::AsyncContextFrame>> 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<v8::Int32>()->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<kj::Rc<IoChannelFactory>(TimerChannel&)> makeChannelFactory =

src/workerd/io/io-context.c++

@@ -1239,31 +1239,38 @@ kj::Own<CacheClient> IoContext::getCacheClient() {
 
 jsg::AsyncContextFrame::StorageScope IoContext::makeAsyncTraceScope(
     Worker::Lock& lock, kj::Maybe<SpanParent> 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<SpanParent> 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<SpanParent> spanParentOverride) {
+  auto& ioContext = IoContext::current();
   jsg::Lock& js = lock;
-  kj::Own<SpanParent> 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<SpanParent>(
-          &const_cast<SpanParent&>(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<SpanParent> 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<tracing::InvocationSpanContext> 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<kj::Own<workerd::WeakRef<BaseTracer>>> tracer;
@@ -1294,9 +1303,9 @@ jsg::AsyncContextFrame::StorageScope IoContext::makeUserAsyncTraceScope(
   auto asyncContext = kj::heap<UserTraceAsyncContext>(
       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<IoOwn<SpanParent>>(js.v8Isolate, handle);
+        auto& state = jsg::unwrapOpaqueRef<TraceScopeState>(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<SpanParent>) {
+            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<IoOwn<UserTraceAsyncContext>>(js.v8Isolate, handle);
+        auto& state = jsg::unwrapOpaqueRef<UserTraceScopeState>(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<UserTraceAsyncContext>) {
+            KJ_IF_SOME(context, asyncContext.tryGet()) {
+              return context.getSpan();
+            }
+          }
         }
       }
     }

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<Id, IoOwn<SpanParent>>;
+  using UserTraceScopeState = kj::OneOf<Id, IoOwn<UserTraceAsyncContext>>;
+
   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<SpanParent> 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<SpanParent> spanParentOverride);
+  jsg::Value makeUserTraceScopeValue(Worker::Lock& lock, kj::Maybe<SpanParent> userSpanOverride);
+
   ThreadContext& thread;
 
   kj::Own<WeakRef> selfRef = kj::refcounted<WeakRef>(kj::Badge<IoContext>(), *this);