Github user zsxwing commented on a diff in the pull request:
https://github.com/apache/spark/pull/516#discussion_r11936630
--- Diff:
core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala ---
@@ -177,6 +177,7 @@ class SparkListenerSuite extends FunSuite with
LocalSparkContext with ShouldMatc
listener.stageInfos.clear()
rdd3.count()
+ assert(sc.listenerBus.waitUntilEmpty(WAIT_TIMEOUT_MILLIS))
--- End diff --
Since `waitUntilEmpty` is only for test, is it possible that moving such
`wait` logic to the SparkListener instances used in tests? E.g.,
```Scala
import java.util.concurrent.{CountDownLatch, TimeUnit}
class SaveStageAndTaskInfo extends SparkListener {
val stageInfos = mutable.Map[StageInfo, Seq[(TaskInfo, TaskMetrics)]]()
var taskInfoMetrics = mutable.Buffer[(TaskInfo, TaskMetrics)]()
val latch = new CountDownLatch(1)
override def onTaskEnd(task: SparkListenerTaskEnd) {
val info = task.taskInfo
val metrics = task.taskMetrics
if (info != null && metrics != null) {
taskInfoMetrics += ((info, metrics))
}
}
override def onStageCompleted(stage: SparkListenerStageCompleted) {
stageInfos(stage.stageInfo) = taskInfoMetrics
taskInfoMetrics = mutable.Buffer.empty
latch.countDown()
}
def waitForCompleted(timeoutMillis: Long) {
latch.await(timeoutMillis, TimeUnit.MILLISECONDS)
}
}
```
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---
