[FLINK-37525] Introduce a overload findModuleFactory to improve the performance of installModules

flink-runtime/src/main/java/org/apache/flink/runtime/security/NoMatchSecurityFactoryException.java

@@ -19,26 +19,29 @@
 package org.apache.flink.runtime.security;
 
 import java.util.List;
+import java.util.Set;
 
 /** Exception for not finding suitable security factories. */
 public class NoMatchSecurityFactoryException extends RuntimeException {
 
+    private Set<String> missingFactoryClasses;
+
     /**
      * Exception for not finding suitable security factories.
      *
      * @param message message that indicates the current matching step
-     * @param factoryClassCanonicalName required factory class
+     * @param factoryClassCanonicalNames required factory classes
      * @param matchingFactories all found factories
      * @param cause the cause
      */
     public NoMatchSecurityFactoryException(
             String message,
-            String factoryClassCanonicalName,
+            String factoryClassCanonicalNames,
             List<?> matchingFactories,
             Throwable cause) {
         super(
                 "Could not find a suitable security factory for '"
-                        + factoryClassCanonicalName
+                        + factoryClassCanonicalNames
                         + "' in the classpath. all matching factories: "
                         + matchingFactories
                         + ". Reason: "
@@ -50,11 +53,16 @@ public NoMatchSecurityFactoryException(
      * Exception for not finding suitable security factories.
      *
      * @param message message that indicates the current matching step
-     * @param factoryClassCanonicalName required factory class
+     * @param factoryClassCanonicalNames required factory classes
      * @param matchingFactories all found factories
      */
     public NoMatchSecurityFactoryException(
-            String message, String factoryClassCanonicalName, List<?> matchingFactories) {
-        this(message, factoryClassCanonicalName, matchingFactories, null);
+            String message, Set<String> factoryClassCanonicalNames, List<?> matchingFactories) {
+        this(message, String.join(",", factoryClassCanonicalNames), matchingFactories, null);
+        this.missingFactoryClasses = factoryClassCanonicalNames;
+    }
+
+    public Set<String> getMissingFactoryClasses() {
+        return missingFactoryClasses;
     }
 }

flink-runtime/src/main/java/org/apache/flink/runtime/security/SecurityFactoryServiceLoader.java

@@ -23,9 +23,11 @@
 import org.apache.flink.util.Preconditions;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.ServiceLoader;
+import java.util.Set;
 
 /**
  * The Service provider discovery for searching suitable security factory.
@@ -35,7 +37,7 @@
  */
 public class SecurityFactoryServiceLoader {
 
-    /** Find a suitable {@link SecurityModuleFactory} based on canonical name. */
+    /** Find a suitable {@link SecurityModuleFactory}s based on canonical name. */
     public static SecurityModuleFactory findModuleFactory(String securityModuleFactoryClass)
             throws NoMatchSecurityFactoryException {
         return findFactoryInternal(
@@ -80,8 +82,52 @@ private static <T> T findFactoryInternal(
             throw new NoMatchSecurityFactoryException(
                     "zero or more than one security factory found",
                     factoryClassCanonicalName,
-                    matchingFactories);
+                    matchingFactories,
+                    null);
         }
         return matchingFactories.get(0);
     }
+
+    /** Find suitable {@link SecurityModuleFactory}s based on canonical names. */
+    public static List<SecurityModuleFactory> findModuleFactory(
+            List<String> securityModuleFactoryClasses) throws NoMatchSecurityFactoryException {
+        return findFactoryInternal(
+                securityModuleFactoryClasses,
+                SecurityModuleFactory.class,
+                SecurityModuleFactory.class.getClassLoader());
+    }
+
+    private static <T> List<T> findFactoryInternal(
+            List<String> factoryClassCanonicalNames, Class<T> factoryClass, ClassLoader classLoader)
+            throws NoMatchSecurityFactoryException {
+        Preconditions.checkArgument(!factoryClassCanonicalNames.isEmpty());
+        factoryClassCanonicalNames.forEach(className -> Preconditions.checkNotNull(className));
+
+        ServiceLoader<T> serviceLoader;
+        if (classLoader != null) {
+            serviceLoader = ServiceLoader.load(factoryClass, classLoader);
+        } else {
+            serviceLoader = ServiceLoader.load(factoryClass);
+        }
+
+        Set<String> factoriesSet = new HashSet<>(factoryClassCanonicalNames);
+        List<T> matchingFactories = new ArrayList<>();
+        Iterator<T> classFactoryIterator = serviceLoader.iterator();
+        classFactoryIterator.forEachRemaining(
+                classFactory -> {
+                    String factoryClassCanonicalName = classFactory.getClass().getCanonicalName();
+                    if (factoriesSet.contains(factoryClassCanonicalName)) {
+                        matchingFactories.add(classFactory);
+                        factoriesSet.remove(factoryClassCanonicalName);
+                    }
+                });
+
+        if (matchingFactories.size() != factoryClassCanonicalNames.size()) {
+            throw new NoMatchSecurityFactoryException(
+                    "zero or more than one security factory found",
+                    factoriesSet,
+                    matchingFactories);
+        }
+        return matchingFactories;
+    }
 }

flink-runtime/src/main/java/org/apache/flink/runtime/security/SecurityUtils.java

@@ -62,14 +62,17 @@ static void installModules(SecurityConfiguration config) throws Exception {
 
         // install the security module factories
         List<SecurityModule> modules = new ArrayList<>();
-        for (String moduleFactoryClass : config.getSecurityModuleFactories()) {
-            SecurityModuleFactory moduleFactory = null;
-            try {
-                moduleFactory = SecurityFactoryServiceLoader.findModuleFactory(moduleFactoryClass);
-            } catch (NoMatchSecurityFactoryException ne) {
-                LOG.error("Unable to instantiate security module factory {}", moduleFactoryClass);
-                throw new IllegalArgumentException("Unable to find module factory class", ne);
-            }
+        List<String> moduleFactoryClasses = config.getSecurityModuleFactories();
+        List<SecurityModuleFactory> moduleFactories = null;
+        try {
+            moduleFactories = SecurityFactoryServiceLoader.findModuleFactory(moduleFactoryClasses);
+        } catch (NoMatchSecurityFactoryException ne) {
+            LOG.error(
+                    "Unable to instantiate security module factories {}",
+                    String.join(",", ne.getMissingFactoryClasses()));
+            throw new IllegalArgumentException("Unable to find module factory classes", ne);
+        }
+        for (SecurityModuleFactory moduleFactory : moduleFactories) {
             SecurityModule module = moduleFactory.createModule(config);
             // can be null if a SecurityModule is not supported in the current environment
             if (module != null) {

flink-test-utils-parent/flink-test-utils/src/main/java/org/apache/flink/test/util/TestingSecurityContext.java

@@ -29,6 +29,7 @@
 
 import javax.security.auth.login.AppConfigurationEntry;
 
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -46,10 +47,11 @@ public static void install(
         SecurityUtils.install(config);
 
         // install dynamic JAAS entries
-        for (String factoryClassName : config.getSecurityModuleFactories()) {
-            SecurityModuleFactory factory =
-                    SecurityFactoryServiceLoader.findModuleFactory(factoryClassName);
-            if (factory instanceof JaasModuleFactory) {
+        List<String> moduleFactoryClasses = config.getSecurityModuleFactories();
+        List<SecurityModuleFactory> moduleFactories =
+                SecurityFactoryServiceLoader.findModuleFactory(moduleFactoryClasses);
+        for (SecurityModuleFactory moduleFactory : moduleFactories) {
+            if (moduleFactory instanceof JaasModuleFactory) {
                 DynamicConfiguration jaasConf =
                         (DynamicConfiguration)
                                 javax.security.auth.login.Configuration.getConfiguration();