SERVER-39146 Refactor Baton

Refactor the baton into regular and networking batons while also
cleaning up the basic baton implementation.

1 files changed, 19 insertions, 1 deletions

diff --git a/src/mongo/db/operation_context.cpp b/src/mongo/db/operation_context.cpp
index 11ca8b8f641..319b803288c 100644
--- a/src/mongo/db/operation_context.cpp
+++ b/src/mongo/db/operation_context.cpp
@@ -77,6 +77,8 @@ MONGO_FAIL_POINT_DEFINE(maxTimeNeverTimeOut);
 // inclusive.
 MONGO_FAIL_POINT_DEFINE(checkForInterruptFail);
 
+const auto kNoWaiterThread = stdx::thread::id();
+
 }  // namespace
 
 OperationContext::OperationContext(Client* client, unsigned int opId)
@@ -261,6 +263,7 @@ StatusWith<stdx::cv_status> OperationContext::waitForConditionOrInterruptNoAsser
         stdx::lock_guard<Client> clientLock(*getClient());
         invariant(!_waitMutex);
         invariant(!_waitCV);
+        invariant(_waitThread == kNoWaiterThread);
         invariant(0 == _numKillers);
 
         // This interrupt check must be done while holding the client lock, so as not to race with a
@@ -271,6 +274,7 @@ StatusWith<stdx::cv_status> OperationContext::waitForConditionOrInterruptNoAsser
         }
         _waitMutex = m.mutex();
         _waitCV = &cv;
+        _waitThread = stdx::this_thread::get_id();
     }
 
     // If the maxTimeNeverTimeOut failpoint is set, behave as though the operation's deadline does
@@ -302,6 +306,7 @@ StatusWith<stdx::cv_status> OperationContext::waitForConditionOrInterruptNoAsser
         if (0 == _numKillers) {
             _waitMutex = nullptr;
             _waitCV = nullptr;
+            _waitThread = kNoWaiterThread;
             return true;
         }
         return false;
@@ -328,7 +333,20 @@ StatusWith<stdx::cv_status> OperationContext::waitForConditionOrInterruptNoAsser
 void OperationContext::markKilled(ErrorCodes::Error killCode) {
     invariant(killCode != ErrorCodes::OK);
     stdx::unique_lock<stdx::mutex> lkWaitMutex;
-    if (_waitMutex) {
+
+    // If we have a _waitMutex, it means this opCtx is currently blocked in
+    // waitForConditionOrInterrupt.
+    //
+    // From there, we also know which thread is actually doing that waiting (it's recorded in
+    // _waitThread).  If that thread isn't our thread, it's necessary to do the regular numKillers
+    // song and dance mentioned in waitForConditionOrInterrupt.
+    //
+    // If it is our thread, we know that we're currently inside that call to
+    // waitForConditionOrInterrupt and are being invoked by a callback run from Baton->run.  And
+    // that means we don't need to deal with the waitMutex and waitCV (because we're running
+    // callbacks, which means run is returning, which means we'll be checking _killCode in the near
+    // future).
+    if (_waitMutex && stdx::this_thread::get_id() != _waitThread) {
         invariant(++_numKillers > 0);
         getClient()->unlock();
         ON_BLOCK_EXIT([this] {