[KYUUBI #7079] Improve performance of AccessRequest initialization

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

@@ -24,14 +24,17 @@ import scala.collection.JavaConverters._
 
 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.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 +48,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]](SparkRangerAdminPlugin, userName, userGroups)
+          setMethod.invoke[Unit](req, roles)
+      }
     } catch {
       case _: Exception =>
     }
@@ -68,17 +70,43 @@ object AccessRequest {
     req
   }
 
+  private val getRolesFromUserAndGroupsMethod: Option[UnboundMethod] =
+    getMethod(
+      SparkRangerAdminPlugin.getClass,
+      "getRolesFromUserAndGroups",
+      classOf[String],
+      classOf[JSet[String]])
+
+  private val setUserRolesMethod: Option[UnboundMethod] =
+    getMethod(classOf[AccessRequest], "setUserRoles", classOf[JSet[String]])
+
   private def getUserGroupsFromUgi(user: UserGroupInformation): JSet[String] = {
     user.getGroupNames.toSet.asJava
   }
 
+  private lazy val getUserStoreEnricherMethod: Option[UnboundMethod] =
+    getMethod(SparkRangerAdminPlugin.getClass, "getUserStoreEnricher")
+
+  private lazy val getRangerUserStoreMethod: Option[UnboundMethod] =
+    getMethod(
+      Class.forName("org.apache.ranger.plugin.contextenricher.RangerUserStoreEnricher"),
+      "getRangerUserStore")
+
+  private lazy val getUserGroupMappingMethod: Option[UnboundMethod] =
+    getMethod(
+      Class.forName("org.apache.ranger.plugin.util.RangerUserStore"),
+      "getUserGroupMapping")
+
   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))
+      getUserStoreEnricherMethod.zip(getRangerUserStoreMethod)
+        .zip(getUserGroupMappingMethod).map {
+          case ((getEnricher, getUserStore), getMapping) =>
+            val enricher = getEnricher.invoke[AnyRef](SparkRangerAdminPlugin)
+            val userStore = getUserStore.invoke[AnyRef](enricher)
+            val userGroupMapping = getMapping.invoke[JHashMap[String, JSet[String]]](userStore)
+            userGroupMapping.get(user.getShortUserName)
+        }.headOption
     } catch {
       case _: NoSuchMethodException =>
         None
@@ -94,4 +122,21 @@ object AccessRequest {
     }
   }
 
+  private def getMethod(
+      clz: Class[_],
+      methodName: String,
+      argClasses: Class[_]*): Option[UnboundMethod] = {
+    try {
+      Some(DynMethods.builder(methodName)
+        .hiddenImpl(clz, argClasses: _*)
+        .impl(clz, argClasses: _*)
+        .buildChecked)
+    } catch {
+      case e: Exception =>
+        logWarning(
+          s"$clz does not have $methodName${argClasses.map(_.getName).mkString("(", ", ", ")")}",
+          e)
+        None
+    }
+  }
 }