protozero: Make ScatteredStreamMemoryDelegate production-ready

Currently, this delegate is only used in testing, but we'd like to use
it e.g. for tracing to a heap buffer before the perfetto service is
available.

This patch makes the delegate compatible with chromium's
HeapScatteredStreamWriterDelegate and renames it to ScatteredHeapBuffer.
In particular, it adds a way to obtain the chunk data without copying
them into a stitched vector.

As a follow-up, we can then remove chromium's HSSWD and replace it with
perfetto's SHB.

Change-Id: I4a4c07f81edd7d9ff5ec8b3f977d72ca1668c09d
diff --git a/Android.bp b/Android.bp
index 749b7af..fb4bd23 100644
--- a/Android.bp
+++ b/Android.bp
@@ -69,7 +69,7 @@
     "src/protozero/message_handle.cc",
     "src/protozero/proto_decoder.cc",
     "src/protozero/proto_field_descriptor.cc",
-    "src/protozero/scattered_stream_memory_delegate.cc",
+    "src/protozero/scattered_heap_buffer.cc",
     "src/protozero/scattered_stream_null_delegate.cc",
     "src/protozero/scattered_stream_writer.cc",
     "src/tracing/core/android_log_config.cc",
@@ -224,7 +224,7 @@
     "src/protozero/message_handle.cc",
     "src/protozero/proto_decoder.cc",
     "src/protozero/proto_field_descriptor.cc",
-    "src/protozero/scattered_stream_memory_delegate.cc",
+    "src/protozero/scattered_heap_buffer.cc",
     "src/protozero/scattered_stream_null_delegate.cc",
     "src/protozero/scattered_stream_writer.cc",
     "src/traced/probes/android_log/android_log_data_source.cc",
@@ -393,7 +393,7 @@
     "src/protozero/message_handle.cc",
     "src/protozero/proto_decoder.cc",
     "src/protozero/proto_field_descriptor.cc",
-    "src/protozero/scattered_stream_memory_delegate.cc",
+    "src/protozero/scattered_heap_buffer.cc",
     "src/protozero/scattered_stream_null_delegate.cc",
     "src/protozero/scattered_stream_writer.cc",
     "src/tracing/core/android_log_config.cc",
@@ -560,7 +560,7 @@
     "src/protozero/message_handle.cc",
     "src/protozero/proto_decoder.cc",
     "src/protozero/proto_field_descriptor.cc",
-    "src/protozero/scattered_stream_memory_delegate.cc",
+    "src/protozero/scattered_heap_buffer.cc",
     "src/protozero/scattered_stream_null_delegate.cc",
     "src/protozero/scattered_stream_writer.cc",
     "src/traced/probes/android_log/android_log_data_source.cc",
@@ -2301,7 +2301,7 @@
     "src/protozero/message_handle.cc",
     "src/protozero/proto_decoder.cc",
     "src/protozero/proto_field_descriptor.cc",
-    "src/protozero/scattered_stream_memory_delegate.cc",
+    "src/protozero/scattered_heap_buffer.cc",
     "src/protozero/scattered_stream_null_delegate.cc",
     "src/protozero/scattered_stream_writer.cc",
     "src/tracing/core/android_log_config.cc",
@@ -2572,7 +2572,7 @@
     "src/protozero/proto_decoder_unittest.cc",
     "src/protozero/proto_field_descriptor.cc",
     "src/protozero/proto_utils_unittest.cc",
-    "src/protozero/scattered_stream_memory_delegate.cc",
+    "src/protozero/scattered_heap_buffer.cc",
     "src/protozero/scattered_stream_null_delegate.cc",
     "src/protozero/scattered_stream_writer.cc",
     "src/protozero/scattered_stream_writer_unittest.cc",
diff --git a/include/perfetto/protozero/BUILD.gn b/include/perfetto/protozero/BUILD.gn
index 06b067f..fef1c07 100644
--- a/include/perfetto/protozero/BUILD.gn
+++ b/include/perfetto/protozero/BUILD.gn
@@ -23,7 +23,7 @@
     "proto_decoder.h",
     "proto_field_descriptor.h",
     "proto_utils.h",
-    "scattered_stream_memory_delegate.h",
+    "scattered_heap_buffer.h",
     "scattered_stream_null_delegate.h",
     "scattered_stream_writer.h",
   ]
diff --git a/include/perfetto/protozero/scattered_heap_buffer.h b/include/perfetto/protozero/scattered_heap_buffer.h
new file mode 100644
index 0000000..a65fb61
--- /dev/null
+++ b/include/perfetto/protozero/scattered_heap_buffer.h
@@ -0,0 +1,89 @@
+/*
+ * Copyright (C) 2017 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 INCLUDE_PERFETTO_PROTOZERO_SCATTERED_HEAP_BUFFER_H_
+#define INCLUDE_PERFETTO_PROTOZERO_SCATTERED_HEAP_BUFFER_H_
+
+#include <memory>
+#include <vector>
+
+#include "perfetto/base/logging.h"
+#include "perfetto/protozero/scattered_stream_writer.h"
+
+namespace protozero {
+
+class ScatteredHeapBuffer : public protozero::ScatteredStreamWriter::Delegate {
+ public:
+  class Slice {
+   public:
+    explicit Slice(size_t size);
+    Slice(Slice&& slice) noexcept;
+    ~Slice();
+
+    inline protozero::ContiguousMemoryRange GetTotalRange() const {
+      return {buffer_.get(), buffer_.get() + size_};
+    }
+
+    inline protozero::ContiguousMemoryRange GetUsedRange() const {
+      return {buffer_.get(), buffer_.get() + size_ - unused_bytes_};
+    }
+
+    uint8_t* start() const { return buffer_.get(); }
+    size_t size() const { return size_; }
+    size_t unused_bytes() const { return unused_bytes_; }
+    void set_unused_bytes(size_t unused_bytes) {
+      PERFETTO_DCHECK(unused_bytes_ <= size_);
+      unused_bytes_ = unused_bytes;
+    }
+
+   private:
+    std::unique_ptr<uint8_t[]> buffer_;
+    const size_t size_;
+    size_t unused_bytes_;
+  };
+
+  ScatteredHeapBuffer(size_t initial_slice_size_bytes = 128,
+                      size_t maximum_slice_size_bytes = 128 * 1024);
+  ~ScatteredHeapBuffer() override;
+
+  // protozero::ScatteredStreamWriter::Delegate implementation.
+  protozero::ContiguousMemoryRange GetNewBuffer() override;
+
+  // Stitch all the slices into a single contiguous buffer.
+  std::vector<uint8_t> StitchSlices();
+
+  const std::vector<Slice>& slices() const { return slices_; }
+
+  void set_writer(protozero::ScatteredStreamWriter* writer) {
+    writer_ = writer;
+  }
+
+  // Update unused_bytes() of the current |Slice| based on the writer's state.
+  void AdjustUsedSizeOfCurrentSlice();
+
+  // Returns the total size the slices occupy in heap memory (including unused).
+  size_t GetTotalSize();
+
+ private:
+  size_t next_slice_size_;
+  const size_t maximum_slice_size_;
+  protozero::ScatteredStreamWriter* writer_ = nullptr;
+  std::vector<Slice> slices_;
+};
+
+}  // namespace protozero
+
+#endif  // INCLUDE_PERFETTO_PROTOZERO_SCATTERED_HEAP_BUFFER_H_
diff --git a/include/perfetto/protozero/scattered_stream_memory_delegate.h b/include/perfetto/protozero/scattered_stream_memory_delegate.h
deleted file mode 100644
index 5910fbd..0000000
--- a/include/perfetto/protozero/scattered_stream_memory_delegate.h
+++ /dev/null
@@ -1,57 +0,0 @@
-/*
- * Copyright (C) 2017 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 INCLUDE_PERFETTO_PROTOZERO_SCATTERED_STREAM_MEMORY_DELEGATE_H_
-#define INCLUDE_PERFETTO_PROTOZERO_SCATTERED_STREAM_MEMORY_DELEGATE_H_
-
-#include <memory>
-#include <vector>
-
-#include "perfetto/base/logging.h"
-#include "perfetto/protozero/scattered_stream_writer.h"
-
-namespace perfetto {
-
-class ScatteredStreamMemoryDelegate
-    : public protozero::ScatteredStreamWriter::Delegate {
- public:
-  explicit ScatteredStreamMemoryDelegate(size_t chunk_size);
-  ~ScatteredStreamMemoryDelegate() override;
-
-  // protozero::ScatteredStreamWriter::Delegate implementation.
-  protozero::ContiguousMemoryRange GetNewBuffer() override;
-
-  // Stitch all the chunks into a single contiguous buffer.
-  std::vector<uint8_t> StitchChunks();
-
-  const std::vector<std::unique_ptr<uint8_t[]>>& chunks() const {
-    return chunks_;
-  }
-
-  void set_writer(protozero::ScatteredStreamWriter* writer) {
-    writer_ = writer;
-  }
-
- private:
-  const size_t chunk_size_;
-  protozero::ScatteredStreamWriter* writer_ = nullptr;
-  std::vector<size_t> chunks_used_size_;
-  std::vector<std::unique_ptr<uint8_t[]>> chunks_;
-};
-
-}  // namespace perfetto
-
-#endif  // INCLUDE_PERFETTO_PROTOZERO_SCATTERED_STREAM_MEMORY_DELEGATE_H_
diff --git a/src/perfetto_cmd/pbtxt_to_pb.cc b/src/perfetto_cmd/pbtxt_to_pb.cc
index cf923d6..c26fa76 100644
--- a/src/perfetto_cmd/pbtxt_to_pb.cc
+++ b/src/perfetto_cmd/pbtxt_to_pb.cc
@@ -29,7 +29,7 @@
 #include "perfetto/base/string_view.h"
 #include "perfetto/protozero/message.h"
 #include "perfetto/protozero/message_handle.h"
-#include "perfetto/protozero/scattered_stream_memory_delegate.h"
+#include "perfetto/protozero/scattered_heap_buffer.h"
 #include "src/perfetto_cmd/perfetto_config.descriptor.h"
 
 namespace perfetto {
@@ -599,7 +599,7 @@
   const DescriptorProto* descriptor = name_to_descriptor[kConfigProtoName];
   PERFETTO_CHECK(descriptor);
 
-  ScatteredStreamMemoryDelegate stream_delegate(base::kPageSize);
+  protozero::ScatteredHeapBuffer stream_delegate(base::kPageSize);
   protozero::ScatteredStreamWriter stream(&stream_delegate);
   stream_delegate.set_writer(&stream);
 
@@ -609,7 +609,7 @@
                           std::move(name_to_descriptor),
                           std::move(name_to_enum));
   Parse(input, &delegate);
-  return stream_delegate.StitchChunks();
+  return stream_delegate.StitchSlices();
 }
 
 }  // namespace perfetto
diff --git a/src/protozero/BUILD.gn b/src/protozero/BUILD.gn
index bfc4fa6..db49c13 100644
--- a/src/protozero/BUILD.gn
+++ b/src/protozero/BUILD.gn
@@ -31,7 +31,7 @@
     "message_handle.cc",
     "proto_decoder.cc",
     "proto_field_descriptor.cc",
-    "scattered_stream_memory_delegate.cc",
+    "scattered_heap_buffer.cc",
     "scattered_stream_null_delegate.cc",
     "scattered_stream_writer.cc",
   ]
diff --git a/src/protozero/proto_decoder_unittest.cc b/src/protozero/proto_decoder_unittest.cc
index 047f2b3..9c5771b 100644
--- a/src/protozero/proto_decoder_unittest.cc
+++ b/src/protozero/proto_decoder_unittest.cc
@@ -21,7 +21,7 @@
 #include "perfetto/base/utils.h"
 #include "perfetto/protozero/message.h"
 #include "perfetto/protozero/proto_utils.h"
-#include "perfetto/protozero/scattered_stream_memory_delegate.h"
+#include "perfetto/protozero/scattered_heap_buffer.h"
 
 namespace protozero {
 namespace {
@@ -33,7 +33,7 @@
 
 TEST(ProtoDecoder, ReadString) {
   Message message;
-  perfetto::ScatteredStreamMemoryDelegate delegate(512);
+  ScatteredHeapBuffer delegate(512, 512);
   ScatteredStreamWriter writer(&delegate);
   delegate.set_writer(&writer);
   message.Reset(&writer);
@@ -41,10 +41,10 @@
   static constexpr char kTestString[] = "test";
   message.AppendString(1, kTestString);
 
-  uint8_t* data = delegate.chunks()[0].get();
-  uint64_t bytes_used = 512 - writer.bytes_available();
+  delegate.AdjustUsedSizeOfCurrentSlice();
+  auto used_range = delegate.slices()[0].GetUsedRange();
 
-  ProtoDecoder decoder(data, bytes_used);
+  ProtoDecoder decoder(used_range.begin, used_range.size());
   ProtoDecoder::Field field = decoder.ReadField();
 
   ASSERT_EQ(field.id, 1u);
diff --git a/src/protozero/scattered_heap_buffer.cc b/src/protozero/scattered_heap_buffer.cc
new file mode 100644
index 0000000..b12e7e4
--- /dev/null
+++ b/src/protozero/scattered_heap_buffer.cc
@@ -0,0 +1,79 @@
+/*
+ * Copyright (C) 2017 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 "perfetto/protozero/scattered_heap_buffer.h"
+
+namespace protozero {
+
+ScatteredHeapBuffer::Slice::Slice(size_t size)
+    : buffer_(std::unique_ptr<uint8_t[]>(new uint8_t[size])),
+      size_(size),
+      unused_bytes_(size) {
+  PERFETTO_DCHECK(size);
+#if PERFETTO_DCHECK_IS_ON()
+  memset(start(), 0xff, size_);
+#endif  // PERFETTO_DCHECK_IS_ON()
+}
+
+ScatteredHeapBuffer::Slice::Slice(Slice&& slice) noexcept = default;
+
+ScatteredHeapBuffer::Slice::~Slice() = default;
+
+ScatteredHeapBuffer::ScatteredHeapBuffer(size_t initial_slice_size_bytes,
+                                         size_t maximum_slice_size_bytes)
+    : next_slice_size_(initial_slice_size_bytes),
+      maximum_slice_size_(maximum_slice_size_bytes) {
+  PERFETTO_DCHECK(next_slice_size_ && maximum_slice_size_);
+  PERFETTO_DCHECK(maximum_slice_size_ >= initial_slice_size_bytes);
+}
+
+ScatteredHeapBuffer::~ScatteredHeapBuffer() = default;
+
+protozero::ContiguousMemoryRange ScatteredHeapBuffer::GetNewBuffer() {
+  PERFETTO_CHECK(writer_);
+  AdjustUsedSizeOfCurrentSlice();
+
+  slices_.emplace_back(next_slice_size_);
+  next_slice_size_ = std::min(maximum_slice_size_, next_slice_size_ * 2);
+  return slices_.back().GetTotalRange();
+}
+
+std::vector<uint8_t> ScatteredHeapBuffer::StitchSlices() {
+  AdjustUsedSizeOfCurrentSlice();
+  std::vector<uint8_t> buffer;
+  size_t i = 0;
+  for (const auto& slice : slices_) {
+    auto used_range = slice.GetUsedRange();
+    buffer.insert(buffer.end(), used_range.begin, used_range.end);
+    i++;
+  }
+  return buffer;
+}
+
+void ScatteredHeapBuffer::AdjustUsedSizeOfCurrentSlice() {
+  if (!slices_.empty())
+    slices_.back().set_unused_bytes(writer_->bytes_available());
+}
+
+size_t ScatteredHeapBuffer::GetTotalSize() {
+  size_t total_size = 0;
+  for (auto& slice : slices_) {
+    total_size += slice.size();
+  }
+  return total_size;
+}
+
+}  // namespace protozero
diff --git a/src/protozero/scattered_stream_memory_delegate.cc b/src/protozero/scattered_stream_memory_delegate.cc
deleted file mode 100644
index 63e078f..0000000
--- a/src/protozero/scattered_stream_memory_delegate.cc
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * Copyright (C) 2017 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 "perfetto/protozero/scattered_stream_memory_delegate.h"
-
-namespace perfetto {
-
-ScatteredStreamMemoryDelegate::ScatteredStreamMemoryDelegate(size_t chunk_size)
-    : chunk_size_(chunk_size) {}
-
-ScatteredStreamMemoryDelegate::~ScatteredStreamMemoryDelegate() {}
-
-protozero::ContiguousMemoryRange ScatteredStreamMemoryDelegate::GetNewBuffer() {
-  PERFETTO_CHECK(writer_);
-  if (!chunks_.empty()) {
-    size_t used = chunk_size_ - writer_->bytes_available();
-    chunks_used_size_.push_back(used);
-  }
-  std::unique_ptr<uint8_t[]> chunk(new uint8_t[chunk_size_]);
-  uint8_t* begin = chunk.get();
-  memset(begin, 0xff, chunk_size_);
-  chunks_.push_back(std::move(chunk));
-  return {begin, begin + chunk_size_};
-}
-
-std::vector<uint8_t> ScatteredStreamMemoryDelegate::StitchChunks() {
-  std::vector<uint8_t> buffer;
-  size_t i = 0;
-  for (const auto& chunk : chunks_) {
-    size_t chunk_size = (i < chunks_used_size_.size())
-                            ? chunks_used_size_[i]
-                            : (chunk_size_ - writer_->bytes_available());
-    PERFETTO_CHECK(chunk_size <= chunk_size_);
-    buffer.insert(buffer.end(), chunk.get(), chunk.get() + chunk_size);
-    i++;
-  }
-  return buffer;
-}
-
-}  // namespace perfetto
diff --git a/src/traced/probes/ftrace/cpu_reader_unittest.cc b/src/traced/probes/ftrace/cpu_reader_unittest.cc
index e9bd27d..fe2ca76 100644
--- a/src/traced/probes/ftrace/cpu_reader_unittest.cc
+++ b/src/traced/probes/ftrace/cpu_reader_unittest.cc
@@ -26,7 +26,7 @@
 #include "perfetto/base/build_config.h"
 #include "perfetto/base/utils.h"
 #include "perfetto/protozero/proto_utils.h"
-#include "perfetto/protozero/scattered_stream_memory_delegate.h"
+#include "perfetto/protozero/scattered_heap_buffer.h"
 #include "perfetto/protozero/scattered_stream_writer.h"
 
 #include "perfetto/trace/ftrace/ftrace_event.pb.h"
@@ -117,7 +117,7 @@
   // on success and nullptr on failure.
   std::unique_ptr<ProtoT> ParseProto() {
     auto bundle = std::unique_ptr<ProtoT>(new ProtoT());
-    std::vector<uint8_t> buffer = delegate_.StitchChunks();
+    std::vector<uint8_t> buffer = delegate_.StitchSlices();
     if (!bundle->ParseFromArray(buffer.data(), static_cast<int>(buffer.size())))
       return nullptr;
     return bundle;
@@ -128,7 +128,7 @@
   ProtoProvider& operator=(const ProtoProvider&) = delete;
 
   size_t chunk_size_;
-  ScatteredStreamMemoryDelegate delegate_;
+  protozero::ScatteredHeapBuffer delegate_;
   protozero::ScatteredStreamWriter stream_;
   ZeroT writer_;
 };
diff --git a/src/tracing/core/trace_writer_for_testing.cc b/src/tracing/core/trace_writer_for_testing.cc
index 3f95c0f..2909a4a 100644
--- a/src/tracing/core/trace_writer_for_testing.cc
+++ b/src/tracing/core/trace_writer_for_testing.cc
@@ -26,7 +26,9 @@
 namespace perfetto {
 
 TraceWriterForTesting::TraceWriterForTesting()
-    : delegate_(static_cast<size_t>(base::kPageSize)), stream_(&delegate_) {
+    : delegate_(static_cast<size_t>(base::kPageSize),
+                static_cast<size_t>(base::kPageSize)),
+      stream_(&delegate_) {
   delegate_.set_writer(&stream_);
   cur_packet_.reset(new protos::pbzero::TracePacket());
   cur_packet_->Finalize();  // To avoid the DCHECK in NewTracePacket().
@@ -45,7 +47,7 @@
 std::unique_ptr<protos::TracePacket> TraceWriterForTesting::ParseProto() {
   PERFETTO_CHECK(cur_packet_->is_finalized());
   auto packet = std::unique_ptr<protos::TracePacket>(new protos::TracePacket());
-  std::vector<uint8_t> buffer = delegate_.StitchChunks();
+  std::vector<uint8_t> buffer = delegate_.StitchSlices();
   if (!packet->ParseFromArray(buffer.data(), static_cast<int>(buffer.size())))
     return nullptr;
   return packet;
diff --git a/src/tracing/core/trace_writer_for_testing.h b/src/tracing/core/trace_writer_for_testing.h
index f667e31..0255def 100644
--- a/src/tracing/core/trace_writer_for_testing.h
+++ b/src/tracing/core/trace_writer_for_testing.h
@@ -17,18 +17,18 @@
 #define SRC_TRACING_CORE_TRACE_WRITER_FOR_TESTING_H_
 
 #include "perfetto/protozero/message_handle.h"
-#include "perfetto/protozero/scattered_stream_memory_delegate.h"
+#include "perfetto/protozero/scattered_heap_buffer.h"
 #include "perfetto/trace/trace_packet.pb.h"
 #include "perfetto/tracing/core/trace_writer.h"
 
 namespace perfetto {
 
 // A specialization of TraceWriter for testing which writes into memory
-// allocated by the ScatteredStreamMemoryDelegate.
+// allocated by the ScatteredHeapBuffer.
 // See //include/perfetto/tracing/core/trace_writer.h for docs.
 class TraceWriterForTesting : public TraceWriter {
  public:
-  // TraceWriterForTesting(const ScatteredStreamMemoryDelegate& delegate);
+  // TraceWriterForTesting(const protozero::ScatteredHeapBuffer& delegate);
   TraceWriterForTesting();
   ~TraceWriterForTesting() override;
 
@@ -45,7 +45,7 @@
   TraceWriterForTesting(const TraceWriterForTesting&) = delete;
   TraceWriterForTesting& operator=(const TraceWriterForTesting&) = delete;
 
-  ScatteredStreamMemoryDelegate delegate_;
+  protozero::ScatteredHeapBuffer delegate_;
   protozero::ScatteredStreamWriter stream_;
 
   // The packet returned via NewTracePacket(). Its owned by this class,