[KYUUBI #7079] Improve performance of AccessRequest initialization

extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/AccessRequest.scala

@@ -21,17 +21,21 @@ import java.util.{HashMap => JHashMap, Set => JSet}
 import java.util.Date
 
 import scala.collection.JavaConverters._
+import scala.util.{Failure, Try}
 
 import org.apache.hadoop.security.UserGroupInformation
 import org.apache.ranger.plugin.policyengine.{RangerAccessRequestImpl, RangerPolicyEngine}
+import org.apache.spark.internal.Logging
 
 import org.apache.kyuubi.plugin.spark.authz.OperationType.OperationType
 import org.apache.kyuubi.plugin.spark.authz.ranger.AccessType._
+import org.apache.kyuubi.util.reflect.DynMethods
+import org.apache.kyuubi.util.reflect.DynMethods.{BoundMethod, UnboundMethod}
 import org.apache.kyuubi.util.reflect.ReflectUtils._
 
 case class AccessRequest private (accessType: AccessType) extends RangerAccessRequestImpl
 
-object AccessRequest {
+object AccessRequest extends Logging {
   def apply(
       resource: AccessResource,
       user: UserGroupInformation,
@@ -45,12 +49,11 @@ object AccessRequest {
     req.setUserGroups(userGroups)
     req.setAction(opType.toString)
     try {
-      val roles = invokeAs[JSet[String]](
-        SparkRangerAdminPlugin,
-        "getRolesFromUserAndGroups",
-        (classOf[String], userName),
-        (classOf[JSet[String]], userGroups))
-      invokeAs[Unit](req, "setUserRoles", (classOf[JSet[String]], roles))
+      (getRolesFromUserAndGroupsMethod zip setUserRolesMethod).foreach {
+        case (getMethod, setMethod) =>
+          val roles = getMethod.invoke[JSet[String]](userName, userGroups)
+          setMethod.bind(req).invoke[Unit](roles)
+      }
     } catch {
       case _: Exception =>
     }
@@ -68,27 +71,92 @@ object AccessRequest {
     req
   }
 
+  private def logWarningForUserGroupRoles(
+      clazzName: String,
+      methodName: String,
+      cause: Throwable): Unit = {
+    logWarning(
+      s"Unable to find method $methodName from class $clazzName. " +
+        s"The UserGroupRoles feature requires Ranger 2.1 or above.",
+      cause)
+  }
+
+  private val getRolesFromUserAndGroupsMethod: Option[BoundMethod] = Try {
+    DynMethods.builder("getRolesFromUserAndGroups")
+      .impl(SparkRangerAdminPlugin.getClass, classOf[String], classOf[JSet[String]])
+      .build(SparkRangerAdminPlugin)
+  }.recoverWith { case rethrow: Throwable =>
+    logWarningForUserGroupRoles("SparkRangerAdminPlugin", "getRolesFromUserAndGroups", rethrow)
+    Failure(rethrow)
+  }.toOption
+
+  private val setUserRolesMethod: Option[UnboundMethod] = Try {
+    DynMethods.builder("setUserRoles")
+      .impl(classOf[AccessRequest], classOf[JSet[String]])
+      .buildChecked()
+  }.recoverWith { case rethrow: Throwable =>
+    logWarningForUserGroupRoles("AccessRequest", "setUserRoles", rethrow)
+    Failure(rethrow)
+  }.toOption
+
   private def getUserGroupsFromUgi(user: UserGroupInformation): JSet[String] = {
     user.getGroupNames.toSet.asJava
   }
 
+  private def logWarningForUserStore(
+      clazzName: String,
+      methodName: String,
+      cause: Throwable): Unit = {
+    val svcType = SparkRangerAdminPlugin.getServiceType
+    val confKey = s"ranger.plugin.$svcType.use.usergroups.from.userstore.enabled"
+    logWarning(
+      s"Unable to find method $methodName from class $clazzName. " +
+        s"The UserStore feature requires Ranger 2.1 or above, " +
+        s"consider diabling '$confKey' if you don't use this feature.",
+      cause)
+  }
+
+  private lazy val getUserStoreEnricherMethod: Option[BoundMethod] = Try {
+    DynMethods.builder("getUserStoreEnricher")
+      .impl(SparkRangerAdminPlugin.getClass)
+      .build(SparkRangerAdminPlugin)
+  }.recoverWith { case rethrow: Throwable =>
+    logWarningForUserStore("SparkRangerAdminPlugin", "getUserStoreEnricher", rethrow)
+    Failure(rethrow)
+  }.toOption
+
+  private lazy val getRangerUserStoreMethod: Option[UnboundMethod] = Try {
+    DynMethods.builder("getRangerUserStore")
+      .impl("org.apache.ranger.plugin.contextenricher.RangerUserStoreEnricher")
+      .build()
+  }.recoverWith { case rethrow: Throwable =>
+    logWarningForUserStore("RangerUserStoreEnricher", "getRangerUserStore", rethrow)
+    Failure(rethrow)
+  }.toOption
+
+  private lazy val getUserGroupMappingMethod: Option[UnboundMethod] = Try {
+    DynMethods.builder("getUserGroupMapping")
+      .impl("org.apache.ranger.plugin.util.RangerUserStore")
+      .build()
+  }.recoverWith { case rethrow: Throwable =>
+    logWarningForUserStore("RangerUserStore", "getUserGroupMapping", rethrow)
+    Failure(rethrow)
+  }.toOption
+
   private def getUserGroupsFromUserStore(user: UserGroupInformation): Option[JSet[String]] = {
-    try {
-      val storeEnricher = invokeAs[AnyRef](SparkRangerAdminPlugin, "getUserStoreEnricher")
-      val userStore = invokeAs[AnyRef](storeEnricher, "getRangerUserStore")
-      val userGroupMapping =
-        invokeAs[JHashMap[String, JSet[String]]](userStore, "getUserGroupMapping")
-      Some(userGroupMapping.get(user.getShortUserName))
-    } catch {
-      case _: NoSuchMethodException =>
-        None
+    (getUserStoreEnricherMethod, getRangerUserStoreMethod, getUserGroupMappingMethod) match {
+      case (Some(getEnricher), Some(getUserStore), Some(getMapping)) =>
+        val enricher = getEnricher.invoke()
+        val userStore = getUserStore.bind(enricher).invoke()
+        val userGroupMapping: JHashMap[String, JSet[String]] = getMapping.bind(userStore).invoke()
+        Some(userGroupMapping.get(user.getShortUserName))
+      case _ => None
     }
   }
 
   private def getUserGroups(user: UserGroupInformation): JSet[String] = {
     if (SparkRangerAdminPlugin.useUserGroupsFromUserStoreEnabled) {
-      getUserGroupsFromUserStore(user)
-        .getOrElse(getUserGroupsFromUgi(user))
+      getUserGroupsFromUserStore(user).getOrElse(getUserGroupsFromUgi(user))
     } else {
       getUserGroupsFromUgi(user)
     }