Use HashItem rather than String key

Signed-off-by: Francisco Javier Tirado Sarti <ftirados@ibm.com>

Author: fjtirado

impl/persistence/api/src/main/java/io/serverlessworkflow/impl/persistence/hashing/HashMappingCoordinator.java

@@ -52,15 +52,14 @@ public boolean isTransient() {
     }
   }
 
-  private record PendingWrite(String key, HashIndex index, BytesWithFlag bytes) {}
-
-  private Map<String, Map<String, Map<HashIndex, BytesWithFlag>>> instanceMaps = new HashMap<>();
+  private record PendingWrite(HashItem key, HashIndex index, BytesWithFlag bytes) {}
 
   private final Supplier<HashIndex> hashIndexSupplier;
   private final LRUCache<String, Map<String, Map<HashIndex, BytesWithFlag>>> mappingInfo;
 
   private final Function<String, Map<String, Map<HashIndex, byte[]>>> retriever;
   private final Consumer<Map<String, List<HashMappingInfo>>> writer;
+  private final Map<String, List<PendingWrite>> pending = new HashMap<>();
 
   protected HashMappingCoordinator(
       LRUCache<String, Map<String, Map<HashIndex, BytesWithFlag>>> mappingInfo,
@@ -73,58 +72,53 @@ protected HashMappingCoordinator(
     this.writer = writer;
   }
 
-  private Map<String, List<PendingWrite>> pending = new HashMap<>();
-
   private HashMappingInfo from(PendingWrite pending) {
     return new HashMappingInfo(pending.key, pending.index, pending.bytes.bytes);
   }
 
   public HashIndex calculateIndex(String instanceId, HashItem item, byte[] bytes) {
-    Map<String, Map<HashIndex, BytesWithFlag>> instanceMap = getCachedInstanceMap(instanceId);
+    List<PendingWrite> list =
+        pending.computeIfAbsent(
+            instanceId,
+            k -> {
+              mappingInfo.pin(k);
+              return new ArrayList<>();
+            });
+    Map<String, Map<HashIndex, BytesWithFlag>> instanceMap = getInstanceMap(instanceId);
     String key = item.key();
     synchronized (instanceMap) {
       Map<HashIndex, BytesWithFlag> duplicateMap =
           instanceMap.computeIfAbsent(key, __ -> new HashMap<>());
-
       for (Entry<HashIndex, BytesWithFlag> entry : duplicateMap.entrySet()) {
         if (Arrays.equals(entry.getValue().bytes, bytes)) {
           if (entry.getValue().isTransient()) {
-            addWrite(instanceId, key, entry.getKey(), entry.getValue());
+            addWrite(list, item, entry.getKey(), entry.getValue());
           }
           return entry.getKey();
         }
       }
       HashIndex index = hashIndexSupplier.get();
       BytesWithFlag bytesWithFlag = new BytesWithFlag(bytes);
       duplicateMap.put(index, bytesWithFlag);
-      addWrite(instanceId, key, index, bytesWithFlag);
+      addWrite(list, item, index, bytesWithFlag);
       return index;
     }
   }
 
-  private void addWrite(String instanceId, String key, HashIndex index, BytesWithFlag bytes) {
-    pending
-        .computeIfAbsent(instanceId, __ -> new ArrayList<>())
-        .add(new PendingWrite(key, index, bytes));
+  private void addWrite(
+      List<PendingWrite> list, HashItem key, HashIndex index, BytesWithFlag bytes) {
+    list.add(new PendingWrite(key, index, bytes));
   }
 
   public Optional<byte[]> readBytes(String instanceId, HashItem item, HashIndex index) {
     Map<String, Map<HashIndex, BytesWithFlag>> instanceMap = getInstanceMap(instanceId);
-    try {
-      synchronized (instanceMap) {
-        return Optional.ofNullable(instanceMap.get(item.key()))
-            .map(m -> m.get(index))
-            .map(bytes -> bytes.bytes);
-      }
-    } finally {
-      mappingInfo.unpin(instanceId);
+    synchronized (instanceMap) {
+      return Optional.ofNullable(instanceMap.get(item.key()))
+          .map(m -> m.get(index))
+          .map(bytes -> bytes.bytes);
     }
   }
 
-  private Map<String, Map<HashIndex, BytesWithFlag>> getCachedInstanceMap(String instanceId) {
-    return instanceMaps.computeIfAbsent(instanceId, k -> getInstanceMap(instanceId));
-  }
-
   private Map<String, Map<HashIndex, BytesWithFlag>> getInstanceMap(String instanceId) {
     return mappingInfo.computeIfAbsent(
         instanceId,
@@ -146,7 +140,6 @@ public void persist() {
     for (Map.Entry<String, List<PendingWrite>> item : pending.entrySet()) {
       String instanceId = item.getKey();
       List<HashMappingInfo> list = new ArrayList<>();
-      result.put(instanceId, list);
       Map<String, Map<HashIndex, BytesWithFlag>> instanceMap = mappingInfo.get(instanceId);
       if (instanceMap != null) {
         synchronized (instanceMap) {
@@ -159,6 +152,9 @@ public void persist() {
                   });
         }
       }
+      if (!list.isEmpty()) {
+        result.put(instanceId, list);
+      }
     }
     writer.accept(result);
   }
@@ -181,14 +177,12 @@ public void afterRollback() {
   }
 
   private void cleanUp() {
-    instanceMaps.keySet().forEach(mappingInfo::unpin);
-    instanceMaps.clear();
+    pending.keySet().forEach(mappingInfo::unpin);
     pending.clear();
   }
 
   public void afterRemove(String instanceId) {
     pending.remove(instanceId);
-    instanceMaps.remove(instanceId);
     mappingInfo.remove(instanceId);
   }
 }

impl/persistence/api/src/main/java/io/serverlessworkflow/impl/persistence/hashing/HashMappingInfo.java

@@ -15,4 +15,4 @@
  */
 package io.serverlessworkflow.impl.persistence.hashing;
 
-public record HashMappingInfo(String key, HashIndex index, byte[] bytes) {}
+public record HashMappingInfo(HashItem item, HashIndex index, byte[] bytes) {}