Merge "Add heap graph references."

7 files changed, 104 insertions, 1 deletions

diff --git a/src/trace_processor/heap_graph_tracker.cc b/src/trace_processor/heap_graph_tracker.cc
index c29b60d94..1f5bb7227 100644
--- a/src/trace_processor/heap_graph_tracker.cc
+++ b/src/trace_processor/heap_graph_tracker.cc
@@ -41,6 +41,11 @@ void HeapGraphTracker::AddInternedTypeName(uint64_t intern_id,
   interned_type_names_.emplace(intern_id, strid);
 }
 
+void HeapGraphTracker::AddInternedFieldName(uint64_t intern_id,
+                                            StringPool::Id strid) {
+  interned_field_names_.emplace(intern_id, strid);
+}
+
 void HeapGraphTracker::SetPacketIndex(uint64_t index) {
   if (prev_index_ != 0 && prev_index_ + 1 != index) {
     PERFETTO_ELOG("Missing packets between %" PRIu64 " and %" PRIu64,
@@ -61,8 +66,48 @@ void HeapGraphTracker::FinalizeProfile() {
     }
     context_->storage->mutable_heap_graph_object_table()->Insert(
         {current_upid_, current_ts_, static_cast<int64_t>(obj.object_id),
-         static_cast<int64_t>(obj.self_size), it->second});
+         static_cast<int64_t>(obj.self_size), -1, it->second});
+    object_id_to_row_.emplace(
+        obj.object_id, context_->storage->heap_graph_object_table().size() - 1);
+  }
+
+  for (const SourceObject& obj : current_objects_) {
+    auto it = object_id_to_row_.find(obj.object_id);
+    if (it == object_id_to_row_.end())
+      continue;
+    int64_t owner_row = it->second;
+
+    int64_t reference_set_id =
+        context_->storage->heap_graph_reference_table().size();
+    for (const SourceObject::Reference& ref : obj.references) {
+      // This is true for unset reference fields.
+      if (ref.owned_object_id == 0)
+        continue;
+
+      it = object_id_to_row_.find(ref.owned_object_id);
+      // This can only happen for an invalid type string id, which is already
+      // reported as an error. Silently continue here.
+      if (it == object_id_to_row_.end())
+        continue;
+
+      int64_t owned_row = it->second;
+      auto field_name_it = interned_field_names_.find(obj.type_id);
+      if (field_name_it == interned_field_names_.end()) {
+        context_->storage->IncrementIndexedStats(
+            stats::heap_graph_invalid_string_id,
+            static_cast<int>(current_upid_));
+        continue;
+      }
+      StringPool::Id field_name = field_name_it->second;
+      context_->storage->mutable_heap_graph_reference_table()->Insert(
+          {reference_set_id, owner_row, owned_row, field_name});
+    }
+    context_->storage->mutable_heap_graph_object_table()
+        ->mutable_reference_set_id()
+        ->Set(static_cast<uint32_t>(owner_row), reference_set_id);
   }
+  interned_field_names_.clear();
+  object_id_to_row_.clear();
   interned_type_names_.clear();
   current_objects_.clear();
   current_upid_ = 0;
diff --git a/src/trace_processor/heap_graph_tracker.h b/src/trace_processor/heap_graph_tracker.h
index 3431fbcd8..c65a7b1ab 100644
--- a/src/trace_processor/heap_graph_tracker.h
+++ b/src/trace_processor/heap_graph_tracker.h
@@ -34,15 +34,23 @@ class TraceProcessorContext;
 class HeapGraphTracker {
  public:
   struct SourceObject {
+    // All ids in this are in the trace iid space, not in the trace processor
+    // id space.
+    struct Reference {
+      uint64_t field_name_id = 0;
+      uint64_t owned_object_id = 0;
+    };
     uint64_t object_id = 0;
     uint64_t self_size = 0;
     uint64_t type_id = 0;
+    std::vector<Reference> references;
   };
 
   explicit HeapGraphTracker(TraceProcessorContext* context);
 
   void AddObject(UniquePid upid, int64_t ts, SourceObject obj);
   void AddInternedTypeName(uint64_t intern_id, StringPool::Id strid);
+  void AddInternedFieldName(uint64_t intern_id, StringPool::Id strid);
   void FinalizeProfile();
   void SetPacketIndex(uint64_t index);
 
@@ -53,6 +61,8 @@ class HeapGraphTracker {
   int64_t current_ts_ = 0;
   std::vector<SourceObject> current_objects_;
   std::map<uint64_t, StringPool::Id> interned_type_names_;
+  std::map<uint64_t, StringPool::Id> interned_field_names_;
+  std::map<uint64_t, int64_t> object_id_to_row_;
   uint64_t prev_index_ = 0;
 };
 
diff --git a/src/trace_processor/proto_trace_parser.cc b/src/trace_processor/proto_trace_parser.cc
index 5fbb11957..8bc2fc96e 100644
--- a/src/trace_processor/proto_trace_parser.cc
+++ b/src/trace_processor/proto_trace_parser.cc
@@ -2730,6 +2730,21 @@ void ProtoTraceParser::ParseHeapGraph(int64_t ts, ConstBytes blob) {
     obj.object_id = object.id();
     obj.self_size = object.self_size();
     obj.type_id = object.type_id();
+    auto ref_field_ids_it = object.reference_field_id();
+    auto ref_object_ids_it = object.reference_object_id();
+    for (; ref_field_ids_it && ref_object_ids_it;
+         ++ref_field_ids_it, ++ref_object_ids_it) {
+      HeapGraphTracker::SourceObject::Reference ref;
+      ref.field_name_id = ref_field_ids_it->as_uint64();
+      ref.owned_object_id = ref_object_ids_it->as_uint64();
+      obj.references.emplace_back(std::move(ref));
+    }
+
+    if (ref_field_ids_it || ref_object_ids_it) {
+      context_->storage->IncrementIndexedStats(stats::heap_graph_missing_packet,
+                                               static_cast<int>(upid));
+      continue;
+    }
     context_->heap_graph_tracker->AddObject(upid, ts, std::move(obj));
   }
   for (auto it = heap_graph.type_names(); it; ++it) {
@@ -2740,6 +2755,14 @@ void ProtoTraceParser::ParseHeapGraph(int64_t ts, ConstBytes blob) {
     context_->heap_graph_tracker->AddInternedTypeName(
         entry.iid(), context_->storage->InternString(str_view));
   }
+  for (auto it = heap_graph.field_names(); it; ++it) {
+    protos::pbzero::InternedString::Decoder entry(it->data(), it->size());
+    const char* str = reinterpret_cast<const char*>(entry.str().data);
+    auto str_view = base::StringView(str, entry.str().size);
+
+    context_->heap_graph_tracker->AddInternedFieldName(
+        entry.iid(), context_->storage->InternString(str_view));
+  }
   if (!heap_graph.continued()) {
     context_->heap_graph_tracker->FinalizeProfile();
   }
diff --git a/src/trace_processor/stats.h b/src/trace_processor/stats.h
index 8387f3727..74eec507f 100644
--- a/src/trace_processor/stats.h
+++ b/src/trace_processor/stats.h
@@ -110,6 +110,7 @@ namespace stats {
   F(json_tokenizer_failure,                   kSingle,  kError,    kTrace),    \
   F(heap_graph_invalid_string_id,             kIndexed, kError,    kTrace),    \
   F(heap_graph_non_finalized_graph,           kSingle,  kError,    kTrace),    \
+  F(heap_graph_malformed_packet,              kIndexed, kError,    kTrace),    \
   F(heap_graph_missing_packet,                kIndexed, kDataLoss, kTrace),    \
   F(heapprofd_buffer_corrupted,               kIndexed, kError,    kTrace),    \
   F(heapprofd_buffer_overran,                 kIndexed, kDataLoss, kTrace),    \
diff --git a/src/trace_processor/tables/profiler_tables.h b/src/trace_processor/tables/profiler_tables.h
index bf70637e4..16133d8eb 100644
--- a/src/trace_processor/tables/profiler_tables.h
+++ b/src/trace_processor/tables/profiler_tables.h
@@ -40,10 +40,21 @@ PERFETTO_TP_TABLE(PERFETTO_TP_SYMBOL_DEF);
   C(int64_t, graph_sample_ts)                              \
   C(int64_t, object_id)                                    \
   C(int64_t, self_size)                                    \
+  C(int64_t, reference_set_id)                             \
   C(StringPool::Id, type_name)
 
 PERFETTO_TP_TABLE(PERFETTO_TP_HEAP_GRAPH_OBJECT_DEF);
 
+#define PERFETTO_TP_HEAP_GRAPH_REFERENCE_DEF(NAME, PARENT, C) \
+  NAME(HeapGraphReferenceTable, "heap_graph_reference")       \
+  PERFETTO_TP_ROOT_TABLE(PARENT, C)                           \
+  C(int64_t, reference_set_id)                                \
+  C(int64_t, owner_id)                                        \
+  C(int64_t, owned_id)                                        \
+  C(StringPool::Id, field_name)
+
+PERFETTO_TP_TABLE(PERFETTO_TP_HEAP_GRAPH_REFERENCE_DEF);
+
 #define PERFETTO_TP_VULKAN_MEMORY_ALLOCATIONS_DEF(NAME, PARENT, C) \
   NAME(VulkanMemoryAllocationsTable, "vulkan_memory_allocations")  \
   PERFETTO_TP_ROOT_TABLE(PARENT, C)                                \
diff --git a/src/trace_processor/trace_processor_impl.cc b/src/trace_processor/trace_processor_impl.cc
index 2086b7492..1492cf7b3 100644
--- a/src/trace_processor/trace_processor_impl.cc
+++ b/src/trace_processor/trace_processor_impl.cc
@@ -376,6 +376,9 @@ TraceProcessorImpl::TraceProcessorImpl(const Config& cfg) {
   DbSqliteTable::RegisterTable(*db_, &storage->heap_graph_object_table(),
                                storage->heap_graph_object_table().table_name());
   DbSqliteTable::RegisterTable(
+      *db_, &storage->heap_graph_reference_table(),
+      storage->heap_graph_reference_table().table_name());
+  DbSqliteTable::RegisterTable(
       *db_, &storage->vulkan_memory_allocations_table(),
       storage->vulkan_memory_allocations_table().table_name());
 }
diff --git a/src/trace_processor/trace_storage.h b/src/trace_processor/trace_storage.h
index 536fcd528..2f6f77e13 100644
--- a/src/trace_processor/trace_storage.h
+++ b/src/trace_processor/trace_storage.h
@@ -1284,6 +1284,14 @@ class TraceStorage {
     return &heap_graph_object_table_;
   }
 
+  const tables::HeapGraphReferenceTable& heap_graph_reference_table() const {
+    return heap_graph_reference_table_;
+  }
+
+  tables::HeapGraphReferenceTable* mutable_heap_graph_reference_table() {
+    return &heap_graph_reference_table_;
+  }
+
   const tables::GpuTrackTable& gpu_track_table() const {
     return gpu_track_table_;
   }
@@ -1405,6 +1413,8 @@ class TraceStorage {
   // Symbol tables (mappings from frames to symbol names)
   tables::SymbolTable symbol_table_{&string_pool_, nullptr};
   tables::HeapGraphObjectTable heap_graph_object_table_{&string_pool_, nullptr};
+  tables::HeapGraphReferenceTable heap_graph_reference_table_{&string_pool_,
+                                                              nullptr};
 
   tables::VulkanMemoryAllocationsTable vulkan_memory_allocations_table_{
       &string_pool_, nullptr};