From 7fd6f71271daf193ce9c5df9c4c2b7a5c8c1eef3 Mon Sep 17 00:00:00 2001
From: Zen <xinjingziranchan@gmail.com>
Date: Fri, 11 Dec 2020 15:49:34 +0800
Subject: [PATCH] ZkClient should close when the session is opened

fixes #260

Squashed commit of the following:

commit ed4fa0cad9f8ad0053035b44a3ca703dd49ba362
Author: zen <xinjingziranchan@gmail.com>
Date:   Fri Dec 11 11:13:32 2020 +0800

    release the zkClient after session established.

commit add46a6acaa19942fcea3175ae253fdae2ef525b
Author: zen <xinjingziranchan@gmail.com>
Date:   Thu Dec 10 14:49:26 2020 +0800

    ZkClient should close when the session is opened.
---
 .../org/apache/kyuubi/session/KyuubiSessionImpl.scala      | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/kyuubi-main/src/main/scala/org/apache/kyuubi/session/KyuubiSessionImpl.scala b/kyuubi-main/src/main/scala/org/apache/kyuubi/session/KyuubiSessionImpl.scala
index 79752b180..8cdbbfa88 100644
--- a/kyuubi-main/src/main/scala/org/apache/kyuubi/session/KyuubiSessionImpl.scala
+++ b/kyuubi-main/src/main/scala/org/apache/kyuubi/session/KyuubiSessionImpl.scala
@@ -17,6 +17,7 @@
 
 package org.apache.kyuubi.session
 
+import java.io.IOException
 import java.nio.charset.StandardCharsets
 import java.util.concurrent.TimeUnit
 
@@ -107,6 +108,11 @@ class KyuubiSessionImpl(
         val Some((host, port)) = getServerHost
         openSession(host, port)
     }
+    try {
+      zkClient.close()
+    } catch {
+      case e: IOException => error("Failed to release the zkClient after session established", e)
+    }
   }
 
   private def openSession(host: String, port: Int): Unit = {
@@ -144,6 +150,5 @@ class KyuubiSessionImpl(
         transport.close()
       }
     }
-    zkClient.close()
   }
 }