From c3114d2ed26825cb5132dd3338df7f2b79d8d8e0 Mon Sep 17 00:00:00 2001
From: Diego <diego@status.im>
Date: Sun, 11 Aug 2024 23:21:46 +0200
Subject: [PATCH] create raceCancell

---
 libp2p/utils/future.nim    |  9 +++++++
 tests/utils/testfuture.nim | 53 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)
 create mode 100644 tests/utils/testfuture.nim

diff --git a/libp2p/utils/future.nim b/libp2p/utils/future.nim
index 61289af0e..4dd5f147e 100644
--- a/libp2p/utils/future.nim
+++ b/libp2p/utils/future.nim
@@ -9,6 +9,7 @@
 
 {.push raises: [].}
 
+import sequtils
 import chronos
 
 type AllFuturesFailedError* = object of CatchableError
@@ -31,3 +32,11 @@ proc anyCompleted*[T](futs: seq[Future[T]]): Future[Future[T]] {.async.} =
 
     let index = requests.find(raceFut)
     requests.del(index)
+
+proc raceCancel*[T](
+    futs: seq[Future[T]]
+): Future[void] {.async: (raises: [ValueError, CancelledError]).} =
+  try:
+    discard await race(futs)
+  finally:
+    await noCancel allFutures(futs.filterIt(not it.finished).mapIt(it.cancelAndWait))
diff --git a/tests/utils/testfuture.nim b/tests/utils/testfuture.nim
new file mode 100644
index 000000000..aa1b60a33
--- /dev/null
+++ b/tests/utils/testfuture.nim
@@ -0,0 +1,53 @@
+{.used.}
+
+# Nim-Libp2p
+# Copyright (c) 2023 Status Research & Development GmbH
+# Licensed under either of
+#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
+#  * MIT license ([LICENSE-MIT](LICENSE-MIT))
+# at your option.
+# This file may not be copied, modified, or distributed except according to
+# those terms.
+
+import ../helpers
+import ../../libp2p/utils/future
+
+suite "Utils Future":
+
+  asyncTest "All Pending Tasks are canceled when returned future is canceled":
+    proc longRunningTaskA() {.async.} =
+      await sleepAsync(10.seconds)
+
+    proc longRunningTaskB() {.async.} =
+      await sleepAsync(10.seconds)
+
+    let
+      futureA = longRunningTaskA()
+      futureB = longRunningTaskB()
+
+    # Both futures should be canceled when raceCancel is called
+    await raceCancel(@[futureA, futureB]).cancelAndWait()
+    check futureA.cancelled
+    check futureB.cancelled
+
+  # Test where one task finishes immediately, leading to the cancellation of the pending task
+  asyncTest "Cancel Pending Tasks When One Completes":
+    proc quickTask() {.async.} =
+      return
+
+    proc slowTask() {.async.} =
+      await sleepAsync(10.seconds)
+
+    let
+      futureQuick = quickTask()
+      futureSlow = slowTask()
+
+    # The quick task finishes, so the slow task should be canceled
+    await raceCancel(@[futureQuick, futureSlow])
+    check futureQuick.finished
+    check futureSlow.cancelled
+
+  asyncTest "raceCancel with AsyncEvent":
+    let asyncEvent = newAsyncEvent()
+
+    await raceCancel(@[asyncEvent.wait()])
\ No newline at end of file