Ahat plugin prototype v2

Android.bp

@@ -16246,6 +16246,7 @@ filegroup {
     name: "perfetto_src_trace_processor_perfetto_sql_intrinsics_functions_functions",
     srcs: [
         "src/trace_processor/perfetto_sql/intrinsics/functions/args.cc",
+        "src/trace_processor/perfetto_sql/intrinsics/functions/art_heap_graph_functions.cc",
         "src/trace_processor/perfetto_sql/intrinsics/functions/base64.cc",
         "src/trace_processor/perfetto_sql/intrinsics/functions/counter_intervals.cc",
         "src/trace_processor/perfetto_sql/intrinsics/functions/create_function.cc",

BUILD

@@ -3278,6 +3278,8 @@ perfetto_filegroup(
     srcs = [
         "src/trace_processor/perfetto_sql/intrinsics/functions/args.cc",
         "src/trace_processor/perfetto_sql/intrinsics/functions/args.h",
+        "src/trace_processor/perfetto_sql/intrinsics/functions/art_heap_graph_functions.cc",
+        "src/trace_processor/perfetto_sql/intrinsics/functions/art_heap_graph_functions.h",
         "src/trace_processor/perfetto_sql/intrinsics/functions/base64.cc",
         "src/trace_processor/perfetto_sql/intrinsics/functions/base64.h",
         "src/trace_processor/perfetto_sql/intrinsics/functions/clock_functions.h",

src/trace_processor/importers/art_hprof/art_heap_graph.h

@@ -48,6 +48,16 @@ class HeapGraph {
     strings_[id] = interned_id;
   }
 
+  void SetObjects(base::FlatHashMap<uint64_t, Object> objs) {
+    objects_ = std::move(objs);
+  }
+  void SetClasses(base::FlatHashMap<uint64_t, ClassDefinition> cls) {
+    classes_ = std::move(cls);
+  }
+  void SetStrings(base::FlatHashMap<uint64_t, StringId> strs) {
+    strings_ = std::move(strs);
+  }
+
   const base::FlatHashMap<uint64_t, Object>& GetObjects() const {
     return objects_;
   }
@@ -61,6 +71,13 @@ class HeapGraph {
   size_t GetStringCount() const { return strings_.size(); }
   uint64_t GetTimestamp() const { return timestamp_; }
 
+  void ClearAll() {
+    objects_.Clear();
+    classes_.Clear();
+    strings_.Clear();
+    heap_id_to_name_.Clear();
+  }
+
   static std::string GetRootTypeName(HprofHeapRootTag root_type_id) {
     switch (root_type_id) {
       case HprofHeapRootTag::kJniGlobal:

src/trace_processor/importers/art_hprof/art_heap_graph_builder.cc

@@ -18,6 +18,7 @@
 #include <cinttypes>
 
 namespace perfetto::trace_processor::art_hprof {
+
 namespace {
 
 // Root type precedence ranking. Lower rank = higher priority.
@@ -85,25 +86,16 @@ void HeapGraphBuilder::PushBlob(TraceBlobView&& blob) {
 }
 
 HeapGraph HeapGraphBuilder::BuildGraph() {
-  // Phase 3: Resolve the heap graph
   resolver_ = std::make_unique<HeapGraphResolver>(context_, header_, objects_,
                                                   classes_, roots_, stats_);
   resolver_->ResolveGraph();
 
   stats_.Write(context_);
   HeapGraph graph(header_.GetTimestamp());
 
-  for (auto it = strings_.GetIterator(); it; ++it) {
-    graph.AddString(it.key(), it.value());
-  }
-
-  for (auto it = classes_.GetIterator(); it; ++it) {
-    graph.AddClass(std::move(it.value()));
-  }
-
-  for (auto it = objects_.GetIterator(); it; ++it) {
-    graph.AddObject(std::move(it.value()));
-  }
+  graph.SetStrings(std::move(strings_));
+  graph.SetClasses(std::move(classes_));
+  graph.SetObjects(std::move(objects_));
 
   return graph;
 }
@@ -203,6 +195,10 @@ bool HeapGraphBuilder::ParseUtf8StringRecord(uint32_t length) {
     return false;
   }
 
+  if (length < header_.GetIdSize()) {
+    return false;
+  }
+
   std::string str;
   if (!iterator_->ReadString(str, length - header_.GetIdSize())) {
     return false;
@@ -663,8 +659,9 @@ bool HeapGraphBuilder::ParsePrimitiveArrayObject() {
 
   size_t type_size = GetFieldTypeSize(element_type, header_.GetIdSize());
 
+  size_t data_length = static_cast<size_t>(element_count) * type_size;
   std::vector<uint8_t> data;
-  if (!iterator_->ReadBytes(data, element_count * type_size)) {
+  if (!iterator_->ReadBytes(data, data_length)) {
     return false;
   }
 

src/trace_processor/importers/art_hprof/art_heap_graph_builder.h

@@ -119,6 +119,7 @@ class HeapGraphResolver {
   void ExtractFieldValues(Object& obj, const ClassDefinition& cls);
   void ExtractPrimitiveArrayValues(Object& obj);
   std::optional<std::string> DecodeJavaString(const Object& string_obj) const;
+  void DecodeJavaStrings();
   const std::vector<Field>& GetClassHierarchyFields(uint64_t class_id);
   void ComputeSelfSizes();
   void CalculateNativeSizes();

src/trace_processor/importers/art_hprof/art_heap_graph_resolver.cc

@@ -14,7 +14,10 @@
  * limitations under the License.
  */
 
+#include <algorithm>
 #include <deque>
+#include <string>
+#include <variant>
 
 #include "perfetto/ext/base/flat_hash_map.h"
 #include "src/trace_processor/importers/art_hprof/art_heap_graph_builder.h"
@@ -79,36 +82,27 @@ HeapGraphResolver::HeapGraphResolver(
 
 void HeapGraphResolver::ResolveGraph() {
   ExtractAllObjectData();
+  field_cache_.Clear();
+  DecodeJavaStrings();
   MarkReachableObjects();
   ComputeSelfSizes();
   CalculateNativeSizes();
 }
 
 void HeapGraphResolver::ExtractAllObjectData() {
-  // Identify classes whose instances need field extraction (for native size
-  // calculation and string decoding). All other instances only need references.
-  base::FlatHashMap<uint64_t, bool> needs_field_extraction;
-  for (auto it = classes_.GetIterator(); it; ++it) {
-    const auto& name = it.value().GetName();
-    if (name == "libcore.util.NativeAllocationRegistry" ||
-        name == kJavaLangString) {
-      needs_field_extraction[it.key()] = true;
-    }
-  }
-
   for (auto it = objects_.GetIterator(); it; ++it) {
     auto& obj = it.value();
     if (obj.GetObjectType() == ObjectType::kInstance ||
         obj.GetObjectType() == ObjectType::kClass) {
       auto cls = classes_.Find(obj.GetClassId());
       if (cls) {
         ExtractObjectReferences(obj, *cls);
-        if (needs_field_extraction.Find(obj.GetClassId())) {
-          ExtractFieldValues(obj, *cls);
-        }
+        ExtractFieldValues(obj, *cls);
       }
     } else if (obj.GetObjectType() == ObjectType::kObjectArray) {
       ExtractArrayElementReferences(obj);
+    } else if (obj.GetObjectType() == ObjectType::kPrimitiveArray) {
+      ExtractPrimitiveArrayValues(obj);
     }
 
     uint64_t obj_id = obj.GetId();
@@ -338,6 +332,8 @@ void HeapGraphResolver::ExtractFieldValues(Object& obj,
 
     obj.AddField(std::move(field));
   }
+
+  obj.ClearRawData();
 }
 
 void HeapGraphResolver::ExtractPrimitiveArrayValues(Object& obj) {
@@ -451,7 +447,7 @@ std::optional<std::string> HeapGraphResolver::DecodeJavaString(
   int32_t count = count_opt.value_or(static_cast<int32_t>(array_len) - offset);
 
   if (offset < 0 || count < 0 ||
-      static_cast<size_t>(offset + count) > array_len)
+      static_cast<size_t>(offset) + static_cast<size_t>(count) > array_len)
     return std::nullopt;
 
   std::string result;
@@ -470,24 +466,45 @@ std::optional<std::string> HeapGraphResolver::DecodeJavaString(
     }
   };
 
-  const auto type = array->GetArrayElementType();
+  const auto& array_data = array->GetArrayDataVariant();
 
-  if (type == FieldType::kByte) {
-    const auto& bytes = array->GetArrayData<uint8_t>();
+  if (const auto* bytes = std::get_if<std::vector<uint8_t>>(&array_data)) {
     for (int32_t i = 0; i < count; ++i)
       result.push_back(
-          static_cast<char>(bytes[static_cast<size_t>(offset + i)]));
+          static_cast<char>((*bytes)[static_cast<size_t>(offset + i)]));
     return result;
-  } else if (type == FieldType::kChar) {
-    const auto& chars = array->GetArrayData<uint16_t>();
+  }
+  if (const auto* chars = std::get_if<std::vector<uint16_t>>(&array_data)) {
     for (int32_t i = 0; i < count; ++i)
-      append_utf8_from_utf16(chars[static_cast<size_t>(offset + i)]);
+      append_utf8_from_utf16((*chars)[static_cast<size_t>(offset + i)]);
     return result;
   }
 
   return std::nullopt;
 }
 
+void HeapGraphResolver::DecodeJavaStrings() {
+  uint64_t string_class_id = 0;
+  for (auto it = classes_.GetIterator(); it; ++it) {
+    if (it.value().GetName() == kJavaLangString) {
+      string_class_id = it.key();
+      break;
+    }
+  }
+  if (string_class_id == 0)
+    return;
+
+  for (auto it = objects_.GetIterator(); it; ++it) {
+    auto& obj = it.value();
+    if (obj.GetClassId() != string_class_id)
+      continue;
+    auto decoded = DecodeJavaString(obj);
+    if (decoded) {
+      obj.SetDecodedString(std::move(*decoded));
+    }
+  }
+}
+
 const std::vector<Field>& HeapGraphResolver::GetClassHierarchyFields(
     uint64_t class_id) {
   auto* cached = field_cache_.Find(class_id);

src/trace_processor/importers/art_hprof/art_hprof_model.h

@@ -103,16 +103,10 @@ class Field {
         value_);
   }
 
-  void SetDecodedString(std::string str) { decoded_string_ = std::move(str); }
-  std::optional<std::string> GetDecodedString() const {
-    return decoded_string_;
-  }
-
  private:
   std::string name_;
   FieldType type_;
   ValueType value_ = std::monostate{};
-  std::optional<std::string> decoded_string_;
 };
 
 struct Reference {
@@ -205,6 +199,16 @@ class Object {
   void SetRawData(std::vector<uint8_t> data) { raw_data_ = std::move(data); }
 
   const std::vector<uint8_t>& GetRawData() const { return raw_data_; }
+  void ClearRawData() { raw_data_ = {}; }
+
+  void ClearParsedData() {
+    fields_ = {};
+    references_ = {};
+    pending_references_ = {};
+    array_elements_ = {};
+    array_data_ = std::monostate{};
+    decoded_string_.reset();
+  }
 
   void AddReference(std::string_view field_name,
                     std::optional<uint64_t> field_class_id,
@@ -254,6 +258,11 @@ class Object {
 
   void AddNativeSize(int64_t size) { native_size_ += size; }
 
+  void SetDecodedString(std::string str) { decoded_string_ = std::move(str); }
+  const std::optional<std::string>& GetDecodedString() const {
+    return decoded_string_;
+  }
+
   void SetSelfSizeOverride(size_t size) { self_size_override_ = size; }
   std::optional<size_t> GetSelfSizeOverride() const {
     return self_size_override_;
@@ -277,6 +286,8 @@ class Object {
     return {};
   }
 
+  const ArrayData& GetArrayDataVariant() const { return array_data_; }
+
   size_t GetArrayElementCount() const {
     return std::visit(
         [](const auto& val) -> size_t {
@@ -308,6 +319,7 @@ class Object {
 
   int64_t native_size_ = 0;
   std::optional<size_t> self_size_override_;
+  std::optional<std::string> decoded_string_;
 
   // Field values
   std::vector<Field> fields_;