diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index 701617cde39..21483ac54e8 100644 --- a/docs/source/contributor-guide/index.md +++ b/docs/source/contributor-guide/index.md @@ -48,6 +48,7 @@ Comet Plugin Overview Arrow FFI JVM Shuffle Native Shuffle +Memory Management ANSI Error Propagation S3 Credential Provider Design ``` diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md new file mode 100644 index 00000000000..04df6bac003 --- /dev/null +++ b/docs/source/contributor-guide/memory_management.md @@ -0,0 +1,477 @@ + + +# Memory Management + +This page describes how memory is budgeted, accounted, and enforced across the JVM/native +boundary. It is aimed at contributors working on memory pools, operators that reserve memory, or +anyone debugging an out-of-memory report. For user-facing tuning advice, see the +[Tuning Guide](../user-guide/latest/tuning.md). + +## Overview + +A Comet executor has to satisfy three separate memory budgets at once, and they are enforced by +three different parties: + +| Budget | Enforced by | What happens when it is exceeded | +| -------------------- | ----------------- | -------------------------------------------------------------------- | +| JVM heap | The JVM | `OutOfMemoryError` in a task; the executor usually survives | +| Spark's memory pools | Spark bookkeeping | A consumer is asked to spill, or a `SparkOutOfMemoryError` is thrown | +| Container RSS | The OS / cgroup | `SIGKILL` of the whole executor process (exit 137, `OOMKilled`) | + +The first two are _accounting_: a running total of bytes that consumers have voluntarily declared. +The third is _physical_: the kernel measures resident pages and does not care what any accounting +layer believes. + +Comet's difficulty is that its allocations are made by Rust code, so they are invisible to the JVM +heap and to Spark's own off-heap accounting, yet they land squarely in container RSS. Comet +therefore maintains its own budget that is meant to shadow the physical one, and the accuracy of +that shadow is the central problem this page is about. + +## Who allocates what + +Enabling Comet does not add one new memory consumer, it adds several, and they are not all +accounted by the same party. This inventory is worth internalizing before reading the rest of the +page: + +| Allocator | Lives in | Bounded by | Visible to Spark? | +| --------------------------------------- | ----------- | ------------------------------------------------------------------- | ----------------- | +| Spark execution + storage (on-heap) | JVM heap | `spark.executor.memory` and the unified memory manager | Yes | +| Spark Tungsten (off-heap) | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | +| Comet native (Rust global allocator) | Native heap | `memory_limit` (see below), enforced only via the memory pool | No | +| Comet JVM Arrow (`CometArrowAllocator`) | Off-heap | **Nothing** — a `RootAllocator(Long.MaxValue)` | No | +| Comet JVM shuffle pages (off-heap mode) | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | +| Comet JVM shuffle pages (on-heap mode) | Off-heap | `spark.comet.shuffle.jvm.memoryFactor * spark.comet.memoryOverhead` | No | + +Three observations follow. + +**Comet's JVM-side Arrow allocator is unbounded and accounted by nobody.** `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 +(`CometNativeArrowSource`), broadcast coalescing, and `CometSparkToColumnarExec`. These are real +off-heap bytes in container RSS that neither Spark's `TaskMemoryManager` nor Comet's native pool nor +the `oom-guard` allocator sees. In practice the volume is modest — a batch at a time per stream — +but there is no ceiling and no backpressure. + +**The JVM shuffle allocator switches accounting model with the memory mode.** +`CometShuffleMemoryAllocator.getInstance` returns `CometUnifiedShuffleMemoryAllocator` when Tungsten +is off-heap, which is a proper Spark `MemoryConsumer` drawing from `spark.memory.offHeap.size`. In +on-heap mode it returns `CometBoundedShuffleMemoryAllocator`, which calls `UnsafeMemoryAllocator` +directly and bounds itself with its own counter. Only the first is arbitrated against Spark's other +consumers. + +**On-heap mode double-counts `spark.comet.memoryOverhead`.** The native pool is sized at +`memory_limit = spark.comet.memoryOverhead`, and the JVM shuffle allocator is _separately_ sized at +`spark.comet.shuffle.jvm.memoryFactor * spark.comet.memoryOverhead`, with the factor defaulting to +`1.0`. They are distinct allocations from the same number, so on-heap Comet can occupy up to roughly +twice `spark.comet.memoryOverhead` in off-heap RSS, before counting `CometArrowAllocator`. Off-heap +mode does not have this problem, which is one more reason it is the recommended configuration. + +## Where Comet's budget comes from + +`CometExecIterator.getMemoryConfig` computes the budget once per executor and passes it across JNI +to `Java_org_apache_comet_Native_createPlan` as `memory_limit` and `memory_limit_per_task`. There +are two paths. + +### Off-heap mode (`spark.memory.offHeap.enabled=true`) + +This is the recommended configuration. Comet shares Spark's off-heap pool rather than asking for a +separate allocation: + +```text +memory_limit = spark.memory.offHeap.size * spark.comet.exec.memoryPool.fraction +memory_limit_per_task = memory_limit * spark.task.cpus / executor_cores +``` + +`spark.comet.exec.memoryPool.fraction` defaults to `1.0`. Lowering it is the current workaround for +Comet's under-accounting (see [The accounting gap](#the-accounting-gap)) — it holds back a slice of +the off-heap pool that Comet is not allowed to reserve, on the assumption that Comet's real usage +overshoots its reservations by roughly that slice. + +### On-heap mode + +Comet asks for a dedicated overhead allocation outside the heap: + +```text +memory_limit = spark.comet.memoryOverhead (default 1024 MiB) +memory_limit_per_task = memory_limit * spark.task.cpus / executor_cores +``` + +On-heap mode is a testing configuration; the pool types it exposes are in the `CATEGORY_TESTING` +group. + +### Resolving the pool type + +`parse_memory_pool_config` (`native/core/src/execution/memory_pools/config.rs`) turns the mode, the +pool-type string, and the two limits into a `MemoryPoolConfig`. Note which limit each pool type is +sized from — this is a common source of confusion: + +| Pool type | Mode | Sized from | Notes | +| ------------------------------------- | -------- | ----------------------- | ------------------------------------------- | +| `fair_unified` (default) | off-heap | `memory_limit` | Delegates to Spark's `TaskMemoryManager` | +| `greedy_unified` | off-heap | n/a (pool size `0`) | Spark owns the limit entirely | +| `greedy_task_shared` | on-heap | `memory_limit_per_task` | Default on-heap pool | +| `fair_spill_task_shared` | on-heap | `memory_limit_per_task` | | +| `greedy` / `fair_spill` | on-heap | `memory_limit_per_task` | Per-plan, not shared across plans in a task | +| `greedy_global` / `fair_spill_global` | on-heap | `memory_limit` | One pool for the whole executor | +| `unbounded` | on-heap | n/a | No limit; testing only | + +## The pool stack + +`create_memory_pool` builds a base pool and `createPlan` then wraps it in decorators. Reading from +the inside out, a Comet plan in the default off-heap configuration sees: + +```text +[LoggingMemoryPool] <- only when spark.comet.debug.memory=true + [RealUsagePool] <- only when the oom-guard build + memoryGuard.enabled + [TaskSharedMemoryPool] <- RAII handle for the per-task registry + [TrackConsumersPool] <- DataFusion; names the top 10 consumers in error messages + [CometFairMemoryPool] <- delegates acquire/release to Spark over JNI +``` + +Each decorator forwards every `MemoryPool` method to its inner pool, so `reserved()` at any level +reports the base pool's number. + +### The unified pools + +`CometUnifiedMemoryPool` and `CometFairMemoryPool` (`unified_pool.rs`, `fair_pool.rs`) are the +bridge to Spark. Their `try_grow` calls `CometTaskMemoryManager.acquireMemory` over JNI, which goes +through Spark's ordinary `TaskMemoryManager`. That means: + +- Comet competes with Spark's own off-heap consumers (Tungsten sorters, `BytesToBytesMap`, and so + on) for the same `spark.memory.offHeap.size`, and Spark's unified memory manager arbitrates. +- Spark can force _Spark's_ consumers to spill to satisfy Comet's request, and vice versa. +- A partial grant (`acquired < additional`) is released immediately and reported as + `ResourcesExhausted`, which is the signal DataFusion uses to spill. + +`CometFairMemoryPool` additionally caps each registered consumer at `pool_size / num_consumers` +before it even asks Spark, which is why it spills earlier than `greedy_unified` but keeps one +operator from starving the others. + +### Task-shared pools and their lifetime + +A single Spark task can run more than one native plan concurrently — a shuffle runs the pre-shuffle +operators and the shuffle writer as separate native execution contexts. If each got its own pool, +the per-task limit would be enforced once per plan rather than once per task. + +`acquire_task_shared_pool` (`task_shared.rs`) keeps a process-wide +`HashMap>`. Plans in the same task upgrade the existing +`Weak` and share one pool; the returned `Arc` is the only lifetime handle, so the registry entry +disappears when the last plan (and its last reservation) drops. There is no explicit release call to +forget, and a `createPlan` that fails partway through cleans up on unwind. + +`TaskSharedMemoryPool::drop` has to handle one race: an `acquire` can observe an expired `Weak` and +insert a replacement before the dying pool reaches the registry lock. The drop therefore compares +pointers and only removes an entry that is still its own. + +## How DataFusion consumes the pool + +Native operators reserve through DataFusion's `MemoryConsumer` / `MemoryReservation` API: + +- `try_grow(n)` may fail. Spillable operators (`ExternalSorter`, the grouped hash aggregate, + sort-merge join) respond to a `ResourcesExhausted` error by spilling to disk and retrying. This is + the only mechanism that turns memory pressure into progress rather than failure. +- `grow(n)` is infallible and panics if the pool refuses. It is used where the caller cannot spill. +- `shrink(n)` returns bytes to the pool. + +An operator that never calls `try_grow` is invisible to the pool no matter how much memory it uses. + +## Crossing the FFI boundary + +Batches move between the JVM and native over the Arrow C Data and C Stream interfaces, which are +zero-copy. Nothing is copied, so the _allocator_ that produced a batch and the _runtime_ that +decides when it dies can be on opposite sides of the boundary. See [Arrow FFI](ffi.md) for the +mechanics; what matters here is who is charged and who controls the lifetime. + +**JVM → native (`ScanExec`).** The JVM allocates the Arrow buffers from a child of +`CometArrowAllocator` and exports the whole per-partition iterator once as an `ArrowArrayStream`. +Native takes ownership by reference through `AlignedArrowStreamReader`. The bytes were allocated by +Java Arrow, so the Rust global allocator never sees them: they are absent from `BALANCE`, absent +from the memory pool, and absent from Spark's `TaskMemoryManager` — but present in container RSS, +and pinned for as long as the native side holds the imported batch. A native operator that buffers +many input batches is therefore pinning JVM-allocated off-heap memory that none of Comet's +accounting can observe. + +**Native → JVM (`CometExecIterator`).** DataFusion produces the batch in Rust, so those bytes _are_ +counted in `BALANCE` and may also be reserved in the pool. The batch is exported as an +`ArrowArray`/`ArrowSchema` pair, the JVM wraps the pointers in `ArrowBuf`s, and the memory is only +freed when the JVM calls `close()` and the release callback runs. The lifetime of native, +pool-charged memory is thus controlled by JVM code. A slow or backed-up JVM consumer keeps +`BALANCE` elevated for memory the native side has logically finished with, which means the guard can +trip on a backlog rather than on genuine native demand. + +The asymmetry is the point: **the direction of data flow determines which accounting layer, if any, +charges for a batch.** Neither direction charges both, and the JVM → native direction charges +nothing at all. + +## The accounting gap + +The pool tracks _declared reservations_. Container RSS counts _pages the process touched_. The two +diverge for several structural reasons: + +- **Undeclared allocations.** Arrow array builders, expression kernels producing intermediate + arrays, decompression buffers, Parquet metadata structures, `object_store` request buffers, and + tokio's own machinery all allocate without reserving. Only operators that were explicitly written + to reserve show up in the pool. +- **Rounding and padding.** Arrow buffers are padded to 64-byte boundaries and builders grow by + doubling, so a reservation of exactly `n` bytes routinely corresponds to more than `n` bytes of + heap. +- **Allocator behavior.** `malloc`-level fragmentation, size-class rounding, and jemalloc's + retained/dirty page cache all add resident bytes that no layer above the allocator can see. + Freeing memory does not necessarily return pages to the OS. +- **Non-Rust allocations.** Memory allocated by C dependencies through libc `malloc`, and anything + `mmap`ed, never passes through Rust's `GlobalAlloc`. +- **FFI-imported buffers.** Batches arriving from the JVM were allocated by Java Arrow, so they + belong to no Comet budget at all while native code holds them (see + [Crossing the FFI boundary](#crossing-the-ffi-boundary)). + +The practical consequence is that `reserved()` is a lower bound on Comet's real footprint, and the +gap is workload-dependent. `spark.comet.exec.memoryPool.fraction` exists purely so operators can +hand-tune a haircut that covers the gap for their workload. + +To measure the gap on a real query, enable tracing with the `jemalloc` feature and compare +`jemalloc_allocated` against the summed `thread_NNN_comet_memory_reserved` values — see +[Tracing](tracing.md#analyzing-memory-usage). + +## What the container sees + +On Kubernetes, Spark sizes the executor pod from `ResourceProfile`: + +```text +pod memory request = pod memory limit + = spark.executor.memory + + spark.executor.memoryOverhead (default max(0.1 * executor.memory, 384 MiB)) + + spark.memory.offHeap.size (when off-heap is enabled) + + pyspark memory (Python applications only) +``` + +Both the request and the limit are set to this same value, so the pod's cgroup `memory.max` is a +hard ceiling on the sum of everything in the container. That cgroup counts, among other things: + +- the JVM heap (`spark.executor.memory`), +- JVM non-heap: metaspace, code cache, thread stacks, GC structures, Netty direct buffers, +- Spark's own off-heap allocations, +- **all of Comet's native allocations**, +- Comet's JVM-side Arrow buffers (`CometArrowAllocator`) and, in on-heap mode, its JVM shuffle pages, +- page cache charged to the cgroup by the container's file I/O, including spill files. + +Only the first and a portion of the third are visible to Spark's accounting. When the total crosses +`memory.max`, the kernel OOM killer kills the process. The failure mode is significantly worse than +a task-level OOM: every task running on that executor dies, every cached block it held is lost and +must be recomputed, and the shuffle files it produced become unavailable to downstream fetches. +Spark's driver sees only `ExecutorLostFailure` with exit code 137. + +Two facts follow that are easy to get wrong: + +1. **`spark.memory.offHeap.size` is part of the pod limit, not extra headroom on top of it.** + Raising Comet's off-heap budget on Kubernetes raises the pod's memory request by the same + amount, so the scheduler will place fewer executors per node rather than silently giving Comet + more room. +2. **`spark.executor.memoryOverhead` is the only slack in the container**, and the JVM's own + non-heap usage already consumes a large part of it. Comet's overshoot beyond its declared + reservations eats into the same allowance. + +YARN behaves analogously — the container size is the same sum, and the NodeManager kills containers +that exceed it — but the kill is done by the NodeManager's monitor rather than the kernel, so it is +somewhat less abrupt. + +## The OOM guard (experimental) + +The layers described so far cannot prevent an OOM kill, because they bound reservations rather than +allocations. The `oom-guard` feature adds allocator-level tracking and two enforcement points built +on it. It is **a prototype and is not in the default cargo feature set**, so released binaries do +not contain it unless it is explicitly enabled at build time — see +[issue #4576](https://github.com/apache/datafusion-comet/issues/4576) and +[PR #4582](https://github.com/apache/datafusion-comet/pull/4582). + +### Building and enabling it + +```shell +cd native && cargo build --release --features oom-guard +``` + +or, through the project Makefile: + +```shell +COMET_FEATURES=oom-guard make release +``` + +```properties +spark.comet.exec.memoryGuard.enabled=true +# optional; defaults to the value of memory_limit above +spark.comet.exec.memoryGuard.size=6g +``` + +Both configurations are inert in a build without the feature. When the feature is absent there is no +allocator wrapper at all, so the default build carries zero per-allocation overhead. + +### Layer 1: the accounting allocator + +`AccountingAllocator` (`memory_pools/oom_guard.rs`) wraps whichever allocator the build selected +— jemalloc, mimalloc, or the system allocator — and is installed as `#[global_allocator]`. On every +`alloc` / `alloc_zeroed` / `dealloc` / `realloc` it adds the `Layout` size delta to a thread-local +`LOCAL_DRIFT`. When a thread's drift exceeds 64 KiB in either direction it is flushed into a single +process-wide `AtomicIsize` called `BALANCE`. Batching keeps the common path to a thread-local +add-and-compare, so only roughly one atomic RMW per 64 KiB of churn reaches the shared cacheline. + +`BALANCE` is therefore the count of **layout bytes currently handed out by the Rust global +allocator, process-wide**. It is not RSS, and it is not per-task. It does capture the undeclared +allocations that the memory pool misses, which is the whole point. + +### Layer 2: the cooperative gate (`RealUsagePool`) + +`RealUsagePool` is a `MemoryPool` decorator. Before delegating a `try_grow(additional)` to the inner +pool it checks the projected real usage: + +```text +if BALANCE + additional > ceiling: + reject with ResourcesExhausted +``` + +where `ceiling` is `memory_limit`. Rejecting _before_ delegating means the inner pool is never +speculatively reserved, so there is nothing to roll back. Because the error is `ResourcesExhausted`, +DataFusion's spilling operators treat it exactly like an ordinary pool rejection: they spill and +retry. This is the layer that is supposed to make Comet react to real usage rather than tracked +reservations, and in principle it removes the need to hand-tune +`spark.comet.exec.memoryPool.fraction`. + +Because `BALANCE` is process-wide but the pool is per-task, a naive gate would punish whichever task +happened to call `try_grow` first after the executor crossed the ceiling. The gate therefore applies +a fair-share test: once over the ceiling, a task is only rejected if its own tracked reservation +would exceed `ceiling / active_tasks`. `active_tasks` is the number of live task-shared pools; +`spark.executor.cores` is the fallback divisor for pool types that keep no task registry. Tasks +under their share are allowed through, and the breaker below is the backstop for the runaway case. + +The gate adds one relaxed atomic load per `try_grow`. + +### Layer 3: the circuit breaker + +The breaker is the last resort. `createPlan` calls `oom_guard::arm(limit)`, and query threads are +"stamped" as eligible to trip it. Stamping happens in two places: `build_runtime` passes +`stamp_current_thread` to tokio's `Builder::on_thread_start`, and `executePlan` stamps the JNI +caller thread directly. Note that tokio runs `on_thread_start` on **every** thread the runtime +spawns, both the multi-thread worker threads and the blocking pool, so the stamped set is wider +than just the workers. When a stamped thread flushes a positive drift that pushes `BALANCE` past +the limit, `AccountingAllocator` raises a typed `OomGuardPanic` via `panic_any` from inside the +allocation call. + +Several details exist to make that survivable: + +- **Only stamped threads panic.** Allocations on threads Comet did not create are still counted but + cannot themselves trip the breaker. +- **`realloc` panics before delegating.** If it panicked after `inner.realloc`, the old block may + already have been freed or moved while the caller still holds the old pointer, and the unwind + would free a dangling pointer. +- **Only one thread may fire per arm cycle.** The breaker CASes `ARMED` from `true` to `false`; + losers bail out before `panic_any`. Several threads dispatching a panic within the same few + milliseconds can abort the process with "failed to initiate panic" instead of unwinding. The next + `createPlan` re-arms. +- **Re-entrancy is handled.** `panic_any` boxes its payload, which allocates and re-enters the + allocator. `ARMED` is already `false` by then, and a thread-local `UNWINDING` flag adds a second + guard in case a concurrent `createPlan` re-arms mid-unwind. + +`executePlan` catches the panic at both execution boundaries (the spawned-task channel path, on both +the producer and the consumer side, and the busy-poll `block_on` path) and maps it to +`DataFusionError::ResourcesExhausted` via `oom_guard::map_panic_to_error`, which also clears the +thread's `UNWINDING` flag — necessary because the JNI caller thread is reused across tasks. + +### What the user sees + +The error reaches Spark as a `CometNativeException`, an ordinary `RuntimeException`. Spark retries +the task up to `spark.task.maxFailures` and then fails the stage. This is a meaningful improvement +over an OOM kill — the executor, its other tasks, and its cached blocks all survive — but it is a +_failure_ path, not a recovery path. A query whose working set genuinely does not fit will now fail +deterministically instead of taking the executor down. + +### Choosing `memoryGuard.size` + +The default sets the breaker's limit equal to the cooperative gate's ceiling (`memory_limit`). The +two layers are intended to order correctly even at the same value, because the gate trips on +projected usage (`BALANCE + additional`) while the breaker trips on actual usage (`BALANCE`), so the +gate should fire first. That ordering holds only when a `try_grow` happens between crossing the +ceiling and the next allocation. Usage that grows purely through undeclared allocations — exactly +the case the guard exists for — reaches the breaker with no cooperative spill attempted. + +Setting `memoryGuard.size` above `memory_limit` therefore gives the gate a real chance to spill +before the breaker fires. On Kubernetes a reasonable target is the slack between the pod limit and +everything else in the container: + +```text +memoryGuard.size ~ pod memory limit + - spark.executor.memory (JVM heap) + - JVM non-heap (metaspace, code cache, thread stacks, direct buffers) + - Spark's own off-heap usage + - a safety margin for allocator fragmentation and page cache +``` + +That is necessarily an estimate. Because `BALANCE` undercounts RSS (see below), the guard should be +given a budget comfortably below the true headroom. + +### Limitations + +These are known and mostly inherent to the prototype: + +- **Layout bytes, not RSS.** `BALANCE` counts what the program asked for, not resident pages. It + misses allocator fragmentation, jemalloc's retained pages, `mmap`ed regions, and allocations made + by C dependencies through libc `malloc`. There is no periodic resync against real jemalloc stats, + so the gap between `BALANCE` and RSS is unbounded and one-directional (RSS is always larger). +- **Per-thread drift is lost when a thread exits.** `LOCAL_DRIFT` is a plain `Cell` with no + TLS destructor, so up to 64 KiB of un-flushed drift is silently discarded each time a thread dies. + Tokio worker threads live for the process lifetime, but blocking-pool threads idle out and churn, + so on a long-lived executor this is a slowly accumulating bias in either direction. +- **Executor-global granularity.** The breaker fires on whichever stamped thread happens to allocate + when the process crosses the limit, which need not be the task responsible for the usage. The + fair-share test in the cooperative gate mitigates this for `try_grow`, but not for the breaker. +- **Not every stamped thread's panic reaches a catch site.** `executePlan` catches on the spawned + channel path and the busy-poll path, which covers the worker threads driving a plan. A panic + raised on a blocking-pool thread inside a `spawn_blocking` task is captured by tokio as a + `JoinError` instead, so it surfaces as a generic failure rather than the intended + `ResourcesExhausted`. +- **Panicking from inside the global allocator** unwinds through code that was mid-allocation. It is + memory-safe in the cases exercised so far, but a guard panic raised while another panic is already + unwinding is a double panic and aborts the process. +- **The FFI boundary is accounted asymmetrically.** Batches imported from the JVM are absent from + `BALANCE` even though native code pins them, so the guard under-reports on scan-heavy plans. + Batches exported to the JVM stay in `BALANCE` until the JVM closes them, so a backed-up consumer + can trip the guard on memory the native side is already done with. Neither is corrected for. +- **The budget is not auto-sized.** The gate reacts to real usage but does not yet adjust the pool + budget or deprecate `spark.comet.exec.memoryPool.fraction`. + +## Debugging memory issues + +| Tool | What it gives you | +| ------------------------------------------------ | -------------------------------------------------------------------------- | +| `spark.comet.debug.memory=true` | `LoggingMemoryPool` logs every register/grow/shrink with the consumer name | +| `spark.comet.explain.native.enabled=true` | Native plan with per-operator metrics, including spill counts | +| [Tracing](tracing.md#analyzing-memory-usage) | `jemalloc_allocated` vs summed pool reservations; the accounting gap | +| `TrackConsumersPool` | Names the top 10 consumers in `ResourcesExhausted` messages (always on) | +| [`thresher`](https://github.com/cetra3/thresher) | Third-party crate that dumps a jemalloc heap profile at a threshold | + +A checklist for triaging an executor OOM kill: + +1. Confirm it is an OOM kill and not a JVM `OutOfMemoryError` — exit code 137 / `OOMKilled` on the + pod, versus a heap dump and a stack trace. +2. Compare `jemalloc_allocated` against the summed pool reservations from a trace. A large excess + points at undeclared native allocations; a small excess points at the budget simply being too + small, or at the JVM side. +3. Check `spark.comet.batchSize` against the schema width. Peak memory scales with + `batch_size * columns`, and wide or deeply nested schemas amplify it. +4. Check whether the operators involved can spill at all. `ShuffledHashJoin` cannot, so + `spark.comet.exec.forceShuffledHashJoin=true` converts a spillable sort-merge join into one that + is not. diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index f0c7735a503..08749357f3f 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -114,6 +114,11 @@ jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"] # Default builds carry zero Delta surface. contrib-delta = ["dep:comet-contrib-delta"] +# Allocator-level OOM circuit breaker. When enabled, the global allocator is +# wrapped to track real allocated bytes and panic an over-budget query-worker +# thread (caught at the task boundary). Off by default; zero overhead when off. +oom-guard = [] + # exclude optional packages from cargo machete verifications [package.metadata.cargo-machete] ignored = ["hdfs-sys", "paste"] diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 692b0d1bccf..7e5640a36a2 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -115,6 +115,8 @@ use crate::execution::spark_config::{ COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, }; +#[cfg(feature = "oom-guard")] +use crate::execution::spark_config::{COMET_MEMORY_GUARD_ENABLED, COMET_MEMORY_GUARD_SIZE}; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use datafusion_comet_proto::spark_operator::operator::OpStruct; use log::{info, warn}; @@ -225,6 +227,8 @@ fn parse_usize_env_var(name: &str) -> Option { fn build_runtime(default_worker_threads: Option) -> Runtime { let mut builder = tokio::runtime::Builder::new_multi_thread(); + #[cfg(feature = "oom-guard")] + builder.on_thread_start(crate::execution::memory_pools::oom_guard::stamp_current_thread); if let Some(n) = parse_usize_env_var("COMET_WORKER_THREADS") { info!("Comet tokio runtime: using COMET_WORKER_THREADS={n}"); builder.worker_threads(n); @@ -484,6 +488,31 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( spark_config.get_u64(COMET_MAX_TEMP_DIRECTORY_SIZE, 100 * 1024 * 1024 * 1024); let logging_memory_pool = spark_config.get_bool(COMET_DEBUG_MEMORY); + #[cfg(feature = "oom-guard")] + { + if spark_config.get_bool(COMET_MEMORY_GUARD_ENABLED) { + // This is the hard, last-resort breaker: a panic on any over-budget + // allocation. It defaults to the same off-heap budget as the cooperative + // real-usage gate below, but the two layers still order correctly because + // the cooperative gate trips on *projected* usage (`balance + additional`) + // while this breaker trips on *actual* usage (`balance`), so cooperative + // spilling is attempted before the breaker fires. Set + // `spark.comet.exec.memoryGuard.size` explicitly above the off-heap budget + // to give the breaker additional headroom (e.g. up to the container RSS + // limit) for a wider spill-before-fail margin. + let default_limit = memory_limit.max(0) as u64; + let limit = spark_config.get_u64(COMET_MEMORY_GUARD_SIZE, default_limit); + if limit == 0 { + warn!( + "spark.comet.exec.memoryGuard.enabled is true but the effective limit \ + is 0 (memory_limit={memory_limit}); the guard will not trip. Set \ + spark.comet.exec.memoryGuard.size explicitly." + ); + } + crate::execution::memory_pools::oom_guard::arm(limit as usize); + } + } + with_trace("createPlan", tracing_enabled, || { // Init JVM classes JVMClasses::init(env); @@ -519,13 +548,38 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( let memory_pool = create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id); - // Register the shared base pool before wrapping it for per-plan debug logging. The - // guard removes the entry if any later plan setup step fails. + // Register the shared base pool before wrapping it for the real-usage gate or + // per-plan debug logging. The guard removes the entry if any later plan setup step + // fails. let rust_thread_id = get_thread_id(); let memory_pool_registration = tracing_enabled.then(|| { ThreadMemoryPoolRegistration::new(rust_thread_id, id, Arc::clone(&memory_pool)) }); + // Cooperative real-usage gate: when the memory guard is enabled, wrap the + // pool so growth is rejected (triggering a spill) once real allocator usage + // plus the request would exceed the process-wide off-heap budget. This is the + // first line of defense and fires before the hard OomGuard breaker armed above + // (projected vs. actual usage; see that comment), so over-budget work spills + // and retries rather than failing the task. + #[cfg(feature = "oom-guard")] + let memory_pool = if spark_config.get_bool(COMET_MEMORY_GUARD_ENABLED) { + let ceiling = memory_limit.max(0) as usize; + // Enable the fair-share guard for pools whose `reserved()` is per-task; + // `executor_cores` is the fallback divisor when no task count is known. + let fair_share = memory_pool_config + .pool_type + .has_per_task_budget() + .then_some(executor_cores); + Arc::new(crate::execution::memory_pools::RealUsagePool::new( + memory_pool, + ceiling, + fair_share, + )) as Arc + } else { + memory_pool + }; + let memory_pool = if logging_memory_pool { Arc::new(LoggingMemoryPool::new(task_attempt_id as u64, memory_pool)) } else { @@ -914,6 +968,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( schema_addrs: JLongArray, ) -> jlong { try_unwrap_or_throw(&e, |env| { + #[cfg(feature = "oom-guard")] + crate::execution::memory_pools::oom_guard::stamp_current_thread(); // Retrieve the query let exec_context = get_execution_context(exec_context); @@ -990,6 +1046,17 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( .await; if let Err(panic) = result { + #[cfg(feature = "oom-guard")] + if let Some(e) = + crate::execution::memory_pools::oom_guard::map_panic_to_error( + panic.as_ref(), + ) + { + // Runs on the tokio worker thread that panicked, so this clears + // that worker's UNWINDING flag (not the blocked JNI caller thread's). + let _ = tx.send(Err(e)).await; + return; + } let msg = match panic.downcast_ref::<&str>() { Some(s) => s.to_string(), None => match panic.downcast_ref::() { @@ -1014,76 +1081,120 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( pull_input_batches(exec_context)?; } - if let Some(rx) = &mut exec_context.batch_receiver { - match rx.blocking_recv() { - Some(Ok(batch)) => { - update_metrics(env, exec_context)?; - return prepare_output( - env, - array_addrs, - schema_addrs, - batch, - exec_context.debug_native, - ); - } - Some(Err(e)) => { - return Err(e.into()); - } - None => { - log_plan_metrics(exec_context, stage_id, partition); - return Ok(-1); + if exec_context.batch_receiver.is_some() { + let recv_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe( + || -> CometResult { + // Scope the rx borrow to just the blocking_recv call so that + // exec_context is free for update_metrics / prepare_output below. + let recv = exec_context + .batch_receiver + .as_mut() + .unwrap() + .blocking_recv(); + match recv { + Some(Ok(batch)) => { + update_metrics(env, exec_context)?; + prepare_output( + env, + array_addrs, + schema_addrs, + batch, + exec_context.debug_native, + ) + } + Some(Err(e)) => Err(e.into()), + None => { + log_plan_metrics(exec_context, stage_id, partition); + Ok(-1) + } + } + }, + )); + + match recv_result { + Ok(r) => return r, + Err(_panic) => { + #[cfg(feature = "oom-guard")] + if let Some(e) = + crate::execution::memory_pools::oom_guard::map_panic_to_error( + _panic.as_ref(), + ) + { + // Drop the receiver so any re-entry re-initializes. + exec_context.batch_receiver = None; + return Err(e.into()); + } + std::panic::resume_unwind(_panic); } } } // ScanExec path: busy-poll to interleave JVM batch pulls with stream polling - get_runtime().block_on(async { - loop { - let next_item = exec_context.stream.as_mut().unwrap().next(); - let poll_output = poll!(next_item); - - // Only check time/tracing every 100 polls to reduce overhead - exec_context.poll_count_since_metrics_check += 1; - if exec_context.poll_count_since_metrics_check >= 100 { - exec_context.poll_count_since_metrics_check = 0; - if let Some(interval) = exec_context.metrics_update_interval { - let now = Instant::now(); - if now - exec_context.metrics_last_update_time >= interval { - update_metrics(env, exec_context)?; - exec_context.metrics_last_update_time = now; + let poll_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + get_runtime().block_on(async { + loop { + let next_item = exec_context.stream.as_mut().unwrap().next(); + let poll_output = poll!(next_item); + + // Only check time/tracing every 100 polls to reduce overhead + exec_context.poll_count_since_metrics_check += 1; + if exec_context.poll_count_since_metrics_check >= 100 { + exec_context.poll_count_since_metrics_check = 0; + if let Some(interval) = exec_context.metrics_update_interval { + let now = Instant::now(); + if now - exec_context.metrics_last_update_time >= interval { + update_metrics(env, exec_context)?; + exec_context.metrics_last_update_time = now; + } + } + if exec_context.tracing_enabled { + log_memory_usage( + &exec_context.tracing_memory_metric_name, + total_reserved_for_thread(exec_context.rust_thread_id) as u64, + ); } } - if exec_context.tracing_enabled { - log_memory_usage( - &exec_context.tracing_memory_metric_name, - total_reserved_for_thread(exec_context.rust_thread_id) as u64, - ); - } - } - match poll_output { - Poll::Ready(Some(output)) => { - return prepare_output( - env, - array_addrs, - schema_addrs, - output?, - exec_context.debug_native, - ); - } - Poll::Ready(None) => { - log_plan_metrics(exec_context, stage_id, partition); - return Ok(-1); - } - Poll::Pending => { - // JNI call to pull batches from JVM into ScanExec operators. - // block_in_place lets tokio move other tasks off this worker - // while we wait for JVM data. - tokio::task::block_in_place(|| pull_input_batches(exec_context))?; + match poll_output { + Poll::Ready(Some(output)) => { + return prepare_output( + env, + array_addrs, + schema_addrs, + output?, + exec_context.debug_native, + ); + } + Poll::Ready(None) => { + log_plan_metrics(exec_context, stage_id, partition); + return Ok(-1); + } + Poll::Pending => { + // JNI call to pull batches from JVM into ScanExec operators. + // block_in_place lets tokio move other tasks off this worker + // while we wait for JVM data. + tokio::task::block_in_place(|| pull_input_batches(exec_context))?; + } } } + }) + })); + + match poll_result { + Ok(r) => r, + Err(_panic) => { + #[cfg(feature = "oom-guard")] + if let Some(e) = crate::execution::memory_pools::oom_guard::map_panic_to_error( + _panic.as_ref(), + ) { + // The block_on future was dropped mid-poll; null the stream so any + // inadvertent re-entry re-initializes rather than polling a half-consumed one. + exec_context.stream = None; + return Err(e.into()); + } + std::panic::resume_unwind(_panic); } - }) + } }); if exec_context.tracing_enabled { diff --git a/native/core/src/execution/memory_pools/config.rs b/native/core/src/execution/memory_pools/config.rs index 312a3604383..556b16a4bc5 100644 --- a/native/core/src/execution/memory_pools/config.rs +++ b/native/core/src/execution/memory_pools/config.rs @@ -30,6 +30,24 @@ pub(crate) enum MemoryPoolType { Unbounded, } +impl MemoryPoolType { + /// True when this pool's `reserved()` reflects a single task's usage, so a + /// per-task fair-share comparison is meaningful. False for process-wide pools + /// whose `reserved()` is the aggregate across tasks. Note the non-shared + /// per-task pools (`Greedy`/`FairSpill`) return true but keep no task registry, + /// so the fair-share divisor falls back to `executor_cores` for them rather + /// than the dynamic active-task count. + #[cfg_attr(not(feature = "oom-guard"), allow(dead_code))] + pub(crate) fn has_per_task_budget(&self) -> bool { + !matches!( + self, + MemoryPoolType::GreedyGlobal + | MemoryPoolType::FairSpillGlobal + | MemoryPoolType::Unbounded + ) + } +} + pub(crate) struct MemoryPoolConfig { pub(crate) pool_type: MemoryPoolType, pub(crate) pool_size: usize, diff --git a/native/core/src/execution/memory_pools/mod.rs b/native/core/src/execution/memory_pools/mod.rs index d7c2911f913..1d87f4cc333 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -18,6 +18,10 @@ mod config; mod fair_pool; pub mod logging_pool; +#[cfg(feature = "oom-guard")] +pub mod oom_guard; +#[cfg(feature = "oom-guard")] +mod real_usage_pool; mod task_shared; mod unified_pool; @@ -32,6 +36,8 @@ use std::sync::Arc; use unified_pool::CometUnifiedMemoryPool; pub(crate) use config::*; +#[cfg(feature = "oom-guard")] +pub(crate) use real_usage_pool::RealUsagePool; pub(crate) use task_shared::*; /// Creates the memory pool for a native plan. diff --git a/native/core/src/execution/memory_pools/oom_guard.rs b/native/core/src/execution/memory_pools/oom_guard.rs new file mode 100644 index 00000000000..3218aeb97c7 --- /dev/null +++ b/native/core/src/execution/memory_pools/oom_guard.rs @@ -0,0 +1,376 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::common::DataFusionError; +use std::alloc::{GlobalAlloc, Layout}; +use std::cell::Cell; +use std::sync::atomic::{AtomicBool, AtomicIsize, AtomicUsize, Ordering}; + +/// Per-thread drift is flushed into the shared balance once it crosses this. +const SETTLE_THRESHOLD: isize = 64 * 1024; + +/// Process-wide outstanding bytes (signed so transient under-settle is fine). +static BALANCE: AtomicIsize = AtomicIsize::new(0); +/// Enforcement limit in bytes; 0 means unset. +static LIMIT: AtomicUsize = AtomicUsize::new(0); +/// Master enforcement gate (single relaxed load on the hot path). +static ARMED: AtomicBool = AtomicBool::new(false); + +thread_local! { + /// Un-flushed per-thread delta. + static LOCAL_DRIFT: Cell = const { Cell::new(0) }; + /// Is this a query-worker thread eligible for enforcement? + static STAMPED: Cell = const { Cell::new(false) }; + /// Set while a guard panic is unwinding this thread, to avoid double-faults. + static UNWINDING: Cell = const { Cell::new(false) }; +} + +/// Payload of the panic raised when an armed, stamped thread exceeds the limit. +#[derive(Debug)] +pub struct OomGuardPanic { + pub balance: usize, + pub limit: usize, +} + +/// Arm the guard with a byte limit. Idempotent. +pub fn arm(limit_bytes: usize) { + LIMIT.store(limit_bytes, Ordering::Relaxed); + ARMED.store(true, Ordering::Relaxed); +} + +/// Disarm the guard (enforcement off; tracking continues cheaply). +#[allow(dead_code)] // used only by tests +pub fn disarm() { + ARMED.store(false, Ordering::Relaxed); +} + +/// Mark the current thread as a query-worker thread eligible for enforcement. +pub fn stamp_current_thread() { + STAMPED.with(|s| s.set(true)); +} + +/// Reset the per-thread unwinding guard after a guard panic has been caught on +/// this thread. Safe to call when not unwinding. The JNI caller thread is +/// reused across tasks, so this must run after catching an OomGuardPanic. +pub fn clear_unwinding() { + UNWINDING.with(|u| u.set(false)); +} + +/// If `panic` is an `OomGuardPanic`, clear this thread's unwinding guard and +/// return the mapped retriable error. Returns `None` for any other panic. +/// Centralizes the downcast + unwinding-reset + error mapping for all catch sites. +pub fn map_panic_to_error(panic: &(dyn std::any::Any + Send)) -> Option { + let g = panic.downcast_ref::()?; + clear_unwinding(); + Some(DataFusionError::ResourcesExhausted(format!( + "Comet OomGuard: native allocation pushed usage to {} bytes, over the limit of {} \ + bytes; failing this task", + g.balance, g.limit + ))) +} + +/// Current process-wide balance in bytes (never reported negative). +pub fn current_balance() -> usize { + BALANCE.load(Ordering::Relaxed).max(0) as usize +} + +/// Record an allocation of `size` bytes; may trip the breaker. +#[inline] +fn record_alloc(size: usize) { + track(size as isize); +} + +/// Record a deallocation of `size` bytes; never trips (credit only). +#[inline] +fn record_dealloc(size: usize) { + track(-(size as isize)); +} + +/// Core tracking + enforcement. Flushes drift; on a debit flush that crosses the +/// limit on an armed, stamped, non-unwinding thread, panics with `OomGuardPanic`. +#[inline] +fn track(delta: isize) { + let new_balance = LOCAL_DRIFT.with(|d| { + let mut drift = d.get(); + let flushed = settle(&mut drift, delta, &BALANCE); + d.set(drift); + flushed + }); + + if delta <= 0 { + return; // credits never enforce + } + let Some(balance) = new_balance else { return }; + if !ARMED.load(Ordering::Relaxed) { + return; + } + if !STAMPED.with(|s| s.get()) { + return; + } + if UNWINDING.with(|u| u.get()) { + return; + } + let limit = LIMIT.load(Ordering::Relaxed); + if should_trip(balance, limit) { + // At most one thread may fire the guard panic per arm cycle. CAS the + // master gate true->false; threads that lose the race bail before + // panic_any. The relaxed load above (line ~121) is not a serialization + // point: several threads can all read ARMED=true and reach here in the + // same tight window. If each then dispatches a panic, Rust's unwind ABI + // can abort the process with "failed to initiate panic" instead of + // unwinding cleanly (observed on the 5-concurrent repro: ~4 threads + // firing within ~10 ms -> exit 133). The guard re-arms on the next + // createPlan. + if ARMED + .compare_exchange(true, false, Ordering::Relaxed, Ordering::Relaxed) + .is_err() + { + return; + } + // panic_any boxes the payload, which re-enters this allocator and calls + // track() again. ARMED is now false so the re-entrant call short-circuits + // at the ARMED check above; setting UNWINDING adds defense in depth in + // case a concurrent createPlan re-arms mid-unwind. + UNWINDING.with(|u| u.set(true)); + std::panic::panic_any(OomGuardPanic { + balance: balance.max(0) as usize, + limit, + }); + } +} + +/// Pure helper: given the current shared balance and a limit, decide whether an +/// armed+stamped thread should trip the breaker. `limit == 0` means "unset". +fn should_trip(balance: isize, limit: usize) -> bool { + limit != 0 && balance > limit.try_into().unwrap_or(isize::MAX) +} + +/// Pure helper: add `delta` to `local_drift`; if it reaches or exceeds `SETTLE_THRESHOLD` +/// in magnitude, flush it into `shared` and return the new shared balance. +/// Otherwise return `None` (nothing flushed). +fn settle(local_drift: &mut isize, delta: isize, shared: &AtomicIsize) -> Option { + *local_drift = local_drift.wrapping_add(delta); + if local_drift.unsigned_abs() >= SETTLE_THRESHOLD as usize { + let flushed = *local_drift; + *local_drift = 0; + let prev = shared.fetch_add(flushed, Ordering::Relaxed); + Some(prev.wrapping_add(flushed)) + } else { + None + } +} + +/// Wraps an inner global allocator, tracking layout bytes for the OomGuard. +pub struct AccountingAllocator { + inner: A, +} + +impl AccountingAllocator { + pub const fn new(inner: A) -> Self { + Self { inner } + } +} + +unsafe impl GlobalAlloc for AccountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = self.inner.alloc(layout); + if !ptr.is_null() { + record_alloc(layout.size()); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + self.inner.dealloc(ptr, layout); + record_dealloc(layout.size()); + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let ptr = self.inner.alloc_zeroed(layout); + if !ptr.is_null() { + record_alloc(layout.size()); + } + ptr + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + // Account for and enforce the size delta BEFORE delegating to the inner + // realloc. If this trips the breaker it panics here, while `ptr` is still + // valid, so the unwind frees it correctly. Panicking *after* inner.realloc + // would be unsound: realloc may have already freed/moved the old block, and + // the caller (which never received the new pointer) would free the dangling + // old pointer on unwind and segfault. Only growth can trip; over-counting on + // a (rare) realloc failure errs on the conservative side for an OOM guard. + // + // Casts and subtraction are safe in practice: a single allocation cannot + // exceed isize::MAX on any real platform, so no wrapping or overflow occurs. + let old = layout.size() as isize; + let new = new_size as isize; + track(new - old); + self.inner.realloc(ptr, layout, new_size) + } +} + +#[cfg(test)] +fn reset_for_test() { + BALANCE.store(0, Ordering::Relaxed); + LIMIT.store(0, Ordering::Relaxed); + ARMED.store(false, Ordering::Relaxed); + LOCAL_DRIFT.with(|d| d.set(0)); + STAMPED.with(|s| s.set(false)); + UNWINDING.with(|u| u.set(false)); +} + +#[cfg(test)] +fn clear_unwinding_for_test() { + UNWINDING.with(|u| u.set(false)); +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Mutex; + + // Serializes tests that mutate the process-global guard state. + static GUARD: Mutex<()> = Mutex::new(()); + + #[test] + fn test_should_trip() { + assert!(!should_trip(100, 0)); // unset limit never trips + assert!(!should_trip(100, 200)); // under limit + assert!(!should_trip(200, 200)); // at limit (strictly greater required) + assert!(should_trip(201, 200)); // over limit + } + + #[test] + fn test_settle_accumulates_then_flushes() { + let shared = AtomicIsize::new(0); + let mut drift = 0isize; + // small allocs below threshold do not flush + assert_eq!(settle(&mut drift, 1024, &shared), None); + assert_eq!(shared.load(Ordering::Relaxed), 0); + // crossing the threshold flushes the accumulated drift + let new_balance = settle(&mut drift, SETTLE_THRESHOLD, &shared); + assert_eq!(new_balance, Some(1024 + SETTLE_THRESHOLD)); + assert_eq!(shared.load(Ordering::Relaxed), 1024 + SETTLE_THRESHOLD); + assert_eq!(drift, 0); // drift reset after flush + } + + #[test] + fn test_settle_flushes_negative_drift() { + let shared = AtomicIsize::new(1_000_000); + let mut drift = 0isize; + assert_eq!( + settle(&mut drift, -SETTLE_THRESHOLD, &shared), + Some(1_000_000 - SETTLE_THRESHOLD) + ); + assert_eq!(drift, 0); + } + + #[test] + fn test_settle_flushes_at_exact_threshold() { + let shared = AtomicIsize::new(0); + let mut drift = 0isize; + assert_eq!( + settle(&mut drift, SETTLE_THRESHOLD, &shared), + Some(SETTLE_THRESHOLD) + ); + assert_eq!(drift, 0); + } + + #[test] + fn test_disarmed_never_trips() { + let _g = GUARD.lock().unwrap_or_else(|e| e.into_inner()); + reset_for_test(); + stamp_current_thread(); + // not armed -> record_alloc must never panic regardless of size + record_alloc(usize::MAX / 2); + record_alloc(usize::MAX / 2); + } + + #[test] + fn test_unstamped_thread_never_trips() { + let _g = GUARD.lock().unwrap_or_else(|e| e.into_inner()); + reset_for_test(); + // arm with a tiny limit relative to current balance, but DO NOT stamp + let limit = current_balance() + 1; + arm(limit); + record_alloc(SETTLE_THRESHOLD as usize * 4); // big enough to flush + disarm(); + } + + #[test] + fn test_stamped_over_budget_trips() { + let _g = GUARD.lock().unwrap_or_else(|e| e.into_inner()); + reset_for_test(); + stamp_current_thread(); + let limit = current_balance() + SETTLE_THRESHOLD as usize; // headroom + arm(limit); + let result = std::panic::catch_unwind(|| { + // exceed the headroom in one flush + record_alloc(SETTLE_THRESHOLD as usize * 4); + }); + disarm(); + clear_unwinding_for_test(); + assert!(result.is_err(), "expected OomGuardPanic"); + let panic = result.unwrap_err(); + assert!( + panic.downcast_ref::().is_some(), + "panic payload should be OomGuardPanic" + ); + } + + // Drives a real heap allocation through the installed AccountingAllocator (only + // wrapped under the `oom-guard` feature) and confirms the guard trips. + #[test] + #[cfg(feature = "oom-guard")] + fn test_real_allocation_trips_guard() { + let _g = GUARD.lock().unwrap_or_else(|e| e.into_inner()); + reset_for_test(); + stamp_current_thread(); + // 8 MiB headroom over the current (noisy) baseline. + let headroom = 8 * 1024 * 1024; + arm(current_balance() + headroom); + + let result = std::panic::catch_unwind(|| { + // Allocate well past the headroom in 1 MiB chunks so a flush crosses the limit. + let mut held: Vec> = Vec::new(); + for _ in 0..64 { + held.push(vec![0u8; 1024 * 1024]); + } + // Touch the data so the allocation cannot be optimized away. + held.iter().map(|v| v.len()).sum::() + }); + + // Disarm BEFORE clearing UNWINDING so no post-catch allocation on this still-armed, + // still-stamped thread can re-trip outside the catch. + disarm(); + clear_unwinding_for_test(); + + assert!( + result.is_err(), + "large allocation on a stamped, armed thread should trip the guard" + ); + assert!( + result + .unwrap_err() + .downcast_ref::() + .is_some(), + "panic payload should be OomGuardPanic" + ); + } +} diff --git a/native/core/src/execution/memory_pools/real_usage_pool.rs b/native/core/src/execution/memory_pools/real_usage_pool.rs new file mode 100644 index 00000000000..cb96eb02a0c --- /dev/null +++ b/native/core/src/execution/memory_pools/real_usage_pool.rs @@ -0,0 +1,365 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::execution::memory_pools::{active_task_count, oom_guard}; +use datafusion::common::{resources_datafusion_err, DataFusionError}; +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use std::sync::Arc; + +/// Source of the current process-wide real allocator usage in bytes. Production +/// wiring uses `oom_guard::current_balance`; tests inject a controllable value. +type BalanceSource = Arc usize + Send + Sync>; + +/// A `MemoryPool` decorator that, on top of the inner pool's tracked-reservation +/// accounting, rejects growth when *real* allocator usage (untracked Arrow / join / +/// kernel bytes included) plus the requested amount would exceed a process-global +/// ceiling. Returning `ResourcesExhausted` lets DataFusion spill and retry. +pub(crate) struct RealUsagePool { + inner: Arc, + /// Process-global real-usage ceiling in bytes; 0 means unset (no gating). + ceiling: usize, + /// Fixed fallback divisor (concurrent-task count) used when the dynamic + /// active-task count is 0. `None` disables the fair-share guard (first-come), + /// used for pools whose `reserved()` is process-wide. + fair_share: Option, + balance_source: BalanceSource, +} + +impl std::fmt::Debug for RealUsagePool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RealUsagePool") + .field("inner", &self.inner) + .field("ceiling", &self.ceiling) + .field("fair_share", &self.fair_share) + .finish_non_exhaustive() + } +} + +impl RealUsagePool { + /// Wrap `inner` with the real-usage gate using the live OomGuard balance. + pub(crate) fn new( + inner: Arc, + ceiling: usize, + fair_share: Option, + ) -> Self { + Self { + inner, + ceiling, + fair_share, + balance_source: Arc::new(oom_guard::current_balance), + } + } + + /// Wrap `inner` with an explicit balance source (test seam). + #[cfg(test)] + fn with_balance_source( + inner: Arc, + ceiling: usize, + fair_share: Option, + balance_source: BalanceSource, + ) -> Self { + Self { + inner, + ceiling, + fair_share, + balance_source, + } + } +} + +/// Per-task fair share of `ceiling` given the number of concurrently active +/// tasks, or `cores_fallback` when the dynamic count is unavailable (0). The +/// divisor is floored at 1 so it is never zero. +fn fair_share_limit(ceiling: usize, active_tasks: usize, cores_fallback: usize) -> usize { + let n = if active_tasks > 0 { + active_tasks + } else { + cores_fallback + }; + ceiling / n.max(1) +} + +/// Given the process is already over the real-usage ceiling, decide whether to +/// reject this task's grow. `None` is first-come (reject whoever hit the ceiling); +/// `Some(s)` rejects only a task whose tracked reservation would exceed its fair +/// share `s`, sparing under-share tasks (the OomGuard breaker backstops runaway +/// cases). +fn should_reject_over_ceiling(reserved: usize, additional: usize, share: Option) -> bool { + match share { + None => true, + Some(s) => reserved.saturating_add(additional) > s, + } +} + +impl std::fmt::Display for RealUsagePool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "RealUsagePool(ceiling={}, inner=", self.ceiling)?; + std::fmt::Display::fmt(self.inner.as_ref(), f)?; + write!(f, ")") + } +} + +impl MemoryPool for RealUsagePool { + fn name(&self) -> &str { + self.inner.name() + } + + fn register(&self, consumer: &MemoryConsumer) { + self.inner.register(consumer) + } + + fn unregister(&self, consumer: &MemoryConsumer) { + self.inner.unregister(consumer) + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional) + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink) + } + + /// Gates growth on real allocator usage before delegating to the inner pool. + /// + /// `additional` is DataFusion's own parameter: the number of extra bytes the calling + /// `MemoryReservation` wants to hold. It is not a configured headroom or buffer — the gate + /// simply projects it onto the current real usage and compares against the ceiling. + fn try_grow( + &self, + reservation: &MemoryReservation, + additional: usize, + ) -> Result<(), DataFusionError> { + // Check the real-usage ceiling before delegating, so an over-budget request is + // rejected without speculatively reserving the inner pool. When the process is + // over the ceiling, the fair-share guard rejects only a task whose own tracked + // reservation exceeds its fair share, sparing innocent small tasks; the OomGuard + // breaker backstops runaway cases. Returning `ResourcesExhausted` lets DataFusion + // spill and retry. + if self.ceiling != 0 && additional != 0 { + let real = (self.balance_source)(); + if real.saturating_add(additional) > self.ceiling { + let share = self + .fair_share + .map(|cores| fair_share_limit(self.ceiling, active_task_count(), cores)); + if should_reject_over_ceiling(self.inner.reserved(), additional, share) { + return Err(resources_datafusion_err!( + "Comet real-usage gate: native usage {real} bytes + requested \ + {additional} bytes exceeds the off-heap budget of {} bytes; \ + spilling/failing this consumer", + self.ceiling + )); + } + } + } + self.inner.try_grow(reservation, additional) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::execution::memory_pool::{GreedyMemoryPool, UnboundedMemoryPool}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + fn fixed_source(bytes: Arc) -> BalanceSource { + Arc::new(move || bytes.load(Ordering::Relaxed)) + } + + #[test] + fn under_ceiling_succeeds_and_delegates() { + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let balance = Arc::new(AtomicUsize::new(100)); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 1000, + None, + fixed_source(balance), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + // real usage 100 + request 100 = 200 <= ceiling 1000 + assert!(pool.try_grow(&reservation, 100).is_ok()); + assert_eq!(inner.reserved(), 100); + } + + #[test] + fn over_ceiling_rejects_without_reserving_inner() { + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let balance = Arc::new(AtomicUsize::new(900)); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 1000, + None, + fixed_source(balance), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + // real usage 900 + request 200 = 1100 > ceiling 1000 -> reject + let result = pool.try_grow(&reservation, 200); + assert!(result.is_err(), "over-ceiling grow should be rejected"); + // inner pool is never touched on rejection, so there is nothing to roll back + assert_eq!(inner.reserved(), 0); + } + + #[test] + fn zero_ceiling_never_gates() { + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let balance = Arc::new(AtomicUsize::new(usize::MAX / 2)); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 0, + None, + fixed_source(balance), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + assert!(pool.try_grow(&reservation, 1024).is_ok()); + assert_eq!(inner.reserved(), 1024); + } + + #[test] + fn shrink_delegates() { + let inner: Arc = Arc::new(UnboundedMemoryPool::default()); + let balance = Arc::new(AtomicUsize::new(0)); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 1_000_000, + None, + fixed_source(balance), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + pool.try_grow(&reservation, 500).unwrap(); + assert_eq!(pool.reserved(), 500); + pool.shrink(&reservation, 200); + assert_eq!(pool.reserved(), 300); + } + + // Drives a real heap allocation through the installed AccountingAllocator (only + // wrapped under the `oom-guard` feature) and confirms the real-usage gate rejects. + // Robust to parallel test noise: other allocations only raise the balance further, + // which can only make the over-ceiling assertion more true. + #[test] + fn real_allocation_trips_real_usage_gate() { + let inner: Arc = Arc::new(UnboundedMemoryPool::default()); + let base = oom_guard::current_balance(); + // 4 MiB headroom over the (noisy) baseline. + let ceiling = base + 4 * 1024 * 1024; + let pool: Arc = + Arc::new(RealUsagePool::new(Arc::clone(&inner), ceiling, None)); + let reservation = MemoryConsumer::new("test").register(&pool); + + // Push real usage ~8 MiB above the baseline, held alive across the check so the + // balance stays elevated. 8 MiB > 64 KiB settle threshold, so it flushes to BALANCE. + let held: Vec = vec![0u8; 8 * 1024 * 1024]; + assert!( + oom_guard::current_balance() > ceiling, + "allocation should push balance over ceiling" + ); + + let result = pool.try_grow(&reservation, 1); + assert!( + result.is_err(), + "real usage over the ceiling should reject the grow" + ); + // Keep `held` alive until after the assertion above. + drop(held); + } + + #[test] + fn fair_share_limit_uses_active_count_when_positive() { + // active count wins over the fallback divisor + assert_eq!(fair_share_limit(1000, 4, 8), 250); + } + + #[test] + fn fair_share_limit_falls_back_when_no_active_tasks() { + assert_eq!(fair_share_limit(1000, 0, 5), 200); + } + + #[test] + fn fair_share_limit_floors_divisor_at_one() { + // active and fallback both zero -> divide by 1, no panic + assert_eq!(fair_share_limit(1000, 0, 0), 1000); + } + + #[test] + fn fair_share_limit_zero_when_ceiling_below_n() { + assert_eq!(fair_share_limit(3, 4, 8), 0); + } + + #[test] + fn should_reject_none_is_first_come() { + assert!(should_reject_over_ceiling(0, 1, None)); + assert!(should_reject_over_ceiling(1000, 0, None)); + } + + #[test] + fn should_reject_some_only_above_share() { + // strictly above share -> reject + assert!(should_reject_over_ceiling(400, 200, Some(500))); + // exactly at share -> allow + assert!(!should_reject_over_ceiling(300, 200, Some(500))); + // below share -> allow + assert!(!should_reject_over_ceiling(100, 100, Some(500))); + } + + #[test] + fn over_ceiling_rejects_task_over_fair_share() { + // ceiling 1000, fallback divisor 2, active count 0 in tests -> fair share 500 + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let balance = Arc::new(AtomicUsize::new(900)); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 1000, + Some(2), + fixed_source(balance), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + // Put this task above its 500-byte fair share. + inner.grow(&reservation, 600); + // Over ceiling (900 + 200 > 1000) AND over fair share (600 + 200 > 500) -> reject. + assert!(pool.try_grow(&reservation, 200).is_err()); + } + + #[test] + fn over_ceiling_spares_task_under_fair_share() { + // ceiling 1000, fallback divisor 2, active count 0 in tests -> fair share 500 + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let balance = Arc::new(AtomicUsize::new(1000)); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 1000, + Some(2), + fixed_source(balance), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + // This task holds only 100, under its 500 fair share. + inner.grow(&reservation, 100); + // Over ceiling (1000 + 50 > 1000) but under fair share (100 + 50 <= 500) -> allowed. + assert!(pool.try_grow(&reservation, 50).is_ok()); + // The grow was delegated to the inner pool. + assert_eq!(inner.reserved(), 150); + } +} diff --git a/native/core/src/execution/memory_pools/task_shared.rs b/native/core/src/execution/memory_pools/task_shared.rs index b5b4da61f9f..787dbf0a01f 100644 --- a/native/core/src/execution/memory_pools/task_shared.rs +++ b/native/core/src/execution/memory_pools/task_shared.rs @@ -23,6 +23,7 @@ use parking_lot::Mutex; use std::collections::hash_map::Entry; use std::collections::HashMap; use std::fmt; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Weak}; /// The memory pools for active task attempts. Weak references let the pool's normal `Arc` @@ -30,6 +31,20 @@ use std::sync::{Arc, Weak}; static TASK_SHARED_MEMORY_POOLS: Lazy>>> = Lazy::new(|| Mutex::new(HashMap::new())); +/// Number of live task-shared memory pools, which is the number of task attempts currently running +/// native plans. The real-usage fair-share guard uses this as the divisor for each task's share of +/// the executor budget. Incremented when a pool is constructed and decremented when it is dropped, +/// so it is exact even across the acquire/drop race that `TaskSharedMemoryPool::drop` handles. +static ACTIVE_TASK_COUNT: AtomicUsize = AtomicUsize::new(0); + +/// Current number of concurrently active task-shared task attempts. Returns 0 when none are active +/// or the configured pool type keeps no task registry, in which case the guard falls back to a +/// fixed divisor. +#[cfg_attr(not(feature = "oom-guard"), allow(dead_code))] +pub(crate) fn active_task_count() -> usize { + ACTIVE_TASK_COUNT.load(Ordering::Relaxed) +} + /// A transparent `MemoryPool` wrapper whose lifetime also controls its registry entry. #[derive(Debug)] struct TaskSharedMemoryPool { @@ -83,6 +98,11 @@ impl MemoryPool for TaskSharedMemoryPool { impl Drop for TaskSharedMemoryPool { fn drop(&mut self) { + // Paired with the increment in `acquire_task_shared_pool`. Both run exactly once per + // pool, so the count never drifts even when the registry entry below belongs to a + // replacement pool rather than this one. + ACTIVE_TASK_COUNT.fetch_sub(1, Ordering::Relaxed); + if let Entry::Occupied(entry) = TASK_SHARED_MEMORY_POOLS.lock().entry(self.task_attempt_id) { // An acquire racing with this drop can replace our expired `Weak` before we obtain the @@ -113,6 +133,7 @@ pub(crate) fn acquire_task_shared_pool( task_attempt_id, inner: create(), }); + ACTIVE_TASK_COUNT.fetch_add(1, Ordering::Relaxed); memory_pool_map.insert(task_attempt_id, Arc::downgrade(&memory_pool)); memory_pool } @@ -122,6 +143,10 @@ mod tests { use super::*; use datafusion::execution::memory_pool::UnboundedMemoryPool; + /// `ACTIVE_TASK_COUNT` is process-wide, so the tests in this module (the only ones that move + /// it) run serially to keep the delta assertions deterministic. + static TEST_LOCK: Mutex<()> = Mutex::new(()); + /// Tests share the process-wide pool map, so each uses its own task attempt id. fn acquire(task_attempt_id: i64) -> Arc { acquire_task_shared_pool(task_attempt_id, || Arc::new(UnboundedMemoryPool::default())) @@ -135,6 +160,7 @@ mod tests { #[test] fn plans_in_the_same_task_share_one_pool() { + let _guard = TEST_LOCK.lock(); let first = acquire(-1001); let second = acquire(-1001); assert!(Arc::ptr_eq(&first, &second)); @@ -142,6 +168,7 @@ mod tests { #[test] fn plans_in_different_tasks_get_different_pools() { + let _guard = TEST_LOCK.lock(); let first = acquire(-1002); let second = acquire(-1003); assert!(!Arc::ptr_eq(&first, &second)); @@ -149,6 +176,7 @@ mod tests { #[test] fn pool_is_removed_only_after_the_last_reference_drops() { + let _guard = TEST_LOCK.lock(); let first = acquire(-1004); let second = acquire(-1004); @@ -164,6 +192,7 @@ mod tests { #[test] fn dropping_the_reference_releases_the_pool() { + let _guard = TEST_LOCK.lock(); // Stands in for `createPlan` failing after the pool was acquired. The ordinary `Arc` drops // on unwind, so no explicit release path is needed. { @@ -175,6 +204,7 @@ mod tests { #[test] fn an_old_pool_does_not_remove_its_replacement() { + let _guard = TEST_LOCK.lock(); let old_pool = acquire(-1006); TASK_SHARED_MEMORY_POOLS.lock().remove(&-1006); let replacement = acquire(-1006); @@ -191,6 +221,7 @@ mod tests { /// pool's `Drop` must not evict the replacement's entry. #[test] fn concurrent_acquire_and_drop_leaves_a_consistent_registry() { + let _guard = TEST_LOCK.lock(); use std::thread; let threads: Vec<_> = (0..8) @@ -215,4 +246,47 @@ mod tests { let _pool = acquire(-1007); assert!(is_registered(-1007)); } + + #[test] + fn active_task_count_tracks_live_pools() { + let _guard = TEST_LOCK.lock(); + let base = active_task_count(); + + let first = acquire(-1008); + assert_eq!(active_task_count(), base + 1); + + // A second plan in the same task shares the pool, so the task is still counted once. + let second = acquire(-1008); + assert_eq!(active_task_count(), base + 1); + + let other_task = acquire(-1009); + assert_eq!(active_task_count(), base + 2); + + drop(first); + assert_eq!(active_task_count(), base + 2); + drop(second); + assert_eq!(active_task_count(), base + 1); + drop(other_task); + assert_eq!(active_task_count(), base); + } + + /// The increment pairs with the pool's `Drop`, not with registry membership, so the count must + /// stay exact through the acquire/drop race where a replacement pool shares a task attempt id. + #[test] + fn active_task_count_survives_a_replaced_registry_entry() { + let _guard = TEST_LOCK.lock(); + let base = active_task_count(); + + let old_pool = acquire(-1010); + TASK_SHARED_MEMORY_POOLS.lock().remove(&-1010); + let replacement = acquire(-1010); + assert_eq!(active_task_count(), base + 2, "both pools are live"); + + // The old pool's drop does not remove the replacement's entry, but must still decrement. + drop(old_pool); + assert_eq!(active_task_count(), base + 1); + + drop(replacement); + assert_eq!(active_task_count(), base); + } } diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs index 55da2c733aa..2067e53d6af 100644 --- a/native/core/src/execution/mod.rs +++ b/native/core/src/execution/mod.rs @@ -25,7 +25,7 @@ pub mod operators; pub(crate) mod planner; pub mod serde; pub use datafusion_comet_shuffle as shuffle; -mod memory_pools; +pub(crate) mod memory_pools; pub(crate) mod sort; pub(crate) mod spark_config; pub(crate) mod spark_plan; diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 4c2811cb5de..fb27d48ef90 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -25,6 +25,10 @@ pub(crate) const COMET_DEBUG_MEMORY: &str = "spark.comet.debug.memory"; pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str = "spark.comet.parquet.rowFilterPushdown.enabled"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; +#[cfg(feature = "oom-guard")] +pub(crate) const COMET_MEMORY_GUARD_ENABLED: &str = "spark.comet.exec.memoryGuard.enabled"; +#[cfg(feature = "oom-guard")] +pub(crate) const COMET_MEMORY_GUARD_SIZE: &str = "spark.comet.exec.memoryGuard.size"; pub(crate) trait SparkConfig { fn get_bool(&self, name: &str) -> bool; diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index b5656dba102..02896d472c2 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -75,18 +75,49 @@ pub mod debug; #[cfg(all( not(target_env = "msvc"), feature = "jemalloc", - not(feature = "mimalloc") + not(feature = "mimalloc"), + not(feature = "oom-guard") ))] #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; #[cfg(all( feature = "mimalloc", - not(all(not(target_env = "msvc"), feature = "jemalloc")) + not(all(not(target_env = "msvc"), feature = "jemalloc")), + not(feature = "oom-guard") ))] #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; +#[cfg(all( + not(target_env = "msvc"), + feature = "jemalloc", + not(feature = "mimalloc"), + feature = "oom-guard" +))] +#[global_allocator] +static GLOBAL: crate::execution::memory_pools::oom_guard::AccountingAllocator = + crate::execution::memory_pools::oom_guard::AccountingAllocator::new(Jemalloc); + +#[cfg(all( + feature = "mimalloc", + not(all(not(target_env = "msvc"), feature = "jemalloc")), + feature = "oom-guard" +))] +#[global_allocator] +static GLOBAL: crate::execution::memory_pools::oom_guard::AccountingAllocator = + crate::execution::memory_pools::oom_guard::AccountingAllocator::new(MiMalloc); + +// oom-guard enabled with system allocator (no mimalloc, and no jemalloc or on MSVC). +#[cfg(all( + feature = "oom-guard", + not(feature = "mimalloc"), + any(target_env = "msvc", not(feature = "jemalloc")) +))] +#[global_allocator] +static GLOBAL: crate::execution::memory_pools::oom_guard::AccountingAllocator = + crate::execution::memory_pools::oom_guard::AccountingAllocator::new(std::alloc::System); + #[no_mangle] pub extern "system" fn Java_org_apache_comet_NativeBase_init( e: EnvUnowned, diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 43f030a7d9d..6a3fb686b3c 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -1001,6 +1001,26 @@ object CometConf extends ShimCometConf { .bytesConf(ByteUnit.BYTE) .createWithDefault(100L * 1024 * 1024 * 1024) // 100 GB + val COMET_EXEC_MEMORY_GUARD_ENABLED: ConfigEntry[Boolean] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.memoryGuard.enabled") + .category(CATEGORY_EXEC) + .doc( + "Experimental. When enabled, Comet tracks real native memory allocations and aborts " + + "an over-budget task with a retriable error instead of risking an executor-wide OOM " + + "kill. Requires a Comet build with the 'oom-guard' native feature; has no effect " + + "on builds without it.") + .booleanConf + .createWithDefault(false) + + val COMET_EXEC_MEMORY_GUARD_SIZE: OptionalConfigEntry[Long] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.memoryGuard.size") + .category(CATEGORY_EXEC) + .doc( + "Experimental. Memory budget for the Comet native OOM guard (accepts sizes like '4g'). " + + "Defaults to the executor off-heap memory size (spark.memory.offHeap.size) when unset.") + .bytesConf(ByteUnit.BYTE) + .createOptional + val COMET_RESPECT_DATAFUSION_CONFIGS: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.respectDataFusionConfigs") .category(CATEGORY_TESTING)