[kyuubi] branch master updated: [KYUUBI #4270] Register shutdown hook for plugin cleanup

[bo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

bowenliang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/kyuubi.git


The following commit(s) were added to refs/heads/master by this push:
     new 0eff3cec5 [KYUUBI #4270] Register shutdown hook for plugin cleanup
0eff3cec5 is described below

commit 0eff3cec535464f08cf6162971b4a80a48536132
Author: liangbowen <li...@gf.com.cn>
AuthorDate: Wed Feb 8 22:56:43 2023 +0800

    [KYUUBI #4270] Register shutdown hook for plugin cleanup
    
    ### _Why are the changes needed?_
    
    to fix #4270.
    - Register shutdown hook for plugin cleanup
    
    ### _How was this patch tested?_
    - [ ] Add some test cases that check the changes thoroughly including negative and positive cases if possible
    
    - [ ] Add screenshots for manual tests if appropriate
    
    - [x] [Run test](https://kyuubi.readthedocs.io/en/master/develop_tools/testing.html#running-tests) locally before make a pull request
    
    Closes #4272 from bowenliang123/authz-shutdown-hook.
    
    Closes #4270
    
    3f376f43 [liangbowen] add log
    6a2018f6 [liangbowen] nit
    8b7ed86c [liangbowen] register plugin's cleanup shutdownHook
    
    Authored-by: liangbowen <li...@gf.com.cn>
    Signed-off-by: liangbowen <li...@gf.com.cn>
---
 .../spark/authz/ranger/RangerSparkExtension.scala  |  2 +-
 .../authz/ranger/SparkRangerAdminPlugin.scala      | 26 ++++++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala
index f4dcb3f9f..5708dfeaf 100644
--- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala
+++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala
@@ -36,7 +36,7 @@ import org.apache.kyuubi.plugin.spark.authz.util.{RuleEliminateMarker, RuleElimi
  *  @since 1.6.0
  */
 class RangerSparkExtension extends (SparkSessionExtensions => Unit) {
-  SparkRangerAdminPlugin.init()
+  SparkRangerAdminPlugin.initialize()
 
   override def apply(v1: SparkSessionExtensions): Unit = {
     v1.injectCheckRule(AuthzConfigurationChecker)
diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/SparkRangerAdminPlugin.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/SparkRangerAdminPlugin.scala
index 7ece55fe5..3d46563aa 100644
--- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/SparkRangerAdminPlugin.scala
+++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/SparkRangerAdminPlugin.scala
@@ -21,8 +21,10 @@ import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
 import scala.collection.mutable.LinkedHashMap
 
+import org.apache.hadoop.util.ShutdownHookManager
 import org.apache.ranger.plugin.policyengine.RangerAccessRequest
 import org.apache.ranger.plugin.service.RangerBasePlugin
+import org.slf4j.LoggerFactory
 
 import org.apache.kyuubi.plugin.spark.authz.AccessControlException
 import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils._
@@ -30,6 +32,7 @@ import org.apache.kyuubi.plugin.spark.authz.util.RangerConfigProvider
 
 object SparkRangerAdminPlugin extends RangerBasePlugin("spark", "sparkSql")
   with RangerConfigProvider {
+  final private val LOG = LoggerFactory.getLogger(getClass)
 
   /**
    * For a Spark SQL query, it may contain 0 or more privilege objects to verify, e.g. a typical
@@ -60,6 +63,29 @@ object SparkRangerAdminPlugin extends RangerBasePlugin("spark", "sparkSql")
     s"ranger.plugin.$getServiceType.use.usergroups.from.userstore.enabled",
     false)
 
+  /**
+   * plugin initialization
+   * with cleanup shutdown hook registered
+   */
+  def initialize(): Unit = {
+    this.init()
+    registerCleanupShutdownHook(this)
+  }
+
+  /**
+   * register shutdown hook for plugin cleanup
+   */
+  private def registerCleanupShutdownHook(plugin: RangerBasePlugin): Unit = {
+    ShutdownHookManager.get().addShutdownHook(
+      () => {
+        if (plugin != null) {
+          LOG.info(s"clean up ranger plugin, appId: ${plugin.getAppId}")
+          this.cleanup()
+        }
+      },
+      Integer.MAX_VALUE)
+  }
+
   def getFilterExpr(req: AccessRequest): Option[String] = {
     val result = evalRowFilterPolicies(req, null)
     Option(result)