Add unit tests

Author: EnricoMi

core/src/test/scala/org/apache/spark/storage/FallbackStorageSuite.scala

@@ -22,10 +22,11 @@ import java.nio.file.Files
 import scala.concurrent.duration._
 import scala.util.Random
 
+import io.netty.buffer.ByteBuf
 import org.apache.hadoop.conf.Configuration
-import org.apache.hadoop.fs.{FSDataInputStream, LocalFileSystem, Path, PositionedReadable, Seekable}
+import org.apache.hadoop.fs.{FileSystem, FSDataInputStream, LocalFileSystem, Path, PositionedReadable, Seekable}
 import org.mockito.{ArgumentMatchers => mc}
-import org.mockito.Mockito.{mock, never, verify, when}
+import org.mockito.Mockito.{mock, never, spy, times, verify, when}
 import org.scalatest.concurrent.Eventually.{eventually, interval, timeout}
 
 import org.apache.spark.{LocalSparkContext, SparkConf, SparkContext, SparkFunSuite, TestUtils}
@@ -110,7 +111,9 @@ class FallbackStorageSuite extends SparkFunSuite with LocalSparkContext {
     intercept[java.io.EOFException] {
       FallbackStorage.read(conf, ShuffleBlockId(1, 1L, 0))
     }
-    FallbackStorage.read(conf, ShuffleBlockId(1, 2L, 0))
+    val readResult = FallbackStorage.read(conf, ShuffleBlockId(1, 2L, 0))
+    assert(readResult.isInstanceOf[FileSystemSegmentManagedBuffer])
+    readResult.createInputStream().close()
   }
 
   test("SPARK-39200: fallback storage APIs - readFully") {
@@ -155,9 +158,49 @@ class FallbackStorageSuite extends SparkFunSuite with LocalSparkContext {
     assert(fallbackStorage.exists(1, ShuffleDataBlockId(1, 2L, NOOP_REDUCE_ID).name))
 
     val readResult = FallbackStorage.read(conf, ShuffleBlockId(1, 2L, 0))
+    assert(readResult.isInstanceOf[FileSystemSegmentManagedBuffer])
     assert(readResult.nioByteBuffer().array().sameElements(content))
   }
 
+  test("SPARK-55469: FileSystemSegmentManagedBuffer reads block data lazily") {
+    withTempDir { dir =>
+      val fs = FileSystem.getLocal(new Configuration())
+      val file = new Path(dir.getAbsolutePath, "file")
+      val data = Array[Byte](1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
+      tryWithResource(fs.create(file)) { os => os.write(data) }
+
+      Seq((0, 4), (1, 2), (4, 4), (7, 2), (8, 0)).foreach { case (offset, length) =>
+        val clue = s"offset: $offset, length: $length"
+
+        // creating the managed buffer does not open the file
+        val mfs = spy(fs)
+        val buf = new FileSystemSegmentManagedBuffer(mfs, file, offset, length)
+        verify(mfs, never()).open(mc.any[Path]())
+        assert(buf.size() === length, clue)
+
+        // creating the input stream opens the file
+        {
+          val bytes = buf.createInputStream().readAllBytes()
+          verify(mfs, times(1)).open(mc.any[Path]())
+          assert(bytes.mkString(",") === data.slice(offset, offset + length).mkString(","), clue)
+        }
+
+        // getting a NIO ByteBuffer opens the file again
+        {
+          val bytes = buf.nioByteBuffer().array()
+          verify(mfs, times(2)).open(mc.any[Path]())
+          assert(bytes.mkString(",") === data.slice(offset, offset + length).mkString(","), clue)
+        }
+
+        // getting a Netty ByteBufs opens the file again and again
+        assert(buf.convertToNetty().asInstanceOf[ByteBuf].release() === length > 0, clue)
+        verify(mfs, times(3)).open(mc.any[Path]())
+        assert(buf.convertToNettyForSsl().asInstanceOf[ByteBuf].release() === length > 0, clue)
+        verify(mfs, times(4)).open(mc.any[Path]())
+      }
+    }
+  }
+
   test("SPARK-34142: fallback storage API - cleanUp app") {
     withTempDir { dir =>
       Seq(true, false).foreach { cleanUp =>
@@ -372,6 +415,7 @@ class FallbackStorageSuite extends SparkFunSuite with LocalSparkContext {
     }
   }
 }
+
 class ReadPartialInputStream(val in: FSDataInputStream) extends InputStream
   with Seekable with PositionedReadable {
   override def read: Int = in.read

core/src/test/scala/org/apache/spark/storage/ShuffleBlockFetcherIteratorSuite.scala

@@ -32,7 +32,7 @@ import scala.concurrent.Future
 import io.netty.util.internal.OutOfDirectMemoryError
 import org.apache.logging.log4j.Level
 import org.mockito.ArgumentMatchers.{any, eq => meq}
-import org.mockito.Mockito.{doThrow, mock, times, verify, when}
+import org.mockito.Mockito.{doThrow, mock, never, times, verify, when}
 import org.mockito.invocation.InvocationOnMock
 import org.mockito.stubbing.Answer
 import org.roaringbitmap.RoaringBitmap
@@ -300,11 +300,11 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
     }
   }
 
-  test("successful 3 local + 4 host local + 2 remote reads") {
+  test("successful 3 local + 4 host local + 2 remote + 2 fallback storage reads") {
     val blockManager = createMockBlockManager()
-    val localBmId = blockManager.blockManagerId
 
     // Make sure blockManager.getBlockData would return the blocks
+    val localBmId = blockManager.blockManagerId
     val localBlocks = Map[BlockId, ManagedBuffer](
       ShuffleBlockId(0, 0, 0) -> createMockManagedBuffer(),
       ShuffleBlockId(0, 1, 0) -> createMockManagedBuffer(),
@@ -334,19 +334,37 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
     // returning local dir for hostLocalBmId
     initHostLocalDirManager(blockManager, hostLocalDirs)
 
+    // Make sure fallback storage blocks would return
+    val fallbackBmId = FallbackStorage.FALLBACK_BLOCK_MANAGER_ID
+    val fallbackBlocks = Map[BlockId, ManagedBuffer](
+      ShuffleBlockId(0, 9, 0) -> createMockManagedBuffer(),
+      ShuffleBlockId(0, 10, 0) -> createMockManagedBuffer())
+    fallbackBlocks.foreach { case (blockId, buf) =>
+      doReturn(buf).when(blockManager).getLocalBlockData(meq(blockId))
+    }
+
     val iterator = createShuffleBlockIteratorWithDefaults(
       Map(
         localBmId -> toBlockList(localBlocks.keys, 1L, 0),
         remoteBmId -> toBlockList(remoteBlocks.keys, 1L, 1),
-        hostLocalBmId -> toBlockList(hostLocalBlocks.keys, 1L, 1)
+        hostLocalBmId -> toBlockList(hostLocalBlocks.keys, 1L, 1),
+        fallbackBmId -> toBlockList(fallbackBlocks.keys, 1L, 1)
       ),
       blockManager = Some(blockManager)
     )
 
-    // 3 local blocks fetched in initialization
-    verify(blockManager, times(3)).getLocalBlockData(any())
+    // 3 local blocks and 2 fallback blocks fetched in initialization
+    verify(blockManager, times(3 + 2)).getLocalBlockData(any())
+
+    // SPARK-55469: but buffer data have never been materialized
+    fallbackBlocks.values.foreach { mockBuf =>
+      verify(mockBuf, never()).nioByteBuffer()
+      verify(mockBuf, never()).createInputStream()
+      verify(mockBuf, never()).convertToNetty()
+      verify(mockBuf, never()).convertToNettyForSsl()
+    }
 
-    val allBlocks = localBlocks ++ remoteBlocks ++ hostLocalBlocks
+    val allBlocks = localBlocks ++ remoteBlocks ++ hostLocalBlocks ++ fallbackBlocks
     for (i <- 0 until allBlocks.size) {
       assert(iterator.hasNext,
         s"iterator should have ${allBlocks.size} elements but actually has $i elements")
@@ -356,14 +374,23 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
       val mockBuf = allBlocks(blockId)
       verifyBufferRelease(mockBuf, inputStream)
     }
+    assert(!iterator.hasNext)
 
     // 4 host-local locks fetched
     verify(blockManager, times(4))
       .getHostLocalShuffleData(any(), meq(Array("local-dir")))
 
-    // 2 remote blocks are read from the same block manager
+    // 2 remote blocks are read from the same block manager in one fetch
     verifyFetchBlocksInvocationCount(1)
     assert(blockManager.hostLocalDirManager.get.getCachedHostLocalDirs.size === 1)
+
+    // SPARK-55469: fallback buffer data have been materialized once
+    fallbackBlocks.values.foreach { mockBuf =>
+      verify(mockBuf, never()).nioByteBuffer()
+      verify(mockBuf, times(1)).createInputStream()
+      verify(mockBuf, never()).convertToNetty()
+      verify(mockBuf, never()).convertToNettyForSsl()
+    }
   }
 
   test("error during accessing host local dirs for executors") {
@@ -451,10 +478,12 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
     assert(!iterator.hasNext)
   }
 
-  test("fetch continuous blocks in batch successful 3 local + 4 host local + 2 remote reads") {
+  test("fetch continuous blocks in batch successful 3 local + 4 host local + 2 remote + " +
+    "2 fallback storage reads") {
     val blockManager = createMockBlockManager()
-    val localBmId = blockManager.blockManagerId
+
     // Make sure blockManager.getBlockData would return the merged block
+    val localBmId = blockManager.blockManagerId
     val localBlocks = Seq[BlockId](
       ShuffleBlockId(0, 0, 0),
       ShuffleBlockId(0, 0, 1),
@@ -465,6 +494,17 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
       doReturn(buf).when(blockManager).getLocalBlockData(meq(blockId))
     }
 
+    // Make sure fallback storage would return the merged block
+    val fallbackBmId = FallbackStorage.FALLBACK_BLOCK_MANAGER_ID
+    val fallbackBlocks = Seq[BlockId](
+      ShuffleBlockId(0, 1, 0),
+      ShuffleBlockId(0, 1, 1))
+    val mergedFallbackBlocks = Map[BlockId, ManagedBuffer](
+      ShuffleBlockBatchId(0, 1, 0, 2) -> createMockManagedBuffer())
+    mergedFallbackBlocks.foreach { case (blockId, buf) =>
+      doReturn(buf).when(blockManager).getLocalBlockData(meq(blockId))
+    }
+
     // Make sure remote blocks would return the merged block
     val remoteBmId = BlockManagerId("test-client-1", "test-client-1", 2)
     val remoteBlocks = Seq[BlockId](
@@ -496,30 +536,49 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
     val iterator = createShuffleBlockIteratorWithDefaults(
       Map(
         localBmId -> toBlockList(localBlocks, 1L, 0),
+        fallbackBmId -> toBlockList(fallbackBlocks, 1L, 1),
         remoteBmId -> toBlockList(remoteBlocks, 1L, 1),
         hostLocalBmId -> toBlockList(hostLocalBlocks.keys, 1L, 1)
       ),
       blockManager = Some(blockManager),
       doBatchFetch = true
     )
 
-    // 3 local blocks batch fetched in initialization
-    verify(blockManager, times(1)).getLocalBlockData(any())
+    // 1 local merge block and 1 fallback merge block fetched in initialization
+    verify(blockManager, times(1 + 1)).getLocalBlockData(any())
 
-    val allBlocks = mergedLocalBlocks ++ mergedRemoteBlocks ++ mergedHostLocalBlocks
-    for (i <- 0 until 3) {
-      assert(iterator.hasNext, s"iterator should have 3 elements but actually has $i elements")
+    // SPARK-55469: but buffer data have never been materialized
+    mergedFallbackBlocks.values.foreach { mockBuf =>
+      verify(mockBuf, never()).nioByteBuffer()
+      verify(mockBuf, never()).createInputStream()
+      verify(mockBuf, never()).convertToNetty()
+      verify(mockBuf, never()).convertToNettyForSsl()
+    }
+
+    val allBlocks = mergedLocalBlocks ++ mergedFallbackBlocks ++ mergedRemoteBlocks ++
+      mergedHostLocalBlocks
+    for (i <- 0 until 4) {
+      assert(iterator.hasNext, s"iterator should have 4 elements but actually has $i elements")
       val (blockId, inputStream) = iterator.next()
       verifyFetchBlocksInvocationCount(1)
       // Make sure we release buffers when a wrapped input stream is closed.
       val mockBuf = allBlocks(blockId)
       verifyBufferRelease(mockBuf, inputStream)
     }
+    assert(!iterator.hasNext)
 
-    // 4 host-local locks fetched
+    // 1 merged host-local locks fetched
     verify(blockManager, times(1))
       .getHostLocalShuffleData(any(), meq(Array("local-dir")))
 
+    // SPARK-55469: merged fallback buffer data have been materialized once
+    mergedFallbackBlocks.values.foreach { mockBuf =>
+      verify(mockBuf, never()).nioByteBuffer()
+      verify(mockBuf, times(1)).createInputStream()
+      verify(mockBuf, never()).convertToNetty()
+      verify(mockBuf, never()).convertToNettyForSsl()
+    }
+
     assert(blockManager.hostLocalDirManager.get.getCachedHostLocalDirs.size === 1)
   }
 
@@ -1051,6 +1110,7 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
       val mockBuf = remoteBlocks(blockId)
       verifyBufferRelease(mockBuf, inputStream)
     }
+    assert(!iterator.hasNext)
 
     // 1st fetch request (contains 1 block) would fail due to Netty OOM
     // 2nd fetch request retry the block of the 1st fetch request
@@ -1091,6 +1151,7 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
         val mockBuf = remoteBlocks(blockId)
         verifyBufferRelease(mockBuf, inputStream)
       }
+      assert(!iterator.hasNext)
 
       // 1st fetch request (contains 3 blocks) would fail on the someone block due to Netty OOM
       // but succeed for the remaining blocks
@@ -2037,9 +2098,9 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
 
   test("SPARK-52395: Fast fail when failed to get host local dirs") {
     val blockManager = createMockBlockManager()
-    val localBmId = blockManager.blockManagerId
 
     // Make sure blockManager.getBlockData would return the blocks
+    val localBmId = blockManager.blockManagerId
     val localBlocks = Map[BlockId, ManagedBuffer](
       ShuffleBlockId(0, 0, 0) -> createMockManagedBuffer(),
       ShuffleBlockId(0, 1, 0) -> createMockManagedBuffer())
@@ -2076,4 +2137,46 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
     assert(iterator.next()._1 === ShuffleBlockId(0, 1, 0))
     assert(!iterator.hasNext)
   }
+
+  test("Fast fail when failed to get fallback storage blocks") {
+    val blockManager = createMockBlockManager()
+
+    // Make sure blockManager.getBlockData would return the blocks
+    val localBmId = blockManager.blockManagerId
+    val localBlocks = Map[BlockId, ManagedBuffer](
+      ShuffleBlockId(0, 0, 0) -> createMockManagedBuffer(),
+      ShuffleBlockId(0, 1, 0) -> createMockManagedBuffer())
+    localBlocks.foreach { case (blockId, buf) =>
+      doReturn(buf).when(blockManager).getLocalBlockData(meq(blockId))
+    }
+
+    // Make sure fallback storage would return the blocks
+    val fallbackBmId = FallbackStorage.FALLBACK_BLOCK_MANAGER_ID
+    val fallbackBlocks = Map[BlockId, ManagedBuffer](
+      ShuffleBlockId(0, 2, 0) -> createMockManagedBuffer(),
+      ShuffleBlockId(0, 3, 0) -> createMockManagedBuffer())
+    fallbackBlocks.take(1).foreach { case (blockId, buf) =>
+      doReturn(buf).when(blockManager).getLocalBlockData(meq(blockId))
+    }
+    fallbackBlocks.takeRight(1).foreach { case (blockId, _) =>
+      doThrow(new RuntimeException("Cannot read from fallback storage"))
+        .when(blockManager).getLocalBlockData(meq(blockId))
+    }
+
+    val iterator = createShuffleBlockIteratorWithDefaults(
+      Map(
+        localBmId -> toBlockList(localBlocks.keys, 1L, 0),
+        fallbackBmId -> toBlockList(fallbackBlocks.keys, 1L, 1)
+      ),
+      blockManager = Some(blockManager)
+    )
+
+    // Fetch failure should be placed in the head of results, exception should be thrown for the
+    // 1st instance.
+    intercept[FetchFailedException] { iterator.next() }
+    assert(iterator.next()._1 === ShuffleBlockId(0, 0, 0))
+    assert(iterator.next()._1 === ShuffleBlockId(0, 1, 0))
+    assert(iterator.next()._1 === ShuffleBlockId(0, 2, 0))
+    assert(!iterator.hasNext)
+  }
 }