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