feat: report JVM Arrow allocations to Spark's memory manager

[andygrove]

Which issue does this PR close?

Part of #5997. This is the first of the four items listed there, and the only one that needs no native changes.

Rationale for this change

CometArrowAllocator is a single process-wide RootAllocator(Long.MaxValue). Child allocators are cut from it for FFI stream export, broadcast coalescing, CometSparkToColumnarExec, the cached batch serializer, codegen output, and the Python runner. Every byte it hands out is real off-heap memory resident in the container, and none of it is visible to Spark's TaskMemoryManager or to Comet's native memory pool. The memory management guide already lists this as an open problem: "unbounded and accounted by nobody".

That makes these bytes a blind spot exactly when it matters, which is when an executor is killed for exceeding its container limit. Gluten solves the same problem by giving each task its own Arrow allocator, whose listener reserves from Spark in fixed-size blocks.

This PR closes the reporting half of the gap. It deliberately does not close the enforcing half: Arrow allocation on these paths cannot fail today, and making it fail is a behavioural change that deserves its own PR once we have numbers on what the real volumes are.

What changes are included in this PR?

The owner of a buffer is the allocator, not the calling thread. Arrow hands AllocationListener nothing but a size, and it reports a release on whichever thread drops the last reference. For anything exported over the C Data Interface that is a Tokio worker with no task context installed, so attributing by TaskContext would silently drop those releases and leave a task charged, batch after batch, for memory it had already freed. Everything below follows from that.

• CometTaskArrowAllocator.forCurrentTask() hands each task a child of CometArrowAllocator carrying its own CometArrowAllocationListener. BaseAllocator.newChildAllocator passes the listener down, so the call sites that cut their own children — the Python runner's stdin/stdout allocators and CometNativeArrowSource — are accounted without knowing anything about this. Every production allocation of a JVM-owned buffer now goes through the accessor.
• CometArrowAllocationListener is one instance per task and never looks at TaskContext. It charges to a MemoryConsumer belonging to its task, so the bytes appear in TaskMemoryManager.showMemoryUsage and are arbitrated against Spark's other off-heap consumers, and both callbacks land on the allocating task whichever thread they arrive on.
• The root allocator keeps no listener, and everything crossing the FFI boundary uses it, in both directions, because native's pool is the authority for bytes native holds. Coming in, BaseAllocator.wrapForeignAllocation reports an imported buffer to the listener at full capacity even though no JVM-side allocation happened. Going out, whichever DataFusion operator retains the batch reserves those buffers itself, through a unified pool that charges the same Spark task, so charging them here as well would reserve the same memory twice and could reject an allocation that fits. NativeUtil, the JVM UDF result in CometUdfBridge together with the codegen output vector it exports, and CometNativeArrowSource.stream therefore all use the root; IPC reads, the cached batch serializer and CometNativeArrowSource.readerBatchIter use the task allocator. The same fallback covers the driver, broadcast coalescing and the cached batch serializer, which have no task to charge, and Comet's on-heap mode, where charging an off-heap consumer would be wrong. This replaces the separate CometImportedArrowAllocator from the first revision, which is no longer needed.
• The split is by allocation site, so it is not exact. A buffer used in the JVM and only later handed to native is still charged on both sides, and a shuffle-read batch feeding a native operator through exportBatch is the common shape rather than a corner case. The allocation site cannot know where the batch ends up. Fixing that means handing reservation ownership over at the boundary, which is a change on both sides of it, so it is recorded in the guide's open problems and left to Improve memory accounting: bound the JVM Arrow allocator and make native reclaim work #5997.
• Reservations are grown and shrunk in whole blocks (1 MiB, a private constant) because Arrow allocates per buffer and acquireExecutionMemory takes an executor-wide lock. The grow request is rounded up to a block multiple, and excess is returned in a single call rather than one per block.
• getUsed and spill are lock-free. Spark calls both while holding the TaskMemoryManager monitor, and the listener holds its own monitor across acquireExecutionMemory, which takes that monitor. A native reservation arriving through CometTaskMemoryManager on a Tokio thread would otherwise deadlock against an Arrow allocation on the same task. CometTaskMemoryManager.NativeMemoryConsumer already reads an AtomicLong for the same reason.
• Neither callback propagates an exception. Arrow's AllocationListener documents that they may not, and BaseAllocator.buffer sets success = true before calling onAllocation, so a throw loses the buffer Arrow has already created and never hands back. Spark's acquisition is fallible in three ways, and only the first is caught by NonFatal: it runs other consumers' spill, and TaskMemoryManager turns a task interrupt into a RuntimeException and an IOException into a SparkOutOfMemoryError; and ExecutionMemoryPool.acquireMemory parks in lock.wait() when a task is below its fair share, so killing a task raises a plain InterruptedException here. An interrupt is caught with the thread's flag re-armed, so the cancellation is left for the task to observe rather than swallowed.
• A failed acquisition adopts whatever Spark already took. acquireExecutionMemory takes its first grant from the pool and only then asks other consumers to spill, so when a spill throws it has already charged the task for bytes it never reports back, and nothing would release them before the final task cleanup. They are adopted, measured as the change in getMemoryConsumptionForThisTask across the call. Spark reports that figure per task rather than per consumer, so both snapshots and the call run as one transaction under the TaskMemoryManager monitor, which is the same reentrant monitor acquireExecutionMemory itself takes and holds for its whole duration. Every acquisition in the task funnels through that method, so no other consumer can take memory between the snapshot and the call and have it adopted here. A release does not take that monitor, so a concurrent release, or a spill that frees some bytes before throwing, can still make the figure too small, which is the safe direction: too small is what would have happened anyway.
• A task allocator is closed only once it has been drained. The process-wide allocator exists precisely because buffers can outlive their task, and Arrow treats closing an allocator that still owns bytes as a leak. At task completion the Spark reservation is dropped and the allocator is closed if it is empty; otherwise it is parked and reaped by a later task once the stragglers are released. Leaving it open indefinitely is not an option, because BaseAllocator keeps every child in a map until it closes.
• spark.comet.arrowAllocator.accounting.enabled (default true) turns the reporting off, in which case forCurrentTask() returns the root and behaviour is exactly what it is on main. It is declared in CometConf so it reaches the generated configuration docs, and read from the SparkConf because the decision is made on executors, where SQLConf does not carry Comet's settings.
• The memory management contributor guide is brought back in line with the code, including the parts of it that described these buffers as invisible to Spark.

Effect on other memory consumers

Worth calling out explicitly for review, because "reports without enforcing" undersells it. The listener cannot fail an allocation of its own, but it does consume budget: acquireExecutionMemory genuinely takes the memory, so once Arrow bytes are reserved, Spark's other off-heap consumers in the same task see correspondingly less headroom and may spill earlier than they do today. This is therefore not a purely observational change.

That is the intended semantics. These bytes were always resident in the container, and the bug being fixed is that nothing counted them. Accounting is enabled by default, because hiding real bytes is the problem this is meant to solve, and spark.comet.arrowAllocator.accounting.enabled=false restores the previous behaviour if a workload regresses.

This also sharpens the case for the second item in #5997. With Arrow bytes now reserved, the one consumer in a Comet task that still cannot return anything when Spark asks for memory is the native one, whose NativeMemoryConsumer.spill() returns 0.

How are these changes tested?

CometArrowAllocationListenerSuite (18 tests) covers the reservation arithmetic, and then the properties the per-task binding exists for:

• a release on a thread with no task context, and a release arriving under a different task, are both charged to the task that allocated
• task completion releases the whole reservation and closes the allocator, driven through markTaskCompleted rather than by calling cleanUpAllAllocatedMemory directly
• a buffer outliving its task parks the allocator until it is released, and a later task reaps it
• a spill that fails with an IOException or an InterruptedIOException neither fails the Arrow allocation nor leaks the buffer Arrow already created
• an InterruptedException is caught and the thread's interrupt flag is left set
• a partial grant lost to a failing spill is adopted, so the task is not left charged for bytes owned by nobody
• getUsed and spill do not wait on the reservation monitor, probed by holding that monitor and calling both from another thread
• concurrent Arrow allocations and native-style reservations both finish
• allocating from the root inside a task is not reported, and a child of a task allocator is reported to that task
• no active task, and on-heap mode, hand back the unaccounted root

CometArrowAllocationListenerBenchmark measures the overhead with accounting on and off across the three buffer shapes Comet produces, and once more with a second thread reserving from the same constrained pool. Numbers are in a comment below.

Locally, CometExecSuite, CometShuffleSuite, CometNativeShuffleSuite, CometCodegenSuite, CometInMemoryCacheSuite, CometNativeColumnarToRowSuite, CometArrowPythonRunnerSuite, UtilsSuite, NativeUtilSuite, CometExecIteratorLifecycleSuite and CometTaskMemoryManagerSuite all pass, and all four Spark profiles compile.

[sunchao]

Correctness

This closes an accounting gap in the process-wide JVM Arrow allocator. Previously its buffers were outside Spark's execution-memory budget. The PR adds a per-task MemoryConsumer, batches reservations in 1 MiB blocks, and introduces a default-enabled switch. It also routes native FFI imports through a no-listener child allocator so wrapping native-owned buffers does not create another JVM reservation. Arrow 18.3.0's listener dispatch supports that separation.

I reviewed 4d7657d4 against 58ab5f61, including the existing stream, cache, broadcast, Python and JVM-UDF ownership paths. Maintained Spark 3.5 and 4.0 confirm that memory acquisition can spill, throw and call consumers while holding the task-memory-manager lock. Three findings need addressing: [P1] reversed lock ordering can deadlock concurrent Arrow/native requests, [P2] releases use the releasing thread's task instead of the allocation owner, and [P2] a spill exception can escape after Arrow creates a buffer and leak that buffer. The inline comments include the triggers and regression cases.

The local component probe compiled a listener byte-for-byte identical to this head and ran against locally cached Spark 4.0.0 JVM artifacts and real Arrow 18.3.0 buffers on JDK 17. It used Spark's TestMemoryManager, an injected failing spill consumer, deterministic thread barriers and a stub for the Comet config key/default accessor. It reproduced all three boundaries. This is component evidence, not an end-to-end Comet/native-JNI run, and those runtime jars are not claimed to contain the maintained Spark commits. Maintained Spark 3.4/4.1 sources were unavailable.

At September 17, 16:15 UTC, CI was still running, with 17 successful checks, three failed lint checks, 13 skips and three checks in progress. The Spark 4.1/JDK 17 build passed, resolving the earlier line-length gate. The Spark 3.4, 3.5 and 4.0 JDK 17 lint logs fail Scalafix because CometImportedArrowAllocator needs the explicit BufferAllocator result type. These logs check out merge 579794d6, whose parents are the exact base/head and whose tree equals the head. Spark execution/expression checks and TPC-DS verification were still in progress, so this is not a complete green product-validation result.

Performance

Block-sized growth and returning excess in one call avoid repeated pool operations while allocations fit in an existing reservation. The null-task fast path and once-per-JVM config lookup also keep unnecessary work out of driver-side allocation. Every accounted buffer still adds a task-context lookup, concurrent-map lookup and reservation monitor acquisition. A short grant causes later allocations to retry Spark, and freeing back to zero means the next small allocation requests another block.

Because this is enabled by default and consumes real off-heap budget, please include an enabled/disabled microbenchmark covering small FFI buffers, representative wide batches and concurrent native reservations under pressure. Report allocation/release throughput, Spark reservation-call counts and spill behavior. The PR contains no timing evidence, and I have not measured a speedup or regression. The ownership and lock-order findings must be fixed before such a comparison is meaningful.

Design

Keeping FFI imports separate from JVM-owned allocation is the right boundary. The behavior change is more than visibility: reservations reduce the headroom available to other Spark consumers and can make them spill earlier, even though a short grant does not reject the Arrow allocation. The updated guide and PR body make this budget impact explicit, and the opt-out makes that behavior reversible. It does not establish a hard native-memory or RSS limit, and native spill support remains outside this PR.

The remaining design requirement is stable ownership across callbacks. A task-bound listener/allocator or equivalent ownership token should survive cross-thread release and become inert after task cleanup. Reading the current TaskContext independently in each callback cannot provide that contract. Please extend the tests to exercise actual task completion, late releases, failed acquisitions and concurrent native/Arrow reservations. The current fixture calls cleanUpAllAllocatedMemory directly and does not execute the registered completion listener.

Abstraction & complexity

A small listener plus one reservation object per task is a reasonable scope, and a private block-size constant avoids adding another tuning knob without evidence. Registering completion callbacks outside computeIfAbsent correctly avoids recursive map updates when Spark invokes a late listener immediately.

The single shared listener currently mixes allocation-owner discovery, task registration, reservation tracking and release attribution. Separating stable owner binding from byte accounting would remove the need to rediscover a task on release and make cleanup behavior easier to test. The usage snapshot should also avoid introducing a second lock order into Spark's existing memory-consumer protocol.

[sunchao]

Correctness

[P1] Could getUsed expose a lock-free snapshot instead of taking the reservation monitor? allocated holds this monitor while entering TaskMemoryManager.acquireExecutionMemory, while Spark 3.5/4.0 hold the task-memory-manager monitor when calling consumers' getUsed during a short grant. A concurrent native reservation through CometTaskMemoryManager can therefore hold the Spark monitor and wait here while the Arrow allocation holds this monitor and waits for Spark. Comet's shared Tokio execution permits those requests to overlap. A two-thread component probe using the unchanged listener and the real Spark task memory manager reproduces the lock cycle. Please establish a consistent lock order and add a bounded concurrent-reservation regression test.

[andygrove]

Yes, and it does now. getUsed reads an AtomicLong and a volatile flag, and spill returns 0 without touching anything, so neither waits on the reservation monitor. That matches what CometTaskMemoryManager.NativeMemoryConsumer already does, and for the same reason. The lock order is now stated on the class: the reservation monitor may be held across acquireExecutionMemory, so nothing Spark can call while holding the TaskMemoryManager monitor is allowed to need it.

The regression test is deterministic rather than a stress loop. It holds the listener's monitor on the test thread and calls getUsed and spill from another one with a timeout, which fails outright if either becomes synchronized again. There is a bounded two-thread Arrow/native reservation test alongside it, but the monitor probe is the one that actually pins the property.

[sunchao]

Correctness

[P2] Can releases be associated with the allocating task rather than the task currently on the releasing thread? The existing JVM UDF path installs a task context, allocates/exports JVM Arrow buffers, then restores or unsets the context before Rust retains the exported result through from_ffi. Its later release callback can run on a Tokio thread without a task context, so this lookup ignores the release while the allocating task is still running. Repeated batches leave reservations charged for memory already freed and can starve the native pool. In a component probe, allocating 1 MiB under task A and closing it without a context leaves A charged 1 MiB with zero Arrow bytes alive. Closing A's buffer under task B instead subtracts B's accounting. The size-only callback cannot recover ownership from the current task. Please bind the listener/accounting to the allocation owner and cover cross-thread release plus release after task completion.

[andygrove]

They can, but not by looking harder at the callback, since Arrow gives it nothing but a size. The fix is to stop asking the thread and make the allocator the owner: CometTaskArrowAllocator.forCurrentTask() now hands each task a child of the root carrying its own listener instance, so onAllocation and onRelease both land on the allocating task no matter which thread drops the last reference. That is the same shape Gluten uses. BaseAllocator.newChildAllocator passes the listener down, so the Python runner and CometNativeArrowSource, which cut their own children, are covered too.

Falling out of that: the root no longer carries a listener at all, so FFI imports are unaccounted simply by using it, and CometImportedArrowAllocator is gone. Off a task and in on-heap mode the accessor returns the root, so those paths behave exactly as they do on main.

The awkward part is lifetime. The process-wide allocator exists because buffers can outlive their task, and Arrow treats closing a non-empty allocator as a leak, so at task completion the reservation is dropped and the allocator is closed only if it is drained; otherwise it is parked and reaped by a later task. It cannot just be left open, because BaseAllocator keeps every child in a map until it closes.

Tests cover a release with no task context, a release under a different task, completion driven through markTaskCompleted rather than cleanUpAllAllocatedMemory, and a buffer outliving its task and being reaped afterwards.

[sunchao]

Correctness

[P2] Please keep potentially throwing Spark acquisition out of Arrow's onAllocation callback. This call can invoke another consumer's spill method, and maintained Spark 3.5/4.0 propagate spill interruption or I/O failure as exceptions. Arrow 18.3.0 sets success = true and creates the buffer before invoking onAllocation, so an exception here escapes without returning or closing that buffer. With the unchanged listener, a 1 MiB pool occupied by a spill consumer that throws InterruptedIOException makes root.buffer(1 MiB) throw RuntimeException while the root still owns 1 MiB that the caller never received. Closing the allocator then reports that leak. Please move fallible work to a safe allocation boundary with rollback, and add a failing-spill/cancellation test that checks both Arrow live bytes and Spark reservations.

[andygrove]

Agreed, and this is worse than "should not throw": Arrow's AllocationListener javadoc says an exception cannot be thrown from onAllocation or onRelease, and BaseAllocator.buffer sets success = true before the call, so a throw skips the finally { releaseBytes } and the buffer is neither returned nor freed.

Rather than move the acquisition to a different boundary, the listener now simply never lets anything out. Both callbacks wrap the memory-manager call and log once. SparkOutOfMemoryError is caught explicitly because it extends OutOfMemoryError and so slips past NonFatal, which is exactly the IOException-from-spill path you describe.

Two tests, one for each way trySpillAndAcquire fails: a consumer holding the only block that throws IOException, and one that throws InterruptedIOException. Both assert the Arrow allocation succeeds, that the allocator reports the buffer's bytes while it is open and zero after closing it, and that nothing was reserved, which is what proves the throw really happened rather than the acquisition quietly succeeding and making the test vacuous.

[andygrove]

You are right, and it is worse than the NonFatal gap alone. ExecutionMemoryPool.acquireMemory parks in lock.wait() when a task is below its fair share, so a kill raises a plain InterruptedException straight out of acquireExecutionMemory, and scala.util.control.NonFatal excludes InterruptedException by name along with the Error subclasses.

It is caught now, but not simply swallowed: the handler re-arms the thread's interrupt flag before returning. Arrow has already created the buffer by the time onAllocation runs, so returning normally is the only option that does not lose it, and re-arming leaves the cancellation for the task to observe at its next interruptible point, which is the only place it can act on it anyway.

The regression test injects the InterruptedException through a failing spill rather than through the pool, since TestMemoryManager never parks, and asserts all three things: the allocation succeeds, the allocator is back to zero after closing the buffer, and Thread.interrupted() is true afterwards.

[sunchao]

Rechecked a2518d97 against 58ab5f61. Since the previous review, the only change is the explicit BufferAllocator result type and import. This fixes the Scalafix failures without changing allocation or release behavior.

All three findings remain: [P1] lock inversion, [P2] release attribution to the wrong task, and [P2] leaked buffers when Spark acquisition throws inside onAllocation. The listener, relevant callers and tests are unchanged. I rechecked the JVM-UDF release path, native reservation concurrency and task completion against maintained Spark 3.5/4.0. The prior component reproductions remain applicable by verified source/dependency identity. No new inline findings.

At September 17, 17:20 UTC, CI is green. The three previously failing lint jobs pass, and the execution log confirms all six listener tests passed within 922 successful tests. The verified checkout is merge d00e8133, with the exact base/head as parents and the same tree as the head. The listener tests still do not exercise concurrent reservation, cross-thread release or failed acquisition. The earlier enabled/disabled benchmark request also remains unanswered. No new local runtime or timing measurements were needed for this type-only change.

[andygrove]

Triage note: #5027 overlaps with this on the UDF output path. It cuts a per-task Arrow child allocator backed by a non-spillable MemoryConsumer for CometScalaUDFCodegen's output vectors, which accounts that one path more precisely than a root-level listener does, and it changes the CometUDF interface to pass the allocator through. Both PRs touch CometUdfBridge.java, and if both land unchanged those output buffers are charged twice.

@peterxcli, worth deciding which layer owns those buffers. My instinct is that the root listener here covers everything by default and the per-task consumer in #5027 should replace it on the UDF path rather than stack on top, but I have not tried writing that.

[andygrove]

Benchmark numbers for the accounting overhead, as asked for. Committed as CometArrowAllocationListenerBenchmark so this is reproducible:

SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.comet.CometArrowAllocationListenerBenchmark

Each case runs the same allocate-then-release loop against a plain RootAllocator and against a task allocator with the listener attached. Both allocators are built once outside the timed body, so this is steady-state cost and not the cost of standing a task up. Apple M3 Max, JDK 17, three runs; the ranges below are across those runs, and the machine noise is wide enough that I would not read much into anything under about 10%.

case	not accounted	accounted	Spark calls per iteration
512 x 128 B (the FFI struct shape)	146-150 ns/buffer	159-169 ns/buffer	1 acquire, 1 release
128 x 64 KiB (a wide batch)	1180-1506 ns/buffer	1424-1457 ns/buffer	8 acquire, 8 release
8 x 1 MiB (every buffer crosses a block)	3646-5036 ns/buffer	3766-5135 ns/buffer	8 acquire, 8 release
8 x 1 MiB, second thread reserving from a 2 MiB pool	3724-5115 ns/buffer	4844-7750 ns/buffer	13 acquire, 2 release
forCurrentTask() vs reading the old val	~0.4 ns (folded away)	3.4-4.8 ns	n/a

Reading it:

• The block batching does what it is meant to. 512 small buffers, which is 64 KiB of peak, reach the memory manager exactly twice in total. That is the NativeUtil path, and it is the one that runs per column per batch.
• Where it costs something is buffers at or above a block, where every one of them is an acquireExecutionMemory plus a releaseExecutionMemory round trip. About 100-850 ns per buffer, which is two executor-wide lock acquisitions.
• Under contention the picture changes in a way worth noting: with only two blocks of pool and another thread holding one, grants come back short and the listener re-asks, so a single iteration makes 13 acquire calls for 8 buffers. That is the worst case in this PR, and it is the one that argues for enforcement being a separate change.
• forCurrentTask() costs about 4 ns, and it is per call site invocation rather than per buffer: once per NativeUtil, once per StreamReader, once per codegen output vector.

On spill behaviour: the accounted consumer never spills, spill returns 0 like NativeMemoryConsumer does, so it is only ever a spill trigger, never a victim that frees anything. What it does do is take budget, which is the point made under "Effect on other memory consumers" in the description: other consumers in the same task can spill earlier than they do today. The failing-spill paths are covered in the suite rather than the benchmark, since what matters there is that the failure does not escape into Arrow.

[sunchao]

Re-reviewed d591b8f2 against 58ab5f61. The per-task allocator fixes the P1 lock inversion and P2 release attribution. Rerunning the original failure probes and current-source controls confirms both fixes. The added benchmark addresses the request for accounting-on/off and contended-pool measurements. Its timings remain author-reported.

The P2 callback-exception finding is partially fixed but remains for plain InterruptedException. Interrupting a thread blocked in Spark's real execution-memory pool reaches lock.wait(). That exception is excluded by NonFatal and escapes onAllocation. With the current listener and real Arrow, the caller receives no buffer while the allocator retains 1 MiB. Releasing the competing reservation instead is a passing control. Could this cancellation boundary preserve Arrow's allocation lifecycle as well as cancellation semantics?

Two additional P2 findings are inline: a partial Spark grant is lost when a later spill throws, and exported JVM-owned UDF buffers retain their JVM charge while a native hash join reserves the same buffers. The former is reproduced through real Spark TaskMemoryManager. The latter combines a real Arrow C Data/Java bridge component check with the pinned DataFusion/Arrow ownership trace. I independently repeated these component checks, but did not run a full Comet JNI query.

At 2026-09-17 20:42 UTC, CI has 22 successful checks, 13 skipped and the Spark 4.1 expressions job still running, with no failures reported. Its checkout is merge 7c0cf107 over 3b942e59, rather than the assigned base. All 17 PR-authored files are identical to the reviewed head, but this is not exact-base CI. Maintained Spark 3.5/4.0 sources were inspected. Required 3.4/4.1 branches are unavailable locally.

[sunchao]

Correctness

[P2] Account for a partial grant when Spark's later spill attempt throws

TaskMemoryManager.acquireExecutionMemory can acquire some bytes from the execution pool and then throw while asking another consumer to spill. In that case it never returns granted, so reserved stays unchanged, while adjustQuietly swallows the exception and lets the Arrow allocation succeed. With a 2 MiB pool, another consumer holding 1 MiB, and a 2 MiB Arrow request, I reproduced this with both IOException and InterruptedIOException: after closing the Arrow buffer and freeing the other consumer, Arrow and reservedBytes are zero but Spark still charges the task 1 MiB. taskCompleted() cannot return that unrecorded grant. It lasts until Spark's final task cleanup, reducing headroom for the rest of the task. Could acquisition failure reconcile or roll back partial grants? A regression with partially available capacity would cover this. The current full-pool failure tests only exercise a zero-byte initial grant.

[andygrove]

Confirmed, and thank you for separating this from the throw itself — I had only fixed the escape, not the accounting behind it. acquireExecutionMemory takes its first grant from the pool before it asks anyone to spill, so on the throw path it has charged the task for bytes it never returns, and since we swallow the exception the task carries on without them for the rest of its life.

Rolling back is not possible from outside, because Spark never tells us how much it took. So acquire adopts them instead, measured as the change in getMemoryConsumptionForThisTask across the call. That is an estimate, but a safe one in the direction that matters: acquireExecutionMemory holds the TaskMemoryManager monitor for its whole duration, so no other consumer in the task can acquire concurrently and inflate the figure, and a concurrent release can only deflate it, which degrades to today's behaviour. The request bounds it from above either way.

The new test is the partially-available case you describe — a two-block pool, one block held by a consumer that throws, and a two-block request. It asserts a block is reserved immediately after the allocation, and that after closing the buffer the task is charged only for the other consumer's block rather than for two.

[sunchao]

Correctness

[P2] Give exported JVM output a single Spark reservation owner

This charges codegen output to the task allocator, but CometUdfBridge exports that vector through C Data and closes only its original reference. Export retains the same buffers and their JVM charge until native code releases them. With a unified native memory pool, the existing DataFusion 55.1.0 hash-join build path then counts and reserves the incoming batch's buffers through Comet's unified pool, charging the same Spark task again. Arrow's foreign-buffer capacity is nonzero. The custom allocation retains its size. I verified the export lifetime with real Arrow C Data export/import and the current Java memory bridge: an unchanged data address retained a 2 MiB JVM reservation while the corresponding 1,064,960-byte native-style reservation also succeeded. This is a component reproduction plus a source trace, not an end-to-end join run. Under a bounded pool, this reduces available headroom twice for one allocation and can reject the join build. Please transfer or otherwise coordinate reservation ownership at export and cover JVM-UDF output feeding a native hash join. This path already exists without unmerged #5027.

[andygrove]

Agreed, and the fix is to stop charging it here rather than to coordinate a handoff, which is not something this PR can do from one side of the boundary.

The rule is now that native's pool is the authority for bytes native holds, in both directions. Imports already used the listener-less root; exports now join them, so NativeUtil, the JVM UDF result and CometNativeArrowSource.stream all allocate from the root, while IPC reads, codegen output, the cached batch serializer and CometNativeArrowSource.readerBatchIter keep the task allocator. Those three sites exist only to hand bytes to native, so the JVM-side charge was pure duplication with no information in it.

The part I want to be explicit about rather than have you find it: this is a split by allocation site, so it is not exact. A buffer used in the JVM and only later handed to native is still charged on both sides, and a shuffle-read batch feeding a native operator through exportBatch is the common shape, not a corner. The allocation site cannot know where the batch ends up. That residual is written into the guide's open problems next to the existing entry about buffer and reservation lifetimes being independent across the boundary, and it is the same problem: which side owns a reservation is decided by where the bytes were allocated rather than by who is holding them. Closing it properly is a change on both sides, so it belongs with the rest of #5997 rather than here.

[andygrove]

Confirmed, and you are right that the previous revision only moved the exporter. CometUdfBridge switched to the root for its import and its Data.exportVector, but the output buffers were allocated before either of those, and CometBatchKernelCodegenOutput.allocateOutput still selected the task allocator. So the JVM charge stood while native reserved the same buffers.

It allocates from the listener-less root now. The only production caller is CometScalaUDFCodegen.evaluate, which is the only CometUDF implementation, and its result is consumed solely by CometUdfBridge, which exports it over the C Data Interface and then closes its own reference. Every buffer this method hands out is therefore allocated to be handed to native, which puts it under the same rule as NativeUtil, the imported inputs and CometNativeArrowSource.stream: native's pool is the authority for bytes native holds, so charging them here reserved the same memory twice.

That should close the 2 MiB JVM reservation you measured alongside the 1,064,960-byte native one, since nothing on the JVM side reserves those buffers any more. The contributor guide listed codegen output under the task allocator and is corrected.

One side effect worth recording: this also unstacks the overlap with #5027. With codegen output no longer charged here, the per-task consumer in that PR can own the UDF output path outright rather than adding a second charge on top of this one.

Latest main is merged in, which clears the conflict that was holding CI to the label check.

[sunchao]

Re-reviewed 054fe9a2 against 58ab5f61. The P2 interruption leak is fixed: rerunning the real Spark pool-wait case returns the Arrow buffer, preserves the interrupt flag and releases the buffer cleanly. The original P2 partial-grant leak also passes with both I/O-failure variants. The new recovery method introduces a different P2 race, described inline: it can adopt and release another consumer's reservation because the before/after snapshots are outside Spark's acquisition lock.

The P2 UDF export ownership finding remains on the original codegen path. CometUdfBridge now uses the root for imports and Data.exportVector, but udf.evaluate(inputs, numRows) still reaches CometScalaUDFCodegen → CometBatchKernelCodegenOutput.allocateOutput, which selects the task allocator at line 91. The export allocator does not transfer the result buffers' owner. I independently reran the C Data/Java bridge check, which still retains the 2 MiB JVM charge alongside the 1,064,960-byte native-style reservation for the same buffers. Could the actual output allocation or ownership handoff be changed? This is the originally reported UDF path, separate from the acknowledged shuffle-read limitation.

At 2026-09-17 23:11:48 UTC, CI is green with 24 successful checks and 14 skips. The execution log confirms all 18 listener tests within 944 successful tests. CI checked out f4c45949 over d0c10f2d, not the assigned base. The listener and its tests match the reviewed head. I independently reran the original failure cases, the new race and its control. These local checks used cached Spark 4.0.0/Arrow 18.3.0 components and real Arrow C Data, without a full Comet JNI query. Maintained Spark 3.5/4.0 sources were checked. Required 3.4/4.1 branches remain unavailable.

[sunchao]

Correctness

[P2] Keep both usage snapshots inside the acquisition transaction

heldBefore is read before acquireExecutionMemory takes the task-manager monitor, and adoptOrphanedGrant reads again after that monitor has been released on the exception path. Another consumer can therefore acquire between either snapshot and the protected call. This lets the listener adopt and later release someone else's reservation, rather than only underestimate its own grant. With the exact current listener and real Spark UnifiedMemoryManager/TaskMemoryManager, I paused after the real before-snapshot, let another OFF_HEAP consumer take the 1 MiB pool, and then let Arrow's zero-grant acquisition trigger that consumer's failing spill. The listener adopted 1 MiB despite receiving no grant. Closing the Arrow buffer reduced Spark's charge to zero while the other consumer still owned 1 MiB. A control that acquired before the snapshot retained the correct 1 MiB charge. Could failure recovery use an atomic, consumer-specific accounting boundary and add this interleaving as a regression? Bounding the result by request does not establish that those bytes belong to this listener.

[andygrove]

Confirmed, and the argument I gave for the old code being safe was wrong in exactly the place you point at: the monitor is held for the whole of acquireExecutionMemory, but heldBefore was read before entering it, so the window was open regardless of what happens inside.

There is no consumer-specific figure to read instead. getMemoryConsumptionForThisTask is the whole task, and this consumer's inherited used stays at zero by design, because it never calls acquireMemory — Arrow has already taken the memory and we are only accounting for it. So what makes the task-wide delta attributable has to be the boundary rather than the quantity: acquire now takes the TaskMemoryManager monitor itself, and both snapshots and the acquisition run inside it. That is the same monitor acquireExecutionMemory takes and holds for its whole duration, spills included, and it is reentrant, so this only widens the window to cover the two reads. Every acquisition in the task funnels through that method, so nothing else in the task can take memory in between and have it adopted here.

What the monitor does not cover is a release, which reaches the pool without it. Another consumer returning memory, or a spill that frees some bytes before throwing, still makes the figure too small. That is the safe direction and degrades to what would have happened anyway, and request still bounds it from above.

The regression is your interleaving, made deterministic rather than timed. A TaskMemoryManager subclass runs a second consumer immediately after the real before-snapshot, which is the pause you described, and then waits until that consumer is either blocked or finished, so the two orderings are told apart by thread state rather than by sleeping. The assertion is that nothing is claimed twice: what the listener adopted, plus what the other two consumers hold, equals what the task holds. I ran it against the unfixed listener first, where it fails with 2099200 adopted against a task total of 2098176 — the interloper's 1024 charged to the listener on top of its own grant — and green afterwards.

The lock order note on the class is updated too, since the listener now takes Spark's monitor explicitly: ours before Spark's, never the reverse, which is why getUsed and spill have to stay lock-free.

[sunchao]

Re-reviewed c2a542d6 against 58ab5f61. The P2 orphan-grant race is fixed: both snapshots now share Spark’s acquisition monitor. I reran the old-head failure and current-head control with real Spark memory managers. The competing consumer now blocks, and closing Arrow releases only its own reservation. The earlier interruption, partial-grant, cross-thread release and task-completion checks still pass.

The P2 UDF export ownership finding remains. The output allocation still selects the task allocator, and the C Data/Java bridge probe still retains the 2 MiB JVM reservation alongside the 1,064,960-byte native-style reservation for the same buffers. The exporter’s root allocator does not transfer ownership. No new findings or duplicate inline comments.

Current-head CI has only the label check; GitHub reports merge conflicts and no merge commit. Local validation used cached Spark 4.0.0/Arrow 18.3.0 components with exact Comet sources, not a full Comet JNI query. Maintained Spark 3.5/4.0 sources were checked. Required 3.4/4.1 branches remain unavailable.