SERVER-23097 Improve killed executor handling in MapReduce

(cherry picked from commit 200b4f971b021f792194489c8ffbc95b9f9cba35)

2 files changed, 24 insertions, 23 deletions

diff --git a/src/mongo/db/commands/mr.cpp b/src/mongo/db/commands/mr.cpp
index 4fd713e359a..d04e91b5db1 100644
--- a/src/mongo/db/commands/mr.cpp
+++ b/src/mongo/db/commands/mr.cpp
@@ -1392,18 +1392,22 @@ public:
                 }
                 std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
 
-                Database* db = scopedAutoDb->getDb();
-                Collection* coll = state.getCollectionOrUassert(db, config.ns);
-                invariant(coll);
-
-                auto statusWithPlanExecutor =
-                    getExecutor(txn, coll, std::move(cq), PlanExecutor::YIELD_AUTO);
-                if (!statusWithPlanExecutor.isOK()) {
-                    uasserted(17239, "Can't get executor for query " + config.filter.toString());
-                    return 0;
-                }
+                unique_ptr<PlanExecutor> exec;
+                {
+                    Database* db = scopedAutoDb->getDb();
+                    Collection* coll = State::getCollectionOrUassert(db, config.ns);
+                    invariant(coll);
+
+                    auto statusWithPlanExecutor =
+                        getExecutor(txn, coll, std::move(cq), PlanExecutor::YIELD_AUTO);
+                    if (!statusWithPlanExecutor.isOK()) {
+                        uasserted(17239,
+                                  "Can't get executor for query " + config.filter.toString());
+                        return 0;
+                    }
 
-                unique_ptr<PlanExecutor> exec = std::move(statusWithPlanExecutor.getValue());
+                    exec = std::move(statusWithPlanExecutor.getValue());
+                }
 
                 Timer mt;
 
@@ -1448,17 +1452,12 @@ public:
                         scopedXact.reset(new ScopedTransaction(txn, MODE_IS));
                         scopedAutoDb.reset(new AutoGetDb(txn, nss.db(), MODE_S));
 
-                        exec->restoreState();
-
-                        // Need to reload the database, in case it was dropped after we
-                        // released the lock
-                        db = scopedAutoDb->getDb();
-                        if (db == NULL) {
-                            // Database was deleted after we freed the lock
-                            StringBuilder sb;
-                            sb << "Database " << nss.db()
-                               << " was deleted in the middle of the reduce job.";
-                            uasserted(28523, sb.str());
+                        if (!exec->restoreState()) {
+                            return appendCommandStatus(
+                                result,
+                                Status(ErrorCodes::OperationFailed,
+                                       str::stream()
+                                           << "Executor killed during mapReduce command"));
                         }
 
                         reduceTime += t.micros();
@@ -1483,6 +1482,8 @@ public:
                 // Record the indexes used by the PlanExecutor.
                 PlanSummaryStats stats;
                 Explain::getSummaryStats(*exec, &stats);
+                Collection* coll = scopedAutoDb->getDb()->getCollection(config.ns);
+                invariant(coll);  // 'exec' hasn't been killed, so collection must be alive.
                 coll->infoCache()->notifyOfQuery(txn, stats.indexesUsed);
                 CurOp::get(txn)->debug().fromMultiPlanner = stats.fromMultiPlanner;
                 CurOp::get(txn)->debug().replanned = stats.replanned;
diff --git a/src/mongo/db/commands/mr.h b/src/mongo/db/commands/mr.h
index 8bc4264794e..04b5f12a661 100644
--- a/src/mongo/db/commands/mr.h
+++ b/src/mongo/db/commands/mr.h
@@ -372,7 +372,7 @@ public:
     void switchMode(bool jsMode);
     void bailFromJS();
 
-    Collection* getCollectionOrUassert(Database* db, StringData ns);
+    static Collection* getCollectionOrUassert(Database* db, StringData ns);
 
     const Config& _config;
     DBDirectClient _db;