From 315a27236c8ea34fb8f96d725892610ddcdf206d Mon Sep 17 00:00:00 2001
From: Eugene Kabanov <eugene.kabanov@status.im>
Date: Thu, 1 Jun 2023 18:09:49 +0300
Subject: [PATCH] Recover cancellation cleanup for AsyncEvent wait(). (#398)

* Recover cancellation cleanup for AsyncEvent wait().

* Address review comments.
---
 chronos/asyncsync.nim | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/chronos/asyncsync.nim b/chronos/asyncsync.nim
index 4a0a13e..a48002c 100644
--- a/chronos/asyncsync.nim
+++ b/chronos/asyncsync.nim
@@ -223,12 +223,15 @@ proc wait*(event: AsyncEvent): Future[void] =
   ## If the internal flag is `true` on entry, return immediately. Otherwise,
   ## block until another task calls `fire()` to set the flag to `true`,
   ## then return.
-  var w = newFuture[void]("AsyncEvent.wait")
+  let retFuture = newFuture[void]("AsyncEvent.wait")
+  proc cancellation(udata: pointer) {.gcsafe, raises: [Defect].} =
+    event.waiters.keepItIf(it != retFuture)
   if not(event.flag):
-    event.waiters.add(w)
+    retFuture.cancelCallback = cancellation
+    event.waiters.add(retFuture)
   else:
-    w.complete()
-  w
+    retFuture.complete()
+  retFuture
 
 proc fire*(event: AsyncEvent) =
   ## Set the internal flag of ``event`` to `true`. All tasks waiting for it