[#772] fix(kerberos): cache proxy user ugi to avoid memory leak

common/src/main/java/org/apache/uniffle/common/security/HadoopSecurityContext.java

@@ -19,17 +19,20 @@
 
 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.annotations.VisibleForTesting;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.uniffle.common.util.JavaUtils;
 import org.apache.uniffle.common.util.ThreadUtils;
 
 public class HadoopSecurityContext implements SecurityContext {
@@ -39,6 +42,12 @@ public class HadoopSecurityContext implements SecurityContext {
   private UserGroupInformation loginUgi;
   private ScheduledExecutorService refreshScheduledExecutor;
 
+  // The cache of proxy user ugi is to avoid creating the different cache keys of
+  // Hadoop filesystem for the same user, scheme and authority, this will cache too
+  // much unnecessary filesystem instances in memory and then cause memory leak.
+  // See details in issue #706
+  private Map<String, UserGroupInformation> proxyUserUgiPool;
+
   public HadoopSecurityContext(
       String krb5ConfPath,
       String keytabFile,
@@ -75,6 +84,7 @@ public HadoopSecurityContext(
         refreshIntervalSec,
         refreshIntervalSec,
         TimeUnit.SECONDS);
+    proxyUserUgiPool = JavaUtils.newConcurrentMap();
   }
 
   private void authRefresh() {
@@ -94,8 +104,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
       );
     }
@@ -113,10 +125,20 @@ private <T> T executeWithUgiWrapper(UserGroupInformation ugi, Callable<T> callab
     return ugi.doAs((PrivilegedExceptionAction<T>) callable::call);
   }
 
+  // Only for tests
+  @VisibleForTesting
+  Map<String, UserGroupInformation> getProxyUserUgiPool() {
+    return proxyUserUgiPool;
+  }
+
   @Override
   public void close() throws IOException {
     if (refreshScheduledExecutor != null) {
       refreshScheduledExecutor.shutdown();
     }
+    if (proxyUserUgiPool != null) {
+      proxyUserUgiPool.clear();
+      proxyUserUgiPool = null;
+    }
   }
 }

common/src/test/java/org/apache/uniffle/common/security/HadoopSecurityContextTest.java

@@ -18,12 +18,16 @@
 package org.apache.uniffle.common.security;
 
 import java.util.concurrent.Callable;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.commons.lang3.StringUtils;
 import org.apache.hadoop.fs.FileStatus;
 import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.security.UserGroupInformation;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import org.apache.uniffle.common.KerberizedHdfsBase;
 
@@ -32,6 +36,7 @@
 import static org.junit.jupiter.api.Assertions.fail;
 
 public class HadoopSecurityContextTest extends KerberizedHdfsBase {
+  private static final Logger LOGGER = LoggerFactory.getLogger(HadoopSecurityContextTest.class);
 
   @BeforeAll
   public static void beforeAll() throws Exception {
@@ -66,13 +71,24 @@ public void testSecuredCallable() throws Exception {
 
       // case3: run by the proxy user
       Path pathWithAlexUser = new Path("/alex/HadoopSecurityContextTest");
+      AtomicReference<UserGroupInformation> ugi1 = new AtomicReference<>();
       context.runSecured("alex", (Callable<Void>) () -> {
+        ugi1.set(UserGroupInformation.getCurrentUser());
         kerberizedHdfs.getFileSystem().mkdirs(pathWithAlexUser);
         return null;
       });
       fileStatus = kerberizedHdfs.getFileSystem().getFileStatus(pathWithAlexUser);
       assertEquals("alex", fileStatus.getOwner());
 
+      // case4: run by the proxy user again, it will always return the same
+      // ugi.
+      AtomicReference<UserGroupInformation> ugi2 = new AtomicReference<>();
+      context.runSecured("alex", (Callable<Void>) () -> {
+        ugi2.set(UserGroupInformation.getCurrentUser());
+        return null;
+      });
+      assertEquals(ugi1.get(), ugi2.get());
+      assertEquals(ugi1.get(), context.getProxyUserUgiPool().get("alex"));
     }
   }