Schedule all receivers at the same time when launching them

zsxwing · zsxwing · commit 59f88874d9c2 · 2015-07-22T22:33:27.000+08:00
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverSchedulingPolicy.scala b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverSchedulingPolicy.scala
@@ -20,6 +20,8 @@ package org.apache.spark.streaming.scheduler
 import scala.collection.Map
 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.
@@ -46,11 +48,80 @@ import scala.collection.mutable
  */
 private[streaming] class ReceiverSchedulingPolicy {
 
+  /**
+   * Try our best to schedule receivers with evenly distributed. However, if the
+   * `preferredLocation`s of receivers are not even, we may not be able to schedule them evenly
+   * because we have to respect them.
+   *
+   * This method is called when we start to launch receivers at the first time.
+   */
+  def scheduleReceivers(
+      receivers: Seq[Receiver[_]], executors: Seq[String]): Map[Int, Seq[String]] = {
+    if (receivers.isEmpty) {
+      return Map.empty
+    }
+
+    require(executors.nonEmpty, "There is no executor up")
+
+    val hostToExecutors = executors.groupBy(_.split(":")(0))
+    val locations = new Array[mutable.ArrayBuffer[String]](receivers.length)
+    val numReceiversOnExecutor = mutable.HashMap[String, Int]()
+    // Set the initial value to 0
+    executors.foreach(numReceiversOnExecutor(_) = 0)
+
+    // Firstly, we need to respect "preferredLocation". So if a receiver has "preferredLocation",
+    // we need to make sure the "preferredLocation" is in the candidate location list.
+    for (i <- 0 until receivers.length) {
+      locations(i) = new mutable.ArrayBuffer[String]()
+      // Note: preferredLocation is host but executors are host:port
+      receivers(i).preferredLocation.foreach { host =>
+        hostToExecutors.get(host) match {
+          case Some(executorsOnHost) =>
+            // preferredLocation is a known host. Select an executor that has the least receivers in
+            // this host
+            val scheduledLocation =
+              executorsOnHost.minBy(executor => numReceiversOnExecutor(executor))
+            locations(i) += scheduledLocation
+            numReceiversOnExecutor(scheduledLocation) =
+              numReceiversOnExecutor(scheduledLocation) + 1
+          case None =>
+            // preferredLocation is an unknown host.
+            // Note: There are two cases:
+            // 1. This executor is not up. But it may be up later.
+            // 2. This executor is dead, or it's not a host in the cluster.
+            // Currently, simply add host to the scheduled locations
+            locations(i) += host
+        }
+      }
+    }
+
+    // For those receivers that don't have preferredLocation, make sure we assign at least one
+    // executor to them.
+    for (scheduledLocations <- locations.filter(_.isEmpty)) {
+      // Select the executor that has the least receivers
+      val (executor, numReceivers) = numReceiversOnExecutor.minBy(_._2)
+      scheduledLocations += executor
+      numReceiversOnExecutor(executor) = numReceivers + 1
+    }
+
+    // Assign idle executors to receivers that have less executors
+    val idleExecutors = numReceiversOnExecutor.filter(_._2 == 0).map(_._1)
+    for (executor <- idleExecutors) {
+      // Assign an idle executor to the receiver that has least locations.
+      val scheduledLocations = locations.minBy(_.size)
+      scheduledLocations += executor
+    }
+
+    receivers.map(_.streamId).zip(locations).toMap
+  }
+
   /**
    * 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 is called when a receiver is registering with ReceiverTracker or is restarting.
    */
-  def scheduleReceiver(
+  def rescheduleReceiver(
       receiverId: Int,
       preferredLocation: Option[String],
       receiverTrackingInfoMap: Map[Int, ReceiverTrackingInfo],
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTracker.scala b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTracker.scala
@@ -62,9 +62,16 @@ private[streaming] case class DeregisterReceiver(streamId: Int, msg: String, err
 private[streaming] sealed trait ReceiverTrackerLocalMessage
 
 /**
- * This message will trigger ReceiverTrackerEndpoint to start a Spark job for the receiver.
+ * This message will trigger ReceiverTrackerEndpoint to restart a Spark job for the receiver.
  */
-private[streaming] case class StartReceiver(receiver: Receiver[_])
+private[streaming] case class RestartReceiver(receiver: Receiver[_])
+  extends ReceiverTrackerLocalMessage
+
+/**
+ * This message is sent to ReceiverTrackerEndpoint when we start to launch Spark jobs for receivers
+ * at the first time.
+ */
+private[streaming] case class StartAllReceivers(receiver: Seq[Receiver[_]])
   extends ReceiverTrackerLocalMessage
 
 /**
@@ -307,8 +314,8 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
 
   private def scheduleReceiver(receiverId: Int): Seq[String] = {
     val preferredLocation = receiverPreferredLocations.getOrElse(receiverId, None)
-    val scheduledLocations = schedulingPolicy.scheduleReceiver(
-      receiverId, preferredLocation, receiverTrackingInfos, getExecutors(ssc))
+    val scheduledLocations = schedulingPolicy.rescheduleReceiver(
+      receiverId, preferredLocation, receiverTrackingInfos, getExecutors)
     updateReceiverScheduledLocations(receiverId, scheduledLocations)
     scheduledLocations
   }
@@ -337,10 +344,14 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
   /**
    * Get the list of executors excluding driver
    */
-  private def getExecutors(ssc: StreamingContext): List[String] = {
-    val executors = ssc.sparkContext.getExecutorMemoryStatus.map(_._1.split(":")(0)).toList
-    val driver = ssc.sparkContext.getConf.get("spark.driver.host")
-    executors.diff(List(driver))
+  private def getExecutors: List[String] = {
+    if (ssc.sc.isLocal) {
+      List("localhost")
+    } else {
+      val executors = ssc.sparkContext.getExecutorMemoryStatus.map(_._1.split(":")(0)).toList
+      val driver = ssc.sparkContext.getConf.get("spark.driver.host")
+      executors.diff(List(driver))
+    }
   }
 
   /**
@@ -355,6 +366,7 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
     if (!ssc.sparkContext.isLocal) {
       ssc.sparkContext.makeRDD(1 to 50, 50).map(x => (x, 1)).reduceByKey(_ + _, 20).collect()
     }
+    assert(getExecutors.nonEmpty)
   }
 
   /**
@@ -370,12 +382,8 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
 
     runDummySparkJob()
 
-    // Distribute the receivers and start them
     logInfo("Starting " + receivers.length + " receivers")
-
-    for (receiver <- receivers) {
-      endpoint.send(StartReceiver(receiver))
-    }
+    endpoint.send(StartAllReceivers)
   }
 
   /** Check if tracker has been marked for starting */
@@ -396,8 +404,22 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
 
     override def receive: PartialFunction[Any, Unit] = {
       // Local messages
-      case StartReceiver(receiver) =>
-        startReceiver(receiver)
+      case StartAllReceivers(receivers) =>
+        val scheduledLocations = schedulingPolicy.scheduleReceivers(receivers, getExecutors)
+        for (receiver <- receivers) {
+          val locations = scheduledLocations(receiver.streamId)
+          updateReceiverScheduledLocations(receiver.streamId, locations)
+          receiverPreferredLocations(receiver.streamId) = receiver.preferredLocation
+          startReceiver(receiver, locations)
+        }
+      case RestartReceiver(receiver) =>
+        val scheduledLocations = schedulingPolicy.rescheduleReceiver(
+          receiver.streamId,
+          receiver.preferredLocation,
+          receiverTrackingInfos,
+          getExecutors)
+        updateReceiverScheduledLocations(receiver.streamId, scheduledLocations)
+        startReceiver(receiver, scheduledLocations)
       case c @ CleanupOldBlocks(cleanupThreshTime) =>
         receiverTrackingInfos.values.flatMap(_.endpoint).foreach(_.send(c))
       // Remote messages
@@ -425,28 +447,22 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
         context.reply(true)
     }
 
-    private def startReceiver(receiver: Receiver[_]): Unit = {
-      val checkpointDirOption = Option(ssc.checkpointDir)
-      val serializableHadoopConf =
-        new SerializableConfiguration(ssc.sparkContext.hadoopConfiguration)
-
-      // Function to start the receiver on the worker node
-      val startReceiverFunc = new StartReceiverFunc(checkpointDirOption, serializableHadoopConf)
-
-      receiverPreferredLocations(receiver.streamId) = receiver.preferredLocation
+    /**
+     * Start a receiver along with its scheduled locations
+     */
+    private def startReceiver(receiver: Receiver[_], scheduledLocations: Seq[String]): Unit = {
       val receiverId = receiver.streamId
-
       if (!isTrackerStarted) {
         onReceiverJobFinish(receiverId)
         return
       }
 
-      val scheduledLocations = schedulingPolicy.scheduleReceiver(
-        receiverId,
-        receiver.preferredLocation,
-        receiverTrackingInfos,
-        getExecutors(ssc))
-      updateReceiverScheduledLocations(receiver.streamId, scheduledLocations)
+      val checkpointDirOption = Option(ssc.checkpointDir)
+      val serializableHadoopConf =
+        new SerializableConfiguration(ssc.sparkContext.hadoopConfiguration)
+
+      // Function to start the receiver on the worker node
+      val startReceiverFunc = new StartReceiverFunc(checkpointDirOption, serializableHadoopConf)
 
       // Create the RDD using the scheduledLocations to run the receiver in a Spark job
       val receiverRDD: RDD[Receiver[_]] =
@@ -465,15 +481,15 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
             onReceiverJobFinish(receiverId)
           } else {
             logInfo(s"Restarting Receiver $receiverId")
-            self.send(StartReceiver(receiver))
+            self.send(RestartReceiver(receiver))
           }
         case Failure(e) =>
           if (!isTrackerStarted) {
             onReceiverJobFinish(receiverId)
           } else {
             logError("Receiver has been stopped. Try to restart it.", e)
             logInfo(s"Restarting Receiver $receiverId")
-            self.send(StartReceiver(receiver))
+            self.send(RestartReceiver(receiver))
           }
       }(submitJobThreadPool)
       logInfo(s"Receiver ${receiver.streamId} started")
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/scheduler/ReceiverSchedulingPolicySuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/scheduler/ReceiverSchedulingPolicySuite.scala
@@ -17,22 +17,26 @@
 
 package org.apache.spark.streaming.scheduler
 
+import scala.collection.mutable
+
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.storage.StorageLevel
+import org.apache.spark.streaming.receiver.Receiver
 
 class ReceiverSchedulingPolicySuite extends SparkFunSuite {
 
   val receiverSchedulingPolicy = new ReceiverSchedulingPolicy
 
   test("empty executors") {
     val scheduledLocations =
-      receiverSchedulingPolicy.scheduleReceiver(0, None, Map.empty, executors = Seq.empty)
+      receiverSchedulingPolicy.rescheduleReceiver(0, None, Map.empty, executors = Seq.empty)
     assert(scheduledLocations === Seq.empty)
   }
 
   test("receiver preferredLocation") {
     val receiverTrackingInfoMap = Map(
       0 -> ReceiverTrackingInfo(0, ReceiverState.INACTIVE, None, None))
-    val scheduledLocations = receiverSchedulingPolicy.scheduleReceiver(
+    val scheduledLocations = receiverSchedulingPolicy.rescheduleReceiver(
       0, Some("host1"), receiverTrackingInfoMap, executors = Seq("host2"))
     assert(scheduledLocations.toSet === Set("host1", "host2"))
   }
@@ -42,7 +46,7 @@ class ReceiverSchedulingPolicySuite extends SparkFunSuite {
     // host3 is idle
     val receiverTrackingInfoMap = Map(
       0 -> ReceiverTrackingInfo(0, ReceiverState.ACTIVE, None, Some("host1")))
-    val scheduledLocations = receiverSchedulingPolicy.scheduleReceiver(
+    val scheduledLocations = receiverSchedulingPolicy.rescheduleReceiver(
       1, None, receiverTrackingInfoMap, executors)
     assert(scheduledLocations.toSet === Set("host2", "host3", "host4", "host5"))
   }
@@ -55,8 +59,82 @@ class ReceiverSchedulingPolicySuite extends SparkFunSuite {
       0 -> ReceiverTrackingInfo(0, ReceiverState.ACTIVE, None, Some("host1")),
       1 -> ReceiverTrackingInfo(1, ReceiverState.SCHEDULED, Some(Seq("host2", "host3")), None),
       2 -> ReceiverTrackingInfo(1, ReceiverState.SCHEDULED, Some(Seq("host1", "host3")), None))
-    val scheduledLocations = receiverSchedulingPolicy.scheduleReceiver(
+    val scheduledLocations = receiverSchedulingPolicy.rescheduleReceiver(
       3, None, receiverTrackingInfoMap, executors)
     assert(scheduledLocations.toSet === Set("host2", "host4", "host5"))
   }
+
+  test("scheduleReceivers should 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)
+    val numReceiversOnExecutor = mutable.HashMap[String, Int]()
+    // There should be 2 receivers running on each executor and each receiver has one location
+    scheduledLocations.foreach { case (receiverId, locations) =>
+      assert(locations.size == 1)
+      numReceiversOnExecutor(locations(0)) = numReceiversOnExecutor.getOrElse(locations(0), 0) + 1
+    }
+    assert(numReceiversOnExecutor === executors.map(_ -> 2).toMap)
+  }
+
+
+  test("scheduleReceivers should 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)
+    val numReceiversOnExecutor = mutable.HashMap[String, Int]()
+    // There should be 1 receiver running on each executor and each receiver has two locations
+    scheduledLocations.foreach { case (receiverId, locations) =>
+      assert(locations.size == 2)
+      locations.foreach { l =>
+        numReceiversOnExecutor(l) = numReceiversOnExecutor.getOrElse(l, 0) + 1
+      }
+    }
+    assert(numReceiversOnExecutor === executors.map(_ -> 1).toMap)
+  }
+
+  test("scheduleReceivers should 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}") ++
+      (10003 until 10006).map(port => s"localhost2:${port}")
+    val scheduledLocations = receiverSchedulingPolicy.scheduleReceivers(receivers, executors)
+    val numReceiversOnExecutor = mutable.HashMap[String, Int]()
+    // There should be 1 receiver running on each executor and each receiver has 1 location
+    scheduledLocations.foreach { case (receiverId, locations) =>
+      assert(locations.size == 1)
+      locations.foreach { l =>
+        numReceiversOnExecutor(l) = numReceiversOnExecutor.getOrElse(l, 0) + 1
+      }
+    }
+    assert(numReceiversOnExecutor === executors.map(_ -> 1).toMap)
+    // Make sure we schedule the receivers to their preferredLocations
+    val locationsForReceiversWithPreferredLocation =
+      scheduledLocations.filter { case (receiverId, locations) => receiverId >= 3 }.flatMap(_._2)
+    // We can simply check the location set because we only know each receiver only has 1 location
+    assert(locationsForReceiversWithPreferredLocation.toSet ===
+      (10000 until 10003).map(port => s"localhost:${port}").toSet)
+  }
+
+  test("scheduleReceivers should return empty if no receiver") {
+    assert(receiverSchedulingPolicy.scheduleReceivers(Seq.empty, Seq("localhost:10000")).isEmpty)
+  }
+}
+
+/**
+ * Dummy receiver implementation
+ */
+private class DummyReceiver(receiverId: Int, host: Option[String] = None)
+  extends Receiver[Int](StorageLevel.MEMORY_ONLY) {
+
+  setReceiverId(receiverId)
+
+  override def onStart(): Unit = {}
+
+  override def onStop(): Unit = {}
+
+  override def preferredLocation: Option[String] = host
 }
