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

Add a local property to indicate if checkpointing all RDDs that are marked with the checkpoint flag, and enable it in Streaming

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #10934 from zsxwing/recursive-checkpoint.

1 files changed, 21 insertions, 0 deletions

diff --git a/core/src/test/scala/org/apache/spark/CheckpointSuite.scala b/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
index 390764ba24..ce35856dce 100644
--- a/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
+++ b/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
@@ -512,6 +512,27 @@ class CheckpointSuite extends SparkFunSuite with RDDCheckpointTester with LocalS
     assert(rdd.isCheckpointedAndMaterialized === true)
     assert(rdd.partitions.size === 0)
   }
+
+  runTest("checkpointAllMarkedAncestors") { reliableCheckpoint: Boolean =>
+    testCheckpointAllMarkedAncestors(reliableCheckpoint, checkpointAllMarkedAncestors = true)
+    testCheckpointAllMarkedAncestors(reliableCheckpoint, checkpointAllMarkedAncestors = false)
+  }
+
+  private def testCheckpointAllMarkedAncestors(
+      reliableCheckpoint: Boolean, checkpointAllMarkedAncestors: Boolean): Unit = {
+    sc.setLocalProperty(RDD.CHECKPOINT_ALL_MARKED_ANCESTORS, checkpointAllMarkedAncestors.toString)
+    try {
+      val rdd1 = sc.parallelize(1 to 10)
+      checkpoint(rdd1, reliableCheckpoint)
+      val rdd2 = rdd1.map(_ + 1)
+      checkpoint(rdd2, reliableCheckpoint)
+      rdd2.count()
+      assert(rdd1.isCheckpointed === checkpointAllMarkedAncestors)
+      assert(rdd2.isCheckpointed === true)
+    } finally {
+      sc.setLocalProperty(RDD.CHECKPOINT_ALL_MARKED_ANCESTORS, null)
+    }
+  }
 }
 
 /** RDD partition that has large serialized size. */