Parse trace ids from samples in trace CPU profiles

As per go/chrome-devtools:better-js-flamecharts, this CL does the
frontend bits to extract the trace id from samples and use it to obtain
the stack trace of the trace event with the same trace id. This stack
trace is used in the consolidation of tracing and sampling data that
creates the JS calls in the thread flamecharts.

Bug: 390155857
Change-Id: I6c21eb9a0b8ecf61422c249d6ebe9fa68091c3eb
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/6218850
Commit-Queue: Andres Olivares <[email protected]>
Reviewed-by: Jack Franklin <[email protected]>

Author: and-oli

front_end/models/cpu_profile/CPUProfileDataModel.ts

@@ -50,6 +50,12 @@ export class CPUProfileDataModel extends ProfileTreeModel {
   profileEndTime: number /* milliseconds */;
   timestamps: number[];
   samples: number[]|undefined;
+  /**
+   * Contains trace ids assigned to samples, if any. Trace ids are
+   * keyed by the sample index in the profile. These are only created
+   * for CPU profiles coming from traces.
+   */
+  traceIds?: Record<string, number>;
   lines?: number[];
   totalHitCount: number;
   profileHead: CPUProfileNode;
@@ -81,6 +87,7 @@ export class CPUProfileDataModel extends ProfileTreeModel {
       this.profileEndTime = profile.endTime / 1000;
       this.timestamps = this.convertTimeDeltas(profile);
     }
+    this.traceIds = profile.traceIds;
     this.samples = profile.samples;
 
     // Lines are available only in profiles coming from tracing.
@@ -547,5 +554,8 @@ export class CPUProfileDataModel extends ProfileTreeModel {
 
 // Format used by profiles coming from traces.
 export type ExtendedProfileNode = Protocol.Profiler.ProfileNode&{parent?: number};
-export type ExtendedProfile =
-    Protocol.Profiler.Profile&{nodes: Protocol.Profiler.ProfileNode[] | ExtendedProfileNode[], lines?: number[]};
+export type ExtendedProfile = Protocol.Profiler.Profile&{
+  nodes: Protocol.Profiler.ProfileNode[] | ExtendedProfileNode[],
+  lines?: number[],
+  traceIds?: Record<string, number>,
+};

front_end/models/trace/handlers/SamplesHandler.test.ts

@@ -255,7 +255,8 @@ describeWithEnvironment('SamplesHandler', function() {
       assert.strictEqual(profileById.size, 1);
       const cpuProfileData = profileById.values().next().value as Trace.Handlers.ModelHandlers.Samples.ProfileData;
       const cpuProfile = cpuProfileData.rawProfile;
-      assert.deepEqual(Object.keys(cpuProfile), ['startTime', 'endTime', 'nodes', 'samples', 'timeDeltas', 'lines']);
+      assert.deepEqual(
+          Object.keys(cpuProfile), ['startTime', 'endTime', 'nodes', 'samples', 'timeDeltas', 'lines', 'traceIds']);
       assert.lengthOf(cpuProfile.nodes, 153);
       assert.strictEqual(cpuProfile.startTime, 287510826176);
       assert.strictEqual(cpuProfile.endTime, 287510847633);

front_end/models/trace/handlers/SamplesHandler.ts

@@ -152,6 +152,7 @@ export function handleEvent(event: Types.Events.Event): void {
     const cdpProfile = profileData.rawProfile;
     const nodesAndSamples: Types.Events.PartialProfile|undefined = event.args?.data?.cpuProfile || {samples: []};
     const samples = nodesAndSamples?.samples || [];
+    const traceIds = event.args?.data?.cpuProfile?.trace_ids || {};
     const nodes: CPUProfile.CPUProfileDataModel.ExtendedProfileNode[] = [];
     for (const n of nodesAndSamples?.nodes || []) {
       const lineNumber = typeof n.callFrame.lineNumber === 'undefined' ? -1 : n.callFrame.lineNumber;
@@ -178,6 +179,7 @@ export function handleEvent(event: Types.Events.Event): void {
     cdpProfile.samples?.push(...samples);
     cdpProfile.timeDeltas?.push(...timeDeltas);
     cdpProfile.lines?.push(...lines);
+    cdpProfile.traceIds = {...(cdpProfile.traceIds || {}), ...traceIds};
     if (cdpProfile.samples && cdpProfile.timeDeltas && cdpProfile.samples.length !== cdpProfile.timeDeltas.length) {
       console.error('Failed to parse CPU profile.');
       return;

front_end/models/trace/helpers/SamplesIntegrator.test.ts

@@ -225,6 +225,74 @@ describeWithEnvironment('SamplesIntegrator', function() {
       assert.strictEqual(constructedCalls[1].ts, 300);
       assert.strictEqual(constructedCalls[1].dur, 200);
     });
+    it('uses the stack in a sample for a trace event when connected via a common trace id ', () => {
+      const traceId = 123;
+      const cdpProfile: CPUProfile.CPUProfileDataModel.ExtendedProfile = {
+        startTime: 0,
+        endTime: 3000,
+        nodes: [
+          {
+            id: 1,
+            hitCount: 0,
+            callFrame: {functionName: '(root)', scriptId, url, lineNumber, columnNumber},
+            children: [2],
+          },
+          {id: 2, callFrame: {functionName: 'foo', scriptId, url, lineNumber, columnNumber}, children: [3, 4]},
+          {id: 3, callFrame: {functionName: 'bar', scriptId, url, lineNumber, columnNumber}},
+          {id: 4, callFrame: {functionName: 'baz', scriptId, url, lineNumber, columnNumber}, children: [5]},
+          {id: 5, callFrame: {functionName: 'sheep', scriptId, url, lineNumber, columnNumber}},
+        ],
+        samples: [3, 3, 3, 5],
+        traceIds: {3: traceId},
+        timeDeltas: new Array(4).fill(100),
+      };
+
+      // Before integrating samples and trace events, the flamechart
+      // looks roughly like:
+      //   |----------------EvaluateScript-----------------|
+      //       |-----------------v8.run--------------------|
+      //                                    |---TimeStamp--|
+      //          | foo |  | foo |  | foo |  | foo |
+      //          | bar |  | bar |  | bar |  | baz |
+      //                                     |sheep|
+
+      // After integrating, it should look as follows. note that the the
+      // stack in the last sample should be used to parent `TimeStamp`,
+      // event if that sample's ts is later, since that event and that
+      // sample are connected by a trace id.
+      //   |----------------EvaluateScript-----------------|
+      //       |-----------------v8.run--------------------|
+      //          |------------------foo-------------------|
+      //          |----------bar--------||--------baz------|
+      //                                 |------sheep------|
+      //                                 |-Timestamp-|
+
+      const parsedProfile = new CPUProfile.CPUProfileDataModel.CPUProfileDataModel(cdpProfile);
+      const integrator = new Trace.Helpers.SamplesIntegrator.SamplesIntegrator(parsedProfile, PROFILE_ID, pid, tid);
+      const evaluateScript = makeCompleteEvent(Trace.Types.Events.Name.EVALUATE_SCRIPT, 0, 500);
+      const v8Run = makeCompleteEvent('v8.run', 10, 490);
+      const consoleTimeStamp =
+          makeCompleteEvent(Trace.Types.Events.Name.CONSOLE_TIME_STAMP, 350, 10) as Trace.Types.Events.ConsoleTimeStamp;
+      consoleTimeStamp.args = {data: {name: 'A timestamp', sampleTraceId: traceId}};
+      const traceEvents = [evaluateScript, v8Run, consoleTimeStamp];
+      const constructedCalls = integrator.buildProfileCalls(traceEvents);
+      assert.lengthOf(constructedCalls, 4);
+      assert.strictEqual(constructedCalls[0].callFrame.functionName, 'foo');
+      assert.strictEqual(constructedCalls[0].ts, 100);
+      assert.strictEqual(constructedCalls[0].dur, 400);  // 100 - 500
+
+      assert.strictEqual(constructedCalls[1].callFrame.functionName, 'bar');
+      assert.strictEqual(constructedCalls[1].ts, 100);
+      assert.strictEqual(constructedCalls[1].dur, 250);  // 100 - 350
+
+      assert.strictEqual(constructedCalls[2].callFrame.functionName, 'baz');
+      assert.strictEqual(constructedCalls[2].ts, 350);
+      assert.strictEqual(constructedCalls[2].dur, 150);  // 350 - 500
+
+      assert.strictEqual(constructedCalls[3].callFrame.functionName, 'sheep');
+      assert.strictEqual(constructedCalls[3].ts, 350);
+      assert.strictEqual(constructedCalls[3].dur, 150);  // 350 - 500
+    });
     it('restarts the call frame stack when a new top level event is encountered', () => {
       // Profile contains the following samples:
       // |a||a||a||a|

front_end/models/trace/helpers/SamplesIntegrator.ts

@@ -7,7 +7,7 @@ import type * as CPUProfile from '../../cpu_profile/cpu_profile.js';
 import * as Types from '../types/types.js';
 
 import {milliToMicro} from './Timing.js';
-import {makeProfileCall, mergeEventsInOrder} from './Trace.js';
+import {extractSampleTraceId, makeProfileCall, mergeEventsInOrder} from './Trace.js';
 
 /**
  * This is a helper that integrates CPU profiling data coming in the
@@ -70,6 +70,12 @@ export class SamplesIntegrator {
    * in the stack before the event came.
    */
   #lockedJsStackDepth: number[] = [];
+  /**
+   * For samples with a trace id, creates a profile call and keeps it
+   * in a record keyed by that id. The value is typed as an union with
+   * undefined to avoid nullish accesses when a key is not present.
+   */
+  #callsByTraceIds: Record<number, Types.Events.SyntheticProfileCall|undefined> = {};
   /**
    * Used to keep track when samples should be integrated even if they
    * are not children of invocation trace events. This is useful in
@@ -243,11 +249,18 @@ export class SamplesIntegrator {
       if (!node) {
         continue;
       }
+      const maybeTraceId = this.#profileModel.traceIds?.[i];
       const call = makeProfileCall(node, this.#profileId, i, timestamp, this.#processId, this.#threadId);
-      calls.push(call);
-
+      // Separate samples with trace ids so that they are only used when
+      // processing the owner event.
+      if (maybeTraceId === undefined) {
+        calls.push(call);
+      } else {
+        this.#callsByTraceIds[maybeTraceId] = call;
+      }
       if (debugModeEnabled) {
-        this.jsSampleEvents.push(this.#makeJSSampleEvent(call, timestamp));
+        const traceId = this.#profileModel.traceIds?.[i];
+        this.jsSampleEvents.push(this.#makeJSSampleEvent(call, timestamp, traceId));
       }
       if (node.id === this.#profileModel.gcNode?.id && prevNode) {
         // GC samples have no stack, so we just put GC node on top of the
@@ -261,7 +274,23 @@ export class SamplesIntegrator {
     return calls;
   }
 
-  #makeProfileCallsForStack(profileCall: Types.Events.SyntheticProfileCall): Types.Events.SyntheticProfileCall[] {
+  /**
+   * Given a synthetic profile call, returns an array of profile calls
+   * representing the stack trace that profile call belongs to based on
+   * its nodeId. The input profile call will be at the top of the
+   * returned stack (last position), meaning that any other frames that
+   * were effectively above it are omitted.
+   * @param profileCall
+   * @param overrideTimeStamp a custom timestamp to use for the returned
+   * profile calls. If not defined, the timestamp of the input
+   * profileCall is used instead. This param is useful for example when
+   * creating the profile calls for a sample with a trace id, since the
+   * timestamp of the corresponding trace event should be used instead
+   * of the sample's.
+   */
+
+  #makeProfileCallsForStack(profileCall: Types.Events.SyntheticProfileCall, overrideTimeStamp?: Types.Timing.Micro):
+      Types.Events.SyntheticProfileCall[] {
     let node = this.#profileModel.nodeById(profileCall.nodeId);
     const isGarbageCollection = node?.id === this.#profileModel.gcNode?.id;
     if (isGarbageCollection) {
@@ -286,7 +315,8 @@ export class SamplesIntegrator {
     // durations
     while (node) {
       callFrames[i--] = makeProfileCall(
-          node, profileCall.profileId, profileCall.sampleIndex, profileCall.ts, this.#processId, this.#threadId);
+          node, profileCall.profileId, profileCall.sampleIndex, overrideTimeStamp ?? profileCall.ts, this.#processId,
+          this.#threadId);
       node = node.parent;
     }
     return callFrames;
@@ -296,7 +326,18 @@ export class SamplesIntegrator {
    * Update tracked stack using this event's call stack.
    */
   #extractStackTrace(event: Types.Events.Event): void {
-    const stackTrace = Types.Events.isProfileCall(event) ? this.#makeProfileCallsForStack(event) : this.#currentJSStack;
+    let stackTrace = this.#currentJSStack;
+    if (Types.Events.isProfileCall(event)) {
+      stackTrace = this.#makeProfileCallsForStack(event);
+    }
+    const traceId = extractSampleTraceId(event);
+    if (traceId !== null) {
+      const maybeCallForTraceId = this.#callsByTraceIds[traceId];
+      if (maybeCallForTraceId) {
+        stackTrace = this.#makeProfileCallsForStack(maybeCallForTraceId, event.ts);
+      }
+    }
+
     SamplesIntegrator.filterStackFrames(stackTrace, this.#engineConfig);
 
     const endTime = event.ts + (event.dur || 0);
@@ -391,13 +432,13 @@ export class SamplesIntegrator {
     this.#currentJSStack.length = depth;
   }
 
-  #makeJSSampleEvent(call: Types.Events.SyntheticProfileCall, timestamp: Types.Timing.Micro):
+  #makeJSSampleEvent(call: Types.Events.SyntheticProfileCall, timestamp: Types.Timing.Micro, traceId?: number):
       Types.Events.SyntheticJSSample {
     const JSSampleEvent: Types.Events.SyntheticJSSample = {
       name: Types.Events.Name.JS_SAMPLE,
       cat: 'devtools.timeline',
       args: {
-        data: {stackTrace: this.#makeProfileCallsForStack(call).map(e => e.callFrame)},
+        data: {traceId, stackTrace: this.#makeProfileCallsForStack(call).map(e => e.callFrame)},
       },
       ph: Types.Events.Phase.INSTANT,
       ts: timestamp,

front_end/models/trace/helpers/Trace.ts

@@ -687,3 +687,10 @@ export function isMatchingCallFrame(
 export function eventContainsTimestamp(event: Types.Events.Event, ts: Types.Timing.Micro): boolean {
   return event.ts <= ts && event.ts + (event.dur || 0) >= ts;
 }
+
+export function extractSampleTraceId(event: Types.Events.Event): number|null {
+  if (Types.Events.isConsoleRunTask(event) || Types.Events.isConsoleTimeStamp(event)) {
+    return event.args?.data?.sampleTraceId || null;
+  }
+  return null;
+}

front_end/models/trace/types/TraceEvents.ts

@@ -174,6 +174,14 @@ export interface ProfileChunk extends Sample {
 export interface PartialProfile {
   nodes?: PartialNode[];
   samples: CallFrameID[];
+  /**
+   * Contains trace ids assigned to samples, if any. Trace ids are
+   * keyed by the sample index in the profile (the keys of the object
+   * are strings containing the numeric index).
+   */
+  /* eslint-disable @typescript-eslint/naming-convention */
+  trace_ids?: Record<string, number>;
+  /* eslint-enable @typescript-eslint/naming-convention */
 }
 
 export interface PartialNode {
@@ -1421,6 +1429,7 @@ export interface ConsoleTimeStamp extends Event {
       track?: string|number,
       trackGroup?: string|number,
       color?: string|number,
+      sampleTraceId?: number,
     },
   };
 }
@@ -1643,7 +1652,8 @@ export interface SyntheticJSSample extends Event {
   name: Name.JS_SAMPLE;
   args: Args&{
     data: ArgsData & {
-      stackTrace: Protocol.Runtime.CallFrame[],
+      // Used to associate a stack sample with a trace event.
+      traceId?: number, stackTrace: Protocol.Runtime.CallFrame[],
     },
   };
   ph: Phase.INSTANT;
@@ -2762,6 +2772,11 @@ export function isJSInvocationEvent(event: Event): boolean {
   return false;
 }
 export interface ConsoleRunTask extends Event {
+  args: Args&{
+    data: ArgsData & {
+      sampleTraceId?: number,
+    },
+  };
   name: Name.V8_CONSOLE_RUN_TASK;
 }