From 4ed0cd6be723c6709a7d1d9a72a5aa5916f6871d Mon Sep 17 00:00:00 2001
From: Eugene Kabanov <eugene.kabanov@status.im>
Date: Tue, 5 Mar 2024 18:34:53 +0200
Subject: [PATCH] Ensure that `OwnCancelSchedule` flag will not be removed from
 `wait()` and `withTimeout()`. (#519)

---
 chronos/internal/asyncfutures.nim |  8 +++++++-
 tests/testfut.nim                 | 26 ++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/chronos/internal/asyncfutures.nim b/chronos/internal/asyncfutures.nim
index de6deba..d084252 100644
--- a/chronos/internal/asyncfutures.nim
+++ b/chronos/internal/asyncfutures.nim
@@ -1410,6 +1410,8 @@ proc withTimeout*[T](fut: Future[T], timeout: Duration): Future[bool] {.
   var
     retFuture = newFuture[bool]("chronos.withTimeout",
                                 {FutureFlag.OwnCancelSchedule})
+      # We set `OwnCancelSchedule` flag, because we going to cancel `retFuture`
+      # manually at proper time.
     moment: Moment
     timer: TimerCallback
     timeouted = false
@@ -1536,6 +1538,8 @@ proc wait*[T](fut: Future[T], timeout = InfiniteDuration): Future[T] =
   ## should return, because it can't be cancelled too.
   var
     retFuture = newFuture[T]("chronos.wait()", {FutureFlag.OwnCancelSchedule})
+      # We set `OwnCancelSchedule` flag, because we going to cancel `retFuture`
+      # manually at proper time.
 
   waitImpl(fut, retFuture, timeout)
 
@@ -1677,6 +1681,8 @@ proc wait*(fut: InternalRaisesFuture, timeout = InfiniteDuration): auto =
     InternalRaisesFutureRaises = E.prepend(CancelledError, AsyncTimeoutError)
 
   let
-    retFuture = newFuture[T]("chronos.wait()", {FutureFlag.OwnCancelSchedule})
+    retFuture = newFuture[T]("chronos.wait()", {OwnCancelSchedule})
+      # We set `OwnCancelSchedule` flag, because we going to cancel `retFuture`
+      # manually at proper time.
 
   waitImpl(fut, retFuture, timeout)
diff --git a/tests/testfut.nim b/tests/testfut.nim
index fc2401d..aee3b15 100644
--- a/tests/testfut.nim
+++ b/tests/testfut.nim
@@ -1594,6 +1594,19 @@ suite "Future[T] behavior test suite":
     discard someFut.tryCancel()
     await someFut
 
+  asyncTest "wait() should allow cancellation test (depends on race())":
+    proc testFoo(): Future[bool] {.async.} =
+      let
+        resFut = sleepAsync(2.seconds).wait(3.seconds)
+        timeFut = sleepAsync(1.seconds)
+        cancelFut = cancelAndWait(resFut)
+      discard await race(cancelFut, timeFut)
+      if cancelFut.finished():
+        return (resFut.cancelled() and cancelFut.completed())
+      false
+
+    check (await testFoo()) == true
+
   asyncTest "withTimeout() cancellation undefined behavior test #1":
     proc testInnerFoo(fooFut: Future[void]): Future[TestFooConnection] {.
          async.} =
@@ -1654,6 +1667,19 @@ suite "Future[T] behavior test suite":
     discard someFut.tryCancel()
     await someFut
 
+  asyncTest "withTimeout() should allow cancellation test (depends on race())":
+    proc testFoo(): Future[bool] {.async.} =
+      let
+        resFut = sleepAsync(2.seconds).withTimeout(3.seconds)
+        timeFut = sleepAsync(1.seconds)
+        cancelFut = cancelAndWait(resFut)
+      discard await race(cancelFut, timeFut)
+      if cancelFut.finished():
+        return (resFut.cancelled() and cancelFut.completed())
+      false
+
+    check (await testFoo()) == true
+
   asyncTest "Cancellation behavior test":
     proc testInnerFoo(fooFut: Future[void]) {.async.} =
       await fooFut