Add JobWaiter.toFuture to avoid blocking threads

Author: zsxwing

core/src/main/scala/org/apache/spark/SparkContext.scala

@@ -17,8 +17,6 @@
 
 package org.apache.spark
 
-import scala.language.implicitConversions
-
 import java.io._
 import java.lang.reflect.Constructor
 import java.net.URI
@@ -30,6 +28,8 @@ import scala.collection.{Map, Set}
 import scala.collection.JavaConversions._
 import scala.collection.generic.Growable
 import scala.collection.mutable.HashMap
+import scala.concurrent.Future
+import scala.language.implicitConversions
 import scala.reflect.{ClassTag, classTag}
 import scala.util.control.NonFatal
 
@@ -1860,13 +1860,14 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
   }
 
   /**
-   * Submit a job for execution and return a FutureJob holding the result.
+   * Submit a job for execution and return a FutureJob holding the result. Return a Future for
+   * monitoring the job success or failure event.
    */
   private[spark] def submitAsyncJob[T, U, R](
       rdd: RDD[T],
       processPartition: (TaskContext, Iterator[T]) => U,
       resultHandler: (Int, U) => Unit,
-      resultFunc: => R): SimpleFutureAction[R] =
+      resultFunc: => R): Future[Unit] =
   {
     assertNotStopped()
     val cleanF = clean(processPartition)
@@ -1879,7 +1880,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
       allowLocal = false,
       resultHandler,
       localProperties.get)
-    new SimpleFutureAction(waiter, resultFunc)
+    waiter.toFuture
   }
 
   /**

core/src/main/scala/org/apache/spark/scheduler/JobWaiter.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.scheduler
 
+import scala.concurrent.{Future, Promise}
+
 /**
  * An object that waits for a DAGScheduler job to complete. As tasks finish, it passes their
  * results to the given handler function.
@@ -28,6 +30,8 @@ private[spark] class JobWaiter[T](
     resultHandler: (Int, T) => Unit)
   extends JobListener {
 
+  private val promise = Promise[Unit]
+
   private var finishedTasks = 0
 
   // Is the job as a whole finished (succeeded or failed)?
@@ -58,13 +62,15 @@ private[spark] class JobWaiter[T](
     if (finishedTasks == totalTasks) {
       _jobFinished = true
       jobResult = JobSucceeded
+      promise.success()
       this.notifyAll()
     }
   }
 
   override def jobFailed(exception: Exception): Unit = synchronized {
     _jobFinished = true
     jobResult = JobFailed(exception)
+    promise.failure(exception)
     this.notifyAll()
   }
 
@@ -74,4 +80,10 @@ private[spark] class JobWaiter[T](
     }
     return jobResult
   }
+
+  /**
+   * Return a Future to monitoring the job success or failure event. You can use this method to
+   * avoid blocking your thread.
+   */
+  def toFuture: Future[Unit] = promise.future
 }

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

@@ -163,7 +163,8 @@ private[streaming] abstract class ReceiverSupervisor(
       stopReceiver("Restarting receiver with delay " + delay + "ms: " + message, error)
       logDebug("Sleeping for " + delay)
       Thread.sleep(delay)
-      if (rescheduleReceiver().contains(host)) {
+      val scheduledLocations = rescheduleReceiver()
+      if (scheduledLocations.isEmpty || scheduledLocations.contains(host)) {
         logInfo("Starting receiver again")
         startReceiver()
         logInfo("Receiver started again")
@@ -174,7 +175,7 @@ private[streaming] abstract class ReceiverSupervisor(
   }
 
   /** Reschedule this receiver and return a candidate executor list */
-  def rescheduleReceiver(): Seq[String]
+  def rescheduleReceiver(): Seq[String] = Seq.empty
 
   /** Check if receiver has been marked for stopping */
   def isReceiverStarted(): Boolean = {

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

@@ -20,6 +20,7 @@ package org.apache.spark.streaming.scheduler
 import scala.collection.mutable
 import scala.util.Random
 
+import org.apache.spark.streaming.StreamingContext
 import org.apache.spark.streaming.scheduler.ReceiverState._
 
 private[streaming] case class ReceiverTrackingInfo(
@@ -31,6 +32,10 @@ private[streaming] case class ReceiverTrackingInfo(
 
 private[streaming] trait ReceiverScheduler {
 
+  /**
+   * Return a candidate executor list to run the receiver. If the list is empty, the caller can run
+   * this receiver in arbitrary executor.
+   */
   def scheduleReceiver(
       receiverId: Int,
       preferredLocation: Option[String],

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

@@ -333,7 +333,7 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
     }
 
     private val submitJobThread = ExecutionContext.fromExecutorService(
-      ThreadUtils.newDaemonCachedThreadPool("streaming-submit-job"))
+      ThreadUtils.newDaemonSingleThreadExecutor("streaming-submit-job"))
 
     /**
      * Get the receivers from the ReceiverInputDStreams, distributes them to the
@@ -390,8 +390,9 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
             }
 
             val self = this
+            val receiverId = receiver.streamId
             val scheduledLocations = scheduler.scheduleReceiver(
-              receiver.streamId,
+              receiverId,
               receiver.preferredLocation,
               getReceiverTrackingInfoMap(),
               getExecutors(ssc))
@@ -411,18 +412,18 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
                 if (stopping) {
                   receiverExitLatch.countDown()
                 } else {
-                  logInfo(s"Restarting Receiver ${receiver.streamId}")
+                  logInfo(s"Restarting Receiver $receiverId")
                   submitJobThread.execute(self)
                 }
               case Failure(e) =>
                 if (stopping) {
                   receiverExitLatch.countDown()
                 } else {
                   logError("Receiver has been stopped. Try to restart it.", e)
-                  logInfo(s"Restarting Receiver ${receiver.streamId}")
+                  logInfo(s"Restarting Receiver $receiverId")
                   submitJobThread.execute(self)
                 }
-            }(submitJobThread)
+            }(ThreadUtils.sameThread)
             logInfo(s"Receiver ${receiver.streamId} started")
           }
         })