Reduce JNI round-trips in the unified memory pools (batching, hysteresis, cheaper call path)

[andygrove]

What / Why

Every DataFusion try_grow / shrink against CometUnifiedMemoryPool or CometFairMemoryPool becomes a JNI round-trip into CometTaskMemoryManager, and there is no batching, no slack, and no hysteresis: memory is acquired in exactly the requested amount and released the moment a reservation shrinks, however small the amount.

This breaks out positions 17 and 18 of #5212 into a standalone issue so the work can be sequenced and benchmarked on its own.

Anatomy of one grow/shrink

native/core/src/execution/memory_pools/unified_pool.rs:64-78

fn acquire_from_spark(&self, additional: usize) -> CometResult<i64> {
    let handle = self.task_memory_manager_handle.as_obj();
    JVMClasses::with_env(|env| unsafe {
        jni_call!(env,
          comet_task_memory_manager(handle).acquire_memory(additional as i64) -> i64)
    })
}

Each call pays:

1. JVMClasses::with_env (native/jni-bridge/src/lib.rs:339-368) — attach_current_thread_guard plus with_local_frame, i.e. a PushLocalFrame/PopLocalFrame pair, even though acquire_memory(long) -> long and release_memory(long) -> void create no local references at all.
2. call_method_unchecked plus check_exception (native/jni-bridge/src/lib.rs:78-100).
3. On the JVM side, internal.acquireExecutionMemory (spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java:62), which takes Spark's executor-wide synchronized MemoryManager lock. This is the part that scales badly — every concurrent task's grow and shrink contends on it.

Callers hit this per batch: DataFusion's aggregate (try_resize per input batch), sort accumulation and repartition, plus Comet's own native shuffle writer, which does try_grow per batch and reservation.free() per flush (native/shuffle/src/partitioners/multi_partition.rs:463, :528) — a full release-then-reacquire cycle on every spill.

Options

A. Chunked acquire with retained slack

Track two numbers instead of one: used (what DataFusion reservations hold) and granted (what Spark has handed us), with the invariant granted >= used.

• try_grow(n): if used + n <= granted, bump used and return — no JNI at all. Otherwise ask Spark for max(need, chunk).
• shrink(n): decrement used, and call Spark only when granted - used exceeds the slack cap, releasing chunk-aligned amounts.
• Release all remaining granted in Drop. Note that today's Drop only warns (unified_pool.rs:81-91), so this has to be added regardless.

One knob, something like spark.comet.exec.memoryPool.acquireChunkSize, defaulting to a few MB with 0 disabling the behaviour; slack cap equal to the chunk size. Steady-state churn within a chunk becomes free, which is exactly the shuffle writer's access pattern.

Two sharp edges:

• acquireExecutionMemory will spill other consumers to satisfy an oversized request, so only round small requests up; pass large ones through at their exact size.
• On a short grant, keep it when grant >= need, otherwise release it and return resources_err as today, so the DataFusion-side spill still triggers. Hoarding memory Spark is actively contending for is worse than the extra round-trip.

Implementation note: the granted/used bookkeeping and the JNI handle should live in one shared type used by both unified_pool.rs and fair_pool.rs, so the fair pool gets the same benefit and position 17 of #5212 (mutex held across the JNI round-trip) can be fixed in the same place.

B. Make the retained slack reclaimable

NativeMemoryConsumer.spill() returns 0 unconditionally (CometTaskMemoryManager.java:110-113), so Spark can never reclaim anything from Comet. Holding slack makes that worse.

Fix: have spill(size, trigger) call back into native — a new @native def releasePoolSlack(taskAttemptId: Long): Long resolved through the existing TASK_SHARED_MEMORY_POOLS registry (memory_pools/task_shared.rs:25) — and return the number of bytes of free slack dropped. No data is spilled, so it is cheap and always safe.

This is what makes hysteresis defensible under memory pressure rather than just shifting pain onto other tasks in the executor. It also partly addresses position 5 of #5212, since Comet's consumer would stop being invisible to Spark's spill-victim ordering.

I would not ship A enabled by default without B.

C. Cheaper per-call path

Add a with_env variant that skips with_local_frame for calls that create no local references. Two fewer JNI transitions per acquire and release, no behavioural change. Worth doing whether or not A lands.

D. Pre-reserve the whole per-task budget

Chunk size equal to memoryLimitPerTask, then run a native GreedyMemoryPool over it and touch Spark only on exhaustion. Nearly eliminates the JNI traffic, but hoards memory from tasks that need little and defeats Spark's dynamic sharing. Reasonable as an opt-in config, not as a default.

E. Reduce call frequency at the source

Round Comet's own reservations up so growth is stepwise rather than per batch; the native shuffle writer is the main offender. Narrower than A, complementary, and needs no new config.

Suggested sequencing

1. Instrumentation. Count acquire/release calls and bytes, and measure cumulative time spent in the JNI round-trip, then report per task. Without this there is no way to show a win — or to tell whether the win is worth the added complexity. This should be its own PR.
2. C — small, unconditional, independent of everything else.
3. A and B together, behind a config, with numbers from position 1 above.
4. E if the shuffle writer still shows up in the counters afterwards.

D only if a workload turns up that genuinely wants it.

Related

• [EPIC] Memory pool and accounting audit sweep #5212 — parent EPIC; this is positions 17 and 18, and B overlaps position 5.
• Investigate adopting DataFusion's allocator-level memory accounting to replace manual memory tuning #4576 — adopting DataFusion's allocator-level accounting would change what triggers these calls, but not what each one costs.

[andygrove]

I ran a single case from CometShuffleBenchmark with the counters from #5384 enabled, and the traffic turns out to come from exactly one consumer, with a much more specific shape than I assumed when I filed this.

Setup

SQL Single INT Shuffle(201 Partition) over 10M rows, native shuffle, release build, spark.memory.offHeap.enabled=true with 4 GB off-heap, default fair_unified pool, local[5]. Stats collected by raising only comet::execution::memory_pools to debug via a custom log4rs.yaml.

Measured traffic

110 tasks (11 iterations x 10 tasks), all identical:

Task 4334 CometFairMemoryPool memory pool stats: acquire(calls=256, requested=25165824 bytes,
  granted=25165824 bytes, short=0), release(calls=1, bytes=25165824), backend_time=432.836µs

Total calls	28,270 (28,160 acquire, 110 release)
Total time inside the round-trip	36.0 ms
Avg per call	1.27 µs (p50 1.23, p90 1.39, max 1.86)
Per task	0.32 ms (p50), 0.51 ms (max)
Short grants	0

Per-task wall time is roughly 182 ms x 5 slots / 10 tasks = 91 ms, so the round-trips are about 0.35% of a shuffle-write task, uncontended. Two runs of the same configuration produced identical call counts and the same 1.27 µs, so the measurement is stable.

The consumer

With spark.comet.debug.memory enabled (note it has to be set through SQLConf so it reaches native via the serialized Comet confs -- a system property does not get there), there is exactly one registrant:

28160 [ShuffleRepartitioner[0]].try_grow
  110 [ShuffleRepartitioner[0]].shrink
  110 [ShuffleRepartitioner[0]].register / unregister

Nothing else touches the pool. The Comet-with-Spark-shuffle and Comet-JVM-shuffle cases in the same benchmark make zero memory pool calls.

Sizes per task:

Size	Calls/task	What it is
32,768	247	96% of all calls -- one 8192-row Int32 Arrow data buffer (8192 x 4)
98,304	2	batches pinning more buffers
163,840 ... 8,421,376	7, one each	a doubling series: 160 KB, 288 KB, 544 KB, 1,056 KB, 2,080 KB, 4,128 KB, 8,224 KB
25,165,824 (shrink)	1	the entire 24 MB released in one call

The site is native/shuffle/src/partitioners/multi_partition.rs:463: buffer_partitioned_batch_may_spill performs exactly one try_grow(mem_growth) per input batch, where mem_growth is count_new_buffers(...) -- the 32 KB steps -- plus the growth of the per-partition index vectors, which is the doubling series as the 201 partition vectors reallocate. The single shrink is the reservation.free() at :528.

What this changes about the plan above

The release side is already effectively batched (one call per task), so hysteresis buys nothing on this workload; the entire cost is on the acquire side, and it is one consumer growing per batch rather than many consumers churning.

That makes option E the sharper tool than option A here: if the repartitioner rounded its reservation up to a step of a megabyte or so and tracked exact usage locally, 256 calls per task would become a handful. It needs no new config, no retained slack, and therefore does not pull in option B. Option A would get the same result for this workload but pays with slack that Spark cannot reclaim.

So I would reorder the sequencing to: instrumentation (#5384) -> C -> E -> A and B only if other workloads show many small consumers rather than one batch-driven one.

Caveat on the ceiling: 0.35% is a shuffle-write task at 5 concurrent tasks with no lock contention. I have not measured aggregate/sort/join-heavy plans, which have several consumers per task and may look quite different.