Mojo: Use new message APIs to reduce copying

mojo/public/cpp/bindings/lib/message.cc

Index: mojo/public/cpp/bindings/lib/message.cc
diff --git a/mojo/public/cpp/bindings/lib/message.cc b/mojo/public/cpp/bindings/lib/message.cc
index e1e8f190df93786d34f16f6342610f8d3ba3b48f..4274cc86cb12695c4d0bebd5ea3664e88d89a71c 100644
--- a/mojo/public/cpp/bindings/lib/message.cc
+++ b/mojo/public/cpp/bindings/lib/message.cc
@@ -9,6 +9,7 @@
 #include <stdlib.h>
 
 #include <algorithm>
+#include <utility>
 
 #include "base/logging.h"
 
@@ -23,7 +24,15 @@ Message::~Message() {
 
 void Message::Initialize(size_t capacity, bool zero_initialized) {
   DCHECK(!buffer_);
-  buffer_.reset(new internal::PickleBuffer(capacity, zero_initialized));
+  buffer_.reset(new internal::MessageBuffer(capacity, zero_initialized));
+}
+
+void Message::InitializeFromMojoMessage(ScopedMessageHandle message,
+                                        uint32_t num_bytes,
+                                        std::vector<Handle>* handles) {
+  DCHECK(!buffer_);
+  buffer_.reset(new internal::MessageBuffer(std::move(message), num_bytes));
+  handles_.swap(*handles);
 }
 
 void Message::MoveTo(Message* destination) {
@@ -38,6 +47,37 @@ void Message::MoveTo(Message* destination) {
   buffer_.reset();
 }
 
+ScopedMessageHandle Message::TakeMojoMessage() {
+  if (handles_.empty())  // Fast path for the common case: No handles.
+    return buffer_->TakeMessage();
+
+  // Allocate a new message with space for the handles, then copy the buffer
+  // contents into it.
+  //
+  // TODO(rockot): We could avoid this copy by extending GetSerializedSize()
+  // behavior to collect handles. It's unoptimized for now because it's much
+  // more common to have messages with no handles.
+  ScopedMessageHandle new_message;
+  MojoResult rv = AllocMessage(
+      data_num_bytes(),
+      handles_.empty() ? nullptr
+                       : reinterpret_cast<const MojoHandle*>(handles_.data()),
+      handles_.size(),
+      MOJO_ALLOC_MESSAGE_FLAG_NONE,
+      &new_message);
+  CHECK_EQ(rv, MOJO_RESULT_OK);
+  handles_.clear();
+
+  void* new_buffer = nullptr;
+  rv = GetMessageBuffer(new_message.get(), &new_buffer);
+  CHECK_EQ(rv, MOJO_RESULT_OK);
+
+  memcpy(new_buffer, data(), data_num_bytes());
+  buffer_.reset();
+
+  return new_message;
+}
+
 void Message::CloseHandles() {
   for (std::vector<Handle>::iterator it = handles_.begin();
        it != handles_.end(); ++it) {
@@ -49,28 +89,32 @@ void Message::CloseHandles() {
 MojoResult ReadMessage(MessagePipeHandle handle, Message* message) {
   MojoResult rv;
 
+  std::vector<Handle> handles;
+  ScopedMessageHandle mojo_message;
   uint32_t num_bytes = 0, num_handles = 0;
-  rv = ReadMessageRaw(handle,
-                      nullptr,
+  rv = ReadMessageNew(handle,
+                      &mojo_message,
                       &num_bytes,
                       nullptr,
                       &num_handles,
                       MOJO_READ_MESSAGE_FLAG_NONE);
-  if (rv != MOJO_RESULT_RESOURCE_EXHAUSTED)
-    return rv;
-
-  message->Initialize(num_bytes, false /* zero_initialized */);
+  if (rv == MOJO_RESULT_RESOURCE_EXHAUSTED) {
+    DCHECK_GT(num_handles, 0u);
+    handles.resize(num_handles);
+    rv = ReadMessageNew(handle,
+                        &mojo_message,
+                        &num_bytes,
+                        reinterpret_cast<MojoHandle*>(handles.data()),
+                        &num_handles,
+                        MOJO_READ_MESSAGE_FLAG_NONE);
+  }
 
-  void* mutable_data = message->buffer()->Allocate(num_bytes);
-  message->mutable_handles()->resize(num_handles);
+  if (rv != MOJO_RESULT_OK)
+    return rv;
 
-  rv = ReadMessageRaw(
-      handle, mutable_data, &num_bytes,
-      message->mutable_handles()->empty()
-          ? nullptr
-          : reinterpret_cast<MojoHandle*>(message->mutable_handles()->data()),
-      &num_handles, MOJO_READ_MESSAGE_FLAG_NONE);
-  return rv;
+  message->InitializeFromMojoMessage(
+      std::move(mojo_message), num_bytes, &handles);
+  return MOJO_RESULT_OK;
 }
 
 }  // namespace mojo