Show heap graphs in UI.

Change-Id: I149a84573716ecc25334c20f730f9a33cf17b9c2
diff --git a/src/trace_processor/importers/proto/heap_graph_tracker.cc b/src/trace_processor/importers/proto/heap_graph_tracker.cc
index f96cf6e..9eb166b 100644
--- a/src/trace_processor/importers/proto/heap_graph_tracker.cc
+++ b/src/trace_processor/importers/proto/heap_graph_tracker.cc
@@ -230,8 +230,7 @@
     base::Optional<uint32_t> parent_id;
     if (node.parent_id != 0)
       parent_id = node_to_row_idx[node.parent_id];
-    const uint32_t depth = node.depth - 1;  // -1 because we do not have the
-                                            // artificial root in the database.
+    const uint32_t depth = node.depth;
 
     tables::ExperimentalFlamegraphNodesTable::Row alloc_row{
         current_ts, current_upid, profile_type, depth,
diff --git a/src/trace_processor/trace_storage.cc b/src/trace_processor/trace_storage.cc
index cbd01b5..74d22bf 100644
--- a/src/trace_processor/trace_storage.cc
+++ b/src/trace_processor/trace_storage.cc
@@ -147,6 +147,14 @@
                            &end_ns);
   DbTableMaybeUpdateMinMax(instant_table_.ts(), &start_ns, &end_ns);
   DbTableMaybeUpdateMinMax(android_log_table_.ts(), &start_ns, &end_ns);
+  DbTableMaybeUpdateMinMax(heap_graph_object_table_.graph_sample_ts(),
+                           &start_ns, &end_ns);
+
+  if (heap_graph_object_table_.row_count() != 0) {
+    // TODO(fmayer): Remove this.
+    // Work around UI bug that the heap marker is not displayed.
+    end_ns += 50000000;  // 50 ms
+  }
 
   if (start_ns == std::numeric_limits<int64_t>::max()) {
     return std::make_pair(0, 0);
diff --git a/test/trace_processor/heap_graph_flamegraph_system-server-heap-graph.out b/test/trace_processor/heap_graph_flamegraph_system-server-heap-graph.out
index 102c246..8f822d1 100644
--- a/test/trace_processor/heap_graph_flamegraph_system-server-heap-graph.out
+++ b/test/trace_processor/heap_graph_flamegraph_system-server-heap-graph.out
@@ -1,11 +1,11 @@
 "id","depth","name","map_name","count","cumulative_count","size","cumulative_size","parent_id"
-0,4294967295,"java.lang.Class<boolean>","JAVA",2,4,240,276,"[NULL]"
-1,0,"java.lang.Object[]","JAVA",1,1,12,12,0
-2,0,"java.lang.String","JAVA",1,1,24,24,0
-3,4294967295,"java.lang.Class<byte>","JAVA",3,4,360,384,"[NULL]"
-4,0,"java.lang.String","JAVA",1,1,24,24,3
-5,4294967295,"java.lang.Class<short>","JAVA",2,3,240,264,"[NULL]"
-6,0,"java.lang.String","JAVA",1,1,24,24,5
-7,4294967295,"java.lang.Class<char>","JAVA",2,3,240,264,"[NULL]"
-8,0,"java.lang.String","JAVA",1,1,24,24,7
-9,4294967295,"java.lang.Class<int>","JAVA",2,3,240,264,"[NULL]"
+0,0,"java.lang.Class<boolean>","JAVA",2,4,240,276,"[NULL]"
+1,1,"java.lang.Object[]","JAVA",1,1,12,12,0
+2,1,"java.lang.String","JAVA",1,1,24,24,0
+3,0,"java.lang.Class<byte>","JAVA",3,4,360,384,"[NULL]"
+4,1,"java.lang.String","JAVA",1,1,24,24,3
+5,0,"java.lang.Class<short>","JAVA",2,3,240,264,"[NULL]"
+6,1,"java.lang.String","JAVA",1,1,24,24,5
+7,0,"java.lang.Class<char>","JAVA",2,3,240,264,"[NULL]"
+8,1,"java.lang.String","JAVA",1,1,24,24,7
+9,0,"java.lang.Class<int>","JAVA",2,3,240,264,"[NULL]"
diff --git a/ui/src/common/actions.ts b/ui/src/common/actions.ts
index c5f7f63..8e43ca5 100644
--- a/ui/src/common/actions.ts
+++ b/ui/src/common/actions.ts
@@ -426,18 +426,26 @@
       },
 
   selectHeapProfile(
-      state: StateDraft, args: {id: number, upid: number, ts: number}): void {
-    state.currentSelection =
-        {kind: 'HEAP_PROFILE', id: args.id, upid: args.upid, ts: args.ts};
+      state: StateDraft,
+      args: {id: number, upid: number, ts: number, type: string}): void {
+    state.currentSelection = {
+      kind: 'HEAP_PROFILE',
+      id: args.id,
+      upid: args.upid,
+      ts: args.ts,
+      type: args.type
+    };
   },
 
   showHeapProfileFlamegraph(
-      state: StateDraft, args: {id: number, upid: number, ts: number}): void {
+      state: StateDraft,
+      args: {id: number, upid: number, ts: number, type: string}): void {
     state.currentHeapProfileFlamegraph = {
       kind: 'HEAP_PROFILE_FLAMEGRAPH',
       id: args.id,
       upid: args.upid,
       ts: args.ts,
+      type: args.type,
     };
   },
 
diff --git a/ui/src/common/state.ts b/ui/src/common/state.ts
index 2b6d2b2..b7f2f40 100644
--- a/ui/src/common/state.ts
+++ b/ui/src/common/state.ts
@@ -163,6 +163,7 @@
   id: number;
   upid: number;
   ts: number;
+  type: string;
 }
 
 export interface HeapProfileFlamegraph {
@@ -170,6 +171,7 @@
   id: number;
   upid: number;
   ts: number;
+  type: string;
   expandedCallsite?: CallsiteInfo;
   viewingOption?: string;
 }
diff --git a/ui/src/controller/heap_profile_controller.ts b/ui/src/controller/heap_profile_controller.ts
index eadb704..2f5d5f7 100644
--- a/ui/src/controller/heap_profile_controller.ts
+++ b/ui/src/controller/heap_profile_controller.ts
@@ -85,7 +85,8 @@
                           selectedHeapProfile.viewingOption :
                           DEFAULT_VIEWING_OPTION,
                       selection.ts,
-                      selectedHeapProfile.upid)
+                      selectedHeapProfile.upid,
+                      selectedHeapProfile.type)
                   .then(flamegraphData => {
                     if (flamegraphData !== undefined && selection &&
                         selection.kind === selectedHeapProfile.kind &&
@@ -119,6 +120,7 @@
       id: heapProfile.id,
       upid: heapProfile.upid,
       ts: heapProfile.ts,
+      type: heapProfile.type,
       expandedCallsite: heapProfile.expandedCallsite,
       viewingOption: heapProfile.viewingOption
     };
@@ -130,6 +132,7 @@
          (this.lastSelectedHeapProfile !== undefined &&
           (this.lastSelectedHeapProfile.id !== selection.id ||
            this.lastSelectedHeapProfile.ts !== selection.ts ||
+           this.lastSelectedHeapProfile.type !== selection.type ||
            this.lastSelectedHeapProfile.upid !== selection.upid ||
            this.lastSelectedHeapProfile.viewingOption !==
                selection.viewingOption ||
@@ -151,8 +154,8 @@
 
 
   async getFlamegraphData(
-      baseKey: string, viewingOption: string, ts: number,
-      upid: number): Promise<CallsiteInfo[]> {
+      baseKey: string, viewingOption: string, ts: number, upid: number,
+      type: string): Promise<CallsiteInfo[]> {
     let currentData: CallsiteInfo[];
     const key = `${baseKey}-${viewingOption}`;
     if (this.flamegraphDatasets.has(key)) {
@@ -163,7 +166,7 @@
       // Collecting data for drawing flamegraph for selected heap profile.
       // Data needs to be in following format:
       // id, name, parent_id, depth, total_size
-      const tableName = await this.prepareViewsAndTables(ts, upid);
+      const tableName = await this.prepareViewsAndTables(ts, upid, type);
       currentData =
           await this.getFlamegraphDataFromTables(tableName, viewingOption);
       this.flamegraphDatasets.set(key, currentData);
@@ -201,7 +204,7 @@
     }
 
     const callsites = await this.args.engine.query(
-        `SELECT id, name, parent_id, depth, cumulative_size,
+        `SELECT id, name, IFNULL(parent_id, -1), depth, cumulative_size,
         cumulative_alloc_size, cumulative_count,
         cumulative_alloc_count, map_name, size from ${tableName} ${orderBy}`);
 
@@ -227,7 +230,7 @@
     return flamegraphData;
   }
 
-  private async prepareViewsAndTables(ts: number, upid: number):
+  private async prepareViewsAndTables(ts: number, upid: number, type: string):
       Promise<string> {
     // Creating unique names for views so we can reuse and not delete them
     // for each marker.
@@ -239,7 +242,7 @@
         select id, name, map_name, parent_id, depth, cumulative_size,
           cumulative_alloc_size, cumulative_count, cumulative_alloc_count,
           size, alloc_size, count, alloc_count
-        from experimental_flamegraph(${ts}, ${upid}, 'native')`);
+        from experimental_flamegraph(${ts}, ${upid}, '${type}')`);
     return tableNameGroupedCallsitesForFlamegraph;
   }
 
diff --git a/ui/src/controller/trace_controller.ts b/ui/src/controller/trace_controller.ts
index f819244..63e45b2 100644
--- a/ui/src/controller/trace_controller.ts
+++ b/ui/src/controller/trace_controller.ts
@@ -403,7 +403,9 @@
     }
 
     const heapProfiles = await engine.query(`
-      select distinct(upid) from heap_profile_allocation`);
+      select distinct(upid) from heap_profile_allocation
+      union
+      select distinct(upid) from heap_graph_object`);
 
     const heapUpids: Set<number> = new Set();
     for (let i = 0; i < heapProfiles.numRecords; i++) {
diff --git a/ui/src/tracks/heap_profile/common.ts b/ui/src/tracks/heap_profile/common.ts
index fc9da62..0a8a16d 100644
--- a/ui/src/tracks/heap_profile/common.ts
+++ b/ui/src/tracks/heap_profile/common.ts
@@ -17,6 +17,7 @@
 
 export interface Data extends TrackData {
   tsStarts: Float64Array;
+  types: string[];
 }
 
 export interface Config {
diff --git a/ui/src/tracks/heap_profile/controller.ts b/ui/src/tracks/heap_profile/controller.ts
index 61ad335..bab0b15 100644
--- a/ui/src/tracks/heap_profile/controller.ts
+++ b/ui/src/tracks/heap_profile/controller.ts
@@ -28,11 +28,23 @@
   async onBoundsChange(start: number, end: number, resolution: number):
       Promise<Data> {
     if (this.config.upid === undefined) {
-      return {start, end, resolution, length: 0, tsStarts: new Float64Array()};
+      return {
+        start,
+        end,
+        resolution,
+        length: 0,
+        tsStarts: new Float64Array(),
+        types: new Array<string>()
+      };
     }
     const result = await this.query(`
-    select distinct(ts) from heap_profile_allocation where upid = ${
-        this.config.upid}`);
+    select * from
+    (select distinct(ts) as ts, 'native' as type from heap_profile_allocation
+     where upid = ${this.config.upid}
+        union
+        select distinct(graph_sample_ts) as ts, 'graph' as type from
+        heap_graph_object
+        where upid = ${this.config.upid}) order by ts`);
     const numRows = +result.numRecords;
     const data: Data = {
       start,
@@ -40,10 +52,12 @@
       resolution,
       length: numRows,
       tsStarts: new Float64Array(numRows),
+      types: new Array<string>(numRows),
     };
 
     for (let row = 0; row < numRows; row++) {
       data.tsStarts[row] = +result.columns[0].longValues![row];
+      data.types[row] = result.columns[1].stringValues![row];
     }
 
     return data;
diff --git a/ui/src/tracks/heap_profile/frontend.ts b/ui/src/tracks/heap_profile/frontend.ts
index 8eaba5e..86d177e 100644
--- a/ui/src/tracks/heap_profile/frontend.ts
+++ b/ui/src/tracks/heap_profile/frontend.ts
@@ -118,10 +118,11 @@
 
     if (index !== -1) {
       const ts = data.tsStarts[index];
+      const type = data.types[index];
       globals.dispatch(Actions.showHeapProfileFlamegraph(
-          {id: index, upid: this.config.upid, ts}));
-      globals.makeSelection(
-          Actions.selectHeapProfile({id: index, upid: this.config.upid, ts}));
+          {id: index, upid: this.config.upid, ts, type}));
+      globals.makeSelection(Actions.selectHeapProfile(
+          {id: index, upid: this.config.upid, ts, type}));
       return true;
     }
     return false;