address comments, and fix the test error

ConeyLiu · ConeyLiu · commit a3e979f5e421 · 2017-05-05T13:47:29.000+08:00
diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
@@ -617,10 +617,14 @@ private[spark] class BlockManager(
   private def getLocations(blockId: BlockId): Seq[BlockManagerId] = {
     val locs = Random.shuffle(master.getLocations(blockId))
     val (preferredLocs, otherLocs) = locs.partition { loc => blockManagerId.host == loc.host }
-    val (sameRackLocs, differentRackLocs) = otherLocs.partition {
-      loc => blockManagerId.topologyInfo == loc.topologyInfo
+    blockManagerId.topologyInfo match {
+      case None => preferredLocs ++ otherLocs
+      case Some(_) =>
+        val (sameRackLocs, differentRackLocs) = otherLocs.partition {
+          loc => blockManagerId.topologyInfo == loc.topologyInfo
+        }
+        preferredLocs ++ sameRackLocs ++ differentRackLocs
     }
-    preferredLocs ++ sameRackLocs ++ differentRackLocs
   }
 
   /**
diff --git a/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala b/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala
@@ -497,7 +497,7 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with BeforeAndAfterE
   }
 
   test("optimize a location order of blocks without topology information") {
-    val localHost = Utils.localHostName()
+    val localHost = "localhost"
     val otherHost = "otherHost"
     val bmMaster = mock(classOf[BlockManagerMaster])
     val bmId1 = BlockManagerId("id1", localHost, 1)
@@ -512,7 +512,7 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with BeforeAndAfterE
   }
 
   test("optimize a location order of blocks with topology information") {
-    val localHost = Utils.localHostName()
+    val localHost = "localhost"
     val otherHost = "otherHost"
     val localRack = "localRack"
     val otherRack = "otherRack"
@@ -527,6 +527,8 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with BeforeAndAfterE
       .thenReturn(Seq(bmId1, bmId2, bmId5, bmId3, bmId4))
 
     val blockManager = makeBlockManager(128, "exec", bmMaster)
+    blockManager.blockManagerId =
+      BlockManagerId(SparkContext.DRIVER_IDENTIFIER, localHost, 1, Some(localRack))
     val getLocations = PrivateMethod[Seq[BlockManagerId]]('getLocations)
     val locations = blockManager invokePrivate getLocations(BroadcastBlockId(0))
     assert(locations.map(_.host) === Seq(localHost, localHost, otherHost, otherHost, otherHost))
