SERVER-30849 _handleARMResultsCallback needs to hold mutex while accessing member variables

2 files changed, 14 insertions, 7 deletions

diff --git a/src/mongo/db/repl/collection_cloner.cpp b/src/mongo/db/repl/collection_cloner.cpp
index 98d658b49f9..ad2987e7603 100644
--- a/src/mongo/db/repl/collection_cloner.cpp
+++ b/src/mongo/db/repl/collection_cloner.cpp
@@ -612,7 +612,7 @@ StatusWith<std::vector<BSONElement>> CollectionCloner::_parseParallelCollectionS
     }
 }
 
-Status CollectionCloner::_bufferNextBatchFromArm() {
+Status CollectionCloner::_bufferNextBatchFromArm(WithLock lock) {
     while (_arm->ready()) {
         auto armResultStatus = _arm->nextReady();
         if (!armResultStatus.getStatus().isOK()) {
@@ -671,13 +671,20 @@ void CollectionCloner::_handleARMResultsCallback(
     }
 
     // Pull the documents from the ARM into a buffer until the entire batch has been processed.
-    auto nextBatchStatus = _bufferNextBatchFromArm();
-    if (!nextBatchStatus.isOK()) {
-        setResultAndCancelRemainingWork(onCompletionGuard, nextBatchStatus);
-        return;
+    bool lastBatch;
+    {
+        UniqueLock lk(_mutex);
+        auto nextBatchStatus = _bufferNextBatchFromArm(lk);
+        if (!nextBatchStatus.isOK()) {
+            onCompletionGuard->setResultAndCancelRemainingWork_inlock(lk, nextBatchStatus);
+            return;
+        }
+
+        // Check if this is the last batch of documents to clone.
+        lastBatch = _arm->remotesExhausted();
     }
 
-    bool lastBatch = _arm->remotesExhausted();
+    // Schedule the next document batch insertion.
     auto&& scheduleResult =
         _scheduleDbWorkFn(stdx::bind(&CollectionCloner::_insertDocumentsCallback,
                                      this,
diff --git a/src/mongo/db/repl/collection_cloner.h b/src/mongo/db/repl/collection_cloner.h
index 222886b9d72..1f8701c3ca5 100644
--- a/src/mongo/db/repl/collection_cloner.h
+++ b/src/mongo/db/repl/collection_cloner.h
@@ -237,7 +237,7 @@ private:
     /**
      * Pull all ready results from the ARM into a buffer to be inserted.
      */
-    Status _bufferNextBatchFromArm();
+    Status _bufferNextBatchFromArm(WithLock lock);
 
     /**
      * Called whenever there is a new batch of documents ready from the 'AsyncResultsMerger'.