SERVER-40053 Add an RAII type to change and restore the timestamp read source

5 files changed, 110 insertions, 18 deletions

diff --git a/src/mongo/db/catalog_raii.cpp b/src/mongo/db/catalog_raii.cpp
index b3fafde57a8..d67b2cbc520 100644
--- a/src/mongo/db/catalog_raii.cpp
+++ b/src/mongo/db/catalog_raii.cpp
@@ -193,4 +193,24 @@ ConcealUUIDCatalogChangesBlock::~ConcealUUIDCatalogChangesBlock() {
     UUIDCatalog::get(_opCtx).onOpenCatalog(_opCtx);
 }
 
+ReadSourceScope::ReadSourceScope(OperationContext* opCtx)
+    : _opCtx(opCtx), _originalReadSource(opCtx->recoveryUnit()->getTimestampReadSource()) {
+
+    if (_originalReadSource == RecoveryUnit::ReadSource::kProvided) {
+        _originalReadTimestamp = *_opCtx->recoveryUnit()->getPointInTimeReadTimestamp();
+    }
+
+    _opCtx->recoveryUnit()->abandonSnapshot();
+    _opCtx->recoveryUnit()->setTimestampReadSource(RecoveryUnit::ReadSource::kUnset);
+}
+
+ReadSourceScope::~ReadSourceScope() {
+    _opCtx->recoveryUnit()->abandonSnapshot();
+    if (_originalReadSource == RecoveryUnit::ReadSource::kProvided) {
+        _opCtx->recoveryUnit()->setTimestampReadSource(_originalReadSource, _originalReadTimestamp);
+    } else {
+        _opCtx->recoveryUnit()->setTimestampReadSource(_originalReadSource);
+    }
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/catalog_raii.h b/src/mongo/db/catalog_raii.h
index 823d923b6bb..730ba4614e8 100644
--- a/src/mongo/db/catalog_raii.h
+++ b/src/mongo/db/catalog_raii.h
@@ -214,4 +214,21 @@ private:
     OperationContext* _opCtx;
 };
 
+/**
+ * RAII type to set and restore the timestamp read source on the recovery unit.
+ *
+ * Snapshot is abandoned in constructor and destructor, so it can only be used before
+ * the recovery unit becomes active or when the existing snapshot is no longer needed.
+ */
+class ReadSourceScope {
+public:
+    ReadSourceScope(OperationContext* opCtx);
+    ~ReadSourceScope();
+
+private:
+    OperationContext* _opCtx;
+    RecoveryUnit::ReadSource _originalReadSource;
+    Timestamp _originalReadTimestamp;
+};
+
 }  // namespace mongo
diff --git a/src/mongo/db/catalog_raii_test.cpp b/src/mongo/db/catalog_raii_test.cpp
index b8f34187e44..a6c90d830ae 100644
--- a/src/mongo/db/catalog_raii_test.cpp
+++ b/src/mongo/db/catalog_raii_test.cpp
@@ -33,12 +33,14 @@
 
 #include <string>
 
+#include "boost/optional/optional_io.hpp"
 #include "mongo/db/catalog/database_holder_mock.h"
 #include "mongo/db/catalog_raii.h"
 #include "mongo/db/client.h"
 #include "mongo/db/concurrency/d_concurrency.h"
 #include "mongo/db/concurrency/lock_state.h"
 #include "mongo/db/service_context_test_fixture.h"
+#include "mongo/db/storage/recovery_unit_noop.h"
 #include "mongo/unittest/unittest.h"
 #include "mongo/util/log.h"
 #include "mongo/util/time_support.h"
@@ -212,5 +214,60 @@ TEST_F(CatalogRAIITestFixture, AutoGetCollectionDeadlineMin) {
         Milliseconds(0));
 }
 
+using ReadSource = RecoveryUnit::ReadSource;
+
+class RecoveryUnitMock : public RecoveryUnitNoop {
+public:
+    void setTimestampReadSource(ReadSource source,
+                                boost::optional<Timestamp> provided = boost::none) override {
+        _source = source;
+        _timestamp = provided;
+    }
+    ReadSource getTimestampReadSource() const override {
+        return _source;
+    };
+    boost::optional<Timestamp> getPointInTimeReadTimestamp() override {
+        return _timestamp;
+    }
+
+private:
+    ReadSource _source = ReadSource::kUnset;
+    boost::optional<Timestamp> _timestamp;
+};
+
+class ReadSourceScopeTest : public ServiceContextTest {
+public:
+    OperationContext* opCtx() {
+        return _opCtx.get();
+    }
+
+protected:
+    void setUp() override;
+
+    ServiceContext::UniqueOperationContext _opCtx;
+};
+
+void ReadSourceScopeTest::setUp() {
+    _opCtx = getClient()->makeOperationContext();
+    _opCtx->setRecoveryUnit(stdx::make_unique<RecoveryUnitMock>(),
+                            WriteUnitOfWork::RecoveryUnitState::kNotInUnitOfWork);
+}
+
+TEST_F(ReadSourceScopeTest, RestoreReadSource) {
+    opCtx()->recoveryUnit()->setTimestampReadSource(ReadSource::kProvided, Timestamp(1, 2));
+    ASSERT_EQ(opCtx()->recoveryUnit()->getTimestampReadSource(), ReadSource::kProvided);
+    ASSERT_EQ(opCtx()->recoveryUnit()->getPointInTimeReadTimestamp(), Timestamp(1, 2));
+    {
+        ReadSourceScope scope(opCtx());
+        ASSERT_EQ(opCtx()->recoveryUnit()->getTimestampReadSource(), ReadSource::kUnset);
+
+        opCtx()->recoveryUnit()->setTimestampReadSource(ReadSource::kLastApplied);
+        ASSERT_EQ(opCtx()->recoveryUnit()->getTimestampReadSource(), ReadSource::kLastApplied);
+        ASSERT_EQ(opCtx()->recoveryUnit()->getPointInTimeReadTimestamp(), boost::none);
+    }
+    ASSERT_EQ(opCtx()->recoveryUnit()->getTimestampReadSource(), ReadSource::kProvided);
+    ASSERT_EQ(opCtx()->recoveryUnit()->getPointInTimeReadTimestamp(), Timestamp(1, 2));
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/repl/transaction_oplog_application.cpp b/src/mongo/db/repl/transaction_oplog_application.cpp
index d5076312631..fea5ff44db4 100644
--- a/src/mongo/db/repl/transaction_oplog_application.cpp
+++ b/src/mongo/db/repl/transaction_oplog_application.cpp
@@ -31,6 +31,7 @@
 
 #include "mongo/db/repl/transaction_oplog_application.h"
 
+#include "mongo/db/catalog_raii.h"
 #include "mongo/db/commands/txn_cmds_gen.h"
 #include "mongo/db/repl/apply_ops.h"
 #include "mongo/db/session_catalog_mongod.h"
@@ -50,20 +51,21 @@ Status _applyTransactionFromOplogChain(OperationContext* opCtx,
     invariant(mode == repl::OplogApplication::Mode::kRecovering ||
               mode == repl::OplogApplication::Mode::kInitialSync);
 
-    // Since the TransactionHistoryIterator uses DBDirectClient, it cannot come with snapshot
-    // isolation.
-    invariant(!opCtx->recoveryUnit()->getPointInTimeReadTimestamp());
+    BSONObj prepareCmd;
+    {
+        // Traverse the oplog chain with its own snapshot and read timestamp.
+        ReadSourceScope readSourceScope(opCtx);
 
-    // Get the corresponding prepareTransaction oplog entry.
-    const auto prepareOpTime = entry.getPrevWriteOpTimeInTransaction();
-    invariant(prepareOpTime);
-    TransactionHistoryIterator iter(prepareOpTime.get());
-    invariant(iter.hasNext());
-    const auto prepareOplogEntry = iter.next(opCtx);
-
-    // Transform prepare command into a normal applyOps command.
-    const auto prepareCmd = prepareOplogEntry.getOperationToApply().removeField("prepare");
+        // Get the corresponding prepareTransaction oplog entry.
+        const auto prepareOpTime = entry.getPrevWriteOpTimeInTransaction();
+        invariant(prepareOpTime);
+        TransactionHistoryIterator iter(prepareOpTime.get());
+        invariant(iter.hasNext());
+        const auto prepareOplogEntry = iter.next(opCtx);
 
+        // Transform prepare command into a normal applyOps command.
+        prepareCmd = prepareOplogEntry.getOperationToApply().removeField("prepare");
+    }
     BSONObjBuilder resultWeDontCareAbout;
     return applyOps(
         opCtx, entry.getNss().db().toString(), prepareCmd, mode, &resultWeDontCareAbout);
diff --git a/src/mongo/db/transaction_participant.cpp b/src/mongo/db/transaction_participant.cpp
index 86ec9d4fa1e..fe4789f2a33 100644
--- a/src/mongo/db/transaction_participant.cpp
+++ b/src/mongo/db/transaction_participant.cpp
@@ -107,12 +107,8 @@ struct ActiveTransactionHistory {
 
 ActiveTransactionHistory fetchActiveTransactionHistory(OperationContext* opCtx,
                                                        const LogicalSessionId& lsid) {
-    // Since we are using DBDirectClient to read the transactions table and the oplog, we should
-    // never be reading from a snapshot, but directly from what is the latest on disk. This
-    // invariant guards against programming errors where the default read concern on the
-    // OperationContext could have been changed to something other than 'local'.
-    invariant(repl::ReadConcernArgs::get(opCtx).getLevel() ==
-              repl::ReadConcernLevel::kLocalReadConcern);
+    // Restore the current timestamp read source after fetching transaction history.
+    ReadSourceScope readSourceScope(opCtx);
 
     ActiveTransactionHistory result;