The YarnClientTrait add clusterClient close (#2906)

* The YarnClientTrait add clusterClient close * Optimized code
apache · Jul 29, 2023 · b71e094 · b71e094
1 parent 99fcf29
commit b71e094
Showing 1 changed file with 10 additions and 10 deletions.
diff --git a/...client-core/src/main/scala/org/apache/streampark/flink/client/trait/YarnClientTrait.scala b/...client-core/src/main/scala/org/apache/streampark/flink/client/trait/YarnClientTrait.scala
@@ -42,8 +42,8 @@ trait YarnClientTrait extends FlinkClientTrait {
       request: R,
       flinkConf: Configuration,
       actionFunc: (JobID, ClusterClient[_]) => O): O = {
-    val jobID = getJobID(request.jobId)
-    val clusterClient = {
+
+    Utils.using {
       flinkConf.safeSet(YarnConfigOptions.APPLICATION_ID, request.clusterId)
       val clusterClientFactory = new YarnClusterClientFactory
       val applicationId = clusterClientFactory.getClusterId(flinkConf)
@@ -53,15 +53,15 @@ trait YarnClientTrait extends FlinkClientTrait {
       }
       val clusterDescriptor = clusterClientFactory.createClusterDescriptor(flinkConf)
       clusterDescriptor.retrieve(applicationId).getClusterClient
+    } {
+      client =>
+        Try(actionFunc(getJobID(request.jobId), client)).recover {
+          case e =>
+            throw new FlinkException(
+              s"[StreamPark] Do ${request.getClass.getSimpleName} for the job ${request.jobId} failed. " +
+                s"detail: ${Utils.stringifyException(e)}");
+        }.get
     }
-    Try {
-      actionFunc(jobID, clusterClient)
-    }.recover {
-      case e =>
-        throw new FlinkException(
-          s"[StreamPark] Do ${request.getClass.getSimpleName} for the job ${request.jobId} failed. " +
-            s"detail: ${Utils.stringifyException(e)}");
-    }.get
   }
 
   override def doTriggerSavepoint(