Improve SparkStatusTracker to also track executor information

Author: cloud-fan

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

@@ -147,8 +147,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
       appName: String,
       sparkHome: String = null,
       jars: Seq[String] = Nil,
-      environment: Map[String, String] = Map()) =
-  {
+      environment: Map[String, String] = Map()) = {
     this(SparkContext.updatedConf(new SparkConf(), master, appName, sparkHome, jars, environment))
   }
 

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

@@ -17,6 +17,8 @@
 
 package org.apache.spark
 
+import org.apache.spark.scheduler.TaskSchedulerImpl
+
 /**
  * Low-level status reporting APIs for monitoring job and stage progress.
  *
@@ -104,4 +106,25 @@ class SparkStatusTracker private[spark] (sc: SparkContext) {
       }
     }
   }
+
+  def getExecutors(): Array[String] = {
+    sc.getExecutorStorageStatus.map { status =>
+      status.blockManagerId.hostPort
+    }
+  }
+
+  def getCacheSizeByExecutors(): Map[String, Long] = {
+    sc.getExecutorStorageStatus.map { status =>
+      status.blockManagerId.hostPort -> status.memUsed
+    }.toMap
+  }
+
+  def getRunningTasksByExecutors(): Map[String, Int] = {
+    val executorIdToRunningTasks: Map[String, Int] =
+      sc.taskScheduler.asInstanceOf[TaskSchedulerImpl].runningTasksByExecutors()
+    sc.getExecutorStorageStatus.map { status =>
+      val bmId = status.blockManagerId
+      bmId.hostPort -> executorIdToRunningTasks.getOrElse(bmId.executorId, 0)
+    }.toMap
+  }
 }

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

@@ -90,6 +90,8 @@ private[spark] class TaskSchedulerImpl(
   // Number of tasks running on each executor
   private val executorIdToTaskCount = new HashMap[String, Int]
 
+  def runningTasksByExecutors(): Map[String, Int] = executorIdToTaskCount.toMap
+
   // The set of executors we have on each host; this is used to compute hostsAlive, which
   // in turn is used to decide when we can attain data locality on a given host
   protected val executorsByHost = new HashMap[String, HashSet[String]]