sunchao commented on code in PR #5027:
URL: https://github.com/apache/datafusion-comet/pull/5027#discussion_r4171966940
##########
spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java:
##########
@@ -299,4 +357,511 @@ private static boolean isOneOf(ValueVector result,
ValueVector[] inputs) {
}
return false;
}
+
+ /**
+ * Moves the result's buffer accounting from the task allocator to the root
allocator and returns
+ * the recorded bytes of the chunks the result owns exclusively (see {@link
#chargedOutputSize}),
+ * which drops their Spark task charge. Neither the ownership transfer nor
the eventual FFI
+ * release is observed by the task allocator's {@link AllocationListener}
(Arrow only notifies the
+ * allocator owning a chunk when the chunk is destroyed), so the charge must
be released here.
+ * When Spark granted less than was recorded, {@link TaskState#release}
settles these bytes
+ * against the shortfall first, so the export never hands Spark back more
than it granted.
+ *
+ * <p>The returned vector shares the original buffers and must be closed by
the caller after
+ * export; the exported FFI array keeps the buffers alive until native
execution releases them.
+ */
+ private static FieldVector transferForExport(
+ TaskState state, BufferAllocator outputAllocator, FieldVector result) {
+ long charged = chargedOutputSize(result, outputAllocator);
+ TransferPair transferPair = result.getTransferPair(result.getField(),
ROOT_ALLOCATOR);
+ transferPair.transfer();
+ state.releaseExportedCharge(charged);
+ return (FieldVector) transferPair.getTo();
+ }
+
+ /**
+ * Bytes of the result's buffers currently accounted against the task
allocator, i.e. the recorded
+ * bytes that move to native ownership on export. {@code getAccountedSize()}
is non-zero only on
+ * the ledger that owns a chunk, so pass-through input buffers (owned by the
root allocator) and
+ * chunks whose ownership already moved on a previous export contribute
nothing.
+ *
+ * <p>Buffers are enumerated recursively from each vector's physical field
buffers rather than
+ * through {@code getBuffers(false)}: that method omits the allocated
buffers of zero-length
+ * children (an empty list's data vector still holds its allocated
capacity), but {@code
+ * TransferPair.transfer()} moves those chunks to the root allocator all the
same, and a charge
+ * not counted here would never be released.
+ *
+ * <p>A chunk is counted only when every live reference to its ledger comes
from the result tree
+ * ({@code getRefCount()} equals the number of distinct result buffers on
that ledger). A chunk
+ * shared with a retained scratch vector (the documented per-task
scratch-buffer contract, e.g. an
+ * aligned {@code splitAndTransfer} slice) keeps its Spark charge: when
native execution releases
+ * the FFI result, Arrow hands ownership back to the surviving scratch
ledger without any listener
+ * callback, and the eventual scratch close fires {@link
TaskState#onRelease}, which must then
+ * release a charge exactly once. Dropping the charge at export as well
would release it twice. If
+ * the scratch side instead closes while native still holds the buffers, the
retained charge is
+ * dropped wholesale at task completion, matching the pre-export accounting
model.
+ */
+ private static long chargedOutputSize(FieldVector result, BufferAllocator
outputAllocator) {
+ Set<ArrowBuf> seenBuffers = Collections.newSetFromMap(new
IdentityHashMap<>());
+ IdentityHashMap<ReferenceManager, Integer> resultRefs = new
IdentityHashMap<>();
+ collectPhysicalBuffers(result, seenBuffers, resultRefs);
+ long charged = 0L;
+ for (Map.Entry<ReferenceManager, Integer> entry : resultRefs.entrySet()) {
+ ReferenceManager referenceManager = entry.getKey();
+ if (referenceManager.getAllocator() == outputAllocator
+ && referenceManager.getRefCount() == entry.getValue()) {
+ charged += referenceManager.getAccountedSize();
+ }
+ }
+ return charged;
+ }
+
+ private static void collectPhysicalBuffers(
+ FieldVector vector,
+ Set<ArrowBuf> seenBuffers,
+ IdentityHashMap<ReferenceManager, Integer> resultRefs) {
+ for (ArrowBuf buf : vector.getFieldBuffers()) {
+ if (seenBuffers.add(buf)) {
+ resultRefs.merge(buf.getReferenceManager(), 1, Integer::sum);
+ }
+ }
+ for (FieldVector child : vector.getChildrenFromFields()) {
+ collectPhysicalBuffers(child, seenBuffers, resultRefs);
+ }
+ }
+
+ /** 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();
+ }
+
+ /** Visible to the focused allocator test in this package. */
+ static Runnable beginTaskEvaluation(TaskContext taskContext) {
+ TaskState state = taskState(taskContext);
+ state.beginEvaluation();
+ return state::finishEvaluation;
+ }
+
+ /** Visible to the focused allocator test in this package. */
+ static FieldVector transferOutputForExport(TaskContext taskContext,
FieldVector result) {
+ TaskState state = taskState(taskContext);
+ return transferForExport(state, state.allocator(), result);
+ }
+
+ private static TaskState taskState(TaskContext taskContext) {
+ TaskState state =
+ TASKS.computeIfAbsent(
+ taskContext,
+ context -> new TaskState(context,
CometTaskContextShim.taskMemoryManager(context)));
+ // Registered outside computeIfAbsent: on an already-completed task with
no listener drain in
+ // progress, Spark invokes the listener on this thread, and the listener
removes this entry,
+ // which the computeIfAbsent mapping function is not allowed to do. That
immediate invocation
+ // marks a straggler's recreated state completed so its beginEvaluation
fails cleanly; during
+ // an active drain the listener is queued instead and cleans the state up
right after.
+ state.ensureCompletionListenerRegistered();
+ return state;
+ }
+
+ /**
+ * Per-task UDF instances, Arrow allocation accounting and non-spillable
Spark memory consumer.
+ * Arrow reaches it through a {@link TaskAllocationListener}, which lets go
of it when the task
+ * retires (see {@link #closeIfIdle}).
+ *
+ * <p>Spark accounting requires off-heap Tungsten memory ({@code
spark.memory.offHeap.enabled}).
+ * Arrow buffers are off-heap, so with on-heap Tungsten memory there is no
matching Spark pool to
+ * charge and accounting is skipped entirely; the allocator still tracks
buffers for cleanup.
+ *
+ * <p>Allocations are recorded, never refused for lack of Spark memory.
{@link #onPreAllocation}
+ * charges Spark whatever it grants and carries the rest as a shortfall.
Refusing would fail the
+ * task at the wrong place: native operators reserve through consumers whose
{@code spill} returns
+ * 0, and they spill only when their own {@code try_grow} fails, so by the
time a task is under
+ * pressure native execution has already filled its share, and a UDF
allocation that asks just
+ * before the native operator downstream would be the one refused, where
that operator could have
+ * spilled instead. Comet's native pools record an infallible {@code grow}
the same way, as
+ * overcommit (see {@code SparkMemory} in {@code spark_memory.rs}).
+ *
+ * <p>Accounting invariant. Let <i>recorded</i> be the bytes passed to
{@link #onPreAllocation}
+ * and not yet returned through {@link #release} (from {@link
TaskAllocationListener#onRelease},
+ * {@link TaskAllocationListener#onFailedAllocation} or {@link
#releaseExportedCharge}). Outside
+ * an allocation in progress, {@code consumer.getUsed() + shortfall ==
recorded}: every recorded
+ * byte is backed by Spark's grant or by the shortfall, never both. A
release repays the shortfall
+ * before handing anything back, so afterwards Spark is charged {@code
min(its previous grant,
+ * recorded)}: for as much of the outstanding memory as it granted, and only
the excess of its
+ * grant over what is still outstanding is returned. Independently, {@link
+ * TaskMemoryConsumer#freeMemory} never hands back more than the consumer
holds. That bound needs
+ * no bookkeeping to be right, so it also holds for chunks Arrow moves into
or out of the task
+ * allocator without a listener callback (ownership transfers), where the
invariant above cannot
+ * see the change. Task completion returns exactly what Spark still has
granted and forgets the
+ * shortfall.
+ *
+ * <p>Lock order: the {@link TaskMemoryManager} monitor, then this {@code
TaskState} monitor, then
+ * Spark's {@code MemoryManager} monitor (taken inside {@code
acquireExecutionMemory} / {@code
+ * releaseExecutionMemory}). {@link #onPreAllocation} takes all three in
that order; every other
+ * path ({@link #release}, which buffer releases and exports go through, and
{@link
+ * #taskCompleted}) takes a suffix: this monitor, then the MemoryManager
monitor via {@code
+ * MemoryConsumer.freeMemory}, which never touches the TaskMemoryManager
monitor (verified against
+ * Spark 3.5 and 4.1). Never acquire the TaskMemoryManager monitor while
holding this monitor.
+ * Spark's {@code acquireExecutionMemory} may call {@code spill()} on other
consumers while
+ * holding the TaskMemoryManager monitor; a spill that releases Arrow
buffers re-enters {@link
+ * #release} in the same TaskMemoryManager-then-TaskState order, so no
inversion arises there
+ * either. The {@link TaskAllocationListener} monitor comes last: {@link
#allocator()} takes it
+ * under this monitor, and nothing takes another of these monitors while
holding it.
+ */
+ private static final class TaskState {
+ private final TaskContext taskContext;
+ private final long taskAttemptId;
+ private final TaskMemoryManager taskMemoryManager;
+ private final TaskMemoryConsumer consumer;
+ private final TaskAllocationListener listener;
+ private final ConcurrentHashMap<String, CometUDF> instances = new
ConcurrentHashMap<>();
+
+ // Arrow updates allocator accounting after onPreAllocation returns.
+ private int evaluationsInFlight;
+ // Recorded bytes Spark did not grant and no release has repaid yet. See
the class comment.
+ private long shortfall;
+ private boolean completed;
+ // Set once the task has completed and nothing is in flight; see
closeIfIdle.
+ private boolean retired;
+ private boolean completionListenerRegistered;
+
+ private TaskState(TaskContext taskContext, TaskMemoryManager
taskMemoryManager) {
+ this.taskContext = taskContext;
+ this.taskAttemptId = taskContext.taskAttemptId();
+ this.taskMemoryManager = taskMemoryManager;
+ this.consumer =
+ taskMemoryManager.getTungstenMemoryMode() == MemoryMode.OFF_HEAP
+ ? new TaskMemoryConsumer(taskMemoryManager)
+ : null;
+ if (consumer == null) {
+ LOG.debug(
+ "JVM UDF Arrow allocations for task {} are not charged to Spark:
Tungsten memory "
+ + "mode is on-heap and Arrow buffers are off-heap",
+ taskAttemptId);
+ }
+ this.listener = new TaskAllocationListener(this);
+ }
+
+ private void ensureCompletionListenerRegistered() {
+ synchronized (this) {
+ if (completionListenerRegistered) {
+ return;
+ }
+ completionListenerRegistered = true;
+ }
+ // Outside the monitor: on an already-completed task Spark invokes the
listener on this
+ // thread, and taskCompleted takes this monitor itself.
+ taskContext.addTaskCompletionListener((TaskCompletionListener) ignored
-> taskCompleted());
+ }
+
+ private synchronized BufferAllocator allocator() {
+ if (completed) {
+ throw new IllegalStateException(
+ "Cannot allocate JVM UDF memory after task " + taskAttemptId + "
completed");
+ }
+ return listener.allocator();
+ }
+
+ private void onPreAllocation(long size) {
+ // Spark's executor cleanup also synchronizes on TaskMemoryManager. Keep
that cleanup from
+ // overtaking an allocation in progress, while leaving this TaskState
monitor free for buffer
+ // releases that can satisfy a blocking acquire.
+ synchronized (taskMemoryManager) {
+ synchronized (this) {
+ if (completed) {
+ // The only refusal, and not for lack of memory: charging a
finished task would leak
+ // the grant, because Spark releases everything the task held
right after the
+ // completion listeners run.
+ throw new OutOfMemoryException(
+ "Cannot allocate " + size + " JVM UDF bytes after task
completion");
+ }
+ }
+
+ // A throw from here is a failure rather than a short grant (another
consumer's spill
+ // failed, or the task is being killed). It propagates, and nothing
has been recorded.
+ long acquired = consumer == null ? size : consumer.acquireMemory(size);
Review Comment:
[P2] [P2] Yield the Tokio worker before blocking on Spark memory
Could we wrap the UDF JNI invocation in `tokio::task::block_in_place`, as
`SparkMemory::ask_spark` already does? This new acquisition can wait inside
`ExecutionMemoryPool.acquireMemory` when another task holds the pool and the
requesting task is below its minimum share. `JvmScalarUdfExpr::evaluate`
invokes the bridge directly on a Tokio worker without that guard. With two
Spark tasks and `COMET_WORKER_THREADS=1`, a waiting UDF therefore prevents
queued native work from releasing the memory it needs, hanging execution.
Accepting short grants does not help because this wait happens before a grant
returns. The expected behavior is for memory-release work to remain runnable
while the allocation waits.
Evidence: A bounded JNI component probe used the exact-head bridge, Spark
4.1.3, Arrow 18.3.0 and Tokio 1.53.1. Task A reserved the entire 64 MiB
off-heap pool through `CometTaskMemoryManager`. Task B invoked
`CometUdfBridge.evaluate` on a single Tokio worker, requesting an 8,192-byte
Arrow output. A release task queued on that runtime could not run during the
750 ms observation window. The captured worker stack showed
`ExecutionMemoryPool.acquireMemory` through `TaskState.onPreAllocation:586`;
external cleanup released A's reservation to terminate the probe. Both controls
completed: wrapping the invocation in `block_in_place`, and allocating the
output from the root allocator as before this PR. Source tracing confirms
`BatchProducer::spawn` polls native streams on Tokio workers and neither
`JvmScalarUdfExpr::evaluate` nor `JVMClasses::with_env` supplies the guard.
Probe sources and output:
`/tmp/comet-5027-42bc-validation-ha_29ln5/TokioUdfProbe.java`,
`tokio-probe/src/lib.rs`, an
d `tokio.log`. This was a component reproduction, not an end-to-end native
query.
--
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]
