Use hostPort as the receiver location rather than host; fix comments and unit tests

Author: zsxwing

streaming/src/main/scala/org/apache/spark/streaming/receiver/ReceiverSupervisor.scala

@@ -24,7 +24,7 @@ import scala.collection.mutable.ArrayBuffer
 import scala.concurrent._
 import scala.util.control.NonFatal
 
-import org.apache.spark.{Logging, SparkConf}
+import org.apache.spark.{SparkEnv, Logging, SparkConf}
 import org.apache.spark.storage.StreamBlockId
 import org.apache.spark.util.{Utils, ThreadUtils}
 
@@ -65,8 +65,6 @@ private[streaming] abstract class ReceiverSupervisor(
   /** State of the receiver */
   @volatile private[streaming] var receiverState = Initialized
 
-  protected val host = Utils.localHostName()
-
   /** Push a single data item to backend data store. */
   def pushSingle(data: Any)
 

streaming/src/main/scala/org/apache/spark/streaming/receiver/ReceiverSupervisorImpl.scala

@@ -46,6 +46,8 @@ private[streaming] class ReceiverSupervisorImpl(
     checkpointDirOption: Option[String]
   ) extends ReceiverSupervisor(receiver, env.conf) with Logging {
 
+  private val hostPort = SparkEnv.get.blockManager.blockManagerId.hostPort
+
   private val receivedBlockHandler: ReceivedBlockHandler = {
     if (WriteAheadLogUtils.enableReceiverLog(env.conf)) {
       if (checkpointDirOption.isEmpty) {
@@ -164,7 +166,7 @@ private[streaming] class ReceiverSupervisorImpl(
 
   override protected def onReceiverStart(): Boolean = {
     val msg = RegisterReceiver(
-      streamId, receiver.getClass.getSimpleName, host, endpoint)
+      streamId, receiver.getClass.getSimpleName, hostPort, endpoint)
     trackerEndpoint.askWithRetry[Boolean](msg)
   }
 

streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverSchedulingPolicy.scala

@@ -22,30 +22,6 @@ import scala.collection.mutable
 
 import org.apache.spark.streaming.receiver.Receiver
 
-/**
- * A ReceiverScheduler trying to balance executors' load. Here is the approach to schedule executors
- * for a receiver.
- * <ol>
- *   <li>
- *     If preferredLocation is set, preferredLocation should be one of the candidate executors.
- *   </li>
- *   <li>
- *     Every executor will be assigned to a weight according to the receivers running or scheduling
- *     on it.
- *     <ul>
- *       <li>
- *         If a receiver is running on an executor, it contributes 1.0 to the executor's weight.
- *       </li>
- *       <li>
- *         If a receiver is scheduled to an executor but has not yet run, it contributes
- *         `1.0 / #candidate_executors_of_this_receiver` to the executor's weight.</li>
- *     </ul>
- *     At last, if there are more than 3 idle executors (weight = 0), returns all idle executors.
- *     Otherwise, we only return 3 best options according to the weights.
- *   </li>
- * </ol>
- *
- */
 private[streaming] class ReceiverSchedulingPolicy {
 
   /**
@@ -61,7 +37,9 @@ private[streaming] class ReceiverSchedulingPolicy {
       return Map.empty
     }
 
-    require(executors.nonEmpty, "There is no executor up")
+    if (executors.isEmpty) {
+      return receivers.map(_.streamId -> Seq.empty).toMap
+    }
 
     val hostToExecutors = executors.groupBy(_.split(":")(0))
     val locations = new Array[mutable.ArrayBuffer[String]](receivers.length)
@@ -119,6 +97,28 @@ private[streaming] class ReceiverSchedulingPolicy {
    * Return a list of candidate executors to run the receiver. If the list is empty, the caller can
    * run this receiver in arbitrary executor.
    *
+   * This method tries to balance executors' load. Here is the approach to schedule executors
+   * for a receiver.
+   * <ol>
+   *   <li>
+   *     If preferredLocation is set, preferredLocation should be one of the candidate executors.
+   *   </li>
+   *   <li>
+   *     Every executor will be assigned to a weight according to the receivers running or
+   *     scheduling on it.
+   *     <ul>
+   *       <li>
+   *         If a receiver is running on an executor, it contributes 1.0 to the executor's weight.
+   *       </li>
+   *       <li>
+   *         If a receiver is scheduled to an executor but has not yet run, it contributes
+   *         `1.0 / #candidate_executors_of_this_receiver` to the executor's weight.</li>
+   *     </ul>
+   *     At last, if there are more than 3 idle executors (weight = 0), returns all idle executors.
+   *     Otherwise, we only return 3 best options according to the weights.
+   *   </li>
+   * </ol>
+   *
    * This method is called when a receiver is registering with ReceiverTracker or is restarting.
    */
   def rescheduleReceiver(

streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTracker.scala

@@ -25,7 +25,7 @@ import scala.language.existentials
 import scala.util.{Failure, Success}
 
 import org.apache.spark.streaming.util.WriteAheadLogUtils
-import org.apache.spark.{TaskContext, Logging, SparkEnv, SparkException}
+import org.apache.spark._
 import org.apache.spark.rdd.RDD
 import org.apache.spark.rpc._
 import org.apache.spark.streaming.{StreamingContext, Time}
@@ -47,7 +47,7 @@ private[streaming] sealed trait ReceiverTrackerMessage
 private[streaming] case class RegisterReceiver(
     streamId: Int,
     typ: String,
-    host: String,
+    hostPort: String,
     receiverEndpoint: RpcEndpointRef
   ) extends ReceiverTrackerMessage
 private[streaming] case class AddBlock(receivedBlockInfo: ReceivedBlockInfo)
@@ -130,12 +130,13 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
 
   /**
    * Track all receivers' information. The key is the receiver id, the value is the receiver info.
+   * It's only accessed in ReceiverTrackerEndpoint.
    */
   private val receiverTrackingInfos = new HashMap[Int, ReceiverTrackingInfo]
 
   /**
    * Store all preferred locations for all receivers. We need this information to schedule
-   * receivers
+   * receivers. It's only accessed in ReceiverTrackerEndpoint.
    */
   private val receiverPreferredLocations = new HashMap[Int, Option[String]]
 
@@ -227,7 +228,7 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
   private def registerReceiver(
       streamId: Int,
       typ: String,
-      host: String,
+      hostPort: String,
       receiverEndpoint: RpcEndpointRef,
       senderAddress: RpcAddress
     ): Boolean = {
@@ -237,19 +238,18 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
 
     if (isTrackerStopping || isTrackerStopped) {
       false
-    } else if (!ssc.sparkContext.isLocal && // We don't need to schedule it in the local mode
-      !scheduleReceiver(streamId).contains(host)) {
+    } else if (!scheduleReceiver(streamId).contains(hostPort)) {
       // Refuse it since it's scheduled to a wrong executor
       false
     } else {
       val name = s"${typ}-${streamId}"
-      val receiverInfo = ReceiverInfo(streamId, name, true, host)
+      val receiverInfo = ReceiverInfo(streamId, name, true, hostPort)
       receiverTrackingInfos.put(streamId,
         ReceiverTrackingInfo(
           streamId,
           ReceiverState.ACTIVE,
           scheduledLocations = None,
-          runningLocation = Some(host),
+          runningLocation = Some(hostPort),
           name = Some(name),
           endpoint = Some(receiverEndpoint)))
       listenerBus.post(StreamingListenerReceiverStarted(receiverInfo))
@@ -344,13 +344,13 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
   /**
    * Get the list of executors excluding driver
    */
-  private def getExecutors: List[String] = {
+  private def getExecutors: Seq[String] = {
     if (ssc.sc.isLocal) {
-      List("localhost")
+      Seq(ssc.sparkContext.env.blockManager.blockManagerId.hostPort)
     } else {
-      val executors = ssc.sparkContext.getExecutorMemoryStatus.map(_._1.split(":")(0)).toList
-      val driver = ssc.sparkContext.getConf.get("spark.driver.host")
-      executors.diff(List(driver))
+      ssc.sparkContext.env.blockManager.master.getMemoryStatus.filter { case (blockManagerId, _) =>
+        blockManagerId.executorId != SparkContext.DRIVER_IDENTIFIER // Ignore the driver location
+      }.map { case (blockManagerId, _) => blockManagerId.hostPort }.toSeq
     }
   }
 
@@ -383,7 +383,7 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
     runDummySparkJob()
 
     logInfo("Starting " + receivers.length + " receivers")
-    endpoint.send(StartAllReceivers)
+    endpoint.send(StartAllReceivers(receivers))
   }
 
   /** Check if tracker has been marked for starting */
@@ -429,9 +429,9 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
       // Remote messages
-      case RegisterReceiver(streamId, typ, host, receiverEndpoint) =>
+      case RegisterReceiver(streamId, typ, hostPort, receiverEndpoint) =>
         val successful =
-          registerReceiver(streamId, typ, host, receiverEndpoint, context.sender.address)
+          registerReceiver(streamId, typ, hostPort, receiverEndpoint, context.sender.address)
         context.reply(successful)
       case AddBlock(receivedBlockInfo) =>
         context.reply(addBlock(receivedBlockInfo))

streaming/src/test/scala/org/apache/spark/streaming/scheduler/ReceiverSchedulingPolicySuite.scala

@@ -27,21 +27,21 @@ class ReceiverSchedulingPolicySuite extends SparkFunSuite {
 
   val receiverSchedulingPolicy = new ReceiverSchedulingPolicy
 
-  test("empty executors") {
+  test("rescheduleReceiver: empty executors") {
     val scheduledLocations =
       receiverSchedulingPolicy.rescheduleReceiver(0, None, Map.empty, executors = Seq.empty)
     assert(scheduledLocations === Seq.empty)
   }
 
-  test("receiver preferredLocation") {
+  test("rescheduleReceiver: receiver preferredLocation") {
     val receiverTrackingInfoMap = Map(
       0 -> ReceiverTrackingInfo(0, ReceiverState.INACTIVE, None, None))
     val scheduledLocations = receiverSchedulingPolicy.rescheduleReceiver(
       0, Some("host1"), receiverTrackingInfoMap, executors = Seq("host2"))
     assert(scheduledLocations.toSet === Set("host1", "host2"))
   }
 
-  test("return all idle executors if more than 3 idle executors") {
+  test("rescheduleReceiver: return all idle executors if more than 3 idle executors") {
     val executors = Seq("host1", "host2", "host3", "host4", "host5")
     // host3 is idle
     val receiverTrackingInfoMap = Map(
@@ -51,7 +51,7 @@ class ReceiverSchedulingPolicySuite extends SparkFunSuite {
     assert(scheduledLocations.toSet === Set("host2", "host3", "host4", "host5"))
   }
 
-  test("return 3 best options if less than 3 idle executors") {
+  test("rescheduleReceiver: return 3 best options if less than 3 idle executors") {
     val executors = Seq("host1", "host2", "host3", "host4", "host5")
     // Weights: host1 = 1.5, host2 = 0.5, host3 = 1.0
     // host4 and host5 are idle
@@ -64,8 +64,8 @@ class ReceiverSchedulingPolicySuite extends SparkFunSuite {
     assert(scheduledLocations.toSet === Set("host2", "host4", "host5"))
   }
 
-  test("scheduleReceivers should schedule receivers evenly " +
-    "when there are more receivers than executors") {
+  test("scheduleReceivers: " +
+    "schedule receivers evenly when there are more receivers than executors") {
     val receivers = (0 until 6).map(new DummyReceiver(_))
     val executors = (10000 until 10003).map(port => s"localhost:${port}")
     val scheduledLocations = receiverSchedulingPolicy.scheduleReceivers(receivers, executors)
@@ -79,8 +79,8 @@ class ReceiverSchedulingPolicySuite extends SparkFunSuite {
   }
 
 
-  test("scheduleReceivers should schedule receivers evenly " +
-    "when there are more executors than receivers") {
+  test("scheduleReceivers: " +
+    "schedule receivers evenly when there are more executors than receivers") {
     val receivers = (0 until 3).map(new DummyReceiver(_))
     val executors = (10000 until 10006).map(port => s"localhost:${port}")
     val scheduledLocations = receiverSchedulingPolicy.scheduleReceivers(receivers, executors)
@@ -95,8 +95,7 @@ class ReceiverSchedulingPolicySuite extends SparkFunSuite {
     assert(numReceiversOnExecutor === executors.map(_ -> 1).toMap)
   }
 
-  test("scheduleReceivers should schedule receivers evenly " +
-    "when the preferredLocations are even") {
+  test("scheduleReceivers: schedule receivers evenly when the preferredLocations are even") {
     val receivers = (0 until 3).map(new DummyReceiver(_)) ++
       (3 until 6).map(new DummyReceiver(_, Some("localhost")))
     val executors = (10000 until 10003).map(port => s"localhost:${port}") ++
@@ -119,9 +118,17 @@ class ReceiverSchedulingPolicySuite extends SparkFunSuite {
       (10000 until 10003).map(port => s"localhost:${port}").toSet)
   }
 
-  test("scheduleReceivers should return empty if no receiver") {
+  test("scheduleReceivers: return empty if no receiver") {
     assert(receiverSchedulingPolicy.scheduleReceivers(Seq.empty, Seq("localhost:10000")).isEmpty)
   }
+
+  test("scheduleReceivers: return empty locations if no executors") {
+    val receivers = (0 until 3).map(new DummyReceiver(_))
+    val scheduledLocations = receiverSchedulingPolicy.scheduleReceivers(receivers, Seq.empty)
+    scheduledLocations.foreach { case (receiverId, locations) =>
+      assert(locations.isEmpty)
+    }
+  }
 }
 
 /**