wip on enum based metrics

core/src/main/java/org/apache/spark/scheduler/MemoryTypes.java

@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.scheduler;
+
+import org.apache.spark.executor.ExecutorMetrics;
+
+public enum MemoryTypes {
+  JvmUsedMemory{
+    @Override
+    long get(ExecutorMetrics em) {
+      return em.jvmUsedMemory();
+    }
+  },
+  OnHeapExecutionMemory {
+    @Override
+    long get(ExecutorMetrics em) {
+      return em.onHeapExecutionMemory();
+    }
+  },
+  OffHeapExecutionMemory {
+    @Override
+    long get(ExecutorMetrics em) {
+      return em.offHeapExecutionMemory();
+    }
+  },
+  OnHeapStorageMemory {
+    @Override
+    long get(ExecutorMetrics em) {
+      return em.onHeapStorageMemory();
+    }
+  },
+  OffHeapStorageMemory {
+    @Override
+    long get(ExecutorMetrics em) {
+      return em.offHeapStorageMemory();
+    }
+  },
+  OnHeapUnifiedMemory {
+    @Override
+    long get(ExecutorMetrics em) {
+      return em.onHeapExecutionMemory() + em.onHeapStorageMemory();
+    }
+  },
+  OffHeapUnifiedMemory {
+    @Override
+    long get(ExecutorMetrics em) {
+      return em.offHeapExecutionMemory() + em.offHeapStorageMemory();
+    }
+  };
+
+  abstract long get(ExecutorMetrics em);
+}

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

@@ -25,36 +25,8 @@ import org.apache.spark.status.api.v1.PeakMemoryMetrics
  * values have been recorded yet.
  */
 private[spark] class PeakExecutorMetrics {
-  private var _jvmUsedHeapMemory = -1L;
-  private var _jvmUsedNonHeapMemory = 0L;
-  private var _onHeapExecutionMemory = 0L
-  private var _offHeapExecutionMemory = 0L
-  private var _onHeapStorageMemory = 0L
-  private var _offHeapStorageMemory = 0L
-  private var _onHeapUnifiedMemory = 0L
-  private var _offHeapUnifiedMemory = 0L
-  private var _directMemory = 0L
-  private var _mappedMemory = 0L
-
-  def jvmUsedHeapMemory: Long = _jvmUsedHeapMemory
-
-  def jvmUsedNonHeapMemory: Long = _jvmUsedNonHeapMemory
-
-  def onHeapExecutionMemory: Long = _onHeapExecutionMemory
-
-  def offHeapExecutionMemory: Long = _offHeapExecutionMemory
-
-  def onHeapStorageMemory: Long = _onHeapStorageMemory
-
-  def offHeapStorageMemory: Long = _offHeapStorageMemory
-
-  def onHeapUnifiedMemory: Long = _onHeapUnifiedMemory
-
-  def offHeapUnifiedMemory: Long = _offHeapUnifiedMemory
-
-  def directMemory: Long = _directMemory
-
-  def mappedMemory: Long = _mappedMemory
+  val metrics = new Array[Long](MemoryTypes.values().length)
+  metrics(0) = -1
 
   /**
    * Compare the specified memory values with the saved peak executor memory
@@ -66,47 +38,13 @@ private[spark] class PeakExecutorMetrics {
   def compareAndUpdate(executorMetrics: ExecutorMetrics): Boolean = {
     var updated: Boolean = false
 
-    if (executorMetrics.jvmUsedHeapMemory > _jvmUsedHeapMemory) {
-      _jvmUsedHeapMemory = executorMetrics.jvmUsedHeapMemory
-      updated = true
-    }
-    if (executorMetrics.jvmUsedNonHeapMemory > _jvmUsedNonHeapMemory) {
-      _jvmUsedNonHeapMemory = executorMetrics.jvmUsedNonHeapMemory
-      updated = true
+    (0 until MemoryTypes.values().length).foreach { metricIdx =>
+      val metricVal = MemoryTypes.values()(metricIdx).get(executorMetrics)
+      if (metricVal > metrics(metricIdx)) {
+        updated = true
+        metrics(metricIdx) = metricVal
+      }
     }
-    if (executorMetrics.onHeapExecutionMemory > _onHeapExecutionMemory) {
-      _onHeapExecutionMemory = executorMetrics.onHeapExecutionMemory
-      updated = true
-    }
-    if (executorMetrics.offHeapExecutionMemory > _offHeapExecutionMemory) {
-      _offHeapExecutionMemory = executorMetrics.offHeapExecutionMemory
-      updated = true
-    }
-    if (executorMetrics.onHeapStorageMemory > _onHeapStorageMemory) {
-      _onHeapStorageMemory = executorMetrics.onHeapStorageMemory
-      updated = true
-    }
-    if (executorMetrics.offHeapStorageMemory > _offHeapStorageMemory) {
-      _offHeapStorageMemory = executorMetrics.offHeapStorageMemory
-      updated = true
-    }
-    if (executorMetrics.onHeapUnifiedMemory > _onHeapUnifiedMemory) {
-      _onHeapUnifiedMemory = executorMetrics.onHeapUnifiedMemory
-      updated = true
-    }
-    if (executorMetrics.offHeapUnifiedMemory > _offHeapUnifiedMemory) {
-      _offHeapUnifiedMemory = executorMetrics.offHeapUnifiedMemory
-      updated = true
-    }
-    if (executorMetrics.directMemory > _directMemory) {
-      _directMemory = executorMetrics.directMemory
-      updated = true
-    }
-    if (executorMetrics.mappedMemory > _mappedMemory) {
-      _mappedMemory = executorMetrics.mappedMemory
-      updated = true
-    }
-
     updated
   }
 
@@ -115,13 +53,18 @@ private[spark] class PeakExecutorMetrics {
    *         values set.
    */
   def getPeakMemoryMetrics: Option[PeakMemoryMetrics] = {
-    if (_jvmUsedHeapMemory < 0) {
+    if (metrics(0) < 0) {
       None
     } else {
-      Some(new PeakMemoryMetrics(_jvmUsedHeapMemory, _jvmUsedNonHeapMemory,
-        _onHeapExecutionMemory, _offHeapExecutionMemory, _onHeapStorageMemory,
-        _offHeapStorageMemory, _onHeapUnifiedMemory, _offHeapUnifiedMemory,
-        _directMemory, _mappedMemory))
+      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}
+    metrics(0) = -1
+  }
 }

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

@@ -25,7 +25,7 @@ import com.fasterxml.jackson.annotation.JsonIgnoreProperties
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize
 
 import org.apache.spark.JobExecutionStatus
-import org.apache.spark.executor.ExecutorMetrics
+import org.apache.spark.scheduler.MemoryTypes
 
 case class ApplicationInfo private[spark](
     id: String,
@@ -108,17 +108,10 @@ class MemoryMetrics private[spark](
     val totalOnHeapStorageMemory: Long,
     val totalOffHeapStorageMemory: Long)
 
-class PeakMemoryMetrics private[spark](
-    val jvmUsedHeapMemory: Long,
-    val jvmUsedNonHeapMemory: Long,
-    val onHeapExecutionMemory: Long,
-    val offHeapExecutionMemory: Long,
-    val onHeapStorageMemory: Long,
-    val offHeapStorageMemory: Long,
-    val onHeapUnifiedMemory: Long,
-    val offHeapUnifiedMemory: Long,
-    val directMemory: Long,
-    val mappedMemory: Long)
+class PeakMemoryMetrics private[spark]() {
+  // TODO special json-ification
+  val metrics = new Array[Long](MemoryTypes.values().length)
+}
 
 class JobData private[spark](
     val jobId: Int,