SERVER-60273 Calculate the estimated size for missing types in sbe::value::getApproximateSize()

25 files changed, 197 insertions, 4 deletions

diff --git a/src/mongo/db/exec/js_function.cpp b/src/mongo/db/exec/js_function.cpp
index 58e3a880b00..f467f48af99 100644
--- a/src/mongo/db/exec/js_function.cpp
+++ b/src/mongo/db/exec/js_function.cpp
@@ -91,4 +91,9 @@ bool JsFunction::runAsPredicate(const BSONObj& obj) const {
     return _scope->getBoolean("__returnValue");
 }
 
+size_t JsFunction::getApproximateSize() const {
+    // The memory pointed to by _scope is owned by the MozJS engine, so we do not account
+    // for the size of that memory here.
+    return sizeof(JsFunction);
+}
 }  // namespace mongo
diff --git a/src/mongo/db/exec/js_function.h b/src/mongo/db/exec/js_function.h
index a658deb81af..268f8a54490 100644
--- a/src/mongo/db/exec/js_function.h
+++ b/src/mongo/db/exec/js_function.h
@@ -41,6 +41,8 @@ public:
 
     bool runAsPredicate(const BSONObj& obj) const;
 
+    size_t getApproximateSize() const;
+
 private:
     std::shared_ptr<Scope> _scope;
     ScriptingFunction _func;
diff --git a/src/mongo/db/exec/sbe/values/slot.cpp b/src/mongo/db/exec/sbe/values/slot.cpp
index e17ebccae70..d827efd1747 100644
--- a/src/mongo/db/exec/sbe/values/slot.cpp
+++ b/src/mongo/db/exec/sbe/values/slot.cpp
@@ -32,6 +32,8 @@
 #include "mongo/db/exec/sbe/values/slot.h"
 
 #include "mongo/bson/util/builder.h"
+#include "mongo/db/exec/js_function.h"
+#include "mongo/db/exec/sbe/values/sort_spec.h"
 #include "mongo/db/storage/key_string.h"
 #include "mongo/util/bufreader.h"
 
@@ -387,6 +389,7 @@ int getApproximateSize(TypeTags tag, Value val) {
         case TypeTags::MinKey:
         case TypeTags::MaxKey:
         case TypeTags::bsonUndefined:
+        case TypeTags::LocalLambda:
             break;
         // There are deep types.
         case TypeTags::NumberDecimal:
@@ -466,17 +469,28 @@ int getApproximateSize(TypeTags tag, Value val) {
             // including the 'length' field itself.
             result += ConstDataView(getRawPointerView(val)).read<LittleEndian<uint32_t>>();
             break;
-        // These types are currently unsupported. Temporarily break in these cases because
-        // calculating the estimated size of a cached SBE plan requires a size value.
-        // TODO: SERVER-60273 Calculate the estimated size for missing types.
-        case TypeTags::LocalLambda:
         case TypeTags::pcreRegex:
+            result += getPcreRegexView(val)->getApproximateSize();
+            break;
         case TypeTags::timeZoneDB:
+            // This type points to a block of memory that it doesn't own, so we don't acccount
+            // for the size of this block of memory here.
+            break;
         case TypeTags::jsFunction:
+            result += getJsFunctionView(val)->getApproximateSize();
+            break;
         case TypeTags::shardFilterer:
+            result += getShardFiltererView(val)->getApproximateSize();
+            break;
         case TypeTags::collator:
+            // This type points to a block of memory that it doesn't own, so we don't acccount
+            // for the size of this block of memory here.
+            break;
         case TypeTags::ftsMatcher:
+            result += getFtsMatcherView(val)->getApproximateSize();
+            break;
         case TypeTags::sortSpec:
+            result += getSortSpecView(val)->getApproximateSize();
             break;
         default:
             MONGO_UNREACHABLE;
diff --git a/src/mongo/db/exec/sbe/values/sort_spec.h b/src/mongo/db/exec/sbe/values/sort_spec.h
index 0082b6b5bf1..39d0e195fb7 100644
--- a/src/mongo/db/exec/sbe/values/sort_spec.h
+++ b/src/mongo/db/exec/sbe/values/sort_spec.h
@@ -57,6 +57,8 @@ public:
         return _collator;
     }
 
+    size_t getApproximateSize() const;
+
 private:
     BtreeKeyGenerator initKeyGen() const;
 
diff --git a/src/mongo/db/exec/sbe/values/value.cpp b/src/mongo/db/exec/sbe/values/value.cpp
index 0f13a52691f..449c6b2a0d7 100644
--- a/src/mongo/db/exec/sbe/values/value.cpp
+++ b/src/mongo/db/exec/sbe/values/value.cpp
@@ -168,6 +168,12 @@ size_t PcreRegex::getNumberCaptures() const {
     return static_cast<size_t>(numCaptures);
 }
 
+size_t PcreRegex::getApproximateSize() const {
+    size_t pcreSize;
+    pcre_fullinfo(_pcrePtr, nullptr, PCRE_INFO_SIZE, &pcreSize);
+    return sizeof(PcreRegex) + _pattern.size() + 1 + _options.size() + 1 + pcreSize;
+}
+
 KeyString::Value SortSpec::generateSortKey(const BSONObj& obj) const {
     KeyStringSet keySet;
     SharedBufferFragmentBuilder allocator(KeyString::HeapBuilder::kHeapAllocatorDefaultBytes);
@@ -208,6 +214,15 @@ BtreeKeyGenerator SortSpec::initKeyGen() const {
     return {std::move(fields), std::move(fixed), isSparse, _collator, version, ordering};
 }
 
+size_t SortSpec::getApproximateSize() const {
+    // _collator points to a block of memory that SortSpec doesn't own, so we don't acccount
+    // for the size of this block of memory here.
+    auto size = sizeof(SortSpec);
+    size += _sortPattern.isOwned() ? _sortPattern.objsize() : 0;
+    size += _keyGen.getApproximateSize() - sizeof(_keyGen);
+    return size;
+}
+
 std::pair<TypeTags, Value> makeCopyJsFunction(const JsFunction& jsFunction) {
     auto ownedJsFunction = bitcastFrom<JsFunction*>(new JsFunction(jsFunction));
     return {TypeTags::jsFunction, ownedJsFunction};
diff --git a/src/mongo/db/exec/sbe/values/value.h b/src/mongo/db/exec/sbe/values/value.h
index 1d085c47a37..6bafc07719c 100644
--- a/src/mongo/db/exec/sbe/values/value.h
+++ b/src/mongo/db/exec/sbe/values/value.h
@@ -905,6 +905,8 @@ public:
 
     size_t getNumberCaptures() const;
 
+    size_t getApproximateSize() const;
+
 private:
     void _compile();
 
diff --git a/src/mongo/db/exec/shard_filterer.h b/src/mongo/db/exec/shard_filterer.h
index 35c747f0b7c..cd32d8b656d 100644
--- a/src/mongo/db/exec/shard_filterer.h
+++ b/src/mongo/db/exec/shard_filterer.h
@@ -70,5 +70,7 @@ public:
      * the shard key.
      */
     virtual const KeyPattern& getKeyPattern() const = 0;
+
+    virtual size_t getApproximateSize() const = 0;
 };
 }  // namespace mongo
diff --git a/src/mongo/db/exec/shard_filterer_impl.cpp b/src/mongo/db/exec/shard_filterer_impl.cpp
index 013fa65ce9a..8f7949f4ef3 100644
--- a/src/mongo/db/exec/shard_filterer_impl.cpp
+++ b/src/mongo/db/exec/shard_filterer_impl.cpp
@@ -88,4 +88,12 @@ const KeyPattern& ShardFiltererImpl::getKeyPattern() const {
     invariant(_keyPattern);
     return _keyPattern->getKeyPattern();
 }
+
+size_t ShardFiltererImpl::getApproximateSize() const {
+    // _collectionFilter contains a shared_ptr to metadata, but it does not own this metadata,
+    // so we don't account for the size of the metadata here.
+    auto size = sizeof(ShardFiltererImpl);
+    size += _keyPattern ? _keyPattern->getApproximateSize() - sizeof(ShardKeyPattern) : 0;
+    return 0;
+}
 }  // namespace mongo
diff --git a/src/mongo/db/exec/shard_filterer_impl.h b/src/mongo/db/exec/shard_filterer_impl.h
index c0b4ef3e0ca..f50428d3768 100644
--- a/src/mongo/db/exec/shard_filterer_impl.h
+++ b/src/mongo/db/exec/shard_filterer_impl.h
@@ -54,6 +54,8 @@ public:
 
     const KeyPattern& getKeyPattern() const override;
 
+    size_t getApproximateSize() const override;
+
 private:
     DocumentBelongsResult keyBelongsToMeHelper(const BSONObj& doc) const;
 
diff --git a/src/mongo/db/exec/shard_filterer_mock.h b/src/mongo/db/exec/shard_filterer_mock.h
index f76bd487e82..d896e257d21 100644
--- a/src/mongo/db/exec/shard_filterer_mock.h
+++ b/src/mongo/db/exec/shard_filterer_mock.h
@@ -63,6 +63,12 @@ public:
         return true;
     }
 
+    size_t getApproximateSize() const override {
+        auto size = sizeof(ConstantFilterMock);
+        size += _pattern.getApproximateSize() - sizeof(KeyPattern);
+        return size;
+    }
+
 private:
     const bool _pass;
     const KeyPattern _pattern;
@@ -105,6 +111,12 @@ public:
         return true;
     }
 
+    size_t getApproximateSize() const override {
+        auto size = sizeof(AllNullShardKeyFilterMock);
+        size += _pattern.getApproximateSize() - sizeof(ShardKeyPattern);
+        return size;
+    }
+
 private:
     const ShardKeyPattern _pattern;
 };
diff --git a/src/mongo/db/fts/fts_matcher.cpp b/src/mongo/db/fts/fts_matcher.cpp
index be9daa5801d..a248d9600d8 100644
--- a/src/mongo/db/fts/fts_matcher.cpp
+++ b/src/mongo/db/fts/fts_matcher.cpp
@@ -141,6 +141,13 @@ bool FTSMatcher::negativePhrasesMatch(const BSONObj& obj) const {
     return true;
 }
 
+size_t FTSMatcher::getApproximateSize() const {
+    auto size = sizeof(FTSMatcher);
+    size += _query.getApproximateSize() - sizeof(_query);
+    size += _spec.getApproximateSize() - sizeof(_spec);
+    return size;
+}
+
 bool FTSMatcher::_phraseMatch(const string& phrase, const BSONObj& obj) const {
     FTSElementIterator it(_spec, obj);
 
diff --git a/src/mongo/db/fts/fts_matcher.h b/src/mongo/db/fts/fts_matcher.h
index 1f3de2c91f6..8698ba0648d 100644
--- a/src/mongo/db/fts/fts_matcher.h
+++ b/src/mongo/db/fts/fts_matcher.h
@@ -83,6 +83,8 @@ public:
         return _spec;
     }
 
+    size_t getApproximateSize() const;
+
 private:
     /**
      * For matching, can we skip the positive term check?  This is done as optimization when
diff --git a/src/mongo/db/fts/fts_query.h b/src/mongo/db/fts/fts_query.h
index 2cfce1326cb..5469f83a313 100644
--- a/src/mongo/db/fts/fts_query.h
+++ b/src/mongo/db/fts/fts_query.h
@@ -97,6 +97,10 @@ public:
      */
     virtual std::unique_ptr<FTSQuery> clone() const = 0;
 
+    virtual size_t getApproximateSize() const {
+        return sizeof(FTSQuery) + _query.size() + 1 + _language.size() + 1;
+    }
+
 private:
     std::string _query;
     std::string _language;
diff --git a/src/mongo/db/fts/fts_query_impl.cpp b/src/mongo/db/fts/fts_query_impl.cpp
index e996a12862c..c5b7ed85c0f 100644
--- a/src/mongo/db/fts/fts_query_impl.cpp
+++ b/src/mongo/db/fts/fts_query_impl.cpp
@@ -206,5 +206,32 @@ BSONObj FTSQueryImpl::toBSON() const {
     bob.append("negatedPhrases", getNegatedPhr());
     return bob.obj();
 }
+
+size_t FTSQueryImpl::getApproximateSize() const {
+    auto computeVectorSize = [](const std::vector<std::string>& v) {
+        size_t size = 0;
+        for (const auto& str : v) {
+            size += sizeof(str) + str.size() + 1;
+        }
+        return size;
+    };
+
+    auto computeSetSize = [](const std::set<std::string>& s) {
+        size_t size = 0;
+        for (const auto& str : s) {
+            size += sizeof(str) + str.size() + 1;
+        }
+        return size;
+    };
+
+    auto size = sizeof(FTSQueryImpl);
+    size += FTSQuery::getApproximateSize() - sizeof(FTSQuery);
+    size += computeSetSize(_positiveTerms);
+    size += computeSetSize(_negatedTerms);
+    size += computeVectorSize(_positivePhrases);
+    size += computeVectorSize(_negatedPhrases);
+    size += computeSetSize(_termsForBounds);
+    return size;
+}
 }  // namespace fts
 }  // namespace mongo
diff --git a/src/mongo/db/fts/fts_query_impl.h b/src/mongo/db/fts/fts_query_impl.h
index 97cdb8388df..06d400b5a19 100644
--- a/src/mongo/db/fts/fts_query_impl.h
+++ b/src/mongo/db/fts/fts_query_impl.h
@@ -75,6 +75,8 @@ public:
      */
     BSONObj toBSON() const;
 
+    size_t getApproximateSize() const final;
+
 private:
     void _addTerms(FTSTokenizer* tokenizer, const std::string& tokens, bool negated);
 
diff --git a/src/mongo/db/fts/fts_spec.cpp b/src/mongo/db/fts/fts_spec.cpp
index 7226ff5bb44..4306cb9e3c5 100644
--- a/src/mongo/db/fts/fts_spec.cpp
+++ b/src/mongo/db/fts/fts_spec.cpp
@@ -512,5 +512,32 @@ StatusWith<BSONObj> FTSSpec::fixSpec(const BSONObj& spec) {
 
     return b.obj();
 }
+
+size_t FTSSpec::getApproximateSize() const {
+    auto computeVectorSize = [](const std::vector<std::string>& v) {
+        size_t size = 0;
+        for (const auto& str : v) {
+            size += sizeof(str) + str.size() + 1;
+        }
+        return size;
+    };
+
+    auto computeWeightsSize = [](const Weights& w) {
+        size_t size = 0;
+        for (const auto& p : w) {
+            size += sizeof(p) + p.first.size() + 1;
+        }
+        return size;
+    };
+
+    // _defaultLanguage is owned by the LanguageRegistry class and may be shared across many
+    // FTSSpec's, so we don't account for the size of _defaultLanguage here.
+    auto size = sizeof(FTSSpec);
+    size += _languageOverrideField.size() + 1;
+    size += computeWeightsSize(_weights);
+    size += computeVectorSize(_extraBefore);
+    size += computeVectorSize(_extraAfter);
+    return size;
+}
 }  // namespace fts
 }  // namespace mongo
diff --git a/src/mongo/db/fts/fts_spec.h b/src/mongo/db/fts/fts_spec.h
index 1061865bcbd..dbc67857b82 100644
--- a/src/mongo/db/fts/fts_spec.h
+++ b/src/mongo/db/fts/fts_spec.h
@@ -122,6 +122,8 @@ public:
         return _textIndexVersion;
     }
 
+    size_t getApproximateSize() const;
+
 private:
     //
     // Helper methods.  Invoked for TEXT_INDEX_VERSION_2 spec objects only.
diff --git a/src/mongo/db/index/btree_key_generator.cpp b/src/mongo/db/index/btree_key_generator.cpp
index 017089822b5..62ec65f8080 100644
--- a/src/mongo/db/index/btree_key_generator.cpp
+++ b/src/mongo/db/index/btree_key_generator.cpp
@@ -276,6 +276,36 @@ void BtreeKeyGenerator::getKeys(SharedBufferFragmentBuilder& pooledBufferBuilder
     }
 }
 
+size_t BtreeKeyGenerator::getApproximateSize() const {
+    auto computePositionalInfoSize = [](const std::vector<PositionalPathInfo>& v) {
+        size_t size = 0;
+        for (const auto& elem : v) {
+            size += elem.getApproximateSize();
+        }
+        return size;
+    };
+
+    // _fieldNames contains pointers to blocks of memory that BtreeKeyGenerator doesn't own,
+    // so we don't account for the sizes of those blocks of memory here. Likewise, _collator
+    // points to a block of memory that BtreeKeyGenerator doesn't own, so we don't acccount
+    // for the size of this block of memory either.
+    auto size = sizeof(BtreeKeyGenerator);
+    size += _fieldNames.size() * sizeof(const char*);
+    size += _nullKeyString.getApproximateSize() - sizeof(_nullKeyString);
+    size += _fixed.size() * sizeof(BSONElement);
+    size += computePositionalInfoSize(_emptyPositionalInfo);
+    size += _pathLengths.size() * sizeof(size_t);
+    return size;
+}
+
+size_t BtreeKeyGenerator::PositionalPathInfo::getApproximateSize() const {
+    // remainingPath points to a block of memory that PositionalPathInfo doesn't own, so we
+    // don't account for the size of this block of memory here.
+    auto size = sizeof(PositionalPathInfo);
+    size += arrayObj.isOwned() ? arrayObj.objsize() : 0;
+    return size;
+}
+
 void BtreeKeyGenerator::_getKeysWithoutArray(SharedBufferFragmentBuilder& pooledBufferBuilder,
                                              const BSONObj& obj,
                                              boost::optional<RecordId> id,
diff --git a/src/mongo/db/index/btree_key_generator.h b/src/mongo/db/index/btree_key_generator.h
index da6bff3185c..2cf55d6a770 100644
--- a/src/mongo/db/index/btree_key_generator.h
+++ b/src/mongo/db/index/btree_key_generator.h
@@ -82,6 +82,8 @@ public:
                  MultikeyPaths* multikeyPaths,
                  boost::optional<RecordId> id = boost::none) const;
 
+    size_t getApproximateSize() const;
+
 private:
     /**
      * Stores info regarding traversal of a positional path. A path through a document is
@@ -103,6 +105,8 @@ private:
             return !positionallyIndexedElt.eoo();
         }
 
+        size_t getApproximateSize() const;
+
         // Stores the array element indexed by position. If the key pattern has no positional
         // element, then this is EOO.
         //
diff --git a/src/mongo/db/keypattern.cpp b/src/mongo/db/keypattern.cpp
index 8e91e561c01..5a64ad8a27a 100644
--- a/src/mongo/db/keypattern.cpp
+++ b/src/mongo/db/keypattern.cpp
@@ -115,4 +115,10 @@ BSONObj KeyPattern::globalMax() const {
     return extendRangeBound(BSONObj(), true);
 }
 
+size_t KeyPattern::getApproximateSize() const {
+    auto size = sizeof(KeyPattern);
+    size += _pattern.isOwned() ? _pattern.objsize() : 0;
+    return size;
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/keypattern.h b/src/mongo/db/keypattern.h
index eb3c9d9da41..986d259eebd 100644
--- a/src/mongo/db/keypattern.h
+++ b/src/mongo/db/keypattern.h
@@ -125,6 +125,8 @@ public:
 
     BSONObj globalMax() const;
 
+    size_t getApproximateSize() const;
+
 private:
     BSONObj _pattern;
 };
diff --git a/src/mongo/db/pipeline/document_source_internal_shard_filter_test.cpp b/src/mongo/db/pipeline/document_source_internal_shard_filter_test.cpp
index 55a01a23a36..6678f433470 100644
--- a/src/mongo/db/pipeline/document_source_internal_shard_filter_test.cpp
+++ b/src/mongo/db/pipeline/document_source_internal_shard_filter_test.cpp
@@ -69,6 +69,10 @@ public:
     bool keyBelongsToMe(const BSONObj& obj) const override {
         MONGO_UNREACHABLE;
     }
+
+    size_t getApproximateSize() const override {
+        MONGO_UNREACHABLE;
+    }
 };
 
 /**
diff --git a/src/mongo/db/pipeline/process_interface/stub_lookup_single_document_process_interface.h b/src/mongo/db/pipeline/process_interface/stub_lookup_single_document_process_interface.h
index 05fbf53349d..40a77f09e14 100644
--- a/src/mongo/db/pipeline/process_interface/stub_lookup_single_document_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/stub_lookup_single_document_process_interface.h
@@ -62,6 +62,10 @@ public:
     bool keyBelongsToMe(const BSONObj& shardKey) const override {
         MONGO_UNREACHABLE;
     }
+
+    size_t getApproximateSize() const override {
+        MONGO_UNREACHABLE;
+    }
 };
 
 /**
diff --git a/src/mongo/db/storage/key_string.cpp b/src/mongo/db/storage/key_string.cpp
index bdefa0c6b24..0a93caa0c9a 100644
--- a/src/mongo/db/storage/key_string.cpp
+++ b/src/mongo/db/storage/key_string.cpp
@@ -2705,6 +2705,12 @@ void Value::serializeWithoutRecordIdLong(BufBuilder& buf) const {
     buf.appendBuf(_buffer.get() + _ksSize, _buffer.size() - _ksSize);  // Serialize TypeBits
 }
 
+size_t Value::getApproximateSize() const {
+    auto size = sizeof(Value);
+    size += !_buffer.isShared() ? SharedBuffer::kHolderSize + _buffer.size() : 0;
+    return size;
+}
+
 template class BuilderBase<Builder>;
 template class BuilderBase<HeapBuilder>;
 template class BuilderBase<PooledBuilder>;
diff --git a/src/mongo/db/storage/key_string.h b/src/mongo/db/storage/key_string.h
index 36da68fc3e4..70f72788995 100644
--- a/src/mongo/db/storage/key_string.h
+++ b/src/mongo/db/storage/key_string.h
@@ -440,6 +440,8 @@ public:
         return _version;
     }
 
+    size_t getApproximateSize() const;
+
 private:
     Version _version;
     // _ksSize is the total length that the KeyString takes up in the buffer.