Add a stop flag and some tests

Author: zsxwing

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

@@ -17,6 +17,7 @@
 
 package org.apache.spark.util
 
+import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.{BlockingQueue, LinkedBlockingDeque}
 
 import scala.util.control.NonFatal
@@ -34,12 +35,14 @@ private[spark] abstract class EventLoop[E](name: String) extends Logging {
 
   private val eventQueue: BlockingQueue[E] = new LinkedBlockingDeque[E]()
 
+  private val stopped = new AtomicBoolean(false)
+
   private val eventThread = new Thread(name) {
     setDaemon(true)
 
     override def run(): Unit = {
       try {
-        while (true) {
+        while (!stopped.get) {
           val event = eventQueue.take()
           try {
             onReceive(event)
@@ -62,16 +65,23 @@ private[spark] abstract class EventLoop[E](name: String) extends Logging {
   }
 
   def start(): Unit = {
+    if (stopped.get) {
+      throw new IllegalStateException(name + " has already been stopped")
+    }
     // Call onStart before starting the event thread to make sure it happens before onReceive
     onStart()
     eventThread.start()
   }
 
   def stop(): Unit = {
-    eventThread.interrupt()
-    eventThread.join()
-    // Call onStop after the event thread exits to make sure onReceive happens before onStop
-    onStop()
+    if (stopped.compareAndSet(false ,true)) {
+      eventThread.interrupt()
+      eventThread.join()
+      // Call onStop after the event thread exits to make sure onReceive happens before onStop
+      onStop()
+    } else {
+      // Keep quiet to allow calling `stop` multiple times.
+    }
   }
 
   /**

core/src/test/scala/org/apache/spark/util/EventLoopSuite.scala

@@ -80,4 +80,81 @@ class EventLoopSuite extends FunSuite {
     }
     eventLoop.stop()
   }
+
+  test("EventLoop: error thrown from onError should not crash the event thread") {
+    val e = new RuntimeException("Oops")
+    val receivedError = new AtomicReference[Throwable]()
+    val eventLoop = new EventLoop[Int]("test") {
+
+      override def onReceive(event: Int): Unit = {
+        throw e
+      }
+
+      override def onError(e: Throwable): Unit = {
+        receivedError.set(e)
+        throw new RuntimeException("Oops")
+      }
+    }
+    eventLoop.start()
+    eventLoop.post(1)
+    eventually(timeout(5 seconds), interval(200 millis)) {
+      assert(e === receivedError.get)
+      assert(eventLoop.isActive)
+    }
+    eventLoop.stop()
+  }
+
+  test("EventLoop: calling stop multiple times should only call onStop once") {
+    var onStopTimes = 0
+    val eventLoop = new EventLoop[Int]("test") {
+
+      override def onReceive(event: Int): Unit = {
+      }
+
+      override def onError(e: Throwable): Unit = {
+      }
+
+      override def onStop(): Unit = {
+        onStopTimes += 1
+      }
+    }
+
+    eventLoop.start()
+
+    eventLoop.stop()
+    eventLoop.stop()
+    eventLoop.stop()
+
+    assert(1 === onStopTimes)
+  }
+
+  test("EventLoop: post event in multiple threads") {
+    var receivedEventsCount = 0
+    val eventLoop = new EventLoop[Int]("test") {
+
+      override def onReceive(event: Int): Unit = {
+        receivedEventsCount += 1
+      }
+
+      override def onError(e: Throwable): Unit = {
+      }
+
+    }
+    eventLoop.start()
+
+    val threadNum = 5
+    val eventsFromEachThread = 100
+    (1 to threadNum).foreach { _ =>
+      new Thread() {
+        override def run(): Unit = {
+          (1 to eventsFromEachThread).foreach(eventLoop.post)
+        }
+      }.start()
+    }
+
+    eventually(timeout(5 seconds), interval(200 millis)) {
+      assert(threadNum * eventsFromEachThread === receivedEventsCount)
+    }
+    eventLoop.stop()
+  }
 }