Changes for spark streaming UI

Author: mubarak

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

@@ -1296,6 +1296,9 @@ object SparkContext extends Logging {
 
   private[spark] val SPARK_UNKNOWN_USER = "<unknown>"
 
+  private[spark] val SPARK_JOB_CALL_SITE_SHORT = "spark.job.callSiteShort"
+  private[spark] val SPARK_JOB_CALL_SITE_LONG = "spark.job.callSiteLong"
+
   implicit object DoubleAccumulatorParam extends AccumulatorParam[Double] {
     def addInPlace(t1: Double, t2: Double): Double = t1 + t2
     def zero(initialValue: Double) = 0.0

core/src/main/scala/org/apache/spark/rdd/RDD.scala

@@ -17,7 +17,7 @@
 
 package org.apache.spark.rdd
 
-import java.util.Random
+import java.util.{Properties, Random}
 
 import scala.collection.{mutable, Map}
 import scala.collection.mutable.ArrayBuffer
@@ -1213,7 +1213,17 @@ abstract class RDD[T: ClassTag](
   private var storageLevel: StorageLevel = StorageLevel.NONE
 
   /** User code that created this RDD (e.g. `textFile`, `parallelize`). */
-  @transient private[spark] val creationSite = Utils.getCallSite
+  @transient private[spark] val creationSite = {
+    val short: String = sc.getLocalProperty("spark.job.callSiteShort")
+    if (short != null) {
+      CallSite(short, sc.getLocalProperty("spark.job.callSiteLong"))
+    } else {
+      val callSite: CallSite = Utils.getCallSite
+      //sc.setLocalProperty("spark.job.callSiteShort", callSite.short)
+      //sc.setLocalProperty("spark.job.callSiteLong", callSite.long)
+      callSite
+    }
+  }
   private[spark] def getCreationSite: String = Option(creationSite).map(_.short).getOrElse("")
 
   private[spark] def elementClassTag: ClassTag[T] = classTag[T]

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

@@ -800,7 +800,10 @@ private[spark] object Utils extends Logging {
    * A regular expression to match classes of the "core" Spark API that we want to skip when
    * finding the call site of a method.
    */
-  private val SPARK_CLASS_REGEX = """^org\.apache\.spark(\.api\.java)?(\.util)?(\.rdd)?\.[A-Z]""".r
+  private val SPARK_CLASS_REGEX = """^org\.apache\.spark(\.api\.java)?(\.util)?(\.rdd)?(\.streaming)?(\.streaming\.dstream)?(\.streaming\.scheduler)?\.[A-Z]""".r
+  private val SCALA_CLASS_REGEX = """^scala(\.util)?(\.collection)?(\.collection\.mutable)?(\.collection\.immutable)?(\.concurrent\.forkjoin)?\.[A-Z]""".r
+  private val AKKA_CLASS_REGEX = """^akka(\.actor)?(\.dispatch)?\.[A-Z]""".r
+  private val JAVA_CLASS_REGEX = """^java(\.util\.concurrent)?(\.lang)?\.[A-Z]""".r
 
   /**
    * When called inside a class in the spark package, returns the name of the user code class
@@ -828,7 +831,10 @@ private[spark] object Utils extends Logging {
 
     for (el <- trace) {
       if (insideSpark) {
-        if (SPARK_CLASS_REGEX.findFirstIn(el.getClassName).isDefined) {
+        if (SPARK_CLASS_REGEX.findFirstIn(el.getClassName).isDefined ||
+            SCALA_CLASS_REGEX.findFirstIn(el.getClassName).isDefined ||
+            AKKA_CLASS_REGEX.findFirstIn(el.getClassName).isDefined ||
+            JAVA_CLASS_REGEX.findFirstIn(el.getClassName).isDefined) {
           lastSparkMethod = if (el.getMethodName == "<init>") {
             // Spark method is a constructor; get its class name
             el.getClassName.substring(el.getClassName.lastIndexOf('.') + 1)
@@ -846,7 +852,7 @@ private[spark] object Utils extends Logging {
         callStack += el.toString
       }
     }
-    val callStackDepth = System.getProperty("spark.callstack.depth", "20").toInt
+    val callStackDepth = System.getProperty("spark.callstack.depth", "10").toInt
     CallSite(
       short = "%s at %s:%s".format(lastSparkMethod, firstUserFile, firstUserLine),
       long = callStack.take(callStackDepth).mkString("\n"))

streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala

@@ -38,7 +38,7 @@ import org.apache.spark.streaming.dstream._
 import org.apache.spark.streaming.receiver.{ActorSupervisorStrategy, ActorReceiver, Receiver}
 import org.apache.spark.streaming.scheduler._
 import org.apache.spark.streaming.ui.StreamingTab
-import org.apache.spark.util.MetadataCleaner
+import org.apache.spark.util.{Utils, MetadataCleaner}
 
 /**
  * Main entry point for Spark Streaming functionality. It provides methods used to create
@@ -112,6 +112,7 @@ class StreamingContext private[streaming] (
     if (isCheckpointPresent) {
       new SparkContext(cp_.sparkConf)
     } else {
+      sc_.setCallSite(Utils.getCallSite.short)
       sc_
     }
   }

streaming/src/main/scala/org/apache/spark/streaming/dstream/DStream.scala

@@ -30,7 +30,7 @@ import org.apache.spark.storage.StorageLevel
 import org.apache.spark.streaming._
 import org.apache.spark.streaming.StreamingContext._
 import org.apache.spark.streaming.scheduler.Job
-import org.apache.spark.util.MetadataCleaner
+import org.apache.spark.util.{Utils, MetadataCleaner}
 
 /**
  * A Discretized Stream (DStream), the basic abstraction in Spark Streaming, is a continuous
@@ -322,6 +322,7 @@ abstract class DStream[T: ClassTag] (
   private[streaming] def generateJob(time: Time): Option[Job] = {
     getOrCompute(time) match {
       case Some(rdd) => {
+        //ssc.sc.setJobGroup("g","d")
         val jobFunc = () => {
           val emptyFunc = { (iterator: Iterator[T]) => {} }
           context.sparkContext.runJob(rdd, emptyFunc)

streaming/src/main/scala/org/apache/spark/streaming/dstream/ForEachDStream.scala

@@ -34,9 +34,12 @@ class ForEachDStream[T: ClassTag] (
 
   override def compute(validTime: Time): Option[RDD[Unit]] = None
 
+  //TODO: where to clear up the threadlocal values?
   override def generateJob(time: Time): Option[Job] = {
     parent.getOrCompute(time) match {
       case Some(rdd) =>
+        parent.ssc.sc.setLocalProperty("spark.job.callSiteShort", rdd.creationSite.short)
+        parent.ssc.sc.setLocalProperty("spark.job.callSiteLong", rdd.creationSite.long)
         val jobFunc = () => {
           foreachFunc(rdd, time)
         }

streaming/src/main/scala/org/apache/spark/streaming/scheduler/JobGenerator.scala

@@ -22,10 +22,11 @@ import org.apache.spark.{SparkException, SparkEnv, Logging}
 import org.apache.spark.streaming.{Checkpoint, Time, CheckpointWriter}
 import org.apache.spark.streaming.util.{ManualClock, RecurringTimer, Clock}
 import scala.util.{Failure, Success, Try}
+import org.apache.spark.util.{CallSite, Utils}
 
 /** Event classes for JobGenerator */
 private[scheduler] sealed trait JobGeneratorEvent
-private[scheduler] case class GenerateJobs(time: Time) extends JobGeneratorEvent
+private[scheduler] case class GenerateJobs(time: Time, callSite: CallSite) extends JobGeneratorEvent
 private[scheduler] case class ClearMetadata(time: Time) extends JobGeneratorEvent
 private[scheduler] case class DoCheckpoint(time: Time) extends JobGeneratorEvent
 private[scheduler] case class ClearCheckpointData(time: Time) extends JobGeneratorEvent
@@ -47,8 +48,10 @@ class JobGenerator(jobScheduler: JobScheduler) extends Logging {
     Class.forName(clockClass).newInstance().asInstanceOf[Clock]
   }
 
+  val callSite: CallSite = Utils.getCallSite
+
   private val timer = new RecurringTimer(clock, ssc.graph.batchDuration.milliseconds,
-    longTime => eventActor ! GenerateJobs(new Time(longTime)), "JobGenerator")
+    longTime => eventActor ! GenerateJobs(new Time(longTime), callSite), "JobGenerator")
 
   // This is marked lazy so that this is initialized after checkpoint duration has been set
   // in the context and the generator has been started.
@@ -162,7 +165,7 @@ class JobGenerator(jobScheduler: JobScheduler) extends Logging {
   private def processEvent(event: JobGeneratorEvent) {
     logDebug("Got event " + event)
     event match {
-      case GenerateJobs(time) => generateJobs(time)
+      case GenerateJobs(time, callSite) => generateJobs(time, callSite)
       case ClearMetadata(time) => clearMetadata(time)
       case DoCheckpoint(time) => doCheckpoint(time)
       case ClearCheckpointData(time) => clearCheckpointData(time)
@@ -216,7 +219,9 @@ class JobGenerator(jobScheduler: JobScheduler) extends Logging {
   }
 
   /** Generate jobs and perform checkpoint for the given `time`.  */
-  private def generateJobs(time: Time) {
+  private def generateJobs(time: Time, callSite: CallSite) {
+    ssc.sc.setLocalProperty("spark.job.callSiteShort", callSite.short)
+    ssc.sc.setLocalProperty("spark.job.callSiteLong", callSite.long)
     SparkEnv.set(ssc.env)
     Try(graph.generateJobs(time)) match {
       case Success(jobs) =>