tp: introduce ForgedTracePacketWriter for writing TracePackets efficiently

This CL introduces a new class to write TracePacket protos efficiently
and get out TraceBlobViews without creating tons of copies. There were
already N instances of this in the codebase and we're now introducing a
new one in #4768.

Instead now we have a single class which works with TraceBlob +
TraceBlobView and is stored in TraceProcessorContext so everyone can use
it. It works in 4MB chunks and reuses them across TracePacket calls
which makes it really efficient now to create these packets with 99% of
cases being totally zero-copy.

Author: LalitMaganti

Android.bp

@@ -15707,6 +15707,7 @@ filegroup {
         "src/trace_processor/importers/proto/chrome_system_probes_module.cc",
         "src/trace_processor/importers/proto/chrome_system_probes_parser.cc",
         "src/trace_processor/importers/proto/default_modules.cc",
+        "src/trace_processor/importers/proto/forged_packet_writer.cc",
         "src/trace_processor/importers/proto/memory_tracker_snapshot_module.cc",
         "src/trace_processor/importers/proto/memory_tracker_snapshot_parser.cc",
         "src/trace_processor/importers/proto/metadata_minimal_module.cc",
@@ -15747,6 +15748,7 @@ filegroup {
     srcs: [
         "src/trace_processor/importers/proto/active_chrome_processes_tracker_unittest.cc",
         "src/trace_processor/importers/proto/chrome_string_lookup_unittest.cc",
+        "src/trace_processor/importers/proto/forged_packet_writer_unittest.cc",
         "src/trace_processor/importers/proto/heap_graph_tracker_unittest.cc",
         "src/trace_processor/importers/proto/jit_tracker_unittest.cc",
         "src/trace_processor/importers/proto/network_trace_module_unittest.cc",

BUILD

@@ -2810,6 +2810,8 @@ perfetto_filegroup(
         "src/trace_processor/importers/proto/chrome_system_probes_parser.h",
         "src/trace_processor/importers/proto/default_modules.cc",
         "src/trace_processor/importers/proto/default_modules.h",
+        "src/trace_processor/importers/proto/forged_packet_writer.cc",
+        "src/trace_processor/importers/proto/forged_packet_writer.h",
         "src/trace_processor/importers/proto/memory_tracker_snapshot_module.cc",
         "src/trace_processor/importers/proto/memory_tracker_snapshot_module.h",
         "src/trace_processor/importers/proto/memory_tracker_snapshot_parser.cc",

src/trace_processor/importers/proto/BUILD.gn

@@ -30,6 +30,8 @@ source_set("minimal") {
     "chrome_system_probes_parser.h",
     "default_modules.cc",
     "default_modules.h",
+    "forged_packet_writer.cc",
+    "forged_packet_writer.h",
     "memory_tracker_snapshot_module.cc",
     "memory_tracker_snapshot_module.h",
     "memory_tracker_snapshot_parser.cc",
@@ -299,6 +301,7 @@ source_set("unittests") {
   sources = [
     "active_chrome_processes_tracker_unittest.cc",
     "chrome_string_lookup_unittest.cc",
+    "forged_packet_writer_unittest.cc",
     "heap_graph_tracker_unittest.cc",
     "jit_tracker_unittest.cc",
     "network_trace_module_unittest.cc",

src/trace_processor/importers/proto/android_probes_module.cc

@@ -26,15 +26,14 @@
 #include "perfetto/base/logging.h"
 #include "perfetto/ext/base/string_view.h"
 #include "perfetto/protozero/field.h"
-#include "perfetto/protozero/scattered_heap_buffer.h"
 #include "perfetto/trace_processor/ref_counted.h"
-#include "perfetto/trace_processor/trace_blob.h"
 #include "protos/perfetto/common/builtin_clock.pbzero.h"
 #include "src/trace_processor/importers/common/clock_tracker.h"
 #include "src/trace_processor/importers/common/event_tracker.h"
 #include "src/trace_processor/importers/common/parser_types.h"
 #include "src/trace_processor/importers/proto/android_probes_parser.h"
 #include "src/trace_processor/importers/proto/android_probes_tracker.h"
+#include "src/trace_processor/importers/proto/forged_packet_writer.h"
 #include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/sorter/trace_sorter.h"
@@ -139,20 +138,20 @@ ModuleResult AndroidProbesModule::TokenizePacket(
         actual_ts = packet_timestamp;
       }
 
-      protozero::HeapBuffered<protos::pbzero::TracePacket> data_packet;
-      // Keep the original timestamp to later extract as an arg; the sorter does
-      // not read this.
-      data_packet->set_timestamp(static_cast<uint64_t>(packet_timestamp));
-
-      auto* power_rails_proto = data_packet->set_power_rails();
-      power_rails_proto->set_session_uuid(evt.session_uuid());
-      auto* energy = power_rails_proto->add_energy_data();
-      energy->set_energy(data.energy());
-      energy->set_index(data.index());
-      energy->set_timestamp_ms(static_cast<uint64_t>(actual_ts / 1000000));
-
-      auto [vec, size] = data_packet.SerializeAsUniquePtr();
-      TraceBlobView tbv(TraceBlob::TakeOwnership(std::move(vec), size));
+      TraceBlobView tbv =
+          context_->forged_packet_writer->WritePacket([&](auto* data_packet) {
+            // Keep the original timestamp to later extract as an arg; the
+            // sorter does not read this.
+            data_packet->set_timestamp(static_cast<uint64_t>(packet_timestamp));
+
+            auto* power_rails_proto = data_packet->set_power_rails();
+            power_rails_proto->set_session_uuid(evt.session_uuid());
+            auto* energy = power_rails_proto->add_energy_data();
+            energy->set_energy(data.energy());
+            energy->set_index(data.index());
+            energy->set_timestamp_ms(
+                static_cast<uint64_t>(actual_ts / 1000000));
+          });
       module_context_->trace_packet_stream->Push(
           actual_ts, TracePacketData{std::move(tbv), state});
     }
@@ -176,36 +175,36 @@ ModuleResult AndroidProbesModule::TokenizePacket(
       }
       int64_t actual_ts = *trace_ts;
 
-      protozero::HeapBuffered<protos::pbzero::TracePacket> data_packet;
-      data_packet->set_timestamp(static_cast<uint64_t>(actual_ts));
-
-      auto* log_pkt = data_packet->set_android_log();
-      auto* log_evt = log_pkt->add_events();
-      log_evt->set_log_id(
-          static_cast<protos::pbzero::AndroidLogId>(evt.log_id()));
-      log_evt->set_pid(evt.pid());
-      log_evt->set_tid(evt.tid());
-      log_evt->set_uid(evt.uid());
-      log_evt->set_timestamp(evt.timestamp());
-      log_evt->set_tag(evt.tag());
-      log_evt->set_prio(
-          static_cast<protos::pbzero::AndroidLogPriority>(evt.prio()));
-      log_evt->set_message(evt.message());
-      for (auto arg_it = evt.args(); arg_it; ++arg_it) {
-        protos::pbzero::AndroidLogPacket::LogEvent::Arg::Decoder arg(*arg_it);
-        auto* new_arg = log_evt->add_args();
-        new_arg->set_name(arg.name());
-        if (arg.has_int_value()) {
-          new_arg->set_int_value(arg.int_value());
-        } else if (arg.has_float_value()) {
-          new_arg->set_float_value(arg.float_value());
-        } else if (arg.has_string_value()) {
-          new_arg->set_string_value(arg.string_value());
-        }
-      }
-
-      auto [vec, size] = data_packet.SerializeAsUniquePtr();
-      TraceBlobView tbv(TraceBlob::TakeOwnership(std::move(vec), size));
+      TraceBlobView tbv =
+          context_->forged_packet_writer->WritePacket([&](auto* data_packet) {
+            data_packet->set_timestamp(static_cast<uint64_t>(actual_ts));
+
+            auto* log_pkt = data_packet->set_android_log();
+            auto* log_evt = log_pkt->add_events();
+            log_evt->set_log_id(
+                static_cast<protos::pbzero::AndroidLogId>(evt.log_id()));
+            log_evt->set_pid(evt.pid());
+            log_evt->set_tid(evt.tid());
+            log_evt->set_uid(evt.uid());
+            log_evt->set_timestamp(evt.timestamp());
+            log_evt->set_tag(evt.tag());
+            log_evt->set_prio(
+                static_cast<protos::pbzero::AndroidLogPriority>(evt.prio()));
+            log_evt->set_message(evt.message());
+            for (auto arg_it = evt.args(); arg_it; ++arg_it) {
+              protos::pbzero::AndroidLogPacket::LogEvent::Arg::Decoder arg(
+                  *arg_it);
+              auto* new_arg = log_evt->add_args();
+              new_arg->set_name(arg.name());
+              if (arg.has_int_value()) {
+                new_arg->set_int_value(arg.int_value());
+              } else if (arg.has_float_value()) {
+                new_arg->set_float_value(arg.float_value());
+              } else if (arg.has_string_value()) {
+                new_arg->set_string_value(arg.string_value());
+              }
+            }
+          });
       module_context_->trace_packet_stream->Push(
           actual_ts, TracePacketData{std::move(tbv), state});
     }

src/trace_processor/importers/proto/app_wakelock_module.cc

@@ -25,14 +25,14 @@
 #include "perfetto/ext/base/murmur_hash.h"
 #include "perfetto/protozero/field.h"
 #include "perfetto/trace_processor/ref_counted.h"
-#include "perfetto/trace_processor/trace_blob.h"
 #include "protos/perfetto/trace/android/app_wakelock_data.pbzero.h"
 #include "protos/perfetto/trace/interned_data/interned_data.pbzero.h"
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
 #include "src/trace_processor/importers/common/parser_types.h"
 #include "src/trace_processor/importers/common/slice_tracker.h"
 #include "src/trace_processor/importers/common/track_compressor.h"
 #include "src/trace_processor/importers/common/tracks.h"
+#include "src/trace_processor/importers/proto/forged_packet_writer.h"
 #include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/sorter/trace_sorter.h"
@@ -93,12 +93,16 @@ ModuleResult AppWakelockModule::TokenizePacket(
       continue;
     }
 
-    packet_buffer_->set_timestamp(static_cast<uint64_t>(real_ts));
-    auto* event = packet_buffer_->set_app_wakelock_bundle();
-    size_t length = static_cast<size_t>(interned->end() - interned->begin());
-    event->set_info()->AppendRawProtoBytes(interned->begin(), length);
-    event->set_acquired(acquired);
-    PushPacketBufferForSort(real_ts, state);
+    TraceBlobView tbv =
+        context_->forged_packet_writer->WritePacket([&](auto* pkt) {
+          pkt->set_timestamp(static_cast<uint64_t>(real_ts));
+          auto* event = pkt->set_app_wakelock_bundle();
+          auto length =
+              static_cast<size_t>(interned->end() - interned->begin());
+          event->set_info()->AppendRawProtoBytes(interned->begin(), length);
+          event->set_acquired(acquired);
+        });
+    PushPacketBufferForSort(real_ts, std::move(tbv), state);
   }
 
   return ModuleResult::Handled();
@@ -159,12 +163,10 @@ void AppWakelockModule::ParseWakelockBundle(int64_t ts, ConstBytes blob) {
 
 void AppWakelockModule::PushPacketBufferForSort(
     int64_t timestamp,
+    TraceBlobView tbv,
     RefPtr<PacketSequenceStateGeneration> state) {
-  auto [vec, size] = packet_buffer_.SerializeAsUniquePtr();
-  TraceBlobView tbv(TraceBlob::TakeOwnership(std::move(vec), size));
   module_context_->trace_packet_stream->Push(
       timestamp, TracePacketData{std::move(tbv), std::move(state)});
-  packet_buffer_.Reset();
 }
 
 }  // namespace perfetto::trace_processor

src/trace_processor/importers/proto/app_wakelock_module.h

@@ -20,8 +20,8 @@
 #include <cstdint>
 
 #include "perfetto/protozero/field.h"
-#include "perfetto/protozero/scattered_heap_buffer.h"
 #include "perfetto/trace_processor/ref_counted.h"
+#include "perfetto/trace_processor/trace_blob_view.h"
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
 #include "src/trace_processor/importers/common/parser_types.h"
 #include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
@@ -54,13 +54,12 @@ class AppWakelockModule : public ProtoImporterModule {
  private:
   void ParseWakelockBundle(int64_t ts, protozero::ConstBytes blob);
 
-  // Helper to simplify pushing a TracePacket to the sorter. The caller fills in
-  // the packet buffer and uses this to push for sorting and reset the buffer.
+  // Pushes a serialized TracePacket to the sorter.
   void PushPacketBufferForSort(int64_t timestamp,
+                               TraceBlobView tbv,
                                RefPtr<PacketSequenceStateGeneration> state);
 
   TraceProcessorContext* context_;
-  protozero::HeapBuffered<protos::pbzero::TracePacket> packet_buffer_;
 
   const StringId arg_flags_;
   const StringId arg_owner_pid_;

src/trace_processor/importers/proto/forged_packet_writer.cc

@@ -0,0 +1,112 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "src/trace_processor/importers/proto/forged_packet_writer.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <utility>
+
+#include "perfetto/base/logging.h"
+#include "perfetto/protozero/contiguous_memory_range.h"
+#include "perfetto/protozero/scattered_stream_writer.h"
+#include "perfetto/public/compiler.h"
+#include "perfetto/trace_processor/ref_counted.h"
+#include "perfetto/trace_processor/trace_blob.h"
+#include "perfetto/trace_processor/trace_blob_view.h"
+
+#include "protos/perfetto/trace/trace_packet.pbzero.h"
+
+namespace perfetto::trace_processor {
+
+ForgedTracePacketWriter::ForgedTracePacketWriter() : writer_(this) {}
+
+ForgedTracePacketWriter::~ForgedTracePacketWriter() = default;
+
+protos::pbzero::TracePacket* ForgedTracePacketWriter::BeginPacket() {
+  if (!slab_ || packet_start_ptr_ >= slab_->data() + slab_->size()) {
+    slab_.reset(new TraceBlob(TraceBlob::Allocate(kSlabSize)));
+    packet_start_ptr_ = slab_->data();
+  }
+
+  protozero::ContiguousMemoryRange range{packet_start_ptr_,
+                                         slab_->data() + slab_->size()};
+  msg_.Reset(&writer_);
+  writer_.Reset(range);
+  slices_.push_back(range);
+  return &msg_;
+}
+
+TraceBlobView ForgedTracePacketWriter::EndPacket() {
+  PERFETTO_CHECK(slices_.size() > 0);
+  msg_.Finalize();
+
+  // Close the last slice with the actual end position of the packet and
+  // keep track of the current offset for the next packet.
+  {
+    auto& s = slices_.back();
+    s.end = writer_.write_ptr();
+    packet_start_ptr_ = s.end;
+  }
+
+  // Common case: packet fits in the single slab. Zero copies.
+  if (PERFETTO_LIKELY(slices_.size() == 1)) {
+    auto s = slices_.back();
+    slices_.clear();
+    return {slab_, static_cast<size_t>(s.begin - slab_->data()), s.size()};
+  }
+
+  PERFETTO_CHECK(overflow_slabs_.size() > 0);
+
+  // Rare: packet spans multiple slabs. Stitch into one contiguous blob.
+  size_t total = 0;
+  for (const auto& s : slices_) {
+    total += s.size();
+  }
+
+  TraceBlob stitched = TraceBlob::Allocate(total);
+  {
+    uint8_t* dst = stitched.data();
+    for (const auto& s : slices_) {
+      memcpy(dst, s.begin, s.size());
+      dst += s.size();
+    }
+  }
+
+  slab_ = std::move(overflow_slabs_.back());
+  overflow_slabs_.clear();
+  slices_.clear();
+  return TraceBlobView(std::move(stitched), 0, total);
+}
+
+protozero::ContiguousMemoryRange ForgedTracePacketWriter::GetNewBuffer() {
+  PERFETTO_CHECK(slices_.size() > 0);
+
+  // Close the current slice and start a new one for the new slab.
+  slices_.back().end = writer_.write_ptr();
+
+  // Allocate a new slab and add it to the overflow slabs.
+  overflow_slabs_.emplace_back(new TraceBlob(TraceBlob::Allocate(kSlabSize)));
+
+  auto& blob = *overflow_slabs_.back();
+  protozero::ContiguousMemoryRange range{blob.data(),
+                                         blob.data() + blob.size()};
+  slices_.push_back(range);
+  return range;
+}
+
+}  // namespace perfetto::trace_processor