Allow use of CPED to store sampling context

Author: szegedi

bindings/profilers/wall.cc

@@ -58,6 +58,8 @@ using namespace v8;
 
 namespace dd {
 
+using ContextPtr = std::shared_ptr<Global<Value>>;
+
 // Maximum number of rounds in the GetV8ToEpochOffset
 static constexpr int MAX_EPOCH_OFFSET_ATTEMPTS = 20;
 
@@ -318,8 +320,7 @@ void SignalHandler::HandleProfilerSignal(int sig,
   auto time_from = Now();
   old_handler(sig, info, context);
   auto time_to = Now();
-  auto async_id = prof->GetAsyncId(isolate);
-  prof->PushContext(time_from, time_to, cpu_time, async_id);
+  prof->PushContext(time_from, time_to, cpu_time, isolate);
 }
 #else
 class SignalHandler {
@@ -513,8 +514,10 @@ WallProfiler::WallProfiler(std::chrono::microseconds samplingPeriod,
                            bool withContexts,
                            bool workaroundV8Bug,
                            bool collectCpuTime,
-                           bool isMainThread)
+                           bool isMainThread,
+                           bool useCPED)
     : samplingPeriod_(samplingPeriod),
+      useCPED_(useCPED),
       includeLines_(includeLines),
       withContexts_(withContexts),
       isMainThread_(isMainThread) {
@@ -529,7 +532,6 @@ WallProfiler::WallProfiler(std::chrono::microseconds samplingPeriod,
     contexts_.reserve(duration * 2 / samplingPeriod);
   }
 
-  curContext_.store(&context1_, std::memory_order_relaxed);
   collectionMode_.store(CollectionMode::kNoCollect, std::memory_order_relaxed);
 
   // TODO: bind to this isolate? Would fix the Dispose(nullptr) issue.
@@ -543,6 +545,9 @@ WallProfiler::WallProfiler(std::chrono::microseconds samplingPeriod,
   jsArray_ = v8::Global<v8::Uint32Array>(isolate, jsArray);
   std::fill(fields_, fields_ + kFieldCount, 0);
 
+  if (useCPED_) {
+    cpedSymbol_ = v8::Global<v8::Symbol>(isolate, v8::Symbol::New(isolate));
+  }
   gcCount = 0;
   isolate->AddGCPrologueCallback(&GCPrologueCallback, this);
   isolate->AddGCEpilogueCallback(&GCEpilogueCallback, this);
@@ -676,13 +681,21 @@ NAN_METHOD(WallProfiler::New) {
           "Include line option is not compatible with contexts.");
     }
 
+    auto useCPEDValue =
+        Nan::Get(arg, Nan::New<v8::String>("useCPED").ToLocalChecked());
+    if (useCPEDValue.IsEmpty() || !useCPEDValue.ToLocalChecked()->IsBoolean()) {
+      return Nan::ThrowTypeError("useCPED must be a boolean.");
+    }
+    bool useCPED = useCPEDValue.ToLocalChecked().As<v8::Boolean>()->Value();
+
     WallProfiler* obj = new WallProfiler(interval,
                                          duration,
                                          lineNumbers,
                                          withContexts,
                                          workaroundV8Bug,
                                          collectCpuTime,
-                                         isMainThread);
+                                         isMainThread,
+                                         useCPED);
     obj->Wrap(info.This());
     info.GetReturnValue().Set(info.This());
   } else {
@@ -995,28 +1008,109 @@ v8::CpuProfiler* WallProfiler::CreateV8CpuProfiler() {
 }
 
 v8::Local<v8::Value> WallProfiler::GetContext(Isolate* isolate) {
-  auto context = *curContext_.load(std::memory_order_relaxed);
+  auto context = GetContextPtr(isolate);
   if (!context) return v8::Undefined(isolate);
   return context->Get(isolate);
 }
 
+class PersistentContextPtr : AtomicContextPtr {
+  Persistent<Object> per;
+
+  void BindLifecycleTo(Isolate* isolate, Local<Object>& obj) {
+    // Register a callback to delete this object when the object is GCed
+    per.Reset(isolate, obj);
+    per.SetWeak(
+        this,
+        [](const WeakCallbackInfo<PersistentContextPtr>& data) {
+          auto& per = data.GetParameter()->per;
+          if (!per.IsEmpty()) {
+            per.ClearWeak();
+            per.Reset();
+          }
+          // Using SetSecondPassCallback as shared_ptr can trigger ~Global and
+          // any V8 API use needs to be in the second pass
+          data.SetSecondPassCallback(
+              [](const WeakCallbackInfo<PersistentContextPtr>& data) {
+                delete data.GetParameter();
+              });
+        },
+        WeakCallbackType::kParameter);
+  }
+
+  friend class WallProfiler;
+};
+
 void WallProfiler::SetContext(Isolate* isolate, Local<Value> value) {
-  // Need to be careful here, because we might be interrupted by a
-  // signal handler that will make use of curContext_.
-  // Update of shared_ptr is not atomic, so instead we use a pointer
-  // (curContext_) that points on two shared_ptr (context1_ and context2_),
-  // update the shared_ptr that is not currently in use and then atomically
-  // update curContext_.
-  auto newCurContext = curContext_.load(std::memory_order_relaxed) == &context1_
-                           ? &context2_
-                           : &context1_;
-  if (!value->IsNullOrUndefined()) {
-    *newCurContext = std::make_shared<Global<Value>>(isolate, value);
+  if (!useCPED_) {
+    curContext_.Set(isolate, value);
+    return;
+  }
+
+  auto cped = isolate->GetContinuationPreservedEmbedderData();
+  // No Node AsyncContextFrame in this continuation yet
+  if (!cped->IsObject()) return;
+
+  auto cpedObj = cped.As<Object>();
+  auto localSymbol = cpedSymbol_.Get(isolate);
+  auto v8Ctx = isolate->GetCurrentContext();
+  auto maybeProfData = cpedObj->Get(v8Ctx, localSymbol);
+  if (maybeProfData.IsEmpty()) return;
+  auto profData = maybeProfData.ToLocalChecked();
+
+  PersistentContextPtr* contextPtr = nullptr;
+  if (profData->IsUndefined()) {
+    contextPtr = new PersistentContextPtr();
+
+    auto maybeSetResult =
+        cpedObj->Set(v8Ctx, localSymbol, External::New(isolate, contextPtr));
+    if (maybeSetResult.IsNothing()) {
+      delete contextPtr;
+      return;
+    }
+    contextPtr->BindLifecycleTo(isolate, cpedObj);
   } else {
-    newCurContext->reset();
+    contextPtr =
+        static_cast<PersistentContextPtr*>(profData.As<External>()->Value());
   }
-  std::atomic_signal_fence(std::memory_order_release);
-  curContext_.store(newCurContext, std::memory_order_relaxed);
+
+  contextPtr->Set(isolate, value);
+}
+
+ContextPtr WallProfiler::GetContextPtrSignalSafe(Isolate* isolate) {
+  if (!useCPED_) {
+    // Not strictly necessary but we can avoid HandleScope creation for this
+    // case.
+    return curContext_.Get();
+  }
+
+  if (gcCount > 0) {
+    return gcContext;
+  } else if (isolate->InContext()) {
+    auto handleScope = HandleScope(isolate);
+    return GetContextPtr(isolate);
+  }
+  // not in a V8 Context
+  return std::shared_ptr<Global<Value>>();
+}
+
+ContextPtr WallProfiler::GetContextPtr(Isolate* isolate) {
+  if (!useCPED_) {
+    return curContext_.Get();
+  }
+
+  auto cped = isolate->GetContinuationPreservedEmbedderData();
+  if (!cped->IsObject()) return std::shared_ptr<Global<Value>>();
+
+  auto cpedObj = cped.As<Object>();
+  auto localSymbol = cpedSymbol_.Get(isolate);
+  auto maybeProfData = cpedObj->Get(isolate->GetCurrentContext(), localSymbol);
+  if (maybeProfData.IsEmpty()) return std::shared_ptr<Global<Value>>();
+  auto profData = maybeProfData.ToLocalChecked();
+
+  if (profData->IsUndefined()) return std::shared_ptr<Global<Value>>();
+
+  return static_cast<PersistentContextPtr*>(profData.As<External>()->Value())
+      ->Get();
 }
 
 NAN_GETTER(WallProfiler::GetContext) {
@@ -1047,14 +1141,16 @@ NAN_METHOD(WallProfiler::Dispose) {
 void WallProfiler::PushContext(int64_t time_from,
                                int64_t time_to,
                                int64_t cpu_time,
-                               int64_t async_id) {
+                               Isolate* isolate) {
   // Be careful this is called in a signal handler context therefore all
   // operations must be async signal safe (in particular no allocations).
   // Our ring buffer avoids allocations.
-  auto context = curContext_.load(std::memory_order_relaxed);
-  std::atomic_signal_fence(std::memory_order_acquire);
   if (contexts_.size() < contexts_.capacity()) {
-    contexts_.push_back({*context, time_from, time_to, cpu_time, async_id});
+    contexts_.push_back({GetContextPtrSignalSafe(isolate),
+                         time_from,
+                         time_to,
+                         cpu_time,
+                         GetAsyncId(isolate)});
     std::atomic_fetch_add_explicit(
         reinterpret_cast<std::atomic<uint32_t>*>(&fields_[kSampleCount]),
         1U,

bindings/profilers/wall.hh

@@ -37,29 +37,53 @@ struct Result {
   std::string msg;
 };
 
+using ContextPtr = std::shared_ptr<v8::Global<v8::Value>>;
+
+class AtomicContextPtr {
+  ContextPtr ptr1;
+  ContextPtr ptr2;
+  std::atomic<ContextPtr*> currentPtr = &ptr1;
+
+  void Set(v8::Isolate* isolate, v8::Local<v8::Value> value) {
+    auto newPtr =
+        currentPtr.load(std::memory_order_relaxed) == &ptr1 ? &ptr2 : &ptr1;
+    if (!value->IsNullOrUndefined()) {
+      *newPtr = std::make_shared<v8::Global<v8::Value>>(isolate, value);
+    } else {
+      newPtr->reset();
+    }
+    std::atomic_signal_fence(std::memory_order_release);
+    currentPtr.store(newPtr, std::memory_order_relaxed);
+  }
+
+  ContextPtr Get() {
+    auto ptr = currentPtr.load(std::memory_order_relaxed);
+    std::atomic_signal_fence(std::memory_order_acquire);
+    return ptr ? *ptr : std::shared_ptr<v8::Global<v8::Value>>();
+  }
+
+  friend class WallProfiler;
+};
+
 class WallProfiler : public Nan::ObjectWrap {
  public:
   enum class CollectionMode { kNoCollect, kPassThrough, kCollectContexts };
 
  private:
   enum Fields { kSampleCount, kFieldCount };
 
-  using ContextPtr = std::shared_ptr<v8::Global<v8::Value>>;
-
   std::chrono::microseconds samplingPeriod_{0};
   v8::CpuProfiler* cpuProfiler_ = nullptr;
-  // TODO: Investigate use of v8::Persistent instead of shared_ptr<Global> to
-  // avoid heap allocation. Need to figure out the right move/copy semantics in
-  // and out of the ring buffer.
 
-  // We're using a pair of shared pointers and an atomic pointer-to-current as
-  // a way to ensure signal safety on update.
-  ContextPtr context1_;
-  ContextPtr context2_;
-  std::atomic<ContextPtr*> curContext_;
+  bool useCPED_ = false;
+  // If we aren't using the CPED, we use a single context ptr stored here.
+  AtomicContextPtr curContext_;
+  // Otherwise we'll use a private symbol to store the context in CPED objects.
+  v8::Global<v8::Symbol> cpedSymbol_;
 
   std::atomic<int> gcCount = 0;
   int64_t gcAsyncId;
+  ContextPtr gcContext;
 
   std::atomic<CollectionMode> collectionMode_;
   std::atomic<uint64_t> noCollectCallCount_;
@@ -105,6 +129,8 @@ class WallProfiler : public Nan::ObjectWrap {
       GENERAL_REGS_ONLY;
 
   bool waitForSignal(uint64_t targetCallCount = 0);
+  ContextPtr GetContextPtr(v8::Isolate* isolate);
+  ContextPtr GetContextPtrSignalSafe(v8::Isolate* isolate);
 
  public:
   /**
@@ -113,21 +139,26 @@ class WallProfiler : public Nan::ObjectWrap {
    * parameter is informative; it is up to the caller to call the Stop method
    * every period. The parameter is used to preallocate data structures that
    * should not be reallocated in async signal safe code.
+   * @param useCPED whether to use the V8 ContinuationPreservingEmbedderData
+   * to store the current sampling context. It can be used if AsyncLocalStorage
+   * uses the AsyncContextFrame implementation (experimental in Node 23, default
+   * in Node 24.)
    */
   explicit WallProfiler(std::chrono::microseconds samplingPeriod,
                         std::chrono::microseconds duration,
                         bool includeLines,
                         bool withContexts,
                         bool workaroundV8bug,
                         bool collectCpuTime,
-                        bool isMainThread);
+                        bool isMainThread,
+                        bool useCPED);
 
   v8::Local<v8::Value> GetContext(v8::Isolate*);
   void SetContext(v8::Isolate*, v8::Local<v8::Value>);
   void PushContext(int64_t time_from,
                    int64_t time_to,
                    int64_t cpu_time,
-                   int64_t async_id);
+                   v8::Isolate* isolate);
   Result StartImpl();
   std::string StartInternal();
   Result StopImpl(bool restart,
@@ -164,6 +195,9 @@ class WallProfiler : public Nan::ObjectWrap {
   void OnGCStart(v8::Isolate* isolate) {
     if (gcCount == 0) {
       gcAsyncId = GetAsyncId(isolate);
+      if (useCPED_) {
+        gcContext = GetContextPtrSignalSafe(v8::Isolate::GetCurrent());
+      }
       gcCount = 1;
     } else {
       ++gcCount;
@@ -173,6 +207,9 @@ class WallProfiler : public Nan::ObjectWrap {
   void OnGCEnd() {
     if (--gcCount == 0) {
       gcAsyncId = -1;
+      if (useCPED_) {
+        gcContext.reset();
+      }
     }
   }
 

ts/src/time-profiler.ts

@@ -91,7 +91,7 @@ export function start(options: TimeProfilerOptions = {}) {
     throw new Error('Wall profiler is already started');
   }
 
-  gProfiler = new TimeProfiler({...options, isMainThread});
+  gProfiler = new TimeProfiler({...options, isMainThread, useCPED: false});
   gSourceMapper = options.sourceMapper;
   gIntervalMicros = options.intervalMicros!;
   gV8ProfilerStuckEventLoopDetected = 0;