viirya commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4250754101
##########
core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala:
##########
@@ -757,6 +757,113 @@ class TaskContextSuite extends SparkFunSuite with
BeforeAndAfter with LocalSpark
assert(invocationOrder === Seq("C", "B", "A", "D"))
}
+ test("SPARK-60113: pre-completion listeners run before all completion
listeners") {
+ val context = TaskContext.empty()
+ val invocationOrder = ArrayBuffer.empty[String]
+
+ def makeLoggingListener(id: String): TaskCompletionListener = new
TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = {
+ invocationOrder += id
+ }
+ }
+ context.addTaskCompletionListener(makeLoggingListener("A"))
+ context.addTaskPreCompletionListener(makeLoggingListener("P1"))
+ context.addTaskPreCompletionListener(makeLoggingListener("P2"))
+ // Added after the pre-completion listeners, e.g. lazily by the thread
they stop. An ordinary
+ // completion listener added here would run before them.
+ context.addTaskCompletionListener(makeLoggingListener("B"))
+
+ context.markTaskCompleted(None)
+ assert(invocationOrder === Seq("P2", "P1", "B", "A"))
+ // Called immediately once the task has completed, like completion
listeners.
+ context.addTaskPreCompletionListener(makeLoggingListener("P3"))
+ assert(invocationOrder === Seq("P2", "P1", "B", "A", "P3"))
+ }
+
+ test("SPARK-60113: completion listeners run even if a pre-completion
listener fails") {
+ val context = TaskContext.empty()
+ val listener = mock(classOf[TaskCompletionListener])
+ context.addTaskCompletionListener(listener)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = throw new
Exception("pre")
+ })
+
+ val e = intercept[TaskCompletionListenerException] {
+ context.markTaskCompleted(None)
+ }
+ assert(e.getMessage.contains("pre"))
+ verify(listener, times(1)).onTaskCompletion(any())
+ assert(context.isFailed())
+ }
+
+ test("SPARK-60113: failures of pre-completion and completion listeners are
both reported") {
+ val context = TaskContext.empty()
+ context.addTaskCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = throw new
Exception("post")
+ })
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = throw new
Exception("pre")
+ })
+
+ val e = intercept[TaskCompletionListenerException] {
+ context.markTaskCompleted(None)
+ }
+ assert(e.getMessage.contains("pre"))
+ assert(e.getMessage.contains("post"))
+ }
+
+ test("SPARK-60113: a pre-completion listener added during completion runs
next") {
+ val context = TaskContext.empty()
+ val invocationOrder = ArrayBuffer.empty[String]
+
+ def makeLoggingListener(id: String): TaskCompletionListener = new
TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = {
+ invocationOrder += id
+ }
+ }
+ context.addTaskCompletionListener(makeLoggingListener("A"))
+ context.addTaskCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = {
+ invocationOrder += "B"
+ // Added from a second thread while this thread runs the completion
listeners, so the
+ // call returns without running it; this thread must still run it.
+ val thread = new Thread(() => {
+ context.addTaskPreCompletionListener(makeLoggingListener("P"))
+ })
+ thread.start()
+ thread.join()
+ assert(invocationOrder === Seq("B"))
+ }
+ })
+
+ context.markTaskCompleted(None)
+ assert(invocationOrder === Seq("B", "P", "A"))
+ }
+
+ test("SPARK-60113: a listener added as completion starts does not take over
the listeners") {
Review Comment:
When I reverted the `markTaskCompleted` claim while keeping the test, it
failed on the first run, so in practice the spin lands in the window well
within 500 iterations. With the claim in place there is no window left between
setting `completed` and claiming the invocation, so there is nothing to block
on there. For the new handoffs, the test with another thread holding the
invocation is deterministic (the holder blocks in a failure listener). The
failure-path window is a few instructions wide in the old code as well, so that
test also relies on contention; it too fails on the first run with the old code.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -1336,9 +1409,9 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
// automatically closed by the JVM, which will cause Python to
receive
// EOF and the reader thread to get IOException.
Thread.currentThread().interrupt()
- case NonFatal(t) =>
+ case t: Throwable =>
Review Comment:
The `setException` path is what keeps the task from waiting on the Python
worker with no cause recorded, so I kept it, and added a comment that the error
is then reported twice (the writer thread's log and the reader's rethrow on the
task thread). `setException` is a field store, and `shutdownOutput()` goes
through `Utils.tryLog`, which catches and logs any failure, so the original
error is still rethrown and logged afterwards.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -128,6 +134,20 @@ private[spark] class TaskContextImpl(
this
}
+ override private[spark] def addTaskPreCompletionListener(
Review Comment:
Done: both now go through a private `addCompletionListenerTo(stack,
listener)`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
