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

5 files changed, 13 insertions, 6 deletions

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;