fixes #23212; Asyncdispatch leaks under --mm:arc (#24556)

Fixes #23212 Inspired by [this chronos PR](https://github.com/status-im/nim-chronos/pull/243) (cherry picked from commit f29234b40f)
2026-06-04 10:54:42 +00:00 · 2024-12-22 03:56:22 -03:00
parent 0c14372a8c
commit ede6540c55
3 changed files with 82 additions and 43 deletions
--- a/lib/pure/asyncmacro.nim
+++ b/lib/pure/asyncmacro.nim
@@ -22,36 +22,39 @@ proc newCallWithLineInfo(fromNode: NimNode; theProc: NimNode, args: varargs[NimN
  result = newCall(theProc, args)
  result.copyLineInfo(fromNode)

-template createCb(retFutureSym, iteratorNameSym,
-                  strName, identName, futureVarCompletions: untyped) =
+type
+  ClosureIt[T] = iterator(f: Future[T]): owned(FutureBase)
+
+template createCb(futTyp, strName, identName, futureVarCompletions: untyped) =
  bind finished
-  var nameIterVar = iteratorNameSym
-  proc identName {.closure, stackTrace: off.} =
+  {.push stackTrace: off.}
+  proc identName(fut: Future[futTyp], it: ClosureIt[futTyp]) {.effectsOf: it.} =
    try:
-      if not nameIterVar.finished:
-        var next = nameIterVar()
+      if not it.finished:
+        var next = it(fut)
        # Continue while the yielded future is already finished.
        while (not next.isNil) and next.finished:
-          next = nameIterVar()
-          if nameIterVar.finished:
+          next = it(fut)
+          if it.finished:
            break

        if next == nil:
-          if not retFutureSym.finished:
+          if not fut.finished:
            let msg = "Async procedure ($1) yielded `nil`, are you await'ing a `nil` Future?"
            raise newException(AssertionDefect, msg % strName)
        else:
          {.gcsafe.}:
-            next.addCallback cast[proc() {.closure, gcsafe.}](identName)
+            next.addCallback(cast[proc() {.closure, gcsafe.}](proc =
+              identName(fut, it)))
    except:
      futureVarCompletions
-      if retFutureSym.finished:
+      if fut.finished:
        # Take a look at tasyncexceptions for the bug which this fixes.
        # That test explains it better than I can here.
        raise
      else:
-        retFutureSym.fail(getCurrentException())
-  identName()
+        fut.fail(getCurrentException())
+  {.pop.}

 proc createFutureVarCompletions(futureVarIdents: seq[NimNode], fromNode: NimNode): NimNode =
  result = newNimNode(nnkStmtList, fromNode)
@@ -68,7 +71,7 @@ proc createFutureVarCompletions(futureVarIdents: seq[NimNode], fromNode: NimNode
      )
    )

-proc processBody(ctx: Context; node, needsCompletionSym, retFutureSym: NimNode, futureVarIdents: seq[NimNode]): NimNode =
+proc processBody(ctx: Context; node, needsCompletionSym, retFutParamSym: NimNode, futureVarIdents: seq[NimNode]): NimNode =
  result = node
  case node.kind
  of nnkReturnStmt:
@@ -80,14 +83,14 @@ proc processBody(ctx: Context; node, needsCompletionSym, retFutureSym: NimNode,
    ctx.hasRet = true
    if node[0].kind == nnkEmpty:
      if ctx.inTry == 0:
-        result.add newCallWithLineInfo(node, newIdentNode("complete"), retFutureSym, newIdentNode("result"))
+        result.add newCallWithLineInfo(node, newIdentNode("complete"), retFutParamSym, newIdentNode("result"))
      else:
        result.add newAssignment(needsCompletionSym, newLit(true))
    else:
-      let x = processBody(ctx, node[0], needsCompletionSym, retFutureSym, futureVarIdents)
+      let x = processBody(ctx, node[0], needsCompletionSym, retFutParamSym, futureVarIdents)
      if x.kind == nnkYieldStmt: result.add x
      elif ctx.inTry == 0:
-        result.add newCallWithLineInfo(node, newIdentNode("complete"), retFutureSym, x)
+        result.add newCallWithLineInfo(node, newIdentNode("complete"), retFutParamSym, x)
      else:
        result.add newAssignment(newIdentNode("result"), x)
        result.add newAssignment(needsCompletionSym, newLit(true))
@@ -100,10 +103,10 @@ proc processBody(ctx: Context; node, needsCompletionSym, retFutureSym: NimNode,
  of nnkTryStmt:
    if result[^1].kind == nnkFinally:
      inc ctx.inTry
-      result[0] = processBody(ctx, result[0], needsCompletionSym, retFutureSym, futureVarIdents)
+      result[0] = processBody(ctx, result[0], needsCompletionSym, retFutParamSym, futureVarIdents)
      dec ctx.inTry
      for i in 1 ..< result.len:
-        result[i] = processBody(ctx, result[i], needsCompletionSym, retFutureSym, futureVarIdents)
+        result[i] = processBody(ctx, result[i], needsCompletionSym, retFutParamSym, futureVarIdents)
      if ctx.inTry == 0 and ctx.hasRet:
        let finallyNode = copyNimNode(result[^1])
        let stmtNode = newNimNode(nnkStmtList)
@@ -111,7 +114,7 @@ proc processBody(ctx: Context; node, needsCompletionSym, retFutureSym: NimNode,
          stmtNode.add child
        stmtNode.add newIfStmt(
          ( needsCompletionSym,
-            newCallWithLineInfo(node, newIdentNode("complete"), retFutureSym,
+            newCallWithLineInfo(node, newIdentNode("complete"), retFutParamSym,
            newIdentNode("result")
            )
          )
@@ -120,10 +123,10 @@ proc processBody(ctx: Context; node, needsCompletionSym, retFutureSym: NimNode,
        result[^1] = finallyNode
    else:
      for i in 0 ..< result.len:
-        result[i] = processBody(ctx, result[i], needsCompletionSym, retFutureSym, futureVarIdents)
+        result[i] = processBody(ctx, result[i], needsCompletionSym, retFutParamSym, futureVarIdents)
  else:
    for i in 0 ..< result.len:
-      result[i] = processBody(ctx, result[i], needsCompletionSym, retFutureSym, futureVarIdents)
+      result[i] = processBody(ctx, result[i], needsCompletionSym, retFutParamSym, futureVarIdents)

  # echo result.repr

@@ -229,29 +232,21 @@ proc asyncSingleProc(prc: NimNode): NimNode =
  # transformation even more complex.
  let body2 = extractDocCommentsAndRunnables(prc.body)

-  # -> var retFuture = newFuture[T]()
-  var retFutureSym = genSym(nskVar, "retFuture")
  var subRetType =
    if returnType.kind == nnkEmpty: newIdentNode("void")
    else: baseType
-  outerProcBody.add(
-    newVarStmt(retFutureSym,
-      newCall(
-        newNimNode(nnkBracketExpr, prc.body).add(
-          newIdentNode("newFuture"),
-          subRetType),
-      newLit(prcName)))) # Get type from return type of this proc
+  let retFutParamSym = genSym(nskParam, "retFutParamSym")

-  # -> iterator nameIter(): FutureBase {.closure.} =
+  # -> iterator nameIter(retFutParam: Future[T]): FutureBase {.closure.} =
  # ->   {.push warning[resultshadowed]: off.}
  # ->   var result: T
  # ->   {.pop.}
  # ->   <proc_body>
-  # ->   complete(retFuture, result)
+  # ->   complete(retFutParam, result)
  var iteratorNameSym = genSym(nskIterator, $prcName & " (Async)")
  var needsCompletionSym = genSym(nskVar, "needsCompletion")
  var ctx = Context()
-  var procBody = processBody(ctx, prc.body, needsCompletionSym, retFutureSym, futureVarIdents)
+  var procBody = processBody(ctx, prc.body, needsCompletionSym, retFutParamSym, futureVarIdents)
  # don't do anything with forward bodies (empty)
  if procBody.kind != nnkEmpty:
    # fix #13899, defer should not escape its original scope
@@ -275,9 +270,11 @@ proc asyncSingleProc(prc: NimNode): NimNode =

      var `needsCompletionSym` = false
    procBody.add quote do:
-      complete(`retFutureSym`, `resultIdent`)
+      complete(`retFutParamSym`, `resultIdent`)

-    var closureIterator = newProc(iteratorNameSym, [quote do: owned(FutureBase)],
+    var retFutureTyp = newNimNode(nnkBracketExpr, prc).add(newIdentNode("Future")).add(subRetType)
+    var retFutureParam = newNimNode(nnkIdentDefs, prc).add(retFutParamSym).add(retFutureTyp).add(newEmptyNode())
+    var closureIterator = newProc(iteratorNameSym, [quote do: owned(FutureBase), retFutureParam],
                                  procBody, nnkIteratorDef)
    closureIterator.pragma = newNimNode(nnkPragma, lineInfoFrom = prc.body)
    closureIterator.addPragma(newIdentNode("closure"))
@@ -287,17 +284,31 @@ proc asyncSingleProc(prc: NimNode): NimNode =
      closureIterator.addPragma(newIdentNode("gcsafe"))
    outerProcBody.add(closureIterator)

-    # -> createCb(retFuture)
+    # -> createCb()
    # NOTE: The NimAsyncContinueSuffix is checked for in asyncfutures.nim to produce
    # friendlier stack traces:
    var cbName = genSym(nskProc, prcName & NimAsyncContinueSuffix)
-    var procCb = getAst createCb(retFutureSym, iteratorNameSym,
-                          newStrLitNode(prcName),
-                          cbName,
-                          createFutureVarCompletions(futureVarIdents, nil)
-                        )
+    var procCb = getAst createCb(
+      subRetType,
+      newStrLitNode(prcName),
+      cbName,
+      createFutureVarCompletions(futureVarIdents, nil)
+    )
    outerProcBody.add procCb

+    # -> var retFuture = newFuture[T]()
+    let retFutureSym = genSym(nskVar, "retFuture")
+    outerProcBody.add(
+      newVarStmt(retFutureSym,
+        newCall(
+          newNimNode(nnkBracketExpr, prc.body).add(
+            newIdentNode("newFuture"),
+            subRetType),
+        newLit(prcName)))) # Get type from return type of this proc
+
+    # -> cb(retFuture, nameIter)
+    outerProcBody.add newCall(cbName, retFutureSym, iteratorNameSym)
+
    # -> return retFuture
    outerProcBody.add newNimNode(nnkReturnStmt, prc.body[^1]).add(retFutureSym)

--- a/tests/arc/tasyncleak.nim
+++ b/tests/arc/tasyncleak.nim
@@ -1,5 +1,5 @@
 discard """
-  outputsub: "(allocCount: 4050, deallocCount: 4048)"
+  outputsub: "(allocCount: 4011, deallocCount: 4009)"
  cmd: "nim c --gc:orc -d:nimAllocStats $file"
 """

--- a/tests/async/t23212.nim
+++ b/tests/async/t23212.nim
@@ -0,0 +1,28 @@
+discard """
+  valgrind: true
+  cmd: '''nim c --mm:arc -d:nimAllocStats -d:useMalloc $file'''
+  output: '''1000'''
+"""
+
+import std/asyncdispatch
+
+var count: int
+
+proc stuff() {.async.} =
+  #echo count, 1
+  await sleepAsync(1)
+  #echo count, 2
+  count.inc
+
+for _ in 0..<1000:
+  asyncCheck stuff()
+
+while hasPendingOperations(): poll()
+
+echo count
+
+setGlobalDispatcher(nil)
+
+import std/importutils
+privateAccess(AllocStats)
+doAssert getAllocStats().allocCount - getAllocStats().deallocCount < 10