SERVER-60181 Avoid pushing insert-only diffs to StringMap when applying docDiff

Also fix usage of reserveBytes when applying docDiff. Reserved bytes need to be claimed.

2 files changed, 37 insertions, 17 deletions

diff --git a/src/mongo/db/update/document_diff_applier.cpp b/src/mongo/db/update/document_diff_applier.cpp
index 946d295c975..ec592ef3363 100644
--- a/src/mongo/db/update/document_diff_applier.cpp
+++ b/src/mongo/db/update/document_diff_applier.cpp
@@ -74,17 +74,17 @@ struct DocumentDiffTables {
         uassert(4728000, str::stream() << "duplicate field name in diff: " << fieldName, inserted);
     }
 
-    // Map from field name to modification for that field.
+    // Map from field name to modification for that field. Not populated for insert only diffs.
     StringDataMap<FieldModification> fieldMap;
 
     // Order in which new fields should be added to the pre image.
     std::vector<BSONElement> fieldsToInsert;
-    std::size_t sizeOfFieldsToInsert = 0;
     // Diff only inserts fields, no deletes or updates
     bool insertOnly = false;
 };
 
-DocumentDiffTables buildObjDiffTables(DocumentDiffReader* reader) {
+DocumentDiffTables buildObjDiffTables(DocumentDiffReader* reader,
+                                      bool mustCheckExistenceForInsertOperations) {
     DocumentDiffTables out;
     out.insertOnly = true;
 
@@ -101,17 +101,19 @@ DocumentDiffTables buildObjDiffTables(DocumentDiffReader* reader) {
         out.insertOnly = false;
     }
 
-    boost::optional<BSONElement> nextInsert;
-    while ((nextInsert = reader->nextInsert())) {
-        out.safeInsert(nextInsert->fieldNameStringData(), Insert{*nextInsert});
-        out.fieldsToInsert.push_back(*nextInsert);
-        out.sizeOfFieldsToInsert += out.fieldsToInsert.back().size();
-    }
-
     for (auto next = reader->nextSubDiff(); next; next = reader->nextSubDiff()) {
         out.safeInsert(next->first, SubDiff{next->second});
         out.insertOnly = false;
     }
+
+    boost::optional<BSONElement> nextInsert;
+    while ((nextInsert = reader->nextInsert())) {
+        if (mustCheckExistenceForInsertOperations || !out.insertOnly) {
+            out.safeInsert(nextInsert->fieldNameStringData(), Insert{*nextInsert});
+        }
+
+        out.fieldsToInsert.push_back(*nextInsert);
+    }
     return out;
 }
 
@@ -193,7 +195,8 @@ int32_t computeDamageOnObject(const BSONObj& preImageRoot,
                               BufBuilder* bufBuilder,
                               size_t offsetRoot,
                               bool mustCheckExistenceForInsertOperations) {
-    const DocumentDiffTables tables = buildObjDiffTables(reader);
+    const DocumentDiffTables tables =
+        buildObjDiffTables(reader, mustCheckExistenceForInsertOperations);
     int32_t diffSize = 0;
     // Reserves four bytes for the total size. Stores the offset instead of the actual pointer in
     // case the buffer grows and invalidates the pointer. Will update the value at the end.
@@ -490,10 +493,10 @@ public:
                            BSONObjBuilder* builder) {
         // First build some tables so we can quickly apply the diff. We shouldn't need to examine
         // the diff again once this is done.
-        const DocumentDiffTables tables = buildObjDiffTables(reader);
+        const DocumentDiffTables tables =
+            buildObjDiffTables(reader, _mustCheckExistenceForInsertOperations);
 
         if (!_mustCheckExistenceForInsertOperations && tables.insertOnly) {
-            builder->bb().reserveBytes(preImage.objsize() + tables.sizeOfFieldsToInsert);
             builder->appendElements(preImage);
             for (auto&& elt : tables.fieldsToInsert) {
                 builder->append(elt);
@@ -718,6 +721,15 @@ ApplyDiffOutput applyDiff(const BSONObj& pre,
     BSONObjBuilder out;
     DiffApplier applier(indexData, mustCheckExistenceForInsertOperations);
     FieldRef path;
+
+    // Use size of pre + diff as an approximation for size needed for post object when the diff is
+    // insert-only. This is an upper bound for the size needed to avoid re-allocations when applying
+    // the diff.
+    const auto estimatedSize = pre.objsize() + diff.objsize();
+    // The interface around reserving bytes is strange. You need to "claim" the reserved bytes for
+    // it to take effect.
+    out.bb().reserveBytes(estimatedSize);
+    out.bb().claimReservedBytes(estimatedSize);
     applier.applyDiffToObject(pre, &path, &reader, &out);
     return {out.obj(), applier.indexesAffected()};
 }
diff --git a/src/mongo/db/update/document_diff_serialization.cpp b/src/mongo/db/update/document_diff_serialization.cpp
index 285da519c48..0f5533e1d06 100644
--- a/src/mongo/db/update/document_diff_serialization.cpp
+++ b/src/mongo/db/update/document_diff_serialization.cpp
@@ -36,6 +36,8 @@
 
 #include "mongo/util/str.h"
 #include "mongo/util/string_map.h"
+#include <boost/container/flat_map.hpp>
+#include <boost/container/static_vector.hpp>
 
 namespace mongo {
 namespace diff_tree {
@@ -511,10 +513,16 @@ DocumentDiffReader::DocumentDiffReader(const Diff& diff) : _diff(diff) {
     static_assert(kInsertSectionFieldName.size() == 1);
     static_assert(kUpdateSectionFieldName.size() == 1);
 
-    const std::map<char, Section> sections{{kDeleteSectionFieldName[0], Section{&_deletes, 1}},
-                                           {kUpdateSectionFieldName[0], Section{&_updates, 2}},
-                                           {kInsertSectionFieldName[0], Section{&_inserts, 3}},
-                                           {kSubDiffSectionFieldPrefix, Section{&_subDiffs, 4}}};
+    // Create a map only using stack memory for this temporary helper map. Make sure to update the
+    // size for the 'static_vector' if changes are made to the number of elements we hold.
+    const boost::container::flat_map<char,
+                                     Section,
+                                     std::less<char>,
+                                     boost::container::static_vector<std::pair<char, Section>, 4>>
+        sections{{kDeleteSectionFieldName[0], Section{&_deletes, 1}},
+                 {kUpdateSectionFieldName[0], Section{&_updates, 2}},
+                 {kInsertSectionFieldName[0], Section{&_inserts, 3}},
+                 {kSubDiffSectionFieldPrefix, Section{&_subDiffs, 4}}};
 
     char prev = 0;
     bool hasSubDiffSections = false;