Blockingly drain all events in LiveListenerBus#stop().

Also, move bus.stop() call to the end of sc.stop(), so that slow
listeners would not affect the cleaning of other states.

Author: concretevitamin

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

@@ -835,13 +835,13 @@ class SparkContext(
     if (dagSchedulerCopy != null) {
       metadataCleaner.cancel()
       dagSchedulerCopy.stop()
-      listenerBus.stop()
       taskScheduler = null
       // TODO: Cache.stop()?
       env.stop()
       SparkEnv.set(null)
       ShuffleMapTask.clearCache()
       ResultTask.clearCache()
+      listenerBus.stop()
       logInfo("Successfully stopped SparkContext")
     } else {
       logInfo("SparkContext already stopped")

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

@@ -37,6 +37,9 @@ private[spark] class LiveListenerBus extends SparkListenerBus with Logging {
   private var queueFullErrorMessageLogged = false
   private var started = false
 
+  private var drained = false
+  private val drainedLock = new Object()
+
   /**
    * Start sending events to attached listeners.
    *
@@ -55,6 +58,10 @@ private[spark] class LiveListenerBus extends SparkListenerBus with Logging {
         while (true) {
           val event = eventQueue.take
           if (event == SparkListenerShutdown) {
+            drainedLock.synchronized {
+              drained = true
+              drainedLock.notify()
+            }
             // Get out of the while loop and shutdown the daemon thread
             return
           }
@@ -92,10 +99,21 @@ private[spark] class LiveListenerBus extends SparkListenerBus with Logging {
     true
   }
 
+  /**
+   * Stop the listener bus; wait until all listener events are processed by the listener bus
+   * thread. The user has to make sure the listeners finish in a reasonable amount of time.
+   */
   def stop() {
     if (!started) {
       throw new IllegalStateException("Attempted to stop a listener bus that has not yet started!")
     }
-    post(SparkListenerShutdown)
+    drainedLock.synchronized {
+      // put post() and wait() in the same synchronized block to ensure wait() happens before
+      // notify()
+      post(SparkListenerShutdown)
+      while (!drained) {
+        drainedLock.wait()
+      }
+    }
   }
 }

core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala

@@ -72,6 +72,20 @@ class SparkListenerSuite extends FunSuite with LocalSparkContext with ShouldMatc
     }
   }
 
+  test("bus.stop() waits for the event queue to completely drain") {
+    val sleepyListener = new SleepyListener(1000)
+    val bus = new LiveListenerBus
+    bus.addListener(sleepyListener)
+    (1 to 5).foreach { _ => bus.post(SparkListenerJobEnd(0, JobSucceeded)) }
+
+    bus.start()
+    // since the handler is just thread sleep, the queue should not drain immediately
+    assert(!bus.waitUntilEmpty(0))
+    bus.stop()
+    // bus.stop() should wait until the event queue is drained, ensuring no events are lost
+    assert(bus.waitUntilEmpty(0))
+  }
+
   test("basic creation of StageInfo") {
     val listener = new SaveStageAndTaskInfo
     sc.addSparkListener(listener)
@@ -282,4 +296,11 @@ class SparkListenerSuite extends FunSuite with LocalSparkContext with ShouldMatc
       startedGettingResultTasks += taskGettingResult.taskInfo.index
     }
   }
+
+  class SleepyListener(val sleepTime: Long) extends SparkListener {
+    override def onJobEnd(job: SparkListenerJobEnd) = {
+      Thread.sleep(sleepTime)
+    }
+  }
+
 }