Move ReceiverTracker.stop into ReceiverTracker.stop

Author: zsxwing

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

@@ -106,7 +106,30 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
     if (isTrackerStarted) {
       // First, stop the receivers
       trackerState = Stopping
-      if (!skipReceiverLaunch) receiverExecutor.stop(graceful)
+      if (!skipReceiverLaunch) {
+        // Send the stop signal to all the receivers
+        endpoint.askWithRetry[Boolean](StopAllReceivers)
+
+        // Wait for the Spark job that runs the receivers to be over
+        // That is, for the receivers to quit gracefully.
+        receiverExecutor.awaitTermination(10000)
+
+        if (graceful) {
+          val pollTime = 100
+          logInfo("Waiting for receiver job to terminate gracefully")
+          while (receiverInfo.nonEmpty || receiverExecutor.running) {
+            Thread.sleep(pollTime)
+          }
+          logInfo("Waited for receiver job to terminate gracefully")
+        }
+
+        // Check if all the receivers have been deregistered or not
+        if (receiverInfo.nonEmpty) {
+          logWarning("Not all of the receivers have deregistered, " + receiverInfo)
+        } else {
+          logInfo("All of the receivers have deregistered successfully")
+        }
+      }
 
       // Finally, stop the endpoint
       ssc.env.rpcEnv.stop(endpoint)
@@ -269,7 +292,7 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
   /** This thread class runs all the receivers on the cluster.  */
   class ReceiverLauncher {
     @transient val env = ssc.env
-    @volatile @transient private var running = false
+    @volatile @transient var running = false
     @transient val thread = new Thread() {
       override def run() {
         try {
@@ -285,31 +308,6 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
       thread.start()
     }
 
-    def stop(graceful: Boolean) {
-      // Send the stop signal to all the receivers
-      endpoint.askWithRetry[Boolean](StopAllReceivers)
-
-      // Wait for the Spark job that runs the receivers to be over
-      // That is, for the receivers to quit gracefully.
-      thread.join(10000)
-
-      if (graceful) {
-        val pollTime = 100
-        logInfo("Waiting for receiver job to terminate gracefully")
-        while (receiverInfo.nonEmpty || running) {
-          Thread.sleep(pollTime)
-        }
-        logInfo("Waited for receiver job to terminate gracefully")
-      }
-
-      // Check if all the receivers have been deregistered or not
-      if (receiverInfo.nonEmpty) {
-        logWarning("Not all of the receivers have deregistered, " + receiverInfo)
-      } else {
-        logInfo("All of the receivers have deregistered successfully")
-      }
-    }
-
     /**
      * Get the list of executors excluding driver
      */
@@ -402,6 +400,13 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
       }
     }
 
+    /**
+     * Wait until the Spark job that runs the receivers is terminated, or return when
+     * `milliseconds` elapses
+     */
+    def awaitTermination(milliseconds: Long): Unit = {
+      thread.join(milliseconds)
+    }
   }
 
   /** Check if tracker has been marked for starting */