SERVER-33354 Add snapshot read support for additional commands

7 files changed, 53 insertions, 42 deletions

diff --git a/src/mongo/db/commands/count_cmd.cpp b/src/mongo/db/commands/count_cmd.cpp
index 85a6e1f02c8..91ba0169b0c 100644
--- a/src/mongo/db/commands/count_cmd.cpp
+++ b/src/mongo/db/commands/count_cmd.cpp
@@ -152,11 +152,8 @@ public:
         // version on initial entry into count.
         auto rangePreserver = CollectionShardingState::get(opCtx, nss)->getMetadata();
 
-        auto statusWithPlanExecutor = getExecutorCount(opCtx,
-                                                       collection,
-                                                       request.getValue(),
-                                                       true,  // explain
-                                                       PlanExecutor::YIELD_AUTO);
+        auto statusWithPlanExecutor =
+            getExecutorCount(opCtx, collection, request.getValue(), true /*explain*/);
         if (!statusWithPlanExecutor.isOK()) {
             return statusWithPlanExecutor.getStatus();
         }
@@ -207,11 +204,8 @@ public:
         // version on initial entry into count.
         auto rangePreserver = CollectionShardingState::get(opCtx, nss)->getMetadata();
 
-        auto statusWithPlanExecutor = getExecutorCount(opCtx,
-                                                       collection,
-                                                       request.getValue(),
-                                                       false,  // !explain
-                                                       PlanExecutor::YIELD_AUTO);
+        auto statusWithPlanExecutor =
+            getExecutorCount(opCtx, collection, request.getValue(), false /*explain*/);
         if (!statusWithPlanExecutor.isOK()) {
             return CommandHelpers::appendCommandStatus(result, statusWithPlanExecutor.getStatus());
         }
diff --git a/src/mongo/db/commands/distinct.cpp b/src/mongo/db/commands/distinct.cpp
index e97d9ee1933..e5cdeba3a60 100644
--- a/src/mongo/db/commands/distinct.cpp
+++ b/src/mongo/db/commands/distinct.cpp
@@ -148,8 +148,8 @@ public:
 
         Collection* const collection = ctx->getCollection();
 
-        auto executor = uassertStatusOK(getExecutorDistinct(
-            opCtx, collection, nss.ns(), &parsedDistinct, PlanExecutor::YIELD_AUTO));
+        auto executor =
+            uassertStatusOK(getExecutorDistinct(opCtx, collection, nss.ns(), &parsedDistinct));
 
         Explain::explainStages(executor.get(), collection, verbosity, out);
         return Status::OK();
@@ -188,8 +188,7 @@ public:
 
         Collection* const collection = ctx->getCollection();
 
-        auto executor = getExecutorDistinct(
-            opCtx, collection, nss.ns(), &parsedDistinct, PlanExecutor::YIELD_AUTO);
+        auto executor = getExecutorDistinct(opCtx, collection, nss.ns(), &parsedDistinct);
         if (!executor.isOK()) {
             return CommandHelpers::appendCommandStatus(result, executor.getStatus());
         }
diff --git a/src/mongo/db/commands/group_cmd.cpp b/src/mongo/db/commands/group_cmd.cpp
index 3441e23a1ea..de3811cc6fa 100644
--- a/src/mongo/db/commands/group_cmd.cpp
+++ b/src/mongo/db/commands/group_cmd.cpp
@@ -134,8 +134,7 @@ private:
         AutoGetCollectionForReadCommand ctx(opCtx, groupRequest.ns);
         Collection* coll = ctx.getCollection();
 
-        auto statusWithPlanExecutor =
-            getExecutorGroup(opCtx, coll, groupRequest, PlanExecutor::YIELD_AUTO);
+        auto statusWithPlanExecutor = getExecutorGroup(opCtx, coll, groupRequest);
         if (!statusWithPlanExecutor.isOK()) {
             return statusWithPlanExecutor.getStatus();
         }
@@ -164,8 +163,7 @@ private:
         AutoGetCollectionForReadCommand ctx(opCtx, groupRequest.ns);
         Collection* coll = ctx.getCollection();
 
-        auto statusWithPlanExecutor =
-            getExecutorGroup(opCtx, coll, groupRequest, PlanExecutor::YIELD_AUTO);
+        auto statusWithPlanExecutor = getExecutorGroup(opCtx, coll, groupRequest);
         if (!statusWithPlanExecutor.isOK()) {
             return CommandHelpers::appendCommandStatus(result, statusWithPlanExecutor.getStatus());
         }
diff --git a/src/mongo/db/commands/parallel_collection_scan.cpp b/src/mongo/db/commands/parallel_collection_scan.cpp
index c040e9e0510..ae3818ab9ea 100644
--- a/src/mongo/db/commands/parallel_collection_scan.cpp
+++ b/src/mongo/db/commands/parallel_collection_scan.cpp
@@ -129,8 +129,15 @@ public:
                 make_unique<MultiIteratorStage>(opCtx, ws.get(), collection);
 
             // Takes ownership of 'ws' and 'mis'.
+            const auto& readConcernArgs = repl::ReadConcernArgs::get(opCtx);
             auto statusWithPlanExecutor = PlanExecutor::make(
-                opCtx, std::move(ws), std::move(mis), collection, PlanExecutor::YIELD_AUTO);
+                opCtx,
+                std::move(ws),
+                std::move(mis),
+                collection,
+                readConcernArgs.getLevel() == repl::ReadConcernLevel::kSnapshotReadConcern
+                    ? PlanExecutor::INTERRUPT_ONLY
+                    : PlanExecutor::YIELD_AUTO);
             invariant(statusWithPlanExecutor.isOK());
             execs.push_back(std::move(statusWithPlanExecutor.getValue()));
         }
@@ -167,6 +174,7 @@ public:
             bucketsBuilder.append(threadResult.obj());
         }
         result.appendArray("cursors", bucketsBuilder.obj());
+        opCtx->setStashedCursor();
 
         return true;
     }
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index f69322c5329..ce5e2b3718e 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -1018,15 +1018,17 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorUpdate(
 //
 
 StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorGroup(
-    OperationContext* opCtx,
-    Collection* collection,
-    const GroupRequest& request,
-    PlanExecutor::YieldPolicy yieldPolicy) {
+    OperationContext* opCtx, Collection* collection, const GroupRequest& request) {
     if (!getGlobalScriptEngine()) {
         return Status(ErrorCodes::BadValue, "server-side JavaScript execution is disabled");
     }
 
     unique_ptr<WorkingSet> ws = make_unique<WorkingSet>();
+    const auto readConcernArgs = repl::ReadConcernArgs::get(opCtx);
+    const auto yieldPolicy =
+        readConcernArgs.getLevel() == repl::ReadConcernLevel::kSnapshotReadConcern
+        ? PlanExecutor::INTERRUPT_ONLY
+        : PlanExecutor::YIELD_AUTO;
 
     if (!collection) {
         // Treat collections that do not exist as empty collections.  Note that the explain
@@ -1269,11 +1271,7 @@ BSONObj getDistinctProjection(const std::string& field) {
 }  // namespace
 
 StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount(
-    OperationContext* opCtx,
-    Collection* collection,
-    const CountRequest& request,
-    bool explain,
-    PlanExecutor::YieldPolicy yieldPolicy) {
+    OperationContext* opCtx, Collection* collection, const CountRequest& request, bool explain) {
     unique_ptr<WorkingSet> ws = make_unique<WorkingSet>();
 
     auto qr = stdx::make_unique<QueryRequest>(request.getNs());
@@ -1296,6 +1294,13 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount(
         return statusWithCQ.getStatus();
     }
     unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
+
+    const auto readConcernArgs = repl::ReadConcernArgs::get(opCtx);
+    const auto yieldPolicy =
+        readConcernArgs.getLevel() == repl::ReadConcernLevel::kSnapshotReadConcern
+        ? PlanExecutor::INTERRUPT_ONLY
+        : PlanExecutor::YIELD_AUTO;
+
     if (!collection) {
         // Treat collections that do not exist as empty collections. Note that the explain
         // reporting machinery always assumes that the root stage for a count operation is
@@ -1477,8 +1482,13 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDistinct(
     OperationContext* opCtx,
     Collection* collection,
     const std::string& ns,
-    ParsedDistinct* parsedDistinct,
-    PlanExecutor::YieldPolicy yieldPolicy) {
+    ParsedDistinct* parsedDistinct) {
+    const auto readConcernArgs = repl::ReadConcernArgs::get(opCtx);
+    const auto yieldPolicy =
+        readConcernArgs.getLevel() == repl::ReadConcernLevel::kSnapshotReadConcern
+        ? PlanExecutor::INTERRUPT_ONLY
+        : PlanExecutor::YIELD_AUTO;
+
     if (!collection) {
         // Treat collections that do not exist as empty collections.
         return PlanExecutor::make(opCtx,
diff --git a/src/mongo/db/query/get_executor.h b/src/mongo/db/query/get_executor.h
index 6a801c1a81b..237f577de7d 100644
--- a/src/mongo/db/query/get_executor.h
+++ b/src/mongo/db/query/get_executor.h
@@ -122,8 +122,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDist
     OperationContext* opCtx,
     Collection* collection,
     const std::string& ns,
-    ParsedDistinct* parsedDistinct,
-    PlanExecutor::YieldPolicy yieldPolicy);
+    ParsedDistinct* parsedDistinct);
 
 /*
  * Get a PlanExecutor for a query executing as part of a count command.
@@ -133,11 +132,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDist
  * executing a count.
  */
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount(
-    OperationContext* opCtx,
-    Collection* collection,
-    const CountRequest& request,
-    bool explain,
-    PlanExecutor::YieldPolicy yieldPolicy);
+    OperationContext* opCtx, Collection* collection, const CountRequest& request, bool explain);
 
 /**
  * Get a PlanExecutor for a delete operation. 'parsedDelete' describes the query predicate
@@ -187,9 +182,6 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorUpda
  * If an executor could not be created, returns a Status indicating why.
  */
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorGroup(
-    OperationContext* opCtx,
-    Collection* collection,
-    const GroupRequest& request,
-    PlanExecutor::YieldPolicy yieldPolicy);
+    OperationContext* opCtx, Collection* collection, const GroupRequest& request);
 
 }  // namespace mongo
diff --git a/src/mongo/db/service_entry_point_common.cpp b/src/mongo/db/service_entry_point_common.cpp
index 64f38e60008..b6b2ccbb8e3 100644
--- a/src/mongo/db/service_entry_point_common.cpp
+++ b/src/mongo/db/service_entry_point_common.cpp
@@ -132,10 +132,12 @@ const StringMap<int> sessionCheckoutWhitelist = {{"aggregate", 1},
 const StringMap<int> readConcernSnapshotWhitelist = {{"aggregate", 1},
                                                      {"count", 1},
                                                      {"delete", 1},
+                                                     {"distinct", 1},
                                                      {"find", 1},
                                                      {"findandmodify", 1},
                                                      {"findAndModify", 1},
                                                      {"geoSearch", 1},
+                                                     {"group", 1},
                                                      {"insert", 1},
                                                      {"parallelCollectionScan", 1},
                                                      {"update", 1}};
@@ -1094,8 +1096,13 @@ DbResponse ServiceEntryPointCommon::handleRequest(OperationContext* opCtx,
     DbMessage dbmsg(m);
 
     Client& c = *opCtx->getClient();
+
     if (c.isInDirectClient()) {
-        invariant(!opCtx->lockState()->inAWriteUnitOfWork());
+        if (!opCtx->getLogicalSessionId() || !opCtx->getTxnNumber() ||
+            opCtx->recoveryUnit()->getReadConcernLevel() !=
+                repl::ReadConcernLevel::kSnapshotReadConcern) {
+            invariant(!opCtx->lockState()->inAWriteUnitOfWork());
+        }
     } else {
         LastError::get(c).startRequest();
         AuthorizationSession::get(c)->startRequest(opCtx);
@@ -1202,6 +1209,8 @@ DbResponse ServiceEntryPointCommon::handleRequest(OperationContext* opCtx,
         // Performance profiling is on
         if (opCtx->lockState()->isReadLocked()) {
             LOG(1) << "note: not profiling because recursive read lock";
+        } else if (c.isInDirectClient()) {
+            LOG(1) << "note: not profiling because we are in DBDirectClient";
         } else if (behaviors.lockedForWriting()) {
             // TODO SERVER-26825: Fix race condition where fsyncLock is acquired post
             // lockedForWriting() call but prior to profile collection lock acquisition.
@@ -1209,6 +1218,7 @@ DbResponse ServiceEntryPointCommon::handleRequest(OperationContext* opCtx,
         } else if (storageGlobalParams.readOnly) {
             LOG(1) << "note: not profiling because server is read-only";
         } else {
+            invariant(!opCtx->lockState()->inAWriteUnitOfWork());
             profile(opCtx, op);
         }
     }