Fix crash in heap profile auto-selection at trace load time (#1435)

https://buganizer.corp.google.com/issues/416190258

When loading a trace with heap profile tracks, we automatically select
the earliest heap profile sample on trace load.
    
We currently assume that this event has id = 0, but this is not
necessarily the case.
    
This change refactors the queries in the heap profile plugin to make
things a lot clearer, as it was getting out of control, and also fixes
the bug so that it works out the necessary event id properly.
    
This change also refactors the selection logic in SelectionManager a
little to make the error message clearer and to ease debugging in the
future.
diff --git a/ui/src/core/selection_manager.ts b/ui/src/core/selection_manager.ts
index 22176aa..524bdab 100644
--- a/ui/src/core/selection_manager.ts
+++ b/ui/src/core/selection_manager.ts
@@ -394,12 +394,25 @@
     opts?: SelectionOpts,
     serializedDetailsPanel?: unknown,
   ) {
-    const details = await this.trackManager
-      .getTrack(trackUri)
-      ?.track.getSelectionDetails?.(eventId);
+    const track = this.trackManager.getTrack(trackUri);
+    if (!track) {
+      throw new Error(
+        `Unable to resolve selection details: Track ${trackUri} not found`,
+      );
+    }
 
+    const trackRenderer = track.track;
+    if (!trackRenderer.getSelectionDetails) {
+      throw new Error(
+        `Unable to resolve selection details: Track ${trackUri} does not support selection details`,
+      );
+    }
+
+    const details = await trackRenderer.getSelectionDetails(eventId);
     if (!exists(details)) {
-      throw new Error('Unable to resolve selection details');
+      throw new Error(
+        `Unable to resolve selection details: Track ${trackUri} returned no details for event ${eventId}`,
+      );
     }
 
     const selection: TrackEventSelection = {
diff --git a/ui/src/plugins/dev.perfetto.HeapProfile/heap_profile_track.ts b/ui/src/plugins/dev.perfetto.HeapProfile/heap_profile_track.ts
index 7adb910..a7d495c 100644
--- a/ui/src/plugins/dev.perfetto.HeapProfile/heap_profile_track.ts
+++ b/ui/src/plugins/dev.perfetto.HeapProfile/heap_profile_track.ts
@@ -39,6 +39,10 @@
         type: STR,
         id: NUM,
       },
+      filter: {
+        col: 'upid',
+        eq: upid,
+      },
     }),
     detailsPanel: (row) => {
       const ts = Time.fromRaw(row.ts);
diff --git a/ui/src/plugins/dev.perfetto.HeapProfile/index.ts b/ui/src/plugins/dev.perfetto.HeapProfile/index.ts
index 55479ad..a245ad7 100644
--- a/ui/src/plugins/dev.perfetto.HeapProfile/index.ts
+++ b/ui/src/plugins/dev.perfetto.HeapProfile/index.ts
@@ -15,120 +15,127 @@
 import {HEAP_PROFILE_TRACK_KIND} from '../../public/track_kinds';
 import {Trace} from '../../public/trace';
 import {PerfettoPlugin} from '../../public/plugin';
-import {LONG, NUM, STR} from '../../trace_processor/query_result';
+import {NUM} from '../../trace_processor/query_result';
 import {createHeapProfileTrack} from './heap_profile_track';
 import {TrackNode} from '../../public/workspace';
 import {createPerfettoTable} from '../../trace_processor/sql_utils';
 import ProcessThreadGroupsPlugin from '../dev.perfetto.ProcessThreadGroups';
+import {Track} from '../../public/track';
 
-function getUriForTrack(upid: number): string {
-  return `/process_${upid}/heap_profile`;
-}
+const EVENT_TABLE_NAME = 'heap_profile_events';
 
 export default class implements PerfettoPlugin {
   static readonly id = 'dev.perfetto.HeapProfile';
   static readonly dependencies = [ProcessThreadGroupsPlugin];
 
-  async onTraceLoad(ctx: Trace): Promise<void> {
-    const it = await ctx.engine.query(`
-      select value from stats
-      where name = 'heap_graph_non_finalized_graph'
-    `);
-    const incomplete = it.firstRow({value: NUM}).value > 0;
+  private readonly trackMap = new Map<number, Track>();
 
-    const result = await ctx.engine.query(`
-      select distinct upid from heap_profile_allocation
-      union
-      select distinct upid from heap_graph_object
+  async onTraceLoad(trace: Trace): Promise<void> {
+    await this.createHeapProfileTable(trace);
+    await this.addProcessTracks(trace);
+
+    trace.onTraceReady.addListener(async () => {
+      await this.selectFirstHeapProfile(trace);
+    });
+  }
+
+  private async createHeapProfileTable(trace: Trace) {
+    await createPerfettoTable(
+      trace.engine,
+      EVENT_TABLE_NAME,
+      `
+        SELECT
+          MIN(id) as id,
+          graph_sample_ts AS ts,
+          upid,
+          0 AS dur,
+          0 AS depth,
+          'graph' AS type
+        FROM heap_graph_object
+        GROUP BY graph_sample_ts, upid
+
+        UNION ALL
+
+        SELECT
+          MIN(id) as id,
+          ts,
+          upid,
+          0 AS dur,
+          0 AS depth,
+          'heap_profile:' || GROUP_CONCAT(DISTINCT heap_name) AS type
+        FROM heap_profile_allocation
+        GROUP BY ts, upid
+      `,
+    );
+  }
+
+  private async addProcessTracks(trace: Trace) {
+    const trackGroupsPlugin = trace.plugins.getPlugin(
+      ProcessThreadGroupsPlugin,
+    );
+    const incomplete = await this.getIncomplete(trace);
+    const result = await trace.engine.query(`
+      SELECT DISTINCT 
+        upid
+      FROM ${EVENT_TABLE_NAME}
     `);
     for (const it = result.iter({upid: NUM}); it.valid(); it.next()) {
       const upid = it.upid;
-      const uri = getUriForTrack(upid);
+      const uri = `/process_${upid}/heap_profile`;
       const title = 'Heap Profile';
-      const tableName = `_heap_profile_${upid}`;
 
-      createPerfettoTable(
-        ctx.engine,
-        tableName,
-        `
-          with
-            heaps as (select group_concat(distinct heap_name) h from heap_profile_allocation where upid = ${upid}),
-            allocation_tses as (select distinct ts from heap_profile_allocation where upid = ${upid}),
-            graph_tses as (select distinct graph_sample_ts from heap_graph_object where upid = ${upid})
-          select
-            *,
-            0 AS dur,
-            0 AS depth
-          from (
-            select
-              (
-                select a.id
-                from heap_profile_allocation a
-                where a.ts = t.ts
-                order by a.id
-                limit 1
-              ) as id,
-              ts,
-              'heap_profile:' || (select h from heaps) AS type
-            from allocation_tses t
-            union all
-            select
-              (
-                select o.id
-                from heap_graph_object o
-                where o.graph_sample_ts = g.graph_sample_ts
-                order by o.id
-                limit 1
-              ) as id,
-              graph_sample_ts AS ts,
-              'graph' AS type
-            from graph_tses g
-          )
-        `,
-      );
-
-      ctx.tracks.registerTrack({
+      const track: Track = {
         uri,
         title,
         tags: {
           kind: HEAP_PROFILE_TRACK_KIND,
           upid,
         },
-        track: createHeapProfileTrack(ctx, uri, tableName, upid, incomplete),
-      });
-      const group = ctx.plugins
-        .getPlugin(ProcessThreadGroupsPlugin)
-        .getGroupForProcess(upid);
-      const track = new TrackNode({uri, title, sortOrder: -30});
-      group?.addChildInOrder(track);
+        track: createHeapProfileTrack(
+          trace,
+          uri,
+          EVENT_TABLE_NAME,
+          upid,
+          incomplete,
+        ),
+      };
+
+      trace.tracks.registerTrack(track);
+      this.trackMap.set(upid, track);
+
+      const group = trackGroupsPlugin.getGroupForProcess(upid);
+      const trackNode = new TrackNode({uri, title, sortOrder: -30});
+      group?.addChildInOrder(trackNode);
     }
-
-    ctx.onTraceReady.addListener(async () => {
-      await selectFirstHeapProfile(ctx);
-    });
   }
-}
 
-async function selectFirstHeapProfile(ctx: Trace) {
-  const query = `
-    select * from (
-      select
-        min(ts) AS ts,
-        'heap_profile:' || group_concat(distinct heap_name) AS type,
-        upid
-      from heap_profile_allocation
-      group by upid
-      union
-      select distinct graph_sample_ts as ts, 'graph' as type, upid
-      from heap_graph_object
-    )
-    order by ts
-    limit 1
-  `;
-  const profile = await ctx.engine.query(query);
-  if (profile.numRows() !== 1) return;
-  const row = profile.firstRow({ts: LONG, type: STR, upid: NUM});
-  const upid = row.upid;
+  private async getIncomplete(trace: Trace): Promise<boolean> {
+    const it = await trace.engine.query(`
+      SELECT value
+      FROM stats
+      WHERE name = 'heap_graph_non_finalized_graph'
+    `);
+    const incomplete = it.firstRow({value: NUM}).value > 0;
+    return incomplete;
+  }
 
-  ctx.selection.selectTrackEvent(getUriForTrack(upid), 0);
+  private async selectFirstHeapProfile(ctx: Trace) {
+    // Select the first sample from each track
+    const result = await ctx.engine.query(`
+        SELECT
+          id,
+          upid
+        FROM ${EVENT_TABLE_NAME}
+        ORDER BY ts
+        LIMIT 1
+      `);
+
+    const iter = result.maybeFirstRow({id: NUM, upid: NUM});
+    if (!iter) return;
+
+    const track = this.trackMap.get(iter.upid);
+    if (!track) return;
+
+    ctx.selection.selectTrackEvent(track.uri, iter.id);
+  }
 }