From f8937c0ae6a1b6afd3d0e4e4dae2afca53cf7390 Mon Sep 17 00:00:00 2001
From: Matthew Russotto <matthew.russotto@mongodb.com>
Date: Wed, 10 May 2023 14:24:57 -0400
Subject: SERVER-76987 Make sure local awaitData notifications are robust to
 WriteConflictExceptions.

---
 src/mongo/db/query/plan_insert_listener.cpp |  1 +
 src/mongo/db/query/plan_insert_listener.h   | 19 ++++++++++++++++---
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/mongo/db/query/plan_insert_listener.cpp b/src/mongo/db/query/plan_insert_listener.cpp
index df947c2e725..35b1729b093 100644
--- a/src/mongo/db/query/plan_insert_listener.cpp
+++ b/src/mongo/db/query/plan_insert_listener.cpp
@@ -123,6 +123,7 @@ void waitForInserts(OperationContext* opCtx,
             planExecutorHangWhileYieldedInWaitForInserts.pauseWhileSet();
         }
     });
+    notifier->doneWaiting(opCtx);
 
     uassertStatusOK(yieldResult);
 }
diff --git a/src/mongo/db/query/plan_insert_listener.h b/src/mongo/db/query/plan_insert_listener.h
index ae6e47b7232..054089fa65d 100644
--- a/src/mongo/db/query/plan_insert_listener.h
+++ b/src/mongo/db/query/plan_insert_listener.h
@@ -49,7 +49,13 @@ public:
     // Performs the necessary work needed for waiting. Should be called prior calling waitUntil().
     virtual void prepareForWait(OperationContext* opCtx) = 0;
 
-    // Blocks the caller until an insert event is fired or the deadline is hit.
+    // Performs any necessary steps after waiting. Should be called after waitUntil().
+    // After calling doneWaiting, the caller must attempt to read the data waited for before
+    // calling prepareForWait and waitUntil again, or a spurious wait may occur.
+    virtual void doneWaiting(OperationContext* opCtx) = 0;
+
+    // Blocks the caller until an insert event is fired or the deadline is hit.  Must be robust
+    // to being called multiple times without an intervening read.
     virtual void waitUntil(OperationContext* opCtx, Date_t deadline) = 0;
 };
 
@@ -61,17 +67,22 @@ public:
 
     void prepareForWait(OperationContext* opCtx) final {
         invariant(_notifier);
+        _currentVersion = _notifier->getVersion();
     }
 
     void waitUntil(OperationContext* opCtx, Date_t deadline) final {
-        auto currentVersion = _notifier->getVersion();
         _notifier->waitUntil(opCtx, _lastEOFVersion, deadline);
-        _lastEOFVersion = currentVersion;
+    }
+
+    void doneWaiting(OperationContext* opCtx) final {
+        _lastEOFVersion = _currentVersion;
     }
 
 private:
     std::shared_ptr<CappedInsertNotifier> _notifier;
     uint64_t _lastEOFVersion = ~uint64_t(0);
+    // This will be initialized by prepareForWait.
+    uint64_t _currentVersion;
 };
 
 // Class used to notify listeners on majority committed point advancement events.
@@ -103,6 +114,8 @@ public:
         });
     }
 
+    void doneWaiting(OperationContext* opCtx) final {}
+
 private:
     repl::OpTime _opTimeToBeMajorityCommitted;
 };
-- 
cgit v1.2.1