[Improvement-17725][TaskPlugin] SagemakerTask resource leak repair

dolphinscheduler-task-plugin/dolphinscheduler-task-sagemaker/src/main/java/org/apache/dolphinscheduler/plugin/task/sagemaker/SagemakerTask.java

@@ -126,26 +126,42 @@ public void submitApplication() throws TaskException {
 
     @Override
     public void cancelApplication() {
-        initPipelineId();
         try {
+            initPipelineId();
             // stop pipeline
             utils.stopPipelineExecution(client, pipelineId);
         } catch (Exception e) {
-            throw new TaskException("cancel application error", e);
+            log.error("SageMaker task cancel application error: {}", e.getMessage(), e);
+            throw new TaskException("SageMaker task cancel application error: " + e.getMessage(), e);
+        } finally {
+            // shutdown client
+            if (client != null) {
+                client.shutdown();
+            }
         }
     }
 
     @Override
     public void trackApplicationStatus() throws TaskException {
-        initPipelineId();
-        // Keep checking the health status
-        exitStatusCode = utils.checkPipelineExecutionStatus(client, pipelineId);
+        try {
+            initPipelineId();
+            // Keep checking the health status
+            exitStatusCode = utils.checkPipelineExecutionStatus(client, pipelineId);
+        } catch (Exception e) {
+            log.error("SageMaker task track application error: {}", e.getMessage(), e);
+            throw new TaskException("SageMaker task track application error: " + e.getMessage(), e);
+        } finally {
+            // shutdown client
+            if (client != null) {
+                client.shutdown();
+            }
+        }
     }
 
     /**
      * init sagemaker applicationId if null
      */
-    private void initPipelineId() {
+    void initPipelineId() {
         if (pipelineId == null) {
             if (StringUtils.isNotEmpty(getAppIds())) {
                 pipelineId = JSONUtils.parseObject(getAppIds(), PipelineUtils.PipelineId.class);

dolphinscheduler-task-plugin/dolphinscheduler-task-sagemaker/src/test/java/org/apache/dolphinscheduler/plugin/task/sagemaker/SagemakerTaskTest.java

@@ -17,13 +17,19 @@
 
 package org.apache.dolphinscheduler.plugin.task.sagemaker;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import org.apache.dolphinscheduler.common.utils.JSONUtils;
 import org.apache.dolphinscheduler.plugin.datasource.api.utils.DataSourceUtils;
 import org.apache.dolphinscheduler.plugin.datasource.sagemaker.param.SagemakerConnectionParam;
+import org.apache.dolphinscheduler.plugin.task.api.TaskException;
 import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext;
 import org.apache.dolphinscheduler.plugin.task.api.parameters.resource.ResourceParametersHelper;
 
@@ -139,4 +145,33 @@ private String buildParameters() {
 
         return JSONUtils.toJsonString(parameters);
     }
+
+    @Test
+    public void testTrackApplicationStatus_InitPipelineIdThrowsException() throws Exception {
+        doThrow(new TaskException("sagemaker applicationID is null")).when(sagemakerTask).initPipelineId();
+
+        TaskException exception = assertThrows(TaskException.class, () -> {
+            sagemakerTask.trackApplicationStatus();
+        });
+
+        assertEquals("sagemaker applicationID is null", exception.getMessage());
+        verify(client, times(1)).shutdown();
+    }
+
+    @Test
+    public void testCancelApplication_InitPipelineIdThrowsException() {
+        // Mock the behavior of initPipelineId to throw an exception
+        doThrow(new TaskException("sagemaker applicationID is null")).when(sagemakerTask).initPipelineId();
+
+        // Call the method under test and expect an exception
+        TaskException exception = assertThrows(TaskException.class, () -> {
+            sagemakerTask.cancelApplication();
+        });
+
+        // Verify the exception message
+        assertEquals("cancel application error", exception.getMessage());
+
+        // Verify that client.shutdown() was called
+        verify(client, times(1)).shutdown();
+    }
 }