SERVER-44121 MultiIndexBlock::drainBackgroundWrites() accepts drain yield policy

7 files changed, 36 insertions, 14 deletions

diff --git a/src/mongo/db/catalog/index_builds_manager.cpp b/src/mongo/db/catalog/index_builds_manager.cpp
index 82172e4b0cf..9dff10efdef 100644
--- a/src/mongo/db/catalog/index_builds_manager.cpp
+++ b/src/mongo/db/catalog/index_builds_manager.cpp
@@ -206,7 +206,7 @@ Status IndexBuildsManager::drainBackgroundWrites(
     IndexBuildInterceptor::DrainYieldPolicy drainYieldPolicy) {
     auto builder = _getBuilder(buildUUID);
 
-    return builder->drainBackgroundWrites(opCtx, readSource);
+    return builder->drainBackgroundWrites(opCtx, readSource, drainYieldPolicy);
 }
 
 Status IndexBuildsManager::finishBuildingPhase(const UUID& buildUUID) {
diff --git a/src/mongo/db/catalog/multi_index_block.cpp b/src/mongo/db/catalog/multi_index_block.cpp
index 20a97d48938..d021b9bad74 100644
--- a/src/mongo/db/catalog/multi_index_block.cpp
+++ b/src/mongo/db/catalog/multi_index_block.cpp
@@ -42,7 +42,6 @@
 #include "mongo/db/catalog/multi_index_block_gen.h"
 #include "mongo/db/client.h"
 #include "mongo/db/concurrency/write_conflict_exception.h"
-#include "mongo/db/index/index_build_interceptor.h"
 #include "mongo/db/index/multikey_paths.h"
 #include "mongo/db/multi_key_path_tracker.h"
 #include "mongo/db/op_observer.h"
@@ -676,8 +675,10 @@ Status MultiIndexBlock::dumpInsertsFromBulk(OperationContext* opCtx,
     return Status::OK();
 }
 
-Status MultiIndexBlock::drainBackgroundWrites(OperationContext* opCtx,
-                                              RecoveryUnit::ReadSource readSource) {
+Status MultiIndexBlock::drainBackgroundWrites(
+    OperationContext* opCtx,
+    RecoveryUnit::ReadSource readSource,
+    IndexBuildInterceptor::DrainYieldPolicy drainYieldPolicy) {
     if (State::kAborted == _getState()) {
         return {ErrorCodes::IndexBuildAborted,
                 str::stream() << "Index build aborted: " << _abortReason
diff --git a/src/mongo/db/catalog/multi_index_block.h b/src/mongo/db/catalog/multi_index_block.h
index 28fc8b8cabc..c5b0b1b1589 100644
--- a/src/mongo/db/catalog/multi_index_block.h
+++ b/src/mongo/db/catalog/multi_index_block.h
@@ -45,6 +45,7 @@
 #include "mongo/db/catalog/index_build_block.h"
 #include "mongo/db/catalog/index_catalog.h"
 #include "mongo/db/index/index_access_method.h"
+#include "mongo/db/index/index_build_interceptor.h"
 #include "mongo/db/record_id.h"
 #include "mongo/platform/mutex.h"
 #include "mongo/util/fail_point.h"
@@ -195,9 +196,9 @@ public:
      *
      * Must not be in a WriteUnitOfWork.
      */
-    Status drainBackgroundWrites(
-        OperationContext* opCtx,
-        RecoveryUnit::ReadSource readSource = RecoveryUnit::ReadSource::kUnset);
+    Status drainBackgroundWrites(OperationContext* opCtx,
+                                 RecoveryUnit::ReadSource readSource,
+                                 IndexBuildInterceptor::DrainYieldPolicy drainYieldPolicy);
 
     /**
      * Check any constraits that may have been temporarily violated during the index build for
diff --git a/src/mongo/db/commands/create_indexes.cpp b/src/mongo/db/commands/create_indexes.cpp
index 5d827ea27b0..93de2e317de 100644
--- a/src/mongo/db/commands/create_indexes.cpp
+++ b/src/mongo/db/commands/create_indexes.cpp
@@ -563,7 +563,10 @@ bool runCreateIndexesForMobile(OperationContext* opCtx,
         // the collection lock.
         ns = collection->ns();
 
-        uassertStatusOK(indexer.drainBackgroundWrites(opCtx));
+        uassertStatusOK(
+            indexer.drainBackgroundWrites(opCtx,
+                                          RecoveryUnit::ReadSource::kUnset,
+                                          IndexBuildInterceptor::DrainYieldPolicy::kYield));
     }
 
     if (MONGO_unlikely(hangAfterIndexBuildFirstDrain.shouldFail())) {
@@ -584,7 +587,10 @@ bool runCreateIndexesForMobile(OperationContext* opCtx,
         // the collection lock.
         ns = collection->ns();
 
-        uassertStatusOK(indexer.drainBackgroundWrites(opCtx));
+        uassertStatusOK(
+            indexer.drainBackgroundWrites(opCtx,
+                                          RecoveryUnit::ReadSource::kUnset,
+                                          IndexBuildInterceptor::DrainYieldPolicy::kNoYield));
     }
 
     if (MONGO_unlikely(hangAfterIndexBuildSecondDrain.shouldFail())) {
@@ -612,7 +618,10 @@ bool runCreateIndexesForMobile(OperationContext* opCtx,
     invariant(CollectionCatalog::get(opCtx).lookupCollectionByNamespace(ns));
 
     // Perform the third and final drain while holding the exclusive collection lock.
-    uassertStatusOK(indexer.drainBackgroundWrites(opCtx));
+    uassertStatusOK(
+        indexer.drainBackgroundWrites(opCtx,
+                                      RecoveryUnit::ReadSource::kUnset,
+                                      IndexBuildInterceptor::DrainYieldPolicy::kNoYield));
 
     // This is required before completion.
     uassertStatusOK(indexer.checkConstraints(opCtx));
diff --git a/src/mongo/db/repl/collection_bulk_loader_impl.cpp b/src/mongo/db/repl/collection_bulk_loader_impl.cpp
index 5df1b865787..944776aee8f 100644
--- a/src/mongo/db/repl/collection_bulk_loader_impl.cpp
+++ b/src/mongo/db/repl/collection_bulk_loader_impl.cpp
@@ -274,7 +274,11 @@ Status CollectionBulkLoaderImpl::commit() {
                     });
             }
 
-            status = _idIndexBlock->drainBackgroundWrites(_opCtx.get());
+            status = _idIndexBlock->drainBackgroundWrites(
+                _opCtx.get(),
+                RecoveryUnit::ReadSource::kUnset,
+                _nss.isSystemDotViews() ? IndexBuildInterceptor::DrainYieldPolicy::kNoYield
+                                        : IndexBuildInterceptor::DrainYieldPolicy::kYield);
             if (!status.isOK()) {
                 return status;
             }
diff --git a/src/mongo/dbtests/querytests.cpp b/src/mongo/dbtests/querytests.cpp
index 97e8ee011db..269d0bad69e 100644
--- a/src/mongo/dbtests/querytests.cpp
+++ b/src/mongo/dbtests/querytests.cpp
@@ -113,7 +113,10 @@ protected:
             wunit.commit();
         }
         uassertStatusOK(indexer.insertAllDocumentsInCollection(&_opCtx, _collection));
-        uassertStatusOK(indexer.drainBackgroundWrites(&_opCtx));
+        uassertStatusOK(
+            indexer.drainBackgroundWrites(&_opCtx,
+                                          RecoveryUnit::ReadSource::kUnset,
+                                          IndexBuildInterceptor::DrainYieldPolicy::kNoYield));
         uassertStatusOK(indexer.checkConstraints(&_opCtx));
         {
             WriteUnitOfWork wunit(&_opCtx);
diff --git a/src/mongo/dbtests/storage_timestamp_tests.cpp b/src/mongo/dbtests/storage_timestamp_tests.cpp
index 4304bfadbc4..36f8bb83236 100644
--- a/src/mongo/dbtests/storage_timestamp_tests.cpp
+++ b/src/mongo/dbtests/storage_timestamp_tests.cpp
@@ -2035,7 +2035,9 @@ public:
                       .getTimestamp(),
                   firstInsert.asTimestamp());
 
-        ASSERT_OK(indexer.drainBackgroundWrites(_opCtx));
+        ASSERT_OK(indexer.drainBackgroundWrites(_opCtx,
+                                                RecoveryUnit::ReadSource::kUnset,
+                                                IndexBuildInterceptor::DrainYieldPolicy::kNoYield));
 
         auto indexCatalog = autoColl.getCollection()->getIndexCatalog();
         const IndexCatalogEntry* buildingIndex = indexCatalog->getEntry(
@@ -2076,7 +2078,9 @@ public:
         const LogicalTime afterSecondInsert = _clock->reserveTicks(1);
         setReplCoordAppliedOpTime(repl::OpTime(afterSecondInsert.asTimestamp(), presentTerm));
 
-        ASSERT_OK(indexer.drainBackgroundWrites(_opCtx));
+        ASSERT_OK(indexer.drainBackgroundWrites(_opCtx,
+                                                RecoveryUnit::ReadSource::kUnset,
+                                                IndexBuildInterceptor::DrainYieldPolicy::kNoYield));
 
         {
             // At time of the second insert, there are un-drained writes.