[Fix #727] Adding persistence support

Signed-off-by: fjtirado <[email protected]>

Author: fjtirado

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

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2020-Present The Serverless Workflow Specification Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.serverlessworkflow.impl.persistence;
+
+import io.serverlessworkflow.impl.WorkflowModel;
+import java.time.Instant;
+
+public record CompletedTaskInfo(
+    Instant instant,
+    WorkflowModel model,
+    WorkflowModel context,
+    Boolean isEndNode,
+    String nextPosition)
+    implements PersistenceTaskInfo {}

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

@@ -32,6 +32,8 @@ public interface PersistenceInstanceWriter extends AutoCloseable {
 
   void resumed(WorkflowContextData workflowContext);
 
+  void taskRetried(WorkflowContextData workflowContext, TaskContextData taskContext);
+
   void taskStarted(WorkflowContextData workflowContext, TaskContextData taskContext);
 
   void taskCompleted(WorkflowContextData workflowContext, TaskContextData taskContext);

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

@@ -15,12 +15,4 @@
  */
 package io.serverlessworkflow.impl.persistence;
 
-import io.serverlessworkflow.impl.WorkflowModel;
-import java.time.Instant;
-
-public record PersistenceTaskInfo(
-    Instant instant,
-    WorkflowModel model,
-    WorkflowModel context,
-    Boolean isEndNode,
-    String nextPosition) {}
+public interface PersistenceTaskInfo {}

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

@@ -0,0 +1,18 @@
+/*
+ * Copyright 2020-Present The Serverless Workflow Specification Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.serverlessworkflow.impl.persistence;
+
+public record RetriedTaskInfo(short retryAttempt) implements PersistenceTaskInfo {}

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

@@ -47,16 +47,20 @@ public CompletableFuture<WorkflowModel> start() {
   @Override
   public void restoreContext(WorkflowContext workflow, TaskContext context) {
     PersistenceTaskInfo taskInfo = info.tasks().remove(context.position().jsonPointer());
-    if (taskInfo != null) {
-      context.output(taskInfo.model());
-      context.completedAt(taskInfo.instant());
+    if (taskInfo instanceof CompletedTaskInfo completedTaskInfo) {
+      context.output(completedTaskInfo.model());
+      context.completedAt(completedTaskInfo.instant());
       context.transition(
           new TransitionInfo(
-              taskInfo.nextPosition() == null
+              completedTaskInfo.nextPosition() == null
                   ? null
-                  : workflow.definition().taskExecutor(taskInfo.nextPosition()),
-              taskInfo.isEndNode()));
-      workflow.context(taskInfo.context());
+                  : workflow.definition().taskExecutor(completedTaskInfo.nextPosition()),
+              completedTaskInfo.isEndNode()));
+      workflow.context(completedTaskInfo.context());
+    } else if (taskInfo instanceof RetriedTaskInfo retriedTaskInfo) {
+      if (context.retryAttempt() == 0) {
+        context.retryAttempt(retriedTaskInfo.retryAttempt());
+      }
     }
   }
 }

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

@@ -18,6 +18,7 @@
 import static io.serverlessworkflow.impl.WorkflowUtils.safeClose;
 
 import io.serverlessworkflow.impl.lifecycle.TaskCompletedEvent;
+import io.serverlessworkflow.impl.lifecycle.TaskRetriedEvent;
 import io.serverlessworkflow.impl.lifecycle.TaskStartedEvent;
 import io.serverlessworkflow.impl.lifecycle.WorkflowCancelledEvent;
 import io.serverlessworkflow.impl.lifecycle.WorkflowCompletedEvent;
@@ -75,6 +76,11 @@ public void onTaskCompleted(TaskCompletedEvent ev) {
     persistenceWriter.taskCompleted(ev.workflowContext(), ev.taskContext());
   }
 
+  @Override
+  public void onTaskRetried(TaskRetriedEvent ev) {
+    persistenceWriter.taskRetried(ev.workflowContext(), ev.taskContext());
+  }
+
   public void close() {
     safeClose(persistenceWriter);
   }

impl/persistence/bigmap/src/main/java/io/serverlessworkflow/impl/persistence/bigmap/BigMapInstanceWriter.java

@@ -68,6 +68,16 @@ public void aborted(WorkflowContextData workflowContext) {
   @Override
   public void taskStarted(WorkflowContextData workflowContext, TaskContextData taskContext) {}
 
+  @Override
+  public void taskRetried(WorkflowContextData workflowContext, TaskContextData taskContext) {
+    doTransaction(
+        t ->
+            t.tasks(key(workflowContext))
+                .put(
+                    taskContext.position().jsonPointer(),
+                    marshallTaskRetried(workflowContext, (TaskContext) taskContext)));
+  }
+
   @Override
   public void taskCompleted(WorkflowContextData workflowContext, TaskContextData taskContext) {
     doTransaction(
@@ -108,5 +118,8 @@ protected void removeProcessInstance(WorkflowContextData workflowContext) {
   protected abstract T marshallTaskCompleted(
       WorkflowContextData workflowContext, TaskContext taskContext);
 
+  protected abstract T marshallTaskRetried(
+      WorkflowContextData workflowContext, TaskContext taskContext);
+
   protected abstract S marshallStatus(WorkflowStatus status);
 }

impl/persistence/bigmap/src/main/java/io/serverlessworkflow/impl/persistence/bigmap/BytesMapInstanceReader.java

@@ -19,7 +19,9 @@
 import io.serverlessworkflow.impl.WorkflowStatus;
 import io.serverlessworkflow.impl.marshaller.WorkflowBufferFactory;
 import io.serverlessworkflow.impl.marshaller.WorkflowInputBuffer;
+import io.serverlessworkflow.impl.persistence.CompletedTaskInfo;
 import io.serverlessworkflow.impl.persistence.PersistenceTaskInfo;
+import io.serverlessworkflow.impl.persistence.RetriedTaskInfo;
 import java.io.ByteArrayInputStream;
 import java.time.Instant;
 
@@ -36,21 +38,43 @@ public BytesMapInstanceReader(
   @Override
   protected PersistenceTaskInfo unmarshallTaskInfo(byte[] taskData) {
     try (WorkflowInputBuffer buffer = factory.input(new ByteArrayInputStream(taskData))) {
-      buffer.readByte(); // version byte not used at the moment
-      Instant date = buffer.readInstant();
-      WorkflowModel model = (WorkflowModel) buffer.readObject();
-      WorkflowModel context = (WorkflowModel) buffer.readObject();
-      Boolean isEndNode = null;
-      String nextPosition = null;
-      isEndNode = buffer.readBoolean();
-      boolean hasNext = buffer.readBoolean();
-      if (hasNext) {
-        nextPosition = buffer.readString();
+      byte version = buffer.readByte();
+      switch (version) {
+        case MarshallingUtils.VERSION_0:
+        default:
+          return readVersion0(buffer);
+        case MarshallingUtils.VERSION_1:
+          return readVersion1(buffer);
       }
-      return new PersistenceTaskInfo(date, model, context, isEndNode, nextPosition);
     }
   }
 
+  private PersistenceTaskInfo readVersion1(WorkflowInputBuffer buffer) {
+
+    TaskStatus taskStatus = buffer.readEnum(TaskStatus.class);
+    switch (taskStatus) {
+      case COMPLETED:
+      default:
+        return readVersion0(buffer);
+      case RETRIED:
+        return new RetriedTaskInfo(buffer.readShort());
+    }
+  }
+
+  private PersistenceTaskInfo readVersion0(WorkflowInputBuffer buffer) {
+    Instant date = buffer.readInstant();
+    WorkflowModel model = (WorkflowModel) buffer.readObject();
+    WorkflowModel context = (WorkflowModel) buffer.readObject();
+    Boolean isEndNode = null;
+    String nextPosition = null;
+    isEndNode = buffer.readBoolean();
+    boolean hasNext = buffer.readBoolean();
+    if (hasNext) {
+      nextPosition = buffer.readString();
+    }
+    return new CompletedTaskInfo(date, model, context, isEndNode, nextPosition);
+  }
+
   @Override
   protected PersistenceInstanceInfo unmarshallInstanceInfo(byte[] instanceData) {
     try (WorkflowInputBuffer buffer = factory.input(new ByteArrayInputStream(instanceData))) {

impl/persistence/bigmap/src/main/java/io/serverlessworkflow/impl/persistence/bigmap/BytesMapInstanceWriter.java

@@ -40,7 +40,8 @@ public BytesMapInstanceWriter(
   protected byte[] marshallTaskCompleted(WorkflowContextData contextData, TaskContext taskContext) {
     ByteArrayOutputStream bytes = new ByteArrayOutputStream();
     try (WorkflowOutputBuffer writer = factory.output(bytes)) {
-      writer.writeByte(MarshallingUtils.VERSION_0);
+      writer.writeByte(MarshallingUtils.VERSION_1);
+      writer.writeEnum(TaskStatus.COMPLETED);
       writer.writeInstant(taskContext.completedAt());
       writeModel(writer, taskContext.output());
       writeModel(writer, contextData.context());
@@ -82,4 +83,16 @@ protected byte[] marshallInstance(WorkflowInstanceData instance) {
   protected void writeModel(WorkflowOutputBuffer writer, WorkflowModel model) {
     writer.writeObject(model);
   }
+
+  @Override
+  protected byte[] marshallTaskRetried(
+      WorkflowContextData workflowContext, TaskContext taskContext) {
+    ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+    try (WorkflowOutputBuffer writer = factory.output(bytes)) {
+      writer.writeByte(MarshallingUtils.VERSION_1);
+      writer.writeEnum(TaskStatus.RETRIED);
+      writer.writeShort(taskContext.retryAttempt());
+    }
+    return bytes.toByteArray();
+  }
 }

impl/persistence/bigmap/src/main/java/io/serverlessworkflow/impl/persistence/bigmap/MarshallingUtils.java

@@ -20,4 +20,5 @@ class MarshallingUtils {
   private MarshallingUtils() {}
 
   public static final byte VERSION_0 = 0;
+  public static final byte VERSION_1 = 1;
 }