WIP: feat(executor): allocator-backed OOM guard behind an opt-in oom-guard feature

[andygrove]

Status: WIP / experimental. This is not ready for review and is just for learning at this stage. This PR is likely to undergo massive changes while I experiment more and keep learning. I will likely close it at some point and open a clean PR instead.

Stacked on #2032. The first commit here is that PR; review only the second. GitHub cannot base a PR on a fork branch, hence the overlap.

Which issue does this PR close?

Relates to #2031. Prior art: apache/datafusion-comet#4582, which explores the same idea for Comet. This diverges from that prototype in one important way — see "Deviations from the Comet prototype".

Rationale for this change

Executor memory accounting relies on voluntary MemoryPool reservations, which miss Arrow buffers, join scratch space, and expression-kernel temporaries. Real native memory therefore runs above what the pool believes is reserved, and the executor process can be OOM-killed.

In Ballista that is far worse than one failed task: a dead executor takes all of its shuffle output with it, so every downstream stage that would have read from it raises FetchPartitionError, cascading into stage rollbacks and map-stage re-runs — potentially across other concurrent jobs.

What changes are included in this PR?

A new oom-guard cargo feature (off by default) that installs a global allocator tracking the bytes actually handed out, and uses that signal in two layers:

• RealUsagePool — a MemoryPool decorator that rejects try_grow with ResourcesExhausted once real live usage plus the request would exceed the budget, so DataFusion spills rather than growing into an OOM. This is the layer that does the useful work.
• MemoryGuardExec — a transparent pass-through ExecutionPlan inserted below each stage's shuffle writer. It checks the budget between batches and fails that one task as a last resort. It debounces (3 consecutive over-budget checks and 100 ms continuously over budget), so a transient spike that the pool gate is already resolving does not shed every task on the executor.

The ceiling is --memory-pool-size. If that flag is unset the guard tracks but never enforces.

Two design points worth reviewer attention

The allocator only tracks; it never enforces. Unwinding out of a GlobalAlloc is documented undefined behaviour (core::alloc::GlobalAlloc: "It's undefined behavior if global allocators unwind... a panic from any of these functions may lead to memory unsafety"). So enforcement happens at safe points — memory-pool growth, and plan poll boundaries — instead. Live bytes are counted in cache-line-padded atomic shards, which keeps the count exact across thread churn without a Drop-carrying thread-local (whose lazy storage would allocate from inside the allocator).

The guard enforces on tracked live bytes, not RSS. RSS is what the OOM killer reads, but allocators do not promptly return freed pages, so a spill that genuinely released memory would leave RSS high and the guard latched on, failing every subsequent task — turning a safety feature into an outage. Failing tasks cannot reclaim memory that is already free. RSS is sampled for observability only, and a warning is logged when it climbs above the limit while tracked usage stays below it, which is precisely the signal that the allocator is holding memory the guard cannot reclaim.

Deviations from the Comet prototype

The Comet prototype raises a panic from inside GlobalAlloc::alloc and catches it with catch_unwind. I did not port that, for two reasons, and both apply to Comet as well:

1. It is UB, per the contract quoted above.
2. It cannot be made reliable anyway: DataFusion spawns internal tokio sub-tasks (RepartitionExec spawns one per input partition, hash-join build side, SortPreservingMerge), and a panic raised inside one of those is swallowed by tokio as a JoinError and never reaches the catch site. Since the breaker disarms itself on trip, it would be permanently disarmed after its first real trip.

Known limitations (please read before reviewing)

• The allocator overhead is UNMEASURED. The guard adds a relaxed fetch_add to every allocation. A TPC-H benchmark (feature on vs off, under both AQE settings) is the next step and is the gate on whether this could ever be default-on. Until then this is a well-tested hypothesis.
• Enforcement is batch-granular. An operator that blows the ceiling within one poll_next is not caught. The pool gate is the first line of defence; the breaker only shrinks the blast radius. This feature reduces the probability of an executor dying; it does not prevent it.
• --memory-pool-size changes meaning in an oom-guard build: the ceiling is compared against all live process bytes (gRPC, tokio, object-store caches), not just query memory, so the pool gates somewhat earlier than the same value implies in a default build. Documented in the CLI help and the tuning guide. A separate --memory-guard-limit may be the better answer; I would take guidance here.
• The cooperative layer is a no-op for the sort-shuffle writer. sort_shuffle/writer.rs does let _ = reservation.try_grow(...), discarding pool rejection in favour of its own absolute byte counter, so RealUsagePool cannot make it spill. The breaker still covers it. Wiring its spill decision to the pool result felt out of scope here.
• No per-task attribution. The balance is process-global; the breaker fails whoever polls next once the debounce elapses, not the task responsible.

Are there any user-facing changes?

Only for those who opt in at build time (cargo build --release --features oom-guard). The default build is unchanged: no allocator is installed and there is no per-allocation overhead. Documented in docs/source/user-guide/tuning-guide.md and cargo-install.md.