SERVER-18930 SERVER-19060 SERVER-19132 SERVER-19133 Implement LegacyReply(Builder) API and up/down conversions

12 files changed, 432 insertions, 133 deletions

diff --git a/src/mongo/rpc/command_reply_builder.cpp b/src/mongo/rpc/command_reply_builder.cpp
index d21358e053a..c02259ad7ab 100644
--- a/src/mongo/rpc/command_reply_builder.cpp
+++ b/src/mongo/rpc/command_reply_builder.cpp
@@ -28,11 +28,13 @@
 
 #include "mongo/platform/basic.h"
 
+#include "mongo/rpc/command_reply_builder.h"
+
 #include <utility>
 
-#include "mongo/rpc/command_reply_builder.h"
 #include "mongo/stdx/memory.h"
 #include "mongo/util/assert_util.h"
+#include "mongo/util/mongoutils/str.h"
 
 namespace mongo {
 namespace rpc {
@@ -42,32 +44,46 @@ CommandReplyBuilder::CommandReplyBuilder() : _message{stdx::make_unique<Message>
 CommandReplyBuilder::CommandReplyBuilder(std::unique_ptr<Message> message)
     : _message{std::move(message)} {}
 
-CommandReplyBuilder& CommandReplyBuilder::setMetadata(BSONObj metadata) {
+CommandReplyBuilder& CommandReplyBuilder::setMetadata(const BSONObj& metadata) {
     invariant(_state == State::kMetadata);
+
     metadata.appendSelfToBufBuilder(_builder);
     _state = State::kCommandReply;
     return *this;
 }
 
-CommandReplyBuilder& CommandReplyBuilder::setRawCommandReply(BSONObj commandReply) {
+CommandReplyBuilder& CommandReplyBuilder::setRawCommandReply(const BSONObj& commandReply) {
     invariant(_state == State::kCommandReply);
+
     commandReply.appendSelfToBufBuilder(_builder);
     _state = State::kOutputDocs;
     return *this;
 }
 
-CommandReplyBuilder& CommandReplyBuilder::addOutputDocs(DocumentRange outputDocs) {
+Status CommandReplyBuilder::addOutputDocs(DocumentRange outputDocs) {
     invariant(_state == State::kOutputDocs);
     auto rangeData = outputDocs.data();
-    _builder.appendBuf(rangeData.data(), rangeData.length());
-    // leave state as is as we can add as many outputDocs as we want.
-    return *this;
+    auto dataSize = rangeData.length();
+    auto hasSpace = _hasSpaceFor(dataSize);
+    if (!hasSpace.isOK()) {
+        return hasSpace;
+    }
+
+    _builder.appendBuf(rangeData.data(), dataSize);
+    return Status::OK();
 }
 
-CommandReplyBuilder& CommandReplyBuilder::addOutputDoc(BSONObj outputDoc) {
+Status CommandReplyBuilder::addOutputDoc(const BSONObj& outputDoc) {
     invariant(_state == State::kOutputDocs);
+
+    auto dataSize = static_cast<std::size_t>(outputDoc.objsize());
+    auto hasSpace = _hasSpaceFor(dataSize);
+    if (!hasSpace.isOK()) {
+        return hasSpace;
+    }
+
     outputDoc.appendSelfToBufBuilder(_builder);
-    return *this;
+    return Status::OK();
 }
 
 ReplyBuilderInterface::State CommandReplyBuilder::getState() const {
@@ -98,8 +114,7 @@ std::unique_ptr<Message> CommandReplyBuilder::done() {
     return std::move(_message);
 }
 
-std::size_t CommandReplyBuilder::availableSpaceForOutputDocs() const {
-    invariant(State::kDone != _state);
+std::size_t CommandReplyBuilder::availableBytes() const {
     int intLen = _builder.len();
     invariant(0 <= intLen);
     std::size_t len = static_cast<std::size_t>(intLen);
@@ -108,5 +123,15 @@ std::size_t CommandReplyBuilder::availableSpaceForOutputDocs() const {
     return mongo::MaxMessageSizeBytes - len - msgHeaderSz;
 }
 
+Status CommandReplyBuilder::_hasSpaceFor(std::size_t dataSize) const {
+    size_t availBytes = availableBytes();
+    if (availBytes < dataSize) {
+        return Status(ErrorCodes::Overflow,
+                      str::stream() << "Not enough space to store " << dataSize << " bytes. Only "
+                                    << availBytes << " bytes are available.");
+    }
+    return Status::OK();
+}
+
 }  // rpc
 }  // mongo
diff --git a/src/mongo/rpc/command_reply_builder.h b/src/mongo/rpc/command_reply_builder.h
index 2433a283e04..92c567f64a6 100644
--- a/src/mongo/rpc/command_reply_builder.h
+++ b/src/mongo/rpc/command_reply_builder.h
@@ -30,7 +30,7 @@
 
 #include <memory>
 
-#include "mongo/base/status_with.h"
+#include "mongo/base/status.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/rpc/document_range.h"
 #include "mongo/rpc/protocol.h"
@@ -56,11 +56,11 @@ public:
      */
     CommandReplyBuilder(std::unique_ptr<Message> message);
 
-    CommandReplyBuilder& setMetadata(BSONObj metadata) final;
-    CommandReplyBuilder& setRawCommandReply(BSONObj commandReply) final;
+    CommandReplyBuilder& setMetadata(const BSONObj& metadata) final;
+    CommandReplyBuilder& setRawCommandReply(const BSONObj& commandReply) final;
 
-    CommandReplyBuilder& addOutputDocs(DocumentRange outputDocs) final;
-    CommandReplyBuilder& addOutputDoc(BSONObj outputDoc) final;
+    Status addOutputDocs(DocumentRange outputDocs) final;
+    Status addOutputDoc(const BSONObj& outputDoc) final;
 
     State getState() const final;
 
@@ -75,13 +75,18 @@ public:
      */
     std::unique_ptr<Message> done() final;
 
-    std::size_t availableSpaceForOutputDocs() const final;
+    std::size_t availableBytes() const final;
 
 private:
+    /**
+     *  Checks if there is enough space in the buffer to store dataSize bytes
+     *  and computes error message if not.
+     */
+    Status _hasSpaceFor(std::size_t dataSize) const;
+
     // Default values are all empty.
     BufBuilder _builder{};
     std::unique_ptr<Message> _message;
-
     State _state{State::kMetadata};
 };
 
diff --git a/src/mongo/rpc/legacy_reply.cpp b/src/mongo/rpc/legacy_reply.cpp
index 62a45ea1286..b516b536503 100644
--- a/src/mongo/rpc/legacy_reply.cpp
+++ b/src/mongo/rpc/legacy_reply.cpp
@@ -33,9 +33,10 @@
 #include <utility>
 #include <tuple>
 
+#include "mongo/rpc/legacy_reply_builder.h"
+#include "mongo/rpc/metadata.h"
 #include "mongo/util/assert_util.h"
 #include "mongo/util/mongoutils/str.h"
-#include "mongo/rpc/metadata.h"
 
 namespace mongo {
 namespace rpc {
@@ -63,9 +64,31 @@ LegacyReply::LegacyReply(const Message* message) : _message(std::move(message))
                           << " expected a value of 0 but got " << qr.getStartingFrom(),
             qr.getStartingFrom() == 0);
 
-    // TODO bson validation
     std::tie(_commandReply, _metadata) =
         uassertStatusOK(rpc::upconvertReplyMetadata(BSONObj(qr.data())));
+
+    // Copy the bson array of documents from the message into
+    // a contiguous area of memory owned by _docBuffer so
+    // DocumentRange can be used to iterate over documents
+    auto cursorElem = _commandReply[LegacyReplyBuilder::kCursorTag];
+    if (cursorElem.eoo())
+        return;
+
+    BSONObj cursorObj = cursorElem.Obj();
+    auto firstBatchElem = cursorObj[LegacyReplyBuilder::kFirstBatchTag];
+    if (firstBatchElem.eoo())
+        return;
+
+    for (BSONObjIterator it(firstBatchElem.Obj()); it.more(); it.next()) {
+        invariant((*it).isABSONObj());
+        BSONObj doc = (*it).Obj();
+        doc.appendSelfToBufBuilder(_docBuffer);
+    }
+    const char* dataBegin = _docBuffer.buf();
+    const char* dataEnd = dataBegin + _docBuffer.len();
+    _outputDocs = DocumentRange(dataBegin, dataEnd);
+
+    return;
 }
 
 const BSONObj& LegacyReply::getMetadata() const {
@@ -77,8 +100,7 @@ const BSONObj& LegacyReply::getCommandReply() const {
 }
 
 DocumentRange LegacyReply::getOutputDocs() const {
-    // return empty range
-    return DocumentRange{};
+    return _outputDocs;
 }
 
 Protocol LegacyReply::getProtocol() const {
diff --git a/src/mongo/rpc/legacy_reply.h b/src/mongo/rpc/legacy_reply.h
index 86258c2642b..9e6fb8b0b79 100644
--- a/src/mongo/rpc/legacy_reply.h
+++ b/src/mongo/rpc/legacy_reply.h
@@ -80,9 +80,10 @@ public:
 private:
     const Message* _message;
 
-    // TODO: SERVER-18236
-    BSONObj _metadata{};
     BSONObj _commandReply{};  // will hold unowned
+    BSONObj _metadata{};
+    BufBuilder _docBuffer{};
+    DocumentRange _outputDocs{};
 };
 
 }  // namespace rpc
diff --git a/src/mongo/rpc/legacy_reply_builder.cpp b/src/mongo/rpc/legacy_reply_builder.cpp
index 4448f1e9810..592e3588b43 100644
--- a/src/mongo/rpc/legacy_reply_builder.cpp
+++ b/src/mongo/rpc/legacy_reply_builder.cpp
@@ -28,51 +28,126 @@
 
 #include "mongo/platform/basic.h"
 
+#include "mongo/rpc/legacy_reply_builder.h"
+
+#include <iterator>
+
 #include "mongo/db/dbmessage.h"
 #include "mongo/db/jsobj.h"
-#include "mongo/rpc/legacy_reply_builder.h"
 #include "mongo/rpc/metadata.h"
 #include "mongo/stdx/memory.h"
 #include "mongo/util/assert_util.h"
+#include "mongo/util/mongoutils/str.h"
 
 namespace mongo {
 namespace rpc {
 
+namespace {
+// Margin of error for availableBytes size estimate
+std::size_t kReservedSize = 1024;
+
+bool isEmptyCommandReply(const BSONObj& bson) {
+    auto cursorElem = bson[LegacyReplyBuilder::kCursorTag];
+    if (cursorElem.eoo())
+        return true;
+
+    BSONObj cursorObj = cursorElem.Obj();
+    auto firstBatchElem = cursorObj[LegacyReplyBuilder::kFirstBatchTag];
+    if (firstBatchElem.eoo())
+        return true;
+
+    BSONObjIterator it(firstBatchElem.Obj());
+
+    return !it.more();
+}
+}  // namespace
+
+const char LegacyReplyBuilder::kCursorTag[] = "cursor";
+const char LegacyReplyBuilder::kFirstBatchTag[] = "firstBatch";
+
 LegacyReplyBuilder::LegacyReplyBuilder() : LegacyReplyBuilder(stdx::make_unique<Message>()) {}
 
 LegacyReplyBuilder::LegacyReplyBuilder(std::unique_ptr<Message> message)
-    : _message{std::move(message)} {
-    _builder.skip(sizeof(QueryResult::Value));
-}
+    : _currentLength{kReservedSize}, _message{std::move(message)} {}
 
 LegacyReplyBuilder::~LegacyReplyBuilder() {}
 
-LegacyReplyBuilder& LegacyReplyBuilder::setMetadata(BSONObj metadata) {
+LegacyReplyBuilder& LegacyReplyBuilder::setMetadata(const BSONObj& metadata) {
     invariant(_state == State::kMetadata);
-    _metadata = std::move(metadata);
+
+    auto dataSize = static_cast<std::size_t>(metadata.objsize());
+
+    _currentLength += dataSize;
+    _metadata = metadata.getOwned();
     _state = State::kCommandReply;
     return *this;
 }
 
-LegacyReplyBuilder& LegacyReplyBuilder::setRawCommandReply(BSONObj commandReply) {
+LegacyReplyBuilder& LegacyReplyBuilder::setRawCommandReply(const BSONObj& commandReply) {
     invariant(_state == State::kCommandReply);
-    BSONObj downconvertedCommandReply = uassertStatusOK(
-        rpc::downconvertReplyMetadata(std::move(commandReply), std::move(_metadata)));
-    downconvertedCommandReply.appendSelfToBufBuilder(_builder);
+
+    auto dataSize = static_cast<std::size_t>(commandReply.objsize());
+
+    _currentLength += dataSize;
+    _commandReply = commandReply.getOwned();
+    _allowAddingOutputDocs = isEmptyCommandReply(_commandReply);
+
     _state = State::kOutputDocs;
     return *this;
 }
 
-LegacyReplyBuilder& LegacyReplyBuilder::addOutputDocs(DocumentRange outputDocs) {
+Status LegacyReplyBuilder::addOutputDocs(DocumentRange docs) {
     invariant(_state == State::kOutputDocs);
-    // no op
-    return *this;
+    invariant(_allowAddingOutputDocs);
+
+    auto dataSize = docs.data().length();
+
+    auto hasSpace = _hasSpaceFor(dataSize);
+    if (!hasSpace.isOK()) {
+        return hasSpace;
+    }
+
+    // The temporary obj is used to address the case when where is not enough space.
+    // BSONArray overhead can not be estimated upfront.
+    std::vector<BSONObj> docsTmp{};
+    std::size_t lenTmp = 0;
+    std::size_t tmpIndex(_currentIndex);
+    for (auto&& it : docs) {
+        docsTmp.emplace_back(it.getOwned());
+        lenTmp += BSONObjBuilder::numStr(++tmpIndex).length() + 2;  // space for storing array index
+    }
+
+    hasSpace = _hasSpaceFor(dataSize + lenTmp);
+    if (!hasSpace.isOK()) {
+        return hasSpace;
+    }
+
+    // vector::insert instead of swap allows to call addOutputDoc(s) multiple times
+    _outputDocs.insert(_outputDocs.end(),
+                       std::make_move_iterator(docsTmp.begin()),
+                       std::make_move_iterator(docsTmp.end()));
+
+    _currentIndex = tmpIndex;
+    _currentLength += lenTmp;
+    _currentLength += dataSize;
+    return Status::OK();
 }
 
-LegacyReplyBuilder& LegacyReplyBuilder::addOutputDoc(BSONObj outputDoc) {
+Status LegacyReplyBuilder::addOutputDoc(const BSONObj& bson) {
     invariant(_state == State::kOutputDocs);
-    // no op
-    return *this;
+    invariant(_allowAddingOutputDocs);
+
+    auto dataSize = static_cast<std::size_t>(bson.objsize());
+    auto hasSpace = _hasSpaceFor(dataSize);
+    if (!hasSpace.isOK()) {
+        return hasSpace;
+    }
+
+    _outputDocs.emplace_back(bson.getOwned());
+    _currentLength += dataSize;
+    _currentLength += BSONObjBuilder::numStr(++_currentIndex).length() + 2;  // storing array index
+
+    return Status::OK();
 }
 
 ReplyBuilderInterface::State LegacyReplyBuilder::getState() const {
@@ -89,36 +164,83 @@ void LegacyReplyBuilder::reset() {
     if (_state == State::kMetadata) {
         return;
     }
-    _builder.reset();
-    _metadata = BSONObj();
     _message = stdx::make_unique<Message>();
+    _currentLength = kReservedSize;
+    _currentIndex = 0U;
     _state = State::kMetadata;
 }
 
 
 std::unique_ptr<Message> LegacyReplyBuilder::done() {
     invariant(_state == State::kOutputDocs);
-    std::unique_ptr<Message> message = stdx::make_unique<Message>();
 
-    QueryResult::View qr = _builder.buf();
+    BSONObj reply = uassertStatusOK(rpc::downconvertReplyMetadata(_commandReply, _metadata));
+
+    BufBuilder bufBuilder;
+    bufBuilder.skip(sizeof(QueryResult::Value));
+
+    if (_allowAddingOutputDocs) {
+        BSONObjBuilder topBuilder(bufBuilder);
+        for (const auto& el : reply) {
+            if (kCursorTag != el.fieldNameStringData()) {
+                topBuilder.append(el);
+                continue;
+            }
+            invariant(el.isABSONObj());
+            BSONObjBuilder curBuilder(topBuilder.subobjStart(kCursorTag));
+            for (const auto& insideEl : el.Obj()) {
+                if (kFirstBatchTag != insideEl.fieldNameStringData()) {
+                    curBuilder.append(insideEl);
+                    continue;
+                }
+                invariant(insideEl.isABSONObj());
+                BSONArrayBuilder arrBuilder(curBuilder.subarrayStart(kFirstBatchTag));
+                for (const auto& doc : _outputDocs) {
+                    arrBuilder.append(doc);
+                }
+                arrBuilder.doneFast();
+            }
+            curBuilder.doneFast();
+        }
+        topBuilder.doneFast();
+    } else {
+        reply.appendSelfToBufBuilder(bufBuilder);
+    }
+
+    auto msgHeaderSz = static_cast<std::size_t>(MsgData::MsgDataHeaderSize);
+
+    invariant(static_cast<std::size_t>(bufBuilder.len()) + msgHeaderSz <=
+              mongo::MaxMessageSizeBytes);
+
+    QueryResult::View qr = bufBuilder.buf();
+
     qr.setResultFlagsToOk();
-    qr.msgdata().setLen(_builder.len());
+    qr.msgdata().setLen(bufBuilder.len());
     qr.msgdata().setOperation(opReply);
     qr.setCursorId(0);
     qr.setStartingFrom(0);
     qr.setNReturned(1);
-    _builder.decouple();
 
-    message->setData(qr.view2ptr(), true);
+    _message->setData(qr.view2ptr(), true);
+    bufBuilder.decouple();
 
     _state = State::kDone;
-    return std::move(message);
+    return std::move(_message);
 }
 
-std::size_t LegacyReplyBuilder::availableSpaceForOutputDocs() const {
-    invariant(State::kDone != _state);
-    // LegacyReplyBuilder currently does not support addOutputDoc(s)
-    return 0u;
+std::size_t LegacyReplyBuilder::availableBytes() const {
+    std::size_t msgHeaderSz = static_cast<std::size_t>(MsgData::MsgDataHeaderSize);
+    return mongo::MaxMessageSizeBytes - _currentLength - msgHeaderSz;
+}
+
+Status LegacyReplyBuilder::_hasSpaceFor(std::size_t dataSize) const {
+    size_t availBytes = availableBytes();
+    if (availBytes < dataSize) {
+        return Status(ErrorCodes::Overflow,
+                      str::stream() << "Not enough space to store " << dataSize << " bytes. Only "
+                                    << availBytes << " bytes are available.");
+    }
+    return Status::OK();
 }
 
 }  // namespace rpc
diff --git a/src/mongo/rpc/legacy_reply_builder.h b/src/mongo/rpc/legacy_reply_builder.h
index 417a6db8492..957265f0611 100644
--- a/src/mongo/rpc/legacy_reply_builder.h
+++ b/src/mongo/rpc/legacy_reply_builder.h
@@ -30,7 +30,7 @@
 
 #include <memory>
 
-#include "mongo/base/status_with.h"
+#include "mongo/base/status.h"
 #include "mongo/bson/util/builder.h"
 #include "mongo/rpc/document_range.h"
 #include "mongo/rpc/protocol.h"
@@ -41,15 +41,18 @@ namespace rpc {
 
 class LegacyReplyBuilder : public ReplyBuilderInterface {
 public:
+    static const char kCursorTag[];
+    static const char kFirstBatchTag[];
+
     LegacyReplyBuilder();
     LegacyReplyBuilder(std::unique_ptr<Message>);
     ~LegacyReplyBuilder() final;
 
-    LegacyReplyBuilder& setMetadata(BSONObj metadata) final;
-    LegacyReplyBuilder& setRawCommandReply(BSONObj commandReply) final;
+    LegacyReplyBuilder& setMetadata(const BSONObj& metadata) final;
+    LegacyReplyBuilder& setRawCommandReply(const BSONObj& commandReply) final;
 
-    LegacyReplyBuilder& addOutputDocs(DocumentRange outputDocs) final;
-    LegacyReplyBuilder& addOutputDoc(BSONObj outputDoc) final;
+    Status addOutputDocs(DocumentRange outputDocs) final;
+    Status addOutputDoc(const BSONObj& outputDoc) final;
 
     State getState() const final;
 
@@ -59,12 +62,25 @@ public:
 
     Protocol getProtocol() const final;
 
-    std::size_t availableSpaceForOutputDocs() const final;
+    std::size_t availableBytes() const final;
 
 private:
-    BufBuilder _builder{};
-    BSONObj _metadata{};
+    /**
+     *  Checks if there is enough space in the buffer to store dataSize bytes
+     *  and computes error message if not.
+     */
+    Status _hasSpaceFor(std::size_t dataSize) const;
+
+    // If _allowAddingOutputDocs is false it enforces the "legacy" semantic
+    // where command results are returned inside commandReply.
+    // In this case calling addOutputDoc(s) will break the invariant.
+    bool _allowAddingOutputDocs{true};
+    BSONObj _commandReply{};
+    std::size_t _currentLength;
+    std::size_t _currentIndex = 0U;
     std::unique_ptr<Message> _message;
+    BSONObj _metadata{};
+    std::vector<BSONObj> _outputDocs{};
     State _state{State::kMetadata};
 };
 
diff --git a/src/mongo/rpc/metadata.cpp b/src/mongo/rpc/metadata.cpp
index e80daabaa15..74d511d5d4e 100644
--- a/src/mongo/rpc/metadata.cpp
+++ b/src/mongo/rpc/metadata.cpp
@@ -123,7 +123,7 @@ StatusWith<LegacyCommandAndFlags> downconvertRequestMetadata(BSONObj cmdObj, BSO
     return std::make_tuple(ssmCommandBob.obj(), std::move(legacyQueryFlags));
 }
 
-StatusWith<CommandReplyWithMetadata> upconvertReplyMetadata(BSONObj legacyReply) {
+StatusWith<CommandReplyWithMetadata> upconvertReplyMetadata(const BSONObj& legacyReply) {
     BSONObjBuilder commandReplyBob;
     BSONObjBuilder metadataBob;
 
@@ -135,7 +135,8 @@ StatusWith<CommandReplyWithMetadata> upconvertReplyMetadata(BSONObj legacyReply)
     return std::make_tuple(commandReplyBob.obj(), metadataBob.obj());
 }
 
-StatusWith<BSONObj> downconvertReplyMetadata(BSONObj commandReply, BSONObj replyMetadata) {
+StatusWith<BSONObj> downconvertReplyMetadata(const BSONObj& commandReply,
+                                             const BSONObj& replyMetadata) {
     BSONObjBuilder legacyCommandReplyBob;
 
     auto downconvertStatus =
diff --git a/src/mongo/rpc/metadata.h b/src/mongo/rpc/metadata.h
index ed2267bbca6..04cc1606d3e 100644
--- a/src/mongo/rpc/metadata.h
+++ b/src/mongo/rpc/metadata.h
@@ -106,13 +106,14 @@ using CommandReplyWithMetadata = std::tuple<BSONObj, BSONObj>;
  * Given a legacy command reply, attempts to strip the metadata from the reply and construct
  * a metadata object.
  */
-StatusWith<CommandReplyWithMetadata> upconvertReplyMetadata(BSONObj legacyReply);
+StatusWith<CommandReplyWithMetadata> upconvertReplyMetadata(const BSONObj& legacyReply);
 
 /**
  * Given a command reply object and an associated metadata object,
  * attempts to construct a legacy command object.
  */
-StatusWith<BSONObj> downconvertReplyMetadata(BSONObj commandReply, BSONObj replyMetadata);
+StatusWith<BSONObj> downconvertReplyMetadata(const BSONObj& commandReply,
+                                             const BSONObj& replyMetadata);
 
 /**
  * A function type for writing request metadata. The function takes a pointer to a
diff --git a/src/mongo/rpc/metadata/sharding_metadata.h b/src/mongo/rpc/metadata/sharding_metadata.h
index fba0fe8f992..048473064b4 100644
--- a/src/mongo/rpc/metadata/sharding_metadata.h
+++ b/src/mongo/rpc/metadata/sharding_metadata.h
@@ -25,6 +25,7 @@
  *    exception statement from all source files in the program, then also delete
  *    it in the license file.
  */
+#pragma once
 
 #include "mongo/db/jsobj.h"
 
diff --git a/src/mongo/rpc/reply_builder_interface.cpp b/src/mongo/rpc/reply_builder_interface.cpp
index 0bf5170d6ee..5949b405a54 100644
--- a/src/mongo/rpc/reply_builder_interface.cpp
+++ b/src/mongo/rpc/reply_builder_interface.cpp
@@ -77,7 +77,7 @@ ReplyBuilderInterface& ReplyBuilderInterface::setCommandReply(StatusWith<BSONObj
 }
 
 ReplyBuilderInterface& ReplyBuilderInterface::setCommandReply(Status nonOKStatus,
-                                                              BSONObj extraErrorInfo) {
+                                                              const BSONObj& extraErrorInfo) {
     invariant(!nonOKStatus.isOK());
     return setRawCommandReply(augmentReplyWithStatus(nonOKStatus, extraErrorInfo));
 }
diff --git a/src/mongo/rpc/reply_builder_interface.h b/src/mongo/rpc/reply_builder_interface.h
index e97e3e7fe9b..c02511be567 100644
--- a/src/mongo/rpc/reply_builder_interface.h
+++ b/src/mongo/rpc/reply_builder_interface.h
@@ -31,7 +31,7 @@
 #include <memory>
 
 #include "mongo/base/disallow_copying.h"
-#include "mongo/base/status_with.h"
+#include "mongo/base/status.h"
 #include "mongo/rpc/protocol.h"
 
 namespace mongo {
@@ -59,13 +59,13 @@ public:
 
     virtual ~ReplyBuilderInterface() = default;
 
-    virtual ReplyBuilderInterface& setMetadata(BSONObj metadata) = 0;
+    virtual ReplyBuilderInterface& setMetadata(const BSONObj& metadata) = 0;
 
     /**
      * Sets the raw command reply. This should probably not be used in favor of the
      * variants that accept a Status or StatusWith.
      */
-    virtual ReplyBuilderInterface& setRawCommandReply(BSONObj reply) = 0;
+    virtual ReplyBuilderInterface& setRawCommandReply(const BSONObj& reply) = 0;
 
     /**
      * Sets the reply for this command. If an engaged StatusWith<BSONObj> is passed, the command
@@ -85,19 +85,21 @@ public:
      * interfacing with legacy code that adds additional data to a failed command reply and
      * its use is discouraged in new code.
      */
-    ReplyBuilderInterface& setCommandReply(Status nonOKStatus, BSONObj extraErrorInfo);
+    ReplyBuilderInterface& setCommandReply(Status nonOKStatus, const BSONObj& extraErrorInfo);
 
     /**
      * Add a range of output documents to the reply. This method can be called multiple times
-     * before calling done().
+     * before calling done(). A non OK status indicates that the message does not have
+     * enough space to store ouput documents.
      */
-    virtual ReplyBuilderInterface& addOutputDocs(DocumentRange outputDocs) = 0;
+    virtual Status addOutputDocs(DocumentRange outputDocs) = 0;
 
     /**
      * Add a single output document to the reply. This method can be called multiple times
-     * before calling done().
+     * before calling done(). A non OK status indicates that the message does not have
+     * enough space to store ouput documents.
      */
-    virtual ReplyBuilderInterface& addOutputDoc(BSONObj outputDoc) = 0;
+    virtual Status addOutputDoc(const BSONObj& outputDoc) = 0;
 
     /**
      * Gets the state of the builder. As the builder will simply crash the process if it is ever
@@ -123,7 +125,7 @@ public:
      * Returns available space in bytes, should be used to verify that the message have enough
      * space for ouput documents.
      */
-    virtual std::size_t availableSpaceForOutputDocs() const = 0;
+    virtual std::size_t availableBytes() const = 0;
 
     /**
      * Writes data then transfers ownership of the message to the caller. The behavior of
diff --git a/src/mongo/rpc/reply_builder_test.cpp b/src/mongo/rpc/reply_builder_test.cpp
index 218e8b24caa..ad50b810989 100644
--- a/src/mongo/rpc/reply_builder_test.cpp
+++ b/src/mongo/rpc/reply_builder_test.cpp
@@ -30,26 +30,118 @@
 
 #include "mongo/bson/util/builder.h"
 #include "mongo/db/jsobj.h"
+#include "mongo/db/json.h"
 #include "mongo/rpc/command_reply.h"
+#include "mongo/rpc/legacy_reply.h"
 #include "mongo/rpc/command_reply_builder.h"
 #include "mongo/rpc/legacy_reply_builder.h"
 #include "mongo/rpc/document_range.h"
 #include "mongo/unittest/unittest.h"
+#include "mongo/unittest/death_test.h"
 
 namespace {
 
 using namespace mongo;
 
-static void _testMaxCommandReply(rpc::ReplyBuilderInterface& replyBuilder);
+void testMaxCommandReply(rpc::ReplyBuilderInterface& replyBuilder);
+
+template <typename T>
+void testRoundTrip(rpc::ReplyBuilderInterface& replyBuilder);
+
+TEST(LegacyReplyBuilder, RoundTrip) {
+    rpc::LegacyReplyBuilder r;
+    testRoundTrip<rpc::LegacyReply>(r);
+}
 
 TEST(CommandReplyBuilder, RoundTrip) {
-    BSONObjBuilder metadataBob{};
-    metadataBob.append("foo", "bar");
-    auto metadata = metadataBob.done();
+    rpc::CommandReplyBuilder r;
+    testRoundTrip<rpc::CommandReply>(r);
+}
 
+BSONObj buildMetadata() {
+    BSONObjBuilder metadataTop{};
+    BSONObjBuilder metadataGle{};
+    metadataGle.append("lastOpTime", Timestamp());
+    metadataGle.append("electionId", OID("5592bee00d21e3aa796e185e"));
+    metadataTop.append("$gleStats", metadataGle.done());
+    return metadataTop.obj();
+}
+
+BSONObj buildEmptyCommand() {
+    const char text[] = "{ ok: 1.0, cursor: { firstBatch: [] } }";
+    mongo::BSONObj obj = mongo::fromjson(text);
+    return obj;
+}
+
+BSONObj buildCommand() {
     BSONObjBuilder commandReplyBob{};
-    commandReplyBob.append("bar", "baz").append("ok", 1.0);
-    auto commandReply = commandReplyBob.done();
+    commandReplyBob.append("ok", 1.0);
+    BSONObjBuilder cursorBuilder;
+    BSONArrayBuilder a(cursorBuilder.subarrayStart("firstBatch"));
+    a.append(BSON("Foo"
+                  << "Bar"));
+    a.done();
+
+    cursorBuilder.appendIntOrLL("id", 1);
+    cursorBuilder.append("ns", "test.$cmd.blah");
+    commandReplyBob.append("cursor", cursorBuilder.done());
+    return commandReplyBob.obj();
+}
+
+TEST(CommandReplyBuilder, MemAccess) {
+    BSONObj metadata = buildMetadata();
+    BSONObj commandReply = buildCommand();
+    rpc::CommandReplyBuilder replyBuilder;
+    replyBuilder.setMetadata(metadata);
+    replyBuilder.setCommandReply(commandReply);
+    auto msg = replyBuilder.done();
+
+    rpc::CommandReply parsed(msg.get());
+
+    ASSERT_EQUALS(parsed.getMetadata(), metadata);
+    ASSERT_EQUALS(parsed.getCommandReply(), commandReply);
+}
+
+TEST(LegacyReplyBuilder, MemAccess) {
+    BSONObj metadata = buildMetadata();
+    BSONObj commandReply = buildEmptyCommand();
+    rpc::LegacyReplyBuilder replyBuilder;
+    replyBuilder.setMetadata(metadata);
+    replyBuilder.setCommandReply(commandReply);
+    auto msg = replyBuilder.done();
+
+    rpc::LegacyReply parsed(msg.get());
+
+    ASSERT_EQUALS(parsed.getMetadata(), metadata);
+    ASSERT_EQUALS(parsed.getCommandReply(), commandReply);
+}
+
+DEATH_TEST(LegacyReplyBuilder, FailureAddingDoc, "Invariant failure _allowAddingOutputDocs") {
+    BSONObj metadata = buildMetadata();
+    BSONObj commandReply = buildCommand();
+    rpc::LegacyReplyBuilder replyBuilder;
+    replyBuilder.setMetadata(metadata);
+    replyBuilder.setCommandReply(commandReply);
+    replyBuilder.addOutputDoc(BSONObj());
+}
+
+DEATH_TEST(LegacyReplyBuilder, FailureAddingDocs, "Invariant failure _allowAddingOutputDocs") {
+    BSONObj metadata = buildMetadata();
+    BSONObj commandReply = buildCommand();
+    rpc::LegacyReplyBuilder replyBuilder;
+    replyBuilder.setMetadata(metadata);
+    replyBuilder.setCommandReply(commandReply);
+    rpc::DocumentRange range;
+    replyBuilder.addOutputDocs(range);
+}
+
+template <typename T>
+void testRoundTrip(rpc::ReplyBuilderInterface& replyBuilder) {
+    auto metadata = buildMetadata();
+    auto commandReply = buildEmptyCommand();
+
+    replyBuilder.setMetadata(metadata);
+    replyBuilder.setCommandReply(commandReply);
 
     BSONObjBuilder outputDoc1Bob{};
     outputDoc1Bob.append("z", "t");
@@ -67,47 +159,64 @@ TEST(CommandReplyBuilder, RoundTrip) {
     outputDoc1.appendSelfToBufBuilder(outputDocs);
     outputDoc2.appendSelfToBufBuilder(outputDocs);
     outputDoc3.appendSelfToBufBuilder(outputDocs);
-
     rpc::DocumentRange outputDocRange{outputDocs.buf(), outputDocs.buf() + outputDocs.len()};
+    replyBuilder.addOutputDocs(outputDocRange);
 
-    rpc::CommandReplyBuilder r;
-
-    auto msg =
-        r.setMetadata(metadata).setCommandReply(commandReply).addOutputDocs(outputDocRange).done();
+    auto msg = replyBuilder.done();
 
-    rpc::CommandReply parsed(msg.get());
+    T parsed(msg.get());
 
     ASSERT_EQUALS(parsed.getMetadata(), metadata);
-    ASSERT_EQUALS(parsed.getCommandReply(), commandReply);
-    // need ostream overloads for ASSERT_EQUALS
     ASSERT_TRUE(parsed.getOutputDocs() == outputDocRange);
 }
 
 TEST(CommandReplyBuilder, MaxCommandReply) {
     rpc::CommandReplyBuilder replyBuilder;
-    _testMaxCommandReply(replyBuilder);
+    testMaxCommandReply(replyBuilder);
 }
 
 TEST(LegacyReplyBuilder, MaxCommandReply) {
     rpc::LegacyReplyBuilder replyBuilder;
-    _testMaxCommandReply(replyBuilder);
+    testMaxCommandReply(replyBuilder);
 }
 
-// verify current functionality - later will need to change
 TEST(LegacyReplyBuilderSpaceTest, DocSize) {
     rpc::LegacyReplyBuilder replyBuilder;
-    replyBuilder.setMetadata(BSONObj()).setCommandReply(BSONObj());
+    auto metadata = buildMetadata();
+    auto commandReply = buildEmptyCommand();
 
-    std::size_t spaceBefore = replyBuilder.availableSpaceForOutputDocs();
-    ASSERT_EQUALS(spaceBefore, 0u);
+    replyBuilder.setMetadata(metadata);
+    replyBuilder.setCommandReply(commandReply);
 
-    BSONObjBuilder docBuilder{};
-    docBuilder.append("foo", "bar");
-    auto doc = docBuilder.done();
+    auto sizeBefore = replyBuilder.availableBytes();
+
+    for (int i = 0; i < 100000; ++i) {
+        BSONObjBuilder docBuilder;
+        docBuilder.append("foo" + std::to_string(i), "bar" + std::to_string(i));
+        auto statusAfter = replyBuilder.addOutputDoc(docBuilder.done());
+        ASSERT_TRUE(statusAfter.isOK());
+    }
+
+    auto sizeAfter = replyBuilder.availableBytes();
+    auto msg = replyBuilder.done();
 
-    replyBuilder.addOutputDoc(doc);  // no-op
-    std::size_t spaceAfter = replyBuilder.availableSpaceForOutputDocs();
-    ASSERT_EQUALS(spaceAfter, 0u);
+    // construct an empty message to compare the estimated size difference with
+    // the actual difference
+    rpc::LegacyReplyBuilder replyBuilder0;
+    replyBuilder0.setMetadata(metadata);
+    replyBuilder0.setCommandReply(commandReply);
+    auto msg0 = replyBuilder0.done();
+
+    QueryResult::View qr0 = msg0->singleData().view2ptr();
+    auto dataLen0 = static_cast<std::size_t>(qr0.msgdata().dataLen());
+
+    QueryResult::View qr = msg->singleData().view2ptr();
+    auto dataLen = static_cast<std::size_t>(qr.msgdata().dataLen());
+
+    // below tests the array space estimates
+    // due to the inaccuracy in size estimation algo the actual size is off by up to 6 bytes
+    // on the large # of documents
+    ASSERT_EQUALS(sizeBefore - sizeAfter, dataLen - dataLen0 + 5);
 }
 
 class CommandReplyBuilderSpaceTest : public mongo::unittest::Test {
@@ -127,9 +236,12 @@ protected:
 
 TEST_F(CommandReplyBuilderSpaceTest, DocSizeEq) {
     rpc::CommandReplyBuilder replyBuilder;
-    replyBuilder.setMetadata(BSONObj()).setCommandReply(BSONObj());
+    auto metadata = buildMetadata();
+    auto commandReply = buildEmptyCommand();
+    replyBuilder.setMetadata(metadata);
+    replyBuilder.setCommandReply(commandReply);
 
-    std::size_t spaceBefore = replyBuilder.availableSpaceForOutputDocs();
+    std::size_t spaceBefore = replyBuilder.availableBytes();
 
     BSONObjBuilder docBuilder{};
     docBuilder.append("foo", "bar");
@@ -137,7 +249,7 @@ TEST_F(CommandReplyBuilderSpaceTest, DocSizeEq) {
     std::size_t docSize = doc.objsize();
 
     replyBuilder.addOutputDoc(doc);
-    std::size_t spaceAfter = replyBuilder.availableSpaceForOutputDocs();
+    std::size_t spaceAfter = replyBuilder.availableBytes();
     ASSERT_EQUALS(spaceBefore - docSize, spaceAfter);
 }
 
@@ -145,9 +257,12 @@ TEST_F(CommandReplyBuilderSpaceTest, DocSizeEq) {
 TEST_F(CommandReplyBuilderSpaceTest, MaxDocSize1) {
     rpc::CommandReplyBuilder replyBuilder;
 
-    replyBuilder.setMetadata(BSONObj()).setCommandReply(BSONObj());
+    auto metadata = buildMetadata();
+    auto commandReply = buildEmptyCommand();
+    replyBuilder.setMetadata(metadata);
+    replyBuilder.setCommandReply(commandReply);
 
-    std::size_t availSpace = replyBuilder.availableSpaceForOutputDocs();
+    std::size_t availSpace = replyBuilder.availableBytes();
 
     while (availSpace > 0u) {
         std::size_t payloadSz =
@@ -158,7 +273,7 @@ TEST_F(CommandReplyBuilderSpaceTest, MaxDocSize1) {
         docBuilder.append("x", payload);
         auto doc = docBuilder.done();
         replyBuilder.addOutputDoc(doc);
-        availSpace = replyBuilder.availableSpaceForOutputDocs();
+        availSpace = replyBuilder.availableBytes();
     }
     auto msg = replyBuilder.done();
     auto sizeUInt = static_cast<std::size_t>(msg->size());
@@ -170,17 +285,12 @@ TEST_F(CommandReplyBuilderSpaceTest, MaxDocSize1) {
 TEST_F(CommandReplyBuilderSpaceTest, MaxDocSize2) {
     rpc::CommandReplyBuilder replyBuilder;
 
-    BSONObjBuilder metadataBuilder{};
-    metadataBuilder.append("foo", "bar");
-    auto metadata = metadataBuilder.done();
-
-    BSONObjBuilder commandReplyBuilder{};
-    commandReplyBuilder.append("oof", "rab");
-    auto commandReply = commandReplyBuilder.done();
-
-    replyBuilder.setMetadata(metadata).setCommandReply(commandReply);
+    auto metadata = buildMetadata();
+    auto commandReply = buildEmptyCommand();
+    replyBuilder.setMetadata(metadata);
+    replyBuilder.setCommandReply(commandReply);
 
-    std::size_t availSpace = replyBuilder.availableSpaceForOutputDocs();
+    std::size_t availSpace = replyBuilder.availableBytes();
 
     while (availSpace > 0u) {
         std::size_t payloadSz =
@@ -191,7 +301,7 @@ TEST_F(CommandReplyBuilderSpaceTest, MaxDocSize2) {
         docBuilder.append("x", payload);
         auto doc = docBuilder.done();
         replyBuilder.addOutputDoc(doc);
-        availSpace = replyBuilder.availableSpaceForOutputDocs();
+        availSpace = replyBuilder.availableBytes();
     }
     auto msg = replyBuilder.done();
     auto sizeUInt = static_cast<std::size_t>(msg->size());
@@ -204,17 +314,12 @@ TEST_F(CommandReplyBuilderSpaceTest, MaxDocSize2) {
 TEST_F(CommandReplyBuilderSpaceTest, MaxDocSize3) {
     rpc::CommandReplyBuilder replyBuilder;
 
-    BSONObjBuilder metadataBuilder{};
-    metadataBuilder.append("foo", "bar");
-    auto metadata = metadataBuilder.done();
-
-    BSONObjBuilder commandReplyBuilder{};
-    commandReplyBuilder.append("oof", "rab");
-    auto commandReply = commandReplyBuilder.done();
-
-    replyBuilder.setMetadata(metadata).setCommandReply(commandReply);
+    auto metadata = buildMetadata();
+    auto commandReply = buildEmptyCommand();
+    replyBuilder.setMetadata(metadata);
+    replyBuilder.setCommandReply(commandReply);
 
-    std::size_t availSpace = replyBuilder.availableSpaceForOutputDocs();
+    std::size_t availSpace = replyBuilder.availableBytes();
 
     BufBuilder docs;
     while (availSpace > 0u) {
@@ -232,21 +337,20 @@ TEST_F(CommandReplyBuilderSpaceTest, MaxDocSize3) {
     replyBuilder.addOutputDocs(docRange);
 
     auto msg = replyBuilder.done();
+
     auto sizeUInt = static_cast<std::size_t>(msg->size());
 
     ASSERT_EQUALS(sizeUInt, mongo::MaxMessageSizeBytes);
 }
 
 // call to addCommandReply
-void _testMaxCommandReply(rpc::ReplyBuilderInterface& replyBuilder) {
+void testMaxCommandReply(rpc::ReplyBuilderInterface& replyBuilder) {
     BSONObjBuilder docBuilder1{};
     docBuilder1.append("x", "");
     auto emptyDoc = docBuilder1.done();
     std::size_t emptyDocSize = emptyDoc.objsize();
 
-    BSONObjBuilder metadataBuilder{};
-    metadataBuilder.append("foo", "bar");
-    auto metadata = metadataBuilder.done();
+    auto metadata = buildMetadata();
     replyBuilder.setMetadata(metadata);
 
     auto payloadSz = static_cast<std::size_t>(mongo::BSONObjMaxUserSize) - emptyDocSize;
@@ -254,9 +358,8 @@ void _testMaxCommandReply(rpc::ReplyBuilderInterface& replyBuilder) {
     BSONObjBuilder commandReplyBuilder{};
     std::string payload = std::string(payloadSz, 'y');
     commandReplyBuilder.append("x", payload);
-    auto commandReply = commandReplyBuilder.done();
+    auto commandReply = commandReplyBuilder.obj();
     ASSERT_EQUALS(commandReply.objsize(), mongo::BSONObjMaxUserSize);
-
     replyBuilder.setCommandReply(commandReply);
 }