(dolphinscheduler) branch dev updated: [Bug] Close SSH session after remote shell finish (#15348)

[chufenggao]

This is an automated email from the ASF dual-hosted git repository.

chufenggao pushed a commit to branch dev
in repository https://gitbox.apache.org/repos/asf/dolphinscheduler.git


The following commit(s) were added to refs/heads/dev by this push:
     new 921253109a [Bug] Close SSH session after remote shell finish (#15348)
921253109a is described below

commit 921253109abcc57590d71c8bac388e4c0a007c0d
Author: Wenjun Ruan <wen...@apache.org>
AuthorDate: Mon Dec 25 13:04:45 2023 +0800

    [Bug] Close SSH session after remote shell finish (#15348)
    
    * Close SSH session after remote shell finish
    
    ---------
    
    Co-authored-by: David Zollo <davidzollo...@gmail.com>
---
 .../plugin/datasource/ssh/SSHUtils.java             |  3 +++
 .../plugin/task/remoteshell/RemoteExecutor.java     | 21 +++++++++++++++++----
 .../plugin/task/remoteshell/RemoteShellTask.java    |  3 ++-
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git 
a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-ssh/src/main/java/org/apache/dolphinscheduler/plugin/datasource/ssh/SSHUtils.java
 
b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-ssh/src/main/java/org/apache/dolphinscheduler/plugin/datasource/ssh/SSHUtils.java
index 8ee8fa79a7..42e1175e2e 100644
--- 
a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-ssh/src/main/java/org/apache/dolphinscheduler/plugin/datasource/ssh/SSHUtils.java
+++ 
b/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;
     }
 }
diff --git 
a/dolphinscheduler-task-plugin/dolphinscheduler-task-remoteshell/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteExecutor.java
 
b/dolphinscheduler-task-plugin/dolphinscheduler-task-remoteshell/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteExecutor.java
index 4de28b7fb5..c590fa9e44 100644
--- 
a/dolphinscheduler-task-plugin/dolphinscheduler-task-remoteshell/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteExecutor.java
+++ 
b/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() {
diff --git 
a/dolphinscheduler-task-plugin/dolphinscheduler-task-remoteshell/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteShellTask.java
 
b/dolphinscheduler-task-plugin/dolphinscheduler-task-remoteshell/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteShellTask.java
index 561392896c..6c55f069e1 100644
--- 
a/dolphinscheduler-task-plugin/dolphinscheduler-task-remoteshell/src/main/java/org/apache/dolphinscheduler/plugin/task/remoteshell/RemoteShellTask.java
+++ 
b/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);