[EXT] Fix repeated batch after stop gracefully

zzcclp · May 19, 2016 · 3d637ab · jerryshao · May 19, 2016 · jerryshao
1 parent ff2590d
commit 3d637ab
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala b/streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala
@@ -32,7 +32,7 @@ import org.apache.spark.streaming.scheduler.JobGenerator
 
 
 private[streaming]
-class Checkpoint(ssc: StreamingContext, val checkpointTime: Time)
+class Checkpoint(ssc: StreamingContext, val checkpointTime: Time, val afterBatchCompletion: Boolean = false)
   extends Logging with Serializable {
   val master = ssc.sc.master
   val framework = ssc.sc.appName

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/JobGenerator.scala b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/JobGenerator.scala
@@ -209,7 +209,14 @@ class JobGenerator(jobScheduler: JobScheduler) extends Logging {
 
     // Batches when the master was down, that is,
     // between the checkpoint and current restart time
-    val checkpointTime = ssc.initialCheckpoint.checkpointTime
+    val checkpointTime = {
+      // if checkpoint after batch completion, start from next batch
+      if (ssc.initialCheckpoint.afterBatchCompletion) {
+        ssc.initialCheckpoint.checkpointTime + batchDuration
+      } else {
+        ssc.initialCheckpoint.checkpointTime
+      }
+    }
     val restartTime = new Time(timer.getRestartTime(graph.zeroTime.milliseconds))
     val downTimes = checkpointTime.until(restartTime, batchDuration)
     logInfo("Batches during down time (" + downTimes.size + " batches): "
@@ -292,7 +299,7 @@ class JobGenerator(jobScheduler: JobScheduler) extends Logging {
     if (shouldCheckpoint && (time - graph.zeroTime).isMultipleOf(ssc.checkpointDuration)) {
       logInfo("Checkpointing graph for time " + time)
       ssc.graph.updateCheckpointData(time)
-      checkpointWriter.write(new Checkpoint(ssc, time), clearCheckpointDataLater)
+      checkpointWriter.write(new Checkpoint(ssc, time, clearCheckpointDataLater), clearCheckpointDataLater)
     }
   }