SERVER-50660 Integrate CancelationTokens with OperationContext

5 files changed, 49 insertions, 13 deletions

diff --git a/src/mongo/util/cancelation.h b/src/mongo/util/cancelation.h
index cd43e1b17e1..a527b11ef1e 100644
--- a/src/mongo/util/cancelation.h
+++ b/src/mongo/util/cancelation.h
@@ -29,14 +29,18 @@
 #pragma once
 
 #include "mongo/util/future.h"
+#include "mongo/util/static_immortal.h"
 
 namespace mongo {
 
 namespace detail {
 
-static const inline Status kCancelNeverCalledOnSourceError{
-    ErrorCodes::CallbackCanceled,
-    "Cancel was never called on the CancelationSource for this token."};
+inline Status getCancelNeverCalledOnSourceError() {
+    static const StaticImmortal<Status> cancelNeverCalledOnSourceError{
+        ErrorCodes::CallbackCanceled,
+        "Cancel was never called on the CancelationSource for this token."};
+    return *cancelNeverCalledOnSourceError;
+}
 
 /**
  * Holds the main state shared between CancelationSource/CancelationToken.
@@ -72,7 +76,7 @@ public:
     void dismiss() {
         State precondition{State::kInit};
         if (_state.compareAndSwap(&precondition, State::kDismissed)) {
-            _cancelationPromise.setError(kCancelNeverCalledOnSourceError);
+            _cancelationPromise.setError(getCancelNeverCalledOnSourceError());
         }
     }
 
diff --git a/src/mongo/util/cancelation_test.cpp b/src/mongo/util/cancelation_test.cpp
index 97829dead4f..7536ed5b320 100644
--- a/src/mongo/util/cancelation_test.cpp
+++ b/src/mongo/util/cancelation_test.cpp
@@ -41,7 +41,7 @@ TEST(CancelTest, CancelSourceDestructorDoesNotCauseCancelation) {
     {
         CancelationSource source;
         source.token().onCancel().unsafeToInlineFuture().getAsync([&ran](Status s) mutable {
-            ASSERT_EQ(s, detail::kCancelNeverCalledOnSourceError);
+            ASSERT_EQ(s, detail::getCancelNeverCalledOnSourceError());
             ran = true;
         });
     }
@@ -181,7 +181,7 @@ TEST(CancelTest, CancelTokenRemembersNotCanceledForNotCanceledSourceEvenAfterSou
     }();
     ASSERT_FALSE(token.isCanceled());
     ASSERT_TRUE(token.onCancel().isReady());
-    ASSERT_EQ(token.onCancel().getNoThrow(), detail::kCancelNeverCalledOnSourceError);
+    ASSERT_EQ(token.onCancel().getNoThrow(), detail::getCancelNeverCalledOnSourceError());
 }
 
 TEST(CancelTest, UncancelableTokenReturnsFalseForIsCanceled) {
@@ -193,7 +193,7 @@ TEST(CancelTest, UncancelableTokenNeverRunsCallbacks) {
     auto token = CancelationToken::uncancelable();
     auto cancelFuture = token.onCancel();
     ASSERT_TRUE(cancelFuture.isReady());
-    ASSERT_EQ(cancelFuture.getNoThrow(), detail::kCancelNeverCalledOnSourceError);
+    ASSERT_EQ(cancelFuture.getNoThrow(), detail::getCancelNeverCalledOnSourceError());
 }
 
 TEST(CancelTest, UncancelableTokenReturnsFalseForIsCancelable) {
diff --git a/src/mongo/util/future_test_utils.h b/src/mongo/util/future_test_utils.h
index f9efe3cf1e5..ecf44d67470 100644
--- a/src/mongo/util/future_test_utils.h
+++ b/src/mongo/util/future_test_utils.h
@@ -61,7 +61,8 @@ class DummyInterruptible final : public Interruptible {
         MONGO_UNREACHABLE;
     }
     Status checkForInterruptNoAssert() noexcept override {
-        MONGO_UNREACHABLE;
+        // Must be implemented because it's called by Interruptible::waitForConditionOrInterrupt.
+        return Status::OK();
     }
     IgnoreInterruptsState pushIgnoreInterrupts() override {
         MONGO_UNREACHABLE;
diff --git a/src/mongo/util/interruptible.h b/src/mongo/util/interruptible.h
index d1b8981adf3..a492a60a87f 100644
--- a/src/mongo/util/interruptible.h
+++ b/src/mongo/util/interruptible.h
@@ -41,6 +41,16 @@
 
 namespace mongo {
 
+namespace interruptible_detail {
+// Helper to release a lock, call a callable, and then reacquire the lock.
+template <typename Callable>
+auto doWithoutLock(BasicLockableAdapter m, Callable&& callable) {
+    m.unlock();
+    ON_BLOCK_EXIT([&] { m.lock(); });
+    return callable();
+}
+}  // namespace interruptible_detail
+
 /**
  * A type which can be used to wait on condition variables with a level triggered one-way interrupt.
  * I.e. after the interrupt is triggered (via some non-public api call) subsequent calls to
@@ -380,6 +390,21 @@ public:
             advanceWaitTimer();
         });
 
+        auto handleInterruptAndAssert = [&](Status status, WakeSpeed speed) {
+            _onWake(latchName, WakeReason::kInterrupt, speed);
+            iassert(std::move(status));
+        };
+
+        auto checkForInterruptWithoutLockAndAssert = [&](WakeSpeed speed) {
+            // We drop the lock before checking for interrupt since checkForInterruptNoAssert can
+            // sometimes try to reacquire the same lock.
+            if (auto status = interruptible_detail::doWithoutLock(
+                    m, [&] { return checkForInterruptNoAssert(); });
+                !status.isOK()) {
+                handleInterruptAndAssert(status, speed);
+            }
+        };
+
         auto waitUntil = [&](Date_t deadline, WakeSpeed speed) -> boost::optional<WakeReason> {
             // If the result of waitForConditionOrInterruptNoAssertUntil() is non-spurious, return
             // a WakeReason. Otherwise, return boost::none
@@ -394,10 +419,13 @@ public:
             }
 
             if (!swResult.isOK()) {
-                _onWake(latchName, WakeReason::kInterrupt, speed);
-                iassert(std::move(swResult));
+                handleInterruptAndAssert(swResult.getStatus(), speed);
             }
 
+            // Check if an interrupt occurred while waiting.
+            checkForInterruptWithoutLockAndAssert(speed);
+
+            // Check the predicate after re-acquiring the lock.
             if (pred()) {
                 _onWake(latchName, WakeReason::kPredicate, speed);
                 return WakeReason::kPredicate;
@@ -412,14 +440,16 @@ public:
         };
 
         auto waitUntilNonSpurious = [&](Date_t deadline, WakeSpeed speed) -> WakeReason {
-            // Check waitUntil() in a loop until it says it has a genuine WakeReason
+            // Check for interrupt before waiting.
+            checkForInterruptWithoutLockAndAssert(speed);
 
+            // Check the predicate after re-acquiring the lock and before waiting.
             if (pred()) {
-                // Check for the predicate first, just in case
                 _onWake(latchName, WakeReason::kPredicate, speed);
                 return WakeReason::kPredicate;
             }
 
+            // Check waitUntil() in a loop until it says it has a genuine WakeReason
             auto maybeWakeReason = waitUntil(deadline, speed);
             while (!maybeWakeReason) {
                 maybeWakeReason = waitUntil(deadline, speed);
diff --git a/src/mongo/util/interruptible_test.cpp b/src/mongo/util/interruptible_test.cpp
index be29ccb88a7..5a6de7c4916 100644
--- a/src/mongo/util/interruptible_test.cpp
+++ b/src/mongo/util/interruptible_test.cpp
@@ -41,7 +41,8 @@ class DummyInterruptible final : public Interruptible {
         MONGO_UNREACHABLE;
     }
     Status checkForInterruptNoAssert() noexcept override {
-        MONGO_UNREACHABLE;
+        // Must be implemented because it's called by Interruptible::waitForConditionOrInterrupt.
+        return Status::OK();
     }
     IgnoreInterruptsState pushIgnoreInterrupts() override {
         MONGO_UNREACHABLE;