SERVER-37750 Optimized $sample stage does not yield

(cherry picked from commit 4177d6d22ab3329a8607bf80a62aa03d4fb2c528)

2 files changed, 17 insertions, 17 deletions

diff --git a/src/mongo/db/pipeline/pipeline_d.cpp b/src/mongo/db/pipeline/pipeline_d.cpp
index 9d49b4ef6cd..9be12b669e6 100644
--- a/src/mongo/db/pipeline/pipeline_d.cpp
+++ b/src/mongo/db/pipeline/pipeline_d.cpp
@@ -455,6 +455,10 @@ private:
  */
 StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> createRandomCursorExecutor(
     Collection* collection, OperationContext* opCtx, long long sampleSize, long long numRecords) {
+    // Verify that we are already under a collection lock. We avoid taking locks ourselves in this
+    // function because double-locking forces any PlanExecutor we create to adopt a NO_YIELD policy.
+    invariant(opCtx->lockState()->isCollectionLockedForMode(collection->ns().ns(), MODE_IS));
+
     double kMaxSampleRatioForRandCursor = 0.05;
     if (sampleSize > numRecords * kMaxSampleRatioForRandCursor || numRecords <= 100) {
         return {nullptr};
@@ -499,22 +503,18 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> createRandomCursorEx
             opCtx, ws.get(), idxIterator.release(), nullptr, collection);
     }
 
-    {
-        AutoGetCollectionForRead autoColl(opCtx, collection->ns());
-
-        // If we're in a sharded environment, we need to filter out documents we don't own.
-        if (ShardingState::get(opCtx)->needCollectionMetadata(opCtx, collection->ns().ns())) {
-            auto shardFilterStage = stdx::make_unique<ShardFilterStage>(
-                opCtx,
-                CollectionShardingState::get(opCtx, collection->ns())->getMetadata(),
-                ws.get(),
-                stage.release());
-            return PlanExecutor::make(opCtx,
-                                      std::move(ws),
-                                      std::move(shardFilterStage),
-                                      collection,
-                                      PlanExecutor::YIELD_AUTO);
-        }
+    // If we're in a sharded environment, we need to filter out documents we don't own.
+    if (ShardingState::get(opCtx)->needCollectionMetadata(opCtx, collection->ns().ns())) {
+        auto shardFilterStage = stdx::make_unique<ShardFilterStage>(
+            opCtx,
+            CollectionShardingState::get(opCtx, collection->ns())->getMetadata(),
+            ws.get(),
+            stage.release());
+        return PlanExecutor::make(opCtx,
+                                  std::move(ws),
+                                  std::move(shardFilterStage),
+                                  collection,
+                                  PlanExecutor::YIELD_AUTO);
     }
 
     return PlanExecutor::make(
diff --git a/src/mongo/shell/shell_options.cpp b/src/mongo/shell/shell_options.cpp
index c85527b6364..ea9987ab6e1 100644
--- a/src/mongo/shell/shell_options.cpp
+++ b/src/mongo/shell/shell_options.cpp
@@ -491,7 +491,7 @@ Status storeMongoShellOptions(const moe::Environment& params,
         if (uriOptions.count("compressors"))
             shellGlobalParams.networkMessageCompressors = uriOptions["compressors"];
     }
-    
+
     if (!shellGlobalParams.networkMessageCompressors.empty()) {
         const auto ret =
             storeMessageCompressionOptions(shellGlobalParams.networkMessageCompressors);