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/08 23:13:55 UTC
(spark) branch master updated: [SPARK-56330][CORE] Add TaskInterruptListener to TaskContext for interrupt notifications
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 a7fb5bbe6f86 [SPARK-56330][CORE] Add TaskInterruptListener to TaskContext for interrupt notifications
a7fb5bbe6f86 is described below
commit a7fb5bbe6f86fc761002a8fd7b0d59b55568a4e9
Author: Ivan Sadikov <iv...@databricks.com>
AuthorDate: Thu Apr 9 08:13:42 2026 +0900
[SPARK-56330][CORE] Add TaskInterruptListener to TaskContext for interrupt notifications
### What changes were proposed in this pull request?
Added `TaskInterruptListener`, a new listener type that fires immediately when a task is interrupted via `markInterrupted`. The new listener follows the same invocation model as `TaskCompletionListener` and `TaskFailureListener` - all listeners run sequentially, exceptions in individual listeners don't suppress other listeners, and listener exceptions cause task failure.
Also fixes `TaskKilledException` to pass the cancellation reason to the `RuntimeException` constructor, so `getMessage()` no longer returns `null`.
### Why are the changes needed?
The polling-based cancellation API (`killTaskIfInterrupted`) cannot support push-style reactions to task interruption. Components that block on I/O or hold locks have no way to react immediately when cancellation is signalled without implementing their own polling thread. `TaskInterruptListener` closes this gap by providing a callback that fires synchronously when `markInterrupted` is called.
### Does this PR introduce _any_ user-facing change?
Yes. Adds a new `DeveloperApi` listener interface `TaskInterruptListener` and a corresponding `addTaskInterruptListener` method on `TaskContext`. Also fixes `TaskKilledException#getMessage()` to return the actual reason string rather than `null`.
### How was this patch tested?
Added a `TaskContextSuite` test that verifies all registered `TaskInterruptListener`s are invoked even when some throw exceptions, and that the interrupt reason is propagated into the resulting `TaskCompletionListenerException` message. Added Java interop coverage in `JavaTaskContextCompileCheck`.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #55151 from sadikovi/SPARK-56330.
Authored-by: Ivan Sadikov <iv...@databricks.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
.../org/apache/spark/BarrierTaskContext.scala | 5 ++
.../main/scala/org/apache/spark/TaskContext.scala | 28 ++++++-
.../scala/org/apache/spark/TaskContextImpl.scala | 86 ++++++++++++++++++----
.../org/apache/spark/TaskKilledException.scala | 2 +-
.../org/apache/spark/util/taskListeners.scala | 10 +++
.../apache/spark/JavaTaskContextCompileCheck.java | 8 ++
.../apache/spark/scheduler/TaskContextSuite.scala | 51 ++++++++++++-
project/MimaExcludes.scala | 4 +-
8 files changed, 176 insertions(+), 18 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/BarrierTaskContext.scala b/core/src/main/scala/org/apache/spark/BarrierTaskContext.scala
index 32c89eb0b092..7f5d1f39d86c 100644
--- a/core/src/main/scala/org/apache/spark/BarrierTaskContext.scala
+++ b/core/src/main/scala/org/apache/spark/BarrierTaskContext.scala
@@ -213,6 +213,11 @@ class BarrierTaskContext private[spark] (
this
}
+ override def addTaskInterruptListener(listener: TaskInterruptListener): this.type = {
+ taskContext.addTaskInterruptListener(listener)
+ this
+ }
+
override def stageId(): Int = taskContext.stageId()
override def stageAttemptNumber(): Int = taskContext.stageAttemptNumber()
diff --git a/core/src/main/scala/org/apache/spark/TaskContext.scala b/core/src/main/scala/org/apache/spark/TaskContext.scala
index ae4a0993a5a1..11f1e1e63f11 100644
--- a/core/src/main/scala/org/apache/spark/TaskContext.scala
+++ b/core/src/main/scala/org/apache/spark/TaskContext.scala
@@ -27,7 +27,7 @@ import org.apache.spark.metrics.source.Source
import org.apache.spark.resource.ResourceInformation
import org.apache.spark.scheduler.Task
import org.apache.spark.shuffle.FetchFailedException
-import org.apache.spark.util.{AccumulatorV2, TaskCompletionListener, TaskFailureListener}
+import org.apache.spark.util.{AccumulatorV2, TaskCompletionListener, TaskFailureListener, TaskInterruptListener}
object TaskContext {
@@ -169,6 +169,32 @@ abstract class TaskContext extends Serializable {
})
}
+ /**
+ * Adds a listener to be executed when the task is interrupted.
+ *
+ * Adding a listener to an already interrupted task will result in that listener being called
+ * immediately.
+ *
+ * There are no ordering guarantees for task interrupt listeners. Listeners are guaranteed to
+ * execute sequentially with other task completion, failure, or interrupt listeners. Listeners may
+ * be invoked concurrently with the task itself.
+ *
+ * Exceptions thrown by the listener will mark the task as failed; they are not propagated from
+ * `markInterrupted` (for example when interruption is delivered on the executor kill thread).
+ */
+ def addTaskInterruptListener(listener: TaskInterruptListener): TaskContext
+
+ /**
+ * Adds a listener to be executed when the task is interrupted (Scala closure form).
+ * Behavior matches `addTaskInterruptListener` with TaskInterruptListener.
+ */
+ def addTaskInterruptListener(f: (TaskContext, String) => Unit): TaskContext = {
+ addTaskInterruptListener(new TaskInterruptListener {
+ override def onTaskInterrupted(context: TaskContext, reason: String): Unit =
+ f(context, reason)
+ })
+ }
+
/** Runs a task with this context, ensuring failure and completion listeners get triggered. */
private[spark] def runTaskWithListeners[T](task: Task[T]): T = {
try {
diff --git a/core/src/main/scala/org/apache/spark/TaskContextImpl.scala b/core/src/main/scala/org/apache/spark/TaskContextImpl.scala
index 3d0adb1f901a..b2dffd795f93 100644
--- a/core/src/main/scala/org/apache/spark/TaskContextImpl.scala
+++ b/core/src/main/scala/org/apache/spark/TaskContextImpl.scala
@@ -73,8 +73,11 @@ private[spark] class TaskContextImpl(
/** List of callback functions to execute when the task fails. */
@transient private val onFailureCallbacks = new Stack[TaskFailureListener]
+ /** List of callback functions to execute when the task is interrupted. */
+ @transient private val onInterruptCallbacks = new Stack[TaskInterruptListener]
+
/**
- * The thread currently executing task completion or failure listeners, if any.
+ * The thread currently executing task completion, failure, or interrupt listeners, if any.
*
* `invokeListeners()` uses this to ensure listeners are called sequentially.
*/
@@ -126,6 +129,22 @@ private[spark] class TaskContextImpl(
this
}
+ 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 there is no such thread, the call to `invokeTaskInterruptListeners()` below will execute
+ // all listeners, including the new listener.
+ onInterruptCallbacks.push(listener)
+ reasonIfKilled
+ }.foreach { reason =>
+ invokeTaskInterruptListeners(reason, new TaskKilledException(reason))
+ }
+ this
+ }
+
override def resourcesJMap(): java.util.Map[String, ResourceInformation] = {
resources.asJava
}
@@ -153,7 +172,11 @@ 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) {
+ invokeListeners(
+ onCompleteCallbacks,
+ "TaskCompletionListener",
+ error,
+ markTaskFailedOnListenerError = true) {
_.onTaskCompletion(this)
}
}
@@ -161,15 +184,35 @@ 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)) {
+ invokeListeners(
+ onFailureCallbacks,
+ "TaskFailureListener",
+ Option(error),
+ markTaskFailedOnListenerError = true) {
_.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)
+ }
+ }
+
private def invokeListeners[T](
listeners: Stack[T],
name: String,
- error: Option[Throwable])(
+ error: Option[Throwable],
+ markTaskFailedOnListenerError: Boolean)(
callback: T => Unit): Unit = {
// This method is subject to two constraints:
//
@@ -212,7 +255,8 @@ private[spark] class TaskContextImpl(
callback(listener)
} catch {
case e: Throwable =>
- // A listener failed. Temporarily clear the listenerInvocationThread and markTaskFailed.
+ // A listener failed. For completion/failure listeners, temporarily clear
+ // listenerInvocationThread and markTaskFailed so nested TaskContext calls can run.
//
// One of the following cases applies (#3 being the interesting one):
//
@@ -246,15 +290,20 @@ 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.
- try {
- listenerInvocationThread = None
- markTaskFailed(e)
- } catch {
- case t: Throwable => e.addSuppressed(t)
- } finally {
- synchronized {
- if (listenerInvocationThread.isEmpty) {
- listenerInvocationThread = Some(Thread.currentThread())
+ //
+ // 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())
+ }
}
}
}
@@ -272,6 +321,15 @@ private[spark] class TaskContextImpl(
private[spark] override def markInterrupted(reason: String): Unit = {
reasonIfKilled = Some(reason)
+ try {
+ invokeTaskInterruptListeners(reason, new TaskKilledException(reason))
+ } catch {
+ case e: TaskCompletionListenerException =>
+ // SPARK-56330: Do not propagate to the executor kill / RPC thread; record the aggregate
+ // failure for the task (not the first listener exception only).
+ logError(log"Exception(s) in TaskInterruptListener(s) after task was interrupted", e)
+ markTaskFailed(e)
+ }
}
private[spark] override def killTaskIfInterrupted(): Unit = {
diff --git a/core/src/main/scala/org/apache/spark/TaskKilledException.scala b/core/src/main/scala/org/apache/spark/TaskKilledException.scala
index 9dbf0d493be1..ccd105fefa5b 100644
--- a/core/src/main/scala/org/apache/spark/TaskKilledException.scala
+++ b/core/src/main/scala/org/apache/spark/TaskKilledException.scala
@@ -24,6 +24,6 @@ import org.apache.spark.annotation.DeveloperApi
* Exception thrown when a task is explicitly killed (i.e., task failure is expected).
*/
@DeveloperApi
-class TaskKilledException(val reason: String) extends RuntimeException {
+class TaskKilledException(val reason: String) extends RuntimeException(reason) {
def this() = this("unknown reason")
}
diff --git a/core/src/main/scala/org/apache/spark/util/taskListeners.scala b/core/src/main/scala/org/apache/spark/util/taskListeners.scala
index 51feccfb8342..39b7dbe66f23 100644
--- a/core/src/main/scala/org/apache/spark/util/taskListeners.scala
+++ b/core/src/main/scala/org/apache/spark/util/taskListeners.scala
@@ -44,6 +44,16 @@ trait TaskFailureListener extends EventListener {
def onTaskFailure(context: TaskContext, error: Throwable): Unit
}
+/**
+ * :: DeveloperApi ::
+ *
+ * Listener providing a callback function to invoke when a task's execution is interrupted.
+ */
+@DeveloperApi
+trait TaskInterruptListener extends EventListener {
+ def onTaskInterrupted(context: TaskContext, reason: String): Unit
+}
+
/**
* Exception thrown when there is an exception in executing the callback in TaskCompletionListener.
diff --git a/core/src/test/java/test/org/apache/spark/JavaTaskContextCompileCheck.java b/core/src/test/java/test/org/apache/spark/JavaTaskContextCompileCheck.java
index 5ce7937c03de..a72e2664c3fb 100644
--- a/core/src/test/java/test/org/apache/spark/JavaTaskContextCompileCheck.java
+++ b/core/src/test/java/test/org/apache/spark/JavaTaskContextCompileCheck.java
@@ -23,6 +23,7 @@ import org.apache.spark.TaskContext;
import org.apache.spark.resource.ResourceInformation;
import org.apache.spark.util.TaskCompletionListener;
import org.apache.spark.util.TaskFailureListener;
+import org.apache.spark.util.TaskInterruptListener;
/**
* Something to make sure that TaskContext can be used in Java.
@@ -37,6 +38,7 @@ public class JavaTaskContextCompileCheck {
tc.addTaskCompletionListener(new JavaTaskCompletionListenerImpl());
tc.addTaskFailureListener(new JavaTaskFailureListenerImpl());
+ tc.addTaskInterruptListener(new JavaTaskInterruptListenerImpl());
tc.attemptNumber();
tc.partitionId();
@@ -77,4 +79,10 @@ public class JavaTaskContextCompileCheck {
}
}
+ static class JavaTaskInterruptListenerImpl implements TaskInterruptListener {
+ @Override
+ public void onTaskInterrupted(TaskContext context, String reason) {
+ }
+ }
+
}
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 d08e75733abf..a232048f445d 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala
@@ -23,7 +23,7 @@ import java.util.concurrent.atomic.AtomicInteger
import scala.collection.mutable.ArrayBuffer
-import org.mockito.ArgumentMatchers.any
+import org.mockito.ArgumentMatchers.{any, eq => mockitoEq}
import org.mockito.Mockito._
import org.scalatest.BeforeAndAfter
@@ -150,6 +150,55 @@ class TaskContextSuite extends SparkFunSuite with BeforeAndAfter with LocalSpark
assert(e.getMessage.contains("exception in task"))
}
+ test("all TaskInterruptListeners should be called even if some fail") {
+ val context = TaskContext.empty()
+ val listener = mock(classOf[TaskInterruptListener])
+ context.addTaskInterruptListener(new TaskInterruptListener {
+ override def onTaskInterrupted(context: TaskContext, reason: String): Unit =
+ throw new Exception("exception in listener1")
+ })
+ context.addTaskInterruptListener(listener)
+ context.addTaskInterruptListener(new TaskInterruptListener {
+ override def onTaskInterrupted(context: TaskContext, reason: String): Unit =
+ throw new Exception("exception in listener3")
+ })
+
+ // Interrupt listener failures mark the task failed but are not thrown from `markInterrupted`
+ // (kill / RPC thread must not see TaskCompletionListenerException).
+ context.markInterrupted("test interrupt")
+ assert(context.isFailed())
+
+ verify(listener, times(1)).onTaskInterrupted(any(), mockitoEq("test interrupt"))
+
+ val taskFailure = context.getTaskFailure.get
+ assert(taskFailure.isInstanceOf[TaskCompletionListenerException])
+ val e = taskFailure.asInstanceOf[TaskCompletionListenerException]
+ assert(e.getMessage.contains("exception in listener1"))
+ assert(e.getMessage.contains("exception in listener3"))
+ assert(e.getMessage.contains("test interrupt"))
+ }
+
+ test("immediately call an interrupt listener if the context is already interrupted") {
+ var invocations = 0
+ var lastReason: String = null
+ val context = TaskContext.empty()
+ context.markInterrupted("already interrupted")
+ context.addTaskInterruptListener(new TaskInterruptListener {
+ override def onTaskInterrupted(context: TaskContext, reason: String): Unit = {
+ invocations += 1
+ lastReason = reason
+ }
+ })
+ assert(invocations == 1)
+ assert(lastReason == "already interrupted")
+ context.addTaskInterruptListener(new TaskInterruptListener {
+ override def onTaskInterrupted(context: TaskContext, reason: String): Unit = {
+ invocations += 1
+ }
+ })
+ assert(invocations == 2)
+ }
+
test("FailureListener throws after task body fails") {
val context = TaskContext.empty()
val listenerCalls = ArrayBuffer.empty[String]
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 599dab36e5db..b5434efee090 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -53,7 +53,9 @@ object MimaExcludes {
// [SPARK-55793][CORE] Add multiple log directories support to SHS
ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.ApplicationAttemptInfo.apply"),
ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.ApplicationAttemptInfo.copy"),
- ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.status.api.v1.ApplicationAttemptInfo$")
+ ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.status.api.v1.ApplicationAttemptInfo$"),
+ // [SPARK-56330][CORE] Add TaskInterruptListener to TaskContext for interrupt notifications
+ ProblemFilters.exclude[ReversedMissingMethodProblem]("org.apache.spark.TaskContext.addTaskInterruptListener")
)
// Exclude rules for 4.1.x from 4.0.0
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org