[Fix #1738] Replace setMetadata by fine grain manipulation. (#1742)

* [Fix #1738] Replace setMetadata by fine grain manipulation.

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

* [Fix #1738] Adding safeguard to avoid overriding keys set by user

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

---------

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

Author: fjtirado

impl/core/src/main/java/io/serverlessworkflow/impl/WorkflowMutableInstance.java

@@ -53,7 +53,7 @@ public class WorkflowMutableInstance implements WorkflowInstance {
   protected AtomicReference<CompletableFuture<WorkflowModel>> futureRef = new AtomicReference<>();
   protected Instant completedAt;
 
-  private Map<String, Object> additionalObjects = new ConcurrentHashMap<>();
+  protected Map<String, Object> additionalObjects = new ConcurrentHashMap<>();
 
   protected final Map<String, Integer> iterationsMap = new ConcurrentHashMap<>();
 
@@ -448,10 +448,6 @@ public void addCancelable(CompletableFuture<?> cancelable) {
     }
   }
 
-  protected void setMetadata(Map<String, Object> metadata) {
-    this.additionalObjects = new ConcurrentHashMap<>(metadata);
-  }
-
   @Override
   public <T> T addMetadataIfAbsent(String key, Supplier<T> supplier) {
     return (T) additionalObjects.computeIfAbsent(key, k -> supplier.get());

impl/persistence/api/src/main/java/io/serverlessworkflow/impl/persistence/PersistenceTaskInfo.java

@@ -15,4 +15,9 @@
  */
 package io.serverlessworkflow.impl.persistence;
 
-public interface PersistenceTaskInfo {}
+import java.util.Map;
+
+public interface PersistenceTaskInfo {
+
+  Map<String, Object> additionalObjects();
+}

impl/persistence/api/src/main/java/io/serverlessworkflow/impl/persistence/RetriedTaskInfo.java

@@ -17,7 +17,7 @@
 
 import java.util.Map;
 
-public record RetriedTaskInfo(int retryAttempt, Map<String, Object> metadata)
+public record RetriedTaskInfo(int retryAttempt, Map<String, Object> additionalObjects)
     implements PersistenceTaskInfo {
   public RetriedTaskInfo(int retryAttempt) {
     this(retryAttempt, Map.of());

impl/persistence/api/src/main/java/io/serverlessworkflow/impl/persistence/WorkflowPersistenceInstance.java

@@ -24,12 +24,16 @@
 import io.serverlessworkflow.impl.WorkflowMutableInstance;
 import io.serverlessworkflow.impl.WorkflowStatus;
 import io.serverlessworkflow.impl.executors.TransitionInfo;
+import java.util.HashSet;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Supplier;
 
 public class WorkflowPersistenceInstance extends WorkflowMutableInstance {
 
   private final PersistenceWorkflowInfo info;
+  private Set<String> userMetaKeys = new HashSet<>();
 
   public static WorkflowInstance of(WorkflowDefinition definition, PersistenceWorkflowInfo info) {
     return definition
@@ -48,7 +52,7 @@ private WorkflowPersistenceInstance(WorkflowDefinition definition, PersistenceWo
               }
             });
     this.startedAt = info.startedAt();
-    setMetadata(info.metadata());
+    additionalObjects.putAll(info.metadata());
   }
 
   @Override
@@ -84,7 +88,6 @@ public void restoreContext(WorkflowContext workflow, TaskContext context) {
                   : workflow.definition().taskExecutor(completedTaskInfo.nextPosition()),
               completedTaskInfo.isEndNode()));
       workflow.context(completedTaskInfo.context());
-      setMetadata(completedTaskInfo.additionalObjects());
     } else if (taskInfo instanceof RetriedTaskInfo retriedTaskInfo) {
       if (context.retryAttempt() == 0) {
         context.retryAttempt(retriedTaskInfo.retryAttempt());
@@ -100,7 +103,44 @@ public void restoreContext(WorkflowContext workflow, TaskContext context) {
         }
         searchContext = tryContext.parent();
       }
-      setMetadata(retriedTaskInfo.metadata());
+    }
+    Set<String> retainKeys;
+    synchronized (userMetaKeys) {
+      retainKeys = new HashSet<>(userMetaKeys);
+      taskInfo
+          .additionalObjects()
+          .forEach(
+              (k, v) -> {
+                if (!retainKeys.contains(k)) {
+                  additionalObjects.put(k, v);
+                  retainKeys.add(k);
+                }
+              });
+      additionalObjects.keySet().retainAll(retainKeys);
+    }
+  }
+
+  @Override
+  public <T> T addMetadataIfAbsent(String key, Supplier<T> supplier) {
+    synchronized (userMetaKeys) {
+      userMetaKeys.add(key);
+      return super.addMetadataIfAbsent(key, supplier);
+    }
+  }
+
+  @Override
+  public void removeMetadata(String key) {
+    synchronized (userMetaKeys) {
+      userMetaKeys.add(key);
+      super.removeMetadata(key);
+    }
+  }
+
+  @Override
+  public <T> Optional<T> removeMetadata(String key, Class<T> clazz) {
+    synchronized (userMetaKeys) {
+      userMetaKeys.add(key);
+      return super.removeMetadata(key, clazz);
     }
   }
 }