This is an automated email from the ASF dual-hosted git repository.
jackylee-ch pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 6fbdb0eb08 [CORE] Run task resource release outside the global
registry lock (#13166)
6fbdb0eb08 is described below
commit 6fbdb0eb0891604ce4f8e6599584cd9009b71f1b
Author: YangJie <[email protected]>
AuthorDate: Wed Oct 21 03:08:57 2026 -0400
[CORE] Run task resource release outside the global registry lock (#13166)
---
.../org/apache/spark/task/TaskResources.scala | 60 ++++++---
.../org/apache/gluten/task/TaskResourceSuite.scala | 140 ++++++++++++++++++++-
2 files changed, 182 insertions(+), 18 deletions(-)
diff --git
a/gluten-core/src/main/scala/org/apache/spark/task/TaskResources.scala
b/gluten-core/src/main/scala/org/apache/spark/task/TaskResources.scala
index e0daeb985c..10a6afc6dd 100644
--- a/gluten-core/src/main/scala/org/apache/spark/task/TaskResources.scala
+++ b/gluten-core/src/main/scala/org/apache/spark/task/TaskResources.scala
@@ -217,28 +217,40 @@ object TaskResources extends TaskListener with Logging {
})
tc.addTaskCompletionListener(new TaskCompletionListener {
override def onTaskCompletion(context: TaskContext): Unit = {
- RESOURCE_REGISTRIES.synchronized {
- val currentTaskRegistries = RESOURCE_REGISTRIES.get(context)
- if (currentTaskRegistries == null) {
+ val currentTaskRegistries = RESOURCE_REGISTRIES.synchronized {
+ val r = RESOURCE_REGISTRIES.get(context)
+ if (r == null) {
throw new IllegalStateException(
"TaskResourceRegistry is not initialized, this should not
happen")
}
- // We should first call `releaseAll` then remove the registries,
because
- // the functions inside registries may register new resource to
registries.
+ r
+ }
+ // Release outside the JVM-global RESOURCE_REGISTRIES lock: the
release callbacks
+ // do native teardown (ms-scale), and holding the lock across them
serializes every
+ // other task's completion and registration. The registry is
per-task with its own
+ // internal lock, so the map is the only state that needs the global
lock.
+ // Take the usage recorder first: releaseAll marks the registry
released, after
+ // which the registry rejects every call, getSharedUsage included.
+ val sharedUsage = currentTaskRegistries.getSharedUsage()
+ try {
+ currentTaskRegistries.releaseAll()
+ } finally {
+ // Remove the map entry even if releaseAll or the metrics update
throws, or the
+ // registry leaks across tasks. Metrics is best-effort so it
cannot mask a
+ // releaseAll failure propagating from the outer try.
try {
- currentTaskRegistries.releaseAll()
+ context
+ .taskMetrics()
+ .incPeakExecutionMemory(sharedUsage.peak())
+ } catch {
+ case NonFatal(e) =>
+ logWarning("Failed to record peak execution memory", e)
} finally {
- // Removing the registry must happen even if the metrics update
throws,
- // otherwise the registry stays reachable and leaks across
tasks. The metrics
- // update itself is best-effort: catch it here so a metrics
failure cannot
- // replace (mask) a releaseAll failure propagating from the
outer try.
- try {
-
context.taskMetrics().incPeakExecutionMemory(registry.getSharedUsage().peak())
- } catch {
- case NonFatal(e) =>
- logWarning("Failed to record peak execution memory", e)
- } finally {
- RESOURCE_REGISTRIES.remove(context)
+ RESOURCE_REGISTRIES.synchronized {
+ // Remove only if the map still points at the registry we
released.
+ if (RESOURCE_REGISTRIES.get(context) eq currentTaskRegistries)
{
+ RESOURCE_REGISTRIES.remove(context)
+ }
}
}
}
@@ -268,8 +280,15 @@ class TaskResourceRegistry extends Logging {
mutable.Map.empty[Int, mutable.LinkedHashSet[TaskResource]]
private var exclusiveLockAcquired: Boolean = false
+ // Set once releaseAll has run. The completion listener releases outside the
global
+ // registry lock, so a thread that fetched this registry before the map
entry is
+ // removed must fail here rather than register a resource nothing will
release.
+ private var released: Boolean = false
private def lock[T](body: => T): T = {
synchronized {
+ if (released) {
+ throw new IllegalStateException("TaskResourceRegistry is already
released")
+ }
if (exclusiveLockAcquired) {
throw new ConcurrentModificationException
}
@@ -304,6 +323,13 @@ class TaskResourceRegistry extends Logging {
/** Release all managed resources according to priority and reversed order */
private[task] def releaseAll(): Unit = lock {
+ // Mark released on every exit: a fatal throwable from a release callback
skips the rest
+ // of the pass, but the completion listener still removes the map entry.
+ try releaseAll0()
+ finally released = true
+ }
+
+ private def releaseAll0(): Unit = {
val failures = mutable.ArrayBuffer.empty[Throwable]
def safeResourceName(resource: TaskResource): String =
try resource.resourceName()
diff --git
a/gluten-core/src/test/scala/org/apache/gluten/task/TaskResourceSuite.scala
b/gluten-core/src/test/scala/org/apache/gluten/task/TaskResourceSuite.scala
index 03ba024b4f..7890c15e0b 100644
--- a/gluten-core/src/test/scala/org/apache/gluten/task/TaskResourceSuite.scala
+++ b/gluten-core/src/test/scala/org/apache/gluten/task/TaskResourceSuite.scala
@@ -16,15 +16,20 @@
*/
package org.apache.gluten.task
+import org.apache.spark.TaskContext
import org.apache.spark.memory.{MemoryConsumer, MemoryMode}
import org.apache.spark.sql.catalyst.plans.SQLHelper
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.task.{TaskResource, TaskResources}
+import org.apache.spark.task.{TaskResource, TaskResourceRegistry,
TaskResources}
import org.apache.spark.util.SparkTaskUtil
import org.scalatest.funsuite.AnyFunSuite
import java.util.UUID
+import java.util.concurrent.TimeUnit
+
+import scala.util.Try
+import scala.util.control.NonFatal
class TaskResourceSuite extends AnyFunSuite with SQLHelper {
test("Run unsafe") {
@@ -119,4 +124,137 @@ class TaskResourceSuite extends AnyFunSuite with
SQLHelper {
// The registry was cleared on the way out, so a fresh task runs cleanly.
TaskResources.runUnsafe(assert(TaskResources.inSparkTask()))
}
+
+ test("Run unsafe - release callbacks run outside the global registry lock") {
+ // On a real executor the release callbacks include JNI native teardown
that
+ // can take milliseconds. Running them under the JVM-global registry lock
+ // serialized every concurrent task's completion and registration on the
+ // slowest callback; a concurrent task start must not block on it.
+ var concurrentTaskStarted = false
+ // Captured on the starter thread; read on this thread only after the join
+ // below confirms the thread terminated (the !isAlive assert gates it),
which
+ // establishes the happens-before to read these plain locals safely.
+ var starterFailure: Option[Throwable] = None
+ TaskResources.runUnsafe {
+ TaskResources.addResource(
+ UUID.randomUUID().toString,
+ new TaskResource {
+ override def release(): Unit = {
+ val starter = new Thread(
+ () => {
+ try {
+ TaskResources.runUnsafe {
+ TaskResources.addResource(
+ UUID.randomUUID().toString,
+ new TaskResource {
+ override def release(): Unit = {}
+ override def resourceName(): String = "concurrent task
resource"
+ }
+ )
+ }
+ concurrentTaskStarted = true
+ } catch {
+ case NonFatal(t) => starterFailure = Some(t)
+ }
+ })
+ starter.start()
+ starter.join(30000)
+ assert(
+ !starter.isAlive,
+ "concurrent task start is still blocked after the release
callback returned")
+ }
+ override def resourceName(): String = "slow releaser"
+ }
+ )
+ }
+ // Surface a genuine unrelated failure with its real stack instead of the
+ // misleading "blocked" message below.
+ starterFailure.foreach(t => throw t)
+ assert(concurrentTaskStarted, "concurrent task start was blocked by the
release pass")
+ }
+
+ test("Run unsafe - registering into a released registry fails") {
+ val outcome = registerDuringRelease(() => ())
+ assert(outcome.isSuccess, outcome)
+ }
+
+ test("Run unsafe - registering after a fatal release failure fails") {
+ // A fatal throwable skips the rest of the release pass, but the
completion listener
+ // still removes the map entry, so the registry must be rejected on this
path too.
+ // scalastyle:off throwerror
+ val outcome = registerDuringRelease(() => throw new LinkageError("fatal
release failure"))
+ // scalastyle:on throwerror
+ assert(outcome.isFailure)
+ val trace = {
+ val writer = new java.io.StringWriter()
+ outcome.failed.get.printStackTrace(new java.io.PrintWriter(writer))
+ writer.toString
+ }
+ assert(trace.contains("fatal release failure"))
+ }
+
+ /**
+ * Runs a task whose release pass lets a thread sharing the task's
TaskContext fetch the registry
+ * and queue on its lock, then ends the pass with `endRelease`. Asserts that
the late registration
+ * was rejected by the released registry, where nothing would ever release
it, and returns the
+ * task's outcome.
+ */
+ private def registerDuringRelease(endRelease: () => Unit): Try[_] = {
+ var lateFailure: Option[Throwable] = None
+ var child: Thread = null
+ val outcome = Try(TaskResources.runUnsafe {
+ val tc = TaskContext.get()
+ TaskResources.addResource(
+ UUID.randomUUID().toString,
+ new TaskResource {
+ override def release(): Unit = {
+ child = new Thread(
+ () => {
+ SparkTaskUtil.setTaskContext(tc)
+ try {
+ TaskResources.addResource(
+ UUID.randomUUID().toString,
+ new TaskResource {
+ override def release(): Unit = {}
+ override def resourceName(): String = "late resource"
+ })
+ } catch {
+ case NonFatal(t) => lateFailure = Some(t)
+ } finally {
+ SparkTaskUtil.unsetTaskContext()
+ }
+ })
+ child.start()
+ awaitBlockedOnRegistry(child)
+ endRelease()
+ }
+ override def resourceName(): String = "releaser"
+ }
+ )
+ })
+ child.join(30000)
+ assert(!child.isAlive)
+ // A missing map entry also fails with IllegalStateException; only the
registry's own
+ // check says "already released".
+ assert(
+ lateFailure.exists(
+ f => f.isInstanceOf[IllegalStateException] &&
f.getMessage.contains("already released")),
+ s"late registration was not rejected by the released registry:
$lateFailure"
+ )
+ outcome
+ }
+
+ // Waits until `t` is blocked entering the registry lock: it already holds
the registry, so
+ // the map entry's removal no longer matters, and it cannot proceed until
the release ends.
+ private def awaitBlockedOnRegistry(t: Thread): Unit = {
+ val registryClass = classOf[TaskResourceRegistry].getName
+ val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(30)
+ while (
+ t.getState != Thread.State.BLOCKED ||
+ !t.getStackTrace.headOption.exists(_.getClassName == registryClass)
+ ) {
+ assert(System.nanoTime() < deadline, "late registration never reached
the registry lock")
+ Thread.sleep(10)
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
