Try to read from cache instead of checkpoint after checkpointing

Author: zsxwing

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

@@ -31,7 +31,8 @@ class CheckpointManager extends Logging {
   /** Keys of RDD partitions that are being checkpointed. */
   private val checkpointingRDDPartitions = new mutable.HashSet[RDDBlockId]
 
-  /** Gets or computes an RDD partition. Used by RDD.iterator() when an RDD is cached. */
+  /** Gets or computes an RDD partition. Used by RDD.iterator() when an RDD is about to be
+    * checkpointed. */
   def getOrCompute[T: ClassTag](
       rdd: RDD[T],
       checkpointData: ReliableRDDCheckpointData[T],
@@ -49,7 +50,7 @@ class CheckpointManager extends Logging {
     } else {
       // Acquire a lock for loading this partition
       // If another thread already holds the lock, wait for it to finish return its results
-      val checkpoint = acquireLockForPartition[T](key, path, conf, context)
+      val checkpoint = acquireLockForPartition[T](rdd, partition, key, context)
       if (checkpoint.isDefined) {
         return new InterruptibleIterator[T](context, checkpoint.get)
       }
@@ -61,7 +62,7 @@ class CheckpointManager extends Logging {
       val computedValues = rdd.computeOrReadCache(partition, context)
       ReliableCheckpointRDD.writeCheckpointFile(
         context, computedValues, checkpointData.cpDir, conf, partition.index)
-      ReliableCheckpointRDD.readCheckpointFile(path, conf, context)
+      rdd.computeOrReadCache(partition, context)
     } finally {
       checkpointingRDDPartitions.synchronized {
         checkpointingRDDPartitions.remove(key)
@@ -78,8 +79,10 @@ class CheckpointManager extends Logging {
    * thread.
    */
   private def acquireLockForPartition[T](
-      id: RDDBlockId, path: Path, conf: Configuration, context: TaskContext): Option[Iterator[T]] =
-  {
+      rdd: RDD[T],
+      partition: Partition,
+      id: RDDBlockId,
+      context: TaskContext): Option[Iterator[T]] = {
     checkpointingRDDPartitions.synchronized {
       if (!checkpointingRDDPartitions.contains(id)) {
         // If the partition is free, acquire its lock to compute its value
@@ -94,7 +97,7 @@ class CheckpointManager extends Logging {
         logInfo(s"Finished waiting for $id")
       }
     }
-    Some(ReliableCheckpointRDD.readCheckpointFile(path, conf, context))
+    Some(rdd.computeOrReadCache(partition, context))
   }
 
 }

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

@@ -81,7 +81,7 @@ class SparkEnv (
   @deprecated("Actor system is no longer supported as of 1.4.0", "1.4.0")
   val actorSystem: ActorSystem = _actorSystem
 
-  private[spark] val checkpointMananger = new CheckpointManager
+  private[spark] val checkpointManager = new CheckpointManager
 
   private[spark] var isStopped = false
   private val pythonWorkers = mutable.HashMap[(String, Map[String, String]), PythonWorkerFactory]()

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

@@ -260,7 +260,7 @@ abstract class RDD[T: ClassTag](
   final def iterator(split: Partition, context: TaskContext): Iterator[T] = {
     if (!isCheckpointedAndMaterialized) {
       if (checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[T]])) {
-        return SparkEnv.get.checkpointMananger.getOrCompute(
+        return SparkEnv.get.checkpointManager.getOrCompute(
           this, checkpointData.get.asInstanceOf[ReliableRDDCheckpointData[T]], split, context)
       }
     }