Update unit tests

Removes "Fast fail when failed to get fallback storage blocks" as fallback storage blocks
are fetched concurrently and fast fail is not guaranteed any more.

Author: EnricoMi

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

@@ -19,6 +19,7 @@ package org.apache.spark.storage
 
 import java.io._
 import java.nio.ByteBuffer
+import java.util
 import java.util.UUID
 import java.util.concurrent.{CompletableFuture, Semaphore}
 import java.util.zip.CheckedInputStream
@@ -153,6 +154,9 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
     val in = mock(classOf[InputStream])
     when(in.read(any())).thenReturn(1)
     when(in.read(any(), any(), any())).thenReturn(1)
+    val buf = ByteBuffer.allocate(size)
+    util.Arrays.fill(buf.array(), 1.byteValue)
+    when(mockManagedBuffer.nioByteBuffer()).thenReturn(buf)
     when(mockManagedBuffer.createInputStream()).thenReturn(in)
     when(mockManagedBuffer.size()).thenReturn(size)
     mockManagedBuffer
@@ -342,7 +346,7 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
       ShuffleBlockId(0, 9, 0) -> createMockManagedBuffer(),
       ShuffleBlockId(0, 10, 0) -> createMockManagedBuffer())
     fallbackBlocks.foreach { case (blockId, buf) =>
-      doReturn(buf).when(blockManager).getLocalBlockData(meq(blockId))
+      doReturn(buf).when(blockManager).getFallbackStorageBlockData(meq(blockId))
     }
 
     val iterator = createShuffleBlockIteratorWithDefaults(
@@ -355,9 +359,11 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
       blockManager = Some(blockManager)
     )
 
-    // 3 local blocks and 2 fallback blocks fetched in initialization
-    verify(blockManager, times(3 + 2)).getLocalBlockData(any())
+    // 3 local blocks fetched in initialization
+    verify(blockManager, times(3)).getLocalBlockData(any())
 
+    // 2 fallback storage blocks fetched in initialization
+    verify(blockManager, times(2)).getFallbackStorageBlockData(any())
     // SPARK-55469: but buffer data have never been materialized
     fallbackBlocks.values.foreach { mockBuf =>
       verify(mockBuf, never()).nioByteBuffer()
@@ -374,7 +380,9 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
 
       // Make sure we release buffers when a wrapped input stream is closed.
       val mockBuf = allBlocks(blockId)
-      verifyBufferRelease(mockBuf, inputStream)
+      if (!fallbackBlocks.contains(blockId)) {
+        verifyBufferRelease(mockBuf, inputStream)
+      }
     }
     assert(!iterator.hasNext)
 
@@ -388,8 +396,8 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
 
     // SPARK-55469: fallback buffer data have been materialized once
     fallbackBlocks.values.foreach { mockBuf =>
-      verify(mockBuf, never()).nioByteBuffer()
-      verify(mockBuf, times(1)).createInputStream()
+      verify(mockBuf, times(1)).nioByteBuffer()
+      verify(mockBuf, never()).createInputStream()
       verify(mockBuf, never()).convertToNetty()
       verify(mockBuf, never()).convertToNettyForSsl()
     }
@@ -504,7 +512,7 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
     val mergedFallbackBlocks = Map[BlockId, ManagedBuffer](
       ShuffleBlockBatchId(0, 1, 0, 2) -> createMockManagedBuffer())
     mergedFallbackBlocks.foreach { case (blockId, buf) =>
-      doReturn(buf).when(blockManager).getLocalBlockData(meq(blockId))
+      doReturn(buf).when(blockManager).getFallbackStorageBlockData(meq(blockId))
     }
 
     // Make sure remote blocks would return the merged block
@@ -546,9 +554,11 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
       doBatchFetch = true
     )
 
-    // 1 local merge block and 1 fallback merge block fetched in initialization
-    verify(blockManager, times(1 + 1)).getLocalBlockData(any())
+    // 1 local merge block fetched in initialization
+    verify(blockManager, times(1)).getLocalBlockData(any())
 
+    // 1 fallback merge block fetched in initialization
+    verify(blockManager, times(1)).getFallbackStorageBlockData(any())
     // SPARK-55469: but buffer data have never been materialized
     mergedFallbackBlocks.values.foreach { mockBuf =>
       verify(mockBuf, never()).nioByteBuffer()
@@ -565,23 +575,27 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite {
       verifyFetchBlocksInvocationCount(1)
       // Make sure we release buffers when a wrapped input stream is closed.
       val mockBuf = allBlocks(blockId)
-      verifyBufferRelease(mockBuf, inputStream)
+      if (!mergedFallbackBlocks.contains(blockId)) {
+        verifyBufferRelease(mockBuf, inputStream)
+      }
     }
     assert(!iterator.hasNext)
 
     // 1 merged host-local locks fetched
     verify(blockManager, times(1))
       .getHostLocalShuffleData(any(), meq(Array("local-dir")))
 
+    // 1 merged remote block is read from the same block manager
+    verifyFetchBlocksInvocationCount(1)
+    assert(blockManager.hostLocalDirManager.get.getCachedHostLocalDirs.size === 1)
+
     // 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, times(1)).nioByteBuffer()
+      verify(mockBuf, never()).createInputStream()
       verify(mockBuf, never()).convertToNetty()
       verify(mockBuf, never()).convertToNettyForSsl()
     }
-
-    assert(blockManager.hostLocalDirManager.get.getCachedHostLocalDirs.size === 1)
   }
 
   test("fetch continuous blocks in batch should respect maxBytesInFlight") {
@@ -2139,46 +2153,4 @@ 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)
-  }
 }