LuciferYang commented on code in PR #49413:
URL: https://github.com/apache/spark/pull/49413#discussion_r1907368576
##########
core/src/main/scala/org/apache/spark/scheduler/Task.scala:
##########
@@ -231,10 +232,23 @@ private[spark] abstract class Task[T](
require(reason != null)
_reasonIfKilled = reason
if (context != null) {
- context.markInterrupted(reason)
- }
- if (interruptThread && taskThread != null) {
- taskThread.interrupt()
+ TaskContext.synchronized {
+ if (context.interruptible()) {
+ context.markInterrupted(reason)
+ if (interruptThread && taskThread != null) {
+ taskThread.interrupt()
+ }
+ } else {
+ val threadToInterrupt = if (interruptThread && taskThread != null) {
Review Comment:
This can be simplified to `if (interruptThread) Option(taskThread) else
None`.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -296,4 +304,42 @@ private[spark] class TaskContextImpl(
private[spark] override def fetchFailed: Option[FetchFailedException] =
_fetchFailedException
private[spark] override def getLocalProperties: Properties = localProperties
+
+
+ override def interruptible(): Boolean =
TaskContext.synchronized(_interruptible)
+
+ override def pendingInterrupt(threadToInterrupt: Option[Thread], reason:
String): Unit = {
+ TaskContext.synchronized {
+ pendingInterruptRequest = Some((threadToInterrupt, reason))
+ }
+ }
+
+ def createResourceUninterruptibly[T <: Closeable](resourceBuilder: => T): T
= {
+
+ @inline def interruptIfRequired(): Unit = {
+ pendingInterruptRequest.foreach { case (threadToInterrupt, reason) =>
+ markInterrupted(reason)
+ threadToInterrupt.foreach(_.interrupt())
+ }
+ killTaskIfInterrupted()
+ }
+
+ TaskContext.synchronized {
+ interruptIfRequired()
+
+ if (_interruptible) {
Review Comment:
Can we skip the condition check and directly assign the value here?
##########
core/src/test/scala/org/apache/spark/JobCancellationSuite.scala:
##########
@@ -712,6 +713,128 @@ class JobCancellationSuite extends SparkFunSuite with
Matchers with BeforeAndAft
assert(executionOfInterruptibleCounter.get() < numElements)
}
+ Seq(true, false).foreach { interruptible =>
+
+ val (hint1, hint2) = if (interruptible) {
+ (" not", "")
+ } else {
+ ("", " not")
+ }
+
+ val testName = s"SPARK-50768:$hint1 use
TaskContext.createResourceUninterruptibly " +
+ s"would$hint2 cause stream leak on task interruption"
+
+ test(testName) {
+ import org.apache.spark.JobCancellationSuite._
+ withTempDir { dir =>
+
+ // `InterruptionSensitiveInputStream` is designed to easily leak the
underlying stream
+ // when task thread interruption happens during its initialization.
+ class InterruptionSensitiveInputStream(fileHint: String) extends
InputStream {
+ private var underlying: InputStream = _
+
+ def initialize(): InputStream = {
+ val in: InputStream = new InputStream {
+
+ open()
+
+ private def dumpFile(typeName: String): Unit = {
+ val file = new File(dir, s"$typeName.$fileHint")
Review Comment:
The `dumpFile` method is best implemented using the `try-with-resources` to
ensure the output stream is closed.
--
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]
