1 files changed, 29 insertions, 27 deletions

diff --git a/src/mongo/db/commands/mr.cpp b/src/mongo/db/commands/mr.cpp
index 4a23c78b7fb..3637948a1e7 100644
--- a/src/mongo/db/commands/mr.cpp
+++ b/src/mongo/db/commands/mr.cpp
@@ -62,7 +62,8 @@
 #include "mongo/db/query/plan_summary_stats.h"
 #include "mongo/db/query/query_planner.h"
 #include "mongo/db/repl/replication_coordinator.h"
-#include "mongo/db/s/collection_sharding_runtime.h"
+#include "mongo/db/s/collection_metadata.h"
+#include "mongo/db/s/collection_sharding_state.h"
 #include "mongo/db/s/sharding_state.h"
 #include "mongo/db/server_options.h"
 #include "mongo/db/service_context.h"
@@ -453,9 +454,6 @@ void State::prepTempCollection() {
                                         << status.code());
             }
             wuow.commit();
-
-            CollectionShardingRuntime::get(_opCtx, _config.incLong)
-                ->setFilteringMetadata(_opCtx, CollectionMetadata());
         });
     }
 
@@ -505,7 +503,6 @@ void State::prepTempCollection() {
 
         CollectionOptions options = finalOptions;
         options.temp = true;
-
         // If a UUID for the final output collection was sent by mongos (i.e., the final output
         // collection is sharded), use the UUID mongos sent when creating the temp collection.
         // When the temp collection is renamed to the final output collection, the UUID will be
@@ -537,9 +534,6 @@ void State::prepTempCollection() {
                 _opCtx, _config.tempNamespace, uuid, *it, false);
         }
         wuow.commit();
-
-        CollectionShardingRuntime::get(_opCtx, _config.tempNamespace)
-            ->setFilteringMetadata(_opCtx, CollectionMetadata());
     });
 }
 
@@ -1432,9 +1426,12 @@ public:
 
         uassert(16149, "cannot run map reduce without the js engine", getGlobalScriptEngine());
 
-        const auto metadata = [&] {
+        // Prevent sharding state from changing during the MR.
+        const auto collMetadata = [&] {
+            // Get metadata before we check our version, to make sure it doesn't increment in the
+            // meantime
             AutoGetCollectionForReadCommand autoColl(opCtx, config.nss);
-            return CollectionShardingState::get(opCtx, config.nss)->getMetadataForOperation(opCtx);
+            return CollectionShardingState::get(opCtx, config.nss)->getMetadata(opCtx);
         }();
 
         bool shouldHaveData = false;
@@ -1502,13 +1499,17 @@ public:
                 const ExtensionsCallbackReal extensionsCallback(opCtx, &config.nss);
 
                 const boost::intrusive_ptr<ExpressionContext> expCtx;
-                auto cq = uassertStatusOKWithContext(
+                auto statusWithCQ =
                     CanonicalQuery::canonicalize(opCtx,
                                                  std::move(qr),
                                                  expCtx,
                                                  extensionsCallback,
-                                                 MatchExpressionParser::kAllowAllSpecialFeatures),
-                    str::stream() << "Can't canonicalize query " << config.filter);
+                                                 MatchExpressionParser::kAllowAllSpecialFeatures);
+                if (!statusWithCQ.isOK()) {
+                    uasserted(17238, "Can't canonicalize query " + config.filter.toString());
+                    return 0;
+                }
+                std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
 
                 unique_ptr<PlanExecutor, PlanExecutor::Deleter> exec;
                 {
@@ -1535,37 +1536,38 @@ public:
 
                 Timer mt;
 
+                // go through each doc
                 BSONObj o;
                 PlanExecutor::ExecState execState;
                 while (PlanExecutor::ADVANCED == (execState = exec->getNext(&o, NULL))) {
-                    o = o.getOwned();  // The object will be accessed outside of collection lock
-
-                    // Check to see if this is a new object we don't own yet because of a chunk
-                    // migration
-                    if (metadata->isSharded()) {
-                        ShardKeyPattern kp(metadata->getKeyPattern());
-                        if (!metadata->keyBelongsToMe(kp.extractShardKeyFromDoc(o))) {
+                    o = o.getOwned();  // we will be accessing outside of the lock
+                    // check to see if this is a new object we don't own yet
+                    // because of a chunk migration
+                    if (collMetadata->isSharded()) {
+                        ShardKeyPattern kp(collMetadata->getKeyPattern());
+                        if (!collMetadata->keyBelongsToMe(kp.extractShardKeyFromDoc(o))) {
                             continue;
                         }
                     }
 
+                    // do map
                     if (config.verbose)
                         mt.reset();
-
                     config.mapper->map(o);
-
                     if (config.verbose)
                         mapTime += mt.micros();
 
-                    // Check if the state accumulated so far needs to be written to a collection.
-                    // This may yield the DB lock temporarily and then acquire it again.
+                    // Check if the state accumulated so far needs to be written to a
+                    // collection. This may yield the DB lock temporarily and then
+                    // acquire it again.
+                    //
                     numInputs++;
                     if (numInputs % 100 == 0) {
                         Timer t;
 
-                        // TODO: As an optimization, we might want to do the save/restore state and
-                        // yield inside the reduceAndSpillInMemoryState method, so it only happens
-                        // if necessary.
+                        // TODO: As an optimization, we might want to do the save/restore
+                        // state and yield inside the reduceAndSpillInMemoryState method, so
+                        // it only happens if necessary.
                         exec->saveState();
 
                         scopedAutoDb.reset();