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

* Close SSH session after remote shell finish --------- Co-authored-by: David Zollo <[email protected]>
apache · Dec 25, 2023 · 9212531 · 9212531
1 parent b73194b
commit 9212531
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 5 deletions.
diff --git a/...asource-ssh/src/main/java/org/apache/dolphinscheduler/plugin/datasource/ssh/SSHUtils.java b/...asource-ssh/src/main/java/org/apache/dolphinscheduler/plugin/datasource/ssh/SSHUtils.java
@@ -23,9 +23,11 @@
 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 static ClientSession getSession(SshClient client, SSHConnectionParam conn
                 throw new Exception("Failed to add public key identity", e);
             }
         }
+        session.setSessionHeartbeat(SessionHeartbeatController.HeartbeatType.RESERVED, Duration.ofSeconds(3));
         return session;
     }
 }
diff --git a/...ell/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteExecutor.java b/...ell/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteExecutor.java
@@ -39,20 +39,21 @@
 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-";
     static final int TRACK_INTERVAL = 5000;
 
     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 @@ private String getRemoteShellHome() {
         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() {

diff --git a/...ll/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteShellTask.java b/...ll/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteShellTask.java
@@ -97,7 +97,8 @@ public void init() {
 
     @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);