SERVER-47952 Shard selects timestamp for one-shard snapshot find

9 files changed, 262 insertions, 82 deletions

diff --git a/src/mongo/s/commands/cluster_command_test_fixture.cpp b/src/mongo/s/commands/cluster_command_test_fixture.cpp
index 0d1a88bf578..81951cd425c 100644
--- a/src/mongo/s/commands/cluster_command_test_fixture.cpp
+++ b/src/mongo/s/commands/cluster_command_test_fixture.cpp
@@ -43,12 +43,19 @@
 #include "mongo/db/logical_time_validator.h"
 #include "mongo/db/vector_clock.h"
 #include "mongo/s/cluster_last_error_info.h"
+#include "mongo/s/query/cluster_cursor_manager.h"
 #include "mongo/util/fail_point.h"
 #include "mongo/util/options_parser/startup_option_init.h"
 #include "mongo/util/tick_source_mock.h"
 
 namespace mongo {
 
+const Timestamp ClusterCommandTestFixture::kAfterClusterTime(50, 2);
+
+const Timestamp ClusterCommandTestFixture::kShardClusterTime(51, 3);
+
+const CursorId ClusterCommandTestFixture::kCursorId(123);
+
 void ClusterCommandTestFixture::setUp() {
     CatalogCacheTestFixture::setUp();
     CatalogCacheTestFixture::setupNShards(numShards);
@@ -80,13 +87,24 @@ void ClusterCommandTestFixture::setUp() {
         "enableStaleVersionAndSnapshotRetriesWithinTransactions");
 }
 
-BSONObj ClusterCommandTestFixture::_makeCmd(BSONObj cmdObj, bool includeAfterClusterTime) {
+void ClusterCommandTestFixture::tearDown() {
+    // Delete any cursors left behind by tests, to avoid invariant in ~ClusterCursorManager.
+    auto opCtx = operationContext();
+    auto cursorManager = Grid::get(opCtx)->getCursorManager();
+    cursorManager->killAllCursors(opCtx);
+}
+
+BSONObj ClusterCommandTestFixture::_makeCmd(BSONObj cmdObj,
+                                            bool startTransaction,
+                                            bool includeAfterClusterTime) {
     BSONObjBuilder bob(cmdObj);
     // Each command runs in a new session.
     bob.append("lsid", makeLogicalSessionIdForTest().toBSON());
-    bob.append("txnNumber", TxnNumber(1));
-    bob.append("autocommit", false);
-    bob.append("startTransaction", true);
+    if (startTransaction) {
+        bob.append("txnNumber", TxnNumber(1));
+        bob.append("autocommit", false);
+        bob.append("startTransaction", true);
+    }
 
     BSONObjBuilder readConcernBob = bob.subobjStart(repl::ReadConcernArgs::kReadConcernFieldName);
     readConcernBob.append("level", "snapshot");
@@ -98,6 +116,14 @@ BSONObj ClusterCommandTestFixture::_makeCmd(BSONObj cmdObj, bool includeAfterClu
     return bob.obj();
 }
 
+BSONObj ClusterCommandTestFixture::_makeTxnCmd(BSONObj cmdObj, bool includeAfterClusterTime) {
+    return _makeCmd(cmdObj, true, includeAfterClusterTime);
+}
+
+BSONObj ClusterCommandTestFixture::_makeNonTxnCmd(BSONObj cmdObj, bool includeAfterClusterTime) {
+    return _makeCmd(cmdObj, false, includeAfterClusterTime);
+}
+
 void ClusterCommandTestFixture::expectReturnsError(ErrorCodes::Error code) {
     onCommandForPoolExecutor([code](const executor::RemoteCommandRequest& request) {
         BSONObjBuilder resBob;
@@ -227,80 +253,108 @@ void ClusterCommandTestFixture::runTxnCommandMaxErrors(BSONObj cmd,
 }
 
 void ClusterCommandTestFixture::testNoErrors(BSONObj targetedCmd, BSONObj scatterGatherCmd) {
-
     // Target one shard.
-    runCommandSuccessful(_makeCmd(targetedCmd), true);
+    runCommandSuccessful(_makeTxnCmd(targetedCmd), true);
 
     // Target all shards.
     if (!scatterGatherCmd.isEmpty()) {
-        runCommandSuccessful(_makeCmd(scatterGatherCmd), false);
+        runCommandSuccessful(_makeTxnCmd(scatterGatherCmd), false);
     }
 }
 
 void ClusterCommandTestFixture::testRetryOnSnapshotError(BSONObj targetedCmd,
                                                          BSONObj scatterGatherCmd) {
     // Target one shard.
-    runTxnCommandOneError(_makeCmd(targetedCmd), ErrorCodes::SnapshotUnavailable, true);
-    runTxnCommandOneError(_makeCmd(targetedCmd), ErrorCodes::SnapshotTooOld, true);
+    runTxnCommandOneError(_makeTxnCmd(targetedCmd), ErrorCodes::SnapshotUnavailable, true);
+    runTxnCommandOneError(_makeTxnCmd(targetedCmd), ErrorCodes::SnapshotTooOld, true);
 
     // Target all shards
     if (!scatterGatherCmd.isEmpty()) {
-        runTxnCommandOneError(_makeCmd(scatterGatherCmd), ErrorCodes::SnapshotUnavailable, false);
-        runTxnCommandOneError(_makeCmd(scatterGatherCmd), ErrorCodes::SnapshotTooOld, false);
+        runTxnCommandOneError(
+            _makeTxnCmd(scatterGatherCmd), ErrorCodes::SnapshotUnavailable, false);
+        runTxnCommandOneError(_makeTxnCmd(scatterGatherCmd), ErrorCodes::SnapshotTooOld, false);
     }
 }
 
 void ClusterCommandTestFixture::testMaxRetriesSnapshotErrors(BSONObj targetedCmd,
                                                              BSONObj scatterGatherCmd) {
     // Target one shard.
-    runTxnCommandMaxErrors(_makeCmd(targetedCmd), ErrorCodes::SnapshotUnavailable, true);
-    runTxnCommandMaxErrors(_makeCmd(targetedCmd), ErrorCodes::SnapshotTooOld, true);
+    runTxnCommandMaxErrors(_makeTxnCmd(targetedCmd), ErrorCodes::SnapshotUnavailable, true);
+    runTxnCommandMaxErrors(_makeTxnCmd(targetedCmd), ErrorCodes::SnapshotTooOld, true);
 
     // Target all shards
     if (!scatterGatherCmd.isEmpty()) {
-        runTxnCommandMaxErrors(_makeCmd(scatterGatherCmd), ErrorCodes::SnapshotUnavailable, false);
-        runTxnCommandMaxErrors(_makeCmd(scatterGatherCmd), ErrorCodes::SnapshotTooOld, false);
+        runTxnCommandMaxErrors(
+            _makeTxnCmd(scatterGatherCmd), ErrorCodes::SnapshotUnavailable, false);
+        runTxnCommandMaxErrors(_makeTxnCmd(scatterGatherCmd), ErrorCodes::SnapshotTooOld, false);
     }
 }
 
-void ClusterCommandTestFixture::testAttachesAtClusterTimeForSnapshotReadConcern(
-    BSONObj targetedCmd, BSONObj scatterGatherCmd) {
+void ClusterCommandTestFixture::testAttachesAtClusterTimeForTxnSnapshotReadConcern(
+    BSONObj targetedCmd, BSONObj scatterGatherCmd, bool createsCursor) {
+
+    // Target one shard.
+    runCommandInspectRequests(_makeTxnCmd(targetedCmd), _containsAtClusterTimeOnly, true);
+    if (createsCursor)
+        _assertCursorReadConcern(true, boost::none);
+
+    // Target all shards.
+    if (!scatterGatherCmd.isEmpty()) {
+        runCommandInspectRequests(_makeTxnCmd(scatterGatherCmd), _containsAtClusterTimeOnly, false);
+        if (createsCursor)
+            _assertCursorReadConcern(false, boost::none);
+    }
+}
 
-    auto containsAtClusterTime = [](const executor::RemoteCommandRequest& request) {
-        ASSERT(!request.cmdObj["readConcern"]["atClusterTime"].eoo());
-    };
+void ClusterCommandTestFixture::testTxnSnapshotReadConcernWithAfterClusterTime(
+    BSONObj targetedCmd, BSONObj scatterGatherCmd, bool createsCursor) {
 
     // Target one shard.
-    runCommandInspectRequests(_makeCmd(targetedCmd), containsAtClusterTime, true);
+    runCommandInspectRequests(_makeTxnCmd(targetedCmd, true), _containsAtClusterTimeOnly, true);
+    if (createsCursor)
+        _assertCursorReadConcern(true, boost::none);
 
     // Target all shards.
     if (!scatterGatherCmd.isEmpty()) {
-        runCommandInspectRequests(_makeCmd(scatterGatherCmd), containsAtClusterTime, false);
+        runCommandInspectRequests(
+            _makeTxnCmd(scatterGatherCmd, true), _containsAtClusterTimeOnly, false);
+        if (createsCursor)
+            _assertCursorReadConcern(false, boost::none);
     }
 }
 
-void ClusterCommandTestFixture::testSnapshotReadConcernWithAfterClusterTime(
-    BSONObj targetedCmd, BSONObj scatterGatherCmd) {
+void ClusterCommandTestFixture::testAttachesAtClusterTimeForNonTxnSnapshotReadConcern(
+    BSONObj targetedCmd, BSONObj scatterGatherCmd, bool createsCursor) {
+
+    // Target one shard.
+    runCommandInspectRequests(_makeNonTxnCmd(targetedCmd), _omitsClusterTime, true);
+    if (createsCursor)
+        _assertCursorReadConcern(true, kShardClusterTime);
 
-    auto containsAtClusterTimeNoAfterClusterTime =
-        [&](const executor::RemoteCommandRequest& request) {
-            ASSERT(!request.cmdObj["readConcern"]["atClusterTime"].eoo());
-            ASSERT(request.cmdObj["readConcern"]["afterClusterTime"].eoo());
+    // Target all shards.
+    if (!scatterGatherCmd.isEmpty()) {
+        runCommandInspectRequests(
+            _makeNonTxnCmd(scatterGatherCmd), _containsAtClusterTimeOnly, false);
+        if (createsCursor)
+            _assertCursorReadConcern(false, kInMemoryLogicalTime.asTimestamp());
+    }
+}
 
-            // The chosen atClusterTime should be greater than or equal to the request's
-            // afterClusterTime.
-            ASSERT_GTE(LogicalTime(request.cmdObj["readConcern"]["atClusterTime"].timestamp()),
-                       LogicalTime(kAfterClusterTime));
-        };
+void ClusterCommandTestFixture::testNonTxnSnapshotReadConcernWithAfterClusterTime(
+    BSONObj targetedCmd, BSONObj scatterGatherCmd, bool createsCursor) {
 
     // Target one shard.
     runCommandInspectRequests(
-        _makeCmd(targetedCmd, true), containsAtClusterTimeNoAfterClusterTime, true);
+        _makeNonTxnCmd(targetedCmd, true), _containsAfterClusterTimeOnly, true);
+    if (createsCursor)
+        _assertCursorReadConcern(true, kShardClusterTime);
 
     // Target all shards.
     if (!scatterGatherCmd.isEmpty()) {
         runCommandInspectRequests(
-            _makeCmd(scatterGatherCmd, true), containsAtClusterTimeNoAfterClusterTime, false);
+            _makeNonTxnCmd(scatterGatherCmd, true), _containsAtClusterTimeOnly, false);
+        if (createsCursor)
+            _assertCursorReadConcern(false, kAfterClusterTime);
     }
 }
 
@@ -310,6 +364,73 @@ void ClusterCommandTestFixture::appendTxnResponseMetadata(BSONObjBuilder& bob) {
     txnResponseMetadata.serialize(&bob);
 }
 
+void ClusterCommandTestFixture::_assertCursorReadConcern(
+    bool isTargeted, boost::optional<Timestamp> expectedAtClusterTime) {
+    auto opCtx = operationContext();
+    auto cursorManager = Grid::get(opCtx)->getCursorManager();
+    auto cursors =
+        cursorManager->getIdleCursors(opCtx, MongoProcessInterface::CurrentOpUserMode::kIncludeAll);
+    ASSERT_EQUALS(cursors.size(), 1);
+    auto cursorId = *cursors[0].getCursorId();
+    auto pinnedCursor = cursorManager->checkOutCursor(
+        kNss, cursorId, opCtx, [](UserNameIterator) { return Status::OK(); });
+
+    ASSERT_OK(pinnedCursor.getStatus());
+    auto readConcern = pinnedCursor.getValue()->getReadConcern();
+    auto nRemotes = pinnedCursor.getValue()->getNumRemotes();
+    ASSERT_EQUALS(nRemotes, isTargeted ? 1 : 2);
+
+    // User supplies no atClusterTime. If not isTargeted then mongos selected a timestamp and called
+    // setArgsAtClusterTime. Else mongos let the shard select atClusterTime. The shard returned it
+    // to mongos, and mongos called setArgsAtClusterTime.
+    if (expectedAtClusterTime) {
+        ASSERT_EQUALS(readConcern->getArgsAtClusterTime()->asTimestamp(), *expectedAtClusterTime);
+        ASSERT_TRUE(readConcern->wasAtClusterTimeSelected());
+    } else {
+        ASSERT_FALSE(readConcern->getArgsAtClusterTime());
+        ASSERT_FALSE(readConcern->wasAtClusterTimeSelected());
+    }
+
+    // Kill cursor on shard(s) and remove from ClusterCursorManager, to ensure we have exactly 1
+    // idle cursor next time this method is called. (Otherwise this test would need a complex
+    // method for determining which cursor to examine.)
+    pinnedCursor.getValue().returnCursor(ClusterCursorManager::CursorState::Exhausted);
+    for (std::size_t i = 0; i < nRemotes; ++i) {
+        onCommandForPoolExecutor([&](const executor::RemoteCommandRequest& request) {
+            ASSERT_EQUALS("killCursors"_sd, request.cmdObj.firstElement().fieldNameStringData());
+            ASSERT_EQUALS(kNss.coll(), request.cmdObj.firstElement().valueStringData());
+            return BSON("ok" << 1);
+        });
+    }
+}
+
+void ClusterCommandTestFixture::_containsSelectedAtClusterTime(
+    const executor::RemoteCommandRequest& request) {
+    ASSERT(!request.cmdObj["readConcern"]["atClusterTime"].eoo());
+    ASSERT(request.cmdObj["readConcern"]["afterClusterTime"].eoo());
+
+    // The chosen atClusterTime should be greater than or equal to the request's afterClusterTime.
+    ASSERT_GTE(LogicalTime(request.cmdObj["readConcern"]["atClusterTime"].timestamp()),
+               LogicalTime(kAfterClusterTime));
+}
+
+void ClusterCommandTestFixture::_containsAtClusterTimeOnly(
+    const executor::RemoteCommandRequest& request) {
+    ASSERT(!request.cmdObj["readConcern"]["atClusterTime"].eoo());
+    ASSERT(request.cmdObj["readConcern"]["afterClusterTime"].eoo());
+}
+
+void ClusterCommandTestFixture::_containsAfterClusterTimeOnly(
+    const executor::RemoteCommandRequest& request) {
+    ASSERT(request.cmdObj["readConcern"]["atClusterTime"].eoo());
+    ASSERT(!request.cmdObj["readConcern"]["afterClusterTime"].eoo());
+}
+
+void ClusterCommandTestFixture::_omitsClusterTime(const executor::RemoteCommandRequest& request) {
+    ASSERT(request.cmdObj["readConcern"]["atClusterTime"].eoo());
+    ASSERT(request.cmdObj["readConcern"]["afterClusterTime"].eoo());
+}
+
 // Satisfies dependency from StoreSASLOPtions.
 MONGO_STARTUP_OPTIONS_STORE(CoreOptions)(InitializerContext*) {
     return Status::OK();
diff --git a/src/mongo/s/commands/cluster_command_test_fixture.h b/src/mongo/s/commands/cluster_command_test_fixture.h
index 9c5406c073b..9c0f3f4d5b7 100644
--- a/src/mongo/s/commands/cluster_command_test_fixture.h
+++ b/src/mongo/s/commands/cluster_command_test_fixture.h
@@ -31,6 +31,7 @@
 
 #include "mongo/platform/basic.h"
 
+#include "mongo/db/cursor_id.h"
 #include "mongo/s/catalog_cache_test_fixture.h"
 #include "mongo/s/commands/strategy.h"
 
@@ -46,10 +47,16 @@ protected:
 
     const LogicalTime kInMemoryLogicalTime = LogicalTime(Timestamp(10, 1));
 
-    const Timestamp kAfterClusterTime = Timestamp(50, 2);
+    static const Timestamp kAfterClusterTime;
+
+    static const Timestamp kShardClusterTime;
+
+    static const CursorId kCursorId;
 
     void setUp() override;
 
+    void tearDown() override;
+
     virtual void expectInspectRequest(int shardIndex, InspectionCallback cb) = 0;
 
     virtual void expectReturnsSuccess(int shardIndex) = 0;
@@ -85,17 +92,34 @@ protected:
     void testMaxRetriesSnapshotErrors(BSONObj targetedCmd, BSONObj scatterGatherCmd = BSONObj());
 
     /**
-     * Verifies that atClusterTime is attached to the given commands.
+     * Verifies that atClusterTime is attached to the given commands in a transaction.
      */
-    void testAttachesAtClusterTimeForSnapshotReadConcern(BSONObj targetedCmd,
-                                                         BSONObj scatterGatherCmd = BSONObj());
+    void testAttachesAtClusterTimeForTxnSnapshotReadConcern(BSONObj targetedCmd,
+                                                            BSONObj scatterGatherCmd = BSONObj(),
+                                                            bool createsCursor = false);
 
     /**
      * Verifies that the chosen atClusterTime is greater than or equal to each command's
-     * afterClusterTime.
+     * afterClusterTime in a transaction.
+     */
+    void testTxnSnapshotReadConcernWithAfterClusterTime(BSONObj targetedCmd,
+                                                        BSONObj scatterGatherCmd = BSONObj(),
+                                                        bool createsCursor = false);
+
+    /**
+     * Verifies that atClusterTime is attached to the given non-transaction snapshot commands.
      */
-    void testSnapshotReadConcernWithAfterClusterTime(BSONObj targetedCmd,
-                                                     BSONObj scatterGatherCmd = BSONObj());
+    void testAttachesAtClusterTimeForNonTxnSnapshotReadConcern(BSONObj targetedCmd,
+                                                               BSONObj scatterGatherCmd = BSONObj(),
+                                                               bool createsCursor = false);
+
+    /**
+     * Verifies that the chosen atClusterTime is greater than or equal to each non-transaction
+     * snapshot command's afterClusterTime.
+     */
+    void testNonTxnSnapshotReadConcernWithAfterClusterTime(BSONObj targetedCmd,
+                                                           BSONObj scatterGatherCmd = BSONObj(),
+                                                           bool createsCursor = false);
 
     /**
      * Appends the metadata shards return on responses to transaction statements, such as the
@@ -105,11 +129,36 @@ protected:
 
 private:
     /**
-     * Makes a new command object from the one given by apppending read concern
-     * snapshot and the appropriate transaction options. If includeAfterClusterTime
-     * is true, also appends afterClusterTime to the read concern.
+     * Makes a new command object from the one given by appending read concern snapshot and the
+     * appropriate transaction options. If includeAfterClusterTime is true, also appends
+     * afterClusterTime to the read concern.
      */
-    BSONObj _makeCmd(BSONObj cmdObj, bool includeAfterClusterTime = false);
+    BSONObj _makeTxnCmd(BSONObj cmdObj, bool includeAfterClusterTime = false);
+
+    /**
+     * Makes a new command object from the one given by appending read concern snapshot. If
+     * includeAfterClusterTime is true, also appends afterClusterTime to the read concern.
+     */
+    BSONObj _makeNonTxnCmd(BSONObj cmdObj, bool includeAfterClusterTime = false);
+
+    /**
+     * Helper method.
+     */
+    BSONObj _makeCmd(BSONObj cmdObj, bool startTransaction, bool includeAfterClusterTime);
+
+    /*
+     * Check that the ClusterCursorManager contains an idle cursor with proper readConcern set.
+     */
+    void _assertCursorReadConcern(bool isTargeted,
+                                  boost::optional<Timestamp> expectedAtClusterTime);
+
+    static void _containsSelectedAtClusterTime(const executor::RemoteCommandRequest& request);
+
+    static void _containsAtClusterTimeOnly(const executor::RemoteCommandRequest& request);
+
+    static void _containsAfterClusterTimeOnly(const executor::RemoteCommandRequest& request);
+
+    static void _omitsClusterTime(const executor::RemoteCommandRequest& request);
 
     // Enables the transaction router to retry within a transaction on stale version and snapshot
     // errors for the duration of each test.
diff --git a/src/mongo/s/commands/cluster_delete_test.cpp b/src/mongo/s/commands/cluster_delete_test.cpp
index 407ce5e78e2..f0fae5cbeb2 100644
--- a/src/mongo/s/commands/cluster_delete_test.cpp
+++ b/src/mongo/s/commands/cluster_delete_test.cpp
@@ -73,11 +73,11 @@ TEST_F(ClusterDeleteTest, NoErrors) {
 }
 
 TEST_F(ClusterDeleteTest, AttachesAtClusterTimeForSnapshotReadConcern) {
-    testAttachesAtClusterTimeForSnapshotReadConcern(kDeleteCmdTargeted, kDeleteCmdScatterGather);
+    testAttachesAtClusterTimeForTxnSnapshotReadConcern(kDeleteCmdTargeted, kDeleteCmdScatterGather);
 }
 
 TEST_F(ClusterDeleteTest, SnapshotReadConcernWithAfterClusterTime) {
-    testSnapshotReadConcernWithAfterClusterTime(kDeleteCmdTargeted, kDeleteCmdScatterGather);
+    testTxnSnapshotReadConcernWithAfterClusterTime(kDeleteCmdTargeted, kDeleteCmdScatterGather);
 }
 
 }  // namespace
diff --git a/src/mongo/s/commands/cluster_distinct_test.cpp b/src/mongo/s/commands/cluster_distinct_test.cpp
index 40f0226e820..d8513f36dc5 100644
--- a/src/mongo/s/commands/cluster_distinct_test.cpp
+++ b/src/mongo/s/commands/cluster_distinct_test.cpp
@@ -80,12 +80,12 @@ TEST_F(ClusterDistinctTest, MaxRetriesSnapshotErrors) {
 }
 
 TEST_F(ClusterDistinctTest, AttachesAtClusterTimeForSnapshotReadConcern) {
-    testAttachesAtClusterTimeForSnapshotReadConcern(kDistinctCmdTargeted,
-                                                    kDistinctCmdScatterGather);
+    testAttachesAtClusterTimeForTxnSnapshotReadConcern(kDistinctCmdTargeted,
+                                                       kDistinctCmdScatterGather);
 }
 
 TEST_F(ClusterDistinctTest, SnapshotReadConcernWithAfterClusterTime) {
-    testSnapshotReadConcernWithAfterClusterTime(kDistinctCmdTargeted, kDistinctCmdScatterGather);
+    testTxnSnapshotReadConcernWithAfterClusterTime(kDistinctCmdTargeted, kDistinctCmdScatterGather);
 }
 
 }  // namespace
diff --git a/src/mongo/s/commands/cluster_find_and_modify_test.cpp b/src/mongo/s/commands/cluster_find_and_modify_test.cpp
index f0a44494591..8f126039fd3 100644
--- a/src/mongo/s/commands/cluster_find_and_modify_test.cpp
+++ b/src/mongo/s/commands/cluster_find_and_modify_test.cpp
@@ -78,11 +78,11 @@ TEST_F(ClusterFindAndModifyTest, MaxRetriesSnapshotErrors) {
 }
 
 TEST_F(ClusterFindAndModifyTest, AttachesAtClusterTimeForSnapshotReadConcern) {
-    testAttachesAtClusterTimeForSnapshotReadConcern(kFindAndModifyCmdTargeted);
+    testAttachesAtClusterTimeForTxnSnapshotReadConcern(kFindAndModifyCmdTargeted);
 }
 
 TEST_F(ClusterFindAndModifyTest, SnapshotReadConcernWithAfterClusterTime) {
-    testSnapshotReadConcernWithAfterClusterTime(kFindAndModifyCmdTargeted);
+    testTxnSnapshotReadConcernWithAfterClusterTime(kFindAndModifyCmdTargeted);
 }
 
 }  // namespace
diff --git a/src/mongo/s/commands/cluster_find_test.cpp b/src/mongo/s/commands/cluster_find_test.cpp
index 8d0dc6792d4..7d9e219961d 100644
--- a/src/mongo/s/commands/cluster_find_test.cpp
+++ b/src/mongo/s/commands/cluster_find_test.cpp
@@ -31,34 +31,23 @@
 
 #include "mongo/db/query/cursor_response.h"
 #include "mongo/s/commands/cluster_command_test_fixture.h"
+#include "mongo/s/query/cluster_cursor_manager.h"
 
 namespace mongo {
 namespace {
 
 class ClusterFindTest : public ClusterCommandTestFixture {
 protected:
-    const BSONObj kFindCmdScatterGather = BSON("find"
-                                               << "coll");
-    const BSONObj kFindCmdTargeted = BSON("find"
-                                          << "coll"
-                                          << "filter" << BSON("_id" << 0));
+    // Batch size 1, so when expectInspectRequest returns one doc, mongos doesn't wait for more.
+    const BSONObj kFindCmdScatterGather = BSON("find" << kNss.coll() << "batchSize" << 1);
+    const BSONObj kFindCmdTargeted =
+        BSON("find" << kNss.coll() << "filter" << BSON("_id" << 0) << "batchSize" << 1);
 
     // The index of the shard expected to receive the response is used to prevent different shards
     // from returning documents with the same shard key. This is expected to be 0 for queries
     // targeting one shard.
     void expectReturnsSuccess(int shardIndex) override {
-        onCommandForPoolExecutor([&](const executor::RemoteCommandRequest& request) {
-            ASSERT_EQ(kNss.coll(), request.cmdObj.firstElement().valueStringData());
-
-            std::vector<BSONObj> batch = {BSON("_id" << shardIndex)};
-            CursorResponse cursorResponse(kNss, CursorId(0), batch);
-
-            BSONObjBuilder bob;
-            bob.appendElementsUnique(
-                cursorResponse.toBSON(CursorResponse::ResponseType::InitialResponse));
-            appendTxnResponseMetadata(bob);
-            return bob.obj();
-        });
+        expectInspectRequest(shardIndex, [](const executor::RemoteCommandRequest&) {});
     }
 
     void expectInspectRequest(int shardIndex, InspectionCallback cb) override {
@@ -68,8 +57,17 @@ protected:
             cb(request);
 
             std::vector<BSONObj> batch = {BSON("_id" << shardIndex)};
-            CursorResponse cursorResponse(kNss, CursorId(0), batch);
-
+            // User supplies no atClusterTime. For single-shard non-transaction snapshot reads,
+            // mongos lets the shard (which this function simulates) select a read timestamp.
+            boost::optional<Timestamp> atClusterTime;
+            if (request.cmdObj["txnNumber"].eoo() && !request.cmdObj["readConcern"].eoo()) {
+                auto rc = request.cmdObj["readConcern"].Obj();
+                if (rc["level"].String() == "snapshot" && rc["atClusterTime"].eoo()) {
+                    atClusterTime = kShardClusterTime;
+                }
+            }
+
+            CursorResponse cursorResponse(kNss, kCursorId, batch, atClusterTime);
             BSONObjBuilder bob;
             bob.appendElementsUnique(
                 cursorResponse.toBSON(CursorResponse::ResponseType::InitialResponse));
@@ -91,12 +89,23 @@ TEST_F(ClusterFindTest, MaxRetriesSnapshotErrors) {
     testMaxRetriesSnapshotErrors(kFindCmdTargeted, kFindCmdScatterGather);
 }
 
-TEST_F(ClusterFindTest, AttachesAtClusterTimeForSnapshotReadConcern) {
-    testAttachesAtClusterTimeForSnapshotReadConcern(kFindCmdTargeted, kFindCmdScatterGather);
+TEST_F(ClusterFindTest, AttachesAtClusterTimeForTransactionSnapshotReadConcern) {
+    testAttachesAtClusterTimeForTxnSnapshotReadConcern(
+        kFindCmdTargeted, kFindCmdScatterGather, true);
+}
+
+TEST_F(ClusterFindTest, TransactionSnapshotReadConcernWithAfterClusterTime) {
+    testTxnSnapshotReadConcernWithAfterClusterTime(kFindCmdTargeted, kFindCmdScatterGather, true);
+}
+
+TEST_F(ClusterFindTest, AttachesAtClusterTimeForNonTransactionSnapshotReadConcern) {
+    testAttachesAtClusterTimeForNonTxnSnapshotReadConcern(
+        kFindCmdTargeted, kFindCmdScatterGather, true);
 }
 
-TEST_F(ClusterFindTest, SnapshotReadConcernWithAfterClusterTime) {
-    testSnapshotReadConcernWithAfterClusterTime(kFindCmdTargeted, kFindCmdScatterGather);
+TEST_F(ClusterFindTest, NonTransactionSnapshotReadConcernWithAfterClusterTime) {
+    testNonTxnSnapshotReadConcernWithAfterClusterTime(
+        kFindCmdTargeted, kFindCmdScatterGather, true);
 }
 
 }  // namespace
diff --git a/src/mongo/s/commands/cluster_insert_test.cpp b/src/mongo/s/commands/cluster_insert_test.cpp
index ff5fb075486..24544eeb8d4 100644
--- a/src/mongo/s/commands/cluster_insert_test.cpp
+++ b/src/mongo/s/commands/cluster_insert_test.cpp
@@ -73,11 +73,11 @@ TEST_F(ClusterInsertTest, NoErrors) {
 }
 
 TEST_F(ClusterInsertTest, AttachesAtClusterTimeForSnapshotReadConcern) {
-    testAttachesAtClusterTimeForSnapshotReadConcern(kInsertCmdTargeted, kInsertCmdScatterGather);
+    testAttachesAtClusterTimeForTxnSnapshotReadConcern(kInsertCmdTargeted, kInsertCmdScatterGather);
 }
 
 TEST_F(ClusterInsertTest, SnapshotReadConcernWithAfterClusterTime) {
-    testSnapshotReadConcernWithAfterClusterTime(kInsertCmdTargeted, kInsertCmdScatterGather);
+    testTxnSnapshotReadConcernWithAfterClusterTime(kInsertCmdTargeted, kInsertCmdScatterGather);
 }
 
 }  // namespace
diff --git a/src/mongo/s/commands/cluster_update_test.cpp b/src/mongo/s/commands/cluster_update_test.cpp
index 53aee707522..9ee788c2a1e 100644
--- a/src/mongo/s/commands/cluster_update_test.cpp
+++ b/src/mongo/s/commands/cluster_update_test.cpp
@@ -77,11 +77,11 @@ TEST_F(ClusterUpdateTest, NoErrors) {
 }
 
 TEST_F(ClusterUpdateTest, AttachesAtClusterTimeForSnapshotReadConcern) {
-    testAttachesAtClusterTimeForSnapshotReadConcern(kUpdateCmdTargeted, kUpdateCmdScatterGather);
+    testAttachesAtClusterTimeForTxnSnapshotReadConcern(kUpdateCmdTargeted, kUpdateCmdScatterGather);
 }
 
 TEST_F(ClusterUpdateTest, SnapshotReadConcernWithAfterClusterTime) {
-    testSnapshotReadConcernWithAfterClusterTime(kUpdateCmdTargeted, kUpdateCmdScatterGather);
+    testTxnSnapshotReadConcernWithAfterClusterTime(kUpdateCmdTargeted, kUpdateCmdScatterGather);
 }
 
 }  // namespace
diff --git a/src/mongo/s/commands/strategy.cpp b/src/mongo/s/commands/strategy.cpp
index 39763246f92..c1c7a18e106 100644
--- a/src/mongo/s/commands/strategy.cpp
+++ b/src/mongo/s/commands/strategy.cpp
@@ -632,7 +632,7 @@ void runCommand(OperationContext* opCtx,
                     }
                     return latestKnownClusterTime.asTimestamp();
                 }();
-                readConcernArgs.setArgsAtClusterTimeForSnapshot(atClusterTime);
+                readConcernArgs.setArgsAtClusterTime(atClusterTime);
             }
 
             replyBuilder->reset();
@@ -822,7 +822,8 @@ void runCommand(OperationContext* opCtx,
 
                     abortGuard.dismiss();
                     continue;
-                } else if (!ReadConcernArgs::get(opCtx).wasAtClusterTimeSelected()) {
+                } else if (auto rc = ReadConcernArgs::get(opCtx);
+                           rc.getArgsAtClusterTime() && !rc.wasAtClusterTimeSelected()) {
                     // Non-transaction snapshot read. The client sent readConcern: {level:
                     // "snapshot", atClusterTime: T}, where T is older than
                     // minSnapshotHistoryWindowInSeconds, retrying won't succeed.