Added more tests to get 100% coverage of the WALBackedBlockRDD

Author: tdas

streaming/src/main/scala/org/apache/spark/streaming/rdd/WriteAheadLogBackedBlockRDD.scala

@@ -187,8 +187,14 @@ class WriteAheadLogBackedBlockRDD[T: ClassTag](
     blockLocations.getOrElse {
       partition.walRecordHandle match {
         case fileSegment: FileBasedWriteAheadLogSegment =>
-          HdfsUtils.getFileSegmentLocations(
-            fileSegment.path, fileSegment.offset, fileSegment.length, hadoopConfig)
+          try {
+            HdfsUtils.getFileSegmentLocations(
+              fileSegment.path, fileSegment.offset, fileSegment.length, hadoopConfig)
+          } catch {
+            case NonFatal(e) =>
+              logError("Error getting WAL file segment locations", e)
+              Seq.empty
+          }
         case _ =>
           Seq.empty
       }

streaming/src/test/scala/org/apache/spark/streaming/rdd/WriteAheadLogBackedBlockRDDSuite.scala

@@ -26,7 +26,7 @@ import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach, FunSuite}
 import org.apache.spark.storage.{BlockId, BlockManager, StorageLevel, StreamBlockId}
 import org.apache.spark.streaming.util.{FileBasedWriteAheadLogSegment, FileBasedWriteAheadLogWriter}
 import org.apache.spark.util.Utils
-import org.apache.spark.{SparkConf, SparkContext}
+import org.apache.spark.{SparkConf, SparkContext, SparkException}
 
 class WriteAheadLogBackedBlockRDDSuite
   extends FunSuite with BeforeAndAfterAll with BeforeAndAfterEach {
@@ -60,48 +60,88 @@ class WriteAheadLogBackedBlockRDDSuite
     System.clearProperty("spark.driver.port")
   }
 
-  test("Read data available in block manager and write ahead log") {
-    testRDD(5, 5)
+  test("Read data available in both block manager and write ahead log") {
+    testRDD(numPartitions = 5, numPartitionsInBM = 5, numPartitionsInWAL = 5)
   }
 
   test("Read data available only in block manager, not in write ahead log") {
-    testRDD(5, 0)
+    testRDD(numPartitions = 5, numPartitionsInBM = 5, numPartitionsInWAL = 0)
   }
 
   test("Read data available only in write ahead log, not in block manager") {
-    testRDD(0, 5)
+    testRDD(numPartitions = 5, numPartitionsInBM = 0, numPartitionsInWAL = 5)
   }
 
-  test("Read data available only in write ahead log, and test storing in block manager") {
-    testRDD(0, 5, testStoreInBM = true)
+  test("Read data with partially available in block manager, and rest in write ahead log") {
+    testRDD(numPartitions = 5, numPartitionsInBM = 3, numPartitionsInWAL = 2)
   }
 
-  test("Read data with partially available in block manager, and rest in write ahead log") {
-    testRDD(3, 2)
+  test("Test isBlockValid skips block fetching from BlockManager") {
+    testRDD(
+      numPartitions = 5, numPartitionsInBM = 5, numPartitionsInWAL = 0, testIsBlockValid = true)
+  }
+
+  test("Test whether RDD is valid after removing blocks from block manager") {
+    testRDD(
+      numPartitions = 5, numPartitionsInBM = 5, numPartitionsInWAL = 5, testBlockRemove = true)
+  }
+
+  test("Test storing of blocks recovered from write ahead log back into block manager") {
+    testRDD(
+      numPartitions = 5, numPartitionsInBM = 0, numPartitionsInWAL = 5, testStoreInBM = true)
   }
 
   /**
    * Test the WriteAheadLogBackedRDD, by writing some partitions of the data to block manager
    * and the rest to a write ahead log, and then reading reading it all back using the RDD.
    * It can also test if the partitions that were read from the log were again stored in
    * block manager.
-   * @param numPartitionsInBM Number of partitions to write to the Block Manager
-   * @param numPartitionsInWAL Number of partitions to write to the Write Ahead Log
-   * @param testStoreInBM Test whether blocks read from log are stored back into block manager
+   *
+   *
+   *
+   * @param numPartitions Number of partitions in RDD
+   * @param numPartitionsInBM Number of partitions to write to the BlockManager.
+   *                          Partitions 0 to (numPartitionsInBM-1) will be written to BlockManager
+   * @param numPartitionsInWAL Number of partitions to write to the Write Ahead Log.
+   *                           Partitions (numPartitions - 1 - numPartitionsInWAL) to
+   *                           (numPartitions - 1) will be written to WAL
+   * @param testIsBlockValid Test whether setting isBlockValid to false skips block fetching
+   * @param testBlockRemove Test whether calling rdd.removeBlock() makes the RDD still usable with
+   *                        reads falling back to the WAL
+   * @param testStoreInBM   Test whether blocks read from log are stored back into block manager
+   *
+   * Example with numPartitions = 5, numPartitionsInBM = 3, and numPartitionsInWAL = 4
+   *
+   *   numPartitionsInBM = 3
+   *   |------------------|
+   *   |                  |
+   *    0       1       2       3       4
+   *           |                         |
+   *           |-------------------------|
+   *              numPartitionsInWAL = 4
    */
   private def testRDD(
-      numPartitionsInBM: Int, numPartitionsInWAL: Int, testStoreInBM: Boolean = false) {
-    val numBlocks = numPartitionsInBM + numPartitionsInWAL
-    val data = Seq.fill(numBlocks, 10)(scala.util.Random.nextString(50))
+      numPartitions: Int,
+      numPartitionsInBM: Int,
+      numPartitionsInWAL: Int,
+      testIsBlockValid: Boolean = false,
+      testBlockRemove: Boolean = false,
+      testStoreInBM: Boolean = false
+    ) {
+    require(numPartitionsInBM <= numPartitions,
+      "Can't put more partitions in BlockManager than that in RDD")
+    require(numPartitionsInWAL <= numPartitions,
+      "Can't put more partitions in write ahead log than that in RDD")
+    val data = Seq.fill(numPartitions, 10)(scala.util.Random.nextString(50))
 
     // Put the necessary blocks in the block manager
-    val blockIds = Array.fill(numBlocks)(StreamBlockId(Random.nextInt(), Random.nextInt()))
+    val blockIds = Array.fill(numPartitions)(StreamBlockId(Random.nextInt(), Random.nextInt()))
     data.zip(blockIds).take(numPartitionsInBM).foreach { case(block, blockId) =>
       blockManager.putIterator(blockId, block.iterator, StorageLevel.MEMORY_ONLY_SER)
     }
 
     // Generate write ahead log record handles
-    val recordHandles = generateFakeRecordHandles(numPartitionsInBM) ++
+    val recordHandles = generateFakeRecordHandles(numPartitions - numPartitionsInWAL) ++
       generateWALRecordHandles(data.takeRight(numPartitionsInWAL),
         blockIds.takeRight(numPartitionsInWAL))
 
@@ -111,7 +151,7 @@ class WriteAheadLogBackedBlockRDDSuite
       "Expected blocks not in BlockManager"
     )
     require(
-      blockIds.takeRight(numPartitionsInWAL).forall(blockManager.get(_).isEmpty),
+      blockIds.takeRight(numPartitions - numPartitionsInBM).forall(blockManager.get(_).isEmpty),
       "Unexpected blocks in BlockManager"
     )
 
@@ -122,16 +162,39 @@ class WriteAheadLogBackedBlockRDDSuite
       "Expected blocks not in write ahead log"
     )
     require(
-      recordHandles.take(numPartitionsInBM).forall(s =>
+      recordHandles.take(numPartitions - numPartitionsInWAL).forall(s =>
         !new File(s.path.stripPrefix("file://")).exists()),
       "Unexpected blocks in write ahead log"
     )
 
     // Create the RDD and verify whether the returned data is correct
     val rdd = new WriteAheadLogBackedBlockRDD[String](sparkContext, blockIds.toArray,
-      recordHandles.toArray, storeInBlockManager = false, storageLevel = StorageLevel.MEMORY_ONLY)
+      recordHandles.toArray, storeInBlockManager = false)
     assert(rdd.collect() === data.flatten)
 
+    // Verify that the block fetching is skipped when isBlockValid is set to false.
+    // This is done by using a RDD whose data is only in memory but is set to skip block fetching
+    // Using that RDD will throw exception, as it skips block fetching even if the blocks are in
+    // in BlockManager.
+    if (testIsBlockValid) {
+      require(numPartitionsInBM === numPartitions, "All partitions must be in BlockManager")
+      require(numPartitionsInWAL === 0, "No partitions must be in WAL")
+      val rdd2 = new WriteAheadLogBackedBlockRDD[String](sparkContext, blockIds.toArray,
+        recordHandles.toArray, isBlockIdValid = Array.fill(blockIds.length)(false))
+      intercept[SparkException] {
+        rdd2.collect()
+      }
+    }
+
+    // Verify that the RDD is not invalid after the blocks are removed and can still read data
+    // from write ahead log
+    if (testBlockRemove) {
+      require(numPartitions === numPartitionsInWAL, "All partitions must be in WAL for this test")
+      require(numPartitionsInBM > 0, "Some partitions must be in BlockManager for this test")
+      rdd.removeBlocks()
+      assert(rdd.collect() === data.flatten)
+    }
+
     if (testStoreInBM) {
       val rdd2 = new WriteAheadLogBackedBlockRDD[String](sparkContext, blockIds.toArray,
         recordHandles.toArray, storeInBlockManager = true, storageLevel = StorageLevel.MEMORY_ONLY)