Add compute() function to ClusterMap

steep-wms · Sep 27, 2024 · 3f845be · 3f845be
1 parent 8ab1b04
commit 3f845be
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 4 deletions.
diff --git a/src/main/kotlin/helper/hazelcast/ClusterMap.kt b/src/main/kotlin/helper/hazelcast/ClusterMap.kt
@@ -20,7 +20,8 @@ interface ClusterMap<K : Any, V : Any> {
   suspend fun put(key: K, value: V): V?
   suspend fun putIfAbsent(key: K, value: V): V?
   suspend fun computeIfAbsent(key: K, mappingFunction: (K) -> V): V?
-  suspend fun computeIfPresent(key: K, remappingFunction: (K, V) -> V): V?
+  suspend fun computeIfPresent(key: K, remappingFunction: (K, V) -> V?): V?
+  suspend fun compute(key: K, remappingFunction: (K, V?) -> V?): V?
   suspend fun delete(key: K)
   suspend fun entries(): Set<Map.Entry<K, V>>
   suspend fun keys(): Set<K>

diff --git a/src/main/kotlin/helper/hazelcast/ClusterMapImpl.kt b/src/main/kotlin/helper/hazelcast/ClusterMapImpl.kt
@@ -46,12 +46,18 @@ class ClusterMapImpl<K : Any, V : Any>(private val map: IMap<K, V>,
     }, false).coAwait()
   }
 
-  override suspend fun computeIfPresent(key: K, remappingFunction: (K, V) -> V): V? {
+  override suspend fun computeIfPresent(key: K, remappingFunction: (K, V) -> V?): V? {
     return vertx.executeBlocking({ p ->
       p.complete(map.computeIfPresent(key, remappingFunction))
     }, false).coAwait()
   }
 
+  override suspend fun compute(key: K, remappingFunction: (K, V?) -> V?): V? {
+    return vertx.executeBlocking({ p ->
+      p.complete(map.compute(key, remappingFunction))
+    }, false).coAwait()
+  }
+
   override suspend fun delete(key: K) {
     return vertx.executeBlocking({ p ->
       p.complete(map.delete(key))

diff --git a/src/test/kotlin/helper/hazelcast/DummyClusterMap.kt b/src/test/kotlin/helper/hazelcast/DummyClusterMap.kt
@@ -81,7 +81,37 @@ class DummyClusterMap<K : Any, V : Any>(name: String, vertx: Vertx) : ClusterMap
     return r
   }
 
-  override suspend fun computeIfPresent(key: K, remappingFunction: (K, V) -> V): V? {
-    return map.computeIfPresent(key, remappingFunction)
+  override suspend fun computeIfPresent(key: K, remappingFunction: (K, V) -> V?): V? {
+    return map.computeIfPresent(key) { _, oldValue ->
+      val newValue = remappingFunction(key, oldValue)
+      if (newValue == null) {
+        for (l in entryRemovedListeners) {
+          context.runOnContext {
+            l(key)
+          }
+        }
+      }
+      newValue
+    }
+  }
+
+  override suspend fun compute(key: K, remappingFunction: (K, V?) -> V?): V? {
+    return map.compute(key) { _, oldValue ->
+      val newValue = remappingFunction(key, oldValue)
+      if (oldValue == null && newValue != null) {
+        for ((l, includeValue) in entryAddedListeners) {
+          context.runOnContext {
+            l(key, if (includeValue) newValue else null)
+          }
+        }
+      } else if (oldValue != null && newValue == null) {
+        for (l in entryRemovedListeners) {
+          context.runOnContext {
+            l(key)
+          }
+        }
+      }
+      newValue
+    }
   }
 }