SERVER-44102 Optimize bsonelement constructor

4 files changed, 69 insertions, 45 deletions

diff --git a/src/mongo/bson/SConscript b/src/mongo/bson/SConscript
index bb20dc7ed68..922048f2538 100644
--- a/src/mongo/bson/SConscript
+++ b/src/mongo/bson/SConscript
@@ -32,9 +32,9 @@ env.CppUnitTest(
 )
 
 env.Benchmark(
-    target='bsonobjbuilder_bm',
+    target='bson_bm',
     source=[
-        'bsonobjbuilder_bm.cpp',
+        'bson_bm.cpp',
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/base',
diff --git a/src/mongo/bson/bsonobjbuilder_bm.cpp b/src/mongo/bson/bson_bm.cpp
index ca344e9ddb6..a4c1eee04c4 100644
--- a/src/mongo/bson/bsonobjbuilder_bm.cpp
+++ b/src/mongo/bson/bson_bm.cpp
@@ -48,6 +48,22 @@ void BM_arrayBuilder(benchmark::State& state) {
     state.SetBytesProcessed(totalBytes);
 }
 
+void BM_arrayLookup(benchmark::State& state) {
+    BSONArrayBuilder builder;
+    auto len = state.range(0);
+    auto totalLen = len * 0;
+    for (auto j = 0; j < len; j++)
+        builder.append(j);
+    BSONObj array = builder.done();
+    for (auto _ : state) {
+        benchmark::ClobberMemory();
+        benchmark::DoNotOptimize(array[len]);
+        totalLen += len;
+    }
+    state.SetItemsProcessed(totalLen);
+}
+
 BENCHMARK(BM_arrayBuilder)->Ranges({{{1}, {100'000}}});
+BENCHMARK(BM_arrayLookup)->Ranges({{{1}, {100'000}}});
 
 }  // namespace mongo
diff --git a/src/mongo/bson/bsonelement.cpp b/src/mongo/bson/bsonelement.cpp
index 707f3042cb4..25ff416444f 100644
--- a/src/mongo/bson/bsonelement.cpp
+++ b/src/mongo/bson/bsonelement.cpp
@@ -739,11 +739,11 @@ MONGO_COMPILER_NOINLINE void msgAssertedBadType [[noreturn]] (const char* data)
 }
 }  // namespace
 
-int BSONElement::computeSize() const {
+int BSONElement::computeSize(int8_t type, const char* elem, int fieldNameSize) {
     enum SizeStyle : uint8_t {
         kFixed,         // Total size is a fixed amount + key length.
         kIntPlusFixed,  // Like Fixed, but also add in the int32 immediately following the key.
-        kRegEx,         // Handled specially.
+        kSpecial,       // Handled specially: RegEx, MinKey, MaxKey.
     };
     struct SizeInfo {
         uint8_t style : 2;
@@ -751,8 +751,8 @@ int BSONElement::computeSize() const {
     };
     MONGO_STATIC_ASSERT(sizeof(SizeInfo) == 1);
 
-    // This table should take 20 bytes. Align to next power of 2 to avoid splitting across cache
-    // lines unnecessarily.
+    // This table should take 32 bytes. Align to that size to avoid splitting across cache lines
+    // unnecessarily.
     static constexpr SizeInfo kSizeInfoTable alignas(32)[] = {
         {SizeStyle::kFixed, 1},          // EOO
         {SizeStyle::kFixed, 9},          // NumberDouble
@@ -765,7 +765,7 @@ int BSONElement::computeSize() const {
         {SizeStyle::kFixed, 2},          // Bool
         {SizeStyle::kFixed, 9},          // Date
         {SizeStyle::kFixed, 1},          // Null
-        {SizeStyle::kRegEx},             // Regex
+        {SizeStyle::kSpecial},           // Regex
         {SizeStyle::kIntPlusFixed, 17},  // DBRef
         {SizeStyle::kIntPlusFixed, 5},   // Code
         {SizeStyle::kIntPlusFixed, 5},   // Symbol
@@ -774,38 +774,47 @@ int BSONElement::computeSize() const {
         {SizeStyle::kFixed, 9},          // Timestamp
         {SizeStyle::kFixed, 9},          // Long
         {SizeStyle::kFixed, 17},         // Decimal
+        {SizeStyle::kSpecial},           // reserved 20
+        {SizeStyle::kSpecial},           // reserved 21
+        {SizeStyle::kSpecial},           // reserved 22
+        {SizeStyle::kSpecial},           // reserved 23
+        {SizeStyle::kSpecial},           // reserved 24
+        {SizeStyle::kSpecial},           // reserved 25
+        {SizeStyle::kSpecial},           // reserved 26
+        {SizeStyle::kSpecial},           // reserved 27
+        {SizeStyle::kSpecial},           // reserved 28
+        {SizeStyle::kSpecial},           // reserved 29
+        {SizeStyle::kSpecial},           // reserved 30
+        {SizeStyle::kSpecial},           // MinKey,  MaxKey
     };
-    MONGO_STATIC_ASSERT((sizeof(kSizeInfoTable) / sizeof(kSizeInfoTable[0])) == JSTypeMax + 1);
-
-    // This is the start of the runtime code for this function. Everything above happens at compile
-    // time. This function attempts to push complex handling of unlikely events out-of-line to
-    // ensure that the common cases never need to spill any registers (at least on x64 with
-    // gcc-5.4), which reduces the function call overhead.
-    int8_t type = *data;
-    if (MONGO_unlikely(type < 0 || type > JSTypeMax)) {
-        if (MONGO_unlikely(type != MinKey && type != MaxKey)) {
-            msgAssertedBadType(data);
-        }
-
-        // MinKey and MaxKey should be treated the same as Null
-        type = jstNULL;
+    MONGO_STATIC_ASSERT(sizeof(kSizeInfoTable) == 32);
+
+    // This function attempts to push complex handling of unlikely events out-of-line to ensure that
+    // the common cases never need to spill any registers, which reduces the function call overhead.
+    // Most invalid types have type != sizeInfoIndex and fall through to the cold path, as do RegEx,
+    // MinKey, MaxKey and the remaining invalid types mapping to SizeStyle::kSpecial.
+    int sizeInfoIndex = type % sizeof(kSizeInfoTable);
+    const auto sizeInfo = kSizeInfoTable[sizeInfoIndex];
+    if (MONGO_likely(type == sizeInfoIndex)) {
+        if (sizeInfo.style == SizeStyle::kFixed)
+            return sizeInfo.bytes + fieldNameSize;
+        if (MONGO_likely(sizeInfo.style == SizeStyle::kIntPlusFixed))
+            return sizeInfo.bytes + fieldNameSize +
+                ConstDataView(elem + fieldNameSize + 1).read<LittleEndian<int32_t>>();
     }
 
-    const auto sizeInfo = kSizeInfoTable[type];
-    if (sizeInfo.style == SizeStyle::kFixed)
-        return sizeInfo.bytes + fieldNameSize();
-    if (MONGO_likely(sizeInfo.style == SizeStyle::kIntPlusFixed))
-        return sizeInfo.bytes + fieldNameSize() + valuestrsize();
-
-    return [this, type]() {
-        // Regex is two c-strings back-to-back.
-        invariant(type == BSONType::RegEx);
-        const char* p = value();
-        size_t len1 = strlen(p);
-        p = p + len1 + 1;
-        size_t len2 = strlen(p);
-        return (len1 + 1 + len2 + 1) + fieldNameSize() + 1;
-    }();
+    // The following code handles all special cases: MinKey, MaxKey, RegEx and invalid types.
+    if (type == MaxKey || type == MinKey)
+        return fieldNameSize + 1;
+    if (type != BSONType::RegEx)
+        msgAssertedBadType(elem);
+
+    // RegEx is two c-strings back-to-back.
+    const char* p = elem + fieldNameSize + 1;
+    size_t len1 = strlen(p);
+    p = p + len1 + 1;
+    size_t len2 = strlen(p);
+    return (len1 + 1 + len2 + 1) + fieldNameSize + 1;
 }
 
 std::string BSONElement::toString(bool includeFieldName, bool full) const {
diff --git a/src/mongo/bson/bsonelement.h b/src/mongo/bson/bsonelement.h
index 892be6df831..a688d40eb60 100644
--- a/src/mongo/bson/bsonelement.h
+++ b/src/mongo/bson/bsonelement.h
@@ -687,13 +687,12 @@ public:
 
     // @param maxLen don't scan more than maxLen bytes
     explicit BSONElement(const char* d) : data(d) {
-        if (eoo()) {
-            fieldNameSize_ = 0;
-            totalSize = 1;
-        } else {
-            fieldNameSize_ = strlen(d + 1 /*skip type*/) + 1 /*include NUL byte*/;
-            totalSize = computeSize();
-        }
+        // While we should skip the type, and add 1 for the terminating null byte, just include
+        // the type byte in the strlen call: the extra byte cancels out. As an extra bonus, this
+        // also handles the EOO case, where the type byte is 0.
+        uint8_t type = *d;
+        fieldNameSize_ = strlen(d);
+        totalSize = computeSize(type, d, fieldNameSize_);
     }
 
     struct CachedSizeTag {};  // Opts in to next constructor.
@@ -714,7 +713,7 @@ public:
                 fieldNameSize_ = fieldNameSize;
             }
             if (totalSize == -1) {
-                this->totalSize = computeSize();
+                this->totalSize = computeSize(*d, d, fieldNameSize_);
             } else {
                 this->totalSize = totalSize;
             }
@@ -762,7 +761,7 @@ private:
     }
 
     // Only called from constructors.
-    int computeSize() const;
+    static int computeSize(int8_t type, const char* data, int fieldNameSize);
 };
 
 inline bool BSONElement::trueValue() const {