From 35a6b9b341d21b68c44c094e5ea42a50d97fad32 Mon Sep 17 00:00:00 2001
From: Fei Wang <fwang12@ebay.com>
Date: Tue, 8 Mar 2022 10:22:26 +0800
Subject: [PATCH] [KYUUBI #2044] Remove authentication thread local objects to
 prevent memory leak

### _Why are the changes needed?_

Remove thread locals to prevent memory leak.

### _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 #2044 from turboFei/remove_thread_local.

Closes #2044

bcce9e31 [Fei Wang] remove thread local

Authored-by: Fei Wang <fwang12@ebay.com>
Signed-off-by: Fei Wang <fwang12@ebay.com>
---
 .../HadoopThriftAuthBridgeServer.scala        | 58 ++++++++++---------
 1 file changed, 32 insertions(+), 26 deletions(-)

diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/HadoopThriftAuthBridgeServer.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/HadoopThriftAuthBridgeServer.scala
index 53d8e1835..f51615669 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/HadoopThriftAuthBridgeServer.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/HadoopThriftAuthBridgeServer.scala
@@ -143,38 +143,44 @@ object HadoopThriftAuthBridgeServer {
           REMOTE_ADDRESS.set(socket.getInetAddress)
           val mechanismName = saslServer.getMechanismName
           USER_AUTH_MECHANISM.set(mechanismName)
-          if (AuthMethod.PLAIN.getMechanismName.equalsIgnoreCase(mechanismName)) {
-            REMOTE_USER.set(endUser)
-            wrapped.process(in, out)
-          } else {
-            if (AuthMethod.TOKEN.getMechanismName.equalsIgnoreCase(mechanismName)) {
-              try {
-                val identifier = SaslRpcServer.getIdentifier(authId, secretMgr)
-                endUser = identifier.getUser.getUserName
-              } catch {
-                case e: InvalidToken => throw new TException(e.getMessage)
-              }
-            }
-            val clientUgi: UserGroupInformation = UserGroupInformation.createRemoteUser(endUser)
-            try {
-              REMOTE_USER.set(clientUgi.getShortUserName)
-              debug(s"SET REMOTE USER: ${REMOTE_USER.get()} from endUser: $clientUgi")
+          try {
+            if (AuthMethod.PLAIN.getMechanismName.equalsIgnoreCase(mechanismName)) {
+              REMOTE_USER.set(endUser)
               wrapped.process(in, out)
-            } catch {
-              case e: RuntimeException => e.getCause match {
-                  case t: TException => throw t
-                  case _ => throw e
+            } else {
+              if (AuthMethod.TOKEN.getMechanismName.equalsIgnoreCase(mechanismName)) {
+                try {
+                  val identifier = SaslRpcServer.getIdentifier(authId, secretMgr)
+                  endUser = identifier.getUser.getUserName
+                } catch {
+                  case e: InvalidToken => throw new TException(e.getMessage)
                 }
-              case e: InterruptedException => throw new RuntimeException(e)
-              case e: IOException => throw new RuntimeException(e)
-            } finally {
+              }
+              val clientUgi: UserGroupInformation = UserGroupInformation.createRemoteUser(endUser)
               try {
-                FileSystem.closeAllForUGI(clientUgi)
+                REMOTE_USER.set(clientUgi.getShortUserName)
+                debug(s"SET REMOTE USER: ${REMOTE_USER.get()} from endUser: $clientUgi")
+                wrapped.process(in, out)
               } catch {
-                case e: IOException =>
-                  error(s"Could not clean up file-system handles for UGI: $clientUgi", e)
+                case e: RuntimeException => e.getCause match {
+                    case t: TException => throw t
+                    case _ => throw e
+                  }
+                case e: InterruptedException => throw new RuntimeException(e)
+                case e: IOException => throw new RuntimeException(e)
+              } finally {
+                try {
+                  FileSystem.closeAllForUGI(clientUgi)
+                } catch {
+                  case e: IOException =>
+                    error(s"Could not clean up file-system handles for UGI: $clientUgi", e)
+                }
               }
             }
+          } finally {
+            REMOTE_USER.remove()
+            REMOTE_ADDRESS.remove()
+            USER_AUTH_MECHANISM.remove()
           }
 
         case _ => throw new TException(s"Unexpected non-SASL transport ${transport.getClass}")