Experimental implementation of AsyncLocalStorage

src/workerd/api/node/async-hooks.c++

@@ -0,0 +1,59 @@
+#include "async-hooks.h"
+#include <workerd/jsg/async-context.h>
+#include <kj/vector.h>
+
+namespace workerd::api::node {
+
+jsg::Ref<AsyncLocalStorage> AsyncLocalStorage::constructor() {
+  return jsg::alloc<AsyncLocalStorage>();
+}
+
+v8::Local<v8::Value> AsyncLocalStorage::run(
+    jsg::Lock& js,
+    v8::Local<v8::Value> store,
+    v8::Local<v8::Function> callback,
+    jsg::Varargs args) {
+
+  jsg::AsyncResource::RunScope runScope(js, js.v8Ref(store));
+
+  kj::Vector<v8::Local<v8::Value>> argv(args.size());
+  for (auto arg : args) {
+    argv.add(arg.getHandle(js));
+  }
+
+  auto context = js.v8Isolate->GetCurrentContext();
+
+  return jsg::check(callback->Call(
+      context,
+      context->Global(),
+      argv.size(),
+      argv.begin()));
+}
+
+v8::Local<v8::Value> AsyncLocalStorage::exit(
+    jsg::Lock& js,
+    v8::Local<v8::Function> callback,
+    jsg::Varargs args) {
+  // Node.js defines exit as running "a function synchronously outside of a context".
+  // It goes on to say that the store is not accessible within the callback or the
+  // asynchronous operations created within the callback. Any getStore() call done
+  // within the callbackfunction will always return undefined... except run() is
+  // called which implicitly enables the context again.
+  //
+  // We do not have to emulate Node.js enable/disable behavior since we are not
+  // implementing Node.js' enterWith/disable methods. We can emulate the correct
+  // behavior simply by running and setting the store value to undefined, which
+  // will propagate correctly.
+  return run(js, v8::Undefined(js.v8Isolate), callback, kj::mv(args));
+}
+
+v8::Local<v8::Value> AsyncLocalStorage::getStore(jsg::Lock& js) {
+  auto& current = jsg::AsyncResource::current(js);
+KJ_DBG(&current);
+  KJ_IF_MAYBE(store, current.store) {
+    return store->getHandle(js);
+  }
+  return v8::Undefined(js.v8Isolate);
+}
+
+}  // namespace workerd::api::node

src/workerd/api/node/async-hooks.h

@@ -0,0 +1,56 @@
+#pragma once
+
+#include <workerd/jsg/jsg.h>
+
+namespace workerd::api::node {
+
+class AsyncLocalStorage final: public jsg::Object {
+  // Implements a subset of the Node.js AsyncLocalStorage API.
+public:
+  AsyncLocalStorage() = default;
+
+  static jsg::Ref<AsyncLocalStorage> constructor();
+
+  v8::Local<v8::Value> run(jsg::Lock& js,
+                           v8::Local<v8::Value> store,
+                           v8::Local<v8::Function> callback,
+                           jsg::Varargs args);
+
+  v8::Local<v8::Value> exit(jsg::Lock& js,
+                           v8::Local<v8::Function> callback,
+                           jsg::Varargs args);
+
+  v8::Local<v8::Value> getStore(jsg::Lock& js);
+
+  inline void enterWith(jsg::Lock&, v8::Local<v8::Value>) {
+    KJ_UNIMPLEMENTED("asyncLocalStorage.enterWith() is not implemented");
+  }
+
+  inline void disable(jsg::Lock&) {
+    KJ_UNIMPLEMENTED("asyncLocalStorage.disable() is not implemented");
+  }
+
+  JSG_RESOURCE_TYPE(AsyncLocalStorage) {
+    JSG_METHOD(run);
+    JSG_METHOD(exit);
+    JSG_METHOD(getStore);
+    JSG_METHOD(enterWith);
+    JSG_METHOD(disable);
+  }
+};
+
+class AsyncHooksModule final: public jsg::Object {
+  // We have no intention of fully-implementing the Node.js async_hooks module.
+  // We provide this because AsyncLocalStorage is exposed via async_hooks in
+  // Node.js.
+public:
+  JSG_RESOURCE_TYPE(AsyncHooksModule) {
+    JSG_NESTED_TYPE(AsyncLocalStorage);
+  }
+};
+
+#define EW_NODE_ASYNCHOOKS_ISOLATE_TYPES       \
+    api::node::AsyncHooksModule,               \
+    api::node::AsyncLocalStorage
+
+}  // namespace workerd::api::node

src/workerd/server/workerd-api.c++

@@ -18,6 +18,7 @@
 #include <workerd/api/r2.h>
 #include <workerd/api/r2-admin.h>
 #include <workerd/api/urlpattern.h>
+#include <workerd/api/node/async-hooks.h>
 #include <workerd/util/thread-scopes.h>
 #include <openssl/sha.h>
 #include <openssl/hmac.h>
@@ -63,6 +64,7 @@ JSG_DECLARE_ISOLATE_TYPE(JsgWorkerdIsolate,
   EW_URL_STANDARD_ISOLATE_TYPES,
   EW_URLPATTERN_ISOLATE_TYPES,
   EW_WEBSOCKET_ISOLATE_TYPES,
+  EW_NODE_ASYNCHOOKS_ISOLATE_TYPES,
 
   jsg::TypeWrapperExtension<PromiseWrapper>,
   jsg::InjectConfiguration<CompatibilityFlags::Reader>,