From e60f10b927dff0c37cbc7bf3c959872d8ffb976c Mon Sep 17 00:00:00 2001
From: hustfeiwang <wangfei3@corp.netease.com>
Date: Tue, 27 Nov 2018 19:45:55 +0800
Subject: [PATCH] KYUUBI-125 fix the error when create session in
 SparkSessionWithUGI

---
 .../scala/yaooqinn/kyuubi/spark/SparkSessionWithUGI.scala | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/kyuubi-server/src/main/scala/yaooqinn/kyuubi/spark/SparkSessionWithUGI.scala b/kyuubi-server/src/main/scala/yaooqinn/kyuubi/spark/SparkSessionWithUGI.scala
index 8dfe8c8ae..0cde33bd8 100644
--- a/kyuubi-server/src/main/scala/yaooqinn/kyuubi/spark/SparkSessionWithUGI.scala
+++ b/kyuubi-server/src/main/scala/yaooqinn/kyuubi/spark/SparkSessionWithUGI.scala
@@ -156,9 +156,11 @@ class SparkSessionWithUGI(
     conf.setAppName(appName)
     configureSparkConf(sessionConf)
     val totalWaitTime: Long = conf.getTimeAsSeconds(BACKEND_SESSTION_INIT_TIMEOUT)
+    var newContextThread: Thread = null
     try {
       KyuubiHadoopUtil.doAs(user) {
-        newContext().start()
+        newContextThread = newContext()
+        newContextThread.start()
         val context =
           Await.result(promisedSparkContext.future, Duration(totalWaitTime, TimeUnit.SECONDS))
         _sparkSession = ReflectUtils.newInstance(
@@ -181,7 +183,9 @@ class SparkSessionWithUGI(
         throw ke
     } finally {
       SparkSessionWithUGI.setFullyConstructed(userName)
-      newContext().join()
+      if (newContextThread != null) {
+        newContextThread.join()
+      }
     }
 
     KyuubiServerMonitor.setListener(userName, new KyuubiServerListener(conf))