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

1 files changed, 20 insertions, 0 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