withClosingOperationLog

Author: bowenliang123

kyuubi-common/src/main/scala/org/apache/kyuubi/operation/AbstractOperation.scala

@@ -17,6 +17,7 @@
 
 package org.apache.kyuubi.operation
 
+import java.io.IOException
 import java.util.concurrent.{Future, ScheduledExecutorService, TimeUnit}
 import java.util.concurrent.locks.ReentrantLock
 
@@ -247,4 +248,19 @@ abstract class AbstractOperation(session: Session) extends Operation with Loggin
     ok.setInfoMessages(hints.asJava)
     ok
   }
+
+  /**
+   * Close the OperationLog, after running the block
+   */
+  def withClosingOperationLog[T](f: => T): T = {
+    try {
+      f
+    } finally {
+      try {
+        getOperationLog.foreach(_.close())
+      } catch {
+        case e: IOException => error(e.getMessage, e)
+      }
+    }
+  }
 }

kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala

@@ -336,7 +336,7 @@ class BatchJobSubmission(
     }
   }
 
-  override def close(): Unit = withLockRequired {
+  override def close(): Unit = withLockRequired(withClosingOperationLog {
     if (!isClosedOrCanceled) {
       MetricsSystem.tracing(_.decCount(MetricRegistry.name(OPERATION_OPEN, opType)))
 
@@ -373,13 +373,7 @@ class BatchJobSubmission(
         }
       }
     }
-
-    try {
-      getOperationLog.foreach(_.close())
-    } catch {
-      case e: IOException => error(e.getMessage, e)
-    }
-  }
+  })
 
   override def cancel(): Unit = {
     throw new IllegalStateException("Use close instead.")