Address a potential race when setting file modification times

Author: JoshRosen

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

@@ -320,6 +320,10 @@ class CheckpointSuite extends TestSuiteBase {
     val testDir = Utils.createTempDir()
     val outputBuffer = new ArrayBuffer[Seq[Int]] with SynchronizedBuffer[Seq[Int]]
 
+    /**
+     * Writes a file named `i` (which contains the number `i`) to the test directory and sets its
+     * modification time to `clock`'s current time.
+     */
     def writeFile(i: Int, clock: ManualClock): Unit = {
       val file = new File(testDir, i.toString)
       Files.write(i + "\n", file, Charsets.UTF_8)
@@ -329,6 +333,9 @@ class CheckpointSuite extends TestSuiteBase {
       assert(file.lastModified() === clock.currentTime())
     }
 
+    /**
+     * Returns ids that identify which files which have been recorded by the file input stream.
+     */
     def recordedFiles(ssc: StreamingContext): Seq[Int] = {
       val fileInputDStream =
         ssc.graph.getInputStreams().head.asInstanceOf[FileInputDStream[_, _, _]]
@@ -369,6 +376,8 @@ class CheckpointSuite extends TestSuiteBase {
         // Create files and advance manual clock to process them
         for (i <- Seq(1, 2, 3)) {
           writeFile(i, clock)
+          // Advance the clock after creating the file to avoid a race when
+          // setting its modification time
           clock.addToTime(batchDuration.milliseconds)
           if (i != 3) {
             // Since we want to shut down while the 3rd batch is processing
@@ -399,6 +408,8 @@ class CheckpointSuite extends TestSuiteBase {
       // Create files while the streaming driver is down
       for (i <- Seq(4, 5, 6)) {
         writeFile(i, clock)
+        // Advance the clock after creating the file to avoid a race when
+        // setting its modification time
         clock.addToTime(batchDuration.milliseconds)
       }
 
@@ -428,6 +439,8 @@ class CheckpointSuite extends TestSuiteBase {
         }
         for ((i, index) <- Seq(7, 8, 9).zipWithIndex) {
           writeFile(i, clock)
+          // Advance the clock after creating the file to avoid a race when
+          // setting its modification time
           clock.addToTime(batchDuration.milliseconds)
           eventually(eventuallyTimeout) {
             assert(batchCounter.getNumCompletedBatches === index + numBatchesAfterRestart + 1)

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

@@ -235,6 +235,7 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter {
   def testFileStream(newFilesOnly: Boolean) {
     val testDir: File = null
     try {
+      val batchDuration = Seconds(2)
       val testDir = Utils.createTempDir()
       // Create a file that exists before the StreamingContext is created:
       val existingFile = new File(testDir, "0")
@@ -245,7 +246,7 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter {
       withStreamingContext(new StreamingContext(conf, batchDuration)) { ssc =>
         val clock = ssc.scheduler.clock.asInstanceOf[ManualClock]
         // This `setTime` call ensures that the clock is past the creation time of `existingFile`
-        clock.setTime(existingFile.lastModified + 1000)
+        clock.setTime(existingFile.lastModified + batchDuration.milliseconds)
         val batchCounter = new BatchCounter(ssc)
         val fileStream = ssc.fileStream[LongWritable, Text, TextInputFormat](
           testDir.toString, (x: Path) => true, newFilesOnly = newFilesOnly).map(_._2.toString)
@@ -254,15 +255,21 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter {
         outputStream.register()
         ssc.start()
 
+        // Advance the clock so that the files are created after StreamingContext starts, but
+        // not enough to trigger a batch
+        clock.addToTime(batchDuration.milliseconds / 2)
+
         // Over time, create files in the directory
         val input = Seq(1, 2, 3, 4, 5)
         input.foreach { i =>
-          clock.addToTime(batchDuration.milliseconds)
           val file = new File(testDir, i.toString)
           Files.write(i + "\n", file, Charset.forName("UTF-8"))
           assert(file.setLastModified(clock.currentTime()))
           assert(file.lastModified === clock.currentTime)
           logInfo("Created file " + file)
+          // Advance the clock after creating the file to avoid a race when
+          // setting its modification time
+          clock.addToTime(batchDuration.milliseconds)
           eventually(eventuallyTimeout) {
             assert(batchCounter.getNumCompletedBatches === i)
           }