trace_events: add trace category enabled tracking

Track state of async_hooks trace event category enablement.
Enable/disable the async_hooks trace event dynamically.

PR-URL: https://github.com/nodejs/node/pull/22128
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: Andreas Madsen <amwebdk@gmail.com>

6 files changed, 81 insertions, 0 deletions

diff --git a/src/bootstrapper.cc b/src/bootstrapper.cc
index 94bb941377..720f89c011 100644
--- a/src/bootstrapper.cc
+++ b/src/bootstrapper.cc
@@ -31,6 +31,12 @@ void RunMicrotasks(const FunctionCallbackInfo<Value>& args) {
   args.GetIsolate()->RunMicrotasks();
 }
 
+void SetupTraceCategoryState(const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+  CHECK(args[0]->IsFunction());
+  env->set_trace_category_state_function(args[0].As<Function>());
+}
+
 void SetupNextTick(const FunctionCallbackInfo<Value>& args) {
   Environment* env = Environment::GetCurrent(args);
   Isolate* isolate = env->isolate();
@@ -117,6 +123,7 @@ void SetupPromises(const FunctionCallbackInfo<Value>& args) {
 // completes so that it can be gc'd as soon as possible.
 void SetupBootstrapObject(Environment* env,
                           Local<Object> bootstrapper) {
+  BOOTSTRAP_METHOD(_setupTraceCategoryState, SetupTraceCategoryState);
   BOOTSTRAP_METHOD(_setupProcessObject, SetupProcessObject);
   BOOTSTRAP_METHOD(_setupNextTick, SetupNextTick);
   BOOTSTRAP_METHOD(_setupPromises, SetupPromises);
diff --git a/src/env-inl.h b/src/env-inl.h
index d65d7e7562..3bca20c81c 100644
--- a/src/env-inl.h
+++ b/src/env-inl.h
@@ -432,6 +432,11 @@ Environment::should_abort_on_uncaught_toggle() {
   return should_abort_on_uncaught_toggle_;
 }
 
+inline AliasedBuffer<uint8_t, v8::Uint8Array>&
+Environment::trace_category_state() {
+  return trace_category_state_;
+}
+
 Environment::ShouldNotAbortOnUncaughtScope::ShouldNotAbortOnUncaughtScope(
     Environment* env)
     : env_(env) {
diff --git a/src/env.cc b/src/env.cc
index f691803039..97405fe4b1 100644
--- a/src/env.cc
+++ b/src/env.cc
@@ -31,6 +31,8 @@ using v8::TryCatch;
 using v8::Value;
 using worker::Worker;
 
+#define kTraceCategoryCount 1
+
 int const Environment::kNodeContextTag = 0x6e6f64;
 void* Environment::kNodeContextTagPtr = const_cast<void*>(
     static_cast<const void*>(&Environment::kNodeContextTag));
@@ -103,6 +105,21 @@ void InitThreadLocalOnce() {
   CHECK_EQ(0, uv_key_create(&Environment::thread_local_env));
 }
 
+void Environment::TrackingTraceStateObserver::UpdateTraceCategoryState() {
+  env_->trace_category_state()[0] =
+      *TRACE_EVENT_API_GET_CATEGORY_GROUP_ENABLED(
+          TRACING_CATEGORY_NODE1(async_hooks));
+
+  Isolate* isolate = env_->isolate();
+  Local<Function> cb = env_->trace_category_state_function();
+  if (cb.IsEmpty())
+    return;
+  TryCatch try_catch(isolate);
+  try_catch.SetVerbose(true);
+  cb->Call(env_->context(), v8::Undefined(isolate),
+           0, nullptr).ToLocalChecked();
+}
+
 Environment::Environment(IsolateData* isolate_data,
                          Local<Context> context,
                          tracing::AgentWriterHandle* tracing_agent_writer)
@@ -118,6 +135,7 @@ Environment::Environment(IsolateData* isolate_data,
       emit_env_nonstring_warning_(true),
       makecallback_cntr_(0),
       should_abort_on_uncaught_toggle_(isolate_, 1),
+      trace_category_state_(isolate_, kTraceCategoryCount),
 #if HAVE_INSPECTOR
       inspector_agent_(new inspector::Agent(this)),
 #endif
@@ -132,6 +150,14 @@ Environment::Environment(IsolateData* isolate_data,
 
   AssignToContext(context, ContextInfo(""));
 
+  if (tracing_agent_writer_ != nullptr) {
+    trace_state_observer_.reset(new TrackingTraceStateObserver(this));
+    v8::TracingController* tracing_controller =
+        tracing_agent_writer_->GetTracingController();
+    if (tracing_controller != nullptr)
+      tracing_controller->AddTraceStateObserver(trace_state_observer_.get());
+  }
+
   destroy_async_id_list_.reserve(512);
   performance_state_.reset(new performance::performance_state(isolate()));
   performance_state_->Mark(
@@ -173,6 +199,13 @@ Environment::~Environment() {
   context()->SetAlignedPointerInEmbedderData(
       ContextEmbedderIndex::kEnvironment, nullptr);
 
+  if (tracing_agent_writer_ != nullptr) {
+    v8::TracingController* tracing_controller =
+        tracing_agent_writer_->GetTracingController();
+    if (tracing_controller != nullptr)
+      tracing_controller->RemoveTraceStateObserver(trace_state_observer_.get());
+  }
+
   delete[] heap_statistics_buffer_;
   delete[] heap_space_statistics_buffer_;
   delete[] http_parser_buffer_;
diff --git a/src/env.h b/src/env.h
index 7a4804cdd4..28aa936c24 100644
--- a/src/env.h
+++ b/src/env.h
@@ -344,6 +344,7 @@ struct PackageConfig {
   V(tick_callback_function, v8::Function)                                     \
   V(timers_callback_function, v8::Function)                                   \
   V(tls_wrap_constructor_function, v8::Function)                              \
+  V(trace_category_state_function, v8::Function)                              \
   V(tty_constructor_template, v8::FunctionTemplate)                           \
   V(udp_constructor_function, v8::Function)                                   \
   V(url_constructor_function, v8::Function)                                   \
@@ -649,6 +650,8 @@ class Environment {
   inline AliasedBuffer<uint32_t, v8::Uint32Array>&
   should_abort_on_uncaught_toggle();
 
+  inline AliasedBuffer<uint8_t, v8::Uint8Array>& trace_category_state();
+
   // The necessary API for async_hooks.
   inline double new_async_id();
   inline double execution_async_id();
@@ -830,6 +833,25 @@ class Environment {
   // This needs to be available for the JS-land setImmediate().
   void ToggleImmediateRef(bool ref);
 
+  class TrackingTraceStateObserver :
+      public v8::TracingController::TraceStateObserver {
+   public:
+    explicit TrackingTraceStateObserver(Environment* env) : env_(env) {}
+
+    void OnTraceEnabled() override {
+      UpdateTraceCategoryState();
+    }
+
+    void OnTraceDisabled() override {
+      UpdateTraceCategoryState();
+    }
+
+   private:
+    void UpdateTraceCategoryState();
+
+    Environment* env_;
+  };
+
   class ShouldNotAbortOnUncaughtScope {
    public:
     explicit inline ShouldNotAbortOnUncaughtScope(Environment* env);
@@ -889,6 +911,10 @@ class Environment {
   AliasedBuffer<uint32_t, v8::Uint32Array> should_abort_on_uncaught_toggle_;
   int should_not_abort_scope_counter_ = 0;
 
+  // Attached to a Uint8Array that tracks the state of trace category
+  AliasedBuffer<uint8_t, v8::Uint8Array> trace_category_state_;
+  std::unique_ptr<TrackingTraceStateObserver> trace_state_observer_;
+
   std::unique_ptr<performance::performance_state> performance_state_;
   std::unordered_map<std::string, uint64_t> performance_marks_;
 
diff --git a/src/node_trace_events.cc b/src/node_trace_events.cc
index 219e142a31..6e70cfce7b 100644
--- a/src/node_trace_events.cc
+++ b/src/node_trace_events.cc
@@ -127,6 +127,10 @@ void Initialize(Local<Object> target,
                   .FromJust();
   target->Set(context, trace,
               binding->Get(context, trace).ToLocalChecked()).FromJust();
+
+  target->Set(context,
+              FIXED_ONE_BYTE_STRING(env->isolate(), "traceCategoryState"),
+              env->trace_category_state().GetJSArray()).FromJust();
 }
 
 }  // namespace node
diff --git a/src/tracing/agent.h b/src/tracing/agent.h
index 045aaef85e..e79480a036 100644
--- a/src/tracing/agent.h
+++ b/src/tracing/agent.h
@@ -51,6 +51,8 @@ class AgentWriterHandle {
 
   inline Agent* agent() { return agent_; }
 
+  inline v8::TracingController* GetTracingController();
+
  private:
   inline AgentWriterHandle(Agent* agent, int id) : agent_(agent), id_(id) {}
 
@@ -155,6 +157,10 @@ void AgentWriterHandle::Disable(const std::set<std::string>& categories) {
   if (agent_ != nullptr) agent_->Disable(id_, categories);
 }
 
+inline v8::TracingController* AgentWriterHandle::GetTracingController() {
+  return agent_ != nullptr ? agent_->GetTracingController() : nullptr;
+}
+
 }  // namespace tracing
 }  // namespace node