change ResumingTasks to ArchivedTasks

st-tech · Feb 15, 2024 · 6f60f54 · 6f60f54
1 parent fdf7462
commit 6f60f54
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 26 deletions.
diff --git a/digdag-core/src/main/java/io/digdag/core/workflow/TaskControl.java b/digdag-core/src/main/java/io/digdag/core/workflow/TaskControl.java
@@ -7,6 +7,7 @@
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.stream.Collectors;
+
 import com.google.common.base.Optional;
 import com.google.common.base.Strings;
 import com.google.common.collect.*;
@@ -56,14 +57,12 @@ public TaskStateCode getState()
 
     public static long addInitialTasksExceptingRootTask(
             TaskControlStore store, long attemptId, long rootTaskId,
-            WorkflowTaskList tasks, List<ResumingTask> resumingTasks, Limits limits)
+            WorkflowTaskList tasks, List<ArchivedTask> archivedTasks, Limits limits)
         throws TaskLimitExceededException
     {
         checkTaskLimit(store, attemptId, tasks, limits);
-        long taskId = addTasks(store, attemptId, rootTaskId,
-                tasks, ImmutableList.of(),
-                false, true, true,
-                resumingTasks);
+        long taskId = addInitialTasks(store, attemptId, rootTaskId, tasks, archivedTasks);
+        List<ResumingTask> resumingTasks = archivedTasks.stream().map(ResumingTask::of).collect(Collectors.toList());
         addResumingTasks(store, attemptId, resumingTasks);
         return taskId;
     }
@@ -322,28 +321,24 @@ private List<ResumingTask> collectResumingTasks(long attemptId, WorkflowTaskList
         return store.getResumingTasksByNamePrefix(attemptId, commonPrefix);
     }
 
-    static List<ResumingTask> buildResumingTaskMap(SessionStore store, long attemptId, List<Long> resumingTaskIds)
+    static List<ArchivedTask> buildResumingTaskMap(SessionStore store, long attemptId, List<Long> resumingTaskIds)
             throws ResourceNotFoundException
     {
         Set<Long> idSet = new HashSet<>(resumingTaskIds);
-        List<ResumingTask> resumingTasks = store
+        List<ArchivedTask> archivedTasks = store
             .getTasksOfAttempt(attemptId)
             .stream()
             .filter(archived -> {
                 if (idSet.remove(archived.getId())) {
-                    if (archived.getState() != TaskStateCode.SUCCESS) {
-                        throw new IllegalResumeException("Resuming non-successful tasks is not allowed: task_id=" + archived.getId());
-                    }
                     return true;
                 }
                 return false;
             })
-            .map(archived -> ResumingTask.of(archived))
             .collect(Collectors.toList());
         if (!idSet.isEmpty()) {
             throw new ResourceNotFoundException("Resuming tasks are not the members of resuming attempt: id list=" + idSet);
         }
-        return resumingTasks;
+        return archivedTasks;
     }
 
     ////

diff --git a/digdag-core/src/main/java/io/digdag/core/workflow/WorkflowExecutor.java b/digdag-core/src/main/java/io/digdag/core/workflow/WorkflowExecutor.java
@@ -21,6 +21,7 @@
 import io.digdag.core.repository.StoredRevision;
 import io.digdag.core.repository.WorkflowDefinition;
 import io.digdag.core.session.ResumingTask;
+import io.digdag.core.session.ArchivedTask;
 import io.digdag.core.session.ParameterUpdate;
 import io.digdag.core.session.Session;
 import io.digdag.core.session.SessionAttempt;
@@ -43,6 +44,8 @@
 import io.digdag.spi.TaskConflictException;
 import io.digdag.spi.TaskNotFoundException;
 import io.digdag.spi.metrics.DigdagMetrics;
+import io.digdag.spi.metrics.DigdagMetrics.Category;
+
 import static io.digdag.spi.metrics.DigdagMetrics.Category;
 import io.digdag.util.RetryControl;
 import org.slf4j.Logger;
@@ -269,21 +272,15 @@ public StoredSessionAttemptWithSession submitTasks(int siteId, AttemptRequest ar
 
         TaskConfig.validateAttempt(attempt);
 
-        List<ResumingTask> resumingTasks;
+        List<ArchivedTask> archivedTasks;
         if (ar.getResumingAttemptId().isPresent()) {
-            WorkflowTask root = tasks.get(0);
-            resumingTasks = TaskControl.buildResumingTaskMap(
+            archivedTasks = TaskControl.buildResumingTaskMap(
                     sm.getSessionStore(siteId),
                     ar.getResumingAttemptId().get(),
                     ar.getResumingTasks());
-            for (ResumingTask resumingTask : resumingTasks) {
-                if (resumingTask.getFullName().equals(root.getFullName())) {
-                    throw new IllegalResumeException("Resuming root task is not allowed");
-                }
-            }
         }
         else {
-            resumingTasks = ImmutableList.of();
+            archivedTasks = ImmutableList.of();
         }
 
         StoredSessionAttemptWithSession stored;
@@ -313,7 +310,7 @@ public StoredSessionAttemptWithSession submitTasks(int siteId, AttemptRequest ar
                         StoredSessionAttemptWithSession.of(siteId, storedSession, storedAttempt);
 
                     try {
-                        storeTasks(store, storedAttemptWithSession, tasks, resumingTasks, ar.getSessionMonitors());
+                        storeTasks(store, storedAttemptWithSession, tasks, archivedTasks, ar.getSessionMonitors());
                     }
                     catch (TaskLimitExceededException ex) {
                         throw new WorkflowTaskLimitExceededException(ex);
@@ -348,21 +345,21 @@ public void storeTasks(
             SessionControlStore store,
             StoredSessionAttemptWithSession storedAttempt,
             WorkflowDefinition def,
-            List<ResumingTask> resumingTasks,
+            List<ArchivedTask> archivedTasks,
             List<SessionMonitor> sessionMonitors)
         throws TaskLimitExceededException
     {
         Workflow workflow = compiler.compile(def.getName(), def.getConfig());
         WorkflowTaskList tasks = workflow.getTasks();
 
-        storeTasks(store, storedAttempt, tasks, resumingTasks, sessionMonitors);
+        storeTasks(store, storedAttempt, tasks, archivedTasks, sessionMonitors);
     }
 
     public void storeTasks(
             SessionControlStore store,
             StoredSessionAttemptWithSession storedAttempt,
             WorkflowTaskList tasks,
-            List<ResumingTask> resumingTasks,
+            List<ArchivedTask> archivedTasks,
             List<SessionMonitor> sessionMonitors)
         throws TaskLimitExceededException
     {
@@ -386,7 +383,7 @@ public void storeTasks(
             store.insertRootTask(storedAttempt.getId(), rootTask, (taskStore, storedTaskId) -> {
                 try {
                     TaskControl.addInitialTasksExceptingRootTask(taskStore, storedAttempt.getId(),
-                            storedTaskId, tasks, resumingTasks, limits);
+                            storedTaskId, tasks, archivedTasks, limits);
                 }
                 catch (TaskLimitExceededException ex) {
                     throw new WorkflowTaskLimitExceededException(ex);