only remove the active listener for an interrupt

Author: squito

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

@@ -85,28 +85,18 @@ private class AsyncEventQueue(
   }
 
   private def dispatch(): Unit = LiveListenerBus.withinListenerThread.withValue(true) {
-    try {
-      var next: SparkListenerEvent = eventQueue.take()
-      while (next != POISON_PILL) {
-        val ctx = processingTime.time()
-        try {
-          super.postToAll(next)
-        } finally {
-          ctx.stop()
-        }
-        eventCount.decrementAndGet()
-        next = eventQueue.take()
+    var next: SparkListenerEvent = eventQueue.take()
+    while (next != POISON_PILL) {
+      val ctx = processingTime.time()
+      try {
+        super.postToAll(next)
+      } finally {
+        ctx.stop()
       }
       eventCount.decrementAndGet()
-    } catch {
-      case ie: InterruptedException =>
-        logInfo(s"Stopping listener queue $name.", ie)
-        stopped.set(true)
-        bus.removeQueue(name)
-        // we're not going to process any more events in this queue, so might as well clear it
-        eventQueue.clear()
-        eventCount.set(0)
+      next = eventQueue.take()
     }
+    eventCount.decrementAndGet()
   }
 
   override protected def getTimer(listener: SparkListenerInterface): Option[Timer] = {
@@ -139,7 +129,11 @@ private class AsyncEventQueue(
       eventCount.incrementAndGet()
       eventQueue.put(POISON_PILL)
     }
-    dispatchThread.join()
+    // this thread might be trying to stop itself as part of error handling -- we can't join
+    // in that case.
+    if (Thread.currentThread() != dispatchThread) {
+      dispatchThread.join()
+    }
   }
 
   def post(event: SparkListenerEvent): Unit = {
@@ -196,6 +190,12 @@ private class AsyncEventQueue(
     true
   }
 
+  override def removeListenerOnError(listener: SparkListenerInterface): Unit = {
+    // the listener failed in an unrecoverably way, we want to remove it from the entire
+    // LiveListenerBus (potentially stopping a queue if its empty)
+    bus.removeListener(listener)
+  }
+
 }
 
 private object AsyncEventQueue {

core/src/main/scala/org/apache/spark/util/ListenerBus.scala

@@ -60,6 +60,15 @@ private[spark] trait ListenerBus[L <: AnyRef, E] extends Logging {
     }
   }
 
+  /**
+   * This can be overriden by subclasses if there is any extra cleanup to do when removing a
+   * listener.  In particular AsyncEventQueues can clean up queues in the LiveListenerBus.
+   */
+  def removeListenerOnError(listener: L): Unit = {
+    removeListener(listener)
+  }
+
+
   /**
    * Post the event to all registered listeners. The `postToAll` caller should guarantee calling
    * `postToAll` in the same thread for all events.
@@ -80,6 +89,11 @@ private[spark] trait ListenerBus[L <: AnyRef, E] extends Logging {
       }
       try {
         doPostEvent(listener, event)
+        if (Thread.interrupted()) {
+          logError(s"Interrupted while posting to ${Utils.getFormattedClassName(listener)}.  " +
+            s"Removing that listener.")
+          removeListenerOnError(listener)
+        }
       } catch {
         case NonFatal(e) if !isIgnorableException(e) =>
           logError(s"Listener ${Utils.getFormattedClassName(listener)} threw an exception", e)

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

@@ -495,25 +495,32 @@ class SparkListenerSuite extends SparkFunSuite with LocalSparkContext with Match
     val bus = new LiveListenerBus(conf)
     val counter1 = new BasicJobCounter()
     val counter2 = new BasicJobCounter()
-    val interruptingListener = new InterruptingListener
+    val interruptingListener1 = new InterruptingListener
+    val interruptingListener2 = new InterruptingListener
     bus.addToSharedQueue(counter1)
-    bus.addToSharedQueue(interruptingListener)
+    bus.addToSharedQueue(interruptingListener1)
     bus.addToStatusQueue(counter2)
-    assert(bus.activeQueues() === Set(SHARED_QUEUE, APP_STATUS_QUEUE))
+    bus.addToEventLogQueue(interruptingListener2)
+    assert(bus.activeQueues() === Set(SHARED_QUEUE, APP_STATUS_QUEUE, EVENT_LOG_QUEUE))
     assert(bus.findListenersByClass[BasicJobCounter]().size === 2)
+    assert(bus.findListenersByClass[InterruptingListener]().size === 2)
 
     bus.start(mockSparkContext, mockMetricsSystem)
 
-    // after we post one event, the shared queue should stop because of the interrupt
+    // after we post one event, both interrupting listeners should get removed, and the
+    // event log queue should be removed
     bus.post(SparkListenerJobEnd(0, jobCompletionTime, JobSucceeded))
     bus.waitUntilEmpty(WAIT_TIMEOUT_MILLIS)
-    assert(bus.activeQueues() === Set(APP_STATUS_QUEUE))
-    assert(bus.findListenersByClass[BasicJobCounter]().size === 1)
+    assert(bus.activeQueues() === Set(SHARED_QUEUE, APP_STATUS_QUEUE))
+    assert(bus.findListenersByClass[BasicJobCounter]().size === 2)
+    assert(bus.findListenersByClass[InterruptingListener]().size === 0)
+    assert(counter1.count === 1)
     assert(counter2.count === 1)
 
     // posting more events should be fine, they'll just get processed from the OK queue.
     (0 until 5).foreach { _ => bus.post(SparkListenerJobEnd(0, jobCompletionTime, JobSucceeded)) }
     bus.waitUntilEmpty(WAIT_TIMEOUT_MILLIS)
+    assert(counter1.count === 6)
     assert(counter2.count === 6)
 
     // Make sure stopping works -- this requires putting a poison pill in all active queues, which