From 493cb1dbfd8afb38365b685f2fa3da9b2be09409 Mon Sep 17 00:00:00 2001
From: cheatfate <eugene.kabanov@status.im>
Date: Tue, 17 Nov 2020 11:59:02 +0200
Subject: [PATCH] Fix behavior of wait() and withTimeout() calls to cancel and
 wait for result of cancelled Future[T]. Add tests.

---
 chronos/asyncloop.nim | 55 ++++++++++++++++++++++++-------------------
 tests/testfut.nim     | 52 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 24 deletions(-)

diff --git a/chronos/asyncloop.nim b/chronos/asyncloop.nim
index 41d495c..d553e77 100644
--- a/chronos/asyncloop.nim
+++ b/chronos/asyncloop.nim
@@ -847,19 +847,23 @@ proc withTimeout*[T](fut: Future[T], timeout: Duration): Future[bool] =
   var retFuture = newFuture[bool]("chronos.`withTimeout`")
   var moment: Moment
   var timer: TimerCallback
+  var cancelling = false
 
   proc continuation(udata: pointer) {.gcsafe.} =
     if not(retFuture.finished()):
-      if not(fut.finished()):
-        # Timer exceeded first.
-        fut.removeCallback(continuation)
-        fut.cancel()
-        retFuture.complete(false)
+      if not(cancelling):
+        if not(fut.finished()):
+          # Timer exceeded first, we going to cancel `fut` and wait until it
+          # not completes.
+          cancelling = true
+          fut.cancel()
+        else:
+          # Future `fut` completed/failed/cancelled first.
+          if not(isNil(timer)):
+            clearTimer(timer)
+          retFuture.complete(true)
       else:
-        # Future `fut` completed/failed/cancelled first.
-        if not isNil(timer):
-          clearTimer(timer)
-        retFuture.complete(true)
+        retFuture.complete(false)
 
   proc cancellation(udata: pointer) {.gcsafe.} =
     if not isNil(timer):
@@ -900,26 +904,29 @@ proc wait*[T](fut: Future[T], timeout = InfiniteDuration): Future[T] =
   var retFuture = newFuture[T]("chronos.wait()")
   var moment: Moment
   var timer: TimerCallback
+  var cancelling = false
 
   proc continuation(udata: pointer) {.gcsafe.} =
     if not(retFuture.finished()):
-      if not(fut.finished()):
-        # Timer exceeded first.
-        fut.removeCallback(continuation)
-        fut.cancel()
-        retFuture.fail(newException(AsyncTimeoutError, "Timeout exceeded!"))
-      else:
-        # Future `fut` completed/failed/cancelled first.
-        if not isNil(timer):
-          clearTimer(timer)
-
-        if fut.failed():
-          retFuture.fail(fut.error)
+      if not(cancelling):
+        if not(fut.finished()):
+          # Timer exceeded first.
+          cancelling = true
+          fut.cancel()
         else:
-          when T is void:
-            retFuture.complete()
+          # Future `fut` completed/failed/cancelled first.
+          if not isNil(timer):
+            clearTimer(timer)
+
+          if fut.failed():
+            retFuture.fail(fut.error)
           else:
-            retFuture.complete(fut.read())
+            when T is void:
+              retFuture.complete()
+            else:
+              retFuture.complete(fut.read())
+      else:
+        retFuture.fail(newException(AsyncTimeoutError, "Timeout exceeded!"))
 
   proc cancellation(udata: pointer) {.gcsafe.} =
     if not isNil(timer):
diff --git a/tests/testfut.nim b/tests/testfut.nim
index 058b981..a96e10c 100644
--- a/tests/testfut.nim
+++ b/tests/testfut.nim
@@ -1067,6 +1067,52 @@ suite "Future[T] behavior test suite":
 
     r10 and r11 and r20 and r21 and r30 and r31
 
+  proc testWithTimeoutCancelAndWait(): bool =
+    proc futureNeverEnds(): Future[void] =
+      newFuture[void]("neverending.future")
+
+    proc futureOneLevelMore() {.async.} =
+      await futureNeverEnds()
+
+    proc testWithTimeout(): Future[bool] {.async.} =
+      var fut = futureOneLevelMore()
+      try:
+        let res = await withTimeout(fut, 100.milliseconds)
+        # Because `fut` is never-ending Future[T], `withTimeout` should return
+        # `false` but it also has to wait until `fut` is cancelled.
+        if not(res) and fut.cancelled():
+          return true
+        else:
+          return false
+      except CatchableError:
+        return false
+
+    waitFor testWithTimeout()
+
+  proc testWaitCancelAndWait(): bool =
+    proc futureNeverEnds(): Future[void] =
+      newFuture[void]("neverending.future")
+
+    proc futureOneLevelMore() {.async.} =
+      await futureNeverEnds()
+
+    proc testWait(): Future[bool] {.async.} =
+      var fut = futureOneLevelMore()
+      try:
+        await wait(fut, 100.milliseconds)
+        return false
+      except AsyncTimeoutError:
+        # Because `fut` is never-ending Future[T], `wait` should raise
+        # `AsyncTimeoutError`, but only after `fut` is cancelled.
+        if fut.cancelled():
+          return true
+        else:
+          return false
+      except CatchableError:
+        return false
+
+    waitFor testWait()
+
   test "Async undefined behavior (#7758) test":
     check test1() == true
   test "Immediately completed asynchronous procedure test":
@@ -1126,3 +1172,9 @@ suite "Future[T] behavior test suite":
 
   test "location test":
     check testSrcLocation() == true
+
+  test "withTimeout(fut) should wait cancellation test":
+    check testWithTimeoutCancelAndWait()
+
+  test "wait(fut) should wait cancellation test":
+    check testWaitCancelAndWait()