SERVER-61757 Add configurable dbCheck batch size

(cherry picked from commit c67b5eb3b1c009f8ed794daae6b01539ca7a23b2)

4 files changed, 104 insertions, 10 deletions

diff --git a/etc/backports_required_for_multiversion_tests.yml b/etc/backports_required_for_multiversion_tests.yml
index a3a01005a02..4249dafd4f1 100644
--- a/etc/backports_required_for_multiversion_tests.yml
+++ b/etc/backports_required_for_multiversion_tests.yml
@@ -187,6 +187,8 @@ all:
     test_file: jstests/sharding/cancel_coordinate_txn_commit_with_tickets_exhausted.js
   - ticket: SERVER-61738
     test_file: jstests/replsets/dbcheck.js
+  - ticket: SERVER-61955
+    test_file: jstests/auth/dbcheck.js
 
 
 suites:
diff --git a/jstests/replsets/dbcheck.js b/jstests/replsets/dbcheck.js
index 54844fb0d59..af3be232c34 100644
--- a/jstests/replsets/dbcheck.js
+++ b/jstests/replsets/dbcheck.js
@@ -46,8 +46,10 @@ function addEnoughForMultipleBatches(collection) {
 
 // Name for a collection which takes multiple batches to check and which shouldn't be modified
 // by any of the tests.
-let multiBatchSimpleCollName = "dbcheck-simple-collection";
-addEnoughForMultipleBatches(replSet.getPrimary().getDB(dbName)[multiBatchSimpleCollName]);
+const multiBatchSimpleCollName = "dbcheck-simple-collection";
+const multiBatchSimpleCollSize = 10000;
+replSet.getPrimary().getDB(dbName)[multiBatchSimpleCollName].insertMany(
+    [...Array(10000).keys()].map(x => ({_id: x})));
 
 function dbCheckCompleted(db) {
     return db.currentOp().inprog.filter(x => x["desc"] == "dbCheck")[0] === undefined;
@@ -293,6 +295,53 @@ function testDbCheckParameters() {
         {dbCheck: multiBatchSimpleCollName, minKey: start, maxKey: end, maxSize: maxSize}));
     awaitDbCheckCompletion(db, multiBatchSimpleCollName, end, maxSize);
     checkEntryBounds(start, start + maxCount);
+
+    const healthlog = db.getSiblingDB('local').system.healthlog;
+    {
+        // Validate custom maxDocsPerBatch
+        clearLog();
+        const maxDocsPerBatch = 100;
+        assert.commandWorked(
+            db.runCommand({dbCheck: multiBatchSimpleCollName, maxDocsPerBatch: maxDocsPerBatch}));
+
+        const healthlog = db.getSiblingDB('local').system.healthlog;
+        assert.soon(function() {
+            const expectedBatches = multiBatchSimpleCollSize / maxDocsPerBatch +
+                (multiBatchSimpleCollSize % maxDocsPerBatch ? 1 : 0);
+            return (healthlog.find({"operation": "dbCheckBatch"}).itcount() == expectedBatches);
+        }, "dbCheck doesn't seem to complete", 60 * 1000);
+
+        assert.eq(
+            db.getSiblingDB('local')
+                .system.healthlog.find({"operation": "dbCheckBatch", "data.count": maxDocsPerBatch})
+                .itcount(),
+            multiBatchSimpleCollSize / maxDocsPerBatch);
+    }
+    {
+        // Validate custom maxBytesPerBatch
+        clearLog();
+        const coll = db.getSiblingDB("maxBytesPerBatch").maxBytesPerBatch;
+
+        // Insert nDocs, each of which being slightly larger than 1MB, and then run dbCheck with
+        // maxBytesPerBatch := 1MB
+        const nDocs = 5;
+        coll.insertMany([...Array(nDocs).keys()].map(x => ({a: 'a'.repeat(1024 * 1024)})));
+        const maxBytesPerBatch = 1024 * 1024;
+        assert.commandWorked(db.getSiblingDB("maxBytesPerBatch").runCommand({
+            dbCheck: coll.getName(),
+            maxBytesPerBatch: maxBytesPerBatch
+        }));
+
+        // Confirm dbCheck logs nDocs batches.
+        assert.soon(function() {
+            return (healthlog.find({"operation": "dbCheckBatch"}).itcount() == nDocs);
+        }, "dbCheck doesn't seem to complete", 60 * 1000);
+
+        assert.eq(db.getSiblingDB('local')
+                      .system.healthlog.find({"operation": "dbCheckBatch", "data.count": 1})
+                      .itcount(),
+                  nDocs);
+    }
 }
 
 testDbCheckParameters();
diff --git a/src/mongo/db/commands/dbcheck.cpp b/src/mongo/db/commands/dbcheck.cpp
index 588c852ab36..627d804f42e 100644
--- a/src/mongo/db/commands/dbcheck.cpp
+++ b/src/mongo/db/commands/dbcheck.cpp
@@ -53,9 +53,6 @@
 namespace mongo {
 
 namespace {
-constexpr uint64_t kBatchDocs = 5'000;
-constexpr uint64_t kBatchBytes = 20'000'000;
-
 
 /**
  * All the information needed to run dbCheck on a single collection.
@@ -67,6 +64,8 @@ struct DbCheckCollectionInfo {
     int64_t maxCount;
     int64_t maxSize;
     int64_t maxRate;
+    int64_t maxDocsPerBatch;
+    int64_t maxBytesPerBatch;
     int64_t maxBatchTimeMillis;
 };
 
@@ -94,9 +93,18 @@ std::unique_ptr<DbCheckRun> singleCollectionRun(OperationContext* opCtx,
     const auto maxCount = invocation.getMaxCount();
     const auto maxSize = invocation.getMaxSize();
     const auto maxRate = invocation.getMaxCountPerSecond();
+    const auto maxDocsPerBatch = invocation.getMaxDocsPerBatch();
+    const auto maxBytesPerBatch = invocation.getMaxBytesPerBatch();
     const auto maxBatchTimeMillis = invocation.getMaxBatchTimeMillis();
-    const auto info =
-        DbCheckCollectionInfo{nss, start, end, maxCount, maxSize, maxRate, maxBatchTimeMillis};
+    const auto info = DbCheckCollectionInfo{nss,
+                                            start,
+                                            end,
+                                            maxCount,
+                                            maxSize,
+                                            maxRate,
+                                            maxDocsPerBatch,
+                                            maxBytesPerBatch,
+                                            maxBatchTimeMillis};
     auto result = std::make_unique<DbCheckRun>();
     result->push_back(info);
     return result;
@@ -113,14 +121,23 @@ std::unique_ptr<DbCheckRun> fullDatabaseRun(OperationContext* opCtx,
 
     const int64_t max = std::numeric_limits<int64_t>::max();
     const auto rate = invocation.getMaxCountPerSecond();
+    const auto maxDocsPerBatch = invocation.getMaxDocsPerBatch();
+    const auto maxBytesPerBatch = invocation.getMaxBytesPerBatch();
     const auto maxBatchTimeMillis = invocation.getMaxBatchTimeMillis();
     auto result = std::make_unique<DbCheckRun>();
     auto perCollectionWork = [&](const Collection* coll) {
         if (!coll->ns().isReplicated()) {
             return true;
         }
-        DbCheckCollectionInfo info{
-            coll->ns(), BSONKey::min(), BSONKey::max(), max, max, rate, maxBatchTimeMillis};
+        DbCheckCollectionInfo info{coll->ns(),
+                                   BSONKey::min(),
+                                   BSONKey::max(),
+                                   max,
+                                   max,
+                                   rate,
+                                   maxDocsPerBatch,
+                                   maxBytesPerBatch,
+                                   maxBatchTimeMillis};
         result->push_back(info);
         return true;
     };
@@ -229,7 +246,7 @@ private:
                 docsInCurrentInterval = 0;
             }
 
-            auto result = _runBatch(info, start, kBatchDocs, kBatchBytes);
+            auto result = _runBatch(info, start, info.maxDocsPerBatch, info.maxBytesPerBatch);
 
             if (_done) {
                 return;
@@ -509,6 +526,8 @@ public:
                "              maxCount: <max number of docs>,\n"
                "              maxSize: <max size of docs>,\n"
                "              maxCountPerSecond: <max rate in docs/sec> } "
+               "              maxDocsPerBatch: <max number of docs/batch> } "
+               "              maxBytesPerBatch: <try to keep a batch within max bytes/batch> } "
                "              maxBatchTimeMillis: <max time processing a batch in milliseconds> } "
                "to check a collection.\n"
                "Invoke with {dbCheck: 1} to check all collections in the database.";
diff --git a/src/mongo/db/repl/dbcheck.idl b/src/mongo/db/repl/dbcheck.idl
index a262aa862d9..9a16e3456e0 100644
--- a/src/mongo/db/repl/dbcheck.idl
+++ b/src/mongo/db/repl/dbcheck.idl
@@ -112,6 +112,18 @@ structs:
       maxCountPerSecond:
         type: safeInt64
         default: "std::numeric_limits<int64_t>::max()"
+      maxDocsPerBatch:
+        type: safeInt64
+        default: 5000
+        validator:
+          gte: 1
+          lte: 10000
+      maxBytesPerBatch:
+        type: safeInt64
+        default: 20 * 1024 * 1024
+        validator:
+          gte: 102400
+          lte: 268435456
       maxBatchTimeMillis:
         type: safeInt64
         default: 1000
@@ -128,6 +140,18 @@ structs:
       maxCountPerSecond:
         type: safeInt64
         default: "std::numeric_limits<int64_t>::max()"
+      maxDocsPerBatch:
+        type: safeInt64
+        default: 5000
+        validator:
+          gte: 1
+          lte: 10000
+      maxBytesPerBatch:
+        type: safeInt64
+        default: 20 * 1024 * 1024
+        validator:
+          gte: 102400
+          lte: 268435456
       maxBatchTimeMillis:
         type: safeInt64
         default: 1000