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

1 files changed, 18 insertions, 10 deletions

diff --git a/src/mongo/db/repl/read_concern_args.h b/src/mongo/db/repl/read_concern_args.h
index ba6004d13bf..2c7aca60282 100644
--- a/src/mongo/db/repl/read_concern_args.h
+++ b/src/mongo/db/repl/read_concern_args.h
@@ -182,13 +182,21 @@ public:
      * Set atClusterTime, clear afterClusterTime. The BSON representation becomes
      * {level: "snapshot", atClusterTime: <ts>}.
      */
-    void setArgsAtClusterTimeForSnapshot(Timestamp ts) {
+    void setArgsAtClusterTime(Timestamp ts) {
         invariant(_level && _level == ReadConcernLevel::kSnapshotReadConcern);
         // Only overwrite a server-selected atClusterTime, not user-supplied.
-        invariant(_atClusterTime.is_initialized() == _atClusterTimeSelected);
-        _afterClusterTime = boost::none;
-        _atClusterTime = LogicalTime(ts);
-        _atClusterTimeSelected = true;
+        invariant(!_clientAtClusterTime);
+        _computedAfterClusterTime = boost::none;
+        _computedAtClusterTime = LogicalTime(ts);
+    }
+
+    /**
+     * If we have selected an atClusterTime for non-transaction snapshot reads, clear it and restore
+     * the atClusterTime and afterClusterTime passed by the client.
+     */
+    void clearArgsAtClusterTime() {
+        _computedAfterClusterTime = _clientAfterClusterTime;
+        _computedAtClusterTime = _clientAtClusterTime;
     }
 
     /**
@@ -196,7 +204,7 @@ public:
      * function returns false if the atClusterTime was specified by the client.
      */
     bool wasAtClusterTimeSelected() const {
-        return _atClusterTimeSelected;
+        return _computedAtClusterTime && _computedAtClusterTime != _clientAtClusterTime;
     }
 
 private:
@@ -213,11 +221,13 @@ private:
     /**
      *  Read data after cluster-wide cluster time.
      */
-    boost::optional<LogicalTime> _afterClusterTime;
+    boost::optional<LogicalTime> _clientAfterClusterTime;
+    boost::optional<LogicalTime> _computedAfterClusterTime;
     /**
      * Read data at a particular cluster time.
      */
-    boost::optional<LogicalTime> _atClusterTime;
+    boost::optional<LogicalTime> _clientAtClusterTime;
+    boost::optional<LogicalTime> _computedAtClusterTime;
     boost::optional<ReadConcernLevel> _level;
 
     /**
@@ -233,8 +243,6 @@ private:
     bool _specified;
 
     ReadWriteConcernProvenance _provenance;
-
-    bool _atClusterTimeSelected = false;
 };
 
 }  // namespace repl