commit 21354445ad2f6a5507845922d4fe6ab14534eb8a
Author: Denis Grebennicov <denis.grebennicov@mongodb.com>
Date:   Sat Mar 26 17:15:32 2022 +0100

    SERVER-57011 cache only dependencies is non-mutable Documents

diff --git a/src/mongo/db/exec/document_value/document.cpp b/src/mongo/db/exec/document_value/document.cpp
index 38fc2b8a758..85abdeb0b07 100644
--- a/src/mongo/db/exec/document_value/document.cpp
+++ b/src/mongo/db/exec/document_value/document.cpp
@@ -29,6 +29,8 @@
 
 #include "mongo/db/exec/document_value/document.h"
 
+#include "mongo/bson/bsontypes.h"
+#include "mongo/db/exec/document_value/document_internal.h"
 #include <boost/functional/hash.hpp>
 
 #include "mongo/bson/bson_depth.h"
@@ -194,16 +196,31 @@ Position DocumentStorage::findFieldInCache(T requested) const {
 template Position DocumentStorage::findFieldInCache<StringData>(StringData field) const;
 template Position DocumentStorage::findFieldInCache<HashedFieldName>(HashedFieldName field) const;
 
+static BSONElement findBSONElement(const BSONObj& bson, const FieldPath& fieldPath, size_t level) {
+    const auto fieldName = fieldPath.getFieldName(level);
+
+    for (auto&& bsonElement : bson) {
+        if (fieldName == bsonElement.fieldNameStringData()) {
+            if (level == fieldPath.getPathLength() - 1) {
+                return bsonElement;
+            } else {
+                return findBSONElement(bsonElement.Obj(), fieldPath, level + 1);
+            }
+        }
+    }
+
+    return BSONElement();
+}
+
 template <typename T>
 Position DocumentStorage::findField(T field, LookupPolicy policy) const {
     if (auto pos = findFieldInCache(field); pos.found() || policy == LookupPolicy::kCacheOnly) {
         return pos;
     }
 
-    for (auto&& bsonElement : _bson) {
-        if (field == bsonElement.fieldNameStringData()) {
-            return const_cast<DocumentStorage*>(this)->constructInCache(bsonElement);
-        }
+    const auto foundElement = findBSONElement(_bson, FieldPath(field.rawData()), 0);
+    if (foundElement.ok()) {
+        return const_cast<DocumentStorage*>(this)->constructInCache(foundElement);
     }
 
     // if we got here, there's no such field
@@ -675,33 +692,11 @@ Document::getNestedFieldNonCaching(const FieldPath& dottedField) const {
     return getNestedFieldNonCachingHelper(dottedField, 0);
 }
 
-static Value getNestedFieldHelper(const Document& doc,
-                                  const FieldPath& fieldNames,
-                                  vector<Position>* positions,
-                                  size_t level) {
-    const auto fieldName = fieldNames.getFieldName(level);
-    const Position pos = doc.positionOf(fieldName);
-
-    if (!pos.found())
-        return Value();
-
-    if (positions)
-        positions->push_back(pos);
-
-    if (level == fieldNames.getPathLength() - 1)
-        return doc.getField(pos);
-
-    Value val = doc.getField(pos);
-    if (val.getType() != Object)
-        return Value();
-
-    return getNestedFieldHelper(val.getDocument(), fieldNames, positions, level + 1);
-}
-
 Value Document::getNestedField(const FieldPath& path, vector<Position>* positions) const {
     fassert(16489, path.getPathLength());
     assertFieldPathLengthOK(path);
-    return getNestedFieldHelper(*this, path, positions, 0);
+    const auto pos = positionOf(path.fullPath());
+    return getField(pos);
 }
 
 size_t Document::getApproximateSizeWithoutBackingBSON() const {
diff --git a/src/mongo/db/exec/document_value/document_value_test.cpp b/src/mongo/db/exec/document_value/document_value_test.cpp
index aee0844b164..f63b86380c8 100644
--- a/src/mongo/db/exec/document_value/document_value_test.cpp
+++ b/src/mongo/db/exec/document_value/document_value_test.cpp
@@ -152,6 +152,39 @@ void appendNestedObject(size_t depth, BSONObjBuilder* builder) {
     }
 }
 
+TEST(DocumentSerialization, ApproximateSizeForNestedDocuments) {
+    std::cout << "ApproximateSizeForNestedDocuments start" << std::endl;
+    std::string largeStr(1024, 'x');
+    const std::string fieldName = "b";
+    const auto nesting = 10;
+    for (size_t i = 0; i < nesting; ++i) {
+        std::cout << "Levels of nesting: " << i << "\n";
+        BSONObj b = BSON(fieldName << largeStr);
+        std::string path = fieldName;
+
+        // Construct our nested BSONObj....
+        for (size_t j = 0; j < i; ++j) {
+            BSONObj subobj = BSON(fieldName << std::move(b));
+            path = fieldName + '.' + path;
+            b = std::move(subobj);
+        }
+
+        auto doc = Document(b);
+        std::cout << "BSONObj " << b.toString() << std::endl;
+        std::cout << "For the path " << path
+                  << ", the reported size BEFORE we cache is: " << doc.getApproximateSize() << "\n";
+        ASSERT_GT(doc.getApproximateSize(), 1024);
+        ASSERT_LT(doc.getApproximateSize(), 1024 * 2);
+
+        // Force 'b' to b cached.
+        ASSERT_VALUE_EQ(doc.getNestedField(path), Value(largeStr));
+        std::cout << "For the path " << path
+                  << ", the reported size AFTER we cache is: " << doc.getApproximateSize() << "\n";
+    }
+
+    std::cout << "ApproximateSizeForNestedDocuments end" << std::endl;
+}
+
 TEST(DocumentSerialization, CanSerializeDocumentExactlyAtDepthLimit) {
     BSONObjBuilder builder;
     appendNestedObject(BSONDepth::getMaxAllowableDepth(), &builder);