Improve memory accounting: bound the JVM Arrow allocator and make native reclaim work

[andygrove]

Background

Comet's memory accounting bounds declared reservations, while the kernel kills on resident pages. The gap between those is documented in the memory management guide and tracked in #4576.

Two recent results narrow the design space:

• Thresholding a cgroup counter does not work. feat: add optional container memory guard based on cgroup usage #5993 sampled memory.current and failed a task before the kernel could kill the executor. It was validated end to end on Kubernetes, and then closed: memory.current includes reclaimable page cache, so it saturates at the container limit on any workload that reads files. In an A/B on TPC-H SF100 Q9 at a 9 GiB pod limit, the guard failed the job while the identical run with the guard off completed in 58.62s, and nothing was ever OOMKilled. Measurements are on Investigate adopting DataFusion's allocator-level memory accounting to replace manual memory tuning #4576.
• Gluten does not solve this either, and that is informative. Its Velox path reads nothing from the OS or cgroup at runtime: no memory.current, no PSI, no RSS, no watchdog. Its whole defence is the accounting tree plus a separately budgeted spark.executor.memoryOverhead sized to max(0.3 * offHeap, 384MiB). Notably, the Bolt backend documents an out-of-tree RSS-driven quota manager whose stated purpose is to "reduce the occurrence of CgroupKill errors", which is direct evidence that accounting alone was not sufficient at scale.

The conclusion is that there is no cheap runtime signal to enforce against. What is available is closing the accounting gaps we already know about, and Gluten has working implementations of most of them. This issue collects four, ordered by cost.

1. Bound the JVM-side Arrow allocator

CometArrowAllocator (spark/src/main/scala/org/apache/comet/package.scala) is a single process-wide new RootAllocator(Long.MaxValue). Child allocators are cut from it for FFI stream export, broadcast coalescing, and CometSparkToColumnarExec. These are real off-heap bytes in container RSS that neither Spark's TaskMemoryManager nor Comet's native pool sees. There is no ceiling and no backpressure.

Gluten keeps the same Long.MAX_VALUE limit but attaches a ManagedAllocationListener to the RootAllocator, so every allocation borrows from a real Spark memory target in fixed-size blocks and an over-limit allocation throws instead of growing silently. It also accumulates leaked bytes at task end rather than ignoring them.

This needs no native changes and is the cheapest item here.

2. Make native reclaim actually return memory

NativeMemoryConsumer.spill() returns 0 (CometTaskMemoryManager.java). Comet registers a consumer so Spark has something to charge, but Spark can never make a native operator release anything. A JVM consumer in the same task can be starved behind native reservations with no way to reclaim them.

Gluten's TreeMemoryConsumer is a real Spark MemoryConsumer whose spill() runs two ordered phases over a tree of targets:

• SHRINK: return unused native pool capacity and push the unreserve back to Spark.
• SPILL: genuine operator spilling to disk.

The SHRINK half is much cheaper to build and closes the starvation hole on its own, because it only needs to hand back capacity that is reserved but unused. Full spill is harder for us than for Gluten, since DataFusion has no MemoryArbitrator equivalent and we would need a per-task registry of spillable native handles plus a JNI reclaim entry point returning bytes actually freed.

Trap worth recording up front. Gluten deliberately does not track the memory manager used during spill against Spark's off-heap pool, with the comment that doing so causes "recursive reservations from Spark off-heap memory pool ... to cause unexpected OOMs". The reclaim path itself needs memory; if reclaiming requires reserving from the pool being relieved, it deadlocks.

3. Validate off-heap settings and size the overhead

Gluten fails fast at driver start if spark.memory.offHeap.enabled is false or the size is trivially small, and auto-sizes spark.executor.memoryOverhead to max(0.3 * offHeap, 384MiB) when the user has not set it, warning if the user set something smaller. The overhead budget is where it deliberately puts everything it does not track.

Comet has no equivalent validation, and spark.comet.exec.memoryPool.fraction currently asks operators to hand-tune a margin instead. This is pure Scala.

4. Round reservations to blocks

Gluten rounds reservations up to whole blocks (8MB by default) so only block-crossing deltas cross into the JVM, damping per-allocation chatter on the JNI path. Small and self-contained.

Out of scope

Fragmentation and allocator overhead. Gluten does not track these either: it uses plain malloc, does not link jemalloc by default, has no malloc_trim or arena tuning, and simply relies on overhead headroom. We are in the same position and this issue does not attempt to change that.

Plan

Item 1 first, as a self-contained PR. The others to follow as separate PRs.