Make 'host' protected; rescheduleReceiver -> getAllowedLocations

Author: zsxwing

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

@@ -64,7 +64,7 @@ private[streaming] abstract class ReceiverSupervisor(
   /** State of the receiver */
   @volatile private[streaming] var receiverState = Initialized
 
-  val host = Utils.localHostName()
+  protected val host = Utils.localHostName()
 
   /** Push a single data item to backend data store. */
   def pushSingle(data: Any)
@@ -163,7 +163,7 @@ private[streaming] abstract class ReceiverSupervisor(
       stopReceiver("Restarting receiver with delay " + delay + "ms: " + message, error)
       logDebug("Sleeping for " + delay)
       Thread.sleep(delay)
-      val scheduledLocations = rescheduleReceiver()
+      val scheduledLocations = getAllowedLocations()
       if (scheduledLocations.isEmpty || scheduledLocations.contains(host)) {
         logInfo("Starting receiver again")
         startReceiver()
@@ -175,7 +175,7 @@ private[streaming] abstract class ReceiverSupervisor(
   }
 
   /** Reschedule this receiver and return a candidate executor list */
-  def rescheduleReceiver(): Seq[String] = Seq.empty
+  def getAllowedLocations(): Seq[String] = Seq.empty
 
   /** Check if receiver has been marked for stopping */
   def isReceiverStarted(): Boolean = {

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

@@ -183,7 +183,7 @@ private[streaming] class ReceiverSupervisorImpl(
     receivedBlockHandler.cleanupOldBlocks(cleanupThreshTime.milliseconds)
   }
 
-  override def rescheduleReceiver(): Seq[String] = {
+  override def getAllowedLocations(): Seq[String] = {
     trackerEndpoint.askWithRetry[Seq[String]](ScheduleReceiver(streamId))
   }
 }

streaming/src/test/scala/org/apache/spark/streaming/ReceiverSuite.scala

@@ -132,7 +132,7 @@ class ReceiverSuite extends TestSuiteBase with Timeouts with Serializable {
   test("restart receiver should consider the scheduled locations") {
     val receiver = new FakeReceiver
     val executor = new FakeReceiverSupervisor(receiver) {
-      override def rescheduleReceiver: Seq[String] = Seq("unknown-host")
+      override def getAllowedLocations: Seq[String] = Seq("unknown-host")
     }
     executor.start()
     receiver.restart("force the receiver restart")