From da234d503b1df4a9510b5505676551dfe810c28b Mon Sep 17 00:00:00 2001
From: Eric <5089238+emizzle@users.noreply.github.com>
Date: Thu, 12 Dec 2024 19:31:51 +0700
Subject: [PATCH] fix(trackedfutures): removes usage of `then` from tracked
 futures (#1032)

- removes usage of `then`, simplifying the logic, and allowing `then` to be removed completely
- updates annotations to reflect that all procs (sync and async) raise no exceptions
---
 codex/utils/trackedfutures.nim | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/codex/utils/trackedfutures.nim b/codex/utils/trackedfutures.nim
index f3fcdb2d..a9753299 100644
--- a/codex/utils/trackedfutures.nim
+++ b/codex/utils/trackedfutures.nim
@@ -1,9 +1,9 @@
-import std/sugar
 import std/tables
 import pkg/chronos
 
 import ../logutils
-import ../utils/then
+
+{.push raises: [].}
 
 type
   TrackedFutures* = ref object
@@ -25,10 +25,10 @@ proc track*[T](self: TrackedFutures, fut: Future[T]): Future[T] =
 
   self.futures[fut.id] = FutureBase(fut)
 
-  fut
-    .then((val: T) => self.removeFuture(fut))
-    .cancelled(() => self.removeFuture(fut))
-    .catch((e: ref CatchableError) => self.removeFuture(fut))
+  proc cb(udata: pointer) =
+    self.removeFuture(fut)
+
+  fut.addCallback(cb)
 
   return fut
 
@@ -38,15 +38,17 @@ proc track*[T, U](future: Future[T], self: U): Future[T] =
   ## `trackedFutures` property.
   self.trackedFutures.track(future)
 
-proc cancelTracked*(self: TrackedFutures) {.async.} =
+proc cancelTracked*(self: TrackedFutures) {.async: (raises: []).} =
   self.cancelling = true
 
   trace "cancelling tracked futures"
 
+  var cancellations: seq[FutureBase]
   for future in self.futures.values:
     if not future.isNil and not future.finished:
-      trace "cancelling tracked future", id = future.id
-      await future.cancelAndWait()
+      cancellations.add future.cancelAndWait()
+
+  await noCancel allFutures cancellations
 
   self.futures.clear()
   self.cancelling = false