From e0056906426de7980bf3d6e9f7d818e5050e8ffb Mon Sep 17 00:00:00 2001
From: Ben <thatbenbierens@gmail.com>
Date: Tue, 28 May 2024 11:27:09 +0200
Subject: [PATCH] Adds tests for concurrent queries. Finishes disposed
 iterators.

---
 datastore/fsds.nim                   |   4 ++
 datastore/leveldb/leveldbds.nim      |   1 +
 datastore/sql/sqliteds.nim           |   1 +
 tests/datastore/querycommontests.nim | 100 +++++++++++++++++++++++++++
 4 files changed, 106 insertions(+)

diff --git a/datastore/fsds.nim b/datastore/fsds.nim
index 7599fce..c9a3213 100644
--- a/datastore/fsds.nim
+++ b/datastore/fsds.nim
@@ -216,6 +216,10 @@ method query*(
     return success (key.some, data)
 
   iter.next = next
+  iter.dispose = proc(): Future[?!void] {.async.} =
+    iter.finished = true
+    return success()
+  
   return success iter
 
 method modifyGet*(
diff --git a/datastore/leveldb/leveldbds.nim b/datastore/leveldb/leveldbds.nim
index 30e30f6..6d6de46 100644
--- a/datastore/leveldb/leveldbds.nim
+++ b/datastore/leveldb/leveldbds.nim
@@ -116,6 +116,7 @@ method query*(
 
   proc dispose(): Future[?!void] {.async.} =
     dbIter.dispose()
+    iter.finished = true
     return success()
 
   iter.next = next
diff --git a/datastore/sql/sqliteds.nim b/datastore/sql/sqliteds.nim
index 8f16f07..a7ca095 100644
--- a/datastore/sql/sqliteds.nim
+++ b/datastore/sql/sqliteds.nim
@@ -325,6 +325,7 @@ method query*(
     discard sqlite3_reset(s)
     discard sqlite3_clear_bindings(s)
     iter.next = nil
+    iter.finished = true
     return success()
 
   iter.next = next
diff --git a/tests/datastore/querycommontests.nim b/tests/datastore/querycommontests.nim
index c3ffd3b..3c26fa5 100644
--- a/tests/datastore/querycommontests.nim
+++ b/tests/datastore/querycommontests.nim
@@ -44,6 +44,106 @@ template queryTests*(ds: Datastore, testLimitsAndOffsets = true, testSortOrder =
 
     (await iter.dispose()).tryGet
 
+  test "Concurrent queries":
+    let
+      q1 = Query.init(key1)
+      q2 = Query.init(key1)
+
+    (await ds.put(key1, val1)).tryGet
+    (await ds.put(key2, val2)).tryGet
+    (await ds.put(key3, val3)).tryGet
+
+    let
+      iter1 = (await ds.query(q1)).tryGet
+      iter2 = (await ds.query(q1)).tryGet
+
+    let one1 = (await iter1.next()).tryGet
+    check not iter1.finished
+
+    let one2 = (await iter2.next()).tryGet
+    check not iter2.finished
+
+    let two1 = (await iter1.next()).tryGet
+    check not iter1.finished
+    let three1 = (await iter1.next()).tryGet
+    check not iter1.finished
+
+    let two2 = (await iter2.next()).tryGet
+    check not iter2.finished
+    let three2 = (await iter2.next()).tryGet
+    check not iter2.finished
+
+    let four1 = (await iter1.next()).tryGet
+    check iter1.finished
+    let four2 = (await iter2.next()).tryGet
+    check iter2.finished
+
+    check:
+      one1[0].get == key1
+      one1[1] == val1
+      two1[0].get == key2
+      two1[1] == val2
+      three1[0].get == key3
+      three1[1] == val3
+      four1[0] == Key.none
+
+      one2[0].get == key1
+      one2[1] == val1
+      two2[0].get == key2
+      two2[1] == val2
+      three2[0].get == key3
+      three2[1] == val3
+      four2[0] == Key.none
+
+    (await iter1.dispose()).tryGet
+    (await iter2.dispose()).tryGet
+
+  test "Concurrent queries - dispose":
+    let
+      q1 = Query.init(key1)
+      q2 = Query.init(key1)
+
+    (await ds.put(key1, val1)).tryGet
+    (await ds.put(key2, val2)).tryGet
+    (await ds.put(key3, val3)).tryGet
+
+    let
+      iter1 = (await ds.query(q1)).tryGet
+      iter2 = (await ds.query(q1)).tryGet
+
+    let one1 = (await iter1.next()).tryGet
+    check not iter1.finished
+
+    let one2 = (await iter2.next()).tryGet
+    check not iter2.finished
+
+    (await iter1.dispose()).tryGet
+    check iter1.finished
+
+    let two2 = (await iter2.next()).tryGet
+    check not iter2.finished
+    let three2 = (await iter2.next()).tryGet
+    check not iter2.finished
+
+    let four2 = (await iter2.next()).tryGet
+    check iter2.finished
+
+    check:
+      one1[0].get == key1
+      one1[1] == val1
+
+      one2[0].get == key1
+      one2[1] == val1
+      two2[0].get == key2
+      two2[1] == val2
+      three2[0].get == key3
+      three2[1] == val3
+      four2[0] == Key.none
+
+    (await iter1.dispose()).tryGet
+    (await iter2.dispose()).tryGet
+
+
   test "Key should query all keys without values":
     let
       q = Query.init(key1, value = false)