Fix asyncdispatch drain behavior (#14820) (#14838)

* Fix asyncdispatch drain behavior (#14820)

* Changed test to use asyncCheck instead of discard after code review (#14820)

* Added some debug statements to help understand what is happening in Azure.

* Removed debug statements and increased timeouts by 1 order of magnitude to account for slow Azure VMs

Co-authored-by: Ray Imber <ray@crankuptheamps.com>

2 files changed, 34 insertions, 4 deletions

diff --git a/lib/pure/asyncdispatch.nim b/lib/pure/asyncdispatch.nim
index a15442109..a0bc96e72 100644
--- a/lib/pure/asyncdispatch.nim
+++ b/lib/pure/asyncdispatch.nim
@@ -1601,11 +1601,16 @@ else:
     p.selector.registerEvent(SelectEvent(ev), data)
 
 proc drain*(timeout = 500) =
-  ## Waits for completion events and processes them. Raises ``ValueError``
+  ## Waits for completion of **all** events and processes them. Raises ``ValueError``
   ## if there are no pending operations. In contrast to ``poll`` this
-  ## processes as many events as are available.
-  if runOnce(timeout) or hasPendingOperations():
-    while hasPendingOperations() and runOnce(timeout): discard
+  ## processes as many events as are available until the timeout has elapsed.
+  var curTimeout = timeout
+  let start = now()
+  while hasPendingOperations():
+    discard runOnce(curTimeout) 
+    curTimeout -= (now() - start).inMilliseconds.int
+    if curTimeout < 0:
+      break
 
 proc poll*(timeout = 500) =
   ## Waits for completion events and processes them. Raises ``ValueError``
diff --git a/tests/async/t14820.nim b/tests/async/t14820.nim
new file mode 100644
index 000000000..359884468
--- /dev/null
+++ b/tests/async/t14820.nim
@@ -0,0 +1,25 @@
+discard """
+output: '''
+iteration: 1
+iteration: 2
+iteration: 3
+iteration: 4
+async done
+iteration: 5
+'''
+"""
+
+import asyncdispatch, times
+
+var done = false
+proc somethingAsync() {.async.} =
+  yield sleepAsync 5000
+  echo "async done"
+  done = true
+  
+asyncCheck somethingAsync()
+var count = 0
+while not done:
+  count += 1
+  drain 1000
+  echo "iteration: ", count