From 32bbbb051c4f13e6e8ab4a501076882ef9b14c9f Mon Sep 17 00:00:00 2001
From: cheatfate <ka@hardcore.kiev.ua>
Date: Fri, 20 Jul 2018 11:58:01 +0300
Subject: [PATCH] Add asynchronous specific exception AsyncError. Add
 AsyncTimeoutError. Add wait[T]() procedure and tests. Bump version to 2.0.2.

---
 asyncdispatch2.nimble                |  2 +-
 asyncdispatch2/asyncloop.nim         | 38 ++++++++++++++++
 asyncdispatch2/transports/common.nim |  2 +-
 tests/testfut.nim                    | 68 ++++++++++++++++++++++++++++
 4 files changed, 108 insertions(+), 2 deletions(-)

diff --git a/asyncdispatch2.nimble b/asyncdispatch2.nimble
index d5d7a93a..8ab7b9f5 100644
--- a/asyncdispatch2.nimble
+++ b/asyncdispatch2.nimble
@@ -1,5 +1,5 @@
 packageName   = "asyncdispatch2"
-version       = "2.0.1"
+version       = "2.0.2"
 author        = "Status Research & Development GmbH"
 description   = "Asyncdispatch2"
 license       = "Apache License 2.0 or MIT"
diff --git a/asyncdispatch2/asyncloop.nim b/asyncdispatch2/asyncloop.nim
index d2112099..eedfbf14 100644
--- a/asyncdispatch2/asyncloop.nim
+++ b/asyncdispatch2/asyncloop.nim
@@ -166,6 +166,11 @@ export asyncfutures2
 # TODO: Check if yielded future is nil and throw a more meaningful exception
 
 type
+  AsyncError* = object of Exception
+    ## Generic async exception
+  AsyncTimeoutError* = object of AsyncError
+    ## Timeout exception
+
   TimerCallback* = object
     finishAt*: uint64
     function*: AsyncCallback
@@ -639,6 +644,39 @@ proc withTimeout*[T](fut: Future[T], timeout: int): Future[bool] =
   fut.addCallback(continuation)
   return retFuture
 
+proc wait*[T](fut: Future[T], timeout = -1): Future[T] =
+  ## Returns a future which will complete once future ``fut`` completes
+  ## or if timeout of ``timeout`` milliseconds has been expired.
+  ## 
+  ## If ``timeout`` is ``-1``, then result future will be completed only
+  ## when ``fut`` become completed.
+  var retFuture = newFuture[T]("asyncdispatch.wait")
+  proc continuation(udata: pointer) {.gcsafe.} =
+    if not retFuture.finished:
+      if isNil(udata):
+        fut.removeCallback(continuation)
+        retFuture.fail(newException(AsyncTimeoutError, ""))
+      else:
+        if not retFuture.finished:
+          if fut.failed:
+            retFuture.fail(fut.error)
+          else:
+            retFuture.complete(fut.read())
+  if timeout == -1:
+    retFuture = fut
+  elif timeout == 0:
+    if fut.finished:
+      if fut.failed:
+        retFuture.fail(fut.error)
+      else:
+        retFuture.complete(fut.read())
+    else:
+      retFuture.fail(newException(AsyncTimeoutError, ""))
+  else:
+    addTimer(fastEpochTime() + uint64(timeout), continuation, nil)
+    fut.addCallback(continuation)
+  return retFuture
+
 include asyncmacro2
 
 proc callSoon(cbproc: CallbackFunc, data: pointer = nil) =
diff --git a/asyncdispatch2/transports/common.nim b/asyncdispatch2/transports/common.nim
index bdd777a4..c3ee4d59 100644
--- a/asyncdispatch2/transports/common.nim
+++ b/asyncdispatch2/transports/common.nim
@@ -82,7 +82,7 @@ else:
       loopFuture*: Future[void]     # Server's main Future
 
 type
-  TransportError* = object of Exception
+  TransportError* = object of AsyncError
     ## Transport's specific exception
   TransportOsError* = object of TransportError
     ## Transport's OS specific exception
diff --git a/tests/testfut.nim b/tests/testfut.nim
index df4c9de7..cca43d3b 100644
--- a/tests/testfut.nim
+++ b/tests/testfut.nim
@@ -18,6 +18,67 @@ proc testFuture2(): Future[int] {.async.} =
 proc testFuture3(): Future[int] {.async.} =
   result = await testFuture2()
 
+proc testFuture4(): Future[int] {.async.} =
+  ## Test for not immediately completed future and timeout = -1
+  result = 0
+  try:
+    var res = await wait(testFuture1(), -1)
+    result = 1
+  except:
+    result = 0
+
+  if result == 0:
+    return
+
+  ## Test for immediately completed future and timeout = -1
+  result = 0
+  try:
+    var res = await wait(testFuture2(), -1)
+    result = 1
+  except:
+    result = 0
+
+  if result == 0:
+    return
+
+  ## Test for not immediately completed future and timeout = 0
+  result = 0
+  try:
+    var res = await wait(testFuture1(), 0)
+  except AsyncTimeoutError:
+    result = 1
+
+  if result == 0:
+    return
+
+  ## Test for immediately completed future and timeout = 0
+  result = 0
+  try:
+    var res = await wait(testFuture2(), 0)
+    result = 1
+  except:
+    result = 0
+
+  if result == 0:
+    return
+
+  ## Test for future which cannot be completed in timeout period
+  result = 0
+  try:
+    var res = await wait(testFuture1(), 50)
+  except AsyncTimeoutError:
+    result = 1
+
+  if result == 0:
+    return
+
+  ## Test for future which will be completed before timeout exceeded.
+  try:
+    var res = await wait(testFuture1(), 150)
+    result = 1
+  except:
+    result = 0
+
 proc test1(): bool =
   var fut = testFuture1()
   poll()
@@ -70,6 +131,11 @@ proc test4(): string =
   if fut.finished:
     result = testResult
 
+proc test5(): bool =
+  var res = waitFor(testFuture4())
+  if res == 1:
+    result = true
+
 when isMainModule:
   suite "Future[T] behavior test suite":
     test "Async undefined behavior (#7758) test":
@@ -80,3 +146,5 @@ when isMainModule:
       check test3() == "12345"
     test "Future[T] callbacks not changing order after removeCallback()":
       check test4() == "1245"
+    test "wait[T]() test":
+      check test5() == true