[Improvement-17724][TaskPlugin] EmrTask resource leak repair

dolphinscheduler-task-plugin/dolphinscheduler-task-emr/src/main/java/org/apache/dolphinscheduler/plugin/task/emr/EmrAddStepsTask.java

@@ -116,6 +116,11 @@ public void trackApplicationStatus() throws TaskException {
             final int exitStatusCode = calculateExitStatusCode(stepStatus);
             setExitStatusCode(exitStatusCode);
             log.info("emr task finished with step status : {}", stepStatus);
+
+            // shutdown emrclient
+            if (emrClient != null) {
+                emrClient.shutdown();
+            }
         }
     }
 
@@ -187,26 +192,37 @@ private StepStatus getStepStatus() {
 
     @Override
     public void cancelApplication() throws TaskException {
-        log.info("trying cancel emr step, taskId:{}, clusterId:{}, stepId:{}",
-                this.taskExecutionContext.getTaskInstanceId(), clusterId, stepId);
-        CancelStepsRequest cancelStepsRequest = new CancelStepsRequest().withClusterId(clusterId).withStepIds(stepId);
-        CancelStepsResult cancelStepsResult = emrClient.cancelSteps(cancelStepsRequest);
+        try {
+            log.info("trying cancel emr step, taskId:{}, clusterId:{}, stepId:{}",
+                    this.taskExecutionContext.getTaskInstanceId(), clusterId, stepId);
+            CancelStepsRequest cancelStepsRequest =
+                    new CancelStepsRequest().withClusterId(clusterId).withStepIds(stepId);
+            CancelStepsResult cancelStepsResult = emrClient.cancelSteps(cancelStepsRequest);
+
+            if (cancelStepsResult == null) {
+                throw new EmrTaskException("cancel emr step failed");
+            }
 
-        if (cancelStepsResult == null) {
-            throw new EmrTaskException("cancel emr step failed");
-        }
+            CancelStepsInfo cancelEmrStepInfo = cancelStepsResult.getCancelStepsInfoList()
+                    .stream()
+                    .filter(cancelStepsInfo -> cancelStepsInfo.getStepId().equals(stepId))
+                    .findFirst()
+                    .orElseThrow(() -> new EmrTaskException("cancel emr step failed"));
 
-        CancelStepsInfo cancelEmrStepInfo = cancelStepsResult.getCancelStepsInfoList()
-                .stream()
-                .filter(cancelStepsInfo -> cancelStepsInfo.getStepId().equals(stepId))
-                .findFirst()
-                .orElseThrow(() -> new EmrTaskException("cancel emr step failed"));
+            if (CancelStepsRequestStatus.FAILED.toString().equals(cancelEmrStepInfo.getStatus())) {
+                throw new EmrTaskException("cancel emr step failed, message:" + cancelEmrStepInfo.getReason());
+            }
 
-        if (CancelStepsRequestStatus.FAILED.toString().equals(cancelEmrStepInfo.getStatus())) {
-            throw new EmrTaskException("cancel emr step failed, message:" + cancelEmrStepInfo.getReason());
+            log.info("the result of cancel emr step is:{}", cancelStepsResult);
+        } catch (EmrTaskException | SdkBaseException e) {
+            log.error("cancel emr step failed", e);
+            throw new TaskException("cancel emr step failed", e);
+        } finally {
+            // shutdown emrclient
+            if (emrClient != null) {
+                emrClient.shutdown();
+            }
         }
-
-        log.info("the result of cancel emr step is:{}", cancelStepsResult);
     }
 
 }

dolphinscheduler-task-plugin/dolphinscheduler-task-emr/src/main/java/org/apache/dolphinscheduler/plugin/task/emr/EmrJobFlowTask.java

@@ -110,6 +110,11 @@ public void trackApplicationStatus() throws TaskException {
             final int exitStatusCode = calculateExitStatusCode(clusterStatus);
             setExitStatusCode(exitStatusCode);
             log.info("emr task finished with cluster status : {}", clusterStatus);
+
+            // shutdown emrclient
+            if (emrClient != null) {
+                emrClient.shutdown();
+            }
         }
     }
 
@@ -180,11 +185,21 @@ private ClusterStatus getClusterStatus() {
 
     @Override
     public void cancelApplication() throws TaskException {
-        log.info("trying terminate job flow, taskId:{}, clusterId:{}", this.taskExecutionContext.getTaskInstanceId(),
-                clusterId);
-        TerminateJobFlowsRequest terminateJobFlowsRequest = new TerminateJobFlowsRequest().withJobFlowIds(clusterId);
-        TerminateJobFlowsResult terminateJobFlowsResult = emrClient.terminateJobFlows(terminateJobFlowsRequest);
-        log.info("the result of terminate job flow is:{}", terminateJobFlowsResult);
+        try {
+            log.info("trying terminate job flow, taskId:{}, clusterId:{}", this.taskExecutionContext.getTaskInstanceId(),
+                    clusterId);
+            TerminateJobFlowsRequest terminateJobFlowsRequest = new TerminateJobFlowsRequest().withJobFlowIds(clusterId);
+            TerminateJobFlowsResult terminateJobFlowsResult = emrClient.terminateJobFlows(terminateJobFlowsRequest);
+            log.info("the result of terminate job flow is:{}", terminateJobFlowsResult);
+        } catch (Exception e) {
+            log.error("emr job flow task cancel error: {}", e.getMessage(), e);
+            throw new TaskException("emr job flow task cancel error: " + e.getMessage(), e);
+        } finally {
+            // shutdown emrclient
+            if (emrClient != null) {
+                emrClient.shutdown();
+            }
+        }
     }
 
 }

dolphinscheduler-task-plugin/dolphinscheduler-task-emr/src/test/java/org/apache/dolphinscheduler/plugin/task/emr/EmrAddStepsTaskTest.java

@@ -21,6 +21,8 @@
 import static org.apache.dolphinscheduler.plugin.task.api.TaskConstants.EXIT_CODE_KILL;
 import static org.apache.dolphinscheduler.plugin.task.api.TaskConstants.EXIT_CODE_SUCCESS;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import org.apache.dolphinscheduler.common.utils.JSONUtils;
 import org.apache.dolphinscheduler.plugin.task.api.TaskCallBack;
@@ -44,6 +46,10 @@
 import com.amazonaws.services.elasticmapreduce.AmazonElasticMapReduce;
 import com.amazonaws.services.elasticmapreduce.model.AddJobFlowStepsResult;
 import com.amazonaws.services.elasticmapreduce.model.AmazonElasticMapReduceException;
+import com.amazonaws.services.elasticmapreduce.model.CancelStepsInfo;
+import com.amazonaws.services.elasticmapreduce.model.CancelStepsRequest;
+import com.amazonaws.services.elasticmapreduce.model.CancelStepsRequestStatus;
+import com.amazonaws.services.elasticmapreduce.model.CancelStepsResult;
 import com.amazonaws.services.elasticmapreduce.model.DescribeStepResult;
 import com.amazonaws.services.elasticmapreduce.model.Step;
 import com.amazonaws.services.elasticmapreduce.model.StepState;
@@ -198,4 +204,34 @@ private EmrParameters buildErrorEmrTaskParameters() {
 
         return emrParameters;
     }
+
+    @Test
+    public void cancelApplication_CancelStepsResultIsNull_ShouldThrowException() {
+        Mockito.when(emrClient.cancelSteps(any(CancelStepsRequest.class))).thenReturn(null);
+
+        TaskException exception = Assertions.assertThrows(TaskException.class, () -> {
+            emrAddStepsTask.cancelApplication();
+        });
+
+        verify(emrClient, times(1)).shutdown();
+    }
+
+    @Test
+    public void cancelApplication_CancelStepsRequestFails_ShouldThrowException() {
+        CancelStepsInfo cancelStepsInfo = new CancelStepsInfo()
+                .withStepId("step-123")
+                .withStatus(CancelStepsRequestStatus.FAILED.toString())
+                .withReason("Test failure");
+
+        CancelStepsResult cancelStepsResult = new CancelStepsResult()
+                .withCancelStepsInfoList(cancelStepsInfo);
+
+        Mockito.when(emrClient.cancelSteps(any(CancelStepsRequest.class))).thenReturn(cancelStepsResult);
+
+        TaskException exception = Assertions.assertThrows(TaskException.class, () -> {
+            emrAddStepsTask.cancelApplication();
+        });
+
+        verify(emrClient, times(1)).shutdown();
+    }
 }

dolphinscheduler-task-plugin/dolphinscheduler-task-emr/src/test/java/org/apache/dolphinscheduler/plugin/task/emr/EmrJobFlowTaskTest.java

@@ -21,6 +21,8 @@
 import static org.apache.dolphinscheduler.plugin.task.api.TaskConstants.EXIT_CODE_KILL;
 import static org.apache.dolphinscheduler.plugin.task.api.TaskConstants.EXIT_CODE_SUCCESS;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import org.apache.dolphinscheduler.common.utils.JSONUtils;
 import org.apache.dolphinscheduler.plugin.task.api.TaskCallBack;
@@ -49,6 +51,8 @@
 import com.amazonaws.services.elasticmapreduce.model.ClusterStatus;
 import com.amazonaws.services.elasticmapreduce.model.DescribeClusterResult;
 import com.amazonaws.services.elasticmapreduce.model.RunJobFlowResult;
+import com.amazonaws.services.elasticmapreduce.model.TerminateJobFlowsRequest;
+import com.amazonaws.services.elasticmapreduce.model.TerminateJobFlowsResult;
 
 @ExtendWith(MockitoExtension.class)
 public class EmrJobFlowTaskTest {
@@ -210,4 +214,17 @@ private String buildEmrTaskParameters() {
 
         return JSONUtils.toJsonString(emrParameters);
     }
+
+    @Test
+    public void testCancelApplication_Success() throws TaskException {
+        TerminateJobFlowsResult terminateJobFlowsResult = Mockito.mock(TerminateJobFlowsResult.class);
+        Mockito.when(emrClient.terminateJobFlows(any(TerminateJobFlowsRequest.class)))
+                .thenReturn(terminateJobFlowsResult);
+        Mockito.doNothing().when(emrClient).shutdown();
+
+        Assertions.assertDoesNotThrow(() -> emrJobFlowTask.cancelApplication());
+
+        verify(emrClient, times(1)).terminateJobFlows(any(TerminateJobFlowsRequest.class));
+        verify(emrClient, times(1)).shutdown();
+    }
 }