Supports fully asynchronous Engine initialization

Author: 韩书航

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

@@ -76,6 +76,7 @@ class ExecuteStatement(
 
   private def waitStatementComplete(): Unit =
     try {
+      executeStatement()
       setState(OperationState.RUNNING)
       var statusResp: TGetOperationStatusResp = null
 
@@ -147,7 +148,18 @@ class ExecuteStatement(
       }
       // see if anymore log could be fetched
       fetchQueryLog()
-    } catch onError()
+    } catch onError().andThen(_ => {
+        if (isClosedOrCanceled && _remoteOpHandle != null) {
+          // close remote operation in case user cancels before engine launched.
+          // In this case, CANCELED/CLOSED -> RUNNING transition will definitely throw Exception.
+          client.cancelOperation(_remoteOpHandle)
+        }
+        if (getSession.sessionManager.getSessionOption(getSession.handle).isEmpty) {
+          // close remote session in case user closes session before engine launched
+          logger.warn(s"Session ${getSession.handle} has been closed, close remote session now")
+          client.closeSession()
+        }
+      })
     finally {
       shutdownTimeoutMonitor()
     }
@@ -168,7 +180,6 @@ class ExecuteStatement(
     if (isTimeoutMonitorEnabled) {
       addTimeoutMonitor(queryTimeout)
     }
-    executeStatement()
     val sessionManager = session.sessionManager
     val asyncOperation: Runnable = () => waitStatementComplete()
     try {

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

@@ -172,7 +172,7 @@ abstract class KyuubiOperation(session: Session) extends AbstractOperation(sessi
   }
 
   override def getResultSetMetadata: TGetResultSetMetadataResp = {
-    if (_remoteOpHandle == null) {
+    if (_remoteOpHandle == null || !hasResultSet) {
       val tColumnDesc = new TColumnDesc()
       tColumnDesc.setColumnName("Result")
       val desc = new TTypeDesc

kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionImpl.scala

@@ -104,7 +104,14 @@ class KyuubiSessionImpl(
   }
 
   @volatile private var _client: KyuubiSyncThriftClient = _
-  def client: KyuubiSyncThriftClient = _client
+  def client: KyuubiSyncThriftClient = {
+    if (sessionConf.get(SESSION_ENGINE_LAUNCH_ASYNC)) {
+      // When `SESSION_ENGINE_LAUNCH_ASYNC` is set to true,
+      // launch engine when client used for the first time.
+      waitForEngineLaunched()
+    }
+    _client
+  }
 
   @volatile private var _engineSessionHandle: SessionHandle = _
 
@@ -180,8 +187,9 @@ class KyuubiSessionImpl(
             _engineSessionHandle =
               engineClient.openSession(protocol, user, passwd, openEngineSessionConf)
             _client = engineClient
-            logSessionInfo(s"Connected to engine [$host:$port]/[${client.engineId.getOrElse("")}]" +
-              s" with ${_engineSessionHandle}]")
+            logSessionInfo(
+              s"Connected to engine [$host:$port]/[${_client.engineId.getOrElse("")}]" +
+                s" with ${_engineSessionHandle}]")
             shouldRetry = false
           } catch {
             case e: TTransportException
@@ -234,12 +242,14 @@ class KyuubiSessionImpl(
 
   override protected def runOperation(operation: Operation): OperationHandle = {
     if (operation != launchEngineOp) {
-      try {
-        waitForEngineLaunched()
-      } catch {
-        case t: Throwable =>
-          operation.close()
-          throw t
+      if (!sessionConf.get(SESSION_ENGINE_LAUNCH_ASYNC)) {
+        try {
+          waitForEngineLaunched()
+        } catch {
+          case t: Throwable =>
+            operation.close()
+            throw t
+        }
       }
       sessionEvent.totalOperations += 1
     }