Merge branch 'squash-pr-6294' into receiver-scheduling

Author: zsxwing

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

@@ -22,6 +22,7 @@ import java.util.concurrent.CountDownLatch
 
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent._
+import scala.util.control.NonFatal
 
 import org.apache.spark.{Logging, SparkConf}
 import org.apache.spark.storage.StreamBlockId
@@ -36,7 +37,7 @@ private[streaming] abstract class ReceiverSupervisor(
     conf: SparkConf
   ) extends Logging {
 
-  /** Enumeration to identify current state of the StreamingContext */
+  /** Enumeration to identify current state of the Receiver */
   object ReceiverState extends Enumeration {
     type CheckpointState = Value
     val Initialized, Started, Stopped = Value
@@ -99,8 +100,8 @@ private[streaming] abstract class ReceiverSupervisor(
   /** Called when supervisor is stopped */
   protected def onStop(message: String, error: Option[Throwable]) { }
 
-  /** Called when receiver is started */
-  protected def onReceiverStart() { }
+  /** Called when receiver is started. Return if the driver accepts us */
+  protected def onReceiverStart(): Boolean = true
 
   /** Called when receiver is stopped */
   protected def onReceiverStop(message: String, error: Option[Throwable]) { }
@@ -123,13 +124,17 @@ private[streaming] abstract class ReceiverSupervisor(
   /** Start receiver */
   def startReceiver(): Unit = synchronized {
     try {
-      logInfo("Starting receiver")
-      receiver.onStart()
-      logInfo("Called receiver onStart")
-      onReceiverStart()
-      receiverState = Started
+      if (onReceiverStart()) {
+        logInfo("Starting receiver")
+        receiverState = Started
+        receiver.onStart()
+        logInfo("Called receiver onStart")
+      } else {
+        // The driver refused us
+        stop("Registered unsuccessfully because the driver refused" + streamId, None)
+      }
     } catch {
-      case t: Throwable =>
+      case NonFatal(t) =>
         stop("Error starting receiver " + streamId, Some(t))
     }
   }
@@ -138,12 +143,19 @@ private[streaming] abstract class ReceiverSupervisor(
   def stopReceiver(message: String, error: Option[Throwable]): Unit = synchronized {
     try {
       logInfo("Stopping receiver with message: " + message + ": " + error.getOrElse(""))
-      receiverState = Stopped
-      receiver.onStop()
-      logInfo("Called receiver onStop")
-      onReceiverStop(message, error)
+      receiverState match {
+        case Initialized =>
+          logWarning("Skip stopping receiver because it has not yet stared")
+        case Started =>
+          receiverState = Stopped
+          receiver.onStop()
+          logInfo("Called receiver onStop")
+          onReceiverStop(message, error)
+        case Stopped =>
+          logWarning("Receiver has been stopped")
+      }
     } catch {
-      case t: Throwable =>
+      case NonFatal(t) =>
         logError("Error stopping receiver " + streamId + t.getStackTraceString)
     }
   }
@@ -177,7 +189,7 @@ private[streaming] abstract class ReceiverSupervisor(
   /** Return a list of candidate executors to run the receiver */
   def getAllowedLocations(): Seq[String] = Seq.empty
 
-  /** Check if receiver has been marked for stopping */
+  /** Check if receiver has been marked for starting */
   def isReceiverStarted(): Boolean = {
     logDebug("state = " + receiverState)
     receiverState == Started

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

@@ -162,7 +162,7 @@ private[streaming] class ReceiverSupervisorImpl(
     env.rpcEnv.stop(endpoint)
   }
 
-  override protected def onReceiverStart() {
+  override protected def onReceiverStart(): Boolean = {
     val msg = RegisterReceiver(
       streamId, receiver.getClass.getSimpleName, host, endpoint)
     trackerEndpoint.askWithRetry[Boolean](msg)

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

@@ -85,15 +85,26 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
   )
   private val listenerBus = ssc.scheduler.listenerBus
 
+  /** Enumeration to identify current state of the ReceiverTracker */
+  object TrackerState extends Enumeration {
+    type CheckpointState = Value
+    val Initialized, Started, Stopping, Stopped = Value
+  }
+  import TrackerState._
+
+  /** State of the tracker. Protected by "trackerStateLock" */
+  private var trackerState = Initialized
+
+  /** "trackerStateLock" is used to protect reading/writing "trackerState" */
+  private val trackerStateLock = new AnyRef
+
   // endpoint is created when generator starts.
   // This not being null means the tracker has been started and not stopped
   private var endpoint: RpcEndpointRef = null
 
   private val schedulingPolicy: ReceiverSchedulingPolicy =
     new LoadBalanceReceiverSchedulingPolicyImpl()
 
-  @volatile private var stopping = false
-
   /**
    * Track receivers' status for scheduling
    */
@@ -107,9 +118,24 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
   /** Use a separate lock to avoid dead-lock */
   private val receiverTrackingInfosLock = new AnyRef
 
+  /** Check if tracker has been marked for starting */
+  private def isTrackerStarted(): Boolean = trackerStateLock.synchronized {
+    trackerState == Started
+  }
+ 
+  /** Check if tracker has been marked for stopping */
+  private def isTrackerStopping(): Boolean = trackerStateLock.synchronized {
+    trackerState == Stopping
+  }
+ 
+  /** Check if tracker has been marked for stopped */
+  private def isTrackerStopped(): Boolean = trackerStateLock.synchronized {
+    trackerState == Stopped
+  }
+
   /** Start the endpoint and receiver execution thread. */
   def start(): Unit = synchronized {
-    if (endpoint != null) {
+    if (isTrackerStarted) {
       throw new SparkException("ReceiverTracker already started")
     }
 
@@ -118,20 +144,29 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
         "ReceiverTracker", new ReceiverTrackerEndpoint(ssc.env.rpcEnv))
       if (!skipReceiverLaunch) receiverExecutor.start()
       logInfo("ReceiverTracker started")
+      trackerStateLock.synchronized {
+        trackerState = Started
+      }
     }
   }
 
   /** Stop the receiver execution thread. */
   def stop(graceful: Boolean): Unit = synchronized {
-    if (!receiverInputStreams.isEmpty && endpoint != null) {
+    if (isTrackerStarted) {
       // First, stop the receivers
+      trackerStateLock.synchronized {
+        trackerState = Stopping
+      }
       if (!skipReceiverLaunch) receiverExecutor.stop(graceful)
 
       // Finally, stop the endpoint
       ssc.env.rpcEnv.stop(endpoint)
       endpoint = null
       receivedBlockTracker.stop()
       logInfo("ReceiverTracker stopped")
+      trackerStateLock.synchronized {
+        trackerState = Stopped
+      }
     }
   }
 
@@ -177,15 +212,24 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
       host: String,
       receiverEndpoint: RpcEndpointRef,
       senderAddress: RpcAddress
-    ) {
+    ): Boolean = {
     if (!receiverInputStreamIds.contains(streamId)) {
       throw new SparkException("Register received for unexpected id " + streamId)
     }
-    receiverInfo(streamId) = ReceiverInfo(
-      streamId, s"${typ}-${streamId}", receiverEndpoint, true, host)
-    updateReceiverRunningLocation(streamId, host)
-    listenerBus.post(StreamingListenerReceiverStarted(receiverInfo(streamId)))
-    logInfo("Registered receiver for stream " + streamId + " from " + senderAddress)
+    trackerStateLock.synchronized {
+      if (isTrackerStopping || isTrackerStopped) {
+        false
+      } else {
+        // When updating "receiverInfo", we should make sure "trackerState" won't be changed at the
+        // same time. Therefore the following line should be in "trackerStateLock.synchronized".
+        receiverInfo(streamId) = ReceiverInfo(
+          streamId, s"${typ}-${streamId}", receiverEndpoint, true, host)
+        updateReceiverRunningLocation(streamId, host)
+        listenerBus.post(StreamingListenerReceiverStarted(receiverInfo(streamId)))
+        logInfo("Registered receiver for stream " + streamId + " from " + senderAddress)
+        true
+      }
+    }
   }
 
   /** Deregister a receiver */
@@ -253,8 +297,9 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
       case RegisterReceiver(streamId, typ, host, receiverEndpoint) =>
-        registerReceiver(streamId, typ, host, receiverEndpoint, context.sender.address)
-        context.reply(true)
+        val successful =
+          registerReceiver(streamId, typ, host, receiverEndpoint, context.sender.address)
+        context.reply(successful)
       case AddBlock(receivedBlockInfo) =>
         context.reply(addBlock(receivedBlockInfo))
       case DeregisterReceiver(streamId, message, error) =>
@@ -285,8 +330,6 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
     }
 
     def stop(graceful: Boolean) {
-      stopping = true
-
       // Send the stop signal to all the receivers
       stopReceivers()
 
@@ -389,7 +432,7 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
       for (receiver <- receivers) {
         submitJobThread.execute(new Runnable {
           override def run(): Unit = {
-            if (stopping) {
+            if (isTrackerStopping()) {
               receiverExitLatch.countDown()
               return
             }
@@ -409,17 +452,17 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
                 ssc.sc.makeRDD(Seq(receiver -> scheduledLocations))
               }
             val future = ssc.sparkContext.submitJob[Receiver[_], Unit, Unit](
-              receiverRDD, startReceiver, (_, _) => Unit, ())
+              receiverRDD, startReceiver, Seq(0), (_, _) => Unit, ())
             future.onComplete {
               case Success(_) =>
-                if (stopping) {
+                if (isTrackerStopping()) {
                   receiverExitLatch.countDown()
                 } else {
                   logInfo(s"Restarting Receiver $receiverId")
                   submitJobThread.execute(self)
                 }
               case Failure(e) =>
-                if (stopping) {
+                if (isTrackerStopping()) {
                   receiverExitLatch.countDown()
                 } else {
                   logError("Receiver has been stopped. Try to restart it.", e)
@@ -431,11 +474,14 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
           }
         })
       }
-      // Wait until all receivers exit
-      receiverExitLatch.await()
-      running = false
-      logInfo("All of the receivers have been terminated")
-      submitJobThread.shutdownNow()
+      try {
+        // Wait until all receivers exit
+        receiverExitLatch.await()
+        logInfo("All of the receivers have been terminated")
+      } finally {
+        running = false
+        submitJobThread.shutdownNow()
+      }
     }
 
     /** Stops the receivers. */

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

@@ -273,6 +273,19 @@ class StreamingContextSuite extends SparkFunSuite with BeforeAndAfter with Timeo
     }
   }
 
+  test("stop gracefully even if a receiver misses StopReceiver") {
+    val conf = new SparkConf().setMaster(master).setAppName(appName)
+    sc = new SparkContext(conf)
+    ssc = new StreamingContext(sc, Milliseconds(100))
+    val input = ssc.receiverStream(new TestReceiver)
+    input.foreachRDD(_ => {})
+    ssc.start()
+    // Call `ssc.stop` at once so that it's possible that the receiver will miss "StopReceiver"
+    failAfter(30000 millis) {
+      ssc.stop(stopSparkContext = true, stopGracefully = true)
+    }
+  }
+
   test("stop slow receiver gracefully") {
     val conf = new SparkConf().setMaster(master).setAppName(appName)
     conf.set("spark.streaming.gracefulStopTimeout", "20000s")