change port to option and some bug fixes

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

@@ -87,7 +87,9 @@ private[spark] class Executor(
 
   private val executorMetrics: ExecutorMetrics = new ExecutorMetrics
   executorMetrics.setHostname(Utils.localHostName)
-  executorMetrics.setPort(env.rpcEnv.address.port)
+  if (env.rpcEnv.address != null) {
+    executorMetrics.setPort(Some(env.rpcEnv.address.port))
+  }
 
   // Whether to load classes in user jars before those in Spark jars
   private val userClassPathFirst = conf.getBoolean("spark.executor.userClassPathFirst", false)

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

@@ -39,11 +39,17 @@ class ExecutorMetrics extends Serializable {
   /**
    * Host's port the executor runs on
    */
-  private var _port: Int = _
-  def port: Int = _port
-  private[spark] def setPort(value: Int) = _port = value
+  private var _port: Option[Int] = None
+  def port: Option[Int] = _port
+  private[spark] def setPort(value: Option[Int]) = _port = value
 
-  private[spark] def hostPort: String = hostname + ":" + port
+  private[spark] def hostPort: String = {
+    val hp = port match {
+      case None => hostname
+      case value => hostname + ":" + value
+    }
+    hp
+  }
 
   private var _transportMetrics: TransportMetrics = new TransportMetrics
   def transportMetrics: TransportMetrics = _transportMetrics
@@ -61,7 +67,7 @@ class ExecutorMetrics extends Serializable {
 object ExecutorMetrics extends Serializable {
   def apply(
       hostName: String,
-      port: Int,
+      port: Option[Int],
       transportMetrics: TransportMetrics): ExecutorMetrics = {
     val execMetrics = new ExecutorMetrics
     execMetrics.setHostname(hostName)

core/src/main/scala/org/apache/spark/ui/memory/MemoryTab.scala

@@ -40,27 +40,32 @@ private[ui] class MemoryTab(parent: SparkUI) extends SparkUITab(parent, "memory"
 class MemoryListener extends SparkListener {
   type ExecutorId = String
   val activeExecutorIdToMem = new HashMap[ExecutorId, MemoryUIInfo]
+  // TODO There might be plenty of removed executors (e.g. Dynamic Allocation Mode), This may use
+  // too much memory.
   val removedExecutorIdToMem = new HashMap[ExecutorId, MemoryUIInfo]
-  // latestExecIdToExecMetrics including all executors that is active and removed.
-  // this may consume a lot of memory when executors are changing frequently, e.g. in dynamical
-  // allocation mode.
+  // A map that maintains the latest metrics of each active executor
   val latestExecIdToExecMetrics = new HashMap[ExecutorId, ExecutorMetrics]
   // activeStagesToMem a map maintains all executors memory information of each stage,
   // the Map type is [(stageId, attemptId), Seq[(executorId, MemoryUIInfo)]
   val activeStagesToMem = new HashMap[(Int, Int), HashMap[ExecutorId, MemoryUIInfo]]
+  // TODO We need to get conf of the retained stages so that we don't need to handle all the
+  // stages since there might be too many completed stages.
   val completedStagesToMem = new HashMap[(Int, Int), HashMap[ExecutorId, MemoryUIInfo]]
 
   override def onExecutorMetricsUpdate(event: SparkListenerExecutorMetricsUpdate): Unit = {
     val executorId = event.execId
     val executorMetrics = event.executorMetrics
     val memoryInfo = activeExecutorIdToMem.getOrElseUpdate(executorId, new MemoryUIInfo)
-    memoryInfo.updateExecutorMetrics(executorMetrics)
+    memoryInfo.updateMemUiInfo(executorMetrics)
     activeStagesToMem.foreach { case (_, stageMemMetrics) =>
-      if (stageMemMetrics.contains(executorId)) {
-        stageMemMetrics.get(executorId).get.updateExecutorMetrics(executorMetrics)
+      // If executor is added in the stage running time, we also update the metrics for the
+      // executor in {{activeStagesToMem}}
+      if (!stageMemMetrics.contains(executorId)) {
+        stageMemMetrics(executorId) = new MemoryUIInfo
       }
+      stageMemMetrics(executorId).updateMemUiInfo(executorMetrics)
     }
-    latestExecIdToExecMetrics.update(executorId, executorMetrics)
+    latestExecIdToExecMetrics(executorId) = executorMetrics
   }
 
   override def onExecutorAdded(event: SparkListenerExecutorAdded): Unit = {
@@ -71,28 +76,39 @@ class MemoryListener extends SparkListener {
   override def onExecutorRemoved(event: SparkListenerExecutorRemoved): Unit = {
     val executorId = event.executorId
     val info = activeExecutorIdToMem.remove(executorId)
+    latestExecIdToExecMetrics.remove(executorId)
     removedExecutorIdToMem.getOrElseUpdate(executorId, info.getOrElse(new MemoryUIInfo))
   }
 
   override def onStageSubmitted(event: SparkListenerStageSubmitted): Unit = {
     val stage = (event.stageInfo.stageId, event.stageInfo.attemptId)
     val memInfoMap = new HashMap[ExecutorId, MemoryUIInfo]
-    activeExecutorIdToMem.foreach(idToMem => memInfoMap.update(idToMem._1, new MemoryUIInfo))
-    activeStagesToMem.update(stage, memInfoMap)
+    activeExecutorIdToMem.map { case (id, _) =>
+      memInfoMap(id) = new MemoryUIInfo
+      val latestExecMetrics = latestExecIdToExecMetrics.get(id)
+      latestExecMetrics match {
+        case None => // Do nothing
+        case Some(metrics) =>
+          memInfoMap(id).updateMemUiInfo(metrics)
+      }
+    }
+    activeStagesToMem(stage) = memInfoMap
   }
 
   override def onStageCompleted(event: SparkListenerStageCompleted): Unit = {
     val stage = (event.stageInfo.stageId, event.stageInfo.attemptId)
+    // We need to refresh {{activeStagesToMem}} with {{activeExecutorIdToMem}} in case the
+    // executor is added in the stage running time and no {{SparkListenerExecutorMetricsUpdate}}
+    // event is updated in this stage.
     activeStagesToMem.get(stage).map { memInfoMap =>
-      activeExecutorIdToMem.foreach { case (executorId, _) =>
-        val memInfo = memInfoMap.getOrElse(executorId, new MemoryUIInfo)
-        latestExecIdToExecMetrics.get(executorId).foreach { prevExecutorMetrics =>
-          memInfo.updateExecutorMetrics(prevExecutorMetrics)
+      activeExecutorIdToMem.foreach { case (executorId, memUiInfo) =>
+        if (!memInfoMap.contains(executorId)) {
+          memInfoMap(executorId) = new MemoryUIInfo
+          memInfoMap(executorId).copyMemUiInfo(memUiInfo)
         }
-        memInfoMap.update(executorId, memInfo)
       }
-      completedStagesToMem.put(stage, activeStagesToMem.remove(stage).get)
     }
+    completedStagesToMem.put(stage, activeStagesToMem.remove(stage).get)
   }
 }
 
@@ -106,13 +122,18 @@ class MemoryUIInfo {
     executorAddress = execInfo.executorHost
   }
 
-  def updateExecutorMetrics(execMetrics: ExecutorMetrics): Unit = {
+  def updateMemUiInfo(execMetrics: ExecutorMetrics): Unit = {
     transportInfo = transportInfo match {
       case Some(transportMemSize) => transportInfo
       case _ => Some(new TransportMemSize)
     }
     executorAddress = execMetrics.hostPort
-    transportInfo.get.updateTransport(execMetrics.transportMetrics)
+    transportInfo.get.updateTransMemSize(execMetrics.transportMetrics)
+  }
+
+  def copyMemUiInfo(memUiInfo: MemoryUIInfo): Unit = {
+    executorAddress = memUiInfo.executorAddress
+    transportInfo.foreach(_.copyTransMemSize(memUiInfo.transportInfo.get))
   }
 }
 
@@ -123,7 +144,7 @@ class TransportMemSize {
   var peakOnHeapSizeTime: MemTime = new MemTime()
   var peakOffHeapSizeTime: MemTime = new MemTime()
 
-  def updateTransport(transportMetrics: TransportMetrics): Unit = {
+  def updateTransMemSize(transportMetrics: TransportMetrics): Unit = {
     val updatedOnHeapSize = transportMetrics.onHeapSize
     val updatedOffHeapSize = transportMetrics.offHeapSize
     val updateTime: Long = transportMetrics.timeStamp
@@ -136,6 +157,15 @@ class TransportMemSize {
       peakOffHeapSizeTime = MemTime(updatedOffHeapSize, updateTime)
     }
   }
+
+  def copyTransMemSize(transMemSize: TransportMemSize): Unit = {
+    onHeapSize = transMemSize.onHeapSize
+    offHeapSize = transMemSize.offHeapSize
+    peakOnHeapSizeTime = MemTime(transMemSize.peakOnHeapSizeTime.memorySize,
+      transMemSize.peakOnHeapSizeTime.timeStamp)
+    peakOffHeapSizeTime = MemTime(transMemSize.peakOffHeapSizeTime.memorySize,
+      transMemSize.peakOffHeapSizeTime.timeStamp)
+  }
 }
 
 @DeveloperApi

core/src/main/scala/org/apache/spark/ui/memory/MemoryTable.scala

@@ -35,9 +35,9 @@ private[ui] class MemTableBase(
     <th>Executor ID</th>
     <th>Address</th>
     <th>Network Memory (on-heap)</th>
-    <th>Network Memory (direct-heap)</th>
+    <th>Network Memory (off-heap)</th>
     <th>Peak Network Memory (on-heap) / Happen Time</th>
-    <th>Peak Network Read (direct-heap) / Happen Time</th>
+    <th>Peak Network Read (off-heap) / Happen Time</th>
   }
 
   def toNodeSeq: Seq[Node] = {

core/src/main/scala/org/apache/spark/ui/memory/StageMemoryPage.scala

@@ -39,7 +39,7 @@ private[ui] class StageMemoryPage(parent: MemoryTab) extends WebUIPage("stage")
 
       val finishedStageToMem = memoryListener.completedStagesToMem
       val content = if (finishedStageToMem.get(stage).isDefined) {
-        val executorIdToMem = finishedStageToMem.get(stage).get.toSeq.sortBy(_._1)
+        val executorIdToMem = finishedStageToMem(stage).toSeq.sortBy(_._1)
         val execMemTable = new MemTableBase(executorIdToMem, memoryListener)
         <h4 id="activeExec">Executors ({executorIdToMem.size})</h4> ++
           execMemTable.toNodeSeq

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

@@ -24,7 +24,7 @@ import scala.collection.Map
 
 import com.fasterxml.jackson.databind.ObjectMapper
 import com.fasterxml.jackson.module.scala.DefaultScalaModule
-import org.json4s.DefaultFormats
+import org.json4s.{JsonAST, DefaultFormats}
 import org.json4s.JsonDSL._
 import org.json4s.JsonAST._
 import org.json4s.jackson.JsonMethods._
@@ -296,7 +296,7 @@ private[spark] object JsonProtocol {
   def executorMetricsToJson(executorMetrics: ExecutorMetrics): JValue = {
     val transportMetrics = transportMetricsToJson(executorMetrics.transportMetrics)
     ("Executor Hostname" -> executorMetrics.hostname) ~
-    ("Executor Port" -> executorMetrics.port) ~
+    ("Executor Port" -> executorMetrics.port.map(new JInt(_)).getOrElse(JNothing)) ~
     ("TransportMetrics" -> transportMetrics)
   }
 
@@ -732,7 +732,7 @@ private[spark] object JsonProtocol {
       return metrics
     }
     metrics.setHostname((json \ "Executor Hostname").extract[String])
-    metrics.setPort((json \ "Executor Port").extract[Int])
+    metrics.setPort(Utils.jsonOption(json \ "Executor Port").map(_.extract[Int]))
     metrics.setTransportMetrics(transportMetricsFromJson(json \ "TransportMetrics"))
     metrics
   }

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

@@ -129,7 +129,7 @@ class EventLoggingListenerSuite extends SparkFunSuite with LocalSparkContext wit
     val eventLogger = new EventLoggingListener("test-memListener", None, testDirPath.toUri(), conf)
     val execId = "exec-1"
     val hostName = "host-1"
-    val port = 80
+    val port: Option[Int] = Some(80)
 
     eventLogger.start()
     eventLogger.onExecutorAdded(SparkListenerExecutorAdded(

core/src/test/scala/org/apache/spark/ui/memory/MemoryListenerSuite.scala

@@ -28,6 +28,7 @@ class MemoryListenerSuite extends SparkFunSuite {
     val listener = new MemoryListener
     val execId1 = "exec-1"
     val host1 = "host-1"
+    val port: Option[Int] = Some(80)
 
     listener.onExecutorAdded(
       SparkListenerExecutorAdded(1L, execId1, new ExecutorInfo(host1, 1, Map.empty)))
@@ -38,13 +39,13 @@ class MemoryListenerSuite extends SparkFunSuite {
 
     // multiple metrics updated in stage 2
     listener.onStageSubmitted(MemoryListenerSuite.createStageStartEvent(2))
-    val execMetrics1 = MemoryListenerSuite.createExecutorMetrics(host1, 80, 2L, 20, 10)
+    val execMetrics1 = MemoryListenerSuite.createExecutorMetrics(host1, port, 2L, 20, 10)
     listener.onExecutorMetricsUpdate(MemoryListenerSuite.createExecutorMetricsUpdateEvent(
       execId1, execMetrics1))
-    val execMetrics2 = MemoryListenerSuite.createExecutorMetrics(host1, 80, 3L, 30, 5)
+    val execMetrics2 = MemoryListenerSuite.createExecutorMetrics(host1, port, 3L, 30, 5)
     listener.onExecutorMetricsUpdate(MemoryListenerSuite.createExecutorMetricsUpdateEvent(
       execId1, execMetrics2))
-    val execMetrics3 = MemoryListenerSuite.createExecutorMetrics(host1, 80, 4L, 15, 15)
+    val execMetrics3 = MemoryListenerSuite.createExecutorMetrics(host1, port, 4L, 15, 15)
     listener.onExecutorMetricsUpdate(MemoryListenerSuite.createExecutorMetricsUpdateEvent(
       execId1, execMetrics3))
     listener.onStageCompleted(MemoryListenerSuite.createStageEndEvent(2))
@@ -87,7 +88,8 @@ class MemoryListenerSuite extends SparkFunSuite {
     val listener = new MemoryListener
     val (execId1, execId2, execId3) = ("exec-1", "exec-2", "exec-3")
     val (host1, host2, host3) = ("host-1", "host-2", "host-3")
-    val (port1, port2, port3) = (80, 80, 80)
+    val (port1, port2, port3): (Option[Int], Option[Int], Option[Int]) =
+      (Some(80), Some(80), Some(80))
 
     // two executors added first
     listener.onExecutorAdded(
@@ -135,7 +137,6 @@ class MemoryListenerSuite extends SparkFunSuite {
 
     assert(listener.activeStagesToMem.isEmpty)
     assert(listener.completedStagesToMem.size === 4)
-    assert(listener.activeExecutorIdToMem.size === listener.latestExecIdToExecMetrics.size)
     assert(listener.removedExecutorIdToMem.size === 1)
 
     listener.onExecutorRemoved(SparkListenerExecutorRemoved(7L, execId1, ""))
@@ -180,7 +181,7 @@ object MemoryListenerSuite extends SparkFunSuite {
 
   def createExecutorMetrics(
       hostname: String,
-      port: Int,
+      port: Option[Int],
       timeStamp: Long,
       onHeapSize: Long,
       offHeapSize: Long): ExecutorMetrics = {

core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala

@@ -86,7 +86,7 @@ class JsonProtocolSuite extends SparkFunSuite {
     val executorMetrics = {
       val execMetrics = new ExecutorMetrics
       execMetrics.setHostname("host-1")
-      execMetrics.setPort(80)
+      execMetrics.setPort(Some(80))
       execMetrics.setTransportMetrics(TransportMetrics(0L, 10, 10))
       execMetrics
     }
@@ -398,7 +398,7 @@ class JsonProtocolSuite extends SparkFunSuite {
     val oldJson = newJson.removeField { case (field, _) => field == "Executor Metrics Updated"}
     val newMetrics = JsonProtocol.executorMetricsUpdateFromJson(oldJson)
     assert(newMetrics.executorMetrics.hostname === "")
-    assert(newMetrics.executorMetrics.port === 0)
+    assert(newMetrics.executorMetrics.port === None)
     assert(newMetrics.executorMetrics.transportMetrics.onHeapSize === 0L)
     assert(newMetrics.executorMetrics.transportMetrics.offHeapSize === 0L)
     assert(newMetrics.executorMetrics.transportMetrics.timeStamp != 0L)