<commit_before>package spark.streaming

import spark.streaming.StreamingContext._
import java.io.File
import runtime.RichInt
import org.scalatest.BeforeAndAfter
import org.apache.commons.io.FileUtils
import collection.mutable.{SynchronizedBuffer, ArrayBuffer}
import util.{Clock, ManualClock}

class CheckpointSuite extends TestSuiteBase with BeforeAndAfter {

  before {
    FileUtils.deleteDirectory(new File(checkpointDir))
  }

  after {
    FileUtils.deleteDirectory(new File(checkpointDir))
  }

  override def framework = "CheckpointSuite"

  override def batchDuration = Milliseconds(500)

  override def checkpointDir = "checkpoint"

  override def checkpointInterval = batchDuration

  override def actuallyWait = true

  test("basic stream+rdd recovery") {

    assert(batchDuration === Milliseconds(500), "batchDuration for this test must be 1 second")
    assert(checkpointInterval === batchDuration, "checkpointInterval for this test much be same as batchDuration")

    System.setProperty("spark.streaming.clock", "spark.streaming.util.ManualClock")

    val stateStreamCheckpointInterval = Seconds(2)

    // this ensure checkpointing occurs at least once
    val firstNumBatches = (stateStreamCheckpointInterval.millis / batchDuration.millis) * 2
    val secondNumBatches = firstNumBatches

    // Setup the streams
    val input = (1 to 10).map(_ => Seq("a")).toSeq
    val operation = (st: DStream[String]) => {
      val updateFunc = (values: Seq[Int], state: Option[RichInt]) => {
        Some(new RichInt(values.foldLeft(0)(_ + _) + state.map(_.self).getOrElse(0)))
      }
      st.map(x => (x, 1))
      .updateStateByKey[RichInt](updateFunc)
      .checkpoint(stateStreamCheckpointInterval)
      .map(t => (t._1, t._2.self))
    }
    var ssc = setupStreams(input, operation)
    var stateStream = ssc.graph.getOutputStreams().head.dependencies.head.dependencies.head

    // Run till a time such that at least one RDD in the stream should have been checkpointed,
    // then check whether some RDD has been checkpointed or not
    ssc.start()
    runStreamsWithRealDelay(ssc, firstNumBatches)
    logInfo("Checkpoint data of state stream = \n[" + stateStream.checkpointData.mkString(",\n") + "]")
    assert(!stateStream.checkpointData.isEmpty, "No checkpointed RDDs in state stream before first failure")
    stateStream.checkpointData.foreach {
      case (time, data) => {
        val file = new File(data.toString)
        assert(file.exists(), "Checkpoint file '" + file +"' for time " + time + " for state stream before first failure does not exist")
      }
    }

    // Run till a further time such that previous checkpoint files in the stream would be deleted
    // and check whether the earlier checkpoint files are deleted
    val checkpointFiles = stateStream.checkpointData.map(x => new File(x._2.toString))
    runStreamsWithRealDelay(ssc, secondNumBatches)
    checkpointFiles.foreach(file => assert(!file.exists, "Checkpoint file '" + file + "' was not deleted"))
    ssc.stop()

    // Restart stream computation using the checkpoint file and check whether
    // checkpointed RDDs have been restored or not
    ssc = new StreamingContext(checkpointDir)
    stateStream = ssc.graph.getOutputStreams().head.dependencies.head.dependencies.head
    logInfo("Restored data of state stream = \n[" + stateStream.generatedRDDs.mkString("\n") + "]")
    assert(!stateStream.generatedRDDs.isEmpty, "No restored RDDs in state stream after recovery from first failure")


    // Run one batch to generate a new checkpoint file and check whether some RDD
    // is present in the checkpoint data or not
    ssc.start()
    runStreamsWithRealDelay(ssc, 1)
    assert(!stateStream.checkpointData.isEmpty, "No checkpointed RDDs in state stream before second failure")
    stateStream.checkpointData.foreach {
      case (time, data) => {
        val file = new File(data.toString)
        assert(file.exists(),
          "Checkpoint file '" + file +"' for time " + time + " for state stream before seconds failure does not exist")
      }
    }
    ssc.stop()

    // Restart stream computation from the new checkpoint file to see whether that file has
    // correct checkpoint data
    ssc = new StreamingContext(checkpointDir)
    stateStream = ssc.graph.getOutputStreams().head.dependencies.head.dependencies.head
    logInfo("Restored data of state stream = \n[" + stateStream.generatedRDDs.mkString("\n") + "]")
    assert(!stateStream.generatedRDDs.isEmpty, "No restored RDDs in state stream after recovery from second failure")

    // Adjust manual clock time as if it is being restarted after a delay
    System.setProperty("spark.streaming.manualClock.jump", (batchDuration.milliseconds * 7).toString)
    ssc.start()
    runStreamsWithRealDelay(ssc, 4)
    ssc.stop()
    System.clearProperty("spark.streaming.manualClock.jump")
  }

  test("map and reduceByKey") {
    testCheckpointedOperation(
      Seq( Seq("a", "a", "b"), Seq("", ""), Seq(), Seq("a", "a", "b"), Seq("", ""), Seq() ),
      (s: DStream[String]) => s.map(x => (x, 1)).reduceByKey(_ + _),
      Seq( Seq(("a", 2), ("b", 1)), Seq(("", 2)), Seq(), Seq(("a", 2), ("b", 1)), Seq(("", 2)), Seq() ),
      3
    )
  }

  test("reduceByKeyAndWindowInv") {
    val n = 10
    val w = 4
    val input = (1 to n).map(_ => Seq("a")).toSeq
    val output = Seq(Seq(("a", 1)), Seq(("a", 2)), Seq(("a", 3))) ++ (1 to (n - w + 1)).map(x => Seq(("a", 4)))
    val operation = (st: DStream[String]) => {
      st.map(x => (x, 1))
        .reduceByKeyAndWindow(_ + _, _ - _, batchDuration * w, batchDuration)
        .checkpoint(Seconds(2))
    }
    for (i <- Seq(2, 3, 4)) {
      testCheckpointedOperation(input, operation, output, i)
    }
  }

  test("updateStateByKey") {
    val input = (1 to 10).map(_ => Seq("a")).toSeq
    val output = (1 to 10).map(x => Seq(("a", x))).toSeq
    val operation = (st: DStream[String]) => {
      val updateFunc = (values: Seq[Int], state: Option[RichInt]) => {
        Some(new RichInt(values.foldLeft(0)(_ + _) + state.map(_.self).getOrElse(0)))
      }
      st.map(x => (x, 1))
        .updateStateByKey[RichInt](updateFunc)
        .checkpoint(Seconds(2))
        .map(t => (t._1, t._2.self))
    }
    for (i <- Seq(2, 3, 4)) {
      testCheckpointedOperation(input, operation, output, i)
    }
  }



  def testCheckpointedOperation[U: ClassManifest, V: ClassManifest](
    input: Seq[Seq[U]],
    operation: DStream[U] => DStream[V],
    expectedOutput: Seq[Seq[V]],
    initialNumBatches: Int
  ) {

    // Current code assumes that:
    // number of inputs = number of outputs = number of batches to be run
    val totalNumBatches = input.size
    val nextNumBatches = totalNumBatches - initialNumBatches
    val initialNumExpectedOutputs = initialNumBatches
    val nextNumExpectedOutputs = expectedOutput.size - initialNumExpectedOutputs

    // Do half the computation (half the number of batches), create checkpoint file and quit

    val ssc = setupStreams[U, V](input, operation)
    val output = runStreams[V](ssc, initialNumBatches, initialNumExpectedOutputs)
    verifyOutput[V](output, expectedOutput.take(initialNumBatches), true)
    Thread.sleep(1000)

    // Restart and complete the computation from checkpoint file
    logInfo(
      "\n-------------------------------------------\n" +
      "        Restarting stream computation          " +
      "\n-------------------------------------------\n"
    )
    val sscNew = new StreamingContext(checkpointDir)
    val outputNew = runStreams[V](sscNew, nextNumBatches, nextNumExpectedOutputs)
    verifyOutput[V](outputNew, expectedOutput.takeRight(nextNumExpectedOutputs), true)
  }

  /**
   * Advances the manual clock on the streaming scheduler by given number of batches.
   * It also wait for the expected amount of time for each batch.
   */


  def runStreamsWithRealDelay(ssc: StreamingContext, numBatches: Long) {
    val clock = ssc.scheduler.clock.asInstanceOf[ManualClock]
    logInfo("Manual clock before advancing = " + clock.time)
    for (i <- 1 to numBatches.toInt) {
      clock.addToTime(batchDuration.milliseconds)
      Thread.sleep(batchDuration.milliseconds)
    }
    logInfo("Manual clock after advancing = " + clock.time)
    Thread.sleep(batchDuration.milliseconds)
  }

}<commit_msg>Speeded up CheckpointSuite<commit_after>package spark.streaming

import spark.streaming.StreamingContext._
import java.io.File
import runtime.RichInt
import org.scalatest.BeforeAndAfter
import org.apache.commons.io.FileUtils
import collection.mutable.{SynchronizedBuffer, ArrayBuffer}
import util.{Clock, ManualClock}

class CheckpointSuite extends TestSuiteBase with BeforeAndAfter {

  before {
    FileUtils.deleteDirectory(new File(checkpointDir))
  }

  after {

    if (ssc != null) ssc.stop()
    FileUtils.deleteDirectory(new File(checkpointDir))
  }

  var ssc: StreamingContext = null

  override def framework = "CheckpointSuite"

  override def batchDuration = Milliseconds(200)

  override def checkpointDir = "checkpoint"

  override def checkpointInterval = batchDuration

  override def actuallyWait = true

  test("basic stream+rdd recovery") {

    assert(batchDuration === Milliseconds(200), "batchDuration for this test must be 1 second")
    assert(checkpointInterval === batchDuration, "checkpointInterval for this test much be same as batchDuration")

    System.setProperty("spark.streaming.clock", "spark.streaming.util.ManualClock")

    val stateStreamCheckpointInterval = Seconds(1)

    // this ensure checkpointing occurs at least once
    val firstNumBatches = (stateStreamCheckpointInterval.millis / batchDuration.millis) * 2
    val secondNumBatches = firstNumBatches

    // Setup the streams
    val input = (1 to 10).map(_ => Seq("a")).toSeq
    val operation = (st: DStream[String]) => {
      val updateFunc = (values: Seq[Int], state: Option[RichInt]) => {
        Some(new RichInt(values.foldLeft(0)(_ + _) + state.map(_.self).getOrElse(0)))
      }
      st.map(x => (x, 1))
      .updateStateByKey[RichInt](updateFunc)
      .checkpoint(stateStreamCheckpointInterval)
      .map(t => (t._1, t._2.self))
    }
    var ssc = setupStreams(input, operation)
    var stateStream = ssc.graph.getOutputStreams().head.dependencies.head.dependencies.head

    // Run till a time such that at least one RDD in the stream should have been checkpointed,
    // then check whether some RDD has been checkpointed or not
    ssc.start()
    runStreamsWithRealDelay(ssc, firstNumBatches)
    logInfo("Checkpoint data of state stream = \n[" + stateStream.checkpointData.mkString(",\n") + "]")
    assert(!stateStream.checkpointData.isEmpty, "No checkpointed RDDs in state stream before first failure")
    stateStream.checkpointData.foreach {
      case (time, data) => {
        val file = new File(data.toString)
        assert(file.exists(), "Checkpoint file '" + file +"' for time " + time + " for state stream before first failure does not exist")
      }
    }

    // Run till a further time such that previous checkpoint files in the stream would be deleted
    // and check whether the earlier checkpoint files are deleted
    val checkpointFiles = stateStream.checkpointData.map(x => new File(x._2.toString))
    runStreamsWithRealDelay(ssc, secondNumBatches)
    checkpointFiles.foreach(file => assert(!file.exists, "Checkpoint file '" + file + "' was not deleted"))
    ssc.stop()

    // Restart stream computation using the checkpoint file and check whether
    // checkpointed RDDs have been restored or not
    ssc = new StreamingContext(checkpointDir)
    stateStream = ssc.graph.getOutputStreams().head.dependencies.head.dependencies.head
    logInfo("Restored data of state stream = \n[" + stateStream.generatedRDDs.mkString("\n") + "]")
    assert(!stateStream.generatedRDDs.isEmpty, "No restored RDDs in state stream after recovery from first failure")


    // Run one batch to generate a new checkpoint file and check whether some RDD
    // is present in the checkpoint data or not
    ssc.start()
    runStreamsWithRealDelay(ssc, 1)
    assert(!stateStream.checkpointData.isEmpty, "No checkpointed RDDs in state stream before second failure")
    stateStream.checkpointData.foreach {
      case (time, data) => {
        val file = new File(data.toString)
        assert(file.exists(),
          "Checkpoint file '" + file +"' for time " + time + " for state stream before seconds failure does not exist")
      }
    }
    ssc.stop()

    // Restart stream computation from the new checkpoint file to see whether that file has
    // correct checkpoint data
    ssc = new StreamingContext(checkpointDir)
    stateStream = ssc.graph.getOutputStreams().head.dependencies.head.dependencies.head
    logInfo("Restored data of state stream = \n[" + stateStream.generatedRDDs.mkString("\n") + "]")
    assert(!stateStream.generatedRDDs.isEmpty, "No restored RDDs in state stream after recovery from second failure")

    // Adjust manual clock time as if it is being restarted after a delay
    System.setProperty("spark.streaming.manualClock.jump", (batchDuration.milliseconds * 7).toString)
    ssc.start()
    runStreamsWithRealDelay(ssc, 4)
    ssc.stop()
    System.clearProperty("spark.streaming.manualClock.jump")
    ssc = null
  }

  test("map and reduceByKey") {
    testCheckpointedOperation(
      Seq( Seq("a", "a", "b"), Seq("", ""), Seq(), Seq("a", "a", "b"), Seq("", ""), Seq() ),
      (s: DStream[String]) => s.map(x => (x, 1)).reduceByKey(_ + _),
      Seq( Seq(("a", 2), ("b", 1)), Seq(("", 2)), Seq(), Seq(("a", 2), ("b", 1)), Seq(("", 2)), Seq() ),
      3
    )
  }

  test("reduceByKeyAndWindowInv") {
    val n = 10
    val w = 4
    val input = (1 to n).map(_ => Seq("a")).toSeq
    val output = Seq(Seq(("a", 1)), Seq(("a", 2)), Seq(("a", 3))) ++ (1 to (n - w + 1)).map(x => Seq(("a", 4)))
    val operation = (st: DStream[String]) => {
      st.map(x => (x, 1))
        .reduceByKeyAndWindow(_ + _, _ - _, batchDuration * w, batchDuration)
        .checkpoint(Seconds(2))
    }
    testCheckpointedOperation(input, operation, output, 3)
  }

  test("updateStateByKey") {
    val input = (1 to 10).map(_ => Seq("a")).toSeq
    val output = (1 to 10).map(x => Seq(("a", x))).toSeq
    val operation = (st: DStream[String]) => {
      val updateFunc = (values: Seq[Int], state: Option[RichInt]) => {
        Some(new RichInt(values.foldLeft(0)(_ + _) + state.map(_.self).getOrElse(0)))
      }
      st.map(x => (x, 1))
        .updateStateByKey[RichInt](updateFunc)
        .checkpoint(Seconds(2))
        .map(t => (t._1, t._2.self))
    }
    testCheckpointedOperation(input, operation, output, 3)
  }



  def testCheckpointedOperation[U: ClassManifest, V: ClassManifest](
    input: Seq[Seq[U]],
    operation: DStream[U] => DStream[V],
    expectedOutput: Seq[Seq[V]],
    initialNumBatches: Int
  ) {

    // Current code assumes that:
    // number of inputs = number of outputs = number of batches to be run
    val totalNumBatches = input.size
    val nextNumBatches = totalNumBatches - initialNumBatches
    val initialNumExpectedOutputs = initialNumBatches
    val nextNumExpectedOutputs = expectedOutput.size - initialNumExpectedOutputs

    // Do half the computation (half the number of batches), create checkpoint file and quit

    ssc = setupStreams[U, V](input, operation)
    val output = runStreams[V](ssc, initialNumBatches, initialNumExpectedOutputs)
    verifyOutput[V](output, expectedOutput.take(initialNumBatches), true)
    Thread.sleep(1000)

    // Restart and complete the computation from checkpoint file
    logInfo(
      "\n-------------------------------------------\n" +
      "        Restarting stream computation          " +
      "\n-------------------------------------------\n"
    )
    ssc = new StreamingContext(checkpointDir)
    val outputNew = runStreams[V](ssc, nextNumBatches, nextNumExpectedOutputs)
    verifyOutput[V](outputNew, expectedOutput.takeRight(nextNumExpectedOutputs), true)
    ssc = null
  }

  /**
   * Advances the manual clock on the streaming scheduler by given number of batches.
   * It also wait for the expected amount of time for each batch.
   */


  def runStreamsWithRealDelay(ssc: StreamingContext, numBatches: Long) {
    val clock = ssc.scheduler.clock.asInstanceOf[ManualClock]
    logInfo("Manual clock before advancing = " + clock.time)
    for (i <- 1 to numBatches.toInt) {
      clock.addToTime(batchDuration.milliseconds)
      Thread.sleep(batchDuration.milliseconds)
    }
    logInfo("Manual clock after advancing = " + clock.time)
    Thread.sleep(batchDuration.milliseconds)
  }

}
{'fraction_non_alphanumeric': 0.1314494840412767, 'fraction_numerical': 0.007439404847612191, 'mean_word_length': 3.3567694720334553, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 47, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4152865', 'n_tokens_mistral': 5213, 'n_tokens_neox': 4959, 'n_words': 1370}