sunchao commented on code in PR #5027:
URL: https://github.com/apache/datafusion-comet/pull/5027#discussion_r3836787391
##########
spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java:
##########
@@ -254,4 +261,152 @@ private static void evaluateInternal(
}
}
}
+
+ /** Visible to the focused allocator test in this package. */
+ static BufferAllocator taskAllocator(TaskContext taskContext) {
+ return taskState(taskContext).allocator();
+ }
+
+ /** Visible to the focused allocator test in this package. */
+ static int taskStateCount() {
+ return TASKS.size();
+ }
+
+ private static TaskState taskState(TaskContext taskContext) {
+ return TASKS.computeIfAbsent(
+ taskContext,
+ context -> {
+ TaskState state = new TaskState(context,
CometTaskContextShim.taskMemoryManager(context));
+ context.addTaskCompletionListener(
+ (TaskCompletionListener) ignored -> state.taskCompleted());
+ return state;
+ });
+ }
+
+ /** Per-task Arrow listener and non-spillable Spark memory consumer. */
+ private static final class TaskState implements AllocationListener {
+ private final TaskContext taskContext;
+ private final long taskAttemptId;
+ private final TaskMemoryConsumer consumer;
+ private final ConcurrentHashMap<String, CometUDF> instances = new
ConcurrentHashMap<>();
+
+ private BufferAllocator allocator;
+ private boolean completed;
+ private boolean closed;
+
+ private TaskState(TaskContext taskContext, TaskMemoryManager
taskMemoryManager) {
+ this.taskContext = taskContext;
+ this.taskAttemptId = taskContext.taskAttemptId();
+ this.consumer =
+ taskMemoryManager.getTungstenMemoryMode() == MemoryMode.OFF_HEAP
+ ? new TaskMemoryConsumer(taskMemoryManager)
+ : null;
+ }
+
+ private synchronized BufferAllocator allocator() {
+ if (completed) {
+ throw new IllegalStateException(
+ "Cannot allocate JVM UDF memory after task " + taskAttemptId + "
completed");
+ }
+ if (allocator == null) {
+ allocator =
+ ROOT_ALLOCATOR.newChildAllocator(
+ "comet-udf-task-" + taskAttemptId, this, 0L, Long.MAX_VALUE);
+ }
+ return allocator;
+ }
+
+ @Override
+ public synchronized void onPreAllocation(long size) {
+ if (completed) {
+ throw new OutOfMemoryException(
+ "Cannot allocate " + size + " JVM UDF bytes after task
completion");
+ }
+ if (consumer != null) {
+ long acquired = consumer.acquireMemory(size);
Review Comment:
**[P2] Keep buffer releases runnable during blocking memory acquisition**
Could we avoid holding the `TaskState` monitor while calling
`consumer.acquireMemory(size)`? Native prefetch can allocate the next UDF
output concurrently with the Spark thread closing the previous output. Under
memory pressure, Spark's execution pool waits here while retaining this
monitor, so `onRelease` cannot return the previous buffer's Spark charge even
when that release would satisfy the pending allocation.
I reproduced this with two real Spark 4.1.3 tasks and a 64 MiB off-heap
pool: another task holds 62 MiB, this task retains 1 MiB of input and 1 MiB of
previous output, and its producer requests another 1 MiB. The producer waits in
`ExecutionMemoryPool.acquireMemory`, while closing the previous output and
completing the task block in `TaskState.onRelease`. Both proceed only after the
other task releases memory. The unchanged `CometTaskMemoryManager` control
releases the previous buffer and completes before the other task frees
anything. Please preserve accounting synchronization without blocking releases
behind the pool wait.
##########
spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java:
##########
@@ -254,4 +261,152 @@ private static void evaluateInternal(
}
}
}
+
+ /** Visible to the focused allocator test in this package. */
+ static BufferAllocator taskAllocator(TaskContext taskContext) {
+ return taskState(taskContext).allocator();
+ }
+
+ /** Visible to the focused allocator test in this package. */
+ static int taskStateCount() {
+ return TASKS.size();
+ }
+
+ private static TaskState taskState(TaskContext taskContext) {
+ return TASKS.computeIfAbsent(
+ taskContext,
+ context -> {
+ TaskState state = new TaskState(context,
CometTaskContextShim.taskMemoryManager(context));
+ context.addTaskCompletionListener(
+ (TaskCompletionListener) ignored -> state.taskCompleted());
+ return state;
+ });
+ }
+
+ /** Per-task Arrow listener and non-spillable Spark memory consumer. */
+ private static final class TaskState implements AllocationListener {
+ private final TaskContext taskContext;
+ private final long taskAttemptId;
+ private final TaskMemoryConsumer consumer;
+ private final ConcurrentHashMap<String, CometUDF> instances = new
ConcurrentHashMap<>();
+
+ private BufferAllocator allocator;
+ private boolean completed;
+ private boolean closed;
+
+ private TaskState(TaskContext taskContext, TaskMemoryManager
taskMemoryManager) {
+ this.taskContext = taskContext;
+ this.taskAttemptId = taskContext.taskAttemptId();
+ this.consumer =
+ taskMemoryManager.getTungstenMemoryMode() == MemoryMode.OFF_HEAP
+ ? new TaskMemoryConsumer(taskMemoryManager)
+ : null;
+ }
+
+ private synchronized BufferAllocator allocator() {
+ if (completed) {
+ throw new IllegalStateException(
+ "Cannot allocate JVM UDF memory after task " + taskAttemptId + "
completed");
+ }
+ if (allocator == null) {
+ allocator =
+ ROOT_ALLOCATOR.newChildAllocator(
+ "comet-udf-task-" + taskAttemptId, this, 0L, Long.MAX_VALUE);
+ }
+ return allocator;
+ }
+
+ @Override
+ public synchronized void onPreAllocation(long size) {
+ if (completed) {
+ throw new OutOfMemoryException(
+ "Cannot allocate " + size + " JVM UDF bytes after task
completion");
+ }
+ if (consumer != null) {
+ long acquired = consumer.acquireMemory(size);
+ if (acquired < size) {
+ if (acquired > 0) {
+ consumer.freeMemory(acquired);
+ }
+ throw new OutOfMemoryException(
+ "Failed to acquire " + size + " JVM UDF bytes from Spark
TaskMemoryManager");
+ }
+ }
+ }
+
+ @Override
+ public synchronized boolean onFailedAllocation(long size,
AllocationOutcome outcome) {
+ if (!completed && consumer != null) {
+ consumer.freeMemory(size);
+ }
+ return false;
+ }
+
+ @Override
+ public void onRelease(long size) {
+ BufferAllocator toClose = null;
+ synchronized (this) {
+ if (!completed && consumer != null) {
+ consumer.freeMemory(size);
+ } else if (completed && allocator.getAllocatedMemory() == 0 &&
!closed) {
+ closed = true;
+ toClose = allocator;
+ }
+ }
+ if (toClose != null) {
+ close(toClose);
+ }
+ }
+
+ private void taskCompleted() {
+ BufferAllocator toClose = null;
+ boolean removeOnly = false;
+ synchronized (this) {
+ completed = true;
+ instances.clear();
+ if (consumer != null && consumer.getUsed() > 0) {
+ // Spark discards all remaining task accounting immediately after
completion listeners.
+ // Release it here; later FFI callbacks only drive allocator cleanup.
+ consumer.freeMemory(consumer.getUsed());
+ }
+ if ((allocator == null || allocator.getAllocatedMemory() == 0) &&
!closed) {
+ closed = true;
+ toClose = allocator;
Review Comment:
**[P2] Fence in-flight allocations before closing the task allocator**
Could we account for admitted/in-flight allocations before treating this
allocator as empty? Arrow 18.3.0 calls `onPreAllocation` before updating
`getAllocatedMemory()`. Completion can therefore observe zero and set `closed =
true` while an allocation has already passed the completion check. If that
allocation finishes before the subsequent `allocator.close()`, Arrow marks the
allocator closed and throws because it still owns a buffer. The bridge removes
the `TASKS` entry, but the child remains permanently attached to
`ROOT_ALLOCATOR`, and later releases cannot repair cleanup.
This overlap is reachable when `CometCollectLimitExec.executeCollect` stops
through Spark's `executeTake` while a native-scan producer is still in a
synchronous JVM UDF call. `releasePlan` drops the receiver/context without
joining that detached producer, so listener ordering alone does not fence the
allocation.
A controlled interleaving against the exact bridge and Arrow versions
reproduced permanent child-allocator and `TaskContext` retention with
assertions disabled. With assertions enabled, the subsequent buffer release
also fails and retains the buffer. Please include in-flight allocations in the
lifecycle check before closing the child.
--
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]
