From dbf9bc6500932c4ef3b882354edfc7503ff02283 Mon Sep 17 00:00:00 2001 From: James M Snell Date: Mon, 30 Jan 2023 16:50:32 -0800 Subject: [PATCH] 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 https://github.com/nodejs/node/pull/46387 lands. --- src/workerd/api/node/async-hooks.c++ | 12 ++++++++++++ src/workerd/api/node/async-hooks.h | 17 +++++++++++++++++ src/workerd/jsg/async-context.c++ | 23 ++++++++++++++++++++++- src/workerd/jsg/async-context.h | 5 +++++ 4 files changed, 56 insertions(+), 1 deletion(-) diff --git a/src/workerd/api/node/async-hooks.c++ b/src/workerd/api/node/async-hooks.c++ index 14b567f644e..bc15ec946ad 100644 --- a/src/workerd/api/node/async-hooks.c++ +++ b/src/workerd/api/node/async-hooks.c++ @@ -57,6 +57,18 @@ v8::Local AsyncLocalStorage::getStore(jsg::Lock& js) { return v8::Undefined(js.v8Isolate); } +v8::Local AsyncLocalStorage::bind(jsg::Lock& js, v8::Local fn) { + KJ_IF_MAYBE(frame, jsg::AsyncContextFrame::current(js)) { + return frame->wrap(js, fn); + } else { + return jsg::AsyncContextFrame::wrapRoot(js, fn); + } +} + +v8::Local AsyncLocalStorage::snapshot(jsg::Lock& js) { + return jsg::AsyncContextFrame::wrapSnapshot(js); +} + namespace { kj::Maybe> tryGetFrameRef(jsg::Lock& js) { return jsg::AsyncContextFrame::current(js).map( diff --git a/src/workerd/api/node/async-hooks.h b/src/workerd/api/node/async-hooks.h index 764caa459e7..42c33bb068e 100644 --- a/src/workerd/api/node/async-hooks.h +++ b/src/workerd/api/node/async-hooks.h @@ -46,6 +46,15 @@ class AsyncLocalStorage final: public jsg::Object { v8::Local getStore(jsg::Lock& js); + static v8::Local bind(jsg::Lock& js, v8::Local 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 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) { 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 { @@ -68,6 +79,8 @@ class AsyncLocalStorage final: public jsg::Object { exit(callback: (...args: TArgs) => R, ...args: TArgs): R; disable(): void; enterWith(store: T): void; + static bind any>(fn: Func): Func; + static snapshot() : ((...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 diff --git a/src/workerd/jsg/async-context.c++ b/src/workerd/jsg/async-context.c++ index 1e1679ccaad..8ff95573290 100644 --- a/src/workerd/jsg/async-context.c++ +++ b/src/workerd/jsg/async-context.c++ @@ -69,6 +69,28 @@ v8::Local AsyncContextFrame::wrap( return wrap(js, fn.getHandle(js), thisArg); } +v8::Local 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& args) { + auto context = js.v8Isolate->GetCurrentContext(); + JSG_REQUIRE(args[0]->IsFunction(), TypeError, "The first argument must be a function"); + auto fn = args[0].As(); + kj::Vector> 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 AsyncContextFrame::wrap( Lock& js, v8::Local fn, @@ -93,7 +115,6 @@ v8::Local AsyncContextFrame::wrap( } AsyncContextFrame::Scope scope(js, *frame.get()); - v8::Local result; return check(function->Call(context, thisArg.getHandle(js), args.Length(), argv.begin())); })); } diff --git a/src/workerd/jsg/async-context.h b/src/workerd/jsg/async-context.h index 8b5fddc0cf6..a53b5631f2c 100644 --- a/src/workerd/jsg/async-context.h +++ b/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 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 wrap( Lock& js, V8Ref& fn, kj::Maybe> thisArg = nullptr);