Merge pull request #331 from cloudflare/jsnell/asynccontext-additiona…

…l-tweaks

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);