[Fix-16789][Task-Plugin] Cancel Flink job through the flink CLI when stopping a task

dolphinscheduler-task-plugin/dolphinscheduler-task-api/src/main/java/org/apache/dolphinscheduler/plugin/task/api/AbstractCommandExecutor.java

@@ -39,13 +39,13 @@
 
 import java.io.BufferedReader;
 import java.io.InputStreamReader;
-import java.lang.reflect.Field;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import lombok.extern.slf4j.Slf4j;
 import io.fabric8.kubernetes.client.dsl.LogWatch;
@@ -56,6 +56,12 @@ public abstract class AbstractCommandExecutor {
     protected volatile Map<String, String> taskOutputParams = new HashMap<>();
     private Process process;
 
+    /**
+     * The future which writes the output of the process into the task log. It is kept so that a caller
+     * which stops the process outside of {@link #run} can wait for the output to be collected.
+     */
+    private volatile CompletableFuture<Void> collectProcessLogFuture;
+
     protected TaskExecutionContext taskRequest;
 
     public AbstractCommandExecutor(TaskExecutionContext taskRequest) {
@@ -101,7 +107,7 @@ public TaskResponse run(IShellInterceptorBuilder iShellInterceptorBuilder,
         process = iShellInterceptor.execute();
 
         // parse process output
-        final CompletableFuture<Void> collectProcessLogFuture = collectProcessLog(this.process);
+        collectProcessLogFuture = collectProcessLog(this.process);
 
         // collect pod log
         final Optional<CompletableFuture<?>> collectPodLogFuture = collectPodLogIfNeeded();
@@ -165,6 +171,44 @@ public Map<String, String> getTaskOutputParams() {
         return taskOutputParams;
     }
 
+    /**
+     * Wait until the output of the process has been written to the task log, at most the given time.
+     * A caller which stops the process outside of {@link #run} uses it to make sure that the tail of
+     * the output is available in the task log.
+     *
+     * @param timeoutMillis the maximum time to wait
+     * @return true if the output was collected, false if the wait timed out
+     */
+    public boolean awaitOutputCollected(long timeoutMillis) {
+        return awaitCollected(collectProcessLogFuture, timeoutMillis);
+    }
+
+    /**
+     * Wait until the given output collection completes, at most the given time. Kept as a static
+     * method so that the waiting itself can be verified without starting a process.
+     *
+     * @param future the output collection, may be null when no process was started
+     * @param timeoutMillis the maximum time to wait
+     * @return true if the output was collected, false if the wait timed out
+     */
+    static boolean awaitCollected(CompletableFuture<Void> future, long timeoutMillis) {
+        if (future == null) {
+            return true;
+        }
+        try {
+            future.get(timeoutMillis, TimeUnit.MILLISECONDS);
+            return true;
+        } catch (TimeoutException e) {
+            return false;
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            return false;
+        } catch (Exception e) {
+            log.error("Collect the process output failed", e);
+            return true;
+        }
+    }
+
     public void cancelApplication() throws InterruptedException {
         if (process == null) {
             return;
@@ -266,18 +310,7 @@ private long getRemainTime() {
      * @return process id
      */
     private int getProcessId(Process process) {
-        int processId = 0;
-
-        try {
-            Field f = process.getClass().getDeclaredField(TaskConstants.PID);
-            f.setAccessible(true);
-
-            processId = f.getInt(process);
-        } catch (Exception e) {
-            log.error("Get task pid failed", e);
-        }
-
-        return processId;
+        return ProcessUtils.getProcessId(process);
     }
 
 }

dolphinscheduler-task-plugin/dolphinscheduler-task-api/src/main/java/org/apache/dolphinscheduler/plugin/task/api/AbstractYarnTask.java

@@ -105,6 +105,18 @@ public List<String> getApplicationIds() throws TaskException {
                 PropertyUtils.getString(APPID_COLLECT, DEFAULT_COLLECT_WAY));
     }
 
+    /**
+     * Wait until the output of the task process has been written to the task log, at most the given
+     * time. A task which stops the process outside of the normal run, for example to cancel a remote
+     * job, uses it to make sure that the tail of the output is available in the task log.
+     *
+     * @param timeoutMillis the maximum time to wait
+     * @return true if the output was collected, false if the wait timed out
+     */
+    protected boolean awaitTaskOutputCollected(long timeoutMillis) {
+        return shellCommandExecutor.awaitOutputCollected(timeoutMillis);
+    }
+
     /**
      * Get the script used to bootstrap the task
      */

dolphinscheduler-task-plugin/dolphinscheduler-task-api/src/main/java/org/apache/dolphinscheduler/plugin/task/api/TaskConstants.java

@@ -33,7 +33,13 @@ private TaskConstants() {
 
     public static final String YARN_APPLICATION_REGEX = "application_\\d+_\\d+";
 
-    public static final String FLINK_APPLICATION_REGEX = "JobID \\w+";
+    /**
+     * The Flink JobID printed by `flink run` ("JobID" followed by the id) or by the Flink SQL
+     * Client ("Job ID:" followed by the id). A Flink JobID is a 32 character hexadecimal string,
+     * which is required here so that an unrelated line containing "Job ID" is not mistaken for a
+     * JobID. The first group is the JobID, which is not the YARN/K8s application id.
+     */
+    public static final String FLINK_JOB_ID_REGEX = "Job\\s*ID:?\\s+([0-9a-fA-F]{32})(?![0-9a-fA-F])";
 
     public static final String DATASOURCE_PASSWORD_REGEX =
             "(?<=((?i)password((\" : \")|(\":\")|(\\\\\":\\\\\")|(=')))).*?(?=((\")|(\\\\\")|(')))";

dolphinscheduler-task-plugin/dolphinscheduler-task-api/src/main/java/org/apache/dolphinscheduler/plugin/task/api/utils/ProcessUtils.java

@@ -42,14 +42,17 @@
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.SystemUtils;
 
+import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.ServiceLoader;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -99,6 +102,12 @@ private ProcessUtils() {
      */
     private static final Pattern PID_PATTERN = Pattern.compile("\\s+");
 
+    /**
+     * The command which lists the pid and the parent pid of every process. It is used when pstree is
+     * not available, otherwise no process tree could be terminated on such a host.
+     */
+    private static final String PS_PID_LIST_COMMAND = "ps -eo pid,ppid";
+
     private static final String SIGINT = "2";
     private static final String SIGTERM = "15";
     private static final String SIGKILL = "9";
@@ -109,9 +118,20 @@ private ProcessUtils() {
      * @return Whether the process was successfully terminated
      */
     public static boolean kill(@NonNull TaskExecutionContext request) {
+        log.info("Begin killing task instance, processId: {}", request.getProcessId());
+        return kill(request.getProcessId(), request.getTenantCode());
+    }
+
+    /**
+     * Terminate the process tree which starts at the given process id, support multi-level signal
+     * processing and fallback strategy.
+     *
+     * @param processId the root process id, 0 means the process already finished
+     * @param tenantCode the tenant which runs the process, may be null
+     * @return Whether the process was successfully terminated
+     */
+    public static boolean kill(int processId, String tenantCode) {
         try {
-            log.info("Begin killing task instance, processId: {}", request.getProcessId());
-            int processId = request.getProcessId();
             if (processId == 0) {
                 log.info("Task instance has already finished, no need to kill");
                 return true;
@@ -121,14 +141,14 @@ public static boolean kill(@NonNull TaskExecutionContext request) {
             List<Integer> pidList = getPidList(processId);
 
             // 1. Try to terminate gracefully `kill -2`
-            boolean gracefulKillSuccess = sendKillSignal(SIGINT, pidList, request.getTenantCode());
+            boolean gracefulKillSuccess = sendKillSignal(SIGINT, pidList, tenantCode);
             if (gracefulKillSuccess) {
                 log.info("Successfully killed process tree by SIGINT, processId: {}", processId);
                 return true;
             }
 
             // 2. Try to terminate gracefully `kill -15`
-            boolean termKillSuccess = sendKillSignal(SIGTERM, pidList, request.getTenantCode());
+            boolean termKillSuccess = sendKillSignal(SIGTERM, pidList, tenantCode);
             if (termKillSuccess) {
                 log.info("Successfully killed process tree by SIGTERM, processId: {}", processId);
                 return true;
@@ -137,7 +157,7 @@ public static boolean kill(@NonNull TaskExecutionContext request) {
             // 3. As a last resort, use `kill -9`
             log.warn("Killing process by SIGINT & SIGTERM failed, using SIGKILL as a last resort for processId: {}",
                     processId);
-            boolean forceKillSuccess = sendKillSignal(SIGKILL, pidList, request.getTenantCode());
+            boolean forceKillSuccess = sendKillSignal(SIGKILL, pidList, tenantCode);
             if (forceKillSuccess) {
                 log.info("Successfully killed process tree by SIGKILL, processId: {}", processId);
             } else {
@@ -146,11 +166,31 @@ public static boolean kill(@NonNull TaskExecutionContext request) {
             return forceKillSuccess;
 
         } catch (Exception e) {
-            log.error("Kill task instance error, processId: {}", request.getProcessId(), e);
+            log.error("Kill task instance error, processId: {}", processId, e);
             return false;
         }
     }
 
+    /**
+     * Get the process id of a process which was started by this JVM.
+     *
+     * @param process the process, may be null
+     * @return the process id, or 0 when it cannot be determined
+     */
+    public static int getProcessId(Process process) {
+        if (process == null) {
+            return 0;
+        }
+        try {
+            Field field = process.getClass().getDeclaredField(TaskConstants.PID);
+            field.setAccessible(true);
+            return field.getInt(process);
+        } catch (Exception e) {
+            log.error("Get process id failed", e);
+            return 0;
+        }
+    }
+
     /**
      * Send a kill signal to a process group
      * @param signal Signal type (SIGINT, SIGTERM, SIGKILL)
@@ -269,14 +309,14 @@ public static List<Integer> getPidList(int processId) throws Exception {
                 rawPidStr = OSUtils.exeCmd(String.format("%s -p %d", TaskConstants.PSTREE, processId));
             }
         } catch (Exception ex) {
-            log.error("Failed to execute 'pstree' command for process ID: {}", processId, ex);
-            throw ex;
+            log.warn("Failed to execute 'pstree' command for process ID: {}, falling back to 'ps'", processId, ex);
+            return getPidListByPs(processId);
         }
 
         String pidsStr = parsePidStr(rawPidStr);
         if (StringUtils.isBlank(pidsStr)) {
-            log.warn("No PIDs found for process: {}", processId);
-            return Collections.emptyList();
+            log.warn("No PIDs found for process: {} with pstree, falling back to 'ps'", processId);
+            return getPidListByPs(processId);
         }
 
         String[] pidArray = PID_PATTERN.split(pidsStr.trim());
@@ -314,6 +354,49 @@ public static List<Integer> getPidList(int processId) throws Exception {
         return pidList;
     }
 
+    /**
+     * Get the descendant process IDs (including the given process) with {@code ps}, which is used
+     * when {@code pstree} is not available.
+     *
+     * @param processId the root process ID
+     * @return list of process IDs; returns empty list if the process does not exist
+     * @throws Exception if command execution fails unexpectedly
+     */
+    private static List<Integer> getPidListByPs(int processId) throws Exception {
+        String rawPidStr = OSUtils.exeCmd(PS_PID_LIST_COMMAND);
+        Map<Integer, List<Integer>> childrenByParentPid = new HashMap<>();
+        Set<Integer> allPids = new HashSet<>();
+        for (String line : rawPidStr.split("\\r?\\n")) {
+            String[] columns = line.trim().split("\\s+");
+            if (columns.length != 2) {
+                continue;
+            }
+            try {
+                int pid = Integer.parseInt(columns[0]);
+                int parentPid = Integer.parseInt(columns[1]);
+                allPids.add(pid);
+                childrenByParentPid.computeIfAbsent(parentPid, key -> new ArrayList<>()).add(pid);
+            } catch (NumberFormatException e) {
+                // the header of ps and any other non numeric line
+            }
+        }
+        if (!allPids.contains(processId)) {
+            log.warn("No PIDs found for process: {}", processId);
+            return Collections.emptyList();
+        }
+        List<Integer> pidList = new ArrayList<>();
+        collectPidAndDescendants(processId, childrenByParentPid, pidList);
+        return pidList;
+    }
+
+    private static void collectPidAndDescendants(int processId, Map<Integer, List<Integer>> childrenByParentPid,
+                                                 List<Integer> pidList) {
+        pidList.add(processId);
+        for (Integer childPid : childrenByParentPid.getOrDefault(processId, Collections.emptyList())) {
+            collectPidAndDescendants(childPid, childrenByParentPid, pidList);
+        }
+    }
+
     public static String parsePidStr(String rawPidStr) {
 
         log.info("prepare to parse pid, raw pid string: {}", rawPidStr);

dolphinscheduler-task-plugin/dolphinscheduler-task-api/src/test/java/org/apache/dolphinscheduler/plugin/task/api/AbstractCommandExecutorTest.java

@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.dolphinscheduler.plugin.task.api;
+
+import java.util.concurrent.CompletableFuture;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class AbstractCommandExecutorTest {
+
+    @Test
+    public void testAwaitCollectedReturnsWhenTheOutputIsCollected() {
+        CompletableFuture<Void> collected = new CompletableFuture<>();
+        collected.complete(null);
+
+        Assertions.assertTrue(AbstractCommandExecutor.awaitCollected(collected, 1_000L));
+    }
+
+    @Test
+    public void testAwaitCollectedWithoutOutputCollection() {
+        // no process was started, there is nothing to wait for
+        Assertions.assertTrue(AbstractCommandExecutor.awaitCollected(null, 1_000L));
+    }
+
+    @Test
+    public void testAwaitCollectedWhenTheCollectionFailed() {
+        // a failed collection has nothing left to wait for, it must not block the caller
+        CompletableFuture<Void> failed = new CompletableFuture<>();
+        failed.completeExceptionally(new IllegalStateException("collect failed"));
+
+        Assertions.assertTrue(AbstractCommandExecutor.awaitCollected(failed, 1_000L));
+    }
+
+    @Test
+    public void testAwaitCollectedTimesOut() {
+        CompletableFuture<Void> pending = new CompletableFuture<>();
+
+        long start = System.currentTimeMillis();
+        boolean collected = AbstractCommandExecutor.awaitCollected(pending, 200L);
+        long elapsed = System.currentTimeMillis() - start;
+
+        Assertions.assertFalse(collected);
+        Assertions.assertTrue(elapsed >= 200L, "the wait must last for the whole timeout, elapsed: " + elapsed + "ms");
+    }
+}