Separate register and start events for receivers.

Author: viirya

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

@@ -93,7 +93,10 @@ 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() { }
 
@@ -117,10 +120,11 @@ private[streaming] abstract class ReceiverSupervisor(
   /** Start receiver */
   def startReceiver(): Unit = synchronized {
     try {
-      onReceiverStart()
+      onReceiverRegister()
       logInfo("Starting receiver")
       receiver.onStart()
       logInfo("Called receiver onStart")
+      onReceiverStart()
       receiverState = Started
     } catch {
       case t: Throwable =>

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

@@ -182,13 +182,20 @@ private[streaming] class ReceiverSupervisorImpl(
     blockGenerator.stop()
     env.actorSystem.stop(actor)
   }
-
-  override protected def onReceiverStart() {
+ 
+  override protected def onReceiverRegister() {
     val msg = RegisterReceiver(
       streamId, receiver.getClass.getSimpleName, Utils.localHostName(), actor)
     val future = trackerActor.ask(msg)(askTimeout)
     Await.result(future, askTimeout)
   }
+ 
+  override protected def onReceiverStart() {
+    val msg = ReceiverStarted(
+      streamId, receiver.getClass.getSimpleName, Utils.localHostName(), actor)
+    val future = trackerActor.ask(msg)(askTimeout)
+    Await.result(future, askTimeout)
+  }
 
   override protected def onReceiverStop(message: String, error: Option[Throwable]) {
     logInfo("Deregistering receiver " + streamId)

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

@@ -38,6 +38,12 @@ private[streaming] case class RegisterReceiver(
     host: String,
     receiverActor: ActorRef
   ) extends ReceiverTrackerMessage
+private[streaming] case class ReceiverStarted(
+    streamId: Int,
+    typ: String,
+    host: String,
+    receiverActor: ActorRef
+ ) extends ReceiverTrackerMessage
 private[streaming] case class AddBlock(receivedBlockInfo: ReceivedBlockInfo)
   extends ReceiverTrackerMessage
 private[streaming] case class ReportError(streamId: Int, message: String, error: String)
@@ -177,10 +183,27 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
     }
     receiverInfo(streamId) = ReceiverInfo(
       streamId, s"${typ}-${streamId}", receiverActor, true, host)
-    listenerBus.post(StreamingListenerReceiverStarted(receiverInfo(streamId)))
+    listenerBus.post(StreamingListenerReceiverRegistered(receiverInfo(streamId)))
     logInfo("Registered receiver for stream " + streamId + " from " + sender.path.address)
   }
-
+ 
+  /** Receiver started */
+  private def receiverStarted(
+      streamId: Int,
+      typ: String,
+      host: String,
+      receiverActor: ActorRef,
+      sender: ActorRef
+    ) {
+    if (!receiverInputStreamIds.contains(streamId)) {
+      throw new SparkException("Start received for unexpected id " + streamId)
+    }
+    receiverInfo(streamId) = ReceiverInfo(
+      streamId, s"${typ}-${streamId}", receiverActor, true, host)
+    listenerBus.post(StreamingListenerReceiverStarted(receiverInfo(streamId)))
+    logInfo("Receiver started for stream " + streamId + " from " + sender.path.address)
+  }
+ 
   /** Deregister a receiver */
   private def deregisterReceiver(streamId: Int, message: String, error: String) {
     val newReceiverInfo = receiverInfo.get(streamId) match {
@@ -238,6 +261,12 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
           registerReceiver(streamId, typ, host, receiverActor, sender)
           sender ! true
         }
+      case ReceiverStarted(streamId, typ, host, receiverActor) =>
+        // Actor stops to accept starting when tracker is stopping
+        if (!isTrackerStopping) {
+          receiverStarted(streamId, typ, host, receiverActor, sender)
+          sender ! true
+        }
       case AddBlock(receivedBlockInfo) =>
         sender ! addBlock(receivedBlockInfo)
       case ReportError(streamId, message, error) =>

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

@@ -37,7 +37,11 @@ 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
@@ -57,7 +61,10 @@ 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,6 +31,8 @@ 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 =>