From 1d523e687e154b26506b52e9d4e148739f53ab20 Mon Sep 17 00:00:00 2001
From: Junfan Zhang <zuston@apache.org>
Date: Tue, 28 Mar 2023 14:22:31 +0800
Subject: [PATCH] [#772] fix(kerberos): cache proxy user ugi to avoid memory
 leak

---
 .../common/security/HadoopSecurityContext.java       | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/common/src/main/java/org/apache/uniffle/common/security/HadoopSecurityContext.java b/common/src/main/java/org/apache/uniffle/common/security/HadoopSecurityContext.java
index 8f32e507a2..b624f707c7 100644
--- a/common/src/main/java/org/apache/uniffle/common/security/HadoopSecurityContext.java
+++ b/common/src/main/java/org/apache/uniffle/common/security/HadoopSecurityContext.java
@@ -19,11 +19,13 @@
 
 import java.io.IOException;
 import java.security.PrivilegedExceptionAction;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 
+import com.google.common.collect.Maps;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.security.UserGroupInformation;
@@ -38,6 +40,7 @@ public class HadoopSecurityContext implements SecurityContext {
 
   private UserGroupInformation loginUgi;
   private ScheduledExecutorService refreshScheduledExecutor;
+  private Map<String, UserGroupInformation> proxyUserUgiPool;
 
   public HadoopSecurityContext(
       String krb5ConfPath,
@@ -75,6 +78,7 @@ public HadoopSecurityContext(
         refreshIntervalSec,
         refreshIntervalSec,
         TimeUnit.SECONDS);
+    proxyUserUgiPool = Maps.newConcurrentMap();
   }
 
   private void authRefresh() {
@@ -94,8 +98,10 @@ public <T> T runSecured(String user, Callable<T> securedCallable) throws Excepti
 
     // Run with the proxy user.
     if (!user.equals(loginUgi.getShortUserName())) {
+      UserGroupInformation proxyUserUgi =
+          proxyUserUgiPool.computeIfAbsent(user, x -> UserGroupInformation.createProxyUser(x, loginUgi));
       return executeWithUgiWrapper(
-          UserGroupInformation.createProxyUser(user, loginUgi),
+          proxyUserUgi,
           securedCallable
       );
     }
@@ -118,5 +124,9 @@ public void close() throws IOException {
     if (refreshScheduledExecutor != null) {
       refreshScheduledExecutor.shutdown();
     }
+    if (proxyUserUgiPool != null) {
+      proxyUserUgiPool.clear();
+      proxyUserUgiPool = null;
+    }
   }
 }