From 33c573cfb268a43f36a3bd3873067ce7ca263ba7 Mon Sep 17 00:00:00 2001
From: Cheng Pan <chengpan@apache.org>
Date: Mon, 17 Oct 2022 16:35:33 +0800
Subject: [PATCH] [KYUUBI #3627] Support vanilla Jetty for Spark packaged by
 sbt

### _Why are the changes needed?_

Fix #3627.

### _How was this patch tested?_
- [ ] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [x] Add screenshots for manual tests if appropriate

manual test, and the class not found error is gone.

- [ ] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request

Closes #3630 from pan3793/jetty.

Closes #3627

bf9e4f9b [Cheng Pan] private
f6c496f6 [Cheng Pan] nit
25a7c8cf [Cheng Pan] [KYUUBI #3627] Support vanilla Jetty for Spark packaged by sbt

Authored-by: Cheng Pan <chengpan@apache.org>
Signed-off-by: Cheng Pan <chengpan@apache.org>
---
 .../scala/org/apache/spark/ui/EngineTab.scala | 34 ++++++++++++++-----
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/ui/EngineTab.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/ui/EngineTab.scala
index e0c6e1aed..b9632f5b6 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/ui/EngineTab.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/ui/EngineTab.scala
@@ -26,6 +26,7 @@ import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.engine.spark.SparkSQLEngine
 import org.apache.kyuubi.engine.spark.events.EngineEventsStore
 import org.apache.kyuubi.service.ServiceState
+import org.apache.kyuubi.util.ClassUtils
 
 /**
  * Note that [[SparkUITab]] is private for Spark
@@ -61,9 +62,10 @@ case class EngineTab(
 
   sparkUI.foreach { ui =>
     try {
-      // Spark shade the jetty package so here we use reflect
+      // Spark shade the jetty package so here we use reflection
+      val sparkServletContextHandlerClz = loadSparkServletContextHandler
       Class.forName("org.apache.spark.ui.SparkUI")
-        .getMethod("attachHandler", classOf[org.sparkproject.jetty.servlet.ServletContextHandler])
+        .getMethod("attachHandler", sparkServletContextHandlerClz)
         .invoke(
           ui,
           Class.forName("org.apache.spark.ui.JettyUtils")
@@ -71,16 +73,32 @@ case class EngineTab(
               "createRedirectHandler",
               classOf[String],
               classOf[String],
-              classOf[(HttpServletRequest) => Unit],
+              classOf[HttpServletRequest => Unit],
               classOf[String],
               classOf[scala.collection.immutable.Set[String]])
             .invoke(null, "/kyuubi/stop", "/kyuubi", handleKillRequest _, "", Set("GET", "POST")))
     } catch {
-      case NonFatal(e) =>
-        warn(
-          "Failed to attach handler using SparkUI, please check the Spark version. " +
-            s"So the config '${KyuubiConf.ENGINE_UI_STOP_ENABLED.key}' does not work.",
-          e)
+      case NonFatal(cause) => reportInstallError(cause)
+      case cause: NoClassDefFoundError => reportInstallError(cause)
+    }
+  }
+
+  private def reportInstallError(cause: Throwable): Unit = {
+    warn(
+      "Failed to attach handler using SparkUI, please check the Spark version. " +
+        s"So the config '${KyuubiConf.ENGINE_UI_STOP_ENABLED.key}' does not work.",
+      cause)
+  }
+
+  private def loadSparkServletContextHandler: Class[_] = {
+    // [KYUUBI #3627]: the official spark release uses the shaded and relocated jetty classes,
+    // but if use sbt to build for testing, e.g. docker image, it still uses vanilla jetty classes.
+    val shaded = "org.sparkproject.jetty.servlet.ServletContextHandler"
+    val vanilla = "org.eclipse.jetty.servlet.ServletContextHandler"
+    if (ClassUtils.classIsLoadable(shaded)) {
+      Class.forName(shaded)
+    } else {
+      Class.forName(vanilla)
     }
   }