Forráskód Böngészése

[Bug] Close SSH session after remote shell finish (#15348)

* Close SSH session after remote shell finish

---------

Co-authored-by: David Zollo <davidzollo365@gmail.com>
Wenjun Ruan 1 éve
szülő
commit
921253109a

+ 3 - 0
dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-ssh/src/main/java/org/apache/dolphinscheduler/plugin/datasource/ssh/SSHUtils.java

@@ -23,9 +23,11 @@ import org.apache.commons.lang3.StringUtils;
 import org.apache.sshd.client.SshClient;
 import org.apache.sshd.client.session.ClientSession;
 import org.apache.sshd.common.config.keys.loader.KeyPairResourceLoader;
+import org.apache.sshd.common.session.SessionHeartbeatController;
 import org.apache.sshd.common.util.security.SecurityUtils;
 
 import java.security.KeyPair;
+import java.time.Duration;
 import java.util.Collection;
 
 public class SSHUtils {
@@ -57,6 +59,7 @@ public class SSHUtils {
                 throw new Exception("Failed to add public key identity", e);
             }
         }
+        session.setSessionHeartbeat(SessionHeartbeatController.HeartbeatType.RESERVED, Duration.ofSeconds(3));
         return session;
     }
 }

+ 17 - 4
dolphinscheduler-task-plugin/dolphinscheduler-task-remoteshell/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteExecutor.java

@@ -39,10 +39,11 @@ import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.Map;
 
+import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
-public class RemoteExecutor {
+public class RemoteExecutor implements AutoCloseable {
 
     static final String REMOTE_SHELL_HOME = "/tmp/dolphinscheduler-remote-shell-%s/";
     static final String STATUS_TAG_MESSAGE = "DOLPHINSCHEDULER-REMOTE-SHELL-TASK-STATUS-";
@@ -50,9 +51,9 @@ public class RemoteExecutor {
 
     protected Map<String, String> taskOutputParams = new HashMap<>();
 
-    SshClient sshClient;
-    ClientSession session;
-    SSHConnectionParam sshConnectionParam;
+    private SshClient sshClient;
+    private ClientSession session;
+    private SSHConnectionParam sshConnectionParam;
 
     public RemoteExecutor(SSHConnectionParam sshConnectionParam) {
 
@@ -205,6 +206,18 @@ public class RemoteExecutor {
         return String.format(REMOTE_SHELL_HOME, sshConnectionParam.getUser());
     }
 
+    @SneakyThrows
+    @Override
+    public void close() {
+        if (session != null && session.isOpen()) {
+            session.close();
+        }
+        if (sshClient != null && sshClient.isStarted()) {
+            sshClient.close();
+        }
+
+    }
+
     static class COMMAND {
 
         private COMMAND() {

+ 2 - 1
dolphinscheduler-task-plugin/dolphinscheduler-task-remoteshell/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteShellTask.java

@@ -97,7 +97,8 @@ public class RemoteShellTask extends AbstractTask {
 
     @Override
     public void handle(TaskCallBack taskCallBack) throws TaskException {
-        try {
+        // add task close method to release resource
+        try (RemoteExecutor executor = remoteExecutor) {
             // construct process
             String localFile = buildCommand();
             int exitCode = remoteExecutor.run(taskId, localFile);