Do not return to scheduler while completing {.async.} proc

chronos/internal/asyncengine.nim

@@ -1158,6 +1158,11 @@ proc callSoon*(cbproc: CallbackFunc, data: pointer) {.
 proc callSoon*(cbproc: CallbackFunc) =
   callSoon(cbproc, nil)
 
+proc internalCallNext*(acb: AsyncCallback) =
+  ## Schedule `acb` to be called immediately after the current callback.
+  ## The callback is called before any other scheduled callbacks and events.
+  getThreadDispatcher().callbacks.addFirst(acb)
+
 proc callIdle*(acb: AsyncCallback) =
   ## Schedule ``cbproc`` to be called when there no pending network events
   ## available.

chronos/internal/asyncfutures.nim

@@ -176,18 +176,22 @@ proc checkFinished(future: FutureBase, loc: ptr SrcLoc) =
   else:
     future.internalLocation[LocationKind.Finish] = loc
 
+proc scheduleCall(
+    callback: AsyncCallback, isFromAsyncMacro: bool) {.raises: [], gcsafe.}
+
 proc finish(fut: FutureBase, state: FutureState, loc: ptr SrcLoc) =
   fut.checkFinished(loc)
 
   fut.internalState = state
   fut.internalCancelcb = nil # release cancellation callback memory
+  let isFromAsyncMacro = not(isNil(fut.internalClosure))
 
   if not(isNil(fut.internalCallback.function)):
-    callSoon(move(fut.internalCallback))
+    scheduleCall(move(fut.internalCallback), isFromAsyncMacro)
 
   for item in fut.internalCallbacks.mitems():
     if not(isNil(item.function)):
-      callSoon(item)
+      scheduleCall(item, isFromAsyncMacro)
     item = default(AsyncCallback) # release memory as early as possible
   fut.internalCallbacks = default(seq[AsyncCallback]) # release seq as well
 
@@ -372,6 +376,13 @@ proc internalContinue(fut: pointer) {.raises: [], gcsafe.} =
   GC_unref(asFut)
   futureContinue(asFut)
 
+proc scheduleCall(
+    callback: AsyncCallback, isFromAsyncMacro: bool) {.raises: [], gcsafe.} =
+  if isFromAsyncMacro and callback.function == internalContinue:
+    internalCallNext(callback)
+  else:
+    callSoon(callback)
+
 proc futureContinue*(fut: FutureBase) {.raises: [], gcsafe.} =
   # This function is responsible for calling the closure iterator generated by
   # the `{.async.}` transformation either until it has completed its iteration

tests/testall.nim

@@ -8,7 +8,7 @@
 import ".."/chronos/config
 
 when (chronosEventEngine in ["epoll", "kqueue"]) or defined(windows):
-  import testmacro, testsync, testsoon, testtime, testfut, testsignal,
+  import testmacro, testsync, testsoon, testunwind, testtime, testfut, testsignal,
          testaddress, testdatagram, teststream, testserver, testbugs, testnet,
          testasyncstream, testhttpserver, testshttpserver, testhttpclient,
          testproc, testratelimit, testfutures, testthreadsync, testasyncsemaphore
@@ -17,10 +17,10 @@ when (chronosEventEngine in ["epoll", "kqueue"]) or defined(windows):
   import testutils
 elif chronosEventEngine == "poll":
   # `poll` engine do not support signals and processes
-  import testmacro, testsync, testsoon, testtime, testfut, testaddress,
-         testdatagram, teststream, testserver, testbugs, testnet,
+  import testmacro, testsync, testsoon, testunwind, testtime, testfut,
+         testaddress, testdatagram, teststream, testserver, testbugs, testnet,
          testasyncstream, testhttpserver, testshttpserver, testhttpclient,
          testratelimit, testfutures, testthreadsync, testasyncsemaphore
 
   # Must be imported last to check for Pending futures
-  import testutils
+  import testutils

tests/testsync.nim

@@ -36,8 +36,8 @@ suite "Asynchronous sync primitives test suite":
     discard testLock(8, lock)
     discard testLock(9, lock)
     lock.release()
-    ## There must be exactly 20 poll() calls
-    for i in 0..<20:
+    ## There must be exactly 10 poll() calls
+    for i in 0..<10:
       poll()
     result = testLockResult
 
@@ -192,9 +192,10 @@ suite "Asynchronous sync primitives test suite":
     var queue = newAsyncQueue[int](1)
     discard task1(queue)
     discard task2(queue)
-    ## There must be exactly 2 poll() calls
-    poll()
-    poll()
+    ## There must be exactly 3 poll() calls
+    poll()  # task1 pops item1
+    poll()  # task2 puts item2
+    poll()  # task1 pops item2
     result = testQueue1Result
 
   const testsCount = 1000
@@ -231,7 +232,6 @@ suite "Asynchronous sync primitives test suite":
     discard task51(queue)
     discard task52(queue)
     poll()
-    poll()
     result = testQueue3Result
 
   proc test6(): bool =

tests/testunwind.nim

@@ -0,0 +1,195 @@
+#                Chronos Test Suite
+#            (c) Copyright 2026-Present
+#         Status Research & Development GmbH
+#
+#              Licensed under either of
+#  Apache License, version 2.0, (LICENSE-APACHEv2)
+#              MIT license (LICENSE-MIT)
+import unittest2
+import ../chronos
+
+{.push raises: [], gcsafe.}
+{.used.}
+
+suite "Stack unwind scheduling test suite":
+  type Trace = seq[string]
+
+  proc runTest(
+      cb: proc (
+          trace: ptr Trace
+        ): Future[void].Raising([CancelledError]) {.raises: [], gcsafe.}
+  ): Trace =
+    var trace: Trace
+    try:
+      waitFor cb(addr trace)
+    except CancelledError:
+      raiseAssert "Not cancelled"
+    trace
+
+  proc competitorCb(udata: pointer) =
+    cast[ptr Trace](udata)[].add "competitor"
+
+  proc observerCb(udata: pointer) =
+    cast[ptr Trace](udata)[].add "observer"
+
+  proc testValueReturn(): Trace =
+    proc producer(
+        trace: ptr Trace): Future[int] {.async: (raises: [CancelledError]).} =
+      await sleepAsync(ZeroDuration)
+      callSoon(competitorCb, trace)
+      trace[].add "producer returns"
+      42
+
+    proc consumer(trace: ptr Trace) {.async: (raises: [CancelledError]).} =
+      let v = await producer(trace)
+      trace[].add "consumer returns " & $v
+
+    runTest consumer
+
+  proc testFailingReturn(): Trace =
+    proc producer(
+        trace: ptr Trace
+    ): Future[int] {.async: (raises: [CancelledError, ValueError]).} =
+      await sleepAsync(ZeroDuration)
+      callSoon(competitorCb, trace)
+      trace[].add "producer raising"
+      raise newException(ValueError, "err")
+
+    proc consumer(trace: ptr Trace) {.async: (raises: [CancelledError]).} =
+      try:
+        discard await producer(trace)
+      except ValueError:
+        trace[].add "consumer caught"
+
+    runTest consumer
+
+  proc testMultiLevel(): Trace =
+    proc bottom(
+        trace: ptr Trace): Future[int] {.async: (raises: [CancelledError]).} =
+      await sleepAsync(ZeroDuration)
+      callSoon(competitorCb, trace)
+      trace[].add "bottom returns"
+      1
+
+    proc mid(
+        trace: ptr Trace): Future[int] {.async: (raises: [CancelledError]).} =
+      let v = await bottom(trace)
+      trace[].add "mid returns"
+      v + 1
+
+    proc top(trace: ptr Trace) {.async: (raises: [CancelledError]).} =
+      let v = await mid(trace)
+      trace[].add "top returns " & $v
+
+    runTest top
+
+  proc testObserverReturn(): Trace =
+    proc producer(
+        trace: ptr Trace): Future[int] {.async: (raises: [CancelledError]).} =
+      await sleepAsync(ZeroDuration)
+      callSoon(competitorCb, trace)
+      trace[].add "producer returns"
+      7
+
+    proc consumer(trace: ptr Trace) {.async: (raises: [CancelledError]).} =
+      let fut = producer(trace)
+      fut.addCallback(observerCb, trace)
+      let v = await fut
+      trace[].add "consumer returns " & $v
+
+    runTest consumer
+
+  proc testObserverRaise(): Trace =
+    proc producer(
+        trace: ptr Trace
+    ): Future[int] {.async: (raises: [CancelledError, ValueError]).} =
+      await sleepAsync(ZeroDuration)
+      callSoon(competitorCb, trace)
+      trace[].add "producer raising"
+      raise newException(ValueError, "err")
+
+    proc consumer(trace: ptr Trace) {.async: (raises: [CancelledError]).} =
+      let fut = producer(trace)
+      fut.addCallback(observerCb, trace)
+      try:
+        discard await fut
+      except ValueError:
+        trace[].add "consumer caught"
+
+    runTest consumer
+
+  proc testManualWakeup(): Trace =
+    let fut = Future[void].Raising([CancelledError]).init()
+
+    proc producer(trace: ptr Trace) {.async: (raises: [CancelledError]).} =
+      await fut
+      trace[].add "producer returns"
+
+    proc consumer(trace: ptr Trace) {.async: (raises: [CancelledError]).} =
+      let w = producer(trace)
+      callSoon(competitorCb, trace)
+      fut.complete()
+      await w
+
+    runTest consumer
+
+  proc testFanout(): Trace =
+    proc produce(
+        trace: ptr Trace): Future[int] {.async: (raises: [CancelledError]).} =
+      await sleepAsync(ZeroDuration)
+      trace[].add "produced"
+      42
+
+    var trace: Trace
+    let shared = produce(addr trace)
+
+    proc subA(
+        trace: ptr Trace): Future[void] {.async: (raises: [CancelledError]).} =
+      discard await shared
+      trace[].add "subA"
+
+    proc subB(
+        trace: ptr Trace): Future[void] {.async: (raises: [CancelledError]).} =
+      discard await shared
+      trace[].add "subB"
+
+    proc strainA(trace: ptr Trace) {.async: (raises: [CancelledError]).} =
+      await subA(trace)
+      trace[].add "strainA"
+
+    proc strainB(trace: ptr Trace) {.async: (raises: [CancelledError]).} =
+      await subB(trace)
+      trace[].add "strainB"
+
+    try:
+      waitFor allFutures(strainA(addr trace), strainB(addr trace))
+    except CancelledError:
+      raiseAssert "Not cancelled"
+    trace
+
+  test "Unwind not interrupted test":
+    check:
+      testValueReturn() ==
+        @["producer returns", "consumer returns 42", "competitor"]
+      testFailingReturn() ==
+        @["producer raising", "consumer caught", "competitor"]
+
+  test "Multi-level unwind not interrupted test":
+    check testMultiLevel() ==
+      @["bottom returns", "mid returns", "top returns 2", "competitor"]
+
+  test "Observer deferred test":
+    check:
+      testObserverReturn() ==
+        @["producer returns", "consumer returns 7", "competitor", "observer"]
+      testObserverRaise() ==
+        @["producer raising", "consumer caught", "competitor", "observer"]
+
+  test "Manual wakeup interruptible test":
+    check testManualWakeup() == @["competitor", "producer returns"]
+
+  test "Fan-out depth-first test":
+    # Each strain unwinds to conclusion before the next strain begins.
+    # Executed in reverse registration order.
+    check testFanout() ==
+      @["produced", "subB", "strainB", "subA", "strainA"]