Merge "Add helper class for heap-based protozero messages"

6 files changed, 70 insertions, 43 deletions

diff --git a/include/perfetto/protozero/field.h b/include/perfetto/protozero/field.h
index cbb48e6ca..9c6da8d07 100644
--- a/include/perfetto/protozero/field.h
+++ b/include/perfetto/protozero/field.h
@@ -19,6 +19,8 @@
 
 #include <stdint.h>
 
+#include <string>
+
 #include "perfetto/base/logging.h"
 #include "perfetto/protozero/contiguous_memory_range.h"
 #include "perfetto/protozero/proto_utils.h"
@@ -34,6 +36,7 @@ struct ConstChars {
   // Allow implicit conversion to perfetto's base::StringView without depending
   // on perfetto/base or viceversa.
   static constexpr bool kConvertibleToStringView = true;
+  std::string ToStdString() const { return std::string(data, size); }
 
   const char* data;
   size_t size;
@@ -112,6 +115,8 @@ class Field {
     return ConstChars{reinterpret_cast<const char*>(data()), size_};
   }
 
+  inline std::string as_std_string() const { return as_string().ToStdString(); }
+
   inline ConstBytes as_bytes() const {
     PERFETTO_DCHECK(!valid() ||
                     type() == proto_utils::ProtoWireType::kLengthDelimited);
diff --git a/include/perfetto/protozero/scattered_heap_buffer.h b/include/perfetto/protozero/scattered_heap_buffer.h
index d88f7e63b..7f3906ed9 100644
--- a/include/perfetto/protozero/scattered_heap_buffer.h
+++ b/include/perfetto/protozero/scattered_heap_buffer.h
@@ -18,6 +18,7 @@
 #define INCLUDE_PERFETTO_PROTOZERO_SCATTERED_HEAP_BUFFER_H_
 
 #include <memory>
+#include <string>
 #include <vector>
 
 #include "perfetto/base/export.h"
@@ -26,6 +27,8 @@
 
 namespace protozero {
 
+class Message;
+
 class PERFETTO_EXPORT ScatteredHeapBuffer
     : public protozero::ScatteredStreamWriter::Delegate {
  public:
@@ -86,6 +89,48 @@ class PERFETTO_EXPORT ScatteredHeapBuffer
   std::vector<Slice> slices_;
 };
 
+// Helper function to create heap-based protozero messages in one line.
+// This is a convenience wrapper, mostly for tests, to avoid having to do:
+//   MyMessage msg;
+//   ScatteredHeapBuffer shb;
+//   ScatteredStreamWriter writer(&shb);
+//   shb.set_writer(&writer);
+//   msg.Reset(&shb);
+// Just to get an easily serializable message. Instead this allows simply:
+//   HeapBuffered<MyMessage> msg;
+//   msg->set_stuff(...);
+//   msg->SerializeAsString();
+template <typename T = ::protozero::Message>
+class HeapBuffered {
+ public:
+  HeapBuffered() : shb_(4096, 4096), writer_(&shb_) {
+    shb_.set_writer(&writer_);
+    msg_.Reset(&writer_);
+  }
+
+  // This can't be neither copied nor moved because Message hands out pointers
+  // to itself when creating submessages.
+  HeapBuffered(const HeapBuffered&) = delete;
+  HeapBuffered& operator=(const HeapBuffered&) = delete;
+  HeapBuffered(HeapBuffered&&) = delete;
+  HeapBuffered& operator=(HeapBuffered&&) = delete;
+
+  T* get() { return &msg_; }
+  T* operator->() { return &msg_; }
+
+  std::vector<uint8_t> SerializeAsArray() { return shb_.StitchSlices(); }
+
+  std::string SerializeAsString() {
+    auto vec = SerializeAsArray();
+    return std::string(reinterpret_cast<const char*>(vec.data()), vec.size());
+  }
+
+ private:
+  ScatteredHeapBuffer shb_;
+  ScatteredStreamWriter writer_;
+  T msg_;
+};
+
 }  // namespace protozero
 
 #endif  // INCLUDE_PERFETTO_PROTOZERO_SCATTERED_HEAP_BUFFER_H_
diff --git a/src/perfetto_cmd/pbtxt_to_pb.cc b/src/perfetto_cmd/pbtxt_to_pb.cc
index 687ae79e8..cc123a93c 100644
--- a/src/perfetto_cmd/pbtxt_to_pb.cc
+++ b/src/perfetto_cmd/pbtxt_to_pb.cc
@@ -661,17 +661,12 @@ std::vector<uint8_t> PbtxtToPb(const std::string& input,
   const DescriptorProto* descriptor = name_to_descriptor[kConfigProtoName];
   PERFETTO_CHECK(descriptor);
 
-  protozero::ScatteredHeapBuffer stream_delegate(base::kPageSize);
-  protozero::ScatteredStreamWriter stream(&stream_delegate);
-  stream_delegate.set_writer(&stream);
-
-  protozero::Message message;
-  message.Reset(&stream);
-  ParserDelegate delegate(descriptor, &message, reporter,
+  protozero::HeapBuffered<protozero::Message> message;
+  ParserDelegate delegate(descriptor, message.get(), reporter,
                           std::move(name_to_descriptor),
                           std::move(name_to_enum));
   Parse(input, &delegate);
-  return stream_delegate.StitchSlices();
+  return message.SerializeAsArray();
 }
 
 }  // namespace perfetto
diff --git a/src/protozero/proto_decoder_unittest.cc b/src/protozero/proto_decoder_unittest.cc
index 16efd9a97..2f5ab5356 100644
--- a/src/protozero/proto_decoder_unittest.cc
+++ b/src/protozero/proto_decoder_unittest.cc
@@ -32,19 +32,12 @@ using ::testing::Invoke;
 using namespace proto_utils;
 
 TEST(ProtoDecoderTest, ReadString) {
-  Message message;
-  ScatteredHeapBuffer delegate(512, 512);
-  ScatteredStreamWriter writer(&delegate);
-  delegate.set_writer(&writer);
-  message.Reset(&writer);
+  HeapBuffered<Message> message;
 
   static constexpr char kTestString[] = "test";
-  message.AppendString(1, kTestString);
-
-  delegate.AdjustUsedSizeOfCurrentSlice();
-  auto used_range = delegate.slices()[0].GetUsedRange();
-
-  TypedProtoDecoder<32, false> decoder(used_range.begin, used_range.size());
+  message->AppendString(1, kTestString);
+  std::vector<uint8_t> proto = message.SerializeAsArray();
+  TypedProtoDecoder<32, false> decoder(proto.data(), proto.size());
 
   const auto& field = decoder.Get(1);
   ASSERT_EQ(field.type(), ProtoWireType::kLengthDelimited);
diff --git a/src/trace_processor/metrics/metrics.cc b/src/trace_processor/metrics/metrics.cc
index ca68c4e66..f8273d6cd 100644
--- a/src/trace_processor/metrics/metrics.cc
+++ b/src/trace_processor/metrics/metrics.cc
@@ -188,12 +188,7 @@ void BuildProto(sqlite3_context* ctx, int argc, sqlite3_value** argv) {
     return;
   }
 
-  protozero::ScatteredHeapBuffer delegate;
-  protozero::ScatteredStreamWriter writer(&delegate);
-  delegate.set_writer(&writer);
-
-  protozero::Message message;
-  message.Reset(&writer);
+  protozero::HeapBuffered<protozero::Message> message;
 
   for (int i = 0; i < argc; i += 2) {
     auto* value = argv[i + 1];
@@ -214,27 +209,27 @@ void BuildProto(sqlite3_context* ctx, int argc, sqlite3_value** argv) {
         return;
       }
       auto* text = reinterpret_cast<const char*>(sqlite3_value_text(value));
-      auto status = BuildProtoRepeatedField(fn_ctx->tp, field, text, &message);
+      auto status =
+          BuildProtoRepeatedField(fn_ctx->tp, field, text, message.get());
       if (!status.ok()) {
         sqlite3_result_error(ctx, status.c_message(), -1);
         return;
       }
     } else {
       auto sql_value = SqlValueFromSqliteValue(value);
-      auto status = AppendValueToMessage(field, sql_value, &message);
+      auto status = AppendValueToMessage(field, sql_value, message.get());
       if (!status.ok()) {
         sqlite3_result_error(ctx, status.c_message(), -1);
         return;
       }
     }
   }
-  message.Finalize();
+  message->Finalize();
 
-  auto slices = delegate.StitchSlices();
-  std::unique_ptr<uint8_t[]> data(static_cast<uint8_t*>(malloc(slices.size())));
-  memcpy(data.get(), slices.data(), slices.size());
-  sqlite3_result_blob(ctx, data.release(), static_cast<int>(slices.size()),
-                      free);
+  std::vector<uint8_t> raw = message.SerializeAsArray();
+  std::unique_ptr<uint8_t[]> data(static_cast<uint8_t*>(malloc(raw.size())));
+  memcpy(data.get(), raw.data(), raw.size());
+  sqlite3_result_blob(ctx, data.release(), static_cast<int>(raw.size()), free);
 }
 
 void RunMetric(sqlite3_context* ctx, int argc, sqlite3_value** argv) {
@@ -299,12 +294,7 @@ util::Status ComputeMetrics(TraceProcessor* tp,
                             const std::vector<SqlMetric>& sql_metrics,
                             const ProtoDescriptor& root_descriptor,
                             std::vector<uint8_t>* metrics_proto) {
-  protozero::ScatteredHeapBuffer delegate;
-  protozero::ScatteredStreamWriter writer(&delegate);
-  delegate.set_writer(&writer);
-
-  protozero::Message metrics_message;
-  metrics_message.Reset(&writer);
+  protozero::HeapBuffered<protozero::Message> metrics_message;
 
   for (const auto& sql_metric : sql_metrics) {
     // If there's no proto to fill in, then we don't need to do a query.
@@ -348,7 +338,7 @@ util::Status ComputeMetrics(TraceProcessor* tp,
 
     const auto& field = root_descriptor.fields()[opt_idx.value()];
     const uint8_t* ptr = static_cast<const uint8_t*>(col.bytes_value);
-    metrics_message.AppendBytes(field.number(), ptr, col.bytes_count);
+    metrics_message->AppendBytes(field.number(), ptr, col.bytes_count);
 
     has_next = it.Next();
     if (has_next)
@@ -358,9 +348,9 @@ util::Status ComputeMetrics(TraceProcessor* tp,
     if (!status.ok())
       return status;
   }
-  metrics_message.Finalize();
+  metrics_message->Finalize();
 
-  *metrics_proto = delegate.StitchSlices();
+  *metrics_proto = metrics_message.SerializeAsArray();
   return util::OkStatus();
 }
 
diff --git a/src/tracing/core/trace_writer_for_testing.h b/src/tracing/core/trace_writer_for_testing.h
index d367c1bec..f7a1a883b 100644
--- a/src/tracing/core/trace_writer_for_testing.h
+++ b/src/tracing/core/trace_writer_for_testing.h
@@ -30,7 +30,6 @@ namespace perfetto {
 // See //include/perfetto/tracing/core/trace_writer.h for docs.
 class TraceWriterForTesting : public TraceWriter {
  public:
-  // TraceWriterForTesting(const protozero::ScatteredHeapBuffer& delegate);
   TraceWriterForTesting();
   ~TraceWriterForTesting() override;