SERVER-44022 Fix race in PrepareConflictTracker and speed up happy path

3 files changed, 16 insertions, 17 deletions

diff --git a/src/mongo/db/prepare_conflict_tracker.cpp b/src/mongo/db/prepare_conflict_tracker.cpp
index ab45568f9dd..7dc8123e282 100644
--- a/src/mongo/db/prepare_conflict_tracker.cpp
+++ b/src/mongo/db/prepare_conflict_tracker.cpp
@@ -40,16 +40,14 @@ bool PrepareConflictTracker::isWaitingOnPrepareConflict() const {
 }
 
 void PrepareConflictTracker::beginPrepareConflict(OperationContext* opCtx) {
-    invariant(_prepareConflictStartTime == 0);
-    _prepareConflictStartTime = opCtx->getServiceContext()->getTickSource()->getTicks();
-
     // Implies that the current read operation is blocked on a prepared transaction.
     _waitOnPrepareConflict.store(true);
+    invariant(_prepareConflictStartTime == 0);
+    _prepareConflictStartTime = opCtx->getServiceContext()->getTickSource()->getTicks();
 }
 
 void PrepareConflictTracker::endPrepareConflict(OperationContext* opCtx) {
-    // This function is called regardless whether there was a prepare conflict.
-    if (_prepareConflictStartTime) {
+    if (_waitOnPrepareConflict.load()) {
         auto tickSource = opCtx->getServiceContext()->getTickSource();
         auto curTick = tickSource->getTicks();
 
@@ -61,16 +59,17 @@ void PrepareConflictTracker::endPrepareConflict(OperationContext* opCtx) {
 
         auto curConflictDuration =
             tickSource->ticksTo<Microseconds>(curTick - _prepareConflictStartTime);
-        _prepareConflictDuration += curConflictDuration;
-    }
-    _prepareConflictStartTime = 0;
+        _prepareConflictDuration.store(_prepareConflictDuration.load() + curConflictDuration);
+        _prepareConflictStartTime = 0;
 
-    // Implies that the current read operation is not blocked on a prepared transaction.
-    _waitOnPrepareConflict.store(false);
+        // Implies that the current read operation is not blocked on a prepared transaction.
+        _waitOnPrepareConflict.store(false);
+    }
+    invariant(_prepareConflictStartTime == 0);
 }
 
 Microseconds PrepareConflictTracker::getPrepareConflictDuration() {
-    return _prepareConflictDuration;
+    return _prepareConflictDuration.load();
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/prepare_conflict_tracker.h b/src/mongo/db/prepare_conflict_tracker.h
index 01fc146d675..2d03fbce2ca 100644
--- a/src/mongo/db/prepare_conflict_tracker.h
+++ b/src/mongo/db/prepare_conflict_tracker.h
@@ -86,7 +86,7 @@ private:
     /**
      * Stores the total amount of time spent blocked on prepare read conflicts.
      */
-    Microseconds _prepareConflictDuration{0};
+    AtomicWord<Microseconds> _prepareConflictDuration{Microseconds(0)};
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_prepare_conflict.h b/src/mongo/db/storage/wiredtiger/wiredtiger_prepare_conflict.h
index 5212a0b19a4..65f2996d4a0 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_prepare_conflict.h
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_prepare_conflict.h
@@ -66,17 +66,17 @@ template <typename F>
 int wiredTigerPrepareConflictRetry(OperationContext* opCtx, F&& f) {
     invariant(opCtx);
 
-    auto recoveryUnit = WiredTigerRecoveryUnit::get(opCtx);
-    int attempts = 1;
-    // If we return from this function, we have either returned successfully or we've returned an
-    // error other than WT_PREPARE_CONFLICT. Reset PrepareConflictTracker accordingly.
-    ON_BLOCK_EXIT([opCtx] { PrepareConflictTracker::get(opCtx).endPrepareConflict(opCtx); });
     // If the failpoint is enabled, don't call the function, just simulate a conflict.
     int ret = MONGO_unlikely(WTPrepareConflictForReads.shouldFail()) ? WT_PREPARE_CONFLICT
                                                                      : WT_READ_CHECK(f());
     if (ret != WT_PREPARE_CONFLICT)
         return ret;
 
+    auto recoveryUnit = WiredTigerRecoveryUnit::get(opCtx);
+    int attempts = 1;
+    // If we return from this function, we have either returned successfully or we've returned an
+    // error other than WT_PREPARE_CONFLICT. Reset PrepareConflictTracker accordingly.
+    ON_BLOCK_EXIT([opCtx] { PrepareConflictTracker::get(opCtx).endPrepareConflict(opCtx); });
     PrepareConflictTracker::get(opCtx).beginPrepareConflict(opCtx);
 
     auto client = opCtx->getClient();