[SPARK-6847][Core][Streaming]Fix stack overflow issue when updateStateByKey is followed by a checkpointed dstream

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

@@ -1535,6 +1535,10 @@ abstract class RDD[T: ClassTag](
 
   private[spark] var checkpointData: Option[RDDCheckpointData[T]] = None
 
+  // Whether recursively checkpoint all RDDs that are marked with the checkpoint flag.
+  private val recursiveCheckpoint =
+    Option(sc.getLocalProperty("spark.checkpoint.recursive")).map(_.toBoolean).getOrElse(false)
+
   /** Returns the first parent RDD */
   protected[spark] def firstParent[U: ClassTag]: RDD[U] = {
     dependencies.head.rdd.asInstanceOf[RDD[U]]
@@ -1578,6 +1582,11 @@ abstract class RDD[T: ClassTag](
       if (!doCheckpointCalled) {
         doCheckpointCalled = true
         if (checkpointData.isDefined) {
+          if (recursiveCheckpoint) {
+            // Checkpoint dependencies first because dependencies will be set to
+            // ReliableCheckpointRDD after checkpointing.
+            dependencies.foreach(_.rdd.doCheckpoint())
+          }
           checkpointData.get.checkpoint()
         } else {
           dependencies.foreach(_.rdd.doCheckpoint())

core/src/test/scala/org/apache/spark/CheckpointSuite.scala

@@ -512,6 +512,21 @@ class CheckpointSuite extends SparkFunSuite with RDDCheckpointTester with LocalS
     assert(rdd.isCheckpointedAndMaterialized === true)
     assert(rdd.partitions.size === 0)
   }
+
+  runTest("recursive RDD checkpoint") { reliableCheckpoint: Boolean =>
+    sc.setLocalProperty("spark.checkpoint.recursive", "true")
+    try {
+      val rdd1 = sc.parallelize(1 to 10)
+      checkpoint(rdd1, reliableCheckpoint)
+      val rdd2 = rdd1.map(_ + 1)
+      checkpoint(rdd2, reliableCheckpoint)
+      rdd2.count()
+      assert(rdd1.isCheckpointed === true)
+      assert(rdd2.isCheckpointed === true)
+    } finally {
+      sc.setLocalProperty("spark.checkpoint.recursive", null)
+    }
+  }
 }
 
 /** RDD partition that has large serialized size. */

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

@@ -243,6 +243,10 @@ class JobGenerator(jobScheduler: JobScheduler) extends Logging {
     // Example: BlockRDDs are created in this thread, and it needs to access BlockManager
     // Update: This is probably redundant after threadlocal stuff in SparkEnv has been removed.
     SparkEnv.set(ssc.env)
+
+    // Enable "spark.checkpoint.recursive" to make sure that all RDDs marked with the checkpoint
+    // flag are all checkpointed to avoid the stack overflow issue. See SPARK-6847
+    ssc.sparkContext.setLocalProperty("spark.checkpoint.recursive", "true")
     Try {
       jobScheduler.receiverTracker.allocateBlocksToBatch(time) // allocate received blocks to batch
       graph.generateJobs(time) // generate jobs using allocated block

streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala

@@ -821,6 +821,33 @@ class CheckpointSuite extends TestSuiteBase with DStreamCheckpointTester
     checkpointWriter.stop()
   }
 
+  test("SPARK-6847: stack overflow when updateStateByKey is followed by a checkpointed dstream") {
+    ssc = new StreamingContext(master, framework, batchDuration)
+    val batchCounter = new BatchCounter(ssc)
+    ssc.checkpoint(checkpointDir)
+    val inputDStream = new CheckpointInputDStream(ssc)
+    val updateFunc = (values: Seq[Int], state: Option[Int]) => {
+      Some(values.sum + state.getOrElse(0))
+    }
+    @volatile var recursiveCheckpoint = false
+    @volatile var rddsBothCheckpointed = false
+    inputDStream.map(i => (i, i)).
+      updateStateByKey[Int](updateFunc).checkpoint(batchDuration).
+      map(i => i).checkpoint(batchDuration).
+      foreachRDD { rdd =>
+        recursiveCheckpoint =
+          Option(rdd.sparkContext.getLocalProperty("spark.checkpoint.recursive")).
+            map(_.toBoolean).getOrElse(false)
+        val stateRDD = rdd.firstParent
+        rdd.count()
+        rddsBothCheckpointed = stateRDD.isCheckpointed && rdd.isCheckpointed
+      }
+    ssc.start()
+    batchCounter.waitUntilBatchesCompleted(1, 10000)
+    assert(recursiveCheckpoint === true)
+    assert(rddsBothCheckpointed === true)
+  }
+
   /**
    * Advances the manual clock on the streaming scheduler by given number of batches.
    * It also waits for the expected amount of time for each batch.