[Fix #1395] persistence helper classes for allcorrelation

Signed-off-by: fjtirado <ftirados@redhat.com>

Author: fjtirado

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

@@ -16,6 +16,7 @@
 package io.serverlessworkflow.impl;
 
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Supplier;
 
 public interface WorkflowInstance extends WorkflowInstanceData {
   CompletableFuture<WorkflowModel> start();
@@ -49,4 +50,6 @@ public interface WorkflowInstance extends WorkflowInstanceData {
   boolean cancel();
 
   boolean resume();
+
+  <T> T addMetadataIfAbsent(String key, Supplier<T> supplier);
 }

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

@@ -16,6 +16,7 @@
 package io.serverlessworkflow.impl;
 
 import java.time.Instant;
+import java.util.Optional;
 
 public interface WorkflowInstanceData {
   String id();
@@ -29,4 +30,6 @@ public interface WorkflowInstanceData {
   WorkflowStatus status();
 
   WorkflowModel context();
+
+  <T> Optional<T> findMetadata(String key, Class<T> objectClass);
 }

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

@@ -350,9 +350,16 @@ public void addCancelable(CompletableFuture<?> cancelable) {
     }
   }
 
-  public <T> T additionalObject(String key, Supplier<T> supplier) {
+  @Override
+  public <T> T addMetadataIfAbsent(String key, Supplier<T> supplier) {
     return (T) additionalObjects.computeIfAbsent(key, k -> supplier.get());
   }
 
+  @Override
+  public <T> Optional<T> findMetadata(String key, Class<T> objectClass) {
+    Object value = additionalObjects.get(key);
+    return objectClass.isInstance(value) ? Optional.of(objectClass.cast(value)) : Optional.empty();
+  }
+
   public void restoreContext(WorkflowContext workflow, TaskContext context) {}
 }

impl/core/src/main/java/io/serverlessworkflow/impl/marshaller/AbstractInputBuffer.java

@@ -15,7 +15,9 @@
  */
 package io.serverlessworkflow.impl.marshaller;
 
+import java.net.URI;
 import java.time.Instant;
+import java.time.OffsetDateTime;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.LinkedHashMap;
@@ -112,6 +114,12 @@ public Object readObject() {
       case CUSTOM:
         return readCustomObject();
 
+      case URI:
+        return URI.create(readString());
+
+      case OFFSET_DATE_TIME:
+        return OffsetDateTime.from(readInstant());
+
       default:
         throw new IllegalStateException("Unsupported type " + type);
     }

impl/core/src/main/java/io/serverlessworkflow/impl/marshaller/AbstractOutputBuffer.java

@@ -16,7 +16,9 @@
 package io.serverlessworkflow.impl.marshaller;
 
 import io.serverlessworkflow.impl.WorkflowModel;
+import java.net.URI;
 import java.time.Instant;
+import java.time.OffsetDateTime;
 import java.util.Collection;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -100,6 +102,12 @@ public WorkflowOutputBuffer writeObject(Object object) {
     } else if (object instanceof Instant value) {
       writeType(Type.INSTANT);
       writeInstant(value);
+    } else if (object instanceof OffsetDateTime time) {
+      writeType(Type.OFFSET_DATE_TIME);
+      writeInstant(time.toInstant());
+    } else if (object instanceof URI uri) {
+      writeType(Type.URI);
+      writeString(uri.toString());
     } else if (object instanceof byte[] bytes) {
       writeType(Type.BYTES);
       writeBytes(bytes);

impl/core/src/main/java/io/serverlessworkflow/impl/marshaller/Type.java

@@ -29,5 +29,7 @@ enum Type {
   MAP,
   COLLECTION,
   NULL,
-  CUSTOM
+  CUSTOM,
+  OFFSET_DATE_TIME,
+  URI
 }

impl/core/src/main/java/io/serverlessworkflow/impl/scheduler/AllStrategyCorrelationInfo.java

@@ -16,15 +16,21 @@
 package io.serverlessworkflow.impl.scheduler;
 
 import io.cloudevents.CloudEvent;
+import io.serverlessworkflow.impl.WorkflowInstance;
 import io.serverlessworkflow.impl.events.EventRegistrationBuilder;
-import java.util.Collection;
+import java.util.Map;
 import java.util.function.Consumer;
 
 public interface AllStrategyCorrelationInfo extends AutoCloseable {
   void correlate(
-      EventRegistrationBuilder reg, CloudEvent event, Consumer<Collection<CloudEvent>> starter);
+      EventRegistrationBuilder reg,
+      CloudEvent event,
+      Consumer<Map<EventRegistrationBuilder, CloudEvent>> starter);
 
   void register(EventRegistrationBuilder reg);
 
+  default void addMetadata(
+      WorkflowInstance instance, Map<EventRegistrationBuilder, CloudEvent> events) {}
+
   default void close() {}
 }

impl/core/src/main/java/io/serverlessworkflow/impl/scheduler/DefaultWorkflowScheduler.java

@@ -16,7 +16,7 @@
 package io.serverlessworkflow.impl.scheduler;
 
 import io.serverlessworkflow.impl.WorkflowDefinition;
-import io.serverlessworkflow.impl.WorkflowModel;
+import io.serverlessworkflow.impl.WorkflowInstance;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
@@ -83,10 +83,10 @@ protected CronResolverIntanceRunner(WorkflowDefinition definition) {
       }
 
       @Override
-      public void accept(WorkflowModel model) {
+      public void accept(WorkflowInstance instance) {
         if (!cancelled.get()) {
           scheduleNext();
-          super.accept(model);
+          super.accept(instance);
         }
       }
     }

impl/core/src/main/java/io/serverlessworkflow/impl/scheduler/InMemoryAllStrategyCorrelationInfo.java

@@ -18,7 +18,6 @@
 import io.cloudevents.CloudEvent;
 import io.serverlessworkflow.impl.events.EventRegistrationBuilder;
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -30,21 +29,23 @@ public class InMemoryAllStrategyCorrelationInfo implements AllStrategyCorrelatio
 
   @Override
   public void correlate(
-      EventRegistrationBuilder reg, CloudEvent event, Consumer<Collection<CloudEvent>> starter) {
-    Collection<CloudEvent> collection = new ArrayList<>();
+      EventRegistrationBuilder reg,
+      CloudEvent event,
+      Consumer<Map<EventRegistrationBuilder, CloudEvent>> starter) {
+    Map<EventRegistrationBuilder, CloudEvent> result = new HashMap<>();
     // to minimize the critical section, conversion is done later, here we are
     // performing just collection, if any
     synchronized (correlatedEvents) {
       correlatedEvents.get(reg).add(event);
-      Collection<List<CloudEvent>> events = correlatedEvents.values();
-      if (satisfyCondition(events)) {
-        for (List<CloudEvent> values : events) {
-          collection.add(values.remove(0));
+      if (satisfyCondition(correlatedEvents)) {
+        for (java.util.Map.Entry<EventRegistrationBuilder, List<CloudEvent>> values :
+            correlatedEvents.entrySet()) {
+          result.put(values.getKey(), values.getValue().remove(0));
         }
       }
     }
-    if (!collection.isEmpty()) {
-      starter.accept(collection);
+    if (!result.isEmpty()) {
+      starter.accept(result);
     }
   }
 
@@ -56,8 +57,8 @@ public void register(EventRegistrationBuilder reg) {
     correlatedEvents.put(reg, new ArrayList<CloudEvent>());
   }
 
-  private boolean satisfyCondition(Collection<List<CloudEvent>> events) {
-    for (List<CloudEvent> values : events) {
+  private boolean satisfyCondition(Map<EventRegistrationBuilder, List<CloudEvent>> events) {
+    for (List<CloudEvent> values : events.values()) {
       if (values.isEmpty()) {
         return false;
       }

impl/core/src/main/java/io/serverlessworkflow/impl/scheduler/ScheduledEventConsumer.java

@@ -19,13 +19,16 @@
 
 import io.cloudevents.CloudEvent;
 import io.serverlessworkflow.impl.WorkflowDefinition;
+import io.serverlessworkflow.impl.WorkflowInstance;
 import io.serverlessworkflow.impl.WorkflowModel;
 import io.serverlessworkflow.impl.WorkflowModelCollection;
 import io.serverlessworkflow.impl.events.EventConsumer;
 import io.serverlessworkflow.impl.events.EventRegistration;
+import io.serverlessworkflow.impl.events.EventRegistrationBuilder;
 import io.serverlessworkflow.impl.events.EventRegistrationBuilderInfo;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Map;
 import java.util.function.Function;
 
 public class ScheduledEventConsumer implements AutoCloseable {
@@ -78,13 +81,15 @@ public ScheduledEventConsumer(
   protected void start(CloudEvent ce) {
     WorkflowModelCollection model = definition.application().modelFactory().createCollection();
     model.add(converter.apply(ce));
-    instanceRunner.accept(model);
+    instanceRunner.accept(definition.instance(model));
   }
 
-  protected void start(Collection<CloudEvent> ces) {
+  protected void start(Map<EventRegistrationBuilder, CloudEvent> ces) {
     WorkflowModelCollection model = definition.application().modelFactory().createCollection();
-    ces.forEach(ce -> model.add(converter.apply(ce)));
-    instanceRunner.accept(model);
+    ces.values().forEach(ce -> model.add(converter.apply(ce)));
+    WorkflowInstance instance = definition.instance(model);
+    allStrategyCorrelationInfo.addMetadata(instance, ces);
+    instanceRunner.accept(instance);
   }
 
   public void close() {