Merge pull request #13628 from meteor/support-async-profiler

Support async profiler actions

tools/static-assets/server/boot.js

@@ -498,12 +498,10 @@ var runMain = Profile("Run main()", async function () {
     global.__METEOR_ASYNC_LOCAL_STORAGE = new AsyncLocalStorage();
   }
 
-  await Profile.run('Server startup', function() {
-    return global.__METEOR_ASYNC_LOCAL_STORAGE.run({}, async () => {
-      await loadServerBundles();
-      await callStartupHooks();
-      await runMain();
-    });
+  await Profile.run('Server startup', async function() {
+    await loadServerBundles();
+    await callStartupHooks();
+    await runMain();
   });
 })().catch(e => {
   console.log('error on boot.js',  e )

tools/tool-env/profile.ts

@@ -161,7 +161,6 @@
 //
 // In both reports the grand total is 600ms.
 
-
 const filter = parseFloat(process.env.METEOR_PROFILE || "100"); // ms
 
 type Stats = {
@@ -230,7 +229,6 @@ function decodeEntryKey(key: string) {
   return key.split('\t');
 }
 
-const globalEntry: string[] = [];
 let running = false;
 
 export function Profile<
@@ -238,47 +236,79 @@ export function Profile<
   TResult,
 >(
   bucketName: string | ((...args: TArgs) => string),
-  f: (...args: TArgs) => TResult,
+  f: (...args: TArgs) => TResult | Promise<TResult>,
 ): typeof f {
   if (! Profile.enabled) {
     return f;
   }
 
   return Object.assign(function profileWrapper(this: any) {
-    if (! running) {
-      return f.apply(this, arguments as any);
+    const args = arguments as unknown as TArgs;
+
+    if (!running) {
+      return f.apply(this, args);
     }
 
-    const name = typeof bucketName === "function"
-      ? bucketName.apply(this, arguments as any)
-      : bucketName;
+    const asyncLocalStorage = global.__METEOR_ASYNC_LOCAL_STORAGE;
+    let existingStore = asyncLocalStorage.getStore();
 
-    // TODO Test with Profile / use __METEOR_ASYNC_LOCAL_STORAGE
-    //const currentStore = asyncLo
-    // const currentEntry = Fiber.current
-    //   ? Fiber.current.profilerEntry || (Fiber.current.profilerEntry = [])
-    //   : globalEntry;
-    const currentEntry = globalEntry;
-
-    currentEntry.push(name);
-    const key = encodeEntryKey(currentEntry);
-    const start = process.hrtime();
-    try {
-      return f.apply(this, arguments as any);
-    } finally {
-      const elapsed = process.hrtime(start);
-      const stats = (bucketStats[key] || (bucketStats[key] = {
-        time: 0.0,
-        count: 0,
-        isOther: false,
-      }));
-      stats.time += (elapsed[0] * 1000 + elapsed[1] / 1000000);
-      stats.count++;
-      currentEntry.pop();
+    // Ensure `existingStore` is an object, initialize it if not present
+    if (!existingStore) {
+      existingStore = { currentEntry: [] };
+      return asyncLocalStorage.run(existingStore, () => runWithContext(bucketName, existingStore, f, this, args));
     }
+
+    // Ensure `existingStore.currentEntry` exists
+    if (!existingStore.currentEntry) {
+      existingStore.currentEntry = [];
+    }
+
+    return runWithContext(bucketName, existingStore, f, this, args);
   }, f) as typeof f;
 }
 
+function runWithContext<TArgs extends any[], TResult>(
+    bucketName: string | ((...args: TArgs) => string),
+    store: { currentEntry: string[]; [key: string]: any }, // Keep other properties intact
+    f: (...args: TArgs) => TResult | Promise<TResult>,
+    context: any,
+    args: IArguments,
+): TResult | Promise<TResult> {
+  const name = typeof bucketName === "function" ? bucketName.apply(context, args) : bucketName;
+  store.currentEntry = [...store.currentEntry || [], name];
+  const key = encodeEntryKey(store.currentEntry);
+  const start = process.hrtime();
+
+  let result: TResult | Promise<TResult>;
+  try {
+    result = f.apply(context, args);
+
+    if (result instanceof Promise) {
+      // Return a promise if async
+      return result.finally(() => finalizeProfiling(key, start, store.currentEntry));
+    }
+
+    // Return directly if sync
+    return result;
+  } finally {
+    if (!(result instanceof Promise)) {
+      finalizeProfiling(key, start, store.currentEntry);
+    }
+  }
+}
+
+function finalizeProfiling(key: string, start: [number, number], currentEntry: string[]) {
+  const elapsed = process.hrtime(start);
+  const stats = (bucketStats[key] || (bucketStats[key] = {
+    time: 0.0,
+    count: 0,
+    isOther: false,
+  }));
+  stats.time += elapsed[0] * 1000 + elapsed[1] / 1_000_000;
+  stats.count++;
+  currentEntry.pop();
+}
+
 export namespace Profile {
   export let enabled = !! process.env.METEOR_PROFILE;