[SPARK-24958][WIP] Report executors' process tree total memory information to heartbeat signals

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

@@ -23,7 +23,7 @@ import java.nio.charset.StandardCharsets
 import java.util.EnumSet
 import java.util.Locale
 
-import scala.collection.mutable.{ArrayBuffer, HashMap}
+import scala.collection.mutable.{ArrayBuffer, Map}
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FileSystem, FSDataOutputStream, Path}
@@ -96,7 +96,7 @@ private[spark] class EventLoggingListener(
   private[scheduler] val logPath = getLogPath(logBaseDir, appId, appAttemptId, compressionCodecName)
 
   // map of (stageId, stageAttempt), to peak executor metrics for the stage
-  private val liveStageExecutorMetrics = HashMap[(Int, Int), HashMap[String, ExecutorMetrics]]()
+  private val liveStageExecutorMetrics = Map.empty[(Int, Int), Map[String, ExecutorMetrics]]
 
   /**
    * Creates the log file in the configured log directory.
@@ -165,7 +165,7 @@ private[spark] class EventLoggingListener(
     if (shouldLogStageExecutorMetrics) {
       // record the peak metrics for the new stage
       liveStageExecutorMetrics.put((event.stageInfo.stageId, event.stageInfo.attemptNumber()),
-        new HashMap[String, ExecutorMetrics]())
+        Map.empty[String, ExecutorMetrics])
     }
   }
 
@@ -338,7 +338,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 HashMap[String, CompressionCodec]
+  private val codecMap = Map.empty[String, CompressionCodec]
 
   /**
    * Write metadata about an event log to the given stream.

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

@@ -706,13 +706,10 @@ private[spark] class AppStatusListener(
     // while reading from the log. SparkListenerStageExecutorMetrics are only processed
     // when reading logs.
     liveExecutors.get(executorMetrics.execId)
-      .orElse(deadExecutors.get(executorMetrics.execId)) match {
-      case Some(exec) =>
-         if (exec.peakExecutorMetrics.compareAndUpdatePeakValues(executorMetrics.executorMetrics)) {
-          maybeUpdate(exec, now)
-        }
-      case None =>
-        logWarning("unable to find executor " + executorMetrics.execId)
+      .orElse(deadExecutors.get(executorMetrics.execId)).map { exec =>
+      if (exec.peakExecutorMetrics.compareAndUpdatePeakValues(executorMetrics.executorMetrics)) {
+        update(exec, now)
+      }
     }
   }
 

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,
-      if (peakExecutorMetrics.isSet()) Some(peakExecutorMetrics) else None)
+      Some(peakExecutorMetrics).filter(_.isSet))
     new ExecutorSummaryWrapper(info)
   }
 }

core/src/main/scala/org/apache/spark/status/api/v1/api.scala

@@ -116,33 +116,27 @@ class MemoryMetrics private[spark](
 
 /** deserializer for peakMemoryMetrics: convert map to ExecutorMetrics */
 private[spark] class ExecutorMetricsJsonDeserializer
-  extends JsonDeserializer[Option[ExecutorMetrics]] {
+    extends JsonDeserializer[Option[ExecutorMetrics]] {
   override def deserialize(
       jsonParser: JsonParser,
       deserializationContext: DeserializationContext): Option[ExecutorMetrics] = {
     val metricsMap = jsonParser.readValueAs[Option[Map[String, Long]]](
       new TypeReference[Option[Map[String, java.lang.Long]]] {})
-    metricsMap match {
-      case Some(metrics) =>
-        Some(new ExecutorMetrics(metrics))
-      case None => None
-    }
+    metricsMap.map(metrics => new ExecutorMetrics(metrics))
   }
 }
 /** serializer for peakMemoryMetrics: convert ExecutorMetrics to map with metric name as key */
 private[spark] class ExecutorMetricsJsonSerializer
-  extends JsonSerializer[Option[ExecutorMetrics]] {
+    extends JsonSerializer[Option[ExecutorMetrics]] {
   override def serialize(
       metrics: Option[ExecutorMetrics],
       jsonGenerator: JsonGenerator,
       serializerProvider: SerializerProvider): Unit = {
-    metrics match {
-      case Some(m) =>
-        val metricsMap = ExecutorMetricType.values.map { metricType =>
+    metrics.foreach { m: ExecutorMetrics =>
+      val metricsMap = ExecutorMetricType.values.map { metricType =>
             metricType.name -> m.getMetricValue(metricType)
-        }.toMap
-        jsonGenerator.writeObject(metricsMap)
-      case None =>
+      }.toMap
+      jsonGenerator.writeObject(metricsMap)
     }
   }
 }

core/src/main/scala/org/apache/spark/util/JsonProtocol.scala

@@ -723,9 +723,8 @@ private[spark] object JsonProtocol {
         (json \ "Accumulator Updates").extract[List[JValue]].map(accumulableInfoFromJson)
       (taskId, stageId, stageAttemptId, updates)
     }
-    val executorUpdates = jsonOption(json \ "Executor Metrics Updated") match {
-      case None => None
-      case Some(executorUpdate) => Some(executorMetricsFromJson(executorUpdate))
+    val executorUpdates = jsonOption(json \ "Executor Metrics Updated").map {
+      executorUpdate => executorMetricsFromJson(executorUpdate)
     }
     SparkListenerExecutorMetricsUpdate(execInfo, accumUpdates, executorUpdates)
   }

core/src/test/scala/org/apache/spark/scheduler/EventLoggingListenerSuite.scala

@@ -272,26 +272,6 @@ class EventLoggingListenerSuite extends SparkFunSuite with LocalSparkContext wit
     val eventLogger = new EventLoggingListener(logName, None, testDirPath.toUri(), conf)
     val listenerBus = new LiveListenerBus(conf)
 
-    // expected StageExecutorMetrics, for the given stage id and executor id
-    val expectedMetricsEvents: Map[(Int, String), SparkListenerStageExecutorMetrics] =
-      Map(
-        ((0, "1"),
-          new SparkListenerStageExecutorMetrics("1", 0, 0,
-            new ExecutorMetrics(Array(5000L, 50L, 4000L, 8000L, 50L, 20L, 50L, 10L, 100L, 30L,
-              70L, 20L)))),
-        ((0, "2"),
-          new SparkListenerStageExecutorMetrics("2", 0, 0,
-            new ExecutorMetrics(Array(7000L, 70L, 4000L, 9000L, 50L, 20L, 10L, 10L, 50L, 30L,
-              80L, 40L)))),
-        ((1, "1"),
-          new SparkListenerStageExecutorMetrics("1", 1, 0,
-            new ExecutorMetrics(Array(7000L, 70L, 3000L, 5000L, 50L, 30L, 60L, 30L, 80L, 55L,
-              50L, 0L)))),
-        ((1, "2"),
-          new SparkListenerStageExecutorMetrics("2", 1, 0,
-            new ExecutorMetrics(Array(7000L, 70L, 5000L, 8000L, 50L, 40L, 10L, 30L, 50L, 60L,
-              40L, 40L)))))
-
     // Events to post.
     val events = Array(
       SparkListenerApplicationStart("executionMetrics", None,
@@ -365,6 +345,26 @@ class EventLoggingListenerSuite extends SparkFunSuite with LocalSparkContext wit
     listenerBus.stop()
     eventLogger.stop()
 
+    // expected StageExecutorMetrics, for the given stage id and executor id
+    val expectedMetricsEvents: Map[(Int, String), SparkListenerStageExecutorMetrics] =
+    Map(
+      ((0, "1"),
+        new SparkListenerStageExecutorMetrics("1", 0, 0,
+          new ExecutorMetrics(Array(5000L, 50L, 4000L, 8000L, 50L, 20L, 50L, 10L, 100L, 30L,
+            70L, 20L)))),
+      ((0, "2"),
+        new SparkListenerStageExecutorMetrics("2", 0, 0,
+          new ExecutorMetrics(Array(7000L, 70L, 4000L, 9000L, 50L, 20L, 10L, 10L, 50L, 30L,
+            80L, 40L)))),
+      ((1, "1"),
+        new SparkListenerStageExecutorMetrics("1", 1, 0,
+          new ExecutorMetrics(Array(7000L, 70L, 3000L, 5000L, 50L, 30L, 60L, 30L, 80L, 55L,
+            50L, 0L)))),
+      ((1, "2"),
+        new SparkListenerStageExecutorMetrics("2", 1, 0,
+          new ExecutorMetrics(Array(7000L, 70L, 5000L, 8000L, 50L, 40L, 10L, 30L, 50L, 60L,
+            40L, 40L)))))
+
     // Verify the log file contains the expected events.
     // Posted events should be logged, except for ExecutorMetricsUpdate events -- these
     // are consolidated, and the peak values for each stage are logged at stage end.

core/src/test/scala/org/apache/spark/scheduler/ReplayListenerSuite.scala

@@ -219,10 +219,7 @@ class ReplayListenerSuite extends SparkFunSuite with BeforeAndAfter with LocalSp
     // Verify the same events are replayed in the same order
     assert(sc.eventLogger.isDefined)
     val originalEvents = sc.eventLogger.get.loggedEvents.filter { e =>
-      JsonProtocol.sparkEventFromJson(e) match {
-        case event: SparkListenerStageExecutorMetrics => false
-        case _ => true
-      }
+      !JsonProtocol.sparkEventFromJson(e).isInstanceOf[SparkListenerStageExecutorMetrics]
     }
     val replayedEvents = eventMonster.loggedEvents
     originalEvents.zip(replayedEvents).foreach { case (e1, e2) =>

core/src/test/scala/org/apache/spark/status/AppStatusListenerSuite.scala

@@ -1293,7 +1293,7 @@ class AppStatusListenerSuite extends SparkFunSuite with BeforeAndAfter {
 
   test("stage executor metrics") {
     // simulate reading in StageExecutorMetrics events from the history log
-    val listener = new AppStatusListener(store, conf, true)
+    val listener = new AppStatusListener(store, conf, false)
     val driver = BlockManagerId(SparkContext.DRIVER_IDENTIFIER, "localhost", 42)
 
     listener.onExecutorAdded(createExecutorAddedEvent(1))