SERVER-44583 M/R Agg: Add explain support for mapReduce command

1 files changed, 45 insertions, 43 deletions

diff --git a/src/mongo/db/pipeline/document_source_writer.h b/src/mongo/db/pipeline/document_source_writer.h
index dd5dd0c37a5..5110582faeb 100644
--- a/src/mongo/db/pipeline/document_source_writer.h
+++ b/src/mongo/db/pipeline/document_source_writer.h
@@ -170,54 +170,56 @@ DocumentSource::GetNextResult DocumentSourceWriter<B>::doGetNext() {
         return GetNextResult::makeEOF();
     }
 
-    if (!_initialized) {
-        // Explain should never try to actually execute any writes. We only ever expect
-        // getNext() to be called for the 'executionStats' and 'allPlansExecution' explain
-        // modes. This assertion should not be triggered for 'queryPlanner' explain, which
-        // is perfectly legal.
-        uassert(51029,
-                "explain of {} is not allowed with verbosity {}"_format(
-                    getSourceName(), ExplainOptions::verbosityString(*pExpCtx->explain)),
-                !pExpCtx->explain);
-        initialize();
-        _initialized = true;
-    }
-
-    BatchedObjects batch;
-    int bufferedBytes = 0;
-
-    auto nextInput = pSource->getNext();
-    for (; nextInput.isAdvanced(); nextInput = pSource->getNext()) {
-        waitWhileFailPointEnabled();
-
-        auto doc = nextInput.releaseDocument();
-        auto [obj, objSize] = makeBatchObject(std::move(doc));
+    // Ignore writes and exhaust input if we are in explain mode.
+    if (pExpCtx->explain) {
+        auto nextInput = pSource->getNext();
+        for (; nextInput.isAdvanced(); nextInput = pSource->getNext()) {
+        }
+        _done = nextInput.getStatus() == GetNextResult::ReturnStatus::kEOF;
+        return nextInput;
+    } else {
+        if (!_initialized) {
+            initialize();
+            _initialized = true;
+        }
 
-        bufferedBytes += objSize;
-        if (!batch.empty() &&
-            (bufferedBytes > BSONObjMaxUserSize || batch.size() >= write_ops::kMaxWriteBatchSize)) {
+        BatchedObjects batch;
+        int bufferedBytes = 0;
+
+        auto nextInput = pSource->getNext();
+        for (; nextInput.isAdvanced(); nextInput = pSource->getNext()) {
+            waitWhileFailPointEnabled();
+
+            auto doc = nextInput.releaseDocument();
+            auto [obj, objSize] = makeBatchObject(std::move(doc));
+
+            bufferedBytes += objSize;
+            if (!batch.empty() &&
+                (bufferedBytes > BSONObjMaxUserSize ||
+                 batch.size() >= write_ops::kMaxWriteBatchSize)) {
+                spill(std::move(batch));
+                batch.clear();
+                bufferedBytes = objSize;
+            }
+            batch.push_back(obj);
+        }
+        if (!batch.empty()) {
             spill(std::move(batch));
             batch.clear();
-            bufferedBytes = objSize;
         }
-        batch.push_back(obj);
-    }
-    if (!batch.empty()) {
-        spill(std::move(batch));
-        batch.clear();
-    }
 
-    switch (nextInput.getStatus()) {
-        case GetNextResult::ReturnStatus::kAdvanced: {
-            MONGO_UNREACHABLE;  // We consumed all advances above.
-        }
-        case GetNextResult::ReturnStatus::kPauseExecution: {
-            return nextInput;  // Propagate the pause.
-        }
-        case GetNextResult::ReturnStatus::kEOF: {
-            _done = true;
-            finalize();
-            return nextInput;
+        switch (nextInput.getStatus()) {
+            case GetNextResult::ReturnStatus::kAdvanced: {
+                MONGO_UNREACHABLE;  // We consumed all advances above.
+            }
+            case GetNextResult::ReturnStatus::kPauseExecution: {
+                return nextInput;  // Propagate the pause.
+            }
+            case GetNextResult::ReturnStatus::kEOF: {
+                _done = true;
+                finalize();
+                return nextInput;
+            }
         }
     }
     MONGO_UNREACHABLE;