From 3c4daf41983747800f15d24f9cc076cfb4db6116 Mon Sep 17 00:00:00 2001
From: Tomasz Bekas <tomasz.bekas@gmail.com>
Date: Wed, 10 Jan 2024 16:54:14 +0100
Subject: [PATCH] Fix review comments

---
 .github/workflows/tests.yml           |  2 +-
 datastore/datastore.nim               |  2 +-
 datastore/defaultimpl.nim             |  2 +-
 datastore/fsds.nim                    | 26 ++++++++++++++++++++------
 tests/datastore/modifycommontests.nim |  8 +++++---
 5 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 8221127..73df27f 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -13,7 +13,7 @@ jobs:
       fail-fast: false
       matrix:
         cache_nonce: [ 1 ]
-        nim_version: [ 1.6.16 ]
+        nim_version: [ 1.6.18 ]
         platform:
           - {
             icon: 🐧,
diff --git a/datastore/datastore.nim b/datastore/datastore.nim
index 6079468..f664ec0 100644
--- a/datastore/datastore.nim
+++ b/datastore/datastore.nim
@@ -13,7 +13,7 @@ push: {.upraises: [].}
 
 type
   BatchEntry* = tuple[key: Key, data: seq[byte]]
-  Function*[T, U] = proc(value: T): U {.upraises: [CatchableError], gcsafe, closure.}
+  Function*[T, U] = proc(value: T): U {.raises: [CatchableError], gcsafe, closure.}
   Modify* = Function[?seq[byte], Future[?seq[byte]]]
   ModifyGet* = Function[?seq[byte], Future[(?seq[byte], seq[byte])]]
 
diff --git a/datastore/defaultimpl.nim b/datastore/defaultimpl.nim
index 05f1661..1562263 100644
--- a/datastore/defaultimpl.nim
+++ b/datastore/defaultimpl.nim
@@ -46,7 +46,7 @@ proc defaultModifyGetImpl*(
   finally:
     lock.release()
 
-method defaultModifyImpl*(
+proc defaultModifyImpl*(
   self: Datastore,
   lock: AsyncLock,
   key: Key,
diff --git a/datastore/fsds.nim b/datastore/fsds.nim
index 038671b..2a58a99 100644
--- a/datastore/fsds.nim
+++ b/datastore/fsds.nim
@@ -1,5 +1,6 @@
 import std/os
 import std/options
+import std/tables
 import std/strutils
 
 import pkg/chronos
@@ -20,7 +21,7 @@ type
     root*: string
     ignoreProtected: bool
     depth: int
-    lock: AsyncLock
+    locks: TableRef[Key, AsyncLock]
 
 proc validDepth*(self: FSDatastore, key: Key): bool =
   key.len <= self.depth
@@ -222,14 +223,26 @@ method query*(
 method modifyGet*(
   self: FSDatastore,
   key: Key,
-  fn: ModifyGet): Future[?!seq[byte]] =
-  defaultModifyGetImpl(self, self.lock, key, fn)
+  fn: ModifyGet): Future[?!seq[byte]] {.async.} =
+  var lock: AsyncLock
+  try:
+    lock = self.locks.mgetOrPut(key, newAsyncLock())
+    return await defaultModifyGetImpl(self, lock, key, fn)
+  finally:
+    if not lock.locked:
+      self.locks.del(key)
 
 method modify*(
   self: FSDatastore,
   key: Key,
-  fn: Modify): Future[?!void] =
-  defaultModifyImpl(self, self.lock, key, fn)
+  fn: Modify): Future[?!void] {.async.} =
+  var lock: AsyncLock
+  try:
+    lock = self.locks.mgetOrPut(key, newAsyncLock())
+    return await defaultModifyImpl(self, lock, key, fn)
+  finally:
+    if not lock.locked:
+      self.locks.del(key)
 
 proc new*(
   T: type FSDatastore,
@@ -250,4 +263,5 @@ proc new*(
     root: root,
     ignoreProtected: ignoreProtected,
     depth: depth,
-    lock: newAsyncLock())
+    locks: newTable[Key, AsyncLock]()
+    )
diff --git a/tests/datastore/modifycommontests.nim b/tests/datastore/modifycommontests.nim
index 1a38514..b7f4871 100644
--- a/tests/datastore/modifycommontests.nim
+++ b/tests/datastore/modifycommontests.nim
@@ -19,7 +19,9 @@ proc modifyTests*(
 
   randomize()
 
-  let processCount = 100
+  let
+    processCount = 100
+    timeout = (1 + processCount div 10).seconds
 
   proc withRandDelay(op: Future[?!void]): Future[void] {.async: (raises: [Exception]).} =
     await sleepAsync(rand(processCount).millis)
@@ -53,7 +55,7 @@ proc modifyTests*(
         return success()
 
     let futs = newSeqWith(processCount, withRandDelay(getIncAndPut()))
-    await allFutures(futs).wait(10.seconds)
+    await allFutures(futs).wait(timeout)
 
     let finalValue = uint64.fromBytes((await ds.get(key)).tryGet)
 
@@ -63,7 +65,7 @@ proc modifyTests*(
     (await ds.put(key, @(0.uint64.toBytes))).tryGet
 
     let futs = newSeqWith(processCount, withRandDelay(ds.modify(key, incAsyncFn)))
-    await allFutures(futs).wait(10.seconds)
+    await allFutures(futs).wait(timeout)
 
     let finalValue = uint64.fromBytes((await ds.get(key)).tryGet)