protozero: let transports read directly into ProtoRingBuffer

Every transport kept a staging buffer, read into it, then handed those bytes
to Append(), which copied them again. BeginWrite()/EndWrite() lets read(2)
deposit them in the ring buffer directly, so they are copied once.

unixd additionally kept a ProtoRingBuffer per connection, re-serialized each
tokenized message back into its TraceProcessorRpcStream framing and pushed it
through a second ring buffer inside Rpc. It now reads straight into Rpc's
tokenizer: a session serves one client at a time, so there is no second byte
stream that could interleave with it, and the per-connection framing buys
nothing.

  out/linux_clang_release/perfetto_benchmarks \
      --benchmark_filter='ProtoRingBufferIngest|ProtoRingBufferDispatch'
  ingest, 4KB reads    142 ns -> 98.3 ns
  ingest, 1MB reads   38.4 us -> 19.2 us
  dispatch             171 ns -> 99.1 ns

traceconv also drops a 64MB read buffer: peak RSS 62.7MB -> 13.6MB.

Author: LalitMaganti

include/perfetto/ext/protozero/proto_ring_buffer.h

@@ -17,6 +17,7 @@
 #ifndef INCLUDE_PERFETTO_EXT_PROTOZERO_PROTO_RING_BUFFER_H_
 #define INCLUDE_PERFETTO_EXT_PROTOZERO_PROTO_RING_BUFFER_H_
 
+#include <stddef.h>
 #include <stdint.h>
 
 #include "perfetto/ext/base/paged_memory.h"
@@ -115,6 +116,14 @@ class RingBufferMessageReader {
   // Will invaildate the pointers previously handed out.
   void Append(const void* data, size_t len);
 
+  // Zero-copy counterpart of Append(), for transports that can deposit their
+  // bytes straight into the ring buffer rather than into a staging buffer of
+  // their own. BeginWrite() reserves `size` bytes at the write cursor; the
+  // caller writes at most that many there (e.g. by passing the pointer to
+  // read(2)) and calls EndWrite() with the number actually written.
+  uint8_t* BeginWrite(size_t size);
+  void EndWrite(size_t size_written);
+
   // If a message can be read, it returns the boundaries of the message
   // (without including the preamble) and advances the read cursor.
   // If no message is available, returns a null range.
@@ -136,6 +145,7 @@ class RingBufferMessageReader {
   bool failed_ = false;  // Set in case of an unrecoverable framing faiulre.
   size_t rd_ = 0;        // Offset of the read cursor in |buf_|.
   size_t wr_ = 0;        // Offset of the write cursor in |buf_|.
+  size_t writable_ = 0;  // Bytes reserved by the last BeginWrite().
 };
 
 class ProtoRingBuffer final : public RingBufferMessageReader {

src/protozero/proto_ring_buffer.cc

@@ -16,6 +16,7 @@
 
 #include "perfetto/ext/protozero/proto_ring_buffer.h"
 
+#include "perfetto/base/compiler.h"
 #include "perfetto/base/logging.h"
 #include "perfetto/ext/base/paged_memory.h"
 #include "perfetto/protozero/proto_utils.h"
@@ -76,40 +77,28 @@ RingBufferMessageReader::RingBufferMessageReader()
     : buf_(perfetto::base::PagedMemory::Allocate(kGrowBytes)) {}
 RingBufferMessageReader::~RingBufferMessageReader() = default;
 
-void RingBufferMessageReader::Append(const void* data_void, size_t data_len) {
-  if (failed_)
-    return;
-  const uint8_t* data = static_cast<const uint8_t*>(data_void);
+uint8_t* RingBufferMessageReader::BeginWrite(size_t data_len) {
+  PERFETTO_CHECK(data_len <= kMaxMsgSize);
   PERFETTO_DCHECK(wr_ <= buf_.size());
   PERFETTO_DCHECK(wr_ >= rd_);
 
+  // Nothing can be tokenized any more, so recycle the whole buffer for the
+  // bytes EndWrite() is about to drop.
+  if (PERFETTO_UNLIKELY(failed_))
+    rd_ = wr_ = 0;
+
   // If the last call to ReadMessage() consumed all the data in the buffer and
   // there are no incomplete messages pending, restart from the beginning rather
   // than keep ringing. This is the most common case.
-  if (rd_ == wr_)
+  if (PERFETTO_LIKELY(rd_ == wr_))
     rd_ = wr_ = 0;
 
-  // The caller is expected to always issue a ReadMessage() after each Append().
-  PERFETTO_CHECK(!fastpath_.valid());
-  if (rd_ == wr_) {
-    auto msg = TryReadMessage(data, data + data_len);
-    if (msg.valid() && msg.end() == (data + data_len)) {
-      // Fastpath: in many cases, the underlying stream will effectively
-      // preserve the atomicity of messages for most small messages.
-      // In this case we can avoid the extra buf_ roundtrip and just pass a
-      // pointer to |data| + (proto preamble len).
-      // The next call to ReadMessage)= will return |fastpath_|.
-      fastpath_ = std::move(msg);
-      return;
-    }
-  }
-
   size_t avail = buf_.size() - wr_;
   if (data_len > avail) {
     // This whole section should be hit extremely rarely.
 
     // Try first just recompacting the buffer by moving everything to the left.
-    // This can happen if we received "a message and a bit" on each Append call
+    // This can happen if we received "a message and a bit" on each write call
     // so we ended pup in a situation like:
     // buf_: [unused space] [msg1 incomplete]
     //                      ^rd_             ^wr_
@@ -133,21 +122,53 @@ void RingBufferMessageReader::Append(const void* data_void, size_t data_len) {
       while (data_len > new_size - wr_)
         new_size += kGrowBytes;
       if (new_size > kMaxMsgSize * 2) {
+        // These bytes can never amount to a message (e.g. a never-ending
+        // varint). |buf_| exceeds kMaxMsgSize by now, so dropping them leaves
+        // room for the write.
         failed_ = true;
-        return;
+        rd_ = wr_ = 0;
+        return static_cast<uint8_t*>(buf_.Get());
       }
       auto new_buf = perfetto::base::PagedMemory::Allocate(new_size);
       memcpy(new_buf.Get(), buf_.Get(), buf_.size());
       buf_ = std::move(new_buf);
-      avail = new_size - wr_;
       // No need to touch rd_ / wr_ cursors.
     }
   }
 
-  // Append the received data at the end of the ring buffer.
-  uint8_t* buf = static_cast<uint8_t*>(buf_.Get());
-  memcpy(&buf[wr_], data, data_len);
-  wr_ += data_len;
+  writable_ = data_len;
+  return static_cast<uint8_t*>(buf_.Get()) + wr_;
+}
+
+void RingBufferMessageReader::EndWrite(size_t size_written) {
+  PERFETTO_CHECK(size_written <= writable_);
+  writable_ = 0;
+  if (PERFETTO_LIKELY(!failed_))
+    wr_ += size_written;
+}
+
+void RingBufferMessageReader::Append(const void* data_void, size_t data_len) {
+  if (failed_)
+    return;
+  const uint8_t* data = static_cast<const uint8_t*>(data_void);
+
+  // The caller is expected to always issue a ReadMessage() after each Append().
+  PERFETTO_CHECK(!fastpath_.valid());
+  if (rd_ == wr_) {
+    auto msg = TryReadMessage(data, data + data_len);
+    if (msg.valid() && msg.end() == (data + data_len)) {
+      // Fastpath: in many cases, the underlying stream will effectively
+      // preserve the atomicity of messages for most small messages.
+      // In this case we can avoid the extra buf_ roundtrip and just pass a
+      // pointer to |data| + (proto preamble len).
+      // The next call to ReadMessage)= will return |fastpath_|.
+      fastpath_ = std::move(msg);
+      return;
+    }
+  }
+
+  memcpy(BeginWrite(data_len), data, data_len);
+  EndWrite(data_len);
 }
 
 RingBufferMessageReader::Message RingBufferMessageReader::ReadMessage() {

src/protozero/proto_ring_buffer_unittest.cc

@@ -19,6 +19,8 @@
 #include <stdint.h>
 #include <sys/types.h>
 
+#include <algorithm>
+#include <cstring>
 #include <list>
 #include <ostream>
 #include <random>
@@ -264,5 +266,28 @@ TEST(RingBufferTest, FixedLengthRingBuffer) {
             "abc");
 }
 
+// A read(2) can come back short, so EndWrite() may report fewer bytes than
+// BeginWrite() reserved. The tests above always write what they reserve.
+TEST_F(ProtoRingBufferTest, PartialWrites) {
+  ProtoRingBuffer buf;
+  auto expected = MakeProtoMessage(/*field_id=*/7, /*len=*/4096);
+
+  buf.BeginWrite(4096);
+  buf.EndWrite(0);  // The socket had nothing for us.
+  EXPECT_FALSE(buf.ReadMessage().valid());
+
+  for (size_t written = 0; written < last_msg_.size();) {
+    // Reserve far more than we intend to write, as a socket reader would.
+    uint8_t* dst = buf.BeginWrite(4096);
+    size_t n = std::min<size_t>(37, last_msg_.size() - written);
+    memcpy(dst, &last_msg_[written], n);
+    buf.EndWrite(n);
+    written += n;
+    if (written < last_msg_.size())
+      ASSERT_FALSE(buf.ReadMessage().valid());
+  }
+  EXPECT_EQ(buf.ReadMessage(), expected);
+}
+
 }  // namespace
 }  // namespace protozero

src/protozero/test/proto_ring_buffer_benchmark.cc

@@ -15,12 +15,147 @@
 #include <benchmark/benchmark.h>
 
 #include <algorithm>
+#include <cstring>
 #include <string>
+#include <vector>
 
+#include "perfetto/base/compiler.h"
 #include "perfetto/ext/base/file_utils.h"
 #include "perfetto/ext/protozero/proto_ring_buffer.h"
+#include "perfetto/protozero/proto_utils.h"
 #include "src/base/test/utils.h"
 
+namespace {
+
+std::string LoadTestTrace() {
+  std::string trace_data;
+  static const char kTestTrace[] = "test/data/example_android_trace_30s.pb";
+  perfetto::base::ReadFile(perfetto::base::GetTestDataPath(kTestTrace),
+                           &trace_data);
+  PERFETTO_CHECK(!trace_data.empty());
+  return trace_data;
+}
+
+// Stands in for the read(2)/recv(2) that brings bytes in from the outside
+// world: like those, it fills a buffer chosen by the caller. That is the one
+// copy no scheme avoids; the benchmarks below differ only in whether the buffer
+// belongs to the transport or to the ring buffer.
+PERFETTO_NO_INLINE void ProduceBytes(void* dst, const void* src, size_t size) {
+  memcpy(dst, src, size);
+}
+
+// Ingestion as the transports used to do it: read into a staging buffer of the
+// transport's own, then hand those bytes to Append(), which copies them again.
+void BM_ProtoRingBufferIngestStaged(benchmark::State& state) {
+  const std::string trace_data = LoadTestTrace();
+  const auto read_size = static_cast<size_t>(state.range(0));
+  std::vector<uint8_t> staging(read_size);
+
+  protozero::ProtoRingBuffer buffer;
+  size_t offset = 0, bytes_ingested = 0, total_packet_size = 0;
+  for (auto _ : state) {
+    size_t n = std::min(read_size, trace_data.size() - offset);
+    ProduceBytes(staging.data(), trace_data.data() + offset, n);
+    buffer.Append(staging.data(), n);
+    // The trace holds whole messages only, so by the time we get back to the
+    // start the buffer has drained and the wrap is seamless.
+    offset = (offset + n) % trace_data.size();
+    bytes_ingested += n;
+    for (;;) {
+      auto msg = buffer.ReadMessage();
+      if (!msg.valid())
+        break;
+      total_packet_size += msg.len;
+    }
+  }
+  benchmark::DoNotOptimize(total_packet_size);
+  state.SetBytesProcessed(static_cast<int64_t>(bytes_ingested));
+}
+
+// Ingestion as they do it now: read straight into the ring buffer, so the bytes
+// are copied exactly once.
+void BM_ProtoRingBufferIngestZeroCopy(benchmark::State& state) {
+  const std::string trace_data = LoadTestTrace();
+  const auto read_size = static_cast<size_t>(state.range(0));
+
+  protozero::ProtoRingBuffer buffer;
+  size_t offset = 0, bytes_ingested = 0, total_packet_size = 0;
+  for (auto _ : state) {
+    size_t n = std::min(read_size, trace_data.size() - offset);
+    ProduceBytes(buffer.BeginWrite(n), trace_data.data() + offset, n);
+    buffer.EndWrite(n);
+    offset = (offset + n) % trace_data.size();
+    bytes_ingested += n;
+    for (;;) {
+      auto msg = buffer.ReadMessage();
+      if (!msg.valid())
+        break;
+      total_packet_size += msg.len;
+    }
+  }
+  benchmark::DoNotOptimize(total_packet_size);
+  state.SetBytesProcessed(static_cast<int64_t>(bytes_ingested));
+}
+
+// unixd used to tokenize each message out of a per-connection ProtoRingBuffer,
+// re-serialize its TraceProcessorRpcStream preamble, and push preamble +
+// payload through a second ProtoRingBuffer inside Rpc, so every message was
+// copied and tokenized twice. It now reads straight into Rpc's tokenizer.
+void RunDispatch(benchmark::State& state, bool reframe) {
+  namespace pu = protozero::proto_utils;
+  const std::string trace_data = LoadTestTrace();
+  constexpr size_t kReadSize = 4096;
+
+  protozero::ProtoRingBuffer conn_buf, rpc_buf;
+  size_t offset = 0, bytes_ingested = 0, total_packet_size = 0;
+  for (auto _ : state) {
+    size_t n = std::min(kReadSize, trace_data.size() - offset);
+    ProduceBytes(conn_buf.BeginWrite(n), trace_data.data() + offset, n);
+    conn_buf.EndWrite(n);
+    offset = (offset + n) % trace_data.size();
+    bytes_ingested += n;
+    for (;;) {
+      auto msg = conn_buf.ReadMessage();
+      if (!msg.valid())
+        break;
+      if (!reframe) {
+        total_packet_size += msg.len;
+        continue;
+      }
+      uint8_t preamble[16];
+      uint8_t* end = preamble;
+      end = pu::WriteVarInt(pu::MakeTagLengthDelimited(1), end);
+      end = pu::WriteVarInt(msg.len, end);
+      rpc_buf.Append(preamble, static_cast<size_t>(end - preamble));
+      rpc_buf.Append(msg.start, msg.len);
+      for (;;) {
+        auto inner = rpc_buf.ReadMessage();
+        if (!inner.valid())
+          break;
+        total_packet_size += inner.len;
+      }
+    }
+  }
+  benchmark::DoNotOptimize(total_packet_size);
+  state.SetBytesProcessed(static_cast<int64_t>(bytes_ingested));
+}
+
+void BM_ProtoRingBufferDispatchReframed(benchmark::State& state) {
+  RunDispatch(state, /*reframe=*/true);
+}
+
+void BM_ProtoRingBufferDispatchDirect(benchmark::State& state) {
+  RunDispatch(state, /*reframe=*/false);
+}
+
+}  // namespace
+
+// 4096 is what the socket transports read; 1MB is what traceconv reads.
+BENCHMARK(BM_ProtoRingBufferIngestStaged)->Arg(4096)->Arg(1024 * 1024);
+BENCHMARK(BM_ProtoRingBufferIngestZeroCopy)->Arg(4096)->Arg(1024 * 1024);
+BENCHMARK(BM_ProtoRingBufferDispatchReframed);
+BENCHMARK(BM_ProtoRingBufferDispatchDirect);
+
 static void BM_ProtoRingBufferReadLargeChunks(benchmark::State& state) {
   std::string trace_data;
   static const char kTestTrace[] = "test/data/example_android_trace_30s.pb";

src/trace_processor/rpc/remote_trace_processor.cc

@@ -291,9 +291,6 @@ base::Status RemoteTraceProcessor::SendStream(
 }
 
 base::Status RemoteTraceProcessor::ReadResponse(std::vector<uint8_t>* out) {
-  // Hoisted out of the loop: ProtoRingBuffer's fastpath can return a message
-  // pointing into the last buffer passed to Append(), so it must stay alive.
-  uint8_t buf[4096];
   for (;;) {
     auto msg = rxbuf_.ReadMessage();
     if (msg.fatal_framing_error)
@@ -302,10 +299,12 @@ base::Status RemoteTraceProcessor::ReadResponse(std::vector<uint8_t>* out) {
       out->assign(msg.start, msg.start + msg.len);
       return base::OkStatus();
     }
-    ASSIGN_OR_RETURN(size_t n, transport_->Recv(buf, sizeof(buf)));
+    constexpr size_t kReadSize = 4096;
+    uint8_t* buf = rxbuf_.BeginWrite(kReadSize);
+    ASSIGN_OR_RETURN(size_t n, transport_->Recv(buf, kReadSize));
+    rxbuf_.EndWrite(n);
     if (n == 0)
       return base::ErrStatus("Session closed the connection");
-    rxbuf_.Append(buf, n);
   }
 }
 

src/trace_processor/rpc/rpc.cc

@@ -253,6 +253,19 @@ void Rpc::ResetTraceProcessorInternal(const Config& config) {
 
 void Rpc::OnRpcRequest(const void* data, size_t len) {
   rxbuf_.Append(data, len);
+  DrainRxBuf();
+}
+
+uint8_t* Rpc::BeginRpcRequest(size_t size) {
+  return rxbuf_.BeginWrite(size);
+}
+
+void Rpc::EndRpcRequest(size_t size_written) {
+  rxbuf_.EndWrite(size_written);
+  DrainRxBuf();
+}
+
+void Rpc::DrainRxBuf() {
   for (;;) {
     auto msg = rxbuf_.ReadMessage();
     if (!msg.valid()) {