TraceBuffer: minor cleanups

Minor cleanups in preparation for next CL.
The major changes introduced by this CL are:
- Intoduce a dedicated PatchList class
- Add uid support to TraceBuffer

Test: perfetto_unittests
Bug: 73612642
Change-Id: Id1c13e569b61e9265b79653598e81aed84064ecd
diff --git a/Android.bp b/Android.bp
index a451202..28e67f8 100644
--- a/Android.bp
+++ b/Android.bp
@@ -3011,6 +3011,7 @@
     "src/tracing/core/id_allocator_unittest.cc",
     "src/tracing/core/packet_stream_validator.cc",
     "src/tracing/core/packet_stream_validator_unittest.cc",
+    "src/tracing/core/patch_list_unittest.cc",
     "src/tracing/core/service_impl.cc",
     "src/tracing/core/service_impl_unittest.cc",
     "src/tracing/core/shared_memory_abi.cc",
diff --git a/include/perfetto/tracing/core/commit_data_request.h b/include/perfetto/tracing/core/commit_data_request.h
index 6a08d40..7974852 100644
--- a/include/perfetto/tracing/core/commit_data_request.h
+++ b/include/perfetto/tracing/core/commit_data_request.h
@@ -101,6 +101,9 @@
 
       const std::string& data() const { return data_; }
       void set_data(const std::string& value) { data_ = value; }
+      void set_data(const void* p, size_t s) {
+        data_.assign(reinterpret_cast<const char*>(p), s);
+      }
 
      private:
       uint32_t offset_ = {};
diff --git a/include/perfetto/tracing/core/shared_memory_abi.h b/include/perfetto/tracing/core/shared_memory_abi.h
index 77c0831..0747f08 100644
--- a/include/perfetto/tracing/core/shared_memory_abi.h
+++ b/include/perfetto/tracing/core/shared_memory_abi.h
@@ -346,6 +346,10 @@
 
     // Begin of the first packet (or packet fragment).
     uint8_t* payload_begin() const { return begin_ + sizeof(ChunkHeader); }
+    size_t payload_size() const {
+      PERFETTO_DCHECK(size_ >= sizeof(ChunkHeader));
+      return size_ - sizeof(ChunkHeader);
+    }
 
     bool is_valid() const { return begin_ && size_; }
 
diff --git a/include/perfetto/tracing/core/trace_packet.h b/include/perfetto/tracing/core/trace_packet.h
index 1491807..cbcef60 100644
--- a/include/perfetto/tracing/core/trace_packet.h
+++ b/include/perfetto/tracing/core/trace_packet.h
@@ -51,8 +51,7 @@
   TracePacket& operator=(TracePacket&&);
 
   // Accesses all the raw slices in the packet, for saving them to file/network.
-  const_iterator begin() const { return slices_.begin(); }
-  const_iterator end() const { return slices_.end(); }
+  const Slices& slices() const { return slices_; }
 
   // Decodes the packet for inline use.
   bool Decode();
@@ -66,6 +65,7 @@
 
   // Mutator, used only by the service and tests.
   void AddSlice(Slice);
+  void AddSlice(const void* start, size_t size);
 
   // Total size of all slices.
   size_t size() const { return size_; }
diff --git a/src/perfetto_cmd/perfetto_cmd.cc b/src/perfetto_cmd/perfetto_cmd.cc
index 946ecec..ee7064a 100644
--- a/src/perfetto_cmd/perfetto_cmd.cc
+++ b/src/perfetto_cmd/perfetto_cmd.cc
@@ -267,7 +267,7 @@
 void PerfettoCmd::OnTraceData(std::vector<TracePacket> packets, bool has_more) {
   PERFETTO_DLOG("Received trace packet, has_more=%d", has_more);
   for (TracePacket& packet : packets) {
-    for (const Slice& slice : packet) {
+    for (const Slice& slice : packet.slices()) {
       uint8_t preamble[16];
       uint8_t* pos = preamble;
       pos = WriteVarInt(
diff --git a/src/tracing/BUILD.gn b/src/tracing/BUILD.gn
index 3a24f6a..486fb20 100644
--- a/src/tracing/BUILD.gn
+++ b/src/tracing/BUILD.gn
@@ -37,6 +37,7 @@
     "core/id_allocator.h",
     "core/packet_stream_validator.cc",
     "core/packet_stream_validator.h",
+    "core/patch_list.h",
     "core/service_impl.cc",
     "core/service_impl.h",
     "core/shared_memory_abi.cc",
@@ -119,6 +120,7 @@
   sources = [
     "core/id_allocator_unittest.cc",
     "core/packet_stream_validator_unittest.cc",
+    "core/patch_list_unittest.cc",
     "core/service_impl_unittest.cc",
     "core/shared_memory_abi_unittest.cc",
     "core/shared_memory_arbiter_impl_unittest.cc",
diff --git a/src/tracing/core/patch_list.h b/src/tracing/core/patch_list.h
new file mode 100644
index 0000000..5cef17d
--- /dev/null
+++ b/src/tracing/core/patch_list.h
@@ -0,0 +1,108 @@
+/*
+ * Copyright (C) 2018 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.
+ */
+
+#ifndef SRC_TRACING_CORE_PATCH_LIST_H_
+#define SRC_TRACING_CORE_PATCH_LIST_H_
+
+#include <array>
+#include <forward_list>
+
+#include "perfetto/base/logging.h"
+#include "perfetto/tracing/core/basic_types.h"
+#include "perfetto/tracing/core/shared_memory_abi.h"
+
+namespace perfetto {
+
+// Used to handle the backfilling of the headers (the |size_field|) of nested
+// messages when a proto is fragmented over several chunks. These patches are
+// sent out-of-band to the tracing service, after having returned the initial
+// chunks of the fragment.
+class Patch {
+ public:
+  using PatchContent = std::array<uint8_t, SharedMemoryABI::kPacketHeaderSize>;
+  Patch(ChunkID c, uint16_t o) : chunk_id(c), offset(o) {}
+  Patch(const Patch&) = default;  // For tests.
+
+  const ChunkID chunk_id;
+  const uint16_t offset;
+  PatchContent size_field{};
+
+  // |size_field| contains a varint. Any varint must start with != 0. Even in
+  // the case we want to encode a size == 0, protozero will write a redundant
+  // varint for that, that is [0x80, 0x80, 0x80, 0x00]. So the first byte is 0
+  // iff we never wrote any varint into that.
+  bool is_patched() const { return size_field[0] != 0; }
+
+  // For tests.
+  bool operator==(const Patch& o) const {
+    return chunk_id == o.chunk_id && offset == o.offset &&
+           size_field == o.size_field;
+  }
+
+ private:
+  Patch& operator=(const Patch&) = default;
+  Patch(Patch&&) noexcept = delete;
+  Patch& operator=(Patch&&) = delete;
+};
+
+// Note: the protozero::Message(s) will take pointers to the |size_field| of
+// these entries. This container must guarantee that the Patch objects are never
+// moved around (i.e. cannot be a vector because of reallocations can change
+// addresses of pre-existing entries).
+class PatchList {
+ public:
+  using ListType = std::forward_list<Patch>;
+  using value_type = ListType::value_type;          // For gtest.
+  using const_iterator = ListType::const_iterator;  // For gtest.
+
+  PatchList() : last_(list_.before_begin()) {}
+
+  Patch* emplace_back(ChunkID chunk_id, uint16_t offset) {
+    PERFETTO_DCHECK(empty() || last_->chunk_id != chunk_id ||
+                    offset >= last_->offset + sizeof(Patch::PatchContent));
+    last_ = list_.emplace_after(last_, chunk_id, offset);
+    return &*last_;
+  }
+
+  void pop_front() {
+    PERFETTO_DCHECK(!list_.empty());
+    list_.pop_front();
+    if (empty())
+      last_ = list_.before_begin();
+  }
+
+  const Patch& front() const {
+    PERFETTO_DCHECK(!list_.empty());
+    return list_.front();
+  }
+
+  const Patch& back() const {
+    PERFETTO_DCHECK(!list_.empty());
+    return *last_;
+  }
+
+  ListType::const_iterator begin() const { return list_.begin(); }
+  ListType::const_iterator end() const { return list_.end(); }
+  bool empty() const { return list_.empty(); }
+
+ private:
+  ListType list_;
+  ListType::iterator last_;
+};
+
+}  // namespace perfetto
+
+#endif  // SRC_TRACING_CORE_PATCH_LIST_H_
diff --git a/src/tracing/core/patch_list_unittest.cc b/src/tracing/core/patch_list_unittest.cc
new file mode 100644
index 0000000..99609b2
--- /dev/null
+++ b/src/tracing/core/patch_list_unittest.cc
@@ -0,0 +1,93 @@
+/*
+ * Copyright (C) 2018 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/tracing/core/patch_list.h"
+
+#include <ostream>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+using testing::ElementsAre;
+
+namespace perfetto {
+
+std::ostream& operator<<(std::ostream& o, const Patch& p);
+std::ostream& operator<<(std::ostream& o, const Patch& p) {
+  o << p.chunk_id << "@" << p.offset << " : {" << std::hex << p.size_field[0]
+    << "," << p.size_field[1] << "," << p.size_field[2] << ","
+    << p.size_field[3] << "}";
+  return o;
+}
+
+namespace {
+
+TEST(PatchListTest, InsertAndRemove) {
+  PatchList pl;
+
+  ASSERT_TRUE(pl.empty());
+
+  pl.emplace_back(ChunkID(5), 50);
+  ASSERT_THAT(pl, ElementsAre(Patch(ChunkID(5), 50)));
+
+  pl.emplace_back(ChunkID(6), 60);
+  ASSERT_THAT(pl, ElementsAre(Patch(ChunkID(5), 50), Patch(ChunkID(6), 60)));
+
+  ASSERT_EQ(pl.front(), Patch(ChunkID(5), 50));
+  ASSERT_EQ(pl.back(), Patch(ChunkID(6), 60));
+
+  pl.pop_front();
+  ASSERT_EQ(pl.front(), Patch(ChunkID(6), 60));
+  pl.emplace_back(ChunkID(7), 70);
+
+  pl.pop_front();
+  ASSERT_EQ(pl.front(), Patch(ChunkID(7), 70));
+  ASSERT_EQ(pl.back(), Patch(ChunkID(7), 70));
+
+  pl.pop_front();
+
+  for (int i = 0; i < 3; i++) {
+    ASSERT_TRUE(pl.empty());
+
+    pl.emplace_back(ChunkID(8), 80);
+    pl.emplace_back(ChunkID(9), 90);
+    ASSERT_THAT(pl, ElementsAre(Patch(ChunkID(8), 80), Patch(ChunkID(9), 90)));
+
+    pl.pop_front();
+    pl.pop_front();
+  }
+}
+
+TEST(PatchListTest, PointerStability) {
+  PatchList pl;
+  const uint8_t* ptrs[10]{};
+  for (uint16_t i = 0; i < 1000; i++) {
+    pl.emplace_back(ChunkID(i), i);
+    if (i >= 1000 - 10)
+      ptrs[i - (1000 - 10)] = &pl.back().size_field[0];
+  }
+
+  for (uint16_t i = 0; i < 1000 - 10; i++)
+    pl.pop_front();
+
+  auto it = pl.begin();
+  for (uint16_t i = 0; it != pl.end(); it++, i++) {
+    EXPECT_EQ(ptrs[i], &it->size_field[0]);
+  }
+}
+
+}  // namespace
+}  // namespace perfetto
diff --git a/src/tracing/core/service_impl.cc b/src/tracing/core/service_impl.cc
index 1e50632..1d61aae 100644
--- a/src/tracing/core/service_impl.cc
+++ b/src/tracing/core/service_impl.cc
@@ -618,7 +618,8 @@
 }
 
 void ServiceImpl::UpdateMemoryGuardrail() {
-#if !PERFETTO_BUILDFLAG(PERFETTO_CHROMIUM_BUILD)
+#if !PERFETTO_BUILDFLAG(PERFETTO_CHROMIUM_BUILD) && \
+    !PERFETTO_BUILDFLAG(PERFETTO_OS_MACOSX)
   uint64_t total_buffer_bytes = 0;
 
   // Sum up all the shared memory buffers.
diff --git a/src/tracing/core/shared_memory_abi_unittest.cc b/src/tracing/core/shared_memory_abi_unittest.cc
index bf6c8ce..78d8476 100644
--- a/src/tracing/core/shared_memory_abi_unittest.cc
+++ b/src/tracing/core/shared_memory_abi_unittest.cc
@@ -111,6 +111,8 @@
           (page_size() - sizeof(SharedMemoryABI::PageHeader)) / num_chunks;
       expected_chunk_size = expected_chunk_size - (expected_chunk_size % 4);
       ASSERT_EQ(expected_chunk_size, chunk.size());
+      ASSERT_EQ(expected_chunk_size - sizeof(SharedMemoryABI::ChunkHeader),
+                chunk.payload_size());
       ASSERT_GT(chunk.begin(), page_start);
       ASSERT_GT(chunk.begin(), last_chunk_begin);
       ASSERT_GE(chunk.begin(), last_chunk_end);
diff --git a/src/tracing/core/trace_buffer.cc b/src/tracing/core/trace_buffer.cc
index d204bcd..6d207d8 100644
--- a/src/tracing/core/trace_buffer.cc
+++ b/src/tracing/core/trace_buffer.cc
@@ -22,13 +22,7 @@
 #include "perfetto/base/logging.h"
 #include "perfetto/protozero/proto_utils.h"
 #include "perfetto/tracing/core/shared_memory_abi.h"
-
-// TODO(primiano): we need some flag to figure out if the packets on the
-// boundary require patching or have already been patched. The current
-// implementation considers all packets eligible to be read once we have all the
-// chunks that compose them.
-
-// TODO(primiano): copy over skyostil@'s trusted_uid logic.
+#include "perfetto/tracing/core/trace_packet.h"
 
 #define TRACE_BUFFER_VERBOSE_LOGGING() 0  // Set to 1 when debugging unittests.
 #if TRACE_BUFFER_VERBOSE_LOGGING()
@@ -108,6 +102,7 @@
 // that the producer is malicious and will change the content of |src|
 // while we execute here. Don't do any processing on it other than memcpy().
 void TraceBuffez::CopyChunkUntrusted(ProducerID producer_id_trusted,
+                                     uid_t producer_uid_trusted,
                                      WriterID writer_id,
                                      ChunkID chunk_id,
                                      uint16_t num_fragments,
@@ -172,14 +167,15 @@
 
   // Now first insert the new chunk. At the end, if necessary, add the padding.
   ChunkMeta::Key key(record);
-  auto it_and_inserted = index_.emplace(
-      key, ChunkMeta(GetChunkRecordAt(wptr_), num_fragments, chunk_flags));
+  auto it_and_inserted =
+      index_.emplace(key, ChunkMeta(GetChunkRecordAt(wptr_), num_fragments,
+                                    chunk_flags, producer_uid_trusted));
   if (PERFETTO_UNLIKELY(!it_and_inserted.second)) {
     // More likely a producer bug, but could also be a malicious producer.
     PERFETTO_DCHECK(suppress_sanity_dchecks_for_testing_);
     index_.erase(it_and_inserted.first);
-    index_.emplace(
-        key, ChunkMeta(GetChunkRecordAt(wptr_), num_fragments, chunk_flags));
+    index_.emplace(key, ChunkMeta(GetChunkRecordAt(wptr_), num_fragments,
+                                  chunk_flags, producer_uid_trusted));
   }
   TRACE_BUFFER_DLOG("  copying @ [%lu - %lu] %zu", wptr_ - begin(),
                     wptr_ - begin() + record_size, record_size);
@@ -381,13 +377,15 @@
     cur = seq_begin;
 }
 
-bool TraceBuffez::ReadNextTracePacket(Slices* slices) {
+bool TraceBuffez::ReadNextTracePacket(TracePacket* packet,
+                                      uid_t* producer_uid) {
   // Note: MoveNext() moves only within the next chunk within the same
   // {ProducerID, WriterID} sequence. Here we want to:
   // - return the next patched+complete packet in the current sequence, if any.
   // - return the first patched+complete packet in the next sequence, if any.
   // - return false if none of the above is found.
   TRACE_BUFFER_DLOG("ReadNextTracePacket()");
+  *producer_uid = -1;  // Just in case we forget to initialize it below.
 #if PERFETTO_DCHECK_IS_ON()
   PERFETTO_DCHECK(!changed_since_last_read_);
 #endif
@@ -415,6 +413,8 @@
       continue;
     }
 
+    const uid_t trusted_uid = chunk_meta.trusted_uid;
+
     // At this point we have a chunk in |chunk_meta| that has not been fully
     // read. We don't know yet whether we have enough data to read the full
     // packet (in the case it's fragmented over several chunks) and we are about
@@ -473,8 +473,10 @@
 
       if (action == kReadOnePacket) {
         // The easy peasy case B.
-        if (PERFETTO_LIKELY(ReadNextPacketInChunk(&chunk_meta, slices)))
+        if (PERFETTO_LIKELY(ReadNextPacketInChunk(&chunk_meta, packet))) {
+          *producer_uid = trusted_uid;
           return true;
+        }
 
         // In extremely rare cases (producer bugged / malicious) the chunk might
         // contain an invalid fragment. In such case we don't want to stall the
@@ -483,9 +485,10 @@
       }
 
       PERFETTO_DCHECK(action == kTryReadAhead);
-      ReadAheadResult ra_res = ReadAhead(slices);
+      ReadAheadResult ra_res = ReadAhead(packet);
       if (ra_res == ReadAheadResult::kSucceededReturnSlices) {
         stats_.fragment_readahead_successes++;
+        *producer_uid = trusted_uid;
         return true;
       }
 
@@ -511,7 +514,7 @@
   }    // for(;;MoveNext()) [iterate over chunks].
 }
 
-TraceBuffez::ReadAheadResult TraceBuffez::ReadAhead(Slices* slices) {
+TraceBuffez::ReadAheadResult TraceBuffez::ReadAhead(TracePacket* packet) {
   static_assert(static_cast<ChunkID>(kMaxChunkID + 1) == 0,
                 "relying on kMaxChunkID to wrap naturally");
   TRACE_BUFFER_DLOG(" readahead start @ chunk %u", read_iter_.chunk_id());
@@ -563,7 +566,7 @@
         // In the unlikely case of a corrupted packet, invalidate the all
         // stitching and move on to the next chunk in the same sequence,
         // if any.
-        packet_corruption |= !ReadNextPacketInChunk(&*read_iter_, slices);
+        packet_corruption |= !ReadNextPacketInChunk(&*read_iter_, packet);
       }
       if (read_iter_.cur == it.cur)
         break;
@@ -572,7 +575,7 @@
     PERFETTO_DCHECK(read_iter_.cur == it.cur);
 
     if (PERFETTO_UNLIKELY(packet_corruption)) {
-      slices->clear();
+      *packet = TracePacket();  // clear.
       return ReadAheadResult::kFailedStayOnSameSequence;
     }
 
@@ -581,7 +584,8 @@
   return ReadAheadResult::kFailedMoveToNextSequence;
 }
 
-bool TraceBuffez::ReadNextPacketInChunk(ChunkMeta* chunk_meta, Slices* slices) {
+bool TraceBuffez::ReadNextPacketInChunk(ChunkMeta* chunk_meta,
+                                        TracePacket* packet) {
   PERFETTO_DCHECK(chunk_meta->num_fragments_read < chunk_meta->num_fragments);
   // TODO DCHECK for chunks that are still awaiting patching.
 
@@ -616,8 +620,8 @@
   if (PERFETTO_UNLIKELY(packet_size == 0))
     return false;
 
-  if (PERFETTO_LIKELY(slices))
-    slices->emplace_back(packet_data, static_cast<size_t>(packet_size));
+  if (PERFETTO_LIKELY(packet))
+    packet->AddSlice(packet_data, static_cast<size_t>(packet_size));
 
   return true;
 }
diff --git a/src/tracing/core/trace_buffer.h b/src/tracing/core/trace_buffer.h
index 1d016f2..2aea918 100644
--- a/src/tracing/core/trace_buffer.h
+++ b/src/tracing/core/trace_buffer.h
@@ -32,6 +32,8 @@
 
 namespace perfetto {
 
+class TracePacket;
+
 // The main buffer, owned by the tracing service, where all the trace data is
 // ultimately stored into. The service will own several instances of this class,
 // at least one per active consumer (as defined in the |buffers| section of
@@ -161,6 +163,7 @@
   // None of the arguments should be trusted, unless otherwise stated. We can
   // trust that |src| points to a valid memory area, but not its contents.
   void CopyChunkUntrusted(ProducerID producer_id_trusted,
+                          uid_t producer_uid_trusted,
                           WriterID writer_id,
                           ChunkID chunk_id,
                           uint16_t num_fragments,
@@ -189,13 +192,14 @@
   // Reads in the TraceBuffer are NOT idempotent.
   void BeginRead();
 
-  // Returns the next packet in the buffer, if any. Returns false if no packets
-  // can be read at this point.
+  // Returns the next packet in the buffer, if any, and the uid of the producer
+  // that wrote it (as passed in the CopyChunkUntrusted() call). Returns false
+  // if no packets can be read at this point.
   // This function returns only complete packets. Specifically:
   // When there is at least one complete packet in the buffer, this function
-  // returns true and populates the |slices| argument with the boundaries of
+  // returns true and populates the TracePacket argument with the boundaries of
   // each fragment for one packet.
-  // The output |slices|.size() will be >= 1 when this function returns true.
+  // TracePacket will have at least one slice when this function returns true.
   // When there are no whole packets eligible to read (e.g. we are still missing
   // fragments) this function returns false.
   // This function guarantees also that packets for a given
@@ -210,9 +214,10 @@
   //   P1, P4, P7, P2, P3, P5, P8, P9, P6
   // But the following is guaranteed to NOT happen:
   //   P1, P5, P7, P4 (P4 cannot come after P5)
-  bool ReadNextTracePacket(Slices* slices);
+  bool ReadNextTracePacket(TracePacket*, uid_t* producer_uid);
 
   const Stats& stats() const { return stats_; }
+  size_t size() const { return size_; }
 
  private:
   friend class TraceBufferTest;
@@ -314,10 +319,11 @@
       ChunkID chunk_id;
     };
 
-    ChunkMeta(ChunkRecord* c, uint16_t p, uint8_t f)
-        : chunk_record{c}, flags{f}, num_fragments{p} {}
+    ChunkMeta(ChunkRecord* c, uint16_t p, uint8_t f, uid_t u)
+        : chunk_record{c}, trusted_uid{u}, flags{f}, num_fragments{p} {}
 
     ChunkRecord* const chunk_record;   // Addr of ChunkRecord within |data_|.
+    const uid_t trusted_uid;           // uid of the producer.
     uint8_t flags = 0;                 // See SharedMemoryABI::flags.
     const uint16_t num_fragments = 0;  // Total number of packet fragments.
     uint16_t num_fragments_read = 0;   // Number of fragments already read.
@@ -327,11 +333,6 @@
     // payload (the 1st fragment starts at |chunk_record| +
     // sizeof(ChunkRecord)).
     uint16_t cur_fragment_offset = 0;
-
-    // If != 0 the last fragment in the chunk cannot be read, even if the
-    // subsequent ChunkID is already available, until a patch is applied through
-    // MaybePatchChunkContents().
-    // TODO(primiano): implement this logic in next CL.
   };
 
   using ChunkMap = std::map<ChunkMeta::Key, ChunkMeta>;
@@ -421,11 +422,12 @@
   void AddPaddingRecord(size_t);
 
   // Look for contiguous fragment of the same packet starting from |read_iter_|.
-  // If a contiguous packet is found, all the fragments are pushed into |slices|
-  // and the function returns kSucceededReturnSlices. If not, the function
-  // returns either kFailedMoveToNextSequence or kFailedStayOnSameSequence,
-  // telling the caller to continue looking for packets.
-  ReadAheadResult ReadAhead(Slices* slices);
+  // If a contiguous packet is found, all the fragments are pushed into
+  // TracePacket and the function returns kSucceededReturnSlices. If not, the
+  // function returns either kFailedMoveToNextSequence or
+  // kFailedStayOnSameSequence, telling the caller to continue looking for
+  // packets.
+  ReadAheadResult ReadAhead(TracePacket*);
 
   // Deletes (by marking the record invalid and removing form the index) all
   // chunks from |wptr_| to |wptr_| + |bytes_to_clear|. Returns the size of the
@@ -443,11 +445,12 @@
   size_t DeleteNextChunksFor(size_t bytes_to_clear);
 
   // Decodes the boundaries of the next packet (or a fragment) pointed by
-  // ChunkMeta and pushes that into |Slices*|. It also increments the
+  // ChunkMeta and pushes that into |TracePacket|. It also increments the
   // |num_fragments_read| counter.
-  // The Slices pointer can be nullptr, in which case the read state is still
-  // advanced.
-  bool ReadNextPacketInChunk(ChunkMeta*, Slices*);
+  // TracePacket can be nullptr, in which case the read state is still advanced.
+  // When TracePacket is not nullptr, ProducerID must also be not null and will
+  // be updated with the ProducerID that originally wrote the chunk.
+  bool ReadNextPacketInChunk(ChunkMeta*, TracePacket*);
 
   void DcheckIsAlignedAndWithinBounds(const uint8_t* ptr) const {
     PERFETTO_DCHECK(ptr >= begin() && ptr <= end() - sizeof(ChunkRecord));
diff --git a/src/tracing/core/trace_buffer_unittest.cc b/src/tracing/core/trace_buffer_unittest.cc
index 0a7f8c1..efbccd4 100644
--- a/src/tracing/core/trace_buffer_unittest.cc
+++ b/src/tracing/core/trace_buffer_unittest.cc
@@ -24,6 +24,7 @@
 #include "perfetto/protozero/proto_utils.h"
 #include "perfetto/tracing/core/basic_types.h"
 #include "perfetto/tracing/core/shared_memory_abi.h"
+#include "perfetto/tracing/core/trace_packet.h"
 #include "src/tracing/core/trace_buffer.h"
 #include "src/tracing/test/fake_packet.h"
 
@@ -67,12 +68,13 @@
         p, w, c, patches.data(), patches.size(), other_patches_pending);
   }
 
-  std::vector<FakePacketFragment> ReadPacket() {
+  std::vector<FakePacketFragment> ReadPacket(uint32_t* uid = nullptr) {
     std::vector<FakePacketFragment> fragments;
-    Slices slices;
-    if (!trace_buffer_->ReadNextTracePacket(&slices))
+    TracePacket packet;
+    uint32_t ignore;
+    if (!trace_buffer_->ReadNextTracePacket(&packet, uid ? uid : &ignore))
       return fragments;
-    for (const Slice& slice : slices)
+    for (const Slice& slice : packet.slices())
       fragments.emplace_back(slice.start, slice.size);
     return fragments;
   }
@@ -606,6 +608,44 @@
   ASSERT_THAT(ReadPacket(), IsEmpty());
 }
 
+TEST_F(TraceBufferTest, Fragments_PreserveUID) {
+  ResetBuffer(4096);
+  CreateChunk(ProducerID(1), WriterID(1), ChunkID(0))
+      .AddPacket(10, 'a')
+      .AddPacket(10, 'b', kContOnNextChunk)
+      .SetUID(11)
+      .CopyIntoTraceBuffer();
+  CreateChunk(ProducerID(2), WriterID(1), ChunkID(0))
+      .AddPacket(10, 'c')
+      .AddPacket(10, 'd')
+      .SetUID(22)
+      .CopyIntoTraceBuffer();
+  CreateChunk(ProducerID(1), WriterID(1), ChunkID(1))
+      .AddPacket(10, 'e', kContFromPrevChunk)
+      .AddPacket(10, 'f')
+      .SetUID(11)
+      .CopyIntoTraceBuffer();
+  trace_buffer()->BeginRead();
+  uid_t uid = -1;
+  ASSERT_THAT(ReadPacket(&uid), ElementsAre(FakePacketFragment(10, 'a')));
+  ASSERT_EQ(11u, uid);
+
+  ASSERT_THAT(ReadPacket(&uid), ElementsAre(FakePacketFragment(10, 'b'),
+                                            FakePacketFragment(10, 'e')));
+  ASSERT_EQ(11u, uid);
+
+  ASSERT_THAT(ReadPacket(&uid), ElementsAre(FakePacketFragment(10, 'f')));
+  ASSERT_EQ(11u, uid);
+
+  ASSERT_THAT(ReadPacket(&uid), ElementsAre(FakePacketFragment(10, 'c')));
+  ASSERT_EQ(22u, uid);
+
+  ASSERT_THAT(ReadPacket(&uid), ElementsAre(FakePacketFragment(10, 'd')));
+  ASSERT_EQ(22u, uid);
+
+  ASSERT_THAT(ReadPacket(), IsEmpty());
+}
+
 // --------------------------
 // Out of band patching tests
 // --------------------------
@@ -746,9 +786,9 @@
       .CopyIntoTraceBuffer();
 
   uint8_t valid_ptr = 0;
-  trace_buffer()->CopyChunkUntrusted(ProducerID(1), WriterID(1), ChunkID(1),
-                                     1 /* num packets */, 0 /* flags*/,
-                                     &valid_ptr, sizeof(valid_ptr));
+  trace_buffer()->CopyChunkUntrusted(
+      ProducerID(1), uid_t(0), WriterID(1), ChunkID(1), 1 /* num packets */,
+      0 /* flags*/, &valid_ptr, sizeof(valid_ptr));
 
   CreateChunk(ProducerID(1), WriterID(1), ChunkID(2))
       .AddPacket(32, 'b')
@@ -828,9 +868,9 @@
   std::vector<uint8_t> chunk;
   chunk.insert(chunk.end(), 128 - sizeof(ChunkRecord), 0xff);
   chunk.back() = 0x7f;
-  trace_buffer()->CopyChunkUntrusted(ProducerID(4), WriterID(1), ChunkID(1),
-                                     1 /* num packets */, 0 /* flags*/,
-                                     chunk.data(), chunk.size());
+  trace_buffer()->CopyChunkUntrusted(ProducerID(4), uid_t(0), WriterID(1),
+                                     ChunkID(1), 1 /* num packets */,
+                                     0 /* flags*/, chunk.data(), chunk.size());
 
   // Add a valid chunk.
   CreateChunk(ProducerID(1), WriterID(1), ChunkID(1))
@@ -853,9 +893,9 @@
   chunk.insert(chunk.end(), 64 * 1024 - sizeof(ChunkRecord) * 2, 0xff);
   chunk.back() = 0x7f;
   for (int i = 0; i < 3; i++) {
-    trace_buffer()->CopyChunkUntrusted(ProducerID(1), WriterID(1), ChunkID(1),
-                                       1 /* num packets */, 0 /* flags*/,
-                                       chunk.data(), chunk.size());
+    trace_buffer()->CopyChunkUntrusted(
+        ProducerID(1), uid_t(0), WriterID(1), ChunkID(1), 1 /* num packets */,
+        0 /* flags*/, chunk.data(), chunk.size());
   }
 
   trace_buffer()->BeginRead();
diff --git a/src/tracing/core/trace_packet.cc b/src/tracing/core/trace_packet.cc
index ee762bf..a0801ad 100644
--- a/src/tracing/core/trace_packet.cc
+++ b/src/tracing/core/trace_packet.cc
@@ -45,4 +45,9 @@
   slices_.push_back(std::move(slice));
 }
 
+void TracePacket::AddSlice(const void* start, size_t size) {
+  size_ += size;
+  slices_.emplace_back(start, size);
+}
+
 }  // namespace perfetto
diff --git a/src/tracing/core/trace_packet_unittest.cc b/src/tracing/core/trace_packet_unittest.cc
index 64f5187..c5d06d6 100644
--- a/src/tracing/core/trace_packet_unittest.cc
+++ b/src/tracing/core/trace_packet_unittest.cc
@@ -30,12 +30,12 @@
   proto.mutable_for_testing()->set_str("string field");
   std::string ser_buf = proto.SerializeAsString();
   TracePacket tp;
-  tp.AddSlice({ser_buf.data(), ser_buf.size()});
-  auto slice = tp.begin();
-  ASSERT_NE(tp.end(), slice);
+  tp.AddSlice(ser_buf.data(), ser_buf.size());
+  auto slice = tp.slices().begin();
+  ASSERT_NE(tp.slices().end(), slice);
   ASSERT_EQ(ser_buf.data(), slice->start);
   ASSERT_EQ(ser_buf.size(), slice->size);
-  ASSERT_EQ(tp.end(), ++slice);
+  ASSERT_EQ(tp.slices().end(), ++slice);
 
   ASSERT_TRUE(tp.Decode());
   ASSERT_TRUE(tp.Decode());  // Decode() should be idempotent.
@@ -65,20 +65,20 @@
   tp.AddSlice({ser_buf.data() + 3 + 5, ser_buf.size() - 3 - 5});
   ASSERT_EQ(ser_buf.size(), tp.size());
 
-  auto slice = tp.begin();
-  ASSERT_NE(tp.end(), slice);
+  auto slice = tp.slices().begin();
+  ASSERT_NE(tp.slices().end(), slice);
   ASSERT_EQ(ser_buf.data(), slice->start);
   ASSERT_EQ(3u, slice->size);
 
-  ASSERT_NE(tp.end(), ++slice);
+  ASSERT_NE(tp.slices().end(), ++slice);
   ASSERT_EQ(ser_buf.data() + 3, slice->start);
   ASSERT_EQ(5u, slice->size);
 
-  ASSERT_NE(tp.end(), ++slice);
+  ASSERT_NE(tp.slices().end(), ++slice);
   ASSERT_EQ(ser_buf.data() + 3 + 5, slice->start);
   ASSERT_EQ(ser_buf.size() - 3 - 5, slice->size);
 
-  ASSERT_EQ(tp.end(), ++slice);
+  ASSERT_EQ(tp.slices().end(), ++slice);
 
   ASSERT_TRUE(tp.Decode());
   ASSERT_NE(nullptr, tp.operator->());
diff --git a/src/tracing/core/trace_writer_impl.cc b/src/tracing/core/trace_writer_impl.cc
index b5933d5..8f78b4e 100644
--- a/src/tracing/core/trace_writer_impl.cc
+++ b/src/tracing/core/trace_writer_impl.cc
@@ -128,16 +128,13 @@
       if (!size_in_current_chunk) {
         auto patch_it = std::find_if(
             patch_list_.begin(), patch_list_.end(),
-            [cur_hdr](const Patch& p) { return p.size_field == cur_hdr; });
+            [cur_hdr](const Patch& p) { return &p.size_field[0] == cur_hdr; });
         PERFETTO_DCHECK(patch_it != patch_list_.end());
       }
 #endif
-      // TODO(primiano): this needs to be adjusted to be the offset within the
-      // payload, not from the start of the chunk header.
-      auto cur_hdr_offset = static_cast<uint16_t>(cur_hdr - cur_chunk_.begin());
-      patch_list_.emplace_front(cur_chunk_id_, cur_hdr_offset);
-      Patch& patch = patch_list_.front();
-      nested_msg->set_size_field(patch.size_field);
+      auto offset = static_cast<uint16_t>(cur_hdr - cur_chunk_.payload_begin());
+      Patch* patch = patch_list_.emplace_back(cur_chunk_id_, offset);
+      nested_msg->set_size_field(&patch->size_field[0]);
       PERFETTO_DLOG("Created new patchlist entry for protobuf nested message");
     }
   }
@@ -177,9 +174,6 @@
   return id_;
 };
 
-TraceWriterImpl::Patch::Patch(uint16_t cid, uint16_t offset)
-    : chunk_id(cid), offset_in_chunk(offset) {}
-
 // Base class ctor/dtor definition.
 TraceWriter::TraceWriter() = default;
 TraceWriter::~TraceWriter() = default;
diff --git a/src/tracing/core/trace_writer_impl.h b/src/tracing/core/trace_writer_impl.h
index e1ea143..18a43b9 100644
--- a/src/tracing/core/trace_writer_impl.h
+++ b/src/tracing/core/trace_writer_impl.h
@@ -17,13 +17,12 @@
 #ifndef SRC_TRACING_CORE_TRACE_WRITER_IMPL_H_
 #define SRC_TRACING_CORE_TRACE_WRITER_IMPL_H_
 
-#include <forward_list>
-
 #include "perfetto/protozero/message_handle.h"
 #include "perfetto/protozero/scattered_stream_writer.h"
 #include "perfetto/tracing/core/basic_types.h"
 #include "perfetto/tracing/core/shared_memory_abi.h"
 #include "perfetto/tracing/core/trace_writer.h"
+#include "src/tracing/core/patch_list.h"
 
 namespace perfetto {
 
@@ -42,21 +41,6 @@
   WriterID writer_id() const override;
 
  private:
-  // Used to handle the backfilling of the headers (the |size_field|) of nested
-  // messages when a proto is fragmented over several chunks. This patchlist is
-  // sent out of band to the tracing service, after having returned the initial
-  // chunks of the fragment.
-  struct Patch {
-    Patch(uint16_t chunk_id, uint16_t offset);
-    Patch(const Patch&) = delete;
-    Patch& operator=(const Patch&) = delete;
-    Patch(Patch&&) noexcept = delete;
-    Patch& operator=(Patch&&) = delete;
-
-    const uint16_t chunk_id;
-    const uint16_t offset_in_chunk;
-    uint8_t size_field[SharedMemoryABI::kPacketHeaderSize] = {};
-  };
   TraceWriterImpl(const TraceWriterImpl&) = delete;
   TraceWriterImpl& operator=(const TraceWriterImpl&) = delete;
 
@@ -105,11 +89,7 @@
   // have to release the ownership of the current Chunk). This list will be
   // later sent out-of-band to the tracing service, who will patch the required
   // chunks, if they are still around.
-  // Note: the Message will take pointers to the |size_field| of these
-  // entries. This container must guarantee that the Patch objects are never
-  // moved around (i.e. cannot be a vector because of reallocations can change
-  // addresses of pre-existing entries).
-  std::forward_list<Patch> patch_list_;
+  PatchList patch_list_;
 };
 
 }  // namespace perfetto
diff --git a/src/tracing/ipc/service/consumer_ipc_service.cc b/src/tracing/ipc/service/consumer_ipc_service.cc
index 75a417f..321e6d3 100644
--- a/src/tracing/ipc/service/consumer_ipc_service.cc
+++ b/src/tracing/ipc/service/consumer_ipc_service.cc
@@ -115,7 +115,7 @@
   for (const TracePacket& trace_packet : trace_packets) {
     std::string* dst = result->add_trace_packets();
     dst->reserve(trace_packet.size());
-    for (const Slice& slice : trace_packet)
+    for (const Slice& slice : trace_packet.slices())
       dst->append(reinterpret_cast<const char*>(slice.start), slice.size);
   }
   read_buffers_response.Resolve(std::move(result));
diff --git a/src/tracing/test/fake_packet.cc b/src/tracing/test/fake_packet.cc
index a79145d..fd5ca05 100644
--- a/src/tracing/test/fake_packet.cc
+++ b/src/tracing/test/fake_packet.cc
@@ -132,8 +132,13 @@
   return *this;
 }
 
+FakeChunk& FakeChunk::SetUID(uid_t u) {
+  uid = u;
+  return *this;
+}
+
 size_t FakeChunk::CopyIntoTraceBuffer() {
-  trace_buffer_->CopyChunkUntrusted(producer_id, writer_id, chunk_id,
+  trace_buffer_->CopyChunkUntrusted(producer_id, uid, writer_id, chunk_id,
                                     num_packets, flags, data.data(),
                                     data.size());
   return data.size() + TraceBuffez::InlineChunkHeaderSize;
diff --git a/src/tracing/test/fake_packet.h b/src/tracing/test/fake_packet.h
index d82bb48..daab5e8 100644
--- a/src/tracing/test/fake_packet.h
+++ b/src/tracing/test/fake_packet.h
@@ -64,6 +64,8 @@
 
   FakeChunk& ClearBytes(size_t offset, size_t len);
 
+  FakeChunk& SetUID(uid_t);
+
   // Returns the full size of the chunk including the ChunkRecord header.
   size_t CopyIntoTraceBuffer();
 
@@ -74,6 +76,7 @@
   ChunkID chunk_id;
   uint8_t flags = 0;
   uint16_t num_packets = 0;
+  uid_t uid = -1;
   std::vector<uint8_t> data;
 };
 
diff --git a/src/tracing/test/tracing_integration_test.cc b/src/tracing/test/tracing_integration_test.cc
index f79a652..29a36c9 100644
--- a/src/tracing/test/tracing_integration_test.cc
+++ b/src/tracing/test/tracing_integration_test.cc
@@ -168,6 +168,7 @@
     sprintf(buf, "evt_%zu", i);
     writer->NewTracePacket()->set_for_testing()->set_str(buf, strlen(buf));
   }
+  writer.reset();
 
   // Allow the service to see the CommitData() before disabling tracing.
   task_runner_->RunUntilIdle();
@@ -185,7 +186,6 @@
   size_t num_pack_rx = 0;
   auto all_packets_rx = task_runner_->CreateCheckpoint("all_packets_rx");
   EXPECT_CALL(consumer, OnTracePackets(_, _))
-      .Times(kNumPackets)
       .WillRepeatedly(
           Invoke([&num_pack_rx, all_packets_rx](
                      std::vector<TracePacket>* packets, bool has_more) {
@@ -195,6 +195,7 @@
               all_packets_rx();
           }));
   task_runner_->RunUntilCheckpoint("all_packets_rx");
+  ASSERT_EQ(kNumPackets, num_pack_rx);
 
   // TODO(primiano): cover FreeBuffers.
 
diff --git a/tools/proto_to_cpp/proto_to_cpp.cc b/tools/proto_to_cpp/proto_to_cpp.cc
index 1b8aa76..565c53f 100644
--- a/tools/proto_to_cpp/proto_to_cpp.cc
+++ b/tools/proto_to_cpp/proto_to_cpp.cc
@@ -328,6 +328,12 @@
       } else {
         p->Print("void set_$n$($t$ value) { $n$_ = value; }\n", "t",
                  GetCppType(field, true), "n", field->lowercase_name());
+        if (field->type() == FieldDescriptor::TYPE_BYTES) {
+          p->Print(
+              "void set_$n$(const void* p, size_t s) { "
+              "$n$_.assign(reinterpret_cast<const char*>(p), s); }\n",
+              "n", field->lowercase_name());
+        }
       }
     } else {  // is_repeated()
       p->Print(