[SPARK-51638][CORE] Fix fetching the remote disk stored RDD blocks via the external shuffle service

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

@@ -863,36 +863,36 @@ class BlockManagerMasterEndpoint(
       blockId: BlockId,
       requesterHost: String): Option[BlockLocationsAndStatus] = {
     val allLocations = Option(blockLocations.get(blockId)).map(_.toSeq).getOrElse(Seq.empty)
-    val hostLocalLocations = allLocations.filter(bmId => bmId.host == requesterHost)
-
     val blockStatusWithBlockManagerId: Option[(BlockStatus, BlockManagerId)] =
       (if (externalShuffleServiceRddFetchEnabled) {
-         // if fetching RDD is enabled from the external shuffle service then first try to find
-         // the block in the external shuffle service of the same host
-         val location = hostLocalLocations.find(_.port == externalShuffleServicePort)
-         location
-           .flatMap(blockStatusByShuffleService.get(_).flatMap(_.get(blockId)))
-           .zip(location)
-       } else {
-         None
-       })
-        .orElse {
-          // if the block is not found via the external shuffle service trying to find it in the
-          // executors running on the same host and persisted on the disk
-          // using flatMap on iterators makes the transformation lazy
-          hostLocalLocations.iterator
-            .flatMap { bmId =>
-              blockManagerInfo.get(bmId).flatMap { blockInfo =>
-                blockInfo.getStatus(blockId).map((_, bmId))
-              }
-            }
-            .find(_._1.storageLevel.useDisk)
-        }
-        .orElse {
-          // if the block cannot be found in the same host search it in all the executors
-          val location = allLocations.headOption
-          location.flatMap(blockManagerInfo.get(_)).flatMap(_.getStatus(blockId)).zip(location)
+        // If fetching disk persisted RDD from the external shuffle service is enabled then first
+        // try to find the block in the external shuffle service preferring the one running on
+        // the same host. This search includes blocks stored on already killed executors as well.
+        val hostLocalLocations = allLocations.find { bmId =>
+          bmId.host == requesterHost && bmId.port == externalShuffleServicePort
         }
+        val location = hostLocalLocations
+          .orElse(allLocations.find(_.port == externalShuffleServicePort))
+        location
+          .flatMap(blockStatusByShuffleService.get(_).flatMap(_.get(blockId)))
+          .zip(location)
+      } else {
+        // trying to find it in the executors running on the same host and persisted on the disk
+        // Implementation detail: using flatMap on iterators makes the transformation lazy.
+        allLocations.find(_.host == requesterHost).iterator
+          .flatMap { bmId =>
+            blockManagerInfo.get(bmId).flatMap { blockInfo =>
+              blockInfo.getStatus(blockId).map((_, bmId))
+            }
+          }
+          .find(_._1.storageLevel.useDisk)
+      })
+      .orElse {
+        // if the block cannot be found in the same host as a disk stored block then extend the
+        // search to all active (not killed) executors and to all storage levels
+        val location = allLocations.headOption
+        location.flatMap(blockManagerInfo.get(_)).flatMap(_.getStatus(blockId)).zip(location)
+      }
     logDebug(s"Identified block: $blockStatusWithBlockManagerId")
     blockStatusWithBlockManagerId
       .map { case (blockStatus: BlockStatus, bmId: BlockManagerId) =>

core/src/test/scala/org/apache/spark/ExternalShuffleServiceSuite.scala

@@ -117,6 +117,7 @@ class ExternalShuffleServiceSuite extends ShuffleSuite with BeforeAndAfterAll wi
       .set(config.SHUFFLE_HOST_LOCAL_DISK_READING_ENABLED, true)
       .set(config.SHUFFLE_SERVICE_FETCH_RDD_ENABLED, true)
       .set(config.EXECUTOR_REMOVE_DELAY.key, "0s")
+      .set(config.DRIVER_BIND_ADDRESS.key, Utils.localHostName())
     sc = new SparkContext("local-cluster[1,1,1024]", "test", confWithRddFetchEnabled)
     sc.env.blockManager.externalShuffleServiceEnabled should equal(true)
     sc.env.blockManager.blockStoreClient.getClass should equal(classOf[ExternalBlockStoreClient])
@@ -173,6 +174,20 @@ class ExternalShuffleServiceSuite extends ShuffleSuite with BeforeAndAfterAll wi
           "external shuffle service port should be contained")
       }
 
+      val locationStatusForLocalHost =
+        sc.env.blockManager.master.getLocationsAndStatus(blockId, Utils.localHostName())
+      assert(locationStatusForLocalHost.isDefined)
+      assert(locationStatusForLocalHost.get.localDirs.isDefined)
+      assert(locationStatusForLocalHost.get.locations.head.executorId == "0")
+      assert(locationStatusForLocalHost.get.locations.head.host == Utils.localHostName())
+
+      val locationStatusForRemoteHost =
+        sc.env.blockManager.master.getLocationsAndStatus(blockId, "<invalid-host>")
+      assert(locationStatusForRemoteHost.isDefined)
+      assert(locationStatusForRemoteHost.get.localDirs.isEmpty)
+      assert(locationStatusForRemoteHost.get.locations.head.executorId == "0")
+      assert(locationStatusForRemoteHost.get.locations.head.host == Utils.localHostName())
+
       assert(sc.env.blockManager.getRemoteValues(blockId).isDefined)
 
       // test unpersist: as executors are killed the blocks will be removed via the shuffle service