diff options
-rw-r--r-- | src/mongo/rpc/command_reply_builder.cpp | 47 | ||||
-rw-r--r-- | src/mongo/rpc/command_reply_builder.h | 19 | ||||
-rw-r--r-- | src/mongo/rpc/legacy_reply.cpp | 30 | ||||
-rw-r--r-- | src/mongo/rpc/legacy_reply.h | 5 | ||||
-rw-r--r-- | src/mongo/rpc/legacy_reply_builder.cpp | 178 | ||||
-rw-r--r-- | src/mongo/rpc/legacy_reply_builder.h | 32 | ||||
-rw-r--r-- | src/mongo/rpc/metadata.cpp | 5 | ||||
-rw-r--r-- | src/mongo/rpc/metadata.h | 5 | ||||
-rw-r--r-- | src/mongo/rpc/metadata/sharding_metadata.h | 1 | ||||
-rw-r--r-- | src/mongo/rpc/reply_builder_interface.cpp | 2 | ||||
-rw-r--r-- | src/mongo/rpc/reply_builder_interface.h | 20 | ||||
-rw-r--r-- | src/mongo/rpc/reply_builder_test.cpp | 221 |
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); } |