[SPARK-50768][CORE] Introduce TaskContext.createResourceUninterruptibly to avoid stream leak by task interruption

core/src/main/scala/org/apache/spark/BarrierTaskContext.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark
 
+import java.io.Closeable
 import java.util.{Properties, TimerTask}
 import java.util.concurrent.{ScheduledThreadPoolExecutor, TimeUnit}
 
@@ -273,6 +274,18 @@ class BarrierTaskContext private[spark] (
   }
 
   override private[spark] def getLocalProperties: Properties = taskContext.getLocalProperties
+
+  override private[spark] def interruptible(): Boolean = taskContext.interruptible()
+
+  override private[spark] def pendingInterrupt(threadToInterrupt: Option[Thread], reason: String)
+  : Unit = {
+    taskContext.pendingInterrupt(threadToInterrupt, reason)
+  }
+
+  override private[spark] def createResourceUninterruptibly[T <: Closeable](resourceBuilder: => T)
+  : T = {
+    taskContext.createResourceUninterruptibly(resourceBuilder)
+  }
 }
 
 @Experimental

core/src/main/scala/org/apache/spark/TaskContext.scala

@@ -17,7 +17,7 @@
 
 package org.apache.spark
 
-import java.io.Serializable
+import java.io.Closeable
 import java.util.Properties
 
 import org.apache.spark.annotation.{DeveloperApi, Evolving, Since}
@@ -305,4 +305,20 @@ abstract class TaskContext extends Serializable {
 
   /** Gets local properties set upstream in the driver. */
   private[spark] def getLocalProperties: Properties
+
+
+  /** Whether the current task is allowed to interrupt. */
+  private[spark] def interruptible(): Boolean
+
+  /**
+   *  Pending the interruption request until the task is able to
+   *  interrupt after creating the resource uninterruptibly.
+   */
+  private[spark] def pendingInterrupt(threadToInterrupt: Option[Thread], reason: String): Unit
+
+  /**
+   * Creating a closeable resource uninterruptibly. A task is not allowed to interrupt in this
+   * state until the resource creation finishes.
+   */
+  private[spark] def createResourceUninterruptibly[T <: Closeable](resourceBuilder: => T): T
 }

core/src/main/scala/org/apache/spark/TaskContextImpl.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark
 
+import java.io.Closeable
 import java.util.{Properties, Stack}
 import javax.annotation.concurrent.GuardedBy
 
@@ -82,6 +83,13 @@ private[spark] class TaskContextImpl(
   // If defined, the corresponding task has been killed and this option contains the reason.
   @volatile private var reasonIfKilled: Option[String] = None
 
+  // The pending interruption request, which is blocked by uninterruptible resource creation.
+  // Should be protected by `TaskContext.synchronized`.
+  private var pendingInterruptRequest: Option[(Option[Thread], String)] = None
+
+  // Whether this task is able to be interrupted. Should be protected by `TaskContext.synchronized`.
+  private var _interruptible = true
+
   // Whether the task has completed.
   private var completed: Boolean = false
 
@@ -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) {
+        _interruptible = false
+      }
+    }
+    try {
+      val resource = resourceBuilder
+      addTaskCompletionListener[Unit](_ => resource.close())
+      resource
+    } finally {
+      TaskContext.synchronized {
+        _interruptible = true
+        interruptIfRequired()
+      }
+    }
+  }
 }

core/src/main/scala/org/apache/spark/scheduler/Task.scala

@@ -22,6 +22,7 @@ import java.util.Properties
 
 import org.apache.spark._
 import org.apache.spark.executor.TaskMetrics
+import org.apache.spark.internal.{Logging, LogKeys, MDC}
 import org.apache.spark.internal.config.APP_CALLER_CONTEXT
 import org.apache.spark.internal.plugin.PluginContainer
 import org.apache.spark.memory.{MemoryMode, TaskMemoryManager}
@@ -70,7 +71,7 @@ private[spark] abstract class Task[T](
     val jobId: Option[Int] = None,
     val appId: Option[String] = None,
     val appAttemptId: Option[String] = None,
-    val isBarrier: Boolean = false) extends Serializable {
+    val isBarrier: Boolean = false) extends Serializable with Logging {
 
   @transient lazy val metrics: TaskMetrics =
     SparkEnv.get.closureSerializer.newInstance().deserialize(ByteBuffer.wrap(serializedTaskMetrics))
@@ -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) {
+            Some(taskThread)
+          } else {
+            None
+          }
+          logInfo(log"Task ${MDC(LogKeys.TASK_ID, context.taskAttemptId())} " +
+            log"is currently not interruptible. ")
+          context.pendingInterrupt(threadToInterrupt, reason)
+        }
+      }
     }
   }
 }

core/src/test/scala/org/apache/spark/JobCancellationSuite.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark
 
+import java.io.{File, FileOutputStream, InputStream, ObjectOutputStream}
 import java.util.concurrent.{Semaphore, TimeUnit}
 import java.util.concurrent.atomic.AtomicInteger
 
@@ -35,7 +36,7 @@ import org.apache.spark.executor.ExecutorExitCode
 import org.apache.spark.internal.config._
 import org.apache.spark.internal.config.Deploy._
 import org.apache.spark.scheduler.{JobFailed, SparkListener, SparkListenerExecutorRemoved, SparkListenerJobEnd, SparkListenerJobStart, SparkListenerStageCompleted, SparkListenerTaskEnd, SparkListenerTaskStart}
-import org.apache.spark.util.ThreadUtils
+import org.apache.spark.util.{ThreadUtils, Utils}
 
 /**
  * Test suite for cancelling running jobs. We run the cancellation tasks for single job action
@@ -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")
+                val out = new FileOutputStream(file)
+                val objOut = new ObjectOutputStream(out)
+                objOut.writeBoolean(true)
+                objOut.close()
+              }
+
+              private def open(): Unit = {
+                dumpFile("open")
+              }
+
+              override def close(): Unit = {
+                dumpFile("close")
+              }
+
+              override def read(): Int = -1
+            }
+
+            // Leave some time for the task to be interrupted during the
+            // creation of `InterruptionSensitiveInputStream`.
+            Thread.sleep(5000)
+
+            underlying = in
+            underlying
+          }
+
+          override def read(): Int = -1
+
+          override def close(): Unit = {
+            if (underlying != null) {
+              underlying.close()
+            }
+          }
+        }
+
+        def createStream(fileHint: String): Unit = {
+          if (interruptible) {
+              Utils.tryInitializeResource {
+                new InterruptionSensitiveInputStream(fileHint)
+              } {
+                _.initialize()
+              }
+          } else {
+            TaskContext.get().createResourceUninterruptibly[java.io.InputStream] {
+              Utils.tryInitializeResource {
+                new InterruptionSensitiveInputStream(fileHint)
+              } {
+                _.initialize()
+              }
+            }
+          }
+        }
+
+        sc = new SparkContext("local[2]", "test interrupt streams")
+
+        sc.addSparkListener(new SparkListener {
+          override def onTaskStart(taskStart: SparkListenerTaskStart): Unit = {
+            // Sleep some time to ensure task has started
+            Thread.sleep(1000)
+            taskStartedSemaphore.release()
+          }
+
+          override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = {
+            if (taskEnd.reason.isInstanceOf[TaskKilled]) {
+              taskCancelledSemaphore.release()
+            }
+          }
+        })
+
+        sc.setLocalProperty(SparkContext.SPARK_JOB_INTERRUPT_ON_CANCEL, "true")
+
+        val fileHint = if (interruptible) "interruptible" else "uninterruptible"
+        val future = sc.parallelize(1 to 100, 1).mapPartitions { _ =>
+          createStream(fileHint)
+          Iterator.single(1)
+        }.collectAsync()
+
+        taskStartedSemaphore.acquire()
+        future.cancel()
+        taskCancelledSemaphore.acquire()
+
+        val fileOpen = new File(dir, s"open.$fileHint")
+        val fileClose = new File(dir, s"close.$fileHint")
+        assert(fileOpen.exists())
+
+        if (interruptible) {
+          // The underlying stream leaks when the stream creation is interruptible.
+          assert(!fileClose.exists())
+        } else {
+          // The underlying stream won't leak when the stream creation is uninterruptible.
+          assert(fileClose.exists())
+        }
+      }
+    }
+  }
+
   def testCount(): Unit = {
     // Cancel before launching any tasks
     {