Addressed review comments

Author: otterc

core/src/main/scala/org/apache/spark/storage/PushBasedFetchHelper.scala

@@ -71,7 +71,7 @@ private class PushBasedFetchHelper(
   }
 
   /**
-   * Returns true if the address is of a local push-merged block. false otherwise.
+   * Returns true if the address is of a push-merged-local block. false otherwise.
    */
   def isLocalPushMergedBlockAddress(address: BlockManagerId): Boolean = {
     isPushMergedShuffleBlockAddress(address) && address.host == blockManager.blockManagerId.host
@@ -190,14 +190,14 @@ private class PushBasedFetchHelper(
     val cachedPushedMergedDirs = hostLocalDirManager.getCachedHostLocalDirsFor(
       SHUFFLE_MERGER_IDENTIFIER)
     if (cachedPushedMergedDirs.isDefined) {
-      logDebug(s"Fetch the local push-merged blocks with cached merged dirs: " +
+      logDebug(s"Fetch the push-merged-local blocks with cached merged dirs: " +
         s"${cachedPushedMergedDirs.get.mkString(", ")}")
       pushMergedLocalBlocks.foreach { blockId =>
         fetchPushMergedLocalBlock(blockId, cachedPushedMergedDirs.get,
           localShuffleMergerBlockMgrId)
       }
     } else {
-      logDebug(s"Asynchronous fetch the local push-merged blocks without cached merged dirs")
+      logDebug(s"Asynchronous fetch the push-merged-local blocks without cached merged dirs")
       hostLocalDirManager.getHostLocalDirs(localShuffleMergerBlockMgrId.host,
         localShuffleMergerBlockMgrId.port, Array(SHUFFLE_MERGER_IDENTIFIER)) {
         case Success(dirs) =>
@@ -211,9 +211,9 @@ private class PushBasedFetchHelper(
                 localShuffleMergerBlockMgrId)
           }
         case Failure(throwable) =>
-          // If we see an exception with getting the local dirs for local push-merged blocks,
+          // If we see an exception with getting the local dirs for push-merged-local blocks,
           // we fallback to fetch the original blocks. We do not report block fetch failure.
-          logWarning(s"Error while fetching the merged dirs for local push-merged " +
+          logWarning(s"Error while fetching the merged dirs for push-merged-local " +
             s"blocks: ${pushMergedLocalBlocks.mkString(", ")}. Fetch the original blocks instead",
             throwable)
           pushMergedLocalBlocks.foreach {
@@ -226,7 +226,7 @@ private class PushBasedFetchHelper(
   }
 
   /**
-   * Fetch a single local push-merged block generated. This can also be executed by the task thread
+   * Fetch a single push-merged-local block generated. This can also be executed by the task thread
    * as well as the netty thread.
    * @param blockId ShuffleBlockId to be fetched
    * @param localDirs Local directories where the push-merged shuffle files are stored
@@ -244,10 +244,10 @@ private class PushBasedFetchHelper(
         localDirs))
     } catch {
       case e: Exception =>
-        // If we see an exception with reading a local push-merged meta, we fallback to
+        // If we see an exception with reading a push-merged-local meta, we fallback to
         // fetch the original blocks. We do not report block fetch failure
         // and will continue with the remaining local block read.
-        logWarning(s"Error occurred while fetching local push-merged meta, " +
+        logWarning(s"Error occurred while fetching push-merged-local meta, " +
           s"prepare to fetch the original blocks", e)
         iterator.addToResultsQueue(
           FallbackOnPushMergedFailureResult(blockId, blockManagerId, 0, isNetworkReqDone = false))
@@ -270,7 +270,7 @@ private class PushBasedFetchHelper(
    * finds more push-merged requests to remote and again updates it with additional requests for
    * original blocks.
    * The fallback happens when:
-   * 1. There is an exception while creating shuffle chunks from local push-merged shuffle block.
+   * 1. There is an exception while creating shuffle chunks from push-merged-local shuffle block.
    *    See fetchLocalBlock.
    * 2. There is a failure when fetching remote shuffle chunks.
    * 3. There is a failure when processing SuccessFetchResult which is for a shuffle chunk
@@ -286,7 +286,7 @@ private class PushBasedFetchHelper(
     val fallbackBlocksByAddr: Iterator[(BlockManagerId, Seq[(BlockId, Long, Int)])] =
       blockId match {
         case shuffleBlockId: ShuffleBlockId =>
-          iterator.incrementNumBlocksToFetch(-1)
+          iterator.decreaseNumBlocksToFetch(1)
           mapOutputTracker.getMapSizesForMergeResult(
             shuffleBlockId.shuffleId, shuffleBlockId.reduceId)
         case _ =>
@@ -311,7 +311,7 @@ private class PushBasedFetchHelper(
             // These blocks were added to numBlocksToFetch so we increment numBlocksProcessed
             blocksProcessed += pendingShuffleChunks.size
           }
-          iterator.incrementNumBlocksToFetch(-blocksProcessed)
+          iterator.decreaseNumBlocksToFetch(blocksProcessed)
           mapOutputTracker.getMapSizesForMergeResult(
             shuffleChunkId.shuffleId, shuffleChunkId.reduceId, chunkBitmap)
       }

core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala

@@ -60,7 +60,7 @@ import org.apache.spark.util.{CompletionIterator, TaskCompletionListener, Utils}
  *                        Note that zero-sized blocks are already excluded, which happened in
  *                        [[org.apache.spark.MapOutputTracker.convertMapStatuses]].
  * @param mapOutputTracker [[MapOutputTracker]] for falling back to fetching the original blocks if
- *                        we fail to fetch shuffle chunks when push based shuffle is enabled.
+ *                         we fail to fetch shuffle chunks when push based shuffle is enabled.
  * @param streamWrapper A function to wrap the returned input stream.
  * @param maxBytesInFlight max size (in bytes) of remote blocks to fetch at any given point.
  * @param maxReqsInFlight max number of remote requests to fetch blocks at any given point.
@@ -371,9 +371,9 @@ final class ShuffleBlockFetcherIterator(
     // blocks.Remote blocks are further split into FetchRequests of size at most maxBytesInFlight
     // in order to limit the amount of data in flight
     val collectedRemoteRequests = new ArrayBuffer[FetchRequest]
-    val hostLocalBlocksCurrentIteration = mutable.LinkedHashSet[(BlockId, Int)]()
     var localBlockBytes = 0L
     var hostLocalBlockBytes = 0L
+    var numHostLocalBlocks = 0
     var pushMergedLocalBlockBytes = 0L
     val prevNumBlocksToFetch = numBlocksToFetch
 
@@ -404,7 +404,7 @@ final class ShuffleBlockFetcherIterator(
         val blocksForAddress =
           mergedBlockInfos.map(info => (info.blockId, info.size, info.mapIndex))
         hostLocalBlocksByExecutor += address -> blocksForAddress
-        hostLocalBlocksCurrentIteration ++= blocksForAddress.map(info => (info._1, info._3))
+        numHostLocalBlocks += blocksForAddress.size
         hostLocalBlockBytes += mergedBlockInfos.map(_.size).sum
       } else {
         val (_, timeCost) = Utils.timeTakenMs[Unit] {
@@ -419,21 +419,22 @@ final class ShuffleBlockFetcherIterator(
       pushMergedLocalBlockBytes
     val blocksToFetchCurrentIteration = numBlocksToFetch - prevNumBlocksToFetch
     assert(blocksToFetchCurrentIteration == localBlocks.size +
-      hostLocalBlocksCurrentIteration.size + numRemoteBlocks + pushMergedLocalBlocks.size,
-      s"The number of non-empty blocks $blocksToFetchCurrentIteration doesn't equal to " +
-        s"the number of local blocks ${localBlocks.size} + " +
-        s"the number of host-local blocks ${hostLocalBlocksCurrentIteration.size} " +
+      numHostLocalBlocks + numRemoteBlocks + pushMergedLocalBlocks.size,
+        s"The number of non-empty blocks $blocksToFetchCurrentIteration doesn't equal to the sum " +
+        s"of the number of local blocks ${localBlocks.size} + " +
+        s"the number of host-local blocks ${numHostLocalBlocks} " +
         s"the number of push-merged-local blocks ${pushMergedLocalBlocks.size} " +
         s"+ the number of remote blocks ${numRemoteBlocks} ")
     logInfo(s"Getting $blocksToFetchCurrentIteration " +
       s"(${Utils.bytesToString(totalBytes)}) non-empty blocks including " +
       s"${localBlocks.size} (${Utils.bytesToString(localBlockBytes)}) local and " +
-      s"${hostLocalBlocksCurrentIteration.size} (${Utils.bytesToString(hostLocalBlockBytes)}) " +
+      s"${numHostLocalBlocks} (${Utils.bytesToString(hostLocalBlockBytes)}) " +
       s"host-local and ${pushMergedLocalBlocks.size} " +
       s"(${Utils.bytesToString(pushMergedLocalBlockBytes)}) " +
-      s"local push-merged and $numRemoteBlocks (${Utils.bytesToString(remoteBlockBytes)}) " +
+      s"push-merged-local and $numRemoteBlocks (${Utils.bytesToString(remoteBlockBytes)}) " +
       s"remote blocks")
-    this.hostLocalBlocks ++= hostLocalBlocksCurrentIteration
+    this.hostLocalBlocks ++= hostLocalBlocksByExecutor.values
+      .flatMap { infos => infos.map(info => (info._1, info._3)) }
     collectedRemoteRequests
   }
 
@@ -883,9 +884,9 @@ final class ShuffleBlockFetcherIterator(
           // We get this result in 3 cases:
           // 1. Failure to fetch the data of a remote shuffle chunk. In this case, the
           //    blockId is a ShuffleBlockChunkId.
-          // 2. Failure to read the local push-merged meta. In this case, the blockId is
+          // 2. Failure to read the push-merged-local meta. In this case, the blockId is
           //    ShuffleBlockId.
-          // 3. Failure to get the local push-merged directories from the ESS. In this case, the
+          // 3. Failure to get the push-merged-local directories from the ESS. In this case, the
           //    blockId is ShuffleBlockId.
           if (pushBasedFetchHelper.isRemotePushMergedBlockAddress(address)) {
             numBlocksInFlightPerAddress(address) = numBlocksInFlightPerAddress(address) - 1
@@ -900,8 +901,8 @@ final class ShuffleBlockFetcherIterator(
           // a SuccessFetchResult or a FailureFetchResult.
           result = null
 
-          case PushMergedLocalMetaFetchResult(shuffleId, reduceId, bitmaps, localDirs, _) =>
-            // Fetch local push-merged shuffle block data as multiple shuffle chunks
+          case PushMergedLocalMetaFetchResult(shuffleId, reduceId, bitmaps, localDirs) =>
+            // Fetch push-merged-local shuffle block data as multiple shuffle chunks
             val shuffleBlockId = ShuffleBlockId(shuffleId, SHUFFLE_PUSH_MAP_ID, reduceId)
             try {
               val bufs: Seq[ManagedBuffer] = blockManager.getLocalMergedBlockData(shuffleBlockId,
@@ -922,17 +923,17 @@ final class ShuffleBlockFetcherIterator(
               }
             } catch {
               case e: Exception =>
-                // If we see an exception with reading local push-merged data, we fallback to
-                // fetch the original blocks. We do not report block fetch failure
+                // If we see an exception with reading push-merged-local index file, we fallback
+                // to fetch the original blocks. We do not report block fetch failure
                 // and will continue with the remaining local block read.
-                logWarning(s"Error occurred while fetching local push-merged data, " +
+                logWarning(s"Error occurred while reading push-merged-local index, " +
                   s"prepare to fetch the original blocks", e)
                 pushBasedFetchHelper.initiateFallbackFetchForPushMergedBlock(
                   shuffleBlockId, pushBasedFetchHelper.localShuffleMergerBlockMgrId)
             }
             result = null
 
-        case PushMergedRemoteMetaFetchResult(shuffleId, reduceId, blockSize, bitmaps, address, _) =>
+        case PushMergedRemoteMetaFetchResult(shuffleId, reduceId, blockSize, bitmaps, address) =>
           // The original meta request is processed so we decrease numBlocksToFetch and
           // numBlocksInFlightPerAddress by 1. We will collect new shuffle chunks request and the
           // count of this is added to numBlocksToFetch in collectFetchReqsFromMergedBlocks.
@@ -946,7 +947,7 @@ final class ShuffleBlockFetcherIterator(
           // Set result to null to force another iteration.
           result = null
 
-        case PushMergedRemoteMetaFailedFetchResult(shuffleId, reduceId, address, _) =>
+        case PushMergedRemoteMetaFailedFetchResult(shuffleId, reduceId, address) =>
           // The original meta request failed so we decrease numBlocksInFlightPerAddress by 1.
           numBlocksInFlightPerAddress(address) = numBlocksInFlightPerAddress(address) - 1
           // If we fail to fetch the meta of a push-merged block, we fall back to fetching the
@@ -1071,8 +1072,8 @@ final class ShuffleBlockFetcherIterator(
     results.put(result)
   }
 
-  private[storage] def incrementNumBlocksToFetch(moreBlocksToFetch: Int): Unit = {
-    numBlocksToFetch += moreBlocksToFetch
+  private[storage] def decreaseNumBlocksToFetch(blocksFetched: Int): Unit = {
+    numBlocksToFetch -= blocksFetched
   }
 
   /**
@@ -1091,7 +1092,7 @@ final class ShuffleBlockFetcherIterator(
       originalLocalBlocks, originalHostLocalBlocksByExecutor, originalMergedLocalBlocks)
     // Add the remote requests into our queue in a random order
     fetchRequests ++= Utils.randomize(originalRemoteReqs)
-    logInfo(s"Started ${originalRemoteReqs.size} fallback remote requests for push-merged")
+    logInfo(s"Created ${originalRemoteReqs.size} fallback remote requests for push-merged")
     // fetch all the fallback blocks that are local.
     fetchLocalBlocks(originalLocalBlocks)
     // Merged local blocks should be empty during fallback
@@ -1246,12 +1247,6 @@ object ShuffleBlockFetcherIterator {
     }
   }
 
-  /**
-   * Dummy shuffle block id to fill into [[PushMergedRemoteMetaFetchResult]] and
-   * [[PushMergedRemoteMetaFailedFetchResult]], to match the [[FetchResult]] trait.
-   */
-  private val DUMMY_SHUFFLE_BLOCK_ID = ShuffleBlockId(-1, -1, -1)
-
   /**
    * This function is used to merged blocks when doBatchFetch is true. Blocks which have the
    * same `mapId` can be merged into one block batch. The block batch is specified by a range
@@ -1436,8 +1431,7 @@ object ShuffleBlockFetcherIterator {
       reduceId: Int,
       blockSize: Long,
       bitmaps: Array[RoaringBitmap],
-      address: BlockManagerId,
-      blockId: BlockId = DUMMY_SHUFFLE_BLOCK_ID) extends FetchResult
+      address: BlockManagerId) extends FetchResult
 
   /**
    * Result of a failure while fetching the meta information for a remote push-merged block.
@@ -1449,11 +1443,10 @@ object ShuffleBlockFetcherIterator {
   private[storage] case class PushMergedRemoteMetaFailedFetchResult(
       shuffleId: Int,
       reduceId: Int,
-      address: BlockManagerId,
-      blockId: BlockId = DUMMY_SHUFFLE_BLOCK_ID) extends FetchResult
+      address: BlockManagerId) extends FetchResult
 
   /**
-   * Result of a successful fetch of meta information for a local push-merged block.
+   * Result of a successful fetch of meta information for a push-merged-local block.
    *
    * @param shuffleId shuffle id.
    * @param reduceId reduce id.
@@ -1464,6 +1457,5 @@ object ShuffleBlockFetcherIterator {
       shuffleId: Int,
       reduceId: Int,
       bitmaps: Array[RoaringBitmap],
-      localDirs: Array[String],
-      blockId: BlockId = DUMMY_SHUFFLE_BLOCK_ID) extends FetchResult
+      localDirs: Array[String]) extends FetchResult
 }

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

@@ -1263,7 +1263,7 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite with PrivateMethodT
     assert(!iterator.hasNext)
   }
 
-  test("SPARK-32922: failure to fetch local push-merged meta should fallback to fetch " +
+  test("SPARK-32922: failure to fetch push-merged-local meta should fallback to fetch " +
     "original shuffle blocks") {
     val blockManager = mock(classOf[BlockManager])
     val localDirs = Array("testPath1", "testPath2")
@@ -1276,7 +1276,7 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite with PrivateMethodT
     verifyLocalBlocksFromFallback(iterator)
   }
 
-  test("SPARK-32922: failure to reading chunkBitmaps of local push-merged meta should " +
+  test("SPARK-32922: failure to reading chunkBitmaps of push-merged-local meta should " +
     "fallback to original shuffle blocks") {
     val blockManager = mock(classOf[BlockManager])
     val localDirs = Array("local-dir")
@@ -1288,7 +1288,7 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite with PrivateMethodT
     verifyLocalBlocksFromFallback(iterator)
   }
 
-  test("SPARK-32922: failure to fetch local push-merged data should fallback to fetch " +
+  test("SPARK-32922: failure to fetch push-merged-local data should fallback to fetch " +
     "original shuffle blocks") {
     val blockManager = mock(classOf[BlockManager])
     val localDirs = Array("testPath1", "testPath2")
@@ -1301,8 +1301,8 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite with PrivateMethodT
     verifyLocalBlocksFromFallback(iterator)
   }
 
-  test("SPARK-32922: failure to fetch local push-merged meta of a single merged block " +
-    "should not drop the fetch of other local push-merged blocks") {
+  test("SPARK-32922: failure to fetch push-merged-local meta of a single merged block " +
+    "should not drop the fetch of other push-merged-local blocks") {
     val blockManager = mock(classOf[BlockManager])
     val localDirs = Array("testPath1", "testPath2")
     prepareForFallbackToLocalBlocks(
@@ -1375,7 +1375,7 @@ class ShuffleBlockFetcherIteratorSuite extends SparkFunSuite with PrivateMethodT
     intercept[FetchFailedException] { iterator.next() }
   }
 
-  test("SPARK-32922: failure to fetch local push-merged block should fallback to fetch " +
+  test("SPARK-32922: failure to fetch push-merged-local block should fallback to fetch " +
     "original shuffle blocks which contain host-local blocks") {
     val blockManager = mock(classOf[BlockManager])
     // BlockManagerId from another executor on the same host