Fix polling (ECS) command executors fail to run tasks on retry (#1620) (

#1632) Signed-off-by: Naoto Yokoyama <builtinnya@gmail.com> Co-authored-by: Naoto Yokoyama <builtinnya@gmail.com>
treasure-data · Sep 3, 2021 · 93da2fc · 93da2fc
1 parent e97e562
commit 93da2fc
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 1 deletion.
diff --git a/digdag-plugin-utils/src/main/java/io/digdag/util/BaseOperator.java b/digdag-plugin-utils/src/main/java/io/digdag/util/BaseOperator.java
@@ -51,9 +51,11 @@ public TaskResult run()
 
             boolean doRetry = retry.evaluate();
             if (doRetry) {
+                // Remove command status so that pollable command executors can run tasks on retry
+                Config nextStateParams = retry.getNextRetryStateParams().remove("commandStatus");
                 throw TaskExecutionException.ofNextPollingWithCause(ex,
                         retry.getNextRetryInterval(),
-                        ConfigElement.copyOf(retry.getNextRetryStateParams()));
+                        ConfigElement.copyOf(nextStateParams));
             }
             else {
                 throw ex;

diff --git a/digdag-plugin-utils/src/test/java/io/digdag/util/BaseOperatorTest.java b/digdag-plugin-utils/src/test/java/io/digdag/util/BaseOperatorTest.java
@@ -1,6 +1,8 @@
 package io.digdag.util;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+
 import io.digdag.client.config.Config;
 import io.digdag.client.config.ConfigElement;
 import io.digdag.client.config.ConfigFactory;
@@ -102,6 +104,37 @@ public TaskResult runTask()
         }
     }
 
+    @Test
+    public void verifyNonPollingRuntimeExceptionRemovesCommandStatusOnRetry()
+            throws Exception
+    {
+        config.set("_retry", 10);
+
+        ObjectNode commandStatus = new ObjectMapper().createObjectNode().put("foo", "bar");
+        state.set("commandStatus", commandStatus);
+
+        RuntimeException ex = new RuntimeException("Command failed");
+
+        BaseOperator op = new BaseOperator(newContext(temporaryFolder.getRoot().toPath(), request))
+        {
+            @Override
+            public TaskResult runTask()
+            {
+                throw ex;
+            }
+        };
+
+        try {
+            op.run();
+            fail();
+        }
+        catch (TaskExecutionException e) {
+            assertThat(e.isError(), is(true));
+            assertThat(e.getCause(), is(ex));
+            assertThat(e.getStateParams(configFactory).get().has("commandStatus"), is(false));
+        }
+    }
+
     private OperatorContext newContext(final Path projectPath, final TaskRequest taskRequest)
     {
         return new OperatorContext()