feat: prototype allocator-level OOM protection (OomGuard breaker + cooperative real-usage accounting)

[andygrove]

Which issue does this PR close?

Relates to #4576. This is an exploratory prototype covering both halves of that issue: the RSS circuit breaker ("OomGuard") and cooperative online accounting that feeds the real allocator balance into DataFusion's MemoryPool. It is not a complete, production-default implementation, so it does not close the issue.

Rationale for this change

Comet's memory accounting relies on voluntary MemoryPool reservations, which miss allocations made by Arrow buffers, join scratch space, and expression kernels. Two problems follow:

1. When real native memory exceeds the container limit, the OS/YARN/Kubernetes kills the entire executor JVM, losing every task and all cached data on it.
2. Because the pool undercounts real usage, users compensate by manually lowering spark.comet.exec.memoryPool.fraction, a crude, workload-specific knob.

This prototype attacks both by tracking the real bytes the global allocator hands out and using that signal in two layers:

• A cooperative gate that rejects pool growth (triggering a DataFusion spill and retry) once real usage plus the request would exceed the off-heap budget. This lets the pool act on real usage rather than tracked reservations alone, so fraction no longer needs tuning.
• A last-resort, executor-global circuit breaker that fails a single task with a retriable ResourcesExhausted error instead of letting the executor get OOM-killed.

The byte-tracking allocator adapts the AccountingAllocator from apache/datafusion#22626 rather than depending on it, since that code lives in DataFusion's test-only sqllogictest crate.

What changes are included in this PR?

Gated behind a new oom-guard cargo feature; the default build is unchanged with zero added per-allocation overhead.

Allocator accounting and circuit breaker:

• native/core/src/execution/memory_pools/oom_guard.rs (new): AccountingAllocator<A> wrapping the inner global allocator; a single process-wide balance with per-thread drift settled at a 64 KiB threshold; arm/disarm/stamp_current_thread/current_balance; a typed OomGuardPanic, raised via panic_any on an armed, stamped thread that crosses the limit, with reentrancy protection so the panic's own boxing allocation does not recurse.
• native/core/src/lib.rs: under the oom-guard feature, installs the wrapper as #[global_allocator] over jemalloc / mimalloc / system; mutually exclusive cfgs leave the default build untouched.

Cooperative real-usage gate:

• native/core/src/execution/memory_pools/real_usage_pool.rs (new): RealUsagePool, a MemoryPool decorator that checks the real allocator balance against a process-global ceiling before delegating growth to the inner pool, returning ResourcesExhausted (which DataFusion catches to spill and retry) when real usage plus the request would exceed the ceiling. It composes around any pool type and adds one relaxed atomic read per try_grow.

Wiring and config:

• native/core/src/execution/jni_api.rs: stamps tokio worker threads (on_thread_start) and the JNI caller thread; arms the guard from config in createPlan; wraps the memory pool in RealUsagePool (ceiling = the off-heap budget) when the guard is enabled; maps OomGuardPanic to DataFusionError::ResourcesExhausted at both execution boundaries (the spawned/channel path, both producer and consumer, and the busy-poll block_on path).
• spark/src/main/scala/org/apache/comet/CometConf.scala: registers spark.comet.exec.memoryGuard.enabled (default false) and spark.comet.exec.memoryGuard.size (optional; defaults to the executor off-heap size). Both layers ride this single switch. They default to the same threshold but order correctly: the cooperative gate trips on projected usage (balance + additional) so it spills first, while the breaker trips on actual usage (balance) as the backstop.

Known limitations / out of scope for this prototype (candidates for follow-ups):

• Executor-global granularity only; no per-task attribution or fairness.
• Only tokio workers and the JNI caller thread are stamped, so allocations on spawn_blocking/IO/other pools are tracked but cannot themselves trip the breaker.
• Layout-byte accounting only; no periodic resync to real jemalloc/mimalloc resident stats.
• The cooperative gate reacts to real usage but does not yet auto-size the pool budget or deprecate spark.comet.exec.memoryPool.fraction.
• JVM end-to-end spill test and benchmark validation are deferred until the feature moves toward a default build.

How are these changes tested?

• Rust unit tests in oom_guard.rs cover the decision/settle helpers, and that the breaker trips only on an armed, stamped thread (disarmed never trips, unstamped never trips).
• Rust unit tests in real_usage_pool.rs cover the cooperative gate: under the ceiling it delegates, over the ceiling it rejects without reserving the inner pool, an unset ceiling never gates, and shrink/reserved/memory_limit delegate to the inner pool.
• End-to-end Rust tests drive real heap allocations through the installed AccountingAllocator: one asserts an OomGuardPanic is raised and caught, and another asserts the cooperative gate rejects an over-budget grow.
• Verified the build and clippy -D warnings across the default, oom-guard, and jemalloc,oom-guard feature combinations.

[avantgardnerio]

Just fixed this bug in DF. You need to panic before the realloc, otherwise the caller still has the old pointer and tries to free it on unwind and segfaults.

[andygrove]

Thanks for the pointer. The current OomGuard sequence panics through std::alloc::handle_alloc_error before returning the resized pointer, so the caller does not observe the new allocation and its unwind uses the original pointer. Let me pin that with a comment near the panic call so a future refactor cannot silently swap the order and re-introduce the double-free.

[cetra3]

We've been using https://github.com/cetra3/thresher ourselves to handle something similar with datafusion. Namely, we log a trace when the threshold is reached with the heap dump of jemalloc, based upon the example

It's helped us so far nail down some gnarly memory allocation related bugs with our compaction process. I'm hoping to expand this to queries soon

[comphead]

how additional would be used/calculated? is it sort of extra buffer?

[andygrove]

additional is the extra headroom the pool grants above real jemalloc usage: on a try_grow call, the guard admits the request if jemalloc_stats.allocated + requested <= configured_pool_limit + additional. It exists so short-lived spikes (e.g. an operator's staging buffer inside a single try_grow) do not trip the breaker when the sustained usage would still fit. Set to 0 when you want a strict cap. I will add a doc comment on the field explaining it.

[comphead]

I'm going to test this approach soon

one thing, if I understand correctly this approach would help with overall accounting difference between DF memory pools and actual resident memory usage. However it wouldn't get us closer to identify threshold for spilling.

Currently it's hard(if even possible) to do fine grained analysis and identify which object/layer allocated resources through this approach. So spilling IMO still relies on DF memory pools unless @avantgardnerio as usual have an elegant memory solution :)

[avantgardnerio]

I think there was talk of tracking which operator was using the memory in the pool? I don't know if anything ever came of it, so I think the most granularity you can have is a MemoryPool per query / context.

AFAIK, all you need to know about when to spill is when allocations would otherwise fail? Which the OomGuard could tell you.

[comphead]

I have a good repro but it would depend on #5314

[andygrove]

I'm closing this rather than continuing to iterate on it.

The honest reason is that it grew into three things at once — allocator-level
byte tracking, a cooperative gate on the memory pool, and a panic-based circuit
breaker — and only the first of those is something I'm confident enough in to
defend. Reviewing it properly against the merged state of main turned up a few
things that undercut the design as described:

The two enforcement layers don't actually layer. The argument in the description
is that the cooperative gate fires before the breaker because it tests projected
usage while the breaker tests actual usage. But the gate only runs inside
try_grow, and the whole premise of the feature is that memory grows through
allocations that never call try_grow. In exactly the case this exists for, the
breaker fires with no spill attempted.

The stamped-thread set is wider than I thought. Tokio runs on_thread_start on
every thread the runtime spawns, including the blocking pool, not just the
workers — so a panic raised inside a spawn_blocking task becomes a JoinError
that reaches neither of the catch sites in executePlan.

LOCAL_DRIFT leaks on thread exit. It's a plain Cell<isize> with no TLS
destructor, so up to 64 KiB of un-flushed drift is discarded every time a thread
dies, and blocking threads churn on a 10s idle timeout. On a long-lived executor
that's a slowly accumulating bias.

And more fundamentally: the tracked balance is layout bytes, not RSS. For a
guard whose job is preventing a container OOM kill, the gap between those two —
fragmentation, jemalloc's retained pages, mmap, anything a C dependency
allocates — is unbounded and always in the dangerous direction.

None of that is wasted. @avantgardnerio's catch that realloc must panic before
delegating, and the finding that concurrent guard panics abort with exit 133
unless serialized, are the kind of thing you only learn by running it, and
they'll carry over.

What I'd like to do instead, in order:

I've opened #5933, a draft documenting the memory model as it actually is today
— the allocator inventory, the FFI ownership asymmetry, why declared
reservations diverge from RSS, and what the container counts. Most of what I
wrote while reviewing this PR turned out to be existing behavior, and it's
useful whether or not any guard ever lands.

Then I'd like to extract just the AccountingAllocator as pure observability:
no gate, no panics, feature-gated, with the balance exposed next to pool
reservations in tracing. That makes the accounting gap measurable, which is what
every memory bug report needs and nobody currently has, and it's small enough to
review properly.

Whether to enforce on that number at all is a decision I'd rather make after
seeing how well it tracks RSS on real workloads, rather than before.

@comphead — you mentioned you had a repro pending on #5314. That would still be
very useful against the observability-only version; it would tell us directly
whether the balance is a trustworthy signal.

The branch stays available for anyone who wants to pull pieces from it.