Copilot commented on code in PR #12487:
URL: https://github.com/apache/gluten/pull/12487#discussion_r3556377527
##########
backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideCache.scala:
##########
@@ -249,4 +262,43 @@ object VeloxBroadcastBuildSideCache
HashJoinBuilder.clearHashTable(key, value.pointer)
}
}
+
+ private def registerLocalCleanupIfNeeded(
+ broadcastHashTableId: String,
+ enableLocalHashTableCleanup: Boolean): Unit = {
+ if (!enableLocalHashTableCleanup || !TaskResources.inSparkTask()) {
+ return
+ }
+
+ val resourceId =
s"broadcast-hash-table-local-cleanup:$broadcastHashTableId"
+ TaskResources.addResourceIfNotRegistered(
+ resourceId,
+ () => {
+ localReferenceCounts
+ .computeIfAbsent(broadcastHashTableId, _ => new AtomicInteger(0))
+ .incrementAndGet()
+
+ new TaskResource {
+ override def release(): Unit =
releaseLocalReference(broadcastHashTableId)
+
+ override def priority(): Int = 200
+
+ override def resourceName(): String = resourceId
+ }
+ }
+ )
+ }
Review Comment:
`localReferenceCounts` is updated concurrently from multiple tasks, but the
increment path isn’t synchronized with `releaseLocalReference`. Because
`releaseLocalReference` uses `remove(key, refCount)` on the same
`AtomicInteger` instance, a concurrent increment can still be lost and the
entry can be invalidated even though another task has already increased the
count. Synchronize the refcount increment with the decrement/invalidate path.
##########
backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideCache.scala:
##########
@@ -249,4 +262,43 @@ object VeloxBroadcastBuildSideCache
HashJoinBuilder.clearHashTable(key, value.pointer)
}
}
+
+ private def registerLocalCleanupIfNeeded(
+ broadcastHashTableId: String,
+ enableLocalHashTableCleanup: Boolean): Unit = {
+ if (!enableLocalHashTableCleanup || !TaskResources.inSparkTask()) {
+ return
+ }
+
+ val resourceId =
s"broadcast-hash-table-local-cleanup:$broadcastHashTableId"
+ TaskResources.addResourceIfNotRegistered(
+ resourceId,
+ () => {
+ localReferenceCounts
+ .computeIfAbsent(broadcastHashTableId, _ => new AtomicInteger(0))
+ .incrementAndGet()
+
+ new TaskResource {
+ override def release(): Unit =
releaseLocalReference(broadcastHashTableId)
+
+ override def priority(): Int = 200
+
+ override def resourceName(): String = resourceId
+ }
+ }
+ )
+ }
+
+ private def releaseLocalReference(broadcastHashTableId: String): Unit = {
+ val refCount = localReferenceCounts.get(broadcastHashTableId)
+ if (refCount == null) {
+ return
+ }
+
+ val remaining = refCount.decrementAndGet()
+ if (remaining <= 0 && localReferenceCounts.remove(broadcastHashTableId,
refCount)) {
+ logInfo(s"Cleaning up broadcast hash table with local fallback:
$broadcastHashTableId")
+ invalidateBroadcastHashtable(broadcastHashTableId)
+ }
+ }
Review Comment:
`releaseLocalReference` performs decrement / conditional map removal / cache
invalidation without synchronization. With concurrent tasks, this can
invalidate the broadcast hash table even if another task is in the middle of
registering/incrementing the same `AtomicInteger` (since `remove(key,
refCount)` only checks object identity). Wrap the whole decrement + invalidate
decision in the same lock used by the increment path so the refcount logic is
consistent.
##########
backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideCache.scala:
##########
@@ -249,4 +262,43 @@ object VeloxBroadcastBuildSideCache
HashJoinBuilder.clearHashTable(key, value.pointer)
}
}
+
+ private def registerLocalCleanupIfNeeded(
+ broadcastHashTableId: String,
+ enableLocalHashTableCleanup: Boolean): Unit = {
+ if (!enableLocalHashTableCleanup || !TaskResources.inSparkTask()) {
Review Comment:
The new task-local cleanup path (via `enableLocalHashTableCleanup`,
TaskResources registration, and refcounted invalidation) changes lifecycle
behavior for broadcast hash tables but isn’t covered by existing tests in the
Velox backend. Adding a regression test that exercises the `executionId ==
null` path and asserts the cache is cleared at task completion would help
prevent reintroducing the memory-leak / outstanding-pool failure this PR
targets.
--
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]
