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

.github/workflows/ci.yml

@@ -13,5 +13,6 @@ jobs:
       test-command: |
         nimble install -y libbacktrace
         nimble test
+        nimble test_sync_continuations
         nimble test_libbacktrace
         nimble examples

chronos.nimble

@@ -89,6 +89,13 @@ task test, "Run all tests":
     if f.startsWith("bench_") and f.endsWith(".nim"):
       build "", f[0..^5]
 
+task test_sync_continuations, "Run all tests with sync continuations":
+  for args in testArguments:
+    # First run tests with `refc` memory manager.
+    run args & " --mm:refc -d:chronosSyncContinuations", "tests/testall"
+    if (NimMajor, NimMinor) >= (2, 2): # ORC on 2.0 is too broken to investigate
+      run args & " --mm:orc -d:chronosSyncContinuations", "tests/testall"
+
 task test_v3_compat, "Run all tests in v3 compatibility mode":
   for args in testArguments:
     if (NimMajor, NimMinor) >= (2, 2):

chronos/config.nim

@@ -34,6 +34,13 @@ const
     ##
     ## `Exception` handling may be removed in future chronos versions.
 
+  chronosSyncContinuations* {.booldefine.}: bool = false
+    ## When enabled, `CallbackFlag.Continuation` callbacks of futures with
+    ## `FutureFlag.SyncContinuations` are scheduled to run before processing
+    ## other already queued events such as timer handlers and I/O.
+    ##
+    ## This feature is experimental and may be removed in future releases.
+
   chronosStrictFutureAccess* {.booldefine.}: bool = defined(chronosPreviewV5)
 
   chronosStackTrace* {.booldefine.}: bool = defined(chronosDebug)

chronos/futures.nim

@@ -26,6 +26,14 @@ type
   # TODO https://github.com/nim-lang/Nim/issues/25976
   CallbackFunc* = proc (arg: pointer) {.gcsafe, raises: [].}
 
+  CallbackFlag* {.pure.} = enum
+    Continuation
+      ## When set and attached to a future with `FutureFlag.SyncContinuations`,
+      ## the callback is scheduled to run before processing other already
+      ## queued events such as timer handlers and I/O.
+      ##
+      ## Only works when the `chronosSyncContinuations` config is enabled.
+
   # Internal type, not part of API
   InternalAsyncCallback* = object
     function*: CallbackFunc
@@ -50,6 +58,13 @@ type
       ## If `cancelCallback` is not set and the future gets cancelled, a
       ## `Defect` will be raised.
 
+    SyncContinuations
+      ## When set, `CallbackFlag.Continuation` callbacks are scheduled
+      ## to run before processing other already queued events such as
+      ## timer handlers and I/O.
+      ##
+      ## Only works when the `chronosSyncContinuations` config is enabled.
+
   FutureFlags* = set[FutureFlag]
 
   InternalFutureBase* = object of RootObj
@@ -58,13 +73,13 @@ type
     # structure may change in the future (haha)
 
     internalLocation*: array[LocationKind, ptr SrcLoc]
-    internalCallback*: InternalAsyncCallback
+    internalContinuation*, internalCallback*: InternalAsyncCallback
       ## The vast majority of futures track a single callback only (the one
       ## installed by `await`) - to avoid allocating a seq (which involves
       ## making a separate allocation with space for several callbacks), we keep
       ## a spot in each future for that first one - the seq below will stay
       ## empty until a second callback is added
-    internalCallbacks*: seq[InternalAsyncCallback]
+    internalContinuations*, internalCallbacks*: seq[InternalAsyncCallback]
     internalCancelcb*: CallbackFunc
     internalChild*: FutureBase
     internalState*: FutureState

chronos/internal/asyncengine.nim

@@ -66,6 +66,7 @@ type
     trackers*: Table[string, TrackerBase]
     counters*: Table[string, TrackerCounter]
     inEventLoop: bool
+    numImmediate: int
 
 proc sentinelCallbackImpl(arg: pointer) {.gcsafe, noreturn.} =
   raiseAssert "Sentinel callback MUST not be scheduled"
@@ -160,10 +161,15 @@ template processCallbacks(loop: untyped) =
   when chronosStrictReentrancy:
     # Process existing callbacks but not those that follow, to allow the network
     # to regain control regularly
-    for _ in 0 ..< len(loop.callbacks):
+    loop.numImmediate = 0
+    var numRemaining = len(loop.callbacks)
+    while numRemaining > 0:
+      dec(numRemaining)
       let callable = loop.callbacks.popFirst()
       if not(isNil(callable.function)):
         callable.function(callable.udata)
+      numRemaining += loop.numImmediate
+      loop.numImmediate = 0
   else:
     while true:
       let callable = loop.callbacks.popFirst()  # len must be > 0 due to sentinel
@@ -1186,20 +1192,29 @@ proc removeTimer*(at: uint64, cb: CallbackFunc, udata: pointer = nil) {.
      inline, deprecated: "Use removeTimer(Duration, cb, udata)".} =
   removeTimer(Moment.init(int64(at), Millisecond), cb, udata)
 
-proc callSoon*(acb: AsyncCallback) =
-  ## Schedule `cbproc` to be called as soon as possible.
-  ## The callback is called when control returns to the event loop.
-  getThreadDispatcher().callbacks.addLast(acb)
+proc callSoon*(acb: AsyncCallback, immediate: static bool = false) =
+  ## Schedule `acb` to be called as soon as possible.
+  ## With `immediate`, no other scheduled callbacks and events are run.
+  ## Otherwise, the callback is called when control returns to the event loop.
+  let loop = getThreadDispatcher()
+  when immediate:
+    loop.callbacks.addFirst(acb)
+    when chronosStrictReentrancy:
+      loop.numImmediate += 1
+  else:
+    loop.callbacks.addLast(acb)
 
-proc callSoon*(cbproc: CallbackFunc, data: pointer) {.
-     gcsafe.} =
+proc callSoon*(
+    cbproc: CallbackFunc, data: pointer,
+    immediate: static bool = false) {.gcsafe.} =
   ## Schedule `cbproc` to be called as soon as possible.
-  ## The callback is called when control returns to the event loop.
+  ## With `immediate`, no other scheduled callbacks and events are run.
+  ## Otherwise, the callback is called when control returns to the event loop.
   doAssert(not isNil(cbproc))
-  callSoon(AsyncCallback(function: cbproc, udata: data))
+  callSoon(AsyncCallback(function: cbproc, udata: data), immediate)
 
-proc callSoon*(cbproc: CallbackFunc) =
-  callSoon(cbproc, nil)
+proc callSoon*(cbproc: CallbackFunc, immediate: static bool = false) =
+  callSoon(cbproc, nil, immediate)
 
 proc callIdle*(acb: AsyncCallback) =
   ## Schedule ``cbproc`` to be called when there no pending network events

chronos/internal/asyncfutures.nim

@@ -109,12 +109,13 @@ template newFuture*[T](fromProc: static[string] = "",
   else:
     newFutureImpl[T](getSrcLocation(fromProc), flags)
 
-template newInternalRaisesFuture*[T, E](fromProc: static[string] = ""): auto =
+template newInternalRaisesFuture*[T, E](
+    fromProc: static[string] = "", flags: static[FutureFlags] = {}): auto =
   ## Creates a new future.
   ##
   ## Specifying ``fromProc``, which is a string specifying the name of the proc
   ## that this future belongs to, is a good habit as it helps with debugging.
-  newInternalRaisesFutureImpl[T, E](getSrcLocation(fromProc), {})
+  newInternalRaisesFutureImpl[T, E](getSrcLocation(fromProc), flags)
 
 template newFutureSeq*[A, B](fromProc: static[string] = ""): FutureSeq[A, B] {.deprecated.} =
   ## Create a new future which can hold/preserve GC sequence until future will
@@ -182,6 +183,17 @@ proc finish(fut: FutureBase, state: FutureState, loc: ptr SrcLoc) =
   fut.internalState = state
   fut.internalCancelcb = nil # release cancellation callback memory
 
+  when chronosSyncContinuations:
+    proc process(callback: var AsyncCallback) =
+      if not(isNil(callback.function)):
+        callSoon(move(callback), immediate = true)
+
+    for i in countdown(fut.internalContinuations.high, 0):
+      process(fut.internalContinuations[i])
+    fut.internalContinuations = default(seq[AsyncCallback])
+
+    process(fut.internalContinuation)
+
   if not(isNil(fut.internalCallback.function)):
     callSoon(move(fut.internalCallback))
 
@@ -298,37 +310,71 @@ proc clearCallbacks(future: FutureBase) =
   future.internalCallback.reset()
   future.internalCallbacks.reset()
 
-proc addCallback*(future: FutureBase, cb: CallbackFunc, udata: pointer) =
+proc addCallback*(
+    future: FutureBase, cb: CallbackFunc, udata: pointer,
+    flags: static set[CallbackFlag] = {}) =
   ## Adds the callbacks proc to be called when the future completes.
   ##
   ## If future has already completed then ``cb`` will be called immediately.
   doAssert(not isNil(cb))
-  if future.finished():
-    callSoon(cb, udata)
+
+  const _: set[CallbackFlag] = flags  # https://github.com/nim-lang/Nim/issues/25938
+  let isContinuation =
+    when chronosSyncContinuations and CallbackFlag.Continuation in flags:
+      FutureFlag.SyncContinuations in future.internalFlags
+    else:
+      false
+
+  if isContinuation:
+    if future.finished():
+      callSoon(cb, udata, immediate = true)
+    else:
+      if isNil(future.internalContinuation.function):
+        future.internalContinuation = AsyncCallback(function: cb, udata: udata)
+      else:
+        future.internalContinuations.add(
+          AsyncCallback(function: cb, udata: udata))
   else:
-    if isNil(future.internalCallback.function):
-      future.internalCallback = AsyncCallback(function: cb, udata: udata)
+    if future.finished():
+      callSoon(cb, udata)
     else:
-      future.internalCallbacks.add AsyncCallback(function: cb, udata: udata)
+      if isNil(future.internalCallback.function):
+        future.internalCallback = AsyncCallback(function: cb, udata: udata)
+      else:
+        future.internalCallbacks.add(
+          AsyncCallback(function: cb, udata: udata))
 
-proc addCallback*(future: FutureBase, cb: CallbackFunc) =
+proc addCallback*(
+    future: FutureBase, cb: CallbackFunc,
+    flags: static set[CallbackFlag] = {}) =
   ## Adds the callbacks proc to be called when the future completes.
   ##
   ## If future has already completed then ``cb`` will be called immediately.
-  future.addCallback(cb, cast[pointer](future))
+  future.addCallback(cb, cast[pointer](future), flags)
 
 proc removeCallback*(future: FutureBase, cb: CallbackFunc,
                      udata: pointer) =
   ## Remove future from list of callbacks - this operation may be slow if there
   ## are many registered callbacks!
   doAssert(not isNil(cb))
+
+  template shouldRemove(callback: AsyncCallback): bool =
+    callback.function == cb and callback.udata == udata
+
   # Make sure to release memory associated with callback, or reference chains
   # may be created!
-  if future.internalCallback.function == cb and future.internalCallback.udata == udata:
+  when chronosSyncContinuations:
+    if shouldRemove(future.internalContinuation):
+      future.internalContinuation.reset()
+
+    future.internalContinuations.keepItIf:
+      not shouldRemove(it)
+
+  if shouldRemove(future.internalCallback):
     future.internalCallback.reset()
 
   future.internalCallbacks.keepItIf:
-    it.function != cb or it.udata != udata
+    not shouldRemove(it)
 
 proc removeCallback*(future: FutureBase, cb: CallbackFunc) =
   future.removeCallback(cb, cast[pointer](future))
@@ -394,7 +440,9 @@ proc futureContinue*(fut: FutureBase) {.raises: [], gcsafe.} =
       # We cannot make progress on `fut` until `next` has finished - schedule
       # `fut` to continue running when that happens
       GC_ref(fut)
-      next.addCallback(CallbackFunc(internalContinue), cast[pointer](fut))
+      next.addCallback(
+        CallbackFunc(internalContinue), cast[pointer](fut),
+        {CallbackFlag.Continuation})
 
       # return here so that we don't remove the closure below
       return

chronos/internal/asyncmacro.nim

@@ -513,7 +513,7 @@ proc asyncSingleProc(prc, params: NimNode): NimNode {.compileTime.} =
     # with:
     #
     # ```nim
-    # let resultFuture = newFuture[T]()
+    # let resultFuture = newFuture[T]({FutureFlag.SyncContinuations})
     # resultFuture.internalClosure = `iteratorNameSym`
     # futureContinue(resultFuture)
     # return resultFuture
@@ -545,7 +545,8 @@ proc asyncSingleProc(prc, params: NimNode): NimNode {.compileTime.} =
     outerProcBody.add(
       newLetStmt(
         retFutureSym,
-        newCall(newFutProc, newLit(prcName))
+        newCall(newFutProc, newLit(prcName), nnkCurly.newTree(
+          newDotExpr(ident "FutureFlag", ident "SyncContinuations")))
       )
     )
 

tests/testall.nim

@@ -10,10 +10,10 @@ import ../chronos/config
 
 import
   ./[
-    testmacro, testsync, testsoon, testtime, testfut, testaddress,
-    testdatagram, teststream, testserver, testbugs, testnet, testasyncstream,
-    testhttpserver, testshttpserver, testhttpclient, testratelimit,
-    testfutures, testthreadsync, testasyncsemaphore,
+    testmacro, testsync, testsoon, testcontinuations, testtime, testfut,
+    testaddress, testdatagram, teststream, testserver, testbugs, testnet,
+    testasyncstream, testhttpserver, testshttpserver, testhttpclient,
+    testratelimit, testfutures, testthreadsync, testasyncsemaphore,
   ]
 
 when (chronosEventEngine in ["epoll", "kqueue"]) or defined(windows):