Support promises in profile method in bun:jsc (#12165)

packages/bun-types/jsc.d.ts

@@ -78,21 +78,7 @@ declare module "bun:jsc" {
    */
   function setTimeZone(timeZone: string): string;
 
-  /**
-   * Run JavaScriptCore's sampling profiler for a particular function
-   *
-   * This is pretty low-level.
-   *
-   * Things to know:
-   * - LLint means "Low Level Interpreter", which is the interpreter that runs before any JIT compilation
-   * - Baseline is the first JIT compilation tier. It's the least optimized, but the fastest to compile
-   * - DFG means "Data Flow Graph", which is the second JIT compilation tier. It has some optimizations, but is slower to compile
-   * - FTL means "Faster Than Light", which is the third JIT compilation tier. It has the most optimizations, but is the slowest to compile
-   */
-  function profile(
-    callback: CallableFunction,
-    sampleInterval?: number,
-  ): {
+  interface SamplingProfile {
     /**
      * A formatted summary of the top functions
      *
@@ -183,7 +169,24 @@ declare module "bun:jsc" {
      * Stack traces of the top functions
      */
     stackTraces: string[];
-  };
+  }
+
+  /**
+   * Run JavaScriptCore's sampling profiler for a particular function
+   *
+   * This is pretty low-level.
+   *
+   * Things to know:
+   * - LLint means "Low Level Interpreter", which is the interpreter that runs before any JIT compilation
+   * - Baseline is the first JIT compilation tier. It's the least optimized, but the fastest to compile
+   * - DFG means "Data Flow Graph", which is the second JIT compilation tier. It has some optimizations, but is slower to compile
+   * - FTL means "Faster Than Light", which is the third JIT compilation tier. It has the most optimizations, but is the slowest to compile
+   */
+  function profile<T extends (...args: any[]) => any>(
+    callback: T,
+    sampleInterval?: number,
+    ...args: Parameters<T>
+  ): ReturnType<T> extends Promise<infer U> ? Promise<SamplingProfile> : SamplingProfile;
 
   /**
    * This returns objects which native code has explicitly protected from being

src/bun.js/modules/BunJSCModule.h

@@ -1,6 +1,11 @@
 #include "_NativeModule.h"
 
 #include "ExceptionOr.h"
+#include "JavaScriptCore/ArgList.h"
+#include "JavaScriptCore/ExceptionScope.h"
+#include "JavaScriptCore/JSCJSValue.h"
+#include "JavaScriptCore/JSGlobalObject.h"
+#include "JavaScriptCore/JSNativeStdFunction.h"
 #include "MessagePort.h"
 #include "SerializedScriptValue.h"
 #include <JavaScriptCore/APICast.h>
@@ -18,6 +23,7 @@
 #include <JavaScriptCore/JSBasePrivate.h>
 #include <JavaScriptCore/JSCInlines.h>
 #include <JavaScriptCore/JSONObject.h>
+#include <JavaScriptCore/JSPromise.h>
 #include <JavaScriptCore/JavaScript.h>
 #include <JavaScriptCore/ObjectConstructor.h>
 #include <JavaScriptCore/SamplingProfiler.h>
@@ -486,6 +492,18 @@ JSC_DEFINE_HOST_FUNCTION(functionRunProfiler, (JSGlobalObject * globalObject,
       vm.ensureSamplingProfiler(WTF::Stopwatch::create());
 
   JSC::JSValue callbackValue = callFrame->argument(0);
+  JSC::JSValue sampleValue = callFrame->argument(1);
+
+  MarkedArgumentBuffer args;
+
+  if (callFrame->argumentCount() > 2) {
+    size_t count = callFrame->argumentCount();
+    args.ensureCapacity(count - 2);
+    for (size_t i = 2; i < count; i++) {
+      args.append(callFrame->argument(i));
+    }
+  }
+
   auto throwScope = DECLARE_THROW_SCOPE(vm);
   if (callbackValue.isUndefinedOrNull() || !callbackValue.isCallable()) {
     throwException(
@@ -496,48 +514,87 @@ JSC_DEFINE_HOST_FUNCTION(functionRunProfiler, (JSGlobalObject * globalObject,
 
   JSC::JSFunction *function = jsCast<JSC::JSFunction *>(callbackValue);
 
-  JSC::JSValue sampleValue = callFrame->argument(1);
   if (sampleValue.isNumber()) {
     unsigned sampleInterval = sampleValue.toUInt32(globalObject);
     samplingProfiler.setTimingInterval(
         Seconds::fromMicroseconds(sampleInterval));
   }
 
+  const auto report = [](JSC::VM &vm,
+                         JSC::JSGlobalObject *globalObject) -> JSC::JSValue {
+    auto throwScope = DECLARE_THROW_SCOPE(vm);
+
+    auto &samplingProfiler = *vm.samplingProfiler();
+    StringPrintStream topFunctions;
+    samplingProfiler.reportTopFunctions(topFunctions);
+
+    StringPrintStream byteCodes;
+    samplingProfiler.reportTopBytecodes(byteCodes);
+
+    JSValue stackTraces = JSONParse(
+        globalObject, samplingProfiler.stackTracesAsJSON()->toJSONString());
+
+    samplingProfiler.shutdown();
+    RETURN_IF_EXCEPTION(throwScope, {});
+
+    JSObject *result =
+        constructEmptyObject(globalObject, globalObject->objectPrototype(), 3);
+    result->putDirect(vm, Identifier::fromString(vm, "functions"_s),
+                      jsString(vm, topFunctions.toString()));
+    result->putDirect(vm, Identifier::fromString(vm, "bytecodes"_s),
+                      jsString(vm, byteCodes.toString()));
+    result->putDirect(vm, Identifier::fromString(vm, "stackTraces"_s),
+                      stackTraces);
+
+    return result;
+  };
+  const auto reportFailure = [](JSC::VM &vm) -> JSC::JSValue {
+    if (auto *samplingProfiler = vm.samplingProfiler()) {
+      samplingProfiler->pause();
+      samplingProfiler->shutdown();
+      samplingProfiler->clearData();
+    }
+
+    return {};
+  };
+
   JSC::CallData callData = JSC::getCallData(function);
-  MarkedArgumentBuffer args;
 
   samplingProfiler.noticeCurrentThreadAsJSCExecutionThread();
   samplingProfiler.start();
-  JSC::call(globalObject, function, callData, JSC::jsUndefined(), args);
-  samplingProfiler.pause();
-  if (throwScope.exception()) {
-    samplingProfiler.shutdown();
-    samplingProfiler.clearData();
-    return JSValue::encode(JSValue{});
+  JSValue returnValue =
+      JSC::call(globalObject, function, callData, JSC::jsUndefined(), args);
+
+  if (returnValue.isEmpty() || throwScope.exception()) {
+    return JSValue::encode(reportFailure(vm));
   }
 
-  StringPrintStream topFunctions;
-  samplingProfiler.reportTopFunctions(topFunctions);
+  if (auto *promise = jsDynamicCast<JSPromise *>(returnValue)) {
+    auto afterOngoingPromiseCapability =
+        JSC::JSPromise::create(vm, globalObject->promiseStructure());
+    RETURN_IF_EXCEPTION(throwScope, {});
 
-  StringPrintStream byteCodes;
-  samplingProfiler.reportTopBytecodes(byteCodes);
+    JSNativeStdFunction *resolve = JSNativeStdFunction::create(
+        vm, globalObject, 0, "resolve"_s,
+        [report](JSGlobalObject *globalObject, CallFrame *callFrame) {
+          return JSValue::encode(JSPromise::resolvedPromise(
+              globalObject, report(globalObject->vm(), globalObject)));
+        });
+    JSNativeStdFunction *reject = JSNativeStdFunction::create(
+        vm, globalObject, 0, "reject"_s,
+        [reportFailure](JSGlobalObject *globalObject, CallFrame *callFrame) {
+          EnsureStillAliveScope error = callFrame->argument(0);
+          auto scope = DECLARE_THROW_SCOPE(globalObject->vm());
+          reportFailure(globalObject->vm());
+          throwException(globalObject, scope, error.value());
+          return JSValue::encode(jsUndefined());
+        });
+    promise->performPromiseThen(globalObject, resolve, reject,
+                                afterOngoingPromiseCapability);
+    return JSValue::encode(afterOngoingPromiseCapability);
+  }
 
-  JSValue stackTraces = JSONParse(
-      globalObject, samplingProfiler.stackTracesAsJSON()->toJSONString());
-
-  samplingProfiler.shutdown();
-  samplingProfiler.clearData();
-
-  JSObject *result =
-      constructEmptyObject(globalObject, globalObject->objectPrototype(), 3);
-  result->putDirect(vm, Identifier::fromString(vm, "functions"_s),
-                    jsString(vm, topFunctions.toString()));
-  result->putDirect(vm, Identifier::fromString(vm, "bytecodes"_s),
-                    jsString(vm, byteCodes.toString()));
-  result->putDirect(vm, Identifier::fromString(vm, "stackTraces"_s),
-                    stackTraces);
-
-  return JSValue::encode(result);
+  return JSValue::encode(report(vm, globalObject));
 }
 
 JSC_DECLARE_HOST_FUNCTION(functionGenerateHeapSnapshotForDebugging);

test/js/bun/jsc/bun-jsc.test.ts

@@ -25,6 +25,7 @@ import {
   drainMicrotasks,
   startRemoteDebugger,
   setTimeZone,
+  profile,
 } from "bun:jsc";
 
 describe("bun:jsc", () => {
@@ -170,4 +171,18 @@ describe("bun:jsc", () => {
     }
     Bun.gc(true);
   });
+
+  it("profile async", async () => {
+    const { promise, resolve } = Promise.withResolvers();
+    const result = await profile(
+      async function hey(arg1: number) {
+        await Bun.sleep(10).then(() => resolve(arguments));
+        return arg1;
+      },
+      1,
+      2,
+    );
+    const input = await promise;
+    expect({ ...input }).toStrictEqual({ "0": 2 });
+  });
 });