Refactored getCallSite and setCallSite to make it simpler. Also added unit test for DStream creation site.

Author: tdas

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

@@ -1024,28 +1024,40 @@ class SparkContext(config: SparkConf) extends Logging {
   }
 
   /**
-   * Support function for API backtraces.
+   * Set the thread-local property for overriding the call sites
+   * of actions and RDDs.
    */
-  def setCallSite(site: String) {
-    setLocalProperty("externalCallSite", site)
+  def setCallSite(shortCallSite: String) {
+    setLocalProperty(CallSite.SHORT_FORM, shortCallSite)
   }
 
   /**
-   * Support function for API backtraces.
+   * Set the thread-local property for overriding the call sites
+   * of actions and RDDs.
+   */
+  private[spark] def setCallSite(callSite: CallSite) {
+    setLocalProperty(CallSite.SHORT_FORM, callSite.shortForm)
+    setLocalProperty(CallSite.LONG_FORM, callSite.longForm)
+  }
+
+  /**
+   * Clear the thread-local property for overriding the call sites
+   * of actions and RDDs.
    */
   def clearCallSite() {
-    setLocalProperty("externalCallSite", null)
+    setLocalProperty(CallSite.SHORT_FORM, null)
+    setLocalProperty(CallSite.LONG_FORM, null)
   }
 
   /**
    * Capture the current user callsite and return a formatted version for printing. If the user
-   * has overridden the call site, this will return the user's version.
+   * has overridden the call site using `setCallSite()`, this will return the user's version.
    */
   private[spark] def getCallSite(): CallSite = {
-    Option(getLocalProperty("externalCallSite")) match {
-      case Some(callSite) => CallSite(callSite, longForm = "")
-      case None => Utils.getCallSite()
-    }
+    Option(getLocalProperty(CallSite.SHORT_FORM)).map { case shortCallSite =>
+      val longCallSite = Option(getLocalProperty(CallSite.LONG_FORM)).getOrElse("")
+      CallSite(shortCallSite, longCallSite)
+    }.getOrElse(Utils.getCallSite())
   }
 
   /**

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

@@ -1220,14 +1220,7 @@ 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 = {
-    val short: String = sc.getLocalProperty(Utils.CALL_SITE_SHORT)
-    if (short != null) {
-      CallSite(short, sc.getLocalProperty(Utils.CALL_SITE_LONG))
-    } else {
-      Utils.getCallSite()
-    }
-  }
+  @transient private[spark] val creationSite = sc.getCallSite()
 
   private[spark] def getCreationSite: String = Option(creationSite).map(_.shortForm).getOrElse("")
 

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

@@ -47,15 +47,17 @@ import org.apache.spark.serializer.{DeserializationStream, SerializationStream,
 /** CallSite represents a place in user code. It can have a short and a long form. */
 private[spark] case class CallSite(shortForm: String, longForm: String)
 
+private[spark] object CallSite {
+  val SHORT_FORM = "callSite.short"
+  val LONG_FORM = "callSite.long"
+}
+
 /**
  * Various utility methods used by Spark.
  */
 private[spark] object Utils extends Logging {
   val random = new Random()
 
-  private[spark] val CALL_SITE_SHORT: String = "callSite.short"
-  private[spark] val CALL_SITE_LONG: String = "callSite.long"
-
   /** Serialize an object using Java serialization */
   def serialize[T](o: T): Array[Byte] = {
     val bos = new ByteArrayOutputStream()
@@ -854,24 +856,27 @@ 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
-  val SCALA_CLASS_REGEX = """^scala""".r
-
-  private def defaultRegexFunc(className: String): Boolean = {
-    SPARK_CLASS_REGEX.findFirstIn(className).isDefined ||
-    SCALA_CLASS_REGEX.findFirstIn(className).isDefined
+  /** Default filtering function for finding call sites using `getCallSite`. */
+  private def defaultCallSiteFilterFunc(className: String): Boolean = {
+    // A regular expression to match classes of the "core" Spark API that we want to skip when
+    // finding the call site of a method.
+    val SPARK_CORE_CLASS_REGEX = """^org\.apache\.spark(\.api\.java)?(\.util)?(\.rdd)?\.[A-Z]""".r
+    val SCALA_CLASS_REGEX = """^scala""".r
+    val isSparkClass = SPARK_CORE_CLASS_REGEX.findFirstIn(className).isDefined
+    val isScalaClass = SCALA_CLASS_REGEX.findFirstIn(className).isDefined
+    // If the class neither belongs to Spark nor is a simple Scala class, then it is a
+    // user-defined class
+    !isSparkClass && !isScalaClass
   }
 
   /**
    * When called inside a class in the spark package, returns the name of the user code class
    * (outside the spark package) that called into Spark, as well as which Spark method they called.
    * This is used, for example, to tell users where in their code each RDD got created.
+   *
+   * @param classFilterFunc Function that returns true if the given class belongs to user code
    */
-  def getCallSite(regexFunc: String => Boolean = defaultRegexFunc(_)): CallSite = {
+  def getCallSite(classFilterFunc: String => Boolean = defaultCallSiteFilterFunc): CallSite = {
     val trace = Thread.currentThread.getStackTrace()
       .filterNot { ste:StackTraceElement =>
         // When running under some profilers, the current stack trace might contain some bogus
@@ -892,7 +897,7 @@ private[spark] object Utils extends Logging {
 
     for (el <- trace) {
       if (insideSpark) {
-        if (regexFunc(el.getClassName)) {
+        if (!classFilterFunc(el.getClassName)) {
             lastSparkMethod = if (el.getMethodName == "<init>") {
             // Spark method is a constructor; get its class name
             el.getClassName.substring(el.getClassName.lastIndexOf('.') + 1)

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

@@ -441,9 +441,7 @@ class StreamingContext private[streaming] (
       throw new SparkException("StreamingContext has already been stopped")
     }
     validate()
-    sc.setCallSite(
-      Utils.getCallSite(org.apache.spark.streaming.util.Utils.streamingRegexFunc).shortForm
-    )
+    sparkContext.setCallSite(DStream.getCallSite())
     scheduler.start()
     state = Started
   }

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

@@ -23,14 +23,15 @@ import java.io.{IOException, ObjectInputStream, ObjectOutputStream}
 import scala.deprecated
 import scala.collection.mutable.HashMap
 import scala.reflect.ClassTag
+import scala.util.matching.Regex
 
 import org.apache.spark.{Logging, SparkException}
 import org.apache.spark.rdd.{BlockRDD, RDD}
 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.{CallSite, Utils, MetadataCleaner}
+import org.apache.spark.util.{CallSite, MetadataCleaner}
 
 /**
  * A Discretized Stream (DStream), the basic abstraction in Spark Streaming, is a continuous
@@ -106,20 +107,8 @@ abstract class DStream[T: ClassTag] (
   /** Return the StreamingContext associated with this DStream */
   def context = ssc
 
-  /* Find the creation callSite */
-  val creationSite = Utils.getCallSite(org.apache.spark.streaming.util.Utils.streamingRegexFunc)
-
-  /* Store the RDD creation callSite in threadlocal */
-  private def setRDDCreationCallSite(callSite: CallSite = creationSite) = {
-    ssc.sparkContext.setLocalProperty(Utils.CALL_SITE_SHORT, callSite.shortForm)
-    ssc.sparkContext.setLocalProperty(Utils.CALL_SITE_LONG, callSite.longForm)
-  }
-
-  /* Return the current callSite */
-  private def getRDDCreationCallSite(): CallSite = {
-    CallSite(ssc.sparkContext.getLocalProperty(Utils.CALL_SITE_SHORT),
-             ssc.sparkContext.getLocalProperty(Utils.CALL_SITE_LONG))
-  }
+  /* Set the creation call site */
+  private[streaming] val creationSite = DStream.getCallSite()
 
   /** Persist the RDDs of this DStream with the given storage level */
   def persist(level: StorageLevel): DStream[T] = {
@@ -264,6 +253,16 @@ abstract class DStream[T: ClassTag] (
     dependencies.foreach(_.setGraph(graph))
   }
 
+  /* Set the custom RDD creation site as this thread's local property. */
+  private def setRDDCreationSite(creationSite: CallSite): Unit = {
+
+  }
+
+  /* Get the custom RDD creation site set as this thread's local property. */
+  private def getRDDCreationSite(): CallSite = {
+    ssc.sparkContext.getCallSite()
+  }
+
   private[streaming] def remember(duration: Duration) {
     if (duration != null && duration > rememberDuration) {
       rememberDuration = duration
@@ -287,47 +286,41 @@ abstract class DStream[T: ClassTag] (
   }
 
   /**
-   * Retrieve a precomputed RDD of this DStream, or computes the RDD. This is an internal
-   * method that should not be called directly.
+   * Get the RDD corresponding to the given time; either retrieve it from cache
+   * or compute-and-cache it.
    */
   private[streaming] def getOrCompute(time: Time): Option[RDD[T]] = {
-    // If this DStream was not initialized (i.e., zeroTime not set), then do it
-    // If RDD was already generated, then retrieve it from HashMap
-    generatedRDDs.get(time) match {
-
-      // If an RDD was already generated and is being reused, then
-      // probably all RDDs in this DStream will be reused and hence should be cached
-      case Some(oldRDD) => Some(oldRDD)
-
-      // if RDD was not generated, and if the time is valid
-      // (based on sliding time of this DStream), then generate the RDD
-      case None => {
-        if (isTimeValid(time)) {
-          val prevCallSite = getRDDCreationCallSite
-          setRDDCreationCallSite()
-          val rddOption = compute(time) match {
-            case Some(newRDD) =>
-              if (storageLevel != StorageLevel.NONE) {
-                newRDD.persist(storageLevel)
-                logInfo("Persisting RDD " + newRDD.id + " for time " +
-                  time + " to " + storageLevel + " at time " + time)
-              }
-              if (checkpointDuration != null &&
-                (time - zeroTime).isMultipleOf(checkpointDuration)) {
-                newRDD.checkpoint()
-                logInfo("Marking RDD " + newRDD.id + " for time " + time +
-                  " for checkpointing at time " + time)
-              }
-              generatedRDDs.put(time, newRDD)
-              Some(newRDD)
-            case None =>
-              return None
+    // If RDD was already generated, then retrieve it from HashMap,
+    // or else compute the RDD
+    generatedRDDs.get(time).orElse {
+      // Compute the RDD if time is valid (e.g. correct time in a sliding window)
+      // of RDD generation, else generate nothing.
+      if (isTimeValid(time)) {
+        // Set the thread-local property for call sites to this DStream's creation site
+        // such that RDDs generated by compute gets that as their creation site.
+        // Note that this `getOrCompute` may get called from another DStream which may have
+        // set its own call site. So we store its call site in a temporary variable,
+        // set this DStream's creation site, generate RDDs and then restore the previous call site.
+        val prevCallSite = ssc.sparkContext.getCallSite()
+        ssc.sparkContext.setCallSite(creationSite)
+        val rddOption = compute(time)
+        ssc.sparkContext.setCallSite(prevCallSite)
+
+        rddOption.foreach { case newRDD =>
+          // Register the generated RDD for caching and checkpointing
+          if (storageLevel != StorageLevel.NONE) {
+            newRDD.persist(storageLevel)
+            logDebug(s"Persisting RDD ${newRDD.id} for time $time to $storageLevel")
           }
-          setRDDCreationCallSite(prevCallSite)
-          return rddOption
-        } else {
-          return None
+          if (checkpointDuration != null && (time - zeroTime).isMultipleOf(checkpointDuration)) {
+            newRDD.checkpoint()
+            logInfo(s"Marking RDD ${newRDD.id} for time $time for checkpointing")
+          }
+          generatedRDDs.put(time, newRDD)
         }
+        rddOption
+      } else {
+        None
       }
     }
   }
@@ -818,3 +811,29 @@ abstract class DStream[T: ClassTag] (
     this
   }
 }
+
+private[streaming] object DStream {
+
+  /** Get the creation site of a DStream from the stack trace of when the DStream is created. */
+  def getCallSite(): CallSite = {
+    val SPARK_CLASS_REGEX = """^org\.apache\.spark""".r
+    val SPARK_STREAMING_TESTCLASS_REGEX = """^org\.apache\.spark\.streaming\.test""".r
+    val SPARK_EXAMPLES_CLASS_REGEX = """^org\.apache\.spark\.examples""".r
+    val SCALA_CLASS_REGEX = """^scala""".r
+
+    /** Filtering function that returns true for classes that belong to a streaming application */
+    def streamingClassFilterFunc(className: String): Boolean = {
+      def doesMatch(r: Regex) = r.findFirstIn(className).isDefined
+      val isSparkClass = doesMatch(SPARK_CLASS_REGEX)
+      val isSparkExampleClass = doesMatch(SPARK_EXAMPLES_CLASS_REGEX)
+      val isSparkStreamingTestClass = doesMatch(SPARK_STREAMING_TESTCLASS_REGEX)
+      val isScalaClass = doesMatch(SCALA_CLASS_REGEX)
+
+      // If the class is a spark example class or a streaming test class then it is considered
+      // as a streaming application class. Otherwise, consider any non-Spark and non-Scala class
+      // as streaming application class.
+      isSparkExampleClass || isSparkStreamingTestClass || !(isSparkClass || isScalaClass)
+    }
+    org.apache.spark.util.Utils.getCallSite(streamingClassFilterFunc)
+  }
+}