You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2026/04/12 22:39:37 UTC
(spark) branch master updated: [SPARK-56330][CORE][FOLLOWUP] Use dedicated loop for TaskInterruptListeners to avoid blocking completion listeners
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 4db309d8fd4b [SPARK-56330][CORE][FOLLOWUP] Use dedicated loop for TaskInterruptListeners to avoid blocking completion listeners
4db309d8fd4b is described below
commit 4db309d8fd4b7e0b1f866dcb40bfda3e79357c6d
Author: Wenchen Fan <we...@databricks.com>
AuthorDate: Mon Apr 13 07:39:27 2026 +0900
[SPARK-56330][CORE][FOLLOWUP] Use dedicated loop for TaskInterruptListeners to avoid blocking completion listeners
### What changes were proposed in this pull request?
Followup to https://github.com/apache/spark/pull/55151.
Changed `invokeTaskInterruptListeners` in `TaskContextImpl` to use a dedicated loop instead of routing through the shared `invokeListeners()` method. Also removed the now-unnecessary `markTaskFailedOnListenerError` parameter from `invokeListeners` since it is always `true` (only completion and failure listeners use this method now).
### Why are the changes needed?
The current implementation routes `TaskInterruptListener`s through `invokeListeners()`, which uses `listenerInvocationThread` to serialize all listener invocations (completion, failure, and interrupt). This creates a problem when `markInterrupted()` is called on the kill thread while completion listeners need to run on the task thread:
1. The kill thread enters `invokeListeners()` for interrupt callbacks and acquires `listenerInvocationThread`.
2. While the interrupt listener runs, the task thread completes and calls `markTaskCompleted()` → `invokeTaskCompletionListeners()` → `invokeListeners()` for completion callbacks.
3. `invokeListeners()` sees `listenerInvocationThread` is held by the kill thread and returns immediately — **silently skipping all completion listeners**.
This causes resource leaks because cleanup logic registered via `addTaskCompletionListener` (e.g., closing file handles, releasing caches, freeing native memory) never executes.
The fix uses a dedicated loop for interrupt listeners with independent serialization (synchronized access to the callback stack, but no `listenerInvocationThread` gate), so interrupt listeners and completion/failure listeners can run on different threads without blocking each other.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added a regression test `"SPARK-56330: task completion during interrupt listener execution"` that:
1. Registers a completion listener and a blocking interrupt listener.
2. Interrupts the task from a separate thread and waits for the interrupt listener to start.
3. Marks the task completed on the main thread while the interrupt listener is still running.
4. Verifies the completion listener was called (it would be silently skipped with the shared `invokeListeners()` approach).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
Closes #55292 from cloud-fan/SPARK-56330-followup.
Authored-by: Wenchen Fan <we...@databricks.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
.../scala/org/apache/spark/TaskContextImpl.scala | 94 +++++++++++-----------
.../apache/spark/scheduler/TaskContextSuite.scala | 35 ++++++++
2 files changed, 84 insertions(+), 45 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/TaskContextImpl.scala b/core/src/main/scala/org/apache/spark/TaskContextImpl.scala
index b2dffd795f93..24202d92601d 100644
--- a/core/src/main/scala/org/apache/spark/TaskContextImpl.scala
+++ b/core/src/main/scala/org/apache/spark/TaskContextImpl.scala
@@ -77,7 +77,7 @@ private[spark] class TaskContextImpl(
@transient private val onInterruptCallbacks = new Stack[TaskInterruptListener]
/**
- * The thread currently executing task completion, failure, or interrupt listeners, if any.
+ * The thread currently executing task completion or failure listeners, if any.
*
* `invokeListeners()` uses this to ensure listeners are called sequentially.
*/
@@ -131,12 +131,11 @@ private[spark] class TaskContextImpl(
override def addTaskInterruptListener(listener: TaskInterruptListener): this.type = {
synchronized {
- // If there is already a thread invoking listeners, adding the new listener to
- // `onInterruptCallbacks` will cause that thread to execute the new listener, and the call to
- // `invokeTaskInterruptListeners()` below will be a no-op.
+ // If another thread is already running `invokeTaskInterruptListeners`, adding the new
+ // listener to `onInterruptCallbacks` will cause that thread to execute it (the loop pops
+ // listeners under the TaskContext lock).
//
- // If there is no such thread, the call to `invokeTaskInterruptListeners()` below will execute
- // all listeners, including the new listener.
+ // Otherwise, `invokeTaskInterruptListeners()` below will execute all listeners.
onInterruptCallbacks.push(listener)
reasonIfKilled
}.foreach { reason =>
@@ -172,11 +171,7 @@ private[spark] class TaskContextImpl(
private def invokeTaskCompletionListeners(error: Option[Throwable]): Unit = {
// It is safe to access the reference to `onCompleteCallbacks` without holding the TaskContext
// lock. `invokeListeners()` acquires the lock before accessing the contents.
- invokeListeners(
- onCompleteCallbacks,
- "TaskCompletionListener",
- error,
- markTaskFailedOnListenerError = true) {
+ invokeListeners(onCompleteCallbacks, "TaskCompletionListener", error) {
_.onTaskCompletion(this)
}
}
@@ -184,35 +179,50 @@ private[spark] class TaskContextImpl(
private def invokeTaskFailureListeners(error: Throwable): Unit = {
// It is safe to access the reference to `onFailureCallbacks` without holding the TaskContext
// lock. `invokeListeners()` acquires the lock before accessing the contents.
- invokeListeners(
- onFailureCallbacks,
- "TaskFailureListener",
- Option(error),
- markTaskFailedOnListenerError = true) {
+ invokeListeners(onFailureCallbacks, "TaskFailureListener", Option(error)) {
_.onTaskFailure(this, error)
}
}
private def invokeTaskInterruptListeners(reason: String, error: Throwable): Unit = {
- // It is safe to access the reference to `onInterruptCallbacks` without holding the TaskContext
- // lock. `invokeListeners()` acquires the lock before accessing the contents.
- // Do not call `markTaskFailed` per listener here: the first failure would win in
- // `failureCauseOpt` and mask the aggregate `TaskCompletionListenerException` that
- // `markInterrupted` records after swallowing the thrown exception (SPARK-56330).
- invokeListeners(
- onInterruptCallbacks,
- "TaskInterruptListener",
- Option(error),
- markTaskFailedOnListenerError = false) {
- _.onTaskInterrupted(this, reason)
+ // Do not use `invokeListeners()`. That method uses `listenerInvocationThread` to serialize
+ // all listener invocations, which would prevent task completion or failure listeners from
+ // running if the task completes or fails while executing an interrupt listener (causing
+ // resource leaks such as task-managed resources not being freed).
+ //
+ // Instead, directly execute the task interrupt listeners with independent serialization.
+ // Exceptions are collected and thrown as a TaskCompletionListenerException so that
+ // `markInterrupted` can catch and record the aggregate failure (SPARK-56330).
+ def getNextListener(): Option[TaskInterruptListener] = synchronized {
+ if (onInterruptCallbacks.empty()) {
+ None
+ } else {
+ Some(onInterruptCallbacks.pop())
+ }
+ }
+ val listenerExceptions = new ArrayBuffer[Throwable](2)
+ var listenerOption: Option[TaskInterruptListener] = None
+ while ({listenerOption = getNextListener(); listenerOption.nonEmpty}) {
+ try {
+ listenerOption.get.onTaskInterrupted(this, reason)
+ } catch {
+ case e: Throwable =>
+ listenerExceptions += e
+ logError(log"Error in TaskInterruptListener", e)
+ }
+ }
+ if (listenerExceptions.nonEmpty) {
+ val exception = new TaskCompletionListenerException(
+ listenerExceptions.map(_.getMessage).toSeq, Option(error))
+ listenerExceptions.foreach(exception.addSuppressed)
+ throw exception
}
}
private def invokeListeners[T](
listeners: Stack[T],
name: String,
- error: Option[Throwable],
- markTaskFailedOnListenerError: Boolean)(
+ error: Option[Throwable])(
callback: T => Unit): Unit = {
// This method is subject to two constraints:
//
@@ -255,8 +265,7 @@ private[spark] class TaskContextImpl(
callback(listener)
} catch {
case e: Throwable =>
- // A listener failed. For completion/failure listeners, temporarily clear
- // listenerInvocationThread and markTaskFailed so nested TaskContext calls can run.
+ // A listener failed. Temporarily clear the listenerInvocationThread and markTaskFailed.
//
// One of the following cases applies (#3 being the interesting one):
//
@@ -290,20 +299,15 @@ private[spark] class TaskContextImpl(
// failed, and now another completion listener has failed. Then our call to
// [[markTaskFailed]] here will have no effect and we simply resume running the
// remaining completion handlers.
- //
- // Task interrupt listeners skip per-listener [[markTaskFailed]]; see
- // [[invokeTaskInterruptListeners]].
- if (markTaskFailedOnListenerError) {
- try {
- listenerInvocationThread = None
- markTaskFailed(e)
- } catch {
- case t: Throwable => e.addSuppressed(t)
- } finally {
- synchronized {
- if (listenerInvocationThread.isEmpty) {
- listenerInvocationThread = Some(Thread.currentThread())
- }
+ try {
+ listenerInvocationThread = None
+ markTaskFailed(e)
+ } catch {
+ case t: Throwable => e.addSuppressed(t)
+ } finally {
+ synchronized {
+ if (listenerInvocationThread.isEmpty) {
+ listenerInvocationThread = Some(Thread.currentThread())
}
}
}
diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala
index a232048f445d..ac8d000d8a44 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala
@@ -199,6 +199,41 @@ class TaskContextSuite extends SparkFunSuite with BeforeAndAfter with LocalSpark
assert(invocations == 2)
}
+ test("SPARK-56330: task completion during interrupt listener execution") {
+ val context = TaskContext.empty()
+ val completionListener = mock(classOf[TaskCompletionListener])
+ context.addTaskCompletionListener(completionListener)
+
+ // Add a task interrupt listener that blocks until released.
+ val interruptListenerStarted = new Semaphore(0)
+ val interruptListenerRelease = new Semaphore(0)
+ context.addTaskInterruptListener(new TaskInterruptListener {
+ override def onTaskInterrupted(context: TaskContext, reason: String): Unit = {
+ interruptListenerStarted.release()
+ interruptListenerRelease.acquire()
+ }
+ })
+
+ // Interrupt the task from a separate thread and wait until the interrupt listener starts.
+ val interruptThread = new Thread(() => context.markInterrupted("test interrupt"))
+ interruptThread.start()
+ interruptListenerStarted.acquire()
+
+ // While the interrupt listener is running on the interrupt thread, mark the task completed
+ // on this thread. With the dedicated interrupt listener loop, this must NOT be blocked.
+ context.markTaskCompleted(None)
+
+ // The completion listener should have been called even though the interrupt listener is still
+ // running. If `invokeListeners()` were shared between interrupt and completion listeners,
+ // the completion listener would be silently skipped because `listenerInvocationThread` would
+ // be held by the interrupt thread.
+ verify(completionListener, times(1)).onTaskCompletion(any())
+
+ // Release the interrupt listener and join the interrupt thread.
+ interruptListenerRelease.release()
+ interruptThread.join()
+ }
+
test("FailureListener throws after task body fails") {
val context = TaskContext.empty()
val listenerCalls = ArrayBuffer.empty[String]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org