SERVER-21459 reserve bytes in ReplyBuilder buffer to avoid excessive reallocing

14 files changed, 54 insertions, 7 deletions

diff --git a/src/mongo/db/commands.h b/src/mongo/db/commands.h
index 5897fb79321..37550c64131 100644
--- a/src/mongo/db/commands.h
+++ b/src/mongo/db/commands.h
@@ -95,6 +95,10 @@ public:
     // collection name or just a database name.
     ResourcePattern parseResourcePattern(const std::string& dbname, const BSONObj& cmdObj) const;
 
+    virtual std::size_t reserveBytesForReply() const {
+        return 0u;
+    }
+
     const std::string name;
 
     /* run the given command
diff --git a/src/mongo/db/commands/distinct.cpp b/src/mongo/db/commands/distinct.cpp
index e20153f7b5d..cf8d0007db0 100644
--- a/src/mongo/db/commands/distinct.cpp
+++ b/src/mongo/db/commands/distinct.cpp
@@ -46,6 +46,7 @@
 #include "mongo/db/instance.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/db/query/explain.h"
+#include "mongo/db/query/find_common.h"
 #include "mongo/db/query/get_executor.h"
 #include "mongo/db/query/query_planner_common.h"
 #include "mongo/util/log.h"
@@ -81,6 +82,10 @@ public:
         return true;
     }
 
+    std::size_t reserveBytesForReply() const override {
+        return FindCommon::kInitReplyBufferSize;
+    }
+
     virtual void addRequiredPrivileges(const std::string& dbname,
                                        const BSONObj& cmdObj,
                                        std::vector<Privilege>* out) {
diff --git a/src/mongo/db/commands/find_cmd.cpp b/src/mongo/db/commands/find_cmd.cpp
index 633bd24ef8b..6ab582e0365 100644
--- a/src/mongo/db/commands/find_cmd.cpp
+++ b/src/mongo/db/commands/find_cmd.cpp
@@ -105,6 +105,10 @@ public:
         return LogicalOp::opQuery;
     }
 
+    std::size_t reserveBytesForReply() const override {
+        return FindCommon::kInitReplyBufferSize;
+    }
+
     /**
      * A find command does not increment the command counter, but rather increments the
      * query counter.
diff --git a/src/mongo/db/commands/geo_near_cmd.cpp b/src/mongo/db/commands/geo_near_cmd.cpp
index f26d4e5e59e..3b5d0ccee6c 100644
--- a/src/mongo/db/commands/geo_near_cmd.cpp
+++ b/src/mongo/db/commands/geo_near_cmd.cpp
@@ -48,6 +48,7 @@
 #include "mongo/db/matcher/expression_geo.h"
 #include "mongo/db/matcher/extensions_callback_real.h"
 #include "mongo/db/query/explain.h"
+#include "mongo/db/query/find_common.h"
 #include "mongo/db/query/get_executor.h"
 #include "mongo/db/range_preserver.h"
 #include "mongo/platform/unordered_map.h"
@@ -75,6 +76,10 @@ public:
         return true;
     }
 
+    std::size_t reserveBytesForReply() const override {
+        return FindCommon::kInitReplyBufferSize;
+    }
+
     void help(stringstream& h) const {
         h << "http://dochub.mongodb.org/core/geo#GeospatialIndexing-geoNearCommand";
     }
diff --git a/src/mongo/db/commands/getmore_cmd.cpp b/src/mongo/db/commands/getmore_cmd.cpp
index 4ba445d1825..882215570f0 100644
--- a/src/mongo/db/commands/getmore_cmd.cpp
+++ b/src/mongo/db/commands/getmore_cmd.cpp
@@ -107,6 +107,13 @@ public:
         return LogicalOp::opGetMore;
     }
 
+    std::size_t reserveBytesForReply() const override {
+        // The extra 1K is an artifact of how we construct batches. We consider a batch to be full
+        // when it exceeds the goal batch size. In the case that we are just below the limit and
+        // then read a large document, the extra 1K helps prevent a final realloc+memcpy.
+        return FindCommon::kMaxBytesToReturnToClientAtOnce + 1024u;
+    }
+
     /**
      * A getMore command increments the getMore counter, not the command counter.
      */
diff --git a/src/mongo/db/commands/group_cmd.cpp b/src/mongo/db/commands/group_cmd.cpp
index 88a9169599f..c279fa458a1 100644
--- a/src/mongo/db/commands/group_cmd.cpp
+++ b/src/mongo/db/commands/group_cmd.cpp
@@ -37,6 +37,7 @@
 #include "mongo/db/db_raii.h"
 #include "mongo/db/exec/group.h"
 #include "mongo/db/exec/working_set_common.h"
+#include "mongo/db/query/find_common.h"
 #include "mongo/db/query/get_executor.h"
 
 namespace mongo {
@@ -71,6 +72,10 @@ private:
         return true;
     }
 
+    std::size_t reserveBytesForReply() const override {
+        return FindCommon::kInitReplyBufferSize;
+    }
+
     virtual void help(std::stringstream& help) const {
         help << "http://dochub.mongodb.org/core/aggregation";
     }
diff --git a/src/mongo/db/commands/mr.cpp b/src/mongo/db/commands/mr.cpp
index 1d9a52067fa..2d050c924f1 100644
--- a/src/mongo/db/commands/mr.cpp
+++ b/src/mongo/db/commands/mr.cpp
@@ -55,6 +55,7 @@
 #include "mongo/db/operation_context_impl.h"
 #include "mongo/db/query/get_executor.h"
 #include "mongo/db/query/query_planner.h"
+#include "mongo/db/query/find_common.h"
 #include "mongo/db/range_preserver.h"
 #include "mongo/db/repl/replication_coordinator_global.h"
 #include "mongo/db/s/collection_metadata.h"
@@ -1251,6 +1252,10 @@ public:
         return true;
     }
 
+    std::size_t reserveBytesForReply() const override {
+        return FindCommon::kInitReplyBufferSize;
+    }
+
     virtual void help(stringstream& help) const {
         help << "Run a map/reduce operation on the server.\n";
         help << "Note this is used for aggregation, not querying, in MongoDB.\n";
diff --git a/src/mongo/db/dbcommands.cpp b/src/mongo/db/dbcommands.cpp
index 09075791f12..59af14b6faf 100644
--- a/src/mongo/db/dbcommands.cpp
+++ b/src/mongo/db/dbcommands.cpp
@@ -1337,7 +1337,7 @@ void Command::execCommand(OperationContext* txn,
 bool Command::run(OperationContext* txn,
                   const rpc::RequestInterface& request,
                   rpc::ReplyBuilderInterface* replyBuilder) {
-    BSONObjBuilder inPlaceReplyBob(replyBuilder->getInPlaceReplyBuilder());
+    BSONObjBuilder inPlaceReplyBob(replyBuilder->getInPlaceReplyBuilder(reserveBytesForReply()));
 
     repl::ReplicationCoordinator* replCoord = repl::getGlobalReplicationCoordinator();
 
diff --git a/src/mongo/db/geo/haystack.cpp b/src/mongo/db/geo/haystack.cpp
index 2caef7bc87e..23c537534c5 100644
--- a/src/mongo/db/geo/haystack.cpp
+++ b/src/mongo/db/geo/haystack.cpp
@@ -43,6 +43,7 @@
 #include "mongo/db/index/index_descriptor.h"
 #include "mongo/db/index_names.h"
 #include "mongo/db/jsobj.h"
+#include "mongo/db/query/find_common.h"
 
 /**
  * Examines all documents in a given radius of a given point.
@@ -74,6 +75,10 @@ public:
         return true;
     }
 
+    std::size_t reserveBytesForReply() const override {
+        return FindCommon::kInitReplyBufferSize;
+    }
+
     virtual void addRequiredPrivileges(const std::string& dbname,
                                        const BSONObj& cmdObj,
                                        std::vector<Privilege>* out) {
diff --git a/src/mongo/rpc/command_reply_builder.cpp b/src/mongo/rpc/command_reply_builder.cpp
index f6f5f7c1ca0..15b70123708 100644
--- a/src/mongo/rpc/command_reply_builder.cpp
+++ b/src/mongo/rpc/command_reply_builder.cpp
@@ -48,14 +48,17 @@ CommandReplyBuilder::CommandReplyBuilder(Message&& message) : _message{std::move
 
 CommandReplyBuilder& CommandReplyBuilder::setRawCommandReply(const BSONObj& commandReply) {
     invariant(_state == State::kCommandReply);
-
     commandReply.appendSelfToBufBuilder(_builder);
     _state = State::kMetadata;
     return *this;
 }
 
-BufBuilder& CommandReplyBuilder::getInPlaceReplyBuilder() {
+BufBuilder& CommandReplyBuilder::getInPlaceReplyBuilder(std::size_t reserveBytes) {
     invariant(_state == State::kCommandReply);
+    // Eagerly allocate reserveBytes bytes.
+    _builder.reserveBytes(reserveBytes);
+    // Claim our reservation immediately so we can actually write data to it.
+    _builder.claimReservedBytes(reserveBytes);
     _state = State::kMetadata;
     return _builder;
 }
diff --git a/src/mongo/rpc/command_reply_builder.h b/src/mongo/rpc/command_reply_builder.h
index 2079ae69790..39cdc678cc2 100644
--- a/src/mongo/rpc/command_reply_builder.h
+++ b/src/mongo/rpc/command_reply_builder.h
@@ -58,7 +58,7 @@ public:
 
 
     CommandReplyBuilder& setRawCommandReply(const BSONObj& commandReply) final;
-    BufBuilder& getInPlaceReplyBuilder() final;
+    BufBuilder& getInPlaceReplyBuilder(std::size_t) final;
 
     CommandReplyBuilder& setMetadata(const BSONObj& metadata) final;
 
diff --git a/src/mongo/rpc/legacy_reply_builder.cpp b/src/mongo/rpc/legacy_reply_builder.cpp
index efa50d8f82c..71ef4f9f5bf 100644
--- a/src/mongo/rpc/legacy_reply_builder.cpp
+++ b/src/mongo/rpc/legacy_reply_builder.cpp
@@ -58,8 +58,12 @@ LegacyReplyBuilder& LegacyReplyBuilder::setRawCommandReply(const BSONObj& comman
     return *this;
 }
 
-BufBuilder& LegacyReplyBuilder::getInPlaceReplyBuilder() {
+BufBuilder& LegacyReplyBuilder::getInPlaceReplyBuilder(std::size_t reserveBytes) {
     invariant(_state == State::kCommandReply);
+    // Eagerly allocate reserveBytes bytes.
+    _builder.reserveBytes(reserveBytes);
+    // Claim our reservation immediately so we can actually write data to it.
+    _builder.claimReservedBytes(reserveBytes);
     _state = State::kMetadata;
     return _builder;
 }
diff --git a/src/mongo/rpc/legacy_reply_builder.h b/src/mongo/rpc/legacy_reply_builder.h
index de229970360..cf5f15c3d88 100644
--- a/src/mongo/rpc/legacy_reply_builder.h
+++ b/src/mongo/rpc/legacy_reply_builder.h
@@ -52,7 +52,7 @@ public:
 
     LegacyReplyBuilder& setRawCommandReply(const BSONObj& commandReply) final;
 
-    BufBuilder& getInPlaceReplyBuilder() final;
+    BufBuilder& getInPlaceReplyBuilder(std::size_t) final;
 
     LegacyReplyBuilder& setMetadata(const BSONObj& metadata) final;
 
diff --git a/src/mongo/rpc/reply_builder_interface.h b/src/mongo/rpc/reply_builder_interface.h
index ec0195df32e..57bf36b0fc0 100644
--- a/src/mongo/rpc/reply_builder_interface.h
+++ b/src/mongo/rpc/reply_builder_interface.h
@@ -70,7 +70,7 @@ public:
     /**
      * Returns a BufBuilder suitable for building a command reply in place.
      */
-    virtual BufBuilder& getInPlaceReplyBuilder() = 0;
+    virtual BufBuilder& getInPlaceReplyBuilder(std::size_t reserveBytes) = 0;
 
     virtual ReplyBuilderInterface& setMetadata(const BSONObj& metadata) = 0;