services/tracing: Add trace packet tokenizer

Add a tokenizer utility which can parse a possibly fragmented stream of
Perfetto trace packet data into complete and well formed trace packets.
This is needed for passing trace data from the mojo tracing service to
the Perfetto client API.

Bug: b/158460267
Change-Id: I893c25699c9d66e7783136362e211955732caa9e
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2279787
Commit-Queue: Sami Kyöstilä <skyostil@chromium.org>
Reviewed-by: Eric Seckler <eseckler@chromium.org>
Reviewed-by: Stephen Nusko <nuskos@chromium.org>
Cr-Commit-Position: refs/heads/master@{#785517}

services/tracing/BUILD.gn

@@ -86,6 +86,7 @@ source_set("tests") {
     "public/cpp/perfetto/producer_test_utils.h",
     "public/cpp/perfetto/task_runner_unittest.cc",
     "public/cpp/perfetto/trace_event_data_source_unittest.cc",
+    "public/cpp/perfetto/trace_packet_tokenizer_unittest.cc",
     "public/cpp/perfetto/traced_value_proto_writer_unittest.cc",
     "public/cpp/stack_sampling/reached_code_data_source_android_unittest.cc",
     "public/cpp/stack_sampling/tracing_sampler_profiler_unittest.cc",

services/tracing/public/cpp/BUILD.gn

@@ -91,6 +91,8 @@ target(tracing_lib_type, "cpp") {
       "perfetto/task_runner.h",
       "perfetto/trace_event_data_source.cc",
       "perfetto/trace_event_data_source.h",
+      "perfetto/trace_packet_tokenizer.cc",
+      "perfetto/trace_packet_tokenizer.h",
       "perfetto/trace_time.cc",
       "perfetto/trace_time.h",
       "perfetto/traced_value_proto_writer.cc",

services/tracing/public/cpp/perfetto/trace_packet_tokenizer.cc

@@ -0,0 +1,99 @@
+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "services/tracing/public/cpp/perfetto/trace_packet_tokenizer.h"
+
+#include "base/check.h"
+#include "base/check_op.h"
+#include "base/macros.h"
+#include "third_party/perfetto/include/perfetto/ext/tracing/core/trace_packet.h"
+
+namespace tracing {
+namespace {
+
+static constexpr uint8_t kPacketTag =
+    protozero::proto_utils::MakeTagLengthDelimited(
+        perfetto::TracePacket::kPacketFieldNumber);
+
+}  // namespace
+
+TracePacketTokenizer::TracePacketTokenizer() = default;
+TracePacketTokenizer::~TracePacketTokenizer() = default;
+TracePacketTokenizer::Packet::Packet() = default;
+TracePacketTokenizer::Packet::~Packet() = default;
+
+std::vector<perfetto::TracePacket> TracePacketTokenizer::Parse(
+    const uint8_t* data,
+    size_t size) {
+  std::vector<perfetto::TracePacket> packets;
+  const uint8_t* data_end = data + size;
+  const uint8_t* packet_ptr = data;
+
+  // Only one fragmented packet can be finalized per call to Parse(), so clear
+  // any previous one.
+  assembled_packet_ = Packet();
+
+  while (packet_ptr < data_end) {
+    // First parse the packet header, i.e., the one byte field tag and the
+    // variable sized packet length field.
+    if (!next_packet_.parsed_size) {
+      // Parse the field tag.
+      auto prev_header_size = next_packet_.header->size();
+      size_t bytes_to_copy = kMaxHeaderSize - prev_header_size;
+      next_packet_.header->insert(
+          next_packet_.header->end(), packet_ptr,
+          std::min(packet_ptr + bytes_to_copy, data_end));
+      DCHECK(next_packet_.header->size() <= kMaxHeaderSize);
+      if (next_packet_.header->size() < kMinHeaderSize) {
+        // Not enough data -- try again later.
+        return packets;
+      }
+      DCHECK_EQ(kPacketTag, next_packet_.header[0]);
+
+      // Parse the size field.
+      const auto* size_begin = &next_packet_.header[1];
+      const auto* size_end = protozero::proto_utils::ParseVarInt(
+          size_begin, &*next_packet_.header->end(), &next_packet_.parsed_size);
+      size_t size_field_size = size_end - size_begin;
+      if (!size_field_size) {
+        // Size field overflows to next chunk. Try again later.
+        return packets;
+      }
+      // Find the start of the packet data after the size field.
+      packet_ptr += sizeof(kPacketTag) + size_field_size - prev_header_size;
+    }
+
+    // We've now parsed the the proto preamble and the size field for our
+    // packet. Let's see if the packet fits completely into this chunk.
+    DCHECK(next_packet_.parsed_size);
+    size_t remaining_size =
+        next_packet_.parsed_size - next_packet_.partial_data->size();
+    if (packet_ptr + remaining_size > data_end) {
+      // Save remaining bytes into overflow buffer and try again later.
+      next_packet_.partial_data->insert(next_packet_.partial_data->end(),
+                                        packet_ptr, data_end);
+      return packets;
+    }
+
+    // The packet is now complete. It can have a slice overflowing from the
+    // previous chunk(s) as well a a slice in the current chunk.
+    packets.emplace_back();
+    if (!next_packet_.partial_data->empty()) {
+      DCHECK(assembled_packet_.partial_data->empty());
+      assembled_packet_ = std::move(next_packet_);
+      packets.back().AddSlice(&assembled_packet_.partial_data[0],
+                              assembled_packet_.partial_data->size());
+    }
+    CHECK_LE(packet_ptr + remaining_size, data_end);
+    packets.back().AddSlice(packet_ptr, remaining_size);
+    packet_ptr += remaining_size;
+
+    // Start a new packet.
+    next_packet_ = Packet();
+  }
+  DCHECK_EQ(packet_ptr, data_end);
+  return packets;
+}
+
+}  // namespace tracing

services/tracing/public/cpp/perfetto/trace_packet_tokenizer.h

@@ -0,0 +1,78 @@
+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef SERVICES_TRACING_PUBLIC_CPP_PERFETTO_TRACE_PACKET_TOKENIZER_H_
+#define SERVICES_TRACING_PUBLIC_CPP_PERFETTO_TRACE_PACKET_TOKENIZER_H_
+
+#include <vector>
+
+#include "base/component_export.h"
+#include "base/containers/stack_container.h"
+#include "third_party/perfetto/include/perfetto/protozero/proto_utils.h"
+
+namespace perfetto {
+class TracePacket;
+}  // namespace perfetto
+
+namespace tracing {
+
+// Converts between a raw stream of perfetto::TracePacket bytes and a tokenized
+// vector of delineated packets.
+//
+// The tracing service provides us with serialized proto bytes, while the
+// Perfetto consumer expects a vector of TracePackets (i.e., without the field
+// number and size header) without partial or truncated packets. To translate
+// between the two, we find the position of each TracePacket and construct a
+// vector pointing to the original data at the corresponding offsets.
+//
+// To make matters more complicated, mojo can split the data chunks
+// arbitrarily, including in the middle of trace packets. To work around
+// this, we tokenize as much data as we can and buffer any unprocessed bytes as
+// long as needed.
+class COMPONENT_EXPORT(TRACING_CPP) TracePacketTokenizer {
+ public:
+  TracePacketTokenizer();
+  ~TracePacketTokenizer();
+
+  // Given a chunk of trace data, tokenize as many trace packets as possible
+  // (could be zero) and return the result. Note that the tokenized packets have
+  // pointers to |data| as well as |this|, so they won't be valid after another
+  // call to Parse().
+  std::vector<perfetto::TracePacket> Parse(const uint8_t* data, size_t size);
+
+  // Returns |true| if there is more data left to be consumed in the tokenizer.
+  bool has_more() const {
+    return !next_packet_.header->empty() || next_packet_.parsed_size ||
+           !next_packet_.partial_data->empty();
+  }
+
+ private:
+  static constexpr size_t kMinHeaderSize = sizeof(uint8_t) + sizeof(uint8_t);
+  static constexpr size_t kMaxHeaderSize =
+      sizeof(uint8_t) + protozero::proto_utils::kMessageLengthFieldSize;
+
+  struct Packet {
+    Packet();
+    ~Packet();
+
+    // Most trace packets are very small, so avoid heap allocations in the
+    // common case where one packet straddles the boundary between chunks.
+    base::StackVector<uint8_t, 64> partial_data;
+
+    uint64_t parsed_size = 0;
+    base::StackVector<uint8_t, kMaxHeaderSize> header;
+  };
+
+  // Packet currently being parsed.
+  Packet next_packet_;
+
+  // Most recently completed packet which spanned multiple chunks. Kept as a
+  // member so that the memory remains valid while the caller is processing the
+  // results.
+  Packet assembled_packet_;
+};
+
+}  // namespace tracing
+
+#endif  // SERVICES_TRACING_PUBLIC_CPP_PERFETTO_TRACE_PACKET_TOKENIZER_H_

services/tracing/public/cpp/perfetto/trace_packet_tokenizer_unittest.cc

@@ -0,0 +1,143 @@
+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "services/tracing/public/cpp/perfetto/trace_packet_tokenizer.h"
+
+#include "testing/gtest/include/gtest/gtest.h"
+#include "third_party/perfetto/include/perfetto/ext/tracing/core/trace_packet.h"
+#include "third_party/perfetto/protos/perfetto/trace/trace.pb.h"
+#include "third_party/perfetto/protos/perfetto/trace/trace_packet.pb.h"
+
+#include <list>
+
+namespace tracing {
+namespace {
+
+constexpr char kTestString[] = "This little packet went to the market";
+
+class TracePacketTokenizerTest : public testing::Test {
+ protected:
+  // Parse a trace chunk using an indirect memory allocation so ASAN can detect
+  // any out-of-bounds reads.
+  std::vector<perfetto::TracePacket> ParseChunk(const uint8_t* data,
+                                                size_t size) {
+    input_chunks_.emplace_back(data, data + size);
+    auto& it = input_chunks_.back();
+    return tokenizer_.Parse(it.data(), it.size());
+  }
+
+  void Reset() {
+    input_chunks_.clear();
+    tokenizer_ = TracePacketTokenizer();
+  }
+
+  TracePacketTokenizer& tokenizer() { return tokenizer_; }
+
+ private:
+  std::list<std::vector<uint8_t>> input_chunks_;
+  TracePacketTokenizer tokenizer_;
+};
+
+TEST_F(TracePacketTokenizerTest, Basic) {
+  perfetto::protos::Trace trace;
+  auto* packet = trace.add_packet();
+  packet->mutable_for_testing()->set_str(kTestString);
+  auto packet_data = trace.SerializeAsString();
+
+  auto packets = ParseChunk(
+      reinterpret_cast<const uint8_t*>(packet_data.data()), packet_data.size());
+  EXPECT_EQ(1u, packets.size());
+  perfetto::protos::TracePacket parsed_packet;
+  EXPECT_TRUE(
+      parsed_packet.ParseFromString(packets[0].GetRawBytesForTesting()));
+  EXPECT_EQ(kTestString, parsed_packet.for_testing().str());
+  EXPECT_FALSE(tokenizer().has_more());
+}
+
+TEST_F(TracePacketTokenizerTest, PartialParse) {
+  perfetto::protos::Trace trace;
+  auto* packet = trace.add_packet();
+  packet->mutable_for_testing()->set_str(kTestString);
+  auto packet_data = trace.SerializeAsString();
+
+  auto packets =
+      ParseChunk(reinterpret_cast<const uint8_t*>(packet_data.data()),
+                 packet_data.size() / 2);
+  EXPECT_TRUE(packets.empty());
+  EXPECT_TRUE(tokenizer().has_more());
+
+  packets = ParseChunk(reinterpret_cast<const uint8_t*>(packet_data.data() +
+                                                        packet_data.size() / 2),
+                       packet_data.size() / 2);
+  EXPECT_EQ(1u, packets.size());
+  perfetto::protos::TracePacket parsed_packet;
+  EXPECT_TRUE(
+      parsed_packet.ParseFromString(packets[0].GetRawBytesForTesting()));
+  EXPECT_EQ(kTestString, parsed_packet.for_testing().str());
+  EXPECT_FALSE(tokenizer().has_more());
+}
+
+TEST_F(TracePacketTokenizerTest, MultiplePackets) {
+  constexpr size_t kPacketCount = 32;
+  perfetto::protos::Trace trace;
+  for (size_t i = 0; i < kPacketCount; i++) {
+    auto* packet = trace.add_packet();
+    packet->set_timestamp(i);
+    packet->mutable_for_testing()->set_str(kTestString);
+  }
+  auto packet_data = trace.SerializeAsString();
+
+  auto packets = ParseChunk(
+      reinterpret_cast<const uint8_t*>(packet_data.data()), packet_data.size());
+  EXPECT_EQ(kPacketCount, packets.size());
+
+  for (size_t i = 0; i < kPacketCount; i++) {
+    perfetto::protos::TracePacket parsed_packet;
+    EXPECT_TRUE(
+        parsed_packet.ParseFromString(packets[i].GetRawBytesForTesting()));
+    EXPECT_EQ(i, parsed_packet.timestamp());
+    EXPECT_EQ(kTestString, parsed_packet.for_testing().str());
+  }
+  EXPECT_FALSE(tokenizer().has_more());
+}
+
+TEST_F(TracePacketTokenizerTest, Fragmentation) {
+  constexpr size_t kPacketCount = 17;
+  perfetto::protos::Trace trace;
+  for (size_t i = 0; i < kPacketCount; i++) {
+    auto* packet = trace.add_packet();
+    packet->set_timestamp(i + 1);
+    packet->mutable_for_testing()->set_str(kTestString);
+  }
+  auto packet_data = trace.SerializeAsString();
+
+  for (size_t chunk_size = 1; chunk_size < packet_data.size(); chunk_size++) {
+    size_t packet_count = 0;
+    for (size_t offset = 0; offset < packet_data.size(); offset += chunk_size) {
+      const auto* chunk_start =
+          reinterpret_cast<const uint8_t*>(packet_data.data()) + offset;
+      const auto* chunk_end =
+          std::min(chunk_start + chunk_size,
+                   reinterpret_cast<const uint8_t*>(&*packet_data.end()));
+      auto packets = ParseChunk(chunk_start, chunk_end - chunk_start);
+      if (packets.empty())
+        continue;
+      packet_count += packets.size();
+
+      for (auto& packet : packets) {
+        perfetto::protos::TracePacket parsed_packet;
+        EXPECT_TRUE(
+            parsed_packet.ParseFromString(packet.GetRawBytesForTesting()));
+        EXPECT_GT(parsed_packet.timestamp(), 0u);
+        EXPECT_EQ(kTestString, parsed_packet.for_testing().str());
+      }
+    }
+    EXPECT_EQ(kPacketCount, packet_count);
+    EXPECT_FALSE(tokenizer().has_more());
+    Reset();
+  }
+}
+
+}  // namespace
+}  // namespace tracing