diff --git a/src/async-wrap.cc b/src/async-wrap.cc index 2dcfdc799720ed..c4777ac827a685 100644 --- a/src/async-wrap.cc +++ b/src/async-wrap.cc @@ -422,6 +422,18 @@ static void SetupHooks(const FunctionCallbackInfo& args) { } +static void EnablePromiseHook(const FunctionCallbackInfo& args) { + Environment* env = Environment::GetCurrent(args); + env->AddPromiseHook(PromiseHook, nullptr); +} + + +static void DisablePromiseHook(const FunctionCallbackInfo& args) { + Environment* env = Environment::GetCurrent(args); + env->RemovePromiseHook(PromiseHook, nullptr); +} + + void AsyncWrap::GetAsyncId(const FunctionCallbackInfo& args) { AsyncWrap* wrap; args.GetReturnValue().Set(-1); @@ -478,6 +490,8 @@ void AsyncWrap::Initialize(Local target, env->SetMethod(target, "popAsyncIds", PopAsyncIds); env->SetMethod(target, "clearIdStack", ClearIdStack); env->SetMethod(target, "addIdToDestroyList", QueueDestroyId); + env->SetMethod(target, "enablePromiseHook", EnablePromiseHook); + env->SetMethod(target, "disablePromiseHook", DisablePromiseHook); v8::PropertyAttribute ReadOnlyDontDelete = static_cast(v8::ReadOnly | v8::DontDelete); diff --git a/src/env.cc b/src/env.cc index 0bf4b573aafd69..c97868a5882b41 100644 --- a/src/env.cc +++ b/src/env.cc @@ -11,6 +11,7 @@ #endif #include +#include namespace node { @@ -178,12 +179,36 @@ void Environment::AtExit(void (*cb)(void* arg), void* arg) { } void Environment::AddPromiseHook(promise_hook_func fn, void* arg) { + auto it = std::find_if( + promise_hooks_.begin(), promise_hooks_.end(), + [&](const PromiseHookCallback& hook) { + return hook.cb_ == fn && hook.arg_ == arg; + }); + CHECK_EQ(it, promise_hooks_.end()); promise_hooks_.push_back(PromiseHookCallback{fn, arg}); + if (promise_hooks_.size() == 1) { isolate_->SetPromiseHook(EnvPromiseHook); } } +bool Environment::RemovePromiseHook(promise_hook_func fn, void* arg) { + auto it = std::find_if( + promise_hooks_.begin(), promise_hooks_.end(), + [&](const PromiseHookCallback& hook) { + return hook.cb_ == fn && hook.arg_ == arg; + }); + + if (it == promise_hooks_.end()) return false; + + promise_hooks_.erase(it); + if (promise_hooks_.empty()) { + isolate_->SetPromiseHook(nullptr); + } + + return true; +} + void Environment::EnvPromiseHook(v8::PromiseHookType type, v8::Local promise, v8::Local parent) { diff --git a/src/env.h b/src/env.h index 2eacf68ef5e93b..2f62663be6ab33 100644 --- a/src/env.h +++ b/src/env.h @@ -654,6 +654,7 @@ class Environment { static const int kContextEmbedderDataIndex = NODE_CONTEXT_EMBEDDER_DATA_INDEX; void AddPromiseHook(promise_hook_func fn, void* arg); + bool RemovePromiseHook(promise_hook_func fn, void* arg); private: inline void ThrowError(v8::Local (*fun)(v8::Local),