Initial experimental implementation of async context tracking

src/workerd/api/global-scope.c++

@@ -7,6 +7,7 @@
 #include <kj/encoding.h>
 
 #include <workerd/api/system-streams.h>
+#include <workerd/jsg/async-context.h>
 #include <workerd/jsg/ser.h>
 #include <workerd/jsg/util.h>
 #include <workerd/io/trace.h>
@@ -479,7 +480,7 @@ v8::Local<v8::String> ServiceWorkerGlobalScope::atob(kj::String data, v8::Isolat
 }
 
 void ServiceWorkerGlobalScope::queueMicrotask(v8::Local<v8::Function> task, v8::Isolate* isolate) {
-  isolate->EnqueueMicrotask(task);
+  isolate->EnqueueMicrotask(jsg::AsyncResource::wrap(jsg::Lock::from(isolate), task));
 }
 
 v8::Local<v8::Value> ServiceWorkerGlobalScope::structuredClone(
@@ -516,7 +517,9 @@ TimeoutId::NumberType ServiceWorkerGlobalScope::setTimeout(
   auto timeoutId = IoContext::current().setTimeoutImpl(
     timeoutIdGenerator,
     /* repeats = */ false,
-    [function = function.addRef(isolate), argv = kj::mv(argv)](jsg::Lock& js) mutable {
+    [function = function.addRef(isolate),
+     argv = kj::mv(argv),
+     resource = jsg::AsyncResource::create(jsg::Lock::from(isolate))](jsg::Lock& js) mutable {
       auto isolate = js.v8Isolate;
       auto context = isolate->GetCurrentContext();
       auto localFunction = function.getHandle(isolate);
@@ -525,6 +528,8 @@ TimeoutId::NumberType ServiceWorkerGlobalScope::setTimeout(
       };
       auto argc = localArgs.size();
 
+      jsg::AsyncResource::Scope scope(js, resource);
+
       // Cast to void to discard the result value.
       (void)jsg::check(localFunction->Call(context, context->Global(), argc, &localArgs.front()));
     },
@@ -547,7 +552,9 @@ TimeoutId::NumberType ServiceWorkerGlobalScope::setInterval(
   auto timeoutId = IoContext::current().setTimeoutImpl(
     timeoutIdGenerator,
     /* repeats = */ true,
-    [function = function.addRef(isolate), argv = kj::mv(argv)](jsg::Lock& js) mutable {
+    [function = function.addRef(isolate),
+     argv = kj::mv(argv),
+     resource = jsg::AsyncResource::create(jsg::Lock::from(isolate))](jsg::Lock& js) mutable {
       auto isolate = js.v8Isolate;
       auto context = isolate->GetCurrentContext();
       auto localFunction = function.getHandle(isolate);
@@ -556,6 +563,8 @@ TimeoutId::NumberType ServiceWorkerGlobalScope::setInterval(
       };
       auto argc = localArgs.size();
 
+      jsg::AsyncResource::Scope scope(js, resource);
+
       // Cast to void to discard the result value.
       (void)jsg::check(localFunction->Call(context, context->Global(), argc, &localArgs.front()));
     },

src/workerd/jsg/async-context.c++

@@ -0,0 +1,191 @@
+#include "async-context.h"
+#include "jsg.h"
+#include "setup.h"
+#include <v8.h>
+
+namespace workerd::jsg {
+
+namespace {
+struct AsyncResourceWrappable final: public Wrappable,
+                                     public AsyncResource {
+  // Used to attach async context to JS objects like Promises.
+  using AsyncResource::AsyncResource;
+
+  static v8::Local<v8::Value> wrap(Lock& js,
+                                   uint64_t id,
+                                   kj::Maybe<AsyncResource&> maybeParent) {
+    auto wrapped = kj::refcounted<AsyncResourceWrappable>(js, id, maybeParent);
+    return wrapped->attachOpaqueWrapper(js.v8Isolate->GetCurrentContext(), false);
+  }
+
+  static kj::Maybe<AsyncResource&> tryUnwrap(v8::Isolate* isolate,
+                                                   v8::Local<v8::Value> handle) {
+    KJ_IF_MAYBE(wrappable, Wrappable::tryUnwrapOpaque(isolate, handle)) {
+      return dynamic_cast<AsyncResource&>(*wrappable);
+    }
+    return nullptr;
+  }
+};
+
+kj::Maybe<Value> propagateStore(Lock& js, kj::Maybe<AsyncResource&>& maybeParent) {
+  KJ_IF_MAYBE(parent, maybeParent) {
+    return parent->store.map([&js](auto& value) mutable -> Value {
+      return value.addRef(js);
+    });
+  }
+  return nullptr;
+}
+
+}  // namespace
+
+AsyncResource::AsyncResource(
+    Lock& js,
+    uint64_t id,
+    kj::Maybe<AsyncResource&> maybeParent)
+    : id(id),
+      parentId(maybeParent.map([](auto& parent) {
+        return parent.id;
+      })),
+      store(propagateStore(js, maybeParent)) {}
+
+AsyncResource& AsyncResource::current(Lock& js) {
+  auto& isolateBase = IsolateBase::from(js.v8Isolate);
+  KJ_ASSERT(!isolateBase.asyncResourceStack.empty());
+  return *isolateBase.asyncResourceStack.front();
+}
+
+AsyncResource AsyncResource::create(Lock& js, kj::Maybe<AsyncResource&> maybeParent) {
+  auto id = IsolateBase::from(js.v8Isolate).getNextAsyncResourceId();
+  KJ_IF_MAYBE(parent, maybeParent) {
+    KJ_ASSERT(id > parent->id);
+    return AsyncResource(js, id, *parent);
+  }
+  return AsyncResource(js, id, current(js));
+}
+
+v8::Local<v8::Function> AsyncResource::wrap(
+    Lock& js,
+    v8::Local<v8::Function> fn,
+    kj::Maybe<AsyncResource&> maybeParent) {
+  auto isolate = js.v8Isolate;
+  auto context = isolate->GetCurrentContext();
+  auto handle = v8::Private::ForApi(isolate, v8StrIntern(isolate, "asyncResource"));
+  if (!fn->HasPrivate(context, handle).FromJust()) {
+    auto id = IsolateBase::from(isolate).getNextAsyncResourceId();
+    auto obj = AsyncResourceWrappable::wrap(js, id,
+        maybeParent.orDefault(AsyncResource::current(js)));
+    KJ_ASSERT(check(fn->SetPrivate(context, handle, obj)));
+  }
+
+  return jsg::check(v8::Function::New(context, [](const v8::FunctionCallbackInfo<v8::Value>& args) {
+    auto isolate = args.GetIsolate();
+    auto context = isolate->GetCurrentContext();
+    auto fn = args.Data().As<v8::Function>();
+    auto handle = v8::Private::ForApi(isolate, v8StrIntern(isolate, "asyncResource"));
+    auto& resource = KJ_ASSERT_NONNULL(AsyncResourceWrappable::tryUnwrap(isolate,
+        check(fn->GetPrivate(context, handle))));
+
+    AsyncResource::Scope scope(jsg::Lock::from(isolate), resource);
+    jsg::check(fn->Call(context, v8::Undefined(isolate), 0, nullptr));
+  }, fn));
+}
+
+AsyncResource::Scope::Scope(Lock& js, AsyncResource& resource)
+    : Scope(js.v8Isolate, resource) {}
+
+AsyncResource::Scope::Scope(v8::Isolate* isolate, AsyncResource& resource)
+    : isolate(IsolateBase::from(isolate)) {
+  this->isolate.pushAsyncResource(resource);
+}
+
+AsyncResource::Scope::~Scope() noexcept(false) {
+  isolate.popAsyncResource();
+}
+
+AsyncResource::RunScope::RunScope(Lock& js, Value store)
+    : RunScope(js, AsyncResource::current(js),
+      kj::mv(store)) {}
+
+AsyncResource::RunScope::RunScope(Lock& js, AsyncResource& resource, Value store)
+    : Scope(js, resource), resource(resource), oldStore(kj::mv(resource.store)) {
+  resource.store = kj::mv(store);
+}
+
+AsyncResource::RunScope::~RunScope() noexcept(false) {
+  resource.store = kj::mv(oldStore);
+}
+
+void IsolateBase::pushAsyncResource(AsyncResource& next) {
+  asyncResourceStack.push_front(&next);
+}
+
+void IsolateBase::popAsyncResource() {
+  asyncResourceStack.pop_front();
+  KJ_ASSERT(!asyncResourceStack.empty(), "the async resource stack was corrupted");
+}
+
+void IsolateBase::promiseHook(v8::PromiseHookType type,
+                              v8::Local<v8::Promise> promise,
+                              v8::Local<v8::Value> parent) {
+  auto isolate = promise->GetIsolate();
+  auto context = isolate->GetCurrentContext();
+  auto& isolateBase = IsolateBase::from(isolate);
+  auto& js = Lock::from(isolate);
+
+  auto handle = v8::Private::ForApi(isolate, v8StrIntern(isolate, "asyncResource"));
+
+  const auto tryGetAsyncResource = [&](v8::Local<v8::Promise> promise)
+      -> kj::Maybe<AsyncResource&> {
+    return AsyncResourceWrappable::tryUnwrap(isolate, check(promise->GetPrivate(context, handle)));
+  };
+
+  const auto createAsyncResource = [&](
+      v8::Local<v8::Promise> promise,
+      kj::Maybe<AsyncResource&> maybeParent)
+        -> AsyncResource& {
+    auto id = IsolateBase::from(isolate).getNextAsyncResourceId();
+    auto obj = AsyncResourceWrappable::wrap(js, id, maybeParent);
+    check(promise->SetPrivate(context, handle, obj));
+    return KJ_ASSERT_NONNULL(tryGetAsyncResource(promise));
+  };
+
+  const auto trackPromise = [&](
+      v8::Local<v8::Promise> promise,
+      v8::Local<v8::Value> parent) -> AsyncResource& {
+    KJ_IF_MAYBE(asyncResource, tryGetAsyncResource(promise)) {
+      return *asyncResource;
+    }
+    kj::Maybe<AsyncResource&> maybeParent = nullptr;
+    if (parent->IsPromise()) {
+      auto parentPromise = parent.As<v8::Promise>();
+      KJ_IF_MAYBE(asyncResource, tryGetAsyncResource(parentPromise)) {
+        maybeParent = *asyncResource;
+      } else {
+        maybeParent = createAsyncResource(parent.As<v8::Promise>(), AsyncResource::current(js));
+      }
+    }
+    return createAsyncResource(promise, maybeParent);
+  };
+
+  switch (type) {
+    case v8::PromiseHookType::kInit: {
+      trackPromise(promise, parent);
+      break;
+    }
+    case v8::PromiseHookType::kBefore: {
+      auto& resource = trackPromise(promise, parent);
+      isolateBase.pushAsyncResource(resource);
+      break;
+    }
+    case v8::PromiseHookType::kAfter: {
+      isolateBase.popAsyncResource();
+      break;
+    }
+    case v8::PromiseHookType::kResolve: {
+      // There's nothing to do here.
+      break;
+    }
+  }
+}
+
+}  // namespace workerd::jsg

src/workerd/jsg/async-context.h

@@ -0,0 +1,58 @@
+#pragma once
+
+#include "jsg.h"
+#include <v8.h>
+
+namespace workerd::jsg {
+
+// Provides for basic internal async context tracking. Eventually, it is expected that
+// this will be provided by V8 assuming that the AsyncContext proposal advances through
+// TC-39. For now, however, we implement a model that is very similar to that implemented
+// by Node.js.
+
+struct AsyncResource {
+  const uint64_t id;
+  const kj::Maybe<uint64_t> parentId;
+  kj::Maybe<Value> store;
+
+  inline explicit AsyncResource() : id(0) {}
+
+  explicit AsyncResource(
+      Lock& js,
+      uint64_t id,
+      kj::Maybe<AsyncResource&> maybeParent = nullptr);
+
+  static AsyncResource& current(Lock& js);
+
+  static AsyncResource create(Lock& js, kj::Maybe<AsyncResource&> maybeParent = nullptr);
+  // Create a new AsyncResource. If maybeParent is not specified, uses the current().
+
+  static v8::Local<v8::Function> wrap(Lock& js, v8::Local<v8::Function> fn,
+                                      kj::Maybe<AsyncResource&> maybeParent = nullptr);
+  // Treats the given JavaScript function as an async resource and returns a wrapper
+  // function that will ensure appropriate propagation of the async context tracking
+  // when the wrapper function is called.
+
+  struct Scope {
+    // AsyncResource::Scope makes the given AsyncResource the current in the
+    // stack until it is destroyed.
+    IsolateBase& isolate;
+    Scope(Lock& js, AsyncResource& resource);
+    Scope(v8::Isolate* isolate, AsyncResource& resource);
+    ~Scope() noexcept(false);
+  };
+
+  struct RunScope: Scope {
+    // RunScope makes the given AsyncResource the current in the stack and
+    // sets the stored context value. Pops the stack and resets the value
+    // when destroyed.
+    AsyncResource& resource;
+    kj::Maybe<Value> oldStore;
+
+    RunScope(Lock& js, AsyncResource& resource, Value store);
+    RunScope(Lock& js, Value store);
+    ~RunScope() noexcept(false);
+  };
+};
+
+}  // namespace workerd::jsg