[SPARK-9026] Refactor SimpleFutureAction.onComplete to not launch separate thread for every callback

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

@@ -108,63 +108,48 @@ trait FutureAction[T] extends Future[T] {
 class SimpleFutureAction[T] private[spark](jobWaiter: JobWaiter[_], resultFunc: => T)
   extends FutureAction[T] {
 
+  // Note: `resultFunc` is a closure which may contain references to state that's updated by the
+  // JobWaiter's result handler function. It should only be evaluated once the job has succeeded.
+
   @volatile private var _cancelled: Boolean = false
+  private[this] val jobWaiterFuture: Future[Unit] = jobWaiter.toFuture
+  private[this] lazy val resultFuncOutput: T = {
+    assert(isCompleted, "resultFunc should only be evaluated after the job has completed")
+    resultFunc
+  }
 
   override def cancel() {
     _cancelled = true
     jobWaiter.cancel()
   }
 
   override def ready(atMost: Duration)(implicit permit: CanAwait): SimpleFutureAction.this.type = {
-    if (!atMost.isFinite()) {
-      awaitResult()
-    } else jobWaiter.synchronized {
-      val finishTime = System.currentTimeMillis() + atMost.toMillis
-      while (!isCompleted) {
-        val time = System.currentTimeMillis()
-        if (time >= finishTime) {
-          throw new TimeoutException
-        } else {
-          jobWaiter.wait(finishTime - time)
-        }
-      }
-    }
+    jobWaiterFuture.ready(atMost)(permit) // Throws exception if the job failed.
     this
   }
 
   @throws(classOf[Exception])
   override def result(atMost: Duration)(implicit permit: CanAwait): T = {
-    ready(atMost)(permit)
-    awaitResult() match {
-      case scala.util.Success(res) => res
-      case scala.util.Failure(e) => throw e
-    }
+    jobWaiterFuture.result(atMost)(permit) // Throws exception if the job failed.
+    resultFuncOutput // This function is safe to evaluate because the job must have succeeded.
   }
 
-  override def onComplete[U](func: (Try[T]) => U)(implicit executor: ExecutionContext) {
-    executor.execute(new Runnable {
-      override def run() {
-        func(awaitResult())
-      }
-    })
+  override def onComplete[U](func: (Try[T]) => U)(implicit executor: ExecutionContext): Unit = {
+    jobWaiterFuture.map { _ => resultFuncOutput }.onComplete(func)
   }
 
   override def isCompleted: Boolean = jobWaiter.jobFinished
 
   override def isCancelled: Boolean = _cancelled
 
   override def value: Option[Try[T]] = {
-    if (jobWaiter.jobFinished) {
-      Some(awaitResult())
-    } else {
+    if (!isCompleted) {
       None
-    }
-  }
-
-  private def awaitResult(): Try[T] = {
-    jobWaiter.awaitResult() match {
-      case JobSucceeded => scala.util.Success(resultFunc)
-      case JobFailed(e: Exception) => scala.util.Failure(e)
+    } else {
+      jobWaiter.awaitResult() match {
+        case JobSucceeded => Some(scala.util.Success(resultFuncOutput))
+        case JobFailed(e) => Some(scala.util.Failure(e))
+      }
     }
   }
 

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

@@ -17,6 +17,9 @@
 
 package org.apache.spark.scheduler
 
+import scala.concurrent.{Future, Promise}
+import scala.util.Success
+
 /**
  * An object that waits for a DAGScheduler job to complete. As tasks finish, it passes their
  * results to the given handler function.
@@ -28,12 +31,18 @@ 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)?
   @volatile
   private var _jobFinished = totalTasks == 0
 
+  if (_jobFinished) {
+    promise.complete(Success(Unit))
+  }
+
   def jobFinished: Boolean = _jobFinished
 
   // If the job is finished, this will be its result. In the case of 0 task jobs (e.g. zero
@@ -58,13 +67,15 @@ private[spark] class JobWaiter[T](
     if (finishedTasks == totalTasks) {
       _jobFinished = true
       jobResult = JobSucceeded
+      promise.trySuccess()
       this.notifyAll()
     }
   }
 
   override def jobFailed(exception: Exception): Unit = synchronized {
     _jobFinished = true
     jobResult = JobFailed(exception)
+    promise.tryFailure(exception)
     this.notifyAll()
   }
 
@@ -74,4 +85,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
 }

core/src/test/scala/org/apache/spark/FutureActionSuite.scala

@@ -17,11 +17,12 @@
 
 package org.apache.spark
 
-import scala.concurrent.Await
+import scala.concurrent.{ExecutionContext, Await}
 import scala.concurrent.duration.Duration
 
 import org.scalatest.{BeforeAndAfter, Matchers}
 
+import org.apache.spark.util.ThreadUtils
 
 class FutureActionSuite
   extends SparkFunSuite
@@ -49,4 +50,20 @@ class FutureActionSuite
     job.jobIds.size should be (2)
   }
 
+  test("simple async action callbacks should not tie up execution context threads (SPARK-9026)") {
+    val rdd = sc.parallelize(1 to 10, 2).map(_ => Thread.sleep(1000 * 1000))
+    val pool = ThreadUtils.newDaemonCachedThreadPool("SimpleFutureActionTest")
+    val executionContext = ExecutionContext.fromExecutorService(pool)
+    val job = rdd.countAsync()
+    try {
+      for (_ <- 1 to 10) {
+        job.onComplete(_ => ())(executionContext)
+        assert(pool.getLargestPoolSize < 10)
+      }
+    } finally {
+      job.cancel()
+      executionContext.shutdownNow()
+    }
+  }
+
 }