Fix validation runtime hanle leak

apache · Nov 15, 2023 · d4c0d53 · d4c0d53
1 parent bc25d07
commit d4c0d53
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/ValidatorApiImpl.scala b/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/ValidatorApiImpl.scala
@@ -17,6 +17,7 @@
 package io.glutenproject.backendsapi.velox
 
 import io.glutenproject.backendsapi.ValidatorApi
+import io.glutenproject.exec.Runtimes
 import io.glutenproject.substrait.plan.PlanNode
 import io.glutenproject.validate.NativePlanValidationInfo
 import io.glutenproject.vectorized.NativePlanEvaluator
@@ -31,8 +32,13 @@ class ValidatorApiImpl extends ValidatorApi {
     doExprValidate(Map(), substraitExprName, expr)
 
   override def doNativeValidateWithFailureReason(plan: PlanNode): NativePlanValidationInfo = {
-    val validator = NativePlanEvaluator.createForValidation()
-    validator.doNativeValidateWithFailureReason(plan.toProtobuf.toByteArray)
+    val tmpRuntime = Runtimes.tmpInstance()
+    try {
+      val validator = NativePlanEvaluator.createForValidation(tmpRuntime)
+      validator.doNativeValidateWithFailureReason(plan.toProtobuf.toByteArray)
+    } finally {
+      tmpRuntime.release()
+    }
   }
 
   override def doSparkPlanValidate(plan: SparkPlan): Boolean = true

diff --git a/gluten-data/src/main/java/io/glutenproject/vectorized/NativePlanEvaluator.java b/gluten-data/src/main/java/io/glutenproject/vectorized/NativePlanEvaluator.java
@@ -45,9 +45,9 @@ public static NativePlanEvaluator create() {
     return new NativePlanEvaluator(Runtimes.contextInstance());
   }
 
-  public static NativePlanEvaluator createForValidation() {
+  public static NativePlanEvaluator createForValidation(Runtime runtime) {
     // Driver side doesn't have context instance of Runtime
-    return new NativePlanEvaluator(Runtimes.tmpInstance());
+    return new NativePlanEvaluator(runtime);
   }
 
   public NativePlanValidationInfo doNativeValidateWithFailureReason(byte[] subPlan) {