From 70a3005ea7893d1065b33ee2dcd6ddfe454db723 Mon Sep 17 00:00:00 2001
From: wforget <643348094@qq.com>
Date: Wed, 20 Apr 2022 09:36:16 +0800
Subject: [PATCH] [KYUUBI #2422] Wrap close session with try-finally

### _Why are the changes needed?_

close #2422

### _How was this patch tested?_
- [ ] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [ ] Add screenshots for manual tests if appropriate

- [ ] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request

Closes #2423 from wForget/KYUUBI-2422.

Closes #2422

94ac9ed1 [wforget] [KYUUBI-2422] Wrap close session with try-finally

Authored-by: wforget <643348094@qq.com>
Signed-off-by: ulysses-you <ulyssesyou@apache.org>
---
 .../scala/org/apache/kyuubi/session/SessionManager.scala   | 7 +++++--
 .../org/apache/kyuubi/session/KyuubiSessionManager.scala   | 7 +++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/session/SessionManager.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/session/SessionManager.scala
index c0025f9b0..ccc263915 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/session/SessionManager.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/session/SessionManager.scala
@@ -118,8 +118,11 @@ abstract class SessionManager(name: String) extends CompositeService(name) {
       throw KyuubiSQLException(s"Invalid $sessionHandle")
     }
     info(s"$sessionHandle is closed, current opening sessions $getOpenSessionCount")
-    session.close()
-    deleteOperationLogSessionDir(sessionHandle)
+    try {
+      session.close()
+    } finally {
+      deleteOperationLogSessionDir(sessionHandle)
+    }
   }
 
   private def deleteOperationLogSessionDir(sessionHandle: SessionHandle): Unit = {
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala
index 404a2d4e0..97ba4fddf 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala
@@ -95,8 +95,11 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) {
 
   override def closeSession(sessionHandle: SessionHandle): Unit = {
     val session = getSession(sessionHandle)
-    super.closeSession(sessionHandle)
-    limiter.foreach(_.decrement(UserIpAddress(session.user, session.ipAddress)))
+    try {
+      super.closeSession(sessionHandle)
+    } finally {
+      limiter.foreach(_.decrement(UserIpAddress(session.user, session.ipAddress)))
+    }
   }
 
   def openBatchSession(