Use a lock to eliminate the race condition when stopping receivers and registering receivers happen at the same time.

Author: zsxwing

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

@@ -93,10 +93,7 @@ private[streaming] abstract class ReceiverSupervisor(
 
   /** Called when supervisor is stopped */
   protected def onStop(message: String, error: Option[Throwable]) { }
- 
-  /** Called when receiver is registered */
-  protected def onReceiverRegister() { }
- 
+
   /** Called when receiver is started */
   protected def onReceiverStart() { }
 
@@ -120,7 +117,6 @@ private[streaming] abstract class ReceiverSupervisor(
   /** Start receiver */
   def startReceiver(): Unit = synchronized {
     try {
-      onReceiverRegister()
       logInfo("Starting receiver")
       receiver.onStart()
       logInfo("Called receiver onStart")

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

@@ -166,23 +166,13 @@ private[streaming] class ReceiverSupervisorImpl(
     blockGenerator.stop()
     env.rpcEnv.stop(endpoint)
   }
- 
-  override protected def onReceiverRegister() {
-    val msg = RegisterReceiver(
-      streamId, receiver.getClass.getSimpleName, Utils.localHostName(), endpoint)
-    val ret = trackerEndpoint.askWithRetry[Boolean](msg)
-    if (!ret) {
-      throw new SparkException("ReceiverTracker is stopping and doesn't accept registeration " +
-        "from receivers.")
-    }
-  }
- 
+
   override protected def onReceiverStart() {
-    val msg = ReceiverStarted(
+    val msg = RegisterReceiver(
       streamId, receiver.getClass.getSimpleName, Utils.localHostName(), endpoint)
-    val ret = trackerEndpoint.askWithRetry[Boolean](msg)
-    if (!ret) {
-      throw new SparkException("ReceiverTracker is stopping and doesn't accept receiver started.")
+    val successful = trackerEndpoint.askWithRetry[Boolean](msg)
+    if (!successful) {
+      stop("Registered unsuccessfully", None)
     }
   }
 

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

@@ -37,12 +37,6 @@ private[streaming] case class RegisterReceiver(
     host: String,
     receiverEndpoint: RpcEndpointRef
   ) extends ReceiverTrackerMessage
-private[streaming] case class ReceiverStarted(
-    streamId: Int,
-    typ: String,
-    host: String,
-    receiverEndpoint: RpcEndpointRef
- ) extends ReceiverTrackerMessage
 private[streaming] case class AddBlock(receivedBlockInfo: ReceivedBlockInfo)
   extends ReceiverTrackerMessage
 private[streaming] case class ReportError(streamId: Int, message: String, error: String)
@@ -83,6 +77,12 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
   /** State of the tracker */
   @volatile private var trackerState = Initialized
 
+  /**
+   * There is a race condition when stopping receivers and registering receivers happen at the same
+   * time. This lock is used to eliminate the race condition. See SPARK-5681.
+   */
+  private val stoppingTrackerLock = 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
@@ -120,8 +120,16 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
   /** Stop the receiver execution thread. */
   def stop(graceful: Boolean): Unit = synchronized {
     if (isTrackerStarted) {
-      trackerState = Stopping
       // First, stop the receivers
+      // acquire "stoppingTrackerLock" so that setting trackerState to "Stopping" and registering
+      // receivers won't happen at the same time
+      stoppingTrackerLock.synchronized {
+        trackerState = Stopping
+        if (!skipReceiverLaunch) {
+          // Send the stop signal to all the receivers
+          receiverExecutor.stopReceivers()
+        }
+      }
       if (!skipReceiverLaunch) receiverExecutor.stop(graceful)
 
       // Finally, stop the endpoint
@@ -175,33 +183,25 @@ 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)
-    listenerBus.post(StreamingListenerReceiverRegistered(receiverInfo(streamId)))
-    logInfo("Registered receiver for stream " + streamId + " from " + senderAddress)
-  }
- 
-  /** Receiver started */
-  private def receiverStarted(
-      streamId: Int,
-      typ: String,
-      host: String,
-      receiverEndpoint: RpcEndpointRef,
-      senderAddress: RpcAddress
-    ) {
-    if (!receiverInputStreamIds.contains(streamId)) {
-      throw new SparkException("Start received for unexpected id " + streamId)
+    // acquire "stoppingTrackerLock" so that setting trackerState to "Stopping" and registering
+    // receivers won't happen at the same time
+    stoppingTrackerLock.synchronized {
+      if (isTrackerStopping || isTrackerStopped) {
+        false
+      } else {
+        receiverInfo(streamId) = ReceiverInfo(
+          streamId, s"${typ}-${streamId}", receiverEndpoint, true, host)
+        listenerBus.post(StreamingListenerReceiverStarted(receiverInfo(streamId)))
+        logInfo("Registered receiver for stream " + streamId + " from " + senderAddress)
+        true
+      }
     }
-    receiverInfo(streamId) = ReceiverInfo(
-      streamId, s"${typ}-${streamId}", receiverEndpoint, true, host)
-    listenerBus.post(StreamingListenerReceiverStarted(receiverInfo(streamId)))
-    logInfo("Receiver started for stream " + streamId + " from " + senderAddress)
   }
- 
+
   /** Deregister a receiver */
   private def deregisterReceiver(streamId: Int, message: String, error: String) {
     val newReceiverInfo = receiverInfo.get(streamId) match {
@@ -267,19 +267,9 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
       case RegisterReceiver(streamId, typ, host, receiverEndpoint) =>
-        if (!isTrackerStopping) {
+        val successful =
           registerReceiver(streamId, typ, host, receiverEndpoint, context.sender.address)
-          context.reply(true)
-        } else {
-          context.reply(false)
-        }
-      case ReceiverStarted(streamId, typ, host, receiverEndpoint) =>
-        if (!isTrackerStopping) {
-          receiverStarted(streamId, typ, host, receiverEndpoint, context.sender.address)
-          context.reply(true)
-        } else {
-          context.reply(false)
-        }
+        context.reply(successful)
       case AddBlock(receivedBlockInfo) =>
         context.reply(addBlock(receivedBlockInfo))
       case DeregisterReceiver(streamId, message, error) =>
@@ -308,9 +298,6 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
     }
 
     def stop(graceful: Boolean) {
-      // Send the stop signal to all the receivers
-      stopReceivers()
-
       // Wait for the Spark job that runs the receivers to be over
       // That is, for the receivers to quit gracefully.
       thread.join(10000)
@@ -385,7 +372,7 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
     }
 
     /** Stops the receivers. */
-    private def stopReceivers() {
+    def stopReceivers() {
       // Signal the receivers to stop
       receiverInfo.values.flatMap { info => Option(info.endpoint)}
                          .foreach { _.send(StopReceiver) }

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

@@ -37,11 +37,7 @@ case class StreamingListenerBatchCompleted(batchInfo: BatchInfo) extends Streami
 
 @DeveloperApi
 case class StreamingListenerBatchStarted(batchInfo: BatchInfo) extends StreamingListenerEvent
- 
-@DeveloperApi
-case class StreamingListenerReceiverRegistered(receiverInfo: ReceiverInfo)
-  extends StreamingListenerEvent
- 
+
 @DeveloperApi
 case class StreamingListenerReceiverStarted(receiverInfo: ReceiverInfo)
   extends StreamingListenerEvent
@@ -61,10 +57,7 @@ case class StreamingListenerReceiverStopped(receiverInfo: ReceiverInfo)
  */
 @DeveloperApi
 trait StreamingListener {
- 
-  /** Called when a receiver has been registered */
-  def onReceiverRegistered(receiverRegistered: StreamingListenerReceiverRegistered) { }
- 
+
   /** Called when a receiver has been started */
   def onReceiverStarted(receiverStarted: StreamingListenerReceiverStarted) { }
 

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

@@ -31,8 +31,6 @@ private[spark] class StreamingListenerBus
 
   override def onPostEvent(listener: StreamingListener, event: StreamingListenerEvent): Unit = {
     event match {
-      case receiverRegistered: StreamingListenerReceiverRegistered =>
-        listener.onReceiverRegistered(receiverRegistered)
       case receiverStarted: StreamingListenerReceiverStarted =>
         listener.onReceiverStarted(receiverStarted)
       case receiverError: StreamingListenerReceiverError =>

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

@@ -224,6 +224,19 @@ class StreamingContextSuite extends FunSuite with BeforeAndAfter with Timeouts w
     }
   }
 
+  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")

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

@@ -117,7 +117,6 @@ class StreamingListenerSuite extends TestSuiteBase with Matchers {
     ssc.start()
     try {
       eventually(timeout(2000 millis), interval(20 millis)) {
-        collector.registeredReceiverStreamIds.size should equal (1)
         collector.startedReceiverStreamIds.size should equal (1)
         collector.startedReceiverStreamIds(0) should equal (0)
         collector.stoppedReceiverStreamIds should have size 1
@@ -162,16 +161,11 @@ class BatchInfoCollector extends StreamingListener {
 
 /** Listener that collects information on processed batches */
 class ReceiverInfoCollector extends StreamingListener {
-  val registeredReceiverStreamIds = new ArrayBuffer[Int] with SynchronizedBuffer[Int]
   val startedReceiverStreamIds = new ArrayBuffer[Int] with SynchronizedBuffer[Int]
   val stoppedReceiverStreamIds = new ArrayBuffer[Int] with SynchronizedBuffer[Int]
   val receiverErrors =
     new ArrayBuffer[(Int, String, String)] with SynchronizedBuffer[(Int, String, String)]
 
-  override def onReceiverRegistered(receiverRegistered: StreamingListenerReceiverRegistered) {
-    registeredReceiverStreamIds += receiverRegistered.receiverInfo.streamId
-  }
-
   override def onReceiverStarted(receiverStarted: StreamingListenerReceiverStarted) {
     startedReceiverStreamIds += receiverStarted.receiverInfo.streamId
   }