Implement runWithContext

bindings/profilers/wall.cc

@@ -1089,6 +1089,7 @@ NAN_MODULE_INIT(WallProfiler::Init) {
   Nan::SetPrototypeMethod(tpl,
                           "v8ProfilerStuckEventLoopDetected",
                           V8ProfilerStuckEventLoopDetected);
+  Nan::SetPrototypeMethod(tpl, "createContextHolder", CreateContextHolder);
 
   Nan::SetAccessor(tpl->InstanceTemplate(),
                    Nan::New("state").ToLocalChecked(),
@@ -1176,23 +1177,29 @@ void WallProfiler::SetContext(Isolate* isolate, Local<Value> value) {
     SignalGuard m(setInProgress_);
     cpedMap->Delete(v8Ctx, localKey).Check();
   } else {
-    auto wrap =
-        wrapObjectTemplate_.Get(isolate)->NewInstance(v8Ctx).ToLocalChecked();
-    // for easy access from JS when cpedKey is an ALS, it can do
-    // als.getStore()?.[0];
-    wrap->Set(v8Ctx, 0, value).Check();
-    auto contextPtr = new PersistentContextPtr(&liveContextPtrs_, wrap);
-    liveContextPtrs_.insert(contextPtr);
-    contextPtr->Set(isolate, value);
-
+    auto contextHolder = CreateContextHolder(isolate, v8Ctx, value);
     SignalGuard m(setInProgress_);
-    cpedMap->Set(v8Ctx, localKey, wrap).ToLocalChecked();
+    cpedMap->Set(v8Ctx, localKey, contextHolder).ToLocalChecked();
   }
 #else
   SetCurrentContextPtr(isolate, value);
 #endif
 }
 
+Local<Object> WallProfiler::CreateContextHolder(Isolate* isolate,
+                                                Local<Context> v8Ctx,
+                                                Local<Value> value) {
+  auto wrap =
+      wrapObjectTemplate_.Get(isolate)->NewInstance(v8Ctx).ToLocalChecked();
+  // for easy access from JS when cpedKey is an ALS, it can do
+  // als.getStore()?.[0];
+  wrap->Set(v8Ctx, 0, value).Check();
+  auto contextPtr = new PersistentContextPtr(&liveContextPtrs_, wrap);
+  liveContextPtrs_.insert(contextPtr);
+  contextPtr->Set(isolate, value);
+  return wrap;
+}
+
 ContextPtr WallProfiler::GetContextPtrSignalSafe(Isolate* isolate) {
   auto isSetInProgress = setInProgress_.load(std::memory_order_relaxed);
   std::atomic_signal_fence(std::memory_order_acquire);
@@ -1279,6 +1286,18 @@ NAN_SETTER(WallProfiler::SetContext) {
   profiler->SetContext(info.GetIsolate(), value);
 }
 
+NAN_METHOD(WallProfiler::CreateContextHolder) {
+  auto profiler = Nan::ObjectWrap::Unwrap<WallProfiler>(info.This());
+  if (!profiler->useCPED()) {
+    return Nan::ThrowTypeError(
+        "CreateContextHolder can only be used with CPED");
+  }
+  auto isolate = info.GetIsolate();
+  auto contextHolder = profiler->CreateContextHolder(
+      isolate, isolate->GetCurrentContext(), info[0]);
+  info.GetReturnValue().Set(contextHolder);
+}
+
 NAN_GETTER(WallProfiler::SharedArrayGetter) {
   auto profiler = Nan::ObjectWrap::Unwrap<WallProfiler>(info.This());
   info.GetReturnValue().Set(profiler->jsArray_.Get(v8::Isolate::GetCurrent()));

bindings/profilers/wall.hh

@@ -143,6 +143,10 @@ class WallProfiler : public Nan::ObjectWrap {
 
   v8::Local<v8::Value> GetContext(v8::Isolate*);
   void SetContext(v8::Isolate*, v8::Local<v8::Value>);
+  v8::Local<v8::Object> CreateContextHolder(v8::Isolate*,
+                                            v8::Local<v8::Context>,
+                                            v8::Local<v8::Value>);
+
   void PushContext(int64_t time_from,
                    int64_t time_to,
                    int64_t cpu_time,
@@ -186,6 +190,7 @@ class WallProfiler : public Nan::ObjectWrap {
   static NAN_MODULE_INIT(Init);
   static NAN_GETTER(GetContext);
   static NAN_SETTER(SetContext);
+  static NAN_METHOD(CreateContextHolder);
   static NAN_GETTER(SharedArrayGetter);
   static NAN_GETTER(GetMetrics);
 };

ts/src/index.ts

@@ -36,6 +36,7 @@ export const time = {
   stop: timeProfiler.stop,
   getContext: timeProfiler.getContext,
   setContext: timeProfiler.setContext,
+  runWithContext: timeProfiler.runWithContext,
   isStarted: timeProfiler.isStarted,
   v8ProfilerStuckEventLoopDetected:
     timeProfiler.v8ProfilerStuckEventLoopDetected,

ts/src/time-profiler.ts

@@ -39,7 +39,7 @@
 type Milliseconds = number;
 
 let gProfiler: InstanceType<typeof TimeProfiler> | undefined;
 let gStore: AsyncLocalStorage<any> | undefined;
 let gSourceMapper: SourceMapper | undefined;
 let gIntervalMicros: Microseconds;
 let gV8ProfilerStuckEventLoopDetected = 0;
@@ -169,6 +169,19 @@
   gProfiler.context = context;
 }
 
+export function runWithContext<R, TArgs extends any[]>(
+  context: object,
+  f: (...args: TArgs) => R,
+  ...args: TArgs
+): R {
+  if (!gProfiler) {
+    throw new Error('Wall profiler is not started');
+  } else if (!gStore) {
+    throw new Error('Can only use runWithContext with AsyncContextFrame');
+  }
+  return gStore.run(gProfiler.createContextHolder(context), f, ...args);
+}
+
 export function getContext() {
   if (!gProfiler) {
     throw new Error('Wall profiler is not started');

ts/test/test-get-value-from-map-profiler.ts

@@ -124,6 +124,27 @@
 
         profiler.dispose();
       });
+
+      it('should work with createContextHolder pattern', () => {
+        // This tests the pattern used by runWithContext where
+        // createContextHolder creates a wrap object that's stored in CPED map
+
+        const als = new AsyncLocalStorage();
+        const profiler = createProfiler(als);
+
+        const context = {label: 'wrapped-context', id: 999};
+
+        // Using als.run mimics what runWithContext does internally
+        als.run(profiler.createContextHolder(context), () => {
+          const retrieved = profiler.context;
+
+          // The wrap object stores context at index 0
+          assert.ok(retrieved !== null && typeof retrieved === 'object');
+          assert.deepStrictEqual(retrieved, context);
+        });
+
+        profiler.dispose();
+      });
     });
 
     describe('multiple context frames', () => {
@@ -179,7 +200,7 @@
   });
 }
 
 function createProfiler(als: AsyncLocalStorage<any>) {
   return new profiler.TimeProfiler({
     intervalMicros: 10000,
     durationMillis: 500,