SERVER-19057 avoid extra buffer copy when creating message in rpc builders

3 files changed, 34 insertions, 16 deletions

diff --git a/src/mongo/rpc/command_reply_builder.cpp b/src/mongo/rpc/command_reply_builder.cpp
index c02259ad7ab..8bc50590333 100644
--- a/src/mongo/rpc/command_reply_builder.cpp
+++ b/src/mongo/rpc/command_reply_builder.cpp
@@ -35,14 +35,17 @@
 #include "mongo/stdx/memory.h"
 #include "mongo/util/assert_util.h"
 #include "mongo/util/mongoutils/str.h"
+#include "mongo/util/net/message.h"
 
 namespace mongo {
 namespace rpc {
 
-CommandReplyBuilder::CommandReplyBuilder() : _message{stdx::make_unique<Message>()} {}
+CommandReplyBuilder::CommandReplyBuilder() : CommandReplyBuilder(stdx::make_unique<Message>()) {}
 
 CommandReplyBuilder::CommandReplyBuilder(std::unique_ptr<Message> message)
-    : _message{std::move(message)} {}
+    : _message{std::move(message)} {
+    _builder.skip(mongo::MsgData::MsgDataHeaderSize);
+}
 
 CommandReplyBuilder& CommandReplyBuilder::setMetadata(const BSONObj& metadata) {
     invariant(_state == State::kMetadata);
@@ -101,15 +104,18 @@ void CommandReplyBuilder::reset() {
         return;
     }
     _builder.reset();
+    _builder.skip(mongo::MsgData::MsgDataHeaderSize);
     _message = stdx::make_unique<Message>();
     _state = State::kMetadata;
 }
 
 std::unique_ptr<Message> CommandReplyBuilder::done() {
     invariant(_state == State::kOutputDocs);
-    // TODO: we can elide a large copy here by transferring the internal buffer of
-    // the BufBuilder to the Message.
-    _message->setData(dbCommandReply, _builder.buf(), _builder.len());
+    MsgData::View msg = _builder.buf();
+    msg.setLen(_builder.len());
+    msg.setOperation(dbCommandReply);
+    _builder.decouple();                      // release ownership from BufBuilder
+    _message->setData(msg.view2ptr(), true);  // transfer ownership to Message
     _state = State::kDone;
     return std::move(_message);
 }
@@ -118,9 +124,8 @@ std::size_t CommandReplyBuilder::availableBytes() const {
     int intLen = _builder.len();
     invariant(0 <= intLen);
     std::size_t len = static_cast<std::size_t>(intLen);
-    std::size_t msgHeaderSz = static_cast<std::size_t>(MsgData::MsgDataHeaderSize);
-    invariant(len + msgHeaderSz <= mongo::MaxMessageSizeBytes);
-    return mongo::MaxMessageSizeBytes - len - msgHeaderSz;
+    invariant(len <= mongo::MaxMessageSizeBytes);
+    return mongo::MaxMessageSizeBytes - len;
 }
 
 Status CommandReplyBuilder::_hasSpaceFor(std::size_t dataSize) const {
diff --git a/src/mongo/rpc/command_request_builder.cpp b/src/mongo/rpc/command_request_builder.cpp
index da280c3054c..25f67c11e7f 100644
--- a/src/mongo/rpc/command_request_builder.cpp
+++ b/src/mongo/rpc/command_request_builder.cpp
@@ -34,16 +34,20 @@
 
 #include "mongo/stdx/memory.h"
 #include "mongo/util/assert_util.h"
+#include "mongo/util/net/message.h"
 
 namespace mongo {
 namespace rpc {
 
-CommandRequestBuilder::CommandRequestBuilder() : _message{stdx::make_unique<Message>()} {}
+CommandRequestBuilder::CommandRequestBuilder()
+    : CommandRequestBuilder(stdx::make_unique<Message>()) {}
 
 CommandRequestBuilder::~CommandRequestBuilder() {}
 
 CommandRequestBuilder::CommandRequestBuilder(std::unique_ptr<Message> message)
-    : _message{std::move(message)} {}
+    : _message{std::move(message)} {
+    _builder.skip(mongo::MsgData::MsgDataHeaderSize);  // Leave room for message header.
+}
 
 CommandRequestBuilder& CommandRequestBuilder::setDatabase(StringData database) {
     invariant(_state == State::kDatabase);
@@ -96,9 +100,11 @@ Protocol CommandRequestBuilder::getProtocol() const {
 
 std::unique_ptr<Message> CommandRequestBuilder::done() {
     invariant(_state == State::kInputDocs);
-    // TODO: we can elide a large copy here by transferring the internal buffer of
-    // the BufBuilder to the Message.
-    _message->setData(dbCommand, _builder.buf(), _builder.len());
+    MsgData::View msg = _builder.buf();
+    msg.setLen(_builder.len());
+    msg.setOperation(dbCommand);
+    _builder.decouple();                      // release ownership from BufBuilder.
+    _message->setData(msg.view2ptr(), true);  // transfer ownership to Message.
     _state = State::kDone;
     return std::move(_message);
 }
diff --git a/src/mongo/rpc/legacy_request_builder.cpp b/src/mongo/rpc/legacy_request_builder.cpp
index 05f1fe79ce3..db333dc4c6a 100644
--- a/src/mongo/rpc/legacy_request_builder.cpp
+++ b/src/mongo/rpc/legacy_request_builder.cpp
@@ -37,16 +37,19 @@
 #include "mongo/rpc/metadata.h"
 #include "mongo/stdx/memory.h"
 #include "mongo/util/assert_util.h"
+#include "mongo/util/net/message.h"
 
 namespace mongo {
 namespace rpc {
 
-LegacyRequestBuilder::LegacyRequestBuilder() : _message{stdx::make_unique<Message>()} {}
+LegacyRequestBuilder::LegacyRequestBuilder() : LegacyRequestBuilder(stdx::make_unique<Message>()) {}
 
 LegacyRequestBuilder::~LegacyRequestBuilder() {}
 
 LegacyRequestBuilder::LegacyRequestBuilder(std::unique_ptr<Message> message)
-    : _message{std::move(message)} {}
+    : _message{std::move(message)} {
+    _builder.skip(mongo::MsgData::MsgDataHeaderSize);
+}
 
 LegacyRequestBuilder& LegacyRequestBuilder::setDatabase(StringData database) {
     invariant(_state == State::kDatabase);
@@ -110,7 +113,11 @@ Protocol LegacyRequestBuilder::getProtocol() const {
 
 std::unique_ptr<Message> LegacyRequestBuilder::done() {
     invariant(_state == State::kInputDocs);
-    _message->setData(dbQuery, _builder.buf(), _builder.len());
+    MsgData::View msg = _builder.buf();
+    msg.setLen(_builder.len());
+    msg.setOperation(dbQuery);
+    _builder.decouple();                      // release ownership from BufBuilder
+    _message->setData(msg.view2ptr(), true);  // transfer ownership to Message
     _state = State::kDone;
     return std::move(_message);
 }