andygrove opened a new pull request, #2033: URL: https://github.com/apache/datafusion-ballista/pull/2033
> **Status: WIP / experimental.** Opened as a draft for early design feedback, not for merge. The feature is **off by default** and I have **not yet measured the allocator overhead** — see "Open questions" below. Please read the limitations section before the code. > > **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`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
