Add AsyncLocalStorage.bind/snapshot

These are proposed (in progress) additions to AsyncLocalStorage
in Node.js that serve a dual purpose:

1. They align the API closer to the expected AsyncContext.wrap
   api (AsyncLocalStorage.bind == AsyncContext.wrap). It uses
   the existing naming from AsyncResource for consistency with
   the existing API.
2. They eliminate the need to use AsyncResource. We will keep
   AsyncResource for backwards compatibility as part of the
   larger Node.js compat story, but these cover all of the
   key use cases of AsyncResource for context tracking.

This should not land until nodejs/node#46387 lands.

Author: jasnell

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

@@ -57,6 +57,18 @@ v8::Local<v8::Value> AsyncLocalStorage::getStore(jsg::Lock& js) {
   return v8::Undefined(js.v8Isolate);
 }
 
+v8::Local<v8::Function> AsyncLocalStorage::bind(jsg::Lock& js, v8::Local<v8::Function> fn) {
+  KJ_IF_MAYBE(frame, jsg::AsyncContextFrame::current(js)) {
+    return frame->wrap(js, fn);
+  } else {
+    return jsg::AsyncContextFrame::wrapRoot(js, fn);
+  }
+}
+
+v8::Local<v8::Function> AsyncLocalStorage::snapshot(jsg::Lock& js) {
+  return jsg::AsyncContextFrame::wrapSnapshot(js);
+}
+
 namespace {
 kj::Maybe<jsg::Ref<jsg::AsyncContextFrame>> tryGetFrameRef(jsg::Lock& js) {
   return jsg::AsyncContextFrame::current(js).map(

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

@@ -46,6 +46,15 @@ class AsyncLocalStorage final: public jsg::Object {
 
   v8::Local<v8::Value> getStore(jsg::Lock& js);
 
+  static v8::Local<v8::Function> bind(jsg::Lock& js, v8::Local<v8::Function> fn);
+  // Binds the given function to the current async context frame such that
+  // whenever the function is called, the bound frame is entered.
+
+  static v8::Local<v8::Function> snapshot(jsg::Lock& js);
+  // Returns a function bound to the current async context frame that calls
+  // the function passed to it as the only argument within that frame.
+  // Equivalent to AsyncLocalStorage.bind((cb, ...args) => cb(...args)).
+
   inline void enterWith(jsg::Lock&, v8::Local<v8::Value>) {
     KJ_UNIMPLEMENTED("asyncLocalStorage.enterWith() is not implemented");
   }
@@ -60,6 +69,8 @@ class AsyncLocalStorage final: public jsg::Object {
     JSG_METHOD(getStore);
     JSG_METHOD(enterWith);
     JSG_METHOD(disable);
+    JSG_STATIC_METHOD(bind);
+    JSG_STATIC_METHOD(snapshot);
 
     if (flags.getNodeJsCompat()) {
       JSG_TS_OVERRIDE(AsyncLocalStorage<T> {
@@ -68,6 +79,8 @@ class AsyncLocalStorage final: public jsg::Object {
         exit<R, TArgs extends any[]>(callback: (...args: TArgs) => R, ...args: TArgs): R;
         disable(): void;
         enterWith(store: T): void;
+        static bind<Func extends (...args: any[]) => any>(fn: Func): Func;
+        static snapshot<R, TArgs extends any[]>() : ((...args: TArgs) => R, ...args: TArgs) => R;
       });
     } else {
       JSG_TS_OVERRIDE(type AsyncLocalStorage = never);
@@ -80,6 +93,10 @@ class AsyncLocalStorage final: public jsg::Object {
 
 
 class AsyncResource final: public jsg::Object {
+  // Note: The AsyncResource class is provided for Node.js backwards compatibility.
+  // The class can be replaced entirely for async context tracking using the
+  // AsyncLocalStorage.bind() and AsyncLocalStorage.snapshot() APIs.
+  //
   // The AsyncResource class is an object that user code can use to define its own
   // async resources for the purpose of storage context propagation. For instance,
   // let's imagine that we have an EventTarget and we want to register two event listeners

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

@@ -69,6 +69,28 @@ v8::Local<v8::Function> AsyncContextFrame::wrap(
   return wrap(js, fn.getHandle(js), thisArg);
 }
 
+v8::Local<v8::Function> AsyncContextFrame::wrapSnapshot(Lock& js) {
+  auto isolate = js.v8Isolate;
+  auto context = isolate->GetCurrentContext();
+
+  return js.wrapReturningFunction(context, JSG_VISITABLE_LAMBDA(
+    (frame = AsyncContextFrame::currentRef(js)),
+    (frame),
+    (Lock& js, const v8::FunctionCallbackInfo<v8::Value>& args) {
+      auto context = js.v8Isolate->GetCurrentContext();
+      JSG_REQUIRE(args[0]->IsFunction(), TypeError, "The first argument must be a function");
+      auto fn = args[0].As<v8::Function>();
+      kj::Vector<v8::Local<v8::Value>> argv(args.Length() - 1);
+      for (int n = 1; n < args.Length(); n++) {
+        argv.add(args[n]);
+      }
+
+      AsyncContextFrame::Scope scope(js, frame);
+      return check(fn->Call(context, context->Global(), argv.size(), argv.begin()));
+    }
+  ));
+}
+
 v8::Local<v8::Function> AsyncContextFrame::wrap(
     Lock& js,
     v8::Local<v8::Function> fn,
@@ -93,7 +115,6 @@ v8::Local<v8::Function> AsyncContextFrame::wrap(
     }
 
     AsyncContextFrame::Scope scope(js, *frame.get());
-    v8::Local<v8::Value> result;
     return check(function->Call(context, thisArg.getHandle(js), args.Length(), argv.begin()));
   }));
 }

src/workerd/jsg/async-context.h

@@ -125,6 +125,11 @@ class AsyncContextFrame final: public Wrappable {
   // Wraps the given JavaScript function such that whenever the wrapper function is called,
   // the root AsyncContextFrame will be entered.
 
+  static v8::Local<v8::Function> wrapSnapshot(Lock& js);
+  // Returns a function that captures the current frame and calls the function passed
+  // in as an argument within that captured context. Equivalent to wrapping a function
+  // with the signature (cb, ...args) => cb(...args).
+
   v8::Local<v8::Function> wrap(
       Lock& js, V8Ref<v8::Function>& fn,
       kj::Maybe<v8::Local<v8::Value>> thisArg = nullptr);