Merge pull request #373 from FlowCI/feature/1498

Feature/1498

Author: gy2006

core/src/main/java/com/flowci/core/agent/domain/ShellIn.java

@@ -44,7 +44,9 @@ public enum ShellType {
 
     private List<String> pwsh;
 
-    private int timeout = 1800;
+    private int timeout = 1800; // from StepNode.timeout
+
+    private int retry; // from StepNode.retry
 
     private Vars<String> inputs;
 

core/src/main/java/com/flowci/core/agent/manager/AgentEventManager.java

@@ -54,20 +54,19 @@ public <T> void writeMessage(String token, ResponseMessage<T> msg) {
 
         try {
             byte[] bytes = objectMapper.writeValueAsBytes(msg);
-            session.sendMessage(new BinaryMessage(bytes));
+            writeMessage(session, bytes);
         } catch (IOException e) {
             log.warn("Unable to write response message for agent {}: {}", token, e.getMessage());
         }
     }
 
-    public void writeMessage(String token, byte[] bytes) throws IOException {
+    public void writeMessage(String token, byte[] bytes) {
         WebSocketSession session = agentSessionStore.get(token);
         if (session == null) {
             log.warn("Agent {} not connected", token);
             return;
         }
-
-        session.sendMessage(new BinaryMessage(bytes));
+        writeMessage(session, bytes);
     }
 
     @Override
@@ -123,12 +122,22 @@ public void afterConnectionClosed(WebSocketSession session, CloseStatus status)
     private <T> void writeMessage(WebSocketSession session, ResponseMessage<T> msg) {
         try {
             byte[] bytes = objectMapper.writeValueAsBytes(msg);
-            session.sendMessage(new BinaryMessage(bytes));
+            writeMessage(session, bytes);
         } catch (IOException e) {
             log.warn(e);
         }
     }
 
+    private void writeMessage(WebSocketSession session, byte[] body) {
+        synchronized (session) {
+            try {
+                session.sendMessage(new BinaryMessage(body));
+            } catch (IOException e) {
+                log.warn(e);
+            }
+        }
+    }
+
     private void onConnected(WebSocketSession session, String token, byte[] body) {
         try {
             AgentInit init = objectMapper.readValue(body, AgentInit.class);

core/src/main/java/com/flowci/core/job/manager/CmdManagerImpl.java

@@ -64,7 +64,8 @@ public CmdIn createShellCmd(Job job, Step step, NodeTree tree) {
                 .setPwsh(linkScript(node, ShellIn.ShellType.PowerShell))
                 .setEnvFilters(linkFilters(node))
                 .setInputs(linkInputs(node).merge(job.getContext(), false))
-                .setTimeout(job.getTimeout());
+                .setTimeout(linkTimeout(node, job.getTimeout()))
+                .setRetry(linkRetry(node, 0));
 
         if (node.hasPlugin()) {
             setPlugin(node.getPlugin(), in);
@@ -112,6 +113,36 @@ private StringVars linkInputs(Node current) {
         return output;
     }
 
+    private Integer linkRetry(StepNode current, Integer defaultRetry) {
+        if (current.hasRetry()) {
+            return current.getRetry();
+        }
+
+        if (current.hasParent()) {
+            Node parent = current.getParent();
+            if (parent instanceof StepNode) {
+                return linkRetry((StepNode) parent, defaultRetry);
+            }
+        }
+
+        return defaultRetry;
+    }
+
+    private Integer linkTimeout(StepNode current, Integer defaultTimeout) {
+        if (current.hasTimeout()) {
+            return current.getTimeout();
+        }
+
+        if (current.hasParent()) {
+            Node parent = current.getParent();
+            if (parent instanceof StepNode) {
+                return linkTimeout((StepNode) parent, defaultTimeout);
+            }
+        }
+
+        return defaultTimeout;
+    }
+
     private Set<String> linkFilters(StepNode current) {
         Set<String> output = new LinkedHashSet<>();
 

core/src/main/java/com/flowci/core/plugin/domain/PluginParser.java

@@ -27,7 +27,6 @@
 import lombok.NoArgsConstructor;
 import lombok.NonNull;
 import org.yaml.snakeyaml.Yaml;
-import sun.jvm.hotspot.oops.ObjectHeap;
 
 import java.io.InputStream;
 import java.util.LinkedList;

core/src/test/java/com/flowci/core/test/job/CmdManagerTest.java

@@ -162,6 +162,8 @@ public void should_handle_step_in_step() throws IOException {
 
         // then: verify step 2 - 1 cmd
         ShellIn cmdStep2_1 = (ShellIn) cmdManager.createShellCmd(job, stepService.get(job.getId(), step2_1.getPathAsString()), tree);
+        Assert.assertEquals(500, cmdStep2_1.getTimeout());
+        Assert.assertEquals(2, cmdStep2_1.getRetry());
 
         // input should be overwrite
         Assert.assertEquals("overwrite-parent", cmdStep2_1.getInputs().get("STEP_2"));
@@ -178,6 +180,8 @@ public void should_handle_step_in_step() throws IOException {
         // then: verify step 2 - 2 cmd
         ShellIn cmdStep2_2 = (ShellIn) cmdManager.createShellCmd(job, stepService.get(job.getId(), step2_2.getPathAsString()), tree);
         Assert.assertEquals("parent", cmdStep2_2.getInputs().get("STEP_2"));
+        Assert.assertEquals(1000, cmdStep2_2.getTimeout());
+        Assert.assertEquals(5, cmdStep2_2.getRetry());
 
         // scripts should be linked
         Assert.assertEquals("echo 2", cmdStep2_2.getBash().get(0));

core/src/test/resources/step-in-step.yml

@@ -19,6 +19,8 @@ steps:
         is_runtime: true
       - image: "mysql"
     bash: "echo 2"
+    timeout: 500
+    retry: 2
 
     steps:
       - name: step-2-1
@@ -28,6 +30,8 @@ steps:
           echo "step-2-1"
 
       - name: step-2-2
+        timeout: 1000
+        retry: 5
         docker:
           image: "redis"
         bash: |

tree/src/main/java/com/flowci/tree/StepNode.java

@@ -36,8 +36,24 @@ public class StepNode extends Node {
      */
     private String plugin;
 
+    /**
+     * Step timeout in seconds
+     */
+    private Integer timeout;
+
+    /**
+     * Num of retry times
+     */
+    private Integer retry; // num of retry
+
+    /**
+     * Env vars export to job context
+     */
     private Set<String> exports = new HashSet<>(0);
 
+    /**
+     * Order in the node tree
+     */
     private int order;
 
     /**
@@ -68,4 +84,14 @@ public boolean hasScript() {
     public boolean isRootStep() {
         return parent instanceof FlowNode;
     }
+
+    @JsonIgnore
+    public boolean hasTimeout() {
+        return timeout != null;
+    }
+
+    @JsonIgnore
+    public boolean hasRetry() {
+        return retry != null;
+    }
 }

tree/src/main/java/com/flowci/tree/yml/StepYml.java

@@ -53,6 +53,10 @@ public class StepYml extends YmlBase<StepNode> {
 
     private String plugin;
 
+    private Integer retry; // num of retry
+
+    private Integer timeout; // timeout in seconds
+
     private List<String> exports = new LinkedList<>();
 
     private boolean allow_failure = false;
@@ -62,6 +66,8 @@ public class StepYml extends YmlBase<StepNode> {
         setEnvs(node.getEnvironments());
         setBash(node.getBash());
         setPwsh(node.getPwsh());
+        setRetry(node.getRetry());
+        setTimeout(node.getTimeout());
         setPlugin(node.getPlugin());
         setAllow_failure(node.isAllowFailure());
     }
@@ -71,6 +77,8 @@ public StepNode toNode(Node parent, int index) {
         node.setCondition(condition);
         node.setBash(bash);
         node.setPwsh(pwsh);
+        node.setRetry(retry);
+        node.setTimeout(timeout);
         node.setPlugin(plugin);
         node.setExports(Sets.newHashSet(exports));
         node.setAllowFailure(allow_failure);

tree/src/test/java/com/flowci/tree/test/YmlParserTest.java

@@ -80,6 +80,7 @@ public void should_get_node_from_yml() {
         Assert.assertEquals("step-1", step1.getName()); // step-1 is default name
         Assert.assertEquals("echo step", step1.getEnv("FLOW_WORKSPACE"));
         Assert.assertEquals("echo step version", step1.getEnv("FLOW_VERSION"));
+        Assert.assertEquals(3600, step1.getTimeout().intValue());
 
         Assert.assertTrue(step1.isAllowFailure());
         Assert.assertEquals("println(FLOW_WORKSPACE)\ntrue\n", step1.getCondition());
@@ -88,6 +89,7 @@ public void should_get_node_from_yml() {
         Assert.assertEquals("step2", step2.getName());
         Assert.assertEquals("echo 2", step2.getBash());
         Assert.assertEquals("echo powershell", step2.getPwsh());
+        Assert.assertNull(step2.getTimeout());
 
         DockerOption dockerOption = step2.getDockers().get(0);
         Assert.assertNotNull(dockerOption);

tree/src/test/resources/flow.yml

@@ -31,6 +31,7 @@ steps:
    envs:
      FLOW_WORKSPACE: "echo step"
      FLOW_VERSION: "echo step version"
+   timeout: 3600
    allow_failure: true
 
  - name: step2