Slight change to pr-784

5 files changed, 10 insertions, 9 deletions

diff --git a/core/src/main/scala/spark/deploy/master/ui/MasterWebUI.scala b/core/src/main/scala/spark/deploy/master/ui/MasterWebUI.scala
index f0a6ffe047..c91e1db9f2 100644
--- a/core/src/main/scala/spark/deploy/master/ui/MasterWebUI.scala
+++ b/core/src/main/scala/spark/deploy/master/ui/MasterWebUI.scala
@@ -59,9 +59,8 @@ class MasterWebUI(val master: Master, requestedPort: Int) extends Logging {
     }
   }
 
-  val metricsHandlers = master.masterMetricsSystem.metricsServlet.map(_.getHandlers)
-    .getOrElse(Array()) ++ master.applicationMetricsSystem.metricsServlet.map(_.getHandlers)
-    .getOrElse(Array())
+  val metricsHandlers = master.masterMetricsSystem.getServletHandlers ++
+    master.applicationMetricsSystem.getServletHandlers
 
   val handlers = metricsHandlers ++ Array[(String, Handler)](
     ("/static", createStaticHandler(MasterWebUI.STATIC_RESOURCE_DIR)),
diff --git a/core/src/main/scala/spark/deploy/worker/ui/WorkerWebUI.scala b/core/src/main/scala/spark/deploy/worker/ui/WorkerWebUI.scala
index b408c63a02..22295069dc 100644
--- a/core/src/main/scala/spark/deploy/worker/ui/WorkerWebUI.scala
+++ b/core/src/main/scala/spark/deploy/worker/ui/WorkerWebUI.scala
@@ -48,7 +48,7 @@ class WorkerWebUI(val worker: Worker, val workDir: File, requestedPort: Option[I
 
   val indexPage = new IndexPage(this)
 
-  val metricsHandlers = worker.metricsSystem.metricsServlet.map(_.getHandlers).getOrElse(Array())
+  val metricsHandlers = worker.metricsSystem.getServletHandlers
 
   val handlers = metricsHandlers ++ Array[(String, Handler)](
     ("/static", createStaticHandler(WorkerWebUI.STATIC_RESOURCE_DIR)),
diff --git a/core/src/main/scala/spark/metrics/MetricsSystem.scala b/core/src/main/scala/spark/metrics/MetricsSystem.scala
index 04c750b17e..4e6c6b26c8 100644
--- a/core/src/main/scala/spark/metrics/MetricsSystem.scala
+++ b/core/src/main/scala/spark/metrics/MetricsSystem.scala
@@ -73,7 +73,10 @@ private[spark] class MetricsSystem private (val instance: String) extends Loggin
   val registry = new MetricRegistry()
 
   // Treat MetricsServlet as a special sink as it should be exposed to add handlers to web ui
-  var metricsServlet: Option[MetricsServlet] = None
+  private var metricsServlet: Option[MetricsServlet] = None
+
+  /** Get any UI handlers used by this metrics system. */
+  def getServletHandlers = metricsServlet.map(_.getHandlers).getOrElse(Array())
 
   metricsConfig.initialize()
   registerSources()
diff --git a/core/src/main/scala/spark/ui/SparkUI.scala b/core/src/main/scala/spark/ui/SparkUI.scala
index 4bcfdeb62b..1fd5a0989e 100644
--- a/core/src/main/scala/spark/ui/SparkUI.scala
+++ b/core/src/main/scala/spark/ui/SparkUI.scala
@@ -45,8 +45,7 @@ private[spark] class SparkUI(sc: SparkContext) extends Logging {
   val exec = new ExecutorsUI(sc)
 
   // Add MetricsServlet handlers by default
-  val metricsServletHandlers = SparkEnv.get.metricsSystem.metricsServlet.map(_.getHandlers)
-    .getOrElse(Array())
+  val metricsServletHandlers = SparkEnv.get.metricsSystem.getServletHandlers
 
   val allHandlers = storage.getHandlers ++ jobs.getHandlers ++ env.getHandlers ++
     exec.getHandlers ++ metricsServletHandlers ++ handlers
diff --git a/core/src/test/scala/spark/metrics/MetricsSystemSuite.scala b/core/src/test/scala/spark/metrics/MetricsSystemSuite.scala
index 35c2ae41e9..dc65ac6994 100644
--- a/core/src/test/scala/spark/metrics/MetricsSystemSuite.scala
+++ b/core/src/test/scala/spark/metrics/MetricsSystemSuite.scala
@@ -34,7 +34,7 @@ class MetricsSystemSuite extends FunSuite with BeforeAndAfter {
 
     assert(sources.length === 0)
     assert(sinks.length === 0)
-    assert(metricsSystem.metricsServlet != None)
+    assert(!metricsSystem.getServletHandlers.isEmpty)
   }
 
   test("MetricsSystem with sources add") {
@@ -44,7 +44,7 @@ class MetricsSystemSuite extends FunSuite with BeforeAndAfter {
 
     assert(sources.length === 0)
     assert(sinks.length === 1)
-    assert(metricsSystem.metricsServlet != None)
+    assert(!metricsSystem.getServletHandlers.isEmpty)
 
     val source = new spark.deploy.master.MasterSource(null)
     metricsSystem.registerSource(source)