SERVER-55760 Fix data races in Cancelation tests for AsyncTry

2 files changed, 69 insertions, 9 deletions

diff --git a/src/mongo/util/future_util.h b/src/mongo/util/future_util.h
index 0d5e45870ec..9fdc53625e6 100644
--- a/src/mongo/util/future_util.h
+++ b/src/mongo/util/future_util.h
@@ -174,8 +174,12 @@ private:
                             executor->sleepFor(delay.getNext(), cancelToken)
                                 .getAsync([this, self, resultPromise = std::move(resultPromise)](
                                               Status s) mutable {
+                                    // Prevent another loop iteration when cancellation happens
+                                    // after loop body
                                     if (s.isOK()) {
                                         runImpl(std::move(resultPromise));
+                                    } else {
+                                        resultPromise.setError(std::move(s));
                                     }
                                 });
                         }
diff --git a/src/mongo/util/future_util_test.cpp b/src/mongo/util/future_util_test.cpp
index d19d8044e4e..3377ea210d2 100644
--- a/src/mongo/util/future_util_test.cpp
+++ b/src/mongo/util/future_util_test.cpp
@@ -423,27 +423,83 @@ TEST_F(AsyncTryUntilTest, AsyncTryUntilCanBeCanceled) {
     ASSERT_EQ(resultFut.getNoThrow(), kCanceledStatus);
 }
 
-TEST_F(AsyncTryUntilTest, AsyncTryUntilWithDelayCanBeCanceled) {
+TEST_F(AsyncTryUntilTest, AsyncTryUntilWithDelayCanBeCanceledWhileLoopBodyIsExecuting) {
     CancellationSource cancelSource;
-    auto resultFut = AsyncTry([] {})
+    unittest::Barrier barrierBeforeCancel{2}, barrierAfterCancel{2};
+    int timesRanCallback = 0;
+    // Arbitrary delay used, enforce only one loop body execution with timesRanCallback
+    auto resultFut = AsyncTry([&] {
+                         timesRanCallback += 1;
+                         barrierBeforeCancel.countDownAndWait();
+                         barrierAfterCancel.countDownAndWait();
+                     })
                          .until([](Status) { return false; })
-                         .withDelayBetweenIterations(Hours(1000))
+                         .withDelayBetweenIterations(Milliseconds(10))
+                         .on(executor(), cancelSource.token());
+    // Enforce cancellation during loop body execution
+    barrierBeforeCancel.countDownAndWait();
+    cancelSource.cancel();
+    barrierAfterCancel.countDownAndWait();
+    ASSERT_EQ(resultFut.getNoThrow(), kCanceledStatus);
+    ASSERT_EQ(timesRanCallback, 1);
+}
+
+TEST_F(AsyncTryUntilTest, AsyncTryUntilWithDelayCanBeCanceledAfterLoopBodyIsDoneExecuting) {
+    CancellationSource cancelSource;
+    unittest::Barrier barrierBeforeCancel{2}, barrierAfterCancel{2};
+    int timesRanCallback = 0;
+    auto resultFut = AsyncTry([&] { timesRanCallback += 1; })
+                         .until([&](Status) {
+                             barrierBeforeCancel.countDownAndWait();
+                             barrierAfterCancel.countDownAndWait();
+                             return false;
+                         })
+                         .withDelayBetweenIterations(Milliseconds(10))
                          .on(executor(), cancelSource.token());
-    // Since the "until" condition is false, and the delay between iterations is very long, the only
-    // way this test should pass without hanging is if the future produced by TaskExecutor::sleepFor
-    // is resolved and set with ErrorCodes::CallbackCanceled well _before_ the deadline.
+    // Enforce cancellation after loop body executes once
+    barrierBeforeCancel.countDownAndWait();
     cancelSource.cancel();
+    barrierAfterCancel.countDownAndWait();
     ASSERT_EQ(resultFut.getNoThrow(), kCanceledStatus);
+    ASSERT_EQ(timesRanCallback, 1);
 }
 
-TEST_F(AsyncTryUntilTest, AsyncTryUntilWithBackoffCanBeCanceled) {
+TEST_F(AsyncTryUntilTest, AsyncTryUntilWithBackoffCanBeCanceledWhileLoopBodyIsExecuting) {
     CancellationSource cancelSource;
-    auto resultFut = AsyncTry([] {})
+    unittest::Barrier barrierBeforeCancel{2}, barrierAfterCancel{2};
+    int timesRanCallback = 0;
+    auto resultFut = AsyncTry([&] {
+                         timesRanCallback += 1;
+                         barrierBeforeCancel.countDownAndWait();
+                         barrierAfterCancel.countDownAndWait();
+                     })
                          .until([](Status) { return false; })
-                         .withBackoffBetweenIterations(TestBackoff{Seconds(10000000)})
+                         .withBackoffBetweenIterations(TestBackoff{Milliseconds(10)})
+                         .on(executor(), cancelSource.token());
+    barrierBeforeCancel.countDownAndWait();
+    cancelSource.cancel();
+    barrierAfterCancel.countDownAndWait();
+    ASSERT_EQ(resultFut.getNoThrow(), kCanceledStatus);
+    ASSERT_EQ(timesRanCallback, 1);
+}
+
+TEST_F(AsyncTryUntilTest, AsyncTryUntilWithBackoffCanBeCanceledAfterLoopBodyIsDoneExecuting) {
+    CancellationSource cancelSource;
+    unittest::Barrier barrierBeforeCancel{2}, barrierAfterCancel{2};
+    int timesRanCallback = 0;
+    auto resultFut = AsyncTry([&] { timesRanCallback += 1; })
+                         .until([&](Status) {
+                             barrierBeforeCancel.countDownAndWait();
+                             barrierAfterCancel.countDownAndWait();
+                             return false;
+                         })
+                         .withBackoffBetweenIterations(TestBackoff{Milliseconds(10)})
                          .on(executor(), cancelSource.token());
+    barrierBeforeCancel.countDownAndWait();
     cancelSource.cancel();
+    barrierAfterCancel.countDownAndWait();
     ASSERT_EQ(resultFut.getNoThrow(), kCanceledStatus);
+    ASSERT_EQ(timesRanCallback, 1);
 }
 
 TEST_F(AsyncTryUntilTest, CanceledTryUntilLoopDoesNotExecuteIfAlreadyCanceled) {