Address comments (move heartbeater from DAGScheduler to SparkContext,…

… move logic for getting

metrics to Heartbeater), and modifiy tests for the new ExecutorMetrics format.

core/src/main/scala/org/apache/spark/Heartbeater.scala

@@ -19,17 +19,26 @@ package org.apache.spark
 
 import java.util.concurrent.TimeUnit
 
+import org.apache.spark.executor.ExecutorMetrics
+import org.apache.spark.internal.Logging
+import org.apache.spark.memory.MemoryManager
+import org.apache.spark.metrics.MetricGetter
 import org.apache.spark.util.{ThreadUtils, Utils}
 
 /**
  * Creates a heartbeat thread which will call the specified reportHeartbeat function at
  * intervals of intervalMs.
  *
+ * @param memoryManager the memory manager for execution and storage memory.
  * @param reportHeartbeat the heartbeat reporting function to call.
  * @param name the thread name for the heartbeater.
  * @param intervalMs the interval between heartbeats.
  */
-private[spark] class Heartbeater(reportHeartbeat: () => Unit, name: String, intervalMs: Long) {
+private[spark] class Heartbeater(
+    memoryManager: MemoryManager,
+    reportHeartbeat: () => Unit,
+    name: String,
+    intervalMs: Long) extends Logging {
   // Executor for the heartbeat task
   private val heartbeater = ThreadUtils.newDaemonSingleThreadScheduledExecutor(name)
 
@@ -49,5 +58,11 @@ private[spark] class Heartbeater(reportHeartbeat: () => Unit, name: String, inte
     heartbeater.shutdown()
     heartbeater.awaitTermination(10, TimeUnit.SECONDS)
   }
+
+  /** Get the current metrics. */
+  def getCurrentMetrics(): ExecutorMetrics = {
+    new ExecutorMetrics(System.currentTimeMillis(),
+      MetricGetter.values.map(_.getMetricValue(memoryManager)).toArray)
+  }
 }
 

core/src/main/scala/org/apache/spark/SparkContext.scala

@@ -213,6 +213,7 @@ class SparkContext(config: SparkConf) extends Logging {
   private var _files: Seq[String] = _
   private var _shutdownHookRef: AnyRef = _
   private var _statusStore: AppStatusStore = _
+  private var _heartbeater: Heartbeater = _
 
   /* ------------------------------------------------------------------------------------- *
    | Accessors and public fields. These provide access to the internal state of the        |
@@ -304,6 +305,11 @@ class SparkContext(config: SparkConf) extends Logging {
     _dagScheduler = ds
   }
 
+  private[spark] def heartbeater: Heartbeater = _heartbeater
+  private[spark] def heartbeater_=(hb: Heartbeater): Unit = {
+    _heartbeater = hb
+  }
+
   /**
    * A unique identifier for the Spark application.
    * Its format depends on the scheduler implementation.
@@ -496,6 +502,11 @@ class SparkContext(config: SparkConf) extends Logging {
     _dagScheduler = new DAGScheduler(this)
     _heartbeatReceiver.ask[Boolean](TaskSchedulerIsSet)
 
+    // create and start the heartbeater for collecting memory metrics
+    _heartbeater = new Heartbeater(env.memoryManager, reportHeartBeat, "driver-heartbeater",
+      conf.getTimeAsMs("spark.executor.heartbeatInterval", "10s"))
+    _heartbeater.start()
+
     // start TaskScheduler after taskScheduler sets DAGScheduler reference in DAGScheduler's
     // constructor
     _taskScheduler.start()
@@ -1922,6 +1933,12 @@ class SparkContext(config: SparkConf) extends Logging {
     Utils.tryLogNonFatalError {
       _eventLogger.foreach(_.stop())
     }
+    if(_heartbeater != null) {
+      Utils.tryLogNonFatalError {
+        _heartbeater.stop()
+      }
+      _heartbeater = null
+    }
     if (_dagScheduler != null) {
       Utils.tryLogNonFatalError {
         _dagScheduler.stop()
@@ -2398,6 +2415,14 @@ class SparkContext(config: SparkConf) extends Logging {
     }
   }
 
+  /** Reports heartbeat metrics for the driver. */
+  private def reportHeartBeat(): Unit = {
+    val driverUpdates = _heartbeater.getCurrentMetrics()
+    val accumUpdates = new Array[(Long, Int, Int, Seq[AccumulableInfo])](0)
+    listenerBus.post(SparkListenerExecutorMetricsUpdate("driver", accumUpdates,
+      Some(driverUpdates)))
+  }
+
   // In order to prevent multiple SparkContexts from being active at the same time, mark this
   // context as having finished construction.
   // NOTE: this must be placed at the end of the SparkContext constructor.

core/src/main/scala/org/apache/spark/executor/Executor.scala

@@ -36,8 +36,7 @@ import org.apache.spark._
 import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config._
-import org.apache.spark.memory.{MemoryManager, SparkOutOfMemoryError, TaskMemoryManager}
-import org.apache.spark.metrics.MetricGetter
+import org.apache.spark.memory.{SparkOutOfMemoryError, TaskMemoryManager}
 import org.apache.spark.rpc.RpcTimeout
 import org.apache.spark.scheduler._
 import org.apache.spark.shuffle.FetchFailedException
@@ -149,8 +148,8 @@ private[spark] class Executor(
   private val runningTasks = new ConcurrentHashMap[Long, TaskRunner]
 
   // Executor for the heartbeat task.
-  private val heartbeater = new Heartbeater(reportHeartBeat, "executor-heartbeater",
-    conf.getTimeAsMs("spark.executor.heartbeatInterval", "10s"))
+  private val heartbeater = new Heartbeater(env.memoryManager, reportHeartBeat,
+    "executor-heartbeater", conf.getTimeAsMs("spark.executor.heartbeatInterval", "10s"))
 
   // must be initialized before running startDriverHeartbeat()
   private val heartbeatReceiverRef =
@@ -789,7 +788,7 @@ private[spark] class Executor(
     val curGCTime = computeTotalGcTime()
 
     // get executor level memory metrics
-    val executorUpdates = Executor.getCurrentExecutorMetrics(env.memoryManager)
+    val executorUpdates = heartbeater.getCurrentMetrics()
 
     for (taskRunner <- runningTasks.values().asScala) {
       if (taskRunner.task != null) {
@@ -827,15 +826,4 @@ private[spark] object Executor {
   // task is fully deserialized. When possible, the TaskContext.getLocalProperty call should be
   // used instead.
   val taskDeserializationProps: ThreadLocal[Properties] = new ThreadLocal[Properties]
-
-  /**
-   * Get the current executor level memory metrics.
-   */
-  def getCurrentExecutorMetrics(memoryManager: MemoryManager): ExecutorMetrics = {
-    val metrics = new ExecutorMetrics(System.currentTimeMillis())
-    MetricGetter.idxAndValues.foreach { case (idx, metric) =>
-      metrics.metrics(idx) = metric.getMetricValue(memoryManager)
-    }
-    metrics
-  }
 }

core/src/main/scala/org/apache/spark/executor/ExecutorMetrics.scala

@@ -29,8 +29,15 @@ import org.apache.spark.metrics.MetricGetter
  *
  * @param timestamp the time the metrics were collected, or -1 for Spark history
  *                  log events which are logged when a stage has completed
+ * @param metrics the array of executor metrics values, order and elements as
+ *                specified in MetricGetter
  */
 @DeveloperApi
-class ExecutorMetrics private[spark] (val timestamp: Long) extends Serializable {
-  val metrics = new Array[Long](MetricGetter.values.length)
+class ExecutorMetrics private[spark] (
+    val timestamp: Long,
+    val metrics: Array[Long]) extends Serializable {
+  if (metrics.length != MetricGetter.values.length) {
+    throw new IllegalArgumentException("invalid metrics length " + metrics.length +
+      " does not equal expected length " + MetricGetter.values.length)
+  }
 }

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

@@ -23,7 +23,7 @@ import org.apache.spark.memory.MemoryManager
 
 sealed trait MetricGetter {
   def getMetricValue(memoryManager: MemoryManager): Long
-  val name = getClass().getName().stripSuffix("$")
+  val name = getClass().getName().stripSuffix("$").split("""\.""").last
 }
 
 abstract class MemoryManagerMetricGetter(f: MemoryManager => Long) extends MetricGetter {
@@ -53,13 +53,19 @@ case object JVMOffHeapMemory extends MetricGetter {
   }
 }
 
-case object OnHeapExecution extends MemoryManagerMetricGetter(_.onHeapExecutionMemoryUsed)
+case object OnHeapExecutionMemory extends MemoryManagerMetricGetter(_.onHeapExecutionMemoryUsed)
 
-case object OffHeapExecution extends MemoryManagerMetricGetter(_.offHeapExecutionMemoryUsed)
+case object OffHeapExecutionMemory extends MemoryManagerMetricGetter(_.offHeapExecutionMemoryUsed)
 
-case object OnHeapStorage extends MemoryManagerMetricGetter(_.onHeapStorageMemoryUsed)
+case object OnHeapStorageMemory extends MemoryManagerMetricGetter(_.onHeapStorageMemoryUsed)
 
-case object OffHeapStorage extends MemoryManagerMetricGetter(_.offHeapStorageMemoryUsed)
+case object OffHeapStorageMemory extends MemoryManagerMetricGetter(_.offHeapStorageMemoryUsed)
+
+case object OnHeapUnifiedMemory extends MemoryManagerMetricGetter(
+  (m => m.onHeapExecutionMemoryUsed + m.onHeapStorageMemoryUsed))
+
+case object OffHeapUnifiedMemory extends MemoryManagerMetricGetter(
+  (m => m.offHeapExecutionMemoryUsed + m.offHeapStorageMemoryUsed))
 
 case object DirectPoolMemory extends MBeanMetricGetter("java.nio:type=BufferPool,name=direct")
 case object MappedPoolMemory extends MBeanMetricGetter("java.nio:type=BufferPool,name=mapped")
@@ -68,10 +74,12 @@ object MetricGetter {
   val values = IndexedSeq(
     JVMHeapMemory,
     JVMOffHeapMemory,
-    OnHeapExecution,
-    OffHeapExecution,
-    OnHeapStorage,
-    OffHeapStorage,
+    OnHeapExecutionMemory,
+    OffHeapExecutionMemory,
+    OnHeapStorageMemory,
+    OffHeapStorageMemory,
+    OnHeapUnifiedMemory,
+    OffHeapUnifiedMemory,
     DirectPoolMemory,
     MappedPoolMemory
   )

core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala

@@ -18,7 +18,6 @@
 package org.apache.spark.scheduler
 
 import java.io.NotSerializableException
-import java.lang.management.ManagementFactory
 import java.util.Properties
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicInteger
@@ -35,7 +34,7 @@ import org.apache.commons.lang3.SerializationUtils
 
 import org.apache.spark._
 import org.apache.spark.broadcast.Broadcast
-import org.apache.spark.executor.{Executor, ExecutorMetrics, TaskMetrics}
+import org.apache.spark.executor.{ExecutorMetrics, TaskMetrics}
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config
 import org.apache.spark.network.util.JavaUtils
@@ -210,10 +209,6 @@ class DAGScheduler(
   private[spark] val eventProcessLoop = new DAGSchedulerEventProcessLoop(this)
   taskScheduler.setDAGScheduler(this)
 
-  /** driver heartbeat for collecting metrics */
-  private val heartbeater: Heartbeater = new Heartbeater(reportHeartBeat, "driver-heartbeater",
-    sc.conf.getTimeAsMs("spark.executor.heartbeatInterval", "10s"))
-
   /**
    * Called by the TaskSetManager to report task's starting.
    */
@@ -1758,20 +1753,9 @@ class DAGScheduler(
     messageScheduler.shutdownNow()
     eventProcessLoop.stop()
     taskScheduler.stop()
-    heartbeater.stop()
-  }
-
-  /** Reports heartbeat metrics for the driver. */
-  private def reportHeartBeat(): Unit = {
-    // get driver memory metrics
-    val driverUpdates = Executor.getCurrentExecutorMetrics(sc.env.memoryManager)
-    val accumUpdates = new Array[(Long, Int, Int, Seq[AccumulableInfo])](0)
-    listenerBus.post(SparkListenerExecutorMetricsUpdate("driver", accumUpdates,
-      Some(driverUpdates)))
-  }
+   }
 
   eventProcessLoop.start()
-  heartbeater.start()
 }
 
 private[scheduler] class DAGSchedulerEventProcessLoop(dagScheduler: DAGScheduler)

core/src/main/scala/org/apache/spark/scheduler/EventLoggingListener.scala

@@ -24,7 +24,7 @@ import java.util.EnumSet
 import java.util.Locale
 
 import scala.collection.mutable
-import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable.{ArrayBuffer, HashMap}
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FileSystem, FSDataOutputStream, Path}
@@ -84,7 +84,7 @@ private[spark] class EventLoggingListener(
   private val compressionCodecName = compressionCodec.map { c =>
     CompressionCodec.getShortName(c.getClass.getName)
   }
-logInfo("spark.eventLog.logExecutorMetricsUpdates.enabled is " + shouldLogExecutorMetricsUpdates)
+
   // Only defined if the file system scheme is not local
   private var hadoopDataStream: Option[FSDataOutputStream] = None
 
@@ -97,8 +97,7 @@ logInfo("spark.eventLog.logExecutorMetricsUpdates.enabled is " + shouldLogExecut
   private[scheduler] val logPath = getLogPath(logBaseDir, appId, appAttemptId, compressionCodecName)
 
   // map of live stages, to peak executor metrics for the stage
-  private val liveStageExecutorMetrics = mutable.HashMap[(Int, Int),
-    mutable.HashMap[String, PeakExecutorMetrics]]()
+  private val liveStageExecutorMetrics = HashMap[(Int, Int), HashMap[String, PeakExecutorMetrics]]()
 
   /**
    * Creates the log file in the configured log directory.
@@ -167,7 +166,7 @@ logInfo("spark.eventLog.logExecutorMetricsUpdates.enabled is " + shouldLogExecut
     if (shouldLogExecutorMetricsUpdates) {
       // record the peak metrics for the new stage
       liveStageExecutorMetrics.put((event.stageInfo.stageId, event.stageInfo.attemptNumber()),
-        new mutable.HashMap[String, PeakExecutorMetrics]())
+        new HashMap[String, PeakExecutorMetrics]())
     }
   }
 
@@ -190,16 +189,15 @@ logInfo("spark.eventLog.logExecutorMetricsUpdates.enabled is " + shouldLogExecut
         liveStageExecutorMetrics.remove((event.stageInfo.stageId, attemptId))
       }
 
-      // log the peak executor metrics for the stage, for each executor
+      // log the peak executor metrics for the stage, for each live executor,
+      // whether or not the executor is running tasks for the stage
       val accumUpdates = new ArrayBuffer[(Long, Int, Int, Seq[AccumulableInfo])]()
       val executorMap = liveStageExecutorMetrics.remove(
         (event.stageInfo.stageId, event.stageInfo.attemptNumber()))
       executorMap.foreach {
        executorEntry => {
           for ((executorId, peakExecutorMetrics) <- executorEntry) {
-            val executorMetrics = new ExecutorMetrics(-1)
-            System.arraycopy(peakExecutorMetrics.metrics, 0, executorMetrics.metrics, 0,
-              peakExecutorMetrics.metrics.size)
+            val executorMetrics = new ExecutorMetrics(-1, peakExecutorMetrics.metrics)
             val executorUpdate = new SparkListenerExecutorMetricsUpdate(
               executorId, accumUpdates, Some(executorMetrics))
             logEvent(executorUpdate)
@@ -346,7 +344,7 @@ private[spark] object EventLoggingListener extends Logging {
   private val LOG_FILE_PERMISSIONS = new FsPermission(Integer.parseInt("770", 8).toShort)
 
   // A cache for compression codecs to avoid creating the same codec many times
-  private val codecMap = new mutable.HashMap[String, CompressionCodec]
+  private val codecMap = new HashMap[String, CompressionCodec]
 
   /**
    * Write metadata about an event log to the given stream.

core/src/main/scala/org/apache/spark/scheduler/PeakExecutorMetrics.scala

@@ -19,7 +19,6 @@ package org.apache.spark.scheduler
 
 import org.apache.spark.executor.ExecutorMetrics
 import org.apache.spark.metrics.MetricGetter
-import org.apache.spark.status.api.v1.PeakMemoryMetrics
 
 /**
  * Records the peak values for executor level metrics. If jvmUsedHeapMemory is -1, then no
@@ -49,20 +48,6 @@ private[spark] class PeakExecutorMetrics {
     updated
   }
 
-  /**
-   * @return None if no peak metrics have been recorded, else PeakMemoryMetrics with the peak
-   *         values set.
-   */
-  def getPeakMemoryMetrics: Option[PeakMemoryMetrics] = {
-    if (metrics(0) < 0) {
-      None
-    } else {
-      val copy = new PeakMemoryMetrics
-      System.arraycopy(this.metrics, 0, copy.metrics, 0, this.metrics.length)
-      Some(copy)
-    }
-  }
-
   /** Clears/resets the saved peak values. */
   def reset(): Unit = {
     (0 until metrics.length).foreach { idx => metrics(idx) = 0}

core/src/main/scala/org/apache/spark/status/LiveEntity.scala

@@ -306,7 +306,7 @@ private class LiveExecutor(val executorId: String, _addTime: Long) extends LiveE
       executorLogs,
       memoryMetrics,
       blacklistedInStages,
-      peakExecutorMetrics.getPeakMemoryMetrics)
+      if (peakExecutorMetrics.metrics(0) == -1) None else Some(peakExecutorMetrics.metrics))
     new ExecutorSummaryWrapper(info)
   }
 }