From 43520c360833744834fca91df7acee9586652068 Mon Sep 17 00:00:00 2001
From: Jaremy Creechley <creechley@gmail.com>
Date: Thu, 28 Sep 2023 14:51:27 -0700
Subject: [PATCH] test task cancel

---
 datastore/threads/threadproxyds.nim   |  5 +++-
 tests/datastore/testthreadproxyds.nim | 37 ++++++++++++++++++++-------
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/datastore/threads/threadproxyds.nim b/datastore/threads/threadproxyds.nim
index 19b5def..e1a7571 100644
--- a/datastore/threads/threadproxyds.nim
+++ b/datastore/threads/threadproxyds.nim
@@ -128,8 +128,10 @@ template dispatchTask*[BT](self: ThreadDatastore[BT],
     trace "Cancelling thread future!", exc = exc.msg
     ctx.setCancelled()
     raise exc
+  except CatchableError as exc:
+    ctx.setCancelled()
+    raise exc
   finally:
-    # echo "signal:CLOSE!"
     discard ctx[].signal.close()
     self.semaphore.release()
 
@@ -359,6 +361,7 @@ method query*[BT](self: ThreadDatastore[BT],
     return success iter
   except CancelledError as exc:
     trace "Cancelling thread future!", exc = exc.msg
+    ctx.setCancelled()
     await iterDispose()
     raise exc
 
diff --git a/tests/datastore/testthreadproxyds.nim b/tests/datastore/testthreadproxyds.nim
index 9799411..c7f825d 100644
--- a/tests/datastore/testthreadproxyds.nim
+++ b/tests/datastore/testthreadproxyds.nim
@@ -14,6 +14,7 @@ import pkg/taskpools
 import pkg/questionable/results
 import pkg/chronicles
 import pkg/threading/smartptrs
+import pkg/threading/atomics
 
 import pkg/datastore/fsds
 import pkg/datastore/sql/sqliteds
@@ -176,7 +177,8 @@ suite "Test ThreadDatastore cancelations":
     var
       signal = ThreadSignalPtr.new().tryGet()
       ms {.global.}: MutexSignal
-      flag {.global.}: int = 0
+      flag {.global.}: Atomic[bool]
+      ready {.global.}: Atomic[bool]
 
     ms.init()
 
@@ -186,12 +188,25 @@ suite "Test ThreadDatastore cancelations":
 
     proc `=destroy`(obj: var TestValue) =
       echo "destroy TestObj!"
-      flag = 10
+      flag.store(true)
+
+    proc wait(flag: var Atomic[bool]) =
+      echo "wait for task to be ready..."
+      defer: echo ""
+      for i in 1..100:
+        stdout.write(".")
+        if flag.load() == true: 
+          return
+        os.sleep(10)
+      raise newException(Defect, "timeout")
 
     proc errorTestTask(ctx: TaskCtx[ThreadTestInt]) {.gcsafe, nimcall.} =
       executeTask(ctx):
+        echo "task:exec"
         discard ctx[].signal.fireSync()
+        ready.store(true)
         ms.wait()
+        echo "ctx:task: ", ctx[]
         (?!ThreadTestInt).ok(default(ThreadTestInt))
 
     proc runTestTask() {.async.} =
@@ -199,17 +214,21 @@ suite "Test ThreadDatastore cancelations":
       let ctx = newTaskCtx(ThreadTestInt, signal=signal)
       dispatchTask(sds, signal):
         sds.tp.spawn errorTestTask(ctx)
-      
-      echo "raise error"
-      raise newException(ValueError, "fake error")
+        ready.wait()
+        echo "raise error"
+        raise newException(ValueError, "fake error")
 
     try:
-      await runTestTask()
+      block:
+        await runTestTask()
     except CatchableError as exc:
       echo "caught: ", $exc
     finally:
       echo "finish"
-      ms.fire()
-      os.sleep(10)
-      check flag == 10
+      check ready.load() == true
+
+      ms.fire()
+      GC_fullCollect()
+      flag.wait()
+      check flag.load() == true