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 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);