Skip to content

Commit

Permalink
Add a stop flag and some tests
Browse files Browse the repository at this point in the history
  • Loading branch information
zsxwing committed Jan 14, 2015
1 parent 37f79c6 commit 227bf33
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 5 deletions.
20 changes: 15 additions & 5 deletions core/src/main/scala/org/apache/spark/util/EventLoop.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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.
}
}

/**
Expand Down
77 changes: 77 additions & 0 deletions core/src/test/scala/org/apache/spark/util/EventLoopSuite.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
}

0 comments on commit 227bf33

Please sign in to comment.