SERVER-46805 Limit memory usage in the second phase of validate and add parameter to configure this limit

(cherry picked from commit e398960773331f3b3afe000f2830c84868aaf9e7)

12 files changed, 208 insertions, 23 deletions

diff --git a/jstests/disk/libs/wt_file_helper.js b/jstests/disk/libs/wt_file_helper.js
index 8b1c29c512c..fc4f1e9aa0a 100644
--- a/jstests/disk/libs/wt_file_helper.js
+++ b/jstests/disk/libs/wt_file_helper.js
@@ -207,3 +207,13 @@ let runWiredTigerTool = function(...args) {
     const cmd = ['wt'].concat(args);
     assert.eq(run.apply(undefined, cmd), 0, "error executing: " + cmd.join(' '));
 };
+
+/**
+ * Stops the given mongod, runs the truncate command on the given uri using the WiredTiger tool, and
+ * starts mongod again on the same path.
+ */
+let truncateUriAndRestartMongod = function(uri, conn) {
+    MongoRunner.stopMongod(conn, null, {skipValidation: true});
+    runWiredTigerTool("-h", conn.dbpath, "truncate", uri);
+    return startMongodOnExistingPath(conn.dbpath, {});
+};
diff --git a/jstests/noPassthrough/validate_memory_limit.js b/jstests/noPassthrough/validate_memory_limit.js
new file mode 100644
index 00000000000..841b247ddaa
--- /dev/null
+++ b/jstests/noPassthrough/validate_memory_limit.js
@@ -0,0 +1,52 @@
+/**
+ * Test that the memory usage of validate is properly limited according to the
+ * maxValidateMemoryUsageMB parameter.
+ *
+ * @tags: [requires_wiredtiger, requires_persistence]
+ */
+(function() {
+"use strict";
+
+load("jstests/disk/libs/wt_file_helper.js");
+
+const kIndexKeyLength = 1024 * 1024;
+
+const baseName = "validate_memory_limit";
+const dbpath = MongoRunner.dataPath + baseName + "/";
+let conn = MongoRunner.runMongod({dbpath: dbpath});
+let coll = conn.getDB("test").getCollection("corrupt");
+
+function corruptIndex() {
+    const uri = getUriForIndex(coll, "_id_");
+    conn = truncateUriAndRestartMongod(uri, conn);
+    coll = conn.getDB("test").getCollection("corrupt");
+}
+
+function checkValidate(errorPrefix, numMissingIndexEntries) {
+    conn.getDB("test").adminCommand({setParameter: 1, maxValidateMemoryUsageMB: 1});
+    const res = coll.validate();
+    assert.commandWorked(res);
+    assert(!res.valid);
+    assert.containsPrefix(errorPrefix, res.errors);
+    assert.eq(res.missingIndexEntries.length, numMissingIndexEntries);
+}
+
+const noneReportedPrefix =
+    "Unable to report index entry inconsistencies due to memory limitations.";
+const notAllReportedPrefix =
+    "Not all index entry inconsistencies are reported due to memory limitations.";
+
+// Insert a document with a key larger than maxValidateMemoryUsageMB so that validate does not
+// report any missing index entries.
+const indexKey = "a".repeat(kIndexKeyLength);
+assert.commandWorked(coll.insert({_id: indexKey}));
+corruptIndex();
+checkValidate(noneReportedPrefix, 0);
+
+// Insert a document with a small key so that validate reports one missing index entry.
+assert.commandWorked(coll.insert({_id: 1}));
+corruptIndex();
+checkValidate(notAllReportedPrefix, 1);
+
+MongoRunner.stopMongod(conn, null, {skipValidation: true});
+})();
\ No newline at end of file
diff --git a/jstests/noPassthrough/validate_with_long_index_name.js b/jstests/noPassthrough/validate_with_long_index_name.js
index 81c9b2dfa87..c61078b6496 100644
--- a/jstests/noPassthrough/validate_with_long_index_name.js
+++ b/jstests/noPassthrough/validate_with_long_index_name.js
@@ -29,16 +29,10 @@ function insertDocsAndBuildIndex() {
     coll.createIndex({a: 1}, {name: indexName});
 }
 
-function truncateAndRestartMongod(uri) {
-    MongoRunner.stopMongod(conn, null, {skipValidation: true});
-    runWiredTigerTool("-h", conn.dbpath, "truncate", uri);
-    conn = startMongodOnExistingPath(dbpath, {});
-    coll = conn.getDB("test").getCollection("corrupt");
-}
-
 insertDocsAndBuildIndex();
 let uri = getUriForIndex(coll, indexName);
-truncateAndRestartMongod(uri);
+conn = truncateUriAndRestartMongod(uri, conn);
+coll = conn.getDB("test").getCollection("corrupt");
 
 const missingIndexEntries = "Detected " + kNumDocs + " missing index entries.";
 const missingSizeLimitations =
@@ -53,7 +47,8 @@ assert.contains(missingSizeLimitations, res.errors);
 coll.drop();
 insertDocsAndBuildIndex();
 uri = getUriForColl(coll);
-truncateAndRestartMongod(uri);
+conn = truncateUriAndRestartMongod(uri, conn);
+coll = conn.getDB("test").getCollection("corrupt");
 
 const extraIndexEntries = "Detected " + 2 * kNumDocs + " extra index entries.";
 const extraSizeLimitations =
diff --git a/jstests/noPassthrough/validation_serverParameters.js b/jstests/noPassthrough/validation_serverParameters.js
index 4b1b7e29dea..825cdbc6259 100644
--- a/jstests/noPassthrough/validation_serverParameters.js
+++ b/jstests/noPassthrough/validation_serverParameters.js
@@ -16,4 +16,15 @@ testNumericServerParameter('maxValidateMBperSec',
                            -1 /*lowerOutOfBounds*/,
                            false /*hasUpperBound*/,
                            "unused" /*upperOutOfBounds*/);
+
+// Valid parameter values are in the range (0, infinity).
+testNumericServerParameter('maxValidateMemoryUsageMB',
+                           true /*isStartupParameter*/,
+                           true /*isRuntimeParameter*/,
+                           200 /*defaultValue*/,
+                           50 /*nonDefaultValidValue*/,
+                           true /*hasLowerBound*/,
+                           0 /*lowerOutOfBounds*/,
+                           false /*hasUpperBound*/,
+                           "unused" /*upperOutOfBounds*/);
 })();
diff --git a/src/mongo/db/catalog/SConscript b/src/mongo/db/catalog/SConscript
index 0fd0767bc5a..f011b9da52f 100644
--- a/src/mongo/db/catalog/SConscript
+++ b/src/mongo/db/catalog/SConscript
@@ -74,9 +74,9 @@ env.Library(
 )
 
 env.Library(
-    target='max_validate_mb_per_sec_idl',
+    target='validate_idl',
     source=[
-        'max_validate_mb_per_sec.idl',
+        'validate.idl',
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/idl/server_parameter',
@@ -378,6 +378,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/transaction',
         'index_build_block',
         'throttle_cursor',
+        'validate_idl',
         'validate_state',
     ],
 )
@@ -406,7 +407,7 @@ env.Library(
     LIBDEPS_PRIVATE=[
         '$BUILD_DIR/mongo/db/curop',
         '$BUILD_DIR/mongo/util/fail_point',
-        'max_validate_mb_per_sec_idl',
+        'validate_idl',
     ],
 )
 
@@ -521,7 +522,7 @@ env.CppUnitTest(
         'index_build_entry_idl',
         'index_builds_manager',
         'index_key_validate',
-        'max_validate_mb_per_sec_idl',
+        'validate_idl',
         'multi_index_block',
         'throttle_cursor',
         'validate_state',
diff --git a/src/mongo/db/catalog/collection_validation.cpp b/src/mongo/db/catalog/collection_validation.cpp
index 027fba090aa..5fc1c4e2b87 100644
--- a/src/mongo/db/catalog/collection_validation.cpp
+++ b/src/mongo/db/catalog/collection_validation.cpp
@@ -192,6 +192,9 @@ void _gatherIndexEntryErrors(OperationContext* opCtx,
                              ValidateResultsMap* indexNsResultsMap,
                              ValidateResults* result) {
     indexConsistency->setSecondPhase();
+    if (!indexConsistency->limitMemoryUsageForSecondPhase(result)) {
+        return;
+    }
 
     LOGV2_OPTIONS(
         20297, {LogComponent::kIndex}, "Starting to traverse through all the document key sets");
diff --git a/src/mongo/db/catalog/index_consistency.cpp b/src/mongo/db/catalog/index_consistency.cpp
index e016a27ebea..07289adf3a9 100644
--- a/src/mongo/db/catalog/index_consistency.cpp
+++ b/src/mongo/db/catalog/index_consistency.cpp
@@ -36,6 +36,7 @@
 #include "mongo/db/catalog/index_consistency.h"
 
 #include "mongo/db/catalog/index_catalog.h"
+#include "mongo/db/catalog/validate_gen.h"
 #include "mongo/db/db_raii.h"
 #include "mongo/db/index/index_descriptor.h"
 #include "mongo/db/storage/storage_debug_util.h"
@@ -77,7 +78,7 @@ IndexInfo::IndexInfo(const IndexDescriptor* descriptor)
 IndexConsistency::IndexConsistency(OperationContext* opCtx,
                                    CollectionValidation::ValidateState* validateState)
     : _validateState(validateState), _firstPhase(true) {
-    _indexKeyCount.resize(kNumHashBuckets);
+    _indexKeyBuckets.resize(kNumHashBuckets);
 
     for (const auto& index : _validateState->getIndexes()) {
         const IndexDescriptor* descriptor = index->descriptor();
@@ -99,8 +100,9 @@ size_t IndexConsistency::getMultikeyMetadataPathCount(IndexInfo* indexInfo) {
 }
 
 bool IndexConsistency::haveEntryMismatch() const {
-    return std::any_of(
-        _indexKeyCount.begin(), _indexKeyCount.end(), [](int count) -> bool { return count; });
+    return std::any_of(_indexKeyBuckets.begin(),
+                       _indexKeyBuckets.end(),
+                       [](const IndexKeyBucket& bucket) -> bool { return bucket.indexKeyCount; });
 }
 
 void IndexConsistency::setSecondPhase() {
@@ -206,7 +208,8 @@ void IndexConsistency::addDocKey(OperationContext* opCtx,
     if (_firstPhase) {
         // During the first phase of validation we only keep track of the count for the document
         // keys encountered.
-        _indexKeyCount[hash]++;
+        _indexKeyBuckets[hash].indexKeyCount++;
+        _indexKeyBuckets[hash].bucketSizeBytes += ks.getSize();
         indexInfo->numRecords++;
 
         if (MONGO_unlikely(_validateState->extraLoggingForTest())) {
@@ -217,7 +220,7 @@ void IndexConsistency::addDocKey(OperationContext* opCtx,
             StorageDebugUtil::printKeyString(
                 recordId, ks, keyPatternBson, keyStringBson, "[validate](record)");
         }
-    } else if (_indexKeyCount[hash]) {
+    } else if (_indexKeyBuckets[hash].indexKeyCount) {
         // Found a document key for a hash bucket that had mismatches.
 
         // Get the documents _id index key.
@@ -249,7 +252,8 @@ void IndexConsistency::addIndexKey(const KeyString::Value& ks,
     if (_firstPhase) {
         // During the first phase of validation we only keep track of the count for the index entry
         // keys encountered.
-        _indexKeyCount[hash]--;
+        _indexKeyBuckets[hash].indexKeyCount--;
+        _indexKeyBuckets[hash].bucketSizeBytes += ks.getSize();
         indexInfo->numKeys++;
 
         if (MONGO_unlikely(_validateState->extraLoggingForTest())) {
@@ -260,7 +264,7 @@ void IndexConsistency::addIndexKey(const KeyString::Value& ks,
             StorageDebugUtil::printKeyString(
                 recordId, ks, keyPatternBson, keyStringBson, "[validate](index)");
         }
-    } else if (_indexKeyCount[hash]) {
+    } else if (_indexKeyBuckets[hash].indexKeyCount) {
         // Found an index key for a bucket that has inconsistencies.
         // If there is a corresponding document key for the index entry key, we remove the key from
         // the '_missingIndexEntries' map. However if there was no document key for the index entry
@@ -286,6 +290,69 @@ void IndexConsistency::addIndexKey(const KeyString::Value& ks,
     }
 }
 
+bool IndexConsistency::limitMemoryUsageForSecondPhase(ValidateResults* result) {
+    invariant(!_firstPhase);
+
+    const uint32_t maxMemoryUsageBytes = maxValidateMemoryUsageMB.load() * 1024 * 1024;
+    const uint64_t totalMemoryNeededBytes =
+        std::accumulate(_indexKeyBuckets.begin(),
+                        _indexKeyBuckets.end(),
+                        0,
+                        [](uint64_t bytes, const IndexKeyBucket& bucket) {
+                            return bucket.indexKeyCount ? bytes + bucket.bucketSizeBytes : bytes;
+                        });
+
+    if (totalMemoryNeededBytes <= maxMemoryUsageBytes) {
+        // The amount of memory we need is under the limit, so no need to do anything else.
+        return true;
+    }
+
+    bool hasNonZeroBucket = false;
+    uint64_t memoryUsedSoFarBytes = 0;
+    uint32_t smallestBucketBytes = std::numeric_limits<uint32_t>::max();
+    // Zero out any nonzero buckets that would put us over maxMemoryUsageBytes.
+    std::for_each(_indexKeyBuckets.begin(), _indexKeyBuckets.end(), [&](IndexKeyBucket& bucket) {
+        if (bucket.indexKeyCount == 0) {
+            return;
+        }
+
+        smallestBucketBytes = std::min(smallestBucketBytes, bucket.bucketSizeBytes);
+        if (bucket.bucketSizeBytes + memoryUsedSoFarBytes > maxMemoryUsageBytes) {
+            // Including this bucket would put us over the memory limit, so zero
+            // this bucket.
+            bucket.indexKeyCount = 0;
+            return;
+        }
+        memoryUsedSoFarBytes += bucket.bucketSizeBytes;
+        hasNonZeroBucket = true;
+    });
+
+    StringBuilder memoryLimitMessage;
+    memoryLimitMessage << "Memory limit for validation is currently set to "
+                       << maxValidateMemoryUsageMB.load()
+                       << "MB and can be configured via the 'maxValidateMemoryUsageMB' parameter.";
+
+    if (!hasNonZeroBucket) {
+        const uint32_t minMemoryNeededMB = (smallestBucketBytes / (1024 * 1024)) + 1;
+        StringBuilder ss;
+        ss << "Unable to report index entry inconsistencies due to memory limitations. Need at "
+              "least "
+           << minMemoryNeededMB << "MB to report at least one index entry inconsistency. "
+           << memoryLimitMessage.str();
+        result->errors.push_back(ss.str());
+        result->valid = false;
+
+        return false;
+    }
+
+    StringBuilder ss;
+    ss << "Not all index entry inconsistencies are reported due to memory limitations. "
+       << memoryLimitMessage.str();
+    result->errors.push_back(ss.str());
+
+    return true;
+}
+
 BSONObj IndexConsistency::_generateInfo(const IndexInfo& indexInfo,
                                         RecordId recordId,
                                         const BSONObj& indexKey,
diff --git a/src/mongo/db/catalog/index_consistency.h b/src/mongo/db/catalog/index_consistency.h
index cb0fdf38958..95a5e0590f6 100644
--- a/src/mongo/db/catalog/index_consistency.h
+++ b/src/mongo/db/catalog/index_consistency.h
@@ -131,7 +131,19 @@ public:
      */
     void addIndexEntryErrors(ValidateResultsMap* indexNsResultsMap, ValidateResults* results);
 
+    /**
+     * Sets up this IndexConsistency object to limit memory usage in the second phase of index
+     * validation. Returns whether the memory limit is sufficient to report at least on index entry
+     * inconsistency and continue with the second phase of validation.
+     */
+    bool limitMemoryUsageForSecondPhase(ValidateResults* result);
+
 private:
+    struct IndexKeyBucket {
+        uint32_t indexKeyCount;
+        uint32_t bucketSizeBytes;
+    };
+
     IndexConsistency() = delete;
 
     CollectionValidation::ValidateState* _validateState;
@@ -147,7 +159,7 @@ private:
     //       than zero, there are too few index entries.
     //     - Similarly, if that count ends up less than zero, there are too many index entries.
 
-    std::vector<uint32_t> _indexKeyCount;
+    std::vector<IndexKeyBucket> _indexKeyBuckets;
 
     // A vector of IndexInfo indexes by index number
     IndexInfoMap _indexesInfo;
diff --git a/src/mongo/db/catalog/throttle_cursor.cpp b/src/mongo/db/catalog/throttle_cursor.cpp
index f94306c11b9..265de074041 100644
--- a/src/mongo/db/catalog/throttle_cursor.cpp
+++ b/src/mongo/db/catalog/throttle_cursor.cpp
@@ -31,7 +31,7 @@
 
 #include "mongo/db/catalog/throttle_cursor.h"
 
-#include "mongo/db/catalog/max_validate_mb_per_sec_gen.h"
+#include "mongo/db/catalog/validate_gen.h"
 #include "mongo/db/curop.h"
 #include "mongo/db/index/index_access_method.h"
 #include "mongo/db/operation_context.h"
diff --git a/src/mongo/db/catalog/throttle_cursor_test.cpp b/src/mongo/db/catalog/throttle_cursor_test.cpp
index 52b5fc853a9..e36f7244a1d 100644
--- a/src/mongo/db/catalog/throttle_cursor_test.cpp
+++ b/src/mongo/db/catalog/throttle_cursor_test.cpp
@@ -35,7 +35,7 @@
 #include "mongo/db/catalog/collection.h"
 #include "mongo/db/catalog/index_catalog.h"
 #include "mongo/db/catalog/index_catalog_entry.h"
-#include "mongo/db/catalog/max_validate_mb_per_sec_gen.h"
+#include "mongo/db/catalog/validate_gen.h"
 #include "mongo/db/db_raii.h"
 #include "mongo/unittest/unittest.h"
 #include "mongo/util/clock_source_mock.h"
diff --git a/src/mongo/db/catalog/max_validate_mb_per_sec.idl b/src/mongo/db/catalog/validate.idl
index 2bfb89a5dbe..c90a9dfbf72 100644
--- a/src/mongo/db/catalog/max_validate_mb_per_sec.idl
+++ b/src/mongo/db/catalog/validate.idl
@@ -39,3 +39,11 @@ server_parameters:
         cpp_vartype: AtomicWord<int>
         validator: { gte: 0 }
         default: 0
+
+    maxValidateMemoryUsageMB:
+        description: "Limits the amount of memory that a single validate command will use."
+        set_at: [ startup, runtime ]
+        cpp_varname: maxValidateMemoryUsageMB
+        cpp_vartype: AtomicWord<int>
+        validator: { gt: 0 }
+        default: 200
diff --git a/src/mongo/shell/assert.js b/src/mongo/shell/assert.js
index f4a87404d9b..5baf8d3cec1 100644
--- a/src/mongo/shell/assert.js
+++ b/src/mongo/shell/assert.js
@@ -306,6 +306,32 @@ assert = (function() {
         }
     };
 
+    assert.containsPrefix = function(prefix, arr, msg) {
+        var wasIn = false;
+        if (typeof (prefix) !== "string") {
+            throw new Error("The first argument to containsPrefix must be a string.");
+        }
+        if (!Array.isArray(arr)) {
+            throw new Error("The second argument to containsPrefix must be an array.");
+        }
+
+        for (let i = 0; i < arr.length; i++) {
+            if (typeof (arr[i]) !== "string") {
+                continue;
+            }
+
+            wasIn = arr[i].startsWith(prefix);
+            if (wasIn) {
+                break;
+            }
+        }
+
+        if (!wasIn) {
+            doassert(_buildAssertionMessage(
+                msg, tojson(prefix) + " was not a prefix in " + tojson(arr)));
+        }
+    };
+
     /*
      * Calls a function 'func' at repeated intervals until either func() returns true
      * or more than 'timeout' milliseconds have elapsed. Throws an exception with