Add a new Receiver scheduling mechanism

Author: zsxwing

core/src/main/scala/org/apache/spark/SparkContext.scala

@@ -1859,6 +1859,29 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
     new SimpleFutureAction(waiter, resultFunc)
   }
 
+  /**
+   * Submit a job for execution and return a FutureJob holding the result.
+   */
+  private[spark] def submitAsyncJob[T, U, R](
+      rdd: RDD[T],
+      processPartition: (TaskContext, Iterator[T]) => U,
+      resultHandler: (Int, U) => Unit,
+      resultFunc: => R): SimpleFutureAction[R] =
+  {
+    assertNotStopped()
+    val cleanF = clean(processPartition)
+    val callSite = getCallSite
+    val waiter = dagScheduler.submitJob(
+      rdd,
+      (context: TaskContext, iter: Iterator[T]) => cleanF(context, iter),
+      0 until rdd.partitions.size,
+      callSite,
+      allowLocal = false,
+      resultHandler,
+      localProperties.get)
+    new SimpleFutureAction(waiter, resultFunc)
+  }
+
   /**
    * Cancel active jobs for the specified group. See [[org.apache.spark.SparkContext.setJobGroup]]
    * for more information.

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

@@ -25,7 +25,7 @@ import scala.concurrent._
 
 import org.apache.spark.{Logging, SparkConf}
 import org.apache.spark.storage.StreamBlockId
-import org.apache.spark.util.ThreadUtils
+import org.apache.spark.util.{Utils, ThreadUtils}
 
 /**
  * Abstract class that is responsible for supervising a Receiver in the worker.
@@ -64,6 +64,8 @@ private[streaming] abstract class ReceiverSupervisor(
   /** State of the receiver */
   @volatile private[streaming] var receiverState = Initialized
 
+  val host = Utils.localHostName()
+
   /** Push a single data item to backend data store. */
   def pushSingle(data: Any)
 
@@ -161,12 +163,19 @@ private[streaming] abstract class ReceiverSupervisor(
       stopReceiver("Restarting receiver with delay " + delay + "ms: " + message, error)
       logDebug("Sleeping for " + delay)
       Thread.sleep(delay)
-      logInfo("Starting receiver again")
-      startReceiver()
-      logInfo("Receiver started again")
+      if (rescheduleReceiver().contains(host)) {
+        logInfo("Starting receiver again")
+        startReceiver()
+        logInfo("Receiver started again")
+      } else {
+        stop("Receiver is scheduled to another executor", None)
+      }
     }(futureExecutionContext)
   }
 
+  /** Reschedule this receiver and return a candidate executor list */
+  def rescheduleReceiver(): Seq[String]
+
   /** Check if receiver has been marked for stopping */
   def isReceiverStarted(): Boolean = {
     logDebug("state = " + receiverState)

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

@@ -30,7 +30,7 @@ import org.apache.spark.storage.StreamBlockId
 import org.apache.spark.streaming.Time
 import org.apache.spark.streaming.scheduler._
 import org.apache.spark.streaming.util.WriteAheadLogUtils
-import org.apache.spark.util.{RpcUtils, Utils}
+import org.apache.spark.util.RpcUtils
 import org.apache.spark.{Logging, SparkEnv, SparkException}
 
 /**
@@ -164,7 +164,7 @@ private[streaming] class ReceiverSupervisorImpl(
 
   override protected def onReceiverStart() {
     val msg = RegisterReceiver(
-      streamId, receiver.getClass.getSimpleName, Utils.localHostName(), endpoint)
+      streamId, receiver.getClass.getSimpleName, host, endpoint)
     trackerEndpoint.askWithRetry[Boolean](msg)
   }
 
@@ -182,4 +182,8 @@ private[streaming] class ReceiverSupervisorImpl(
     logDebug(s"Cleaning up blocks older then $cleanupThreshTime")
     receivedBlockHandler.cleanupOldBlocks(cleanupThreshTime.milliseconds)
   }
+
+  override def rescheduleReceiver(): Seq[String] = {
+    trackerEndpoint.askWithRetry[Seq[String]](ScheduleReceiver(streamId))
+  }
 }

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

@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.scheduler
+
+import scala.collection.mutable
+import scala.util.Random
+
+import org.apache.spark.streaming.scheduler.ReceiverState._
+
+private[streaming] case class ReceiverTrackingInfo(
+    receiverId: Int,
+    state: ReceiverState,
+    preferredLocation: Option[String],
+    scheduledLocations: Option[Seq[String]],
+    runningLocation: Option[String])
+
+private[streaming] trait ReceiverScheduler {
+
+  def scheduleReceiver(
+      receiverId: Int,
+      preferredLocation: Option[String],
+      receiverTrackingInfoMap: Map[Int, ReceiverTrackingInfo],
+      executors: Seq[String]): Seq[String]
+}
+
+/**
+ * A ReceiverScheduler trying to balance executors' load.
+ */
+private[streaming] class LoadBalanceReceiverSchedulerImpl extends ReceiverScheduler {
+
+  def scheduleReceiver(
+      receiverId: Int,
+      preferredLocation: Option[String],
+      receiverTrackingInfoMap: Map[Int, ReceiverTrackingInfo],
+      executors: Seq[String]): Seq[String] = {
+    if (executors.isEmpty) {
+      return Seq.empty
+    }
+
+    // Always try to schedule to the preferred locations
+    val locations = mutable.Set[String]()
+    locations ++= preferredLocation
+
+    val executorWeights = receiverTrackingInfoMap.filter { case (id, _) =>
+      // Ignore the receiver to be scheduled. It may be still running.
+      id != receiverId
+    }.values.flatMap { receiverTrackingInfo =>
+      receiverTrackingInfo.state match {
+        case ReceiverState.INACTIVE => Nil
+        case ReceiverState.SCHEDULED =>
+          val scheduledLocations = receiverTrackingInfo.scheduledLocations.get
+          // The probability that a scheduled receiver will run in an executor is
+          // 1.0 / scheduledLocations.size
+          scheduledLocations.map(location => (location, 1.0 / scheduledLocations.size))
+        case ReceiverState.ACTIVE => Seq(receiverTrackingInfo.runningLocation.get -> 1.0)
+      }
+    }.groupBy(_._1).mapValues(_.map(_._2).sum) // Sum weights for each executor
+
+    val idleExecutors = (executors.toSet -- executorWeights.keys).toSeq
+    if (idleExecutors.nonEmpty) {
+      // If there are idle executors, randomly select one
+      locations += idleExecutors(Random.nextInt(idleExecutors.size))
+    } else {
+      // Use the executor that runs the least receivers
+      locations += executorWeights.minBy(_._2)._1
+    }
+    locations.toSeq
+  }
+}