perfetto_hprof: be smarter about splitting packets

We now keep track of how much we have already written.

Test: grab profile and look at in ui.perfetto.dev.

Bug: 149378918
Bug: 143874090
Change-Id: I76c491fe3738611d012294169ee7c77541a6d8cb
diff --git a/perfetto_hprof/perfetto_hprof.cc b/perfetto_hprof/perfetto_hprof.cc
index 0820aa6..efd105b 100644
--- a/perfetto_hprof/perfetto_hprof.cc
+++ b/perfetto_hprof/perfetto_hprof.cc
@@ -59,7 +59,11 @@
 
 constexpr int kJavaHeapprofdSignal = __SIGRTMIN + 6;
 constexpr time_t kWatchdogTimeoutSec = 120;
-constexpr size_t kObjectsPerPacket = 100;
+// This needs to be lower than the maximum acceptable chunk size, because this
+// is checked *before* writing another submessage. We conservatively assume
+// submessages can be up to 100k here for a 500k chunk size.
+// DropBox has a 500k chunk limit, and each chunk needs to parse as a proto.
+constexpr uint32_t kPacketSizeThreshold = 400000;
 constexpr char kByte[1] = {'x'};
 static art::Mutex& GetStateMutex() {
   static art::Mutex state_mutex("perfetto_hprof_state_mutex", art::LockLevel::kGenericBottomLock);
@@ -221,24 +225,38 @@
 class Writer {
  public:
   Writer(pid_t parent_pid, JavaHprofDataSource::TraceContext* ctx, uint64_t timestamp)
-      : parent_pid_(parent_pid), ctx_(ctx), timestamp_(timestamp) {}
+      : parent_pid_(parent_pid), ctx_(ctx), timestamp_(timestamp),
+        last_written_(ctx_->written()) {}
+
+  // Return whether the next call to GetHeapGraph will create a new TracePacket.
+  bool will_create_new_packet() {
+    return !heap_graph_ || ctx_->written() - last_written_ > kPacketSizeThreshold;
+  }
 
   perfetto::protos::pbzero::HeapGraph* GetHeapGraph() {
-    if (!heap_graph_ || ++objects_written_ % kObjectsPerPacket == 0) {
-      if (heap_graph_) {
-        heap_graph_->set_continued(true);
-      }
-      Finalize();
-
-      trace_packet_ = ctx_->NewTracePacket();
-      trace_packet_->set_timestamp(timestamp_);
-      heap_graph_ = trace_packet_->set_heap_graph();
-      heap_graph_->set_pid(parent_pid_);
-      heap_graph_->set_index(index_++);
+    if (will_create_new_packet()) {
+      CreateNewHeapGraph();
     }
     return heap_graph_;
   }
 
+  void CreateNewHeapGraph() {
+    if (heap_graph_) {
+      heap_graph_->set_continued(true);
+    }
+    Finalize();
+
+    uint64_t written = ctx_->written();
+
+    trace_packet_ = ctx_->NewTracePacket();
+    trace_packet_->set_timestamp(timestamp_);
+    heap_graph_ = trace_packet_->set_heap_graph();
+    heap_graph_->set_pid(parent_pid_);
+    heap_graph_->set_index(index_++);
+
+    last_written_ = written;
+  }
+
   void Finalize() {
     if (trace_packet_) {
       trace_packet_->Finalize();
@@ -253,12 +271,13 @@
   JavaHprofDataSource::TraceContext* const ctx_;
   const uint64_t timestamp_;
 
+  uint64_t last_written_ = 0;
+
   perfetto::DataSource<JavaHprofDataSource>::TraceContext::TracePacketHandle
       trace_packet_;
   perfetto::protos::pbzero::HeapGraph* heap_graph_ = nullptr;
 
   uint64_t index_ = 0;
-  size_t objects_written_ = 0;
 };
 
 class ReferredObjectsFinder {
@@ -429,8 +448,15 @@
               perfetto::protos::pbzero::HeapGraphRoot* root_proto =
                 writer.GetHeapGraph()->add_roots();
               root_proto->set_root_type(ToProtoType(root_type));
-              for (art::mirror::Object* obj : children)
+              for (art::mirror::Object* obj : children) {
+                if (writer.will_create_new_packet()) {
+                  root_proto->set_object_ids(*object_ids);
+                  object_ids->Reset();
+                  root_proto = writer.GetHeapGraph()->add_roots();
+                  root_proto->set_root_type(ToProtoType(root_type));
+                }
                 object_ids->Append(reinterpret_cast<uintptr_t>(obj));
+              }
               root_proto->set_object_ids(*object_ids);
               object_ids->Reset();
             }