Merge existing registry with default one or configure default metric registry (apache-spark-on-k8s#214)

Author: robert3005

core/src/main/scala/org/apache/spark/metrics/MetricsSystem.scala

@@ -22,7 +22,7 @@ import java.util.concurrent.TimeUnit
 
 import scala.collection.mutable
 
-import com.codahale.metrics.{Metric, MetricFilter, MetricRegistry}
+import com.codahale.metrics._
 import org.eclipse.jetty.servlet.ServletContextHandler
 
 import org.apache.spark.{SecurityManager, SparkConf}
@@ -70,14 +70,15 @@ import org.apache.spark.util.Utils
 private[spark] class MetricsSystem private (
     val instance: String,
     conf: SparkConf,
-    securityMgr: SecurityManager)
+    securityMgr: SecurityManager,
+    registry: MetricRegistry)
   extends Logging {
 
   private[this] val metricsConfig = new MetricsConfig(conf)
 
   private val sinks = new mutable.ArrayBuffer[Sink]
-  private val sources = new mutable.ArrayBuffer[Source]
-  private val registry = new MetricRegistry()
+  private val sourceToListeners = new mutable.HashMap[Source, MetricRegistryListener]
+  private val defaultListener = new MetricsSystemListener("")
 
   private var running: Boolean = false
 
@@ -99,13 +100,17 @@ private[spark] class MetricsSystem private (
     running = true
     StaticSources.allSources.foreach(registerSource)
     registerSources()
+    SharedMetricRegistries.getDefault.addListener(defaultListener)
     registerSinks()
-    sinks.foreach(_.start)
+    sinks.foreach(_.start())
   }
 
   def stop() {
     if (running) {
-      sinks.foreach(_.stop)
+      sinks.foreach(_.stop())
+      sourceToListeners.keySet.foreach(deregisterSource)
+      sourceToListeners.clear()
+      SharedMetricRegistries.getDefault.removeListener(defaultListener)
     } else {
       logWarning("Stopping a MetricsSystem that is not running")
     }
@@ -151,26 +156,36 @@ private[spark] class MetricsSystem private (
   }
 
   def getSourcesByName(sourceName: String): Seq[Source] =
-    sources.filter(_.sourceName == sourceName)
+    sourceToListeners.keySet.filter(_.sourceName == sourceName).toSeq
 
   def registerSource(source: Source) {
-    sources += source
     try {
-      val regName = buildRegistryName(source)
-      registry.register(regName, source.metricRegistry)
+      val listener = new MetricsSystemListener(buildRegistryName(source))
+      source.metricRegistry.addListener(listener)
+      sourceToListeners += source -> listener
     } catch {
       case e: IllegalArgumentException => logInfo("Metrics already registered", e)
     }
   }
 
-  def removeSource(source: Source) {
-    sources -= source
+  def deregisterSource(source: Source) {
     val regName = buildRegistryName(source)
     registry.removeMatching(new MetricFilter {
       def matches(name: String, metric: Metric): Boolean = name.startsWith(regName)
     })
+    sourceToListeners.get(source).foreach(source.metricRegistry.removeListener)
   }
 
+  def removeSource(source: Source): Unit = {
+    deregisterSource(source)
+    sourceToListeners.remove(source)
+  }
+
+  def getSources: Seq[Source] =
+    sourceToListeners.keySet.to[collection.immutable.Seq]
+
+  def getSinks: Seq[Sink] = sinks.to[collection.immutable.Seq]
+
   private def registerSources() {
     val instConfig = metricsConfig.getInstance(instance)
     val sourceConfigs = metricsConfig.subProperties(instConfig, MetricsSystem.SOURCE_REGEX)
@@ -211,6 +226,41 @@ private[spark] class MetricsSystem private (
       }
     }
   }
+
+  private[spark] class MetricsSystemListener(prefix: String)
+      extends MetricRegistryListener {
+    def metricName(name: String): String = MetricRegistry.name(prefix, name)
+
+    override def onHistogramAdded(name: String, histogram: Histogram): Unit =
+      registry.register(metricName(name), histogram)
+
+    override def onCounterAdded(name: String, counter: Counter): Unit =
+      registry.register(metricName(name), counter)
+
+    override def onHistogramRemoved(name: String): Unit =
+      registry.remove(metricName(name))
+
+    override def onGaugeRemoved(name: String): Unit =
+      registry.remove(metricName(name))
+
+    override def onMeterRemoved(name: String): Unit =
+      registry.remove(metricName(name))
+
+    override def onTimerAdded(name: String, timer: Timer): Unit =
+      registry.register(metricName(name), timer)
+
+    override def onCounterRemoved(name: String): Unit =
+      registry.remove(metricName(name))
+
+    override def onGaugeAdded(name: String, gauge: Gauge[_]): Unit =
+      registry.register(metricName(name), gauge)
+
+    override def onTimerRemoved(name: String): Unit =
+      registry.remove(metricName(name))
+
+    override def onMeterAdded(name: String, meter: Meter): Unit =
+      registry.register(metricName(name), meter)
+  }
 }
 
 private[spark] object MetricsSystem {
@@ -220,6 +270,10 @@ private[spark] object MetricsSystem {
   private[this] val MINIMAL_POLL_UNIT = TimeUnit.SECONDS
   private[this] val MINIMAL_POLL_PERIOD = 1
 
+  scala.util.control.Exception.ignoring(classOf[IllegalStateException]) {
+    SharedMetricRegistries.setDefault("spark", new MetricRegistry())
+  }
+
   def checkMinimalPollingPeriod(pollUnit: TimeUnit, pollPeriod: Int) {
     val period = MINIMAL_POLL_UNIT.convert(pollPeriod, pollUnit)
     if (period < MINIMAL_POLL_PERIOD) {
@@ -230,6 +284,14 @@ private[spark] object MetricsSystem {
 
   def createMetricsSystem(
       instance: String, conf: SparkConf, securityMgr: SecurityManager): MetricsSystem = {
-    new MetricsSystem(instance, conf, securityMgr)
+    new MetricsSystem(instance, conf, securityMgr, new MetricRegistry())
+  }
+
+  def createMetricsSystem(
+     instance: String,
+     conf: SparkConf,
+     securityMgr: SecurityManager,
+     registry: MetricRegistry): MetricsSystem = {
+    new MetricsSystem(instance, conf, securityMgr, registry)
   }
 }

core/src/test/scala/org/apache/spark/metrics/MetricsSystemSuite.scala

@@ -17,9 +17,7 @@
 
 package org.apache.spark.metrics
 
-import scala.collection.mutable.ArrayBuffer
-
-import com.codahale.metrics.MetricRegistry
+import com.codahale.metrics.{Gauge, MetricRegistry}
 import org.scalatest.{BeforeAndAfter, PrivateMethodTester}
 
 import org.apache.spark.{SecurityManager, SparkConf, SparkFunSuite}
@@ -41,27 +39,23 @@ class MetricsSystemSuite extends SparkFunSuite with BeforeAndAfter with PrivateM
   test("MetricsSystem with default config") {
     val metricsSystem = MetricsSystem.createMetricsSystem("default", conf, securityMgr)
     metricsSystem.start()
-    val sources = PrivateMethod[ArrayBuffer[Source]]('sources)
-    val sinks = PrivateMethod[ArrayBuffer[Source]]('sinks)
 
-    assert(metricsSystem.invokePrivate(sources()).length === StaticSources.allSources.length)
-    assert(metricsSystem.invokePrivate(sinks()).length === 0)
+    assert(metricsSystem.getSources.length === StaticSources.allSources.length)
+    assert(metricsSystem.getSinks.length === 0)
     assert(metricsSystem.getServletHandlers.nonEmpty)
   }
 
   test("MetricsSystem with sources add") {
     val metricsSystem = MetricsSystem.createMetricsSystem("test", conf, securityMgr)
     metricsSystem.start()
-    val sources = PrivateMethod[ArrayBuffer[Source]]('sources)
-    val sinks = PrivateMethod[ArrayBuffer[Source]]('sinks)
 
-    assert(metricsSystem.invokePrivate(sources()).length === StaticSources.allSources.length)
-    assert(metricsSystem.invokePrivate(sinks()).length === 1)
+    assert(metricsSystem.getSources.length === StaticSources.allSources.length)
+    assert(metricsSystem.getSinks.length === 1)
     assert(metricsSystem.getServletHandlers.nonEmpty)
 
     val source = new MasterSource(null)
     metricsSystem.registerSource(source)
-    assert(metricsSystem.invokePrivate(sources()).length === StaticSources.allSources.length + 1)
+    assert(metricsSystem.getSources.length === StaticSources.allSources.length + 1)
   }
 
   test("MetricsSystem with Driver instance") {
@@ -268,4 +262,24 @@ class MetricsSystemSuite extends SparkFunSuite with BeforeAndAfter with PrivateM
     assert(metricName === source.sourceName)
   }
 
+  test("MetricsSystem registers dynamically added metrics") {
+    val registry = new MetricRegistry()
+    val source = new Source {
+      override val sourceName = "dummySource"
+      override val metricRegistry = new MetricRegistry()
+    }
+
+    val instanceName = "testInstance"
+    val metricsSystem = MetricsSystem.createMetricsSystem(
+      instanceName, conf, securityMgr, registry)
+    metricsSystem.registerSource(source)
+    assert(!registry.getNames.contains("dummySource.newMetric"), "Metric shouldn't be registered")
+
+    source.metricRegistry.register("newMetric", new Gauge[Integer] {
+      override def getValue: Integer = 1
+    })
+    assert(registry.getNames.contains("dummySource.newMetric"),
+      "Metric should have been registered")
+  }
+
 }

pom.xml

@@ -742,11 +742,6 @@
         <artifactId>jackson-core</artifactId>
         <version>${fasterxml.jackson.version}</version>
       </dependency>
-      <dependency>
-        <groupId>com.fasterxml.jackson.core</groupId>
-        <artifactId>jackson-core</artifactId>
-        <version>${fasterxml.jackson.version}</version>
-      </dependency>
       <dependency>
         <groupId>com.fasterxml.jackson.core</groupId>
         <artifactId>jackson-databind</artifactId>
@@ -757,11 +752,6 @@
         <artifactId>jackson-annotations</artifactId>
         <version>${fasterxml.jackson.version}</version>
       </dependency>
-      <dependency>
-        <groupId>com.fasterxml.jackson.dataformat</groupId>
-        <artifactId>jackson-dataformat-yaml</artifactId>
-        <version>${fasterxml.jackson.version}</version>
-      </dependency>
       <!-- Guava is excluded because of SPARK-6149.  The Guava version referenced in this module is
            15.0, which causes runtime incompatibility issues. -->
       <dependency>
@@ -820,11 +810,6 @@
         <artifactId>jackson-datatype-jsr310</artifactId>
         <version>${fasterxml.jackson.version}</version>
       </dependency>
-      <dependency>
-        <groupId>com.fasterxml.jackson.datatype</groupId>
-        <artifactId>jackson-datatype-guava</artifactId>
-        <version>${fasterxml.jackson.version}</version>
-      </dependency>
       <dependency>
         <groupId>org.glassfish.jersey.core</groupId>
         <artifactId>jersey-server</artifactId>

streaming/src/test/scala/org/apache/spark/streaming/StreamingContextSuite.scala

@@ -377,13 +377,13 @@ class StreamingContextSuite extends SparkFunSuite with BeforeAndAfter with Timeo
     addInputStream(ssc).register()
     ssc.start()
 
-    val sources = StreamingContextSuite.getSources(ssc.env.metricsSystem)
+    val sources = ssc.env.metricsSystem.getSources
     val streamingSource = StreamingContextSuite.getStreamingSource(ssc)
     assert(sources.contains(streamingSource))
     assert(ssc.getState() === StreamingContextState.ACTIVE)
 
     ssc.stop()
-    val sourcesAfterStop = StreamingContextSuite.getSources(ssc.env.metricsSystem)
+    val sourcesAfterStop = ssc.env.metricsSystem.getSources
     val streamingSourceAfterStop = StreamingContextSuite.getStreamingSource(ssc)
     assert(ssc.getState() === StreamingContextState.STOPPED)
     assert(!sourcesAfterStop.contains(streamingSourceAfterStop))
@@ -979,10 +979,6 @@ package object testPackage extends Assertions {
  * This includes methods to access private methods and fields in StreamingContext and MetricsSystem
  */
 private object StreamingContextSuite extends PrivateMethodTester {
-  private val _sources = PrivateMethod[ArrayBuffer[Source]]('sources)
-  private def getSources(metricsSystem: MetricsSystem): ArrayBuffer[Source] = {
-    metricsSystem.invokePrivate(_sources())
-  }
   private val _streamingSource = PrivateMethod[StreamingSource]('streamingSource)
   private def getStreamingSource(streamingContext: StreamingContext): StreamingSource = {
     streamingContext.invokePrivate(_streamingSource())