From 0f9cbaa225b82ffb6e6e9d79c4235004cb30e0d9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 18 Sep 2026 09:33:24 -0600 Subject: [PATCH 1/2] docs: explain allocator hazards and diagram where memory is allocated The memory management guide described the layout of Comet's allocators but not the constraints that keep them correct, so the rules that CometTaskMemoryManager already follows were implicit and easy to break by accident. - Add a constraints subsection covering why getUsed and spill must stay lock-free, why NonFatal does not contain an acquireExecutionMemory call, why Spark exposes no per-consumer usage figure, how a partial grant is stranded when a later spill throws, and why a consumer whose spill returns zero takes budget it can never give back. - Add two mermaid diagrams: which allocator each JVM call site uses and who is charged for the bytes, and everything the pod cgroup counts toward RSS grouped by who accounts for it. - Distinguish off-heap from native heap. Both sit outside the JVM heap and both count toward RSS, but only the first is allocated by JVM code and reportable to TaskMemoryManager, and the budget is shared even though the memory is not. - Name Comet's actual C dependencies, and the codecs that are pure Rust despite their names, so the non-Rust allocation bullet is actionable. - Describe spark.comet.exec.memoryPool.fraction as a margin rather than a haircut, in both places it appears. The tuning guide said only that Comet's memory accounting "isn't 100% accurate", which gave the reader nothing to act on. Say what the pool actually tracks, which is memory an operator explicitly reserves, name those cases, and list what goes uncounted, so the reason for lowering memoryPool.fraction is concrete. Also fix two places where off-heap and native heap were conflated: an ffi.md lifecycle table whose column header read Off-heap/Native for a direction in which nothing is off-heap, an ArrowBuf annotated as off-heap inside a JVM heap box, and a tuning.md sentence that read as though Comet's own allocations come from JVM off-heap memory. --- docs/source/contributor-guide/ffi.md | 4 +- .../contributor-guide/memory_management.md | 130 +++++++++++++++++- docs/source/user-guide/latest/tuning.md | 27 +++- 3 files changed, 149 insertions(+), 12 deletions(-) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index f45fb9beafd..c1f53b138ca 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -132,7 +132,7 @@ JVM Heap: Native Memory: │ ColumnarBatch │ │ FFI_ArrowArray │ │ ┌──────────────┐ │ │ ┌──────────────┐ │ │ │ ArrowBuf │─┼──────────────>│ │ buffers[0] │ │ -│ │ (off-heap) │ │ │ │ (pointer) │ │ +│ │ (handle) │ │ │ │ (pointer) │ │ │ └──────────────┘ │ │ └──────────────┘ │ └──────────────────┘ └──────────────────┘ │ │ @@ -253,7 +253,7 @@ pub extern "system" fn Java_..._exportVector( ### Wrapper Object Lifecycle (Native → JVM) ``` -Time Native Memory JVM Heap Off-heap/Native +Time Native Memory JVM Heap Data location ──────────────────────────────────────────────────────────────────────── t0 RecordBatch produced - Data in native in DataFusion diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index c252afedc45..91b7403d10f 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -72,7 +72,26 @@ page: | Comet JVM Arrow (`CometArrowAllocator`) | Off-heap | **Nothing**: a `RootAllocator(Long.MaxValue)` | No | | Comet JVM shuffle pages | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | -Two observations follow. +Several observations follow. + +**"Off-heap" and "native heap" are not the same thing.** Both sit outside the JVM heap and both +count toward container RSS, which is what makes them easy to conflate, but they differ in who +allocates the bytes and in who is able to see them. + +- **Off-heap** is allocated by JVM code: `sun.misc.Unsafe` for Spark's Tungsten pages, and Arrow's + Unsafe-backed allocator for Comet's Java Arrow buffers. A JVM-side allocator makes the call, so a + JVM-side consumer is in a position to report it, and `spark.memory.offHeap.size` together with + `TaskMemoryManager` arbitrates it. +- **Native heap** is allocated by Rust through its global allocator, in the same process. No JVM + code is involved and no JVM metric counts it, so the only layer that sees any of it is Comet's own + memory pool, and then only the portion that operators explicitly reserve. + +**The budget is shared even though the memory is not.** `memory_limit` is derived from +`spark.memory.offHeap.size` (see [Where Comet's budget comes from](#where-comets-budget-comes-from)), +and in both off-heap pool types every native reservation is forwarded to Spark's off-heap execution +pool over JNI. A native allocation therefore consumes the same accounting budget as a Tungsten page +while occupying entirely different memory. Raising `spark.memory.offHeap.size` raises the ceiling +for both at once, and raises the pod's memory request by the same amount. **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 @@ -87,6 +106,78 @@ returns `CometUnifiedShuffleMemoryAllocator`, a Spark `MemoryConsumer` drawing f `spark.memory.offHeap.size`, so shuffle pages are arbitrated against Spark's other consumers in the same task like any other allocation. +Which allocator each call site uses, and who ends up charged for the bytes: + +```mermaid +flowchart LR + subgraph SITES["JVM Arrow allocation sites"] + NU["NativeUtil
FFI structs, imports, exports"] + UDF["CometUdfBridge
JVM UDF inputs and result"] + CGO["CometBatchKernelCodegenOutput
codegen UDF output"] + SR["StreamReader
shuffle and IPC reads"] + NAS["CometNativeArrowSource
stream and readerBatchIter"] + CACHE["ArrowCachedBatchSerializer"] + PY["CometArrowPythonRunnerBase"] + BC["Utils broadcast-coalesce"] + end + + ROOT["CometArrowAllocator
RootAllocator, no limit
no allocation listener"] + SHUF["Comet JVM shuffle pages
CometUnifiedShuffleMemoryAllocator"] + NPOOL["Comet native memory pool
declared reservations only"] + TMM["Spark off-heap execution pool
TaskMemoryManager"] + NOBODY["accounted by nobody"] + + NU --> ROOT + UDF --> ROOT + CGO --> ROOT + SR --> ROOT + NAS --> ROOT + CACHE --> ROOT + PY --> ROOT + BC --> ROOT + ROOT --> NOBODY + SHUF --> TMM + NPOOL -->|"CometTaskMemoryManager over JNI"| TMM +``` + +### Constraints on a Comet memory consumer + +`CometTaskMemoryManager` is the one place where Comet code acts as a Spark `MemoryConsumer`, and the +rules below are why it is written the way it is. Each is easy to break by accident. + +**`getUsed` and `spill` must stay lock-free.** Spark calls both while already holding the +`TaskMemoryManager` monitor, which is why `NativeMemoryConsumer.getUsed` reads an `AtomicLong` and +`spill` returns without touching guarded state. A consumer that took a monitor of its own in either +method would risk a deadlock: a native reservation arriving over JNI on a Comet Tokio worker holds +Spark's monitor and would wait for the consumer's, while whatever held the consumer's waited for +Spark's. For the same reason, nothing reachable from a `spill` callback may allocate memory that +routes back through the same consumer. + +**`scala.util.control.NonFatal` does not contain an acquisition.** `acquireExecutionMemory` fails in +three ways and `NonFatal` catches only the first. It runs other consumers' `spill`, where +`TaskMemoryManager` turns an interrupted spill into a `RuntimeException` and an `IOException` into a +`SparkOutOfMemoryError`, which extends `OutOfMemoryError` and is therefore an `Error`. Separately, +`ExecutionMemoryPool.acquireMemory` parks in `lock.wait()` when a task is below its fair share, so +killing a task raises a plain `InterruptedException`, which `NonFatal` excludes by name. Comet +reaches this method from native over JNI, so whatever escapes crosses the JNI boundary. + +**Spark exposes no per-consumer usage figure.** `TaskMemoryManager.getMemoryConsumptionForThisTask` +is task-wide, and a consumer that reaches `acquireExecutionMemory` directly rather than through +`MemoryConsumer.acquireMemory` keeps an inherited `used` of zero. That is why `NativeMemoryConsumer` +overrides `getUsed` to report Comet's own tally instead: without it, Spark's spill-victim ordering +and `showMemoryUsage` would believe the consumer held nothing. + +**A partial grant can be stranded.** `acquireExecutionMemory` takes its first grant from the pool +before asking other consumers to spill, so when a spill throws, the task has been charged for bytes +the call never returns. Nothing releases them until Spark's final task cleanup, so they are headroom +nobody can use for the rest of the task. Any caller that swallows the exception has to reconcile +that grant, and the only figure available for doing so is the task-wide one above. + +**A consumer whose `spill` returns zero takes budget it can never give back.** +`NativeMemoryConsumer.spill` returns `0`, so Spark can select it as a spill victim and reclaim +nothing from it; it is only ever a spill trigger. The bytes it holds are real, so other consumers in +the same task see correspondingly less headroom and can spill earlier than they otherwise would. + ## Where Comet's budget comes from `CometExecIterator.getMemoryConfig` computes the budget once per executor and passes it across JNI @@ -253,7 +344,15 @@ diverge for several structural reasons: 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`. + `mmap`ed, never passes through Rust's `GlobalAlloc`, so neither the memory pool nor the + `jemalloc_allocated` metric sees it. In a default build the C dependencies are libzstd + (`zstd-sys`, behind the Parquet `zstd` codec), libhdfs (`hdfs-sys`, pulled in by the default + `hdfs-opendal` feature), and the TLS stack used for cloud object stores (`aws-lc-sys`). Building + with the `jemalloc` or `mimalloc` feature adds the allocator itself (`tikv-jemalloc-sys`, + `libmimalloc-sys`). It is worth knowing which dependencies are _not_ C, because several names + suggest otherwise: the other Parquet codecs are pure Rust in this build, `snap` for Snappy, + `lz4_flex` for LZ4 and `zlib-rs` for gzip, as is `libbz2-rs-sys` despite its name, so those + allocations do pass through `GlobalAlloc` and are counted. - **Batches in flight across the FFI boundary.** Reservations stop at the operator that made them. Imported JVM batches are reserved only while a reserving operator holds them, and exported native batches have usually been released by the time the JVM receives them yet stay resident until the @@ -261,7 +360,7 @@ diverge for several structural reasons: 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. +hand-tune a margin 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 @@ -289,6 +388,29 @@ hard ceiling on the sum of everything in the container. That cgroup counts, amon - Comet's JVM-side Arrow buffers (`CometArrowAllocator`), - page cache charged to the cgroup by the container's file I/O, including spill files. +Everything the cgroup counts, and who accounts for each part: + +```mermaid +flowchart TB + subgraph CG["pod cgroup memory.max, kernel OOM kill above this"] + subgraph SEEN["visible to Spark's accounting"] + HEAP["JVM heap
execution and storage
spark.executor.memory"] + TUNG["Spark Tungsten off-heap
TaskMemoryManager"] + SHUFP["Comet JVM shuffle pages
CometUnifiedShuffleMemoryAllocator"] + end + subgraph DECL["declared to Comet's native pool only"] + NATRES["Comet native heap
operators that call try_grow"] + end + subgraph NONE["accounted by nobody"] + NATUND["Comet native heap, undeclared
kernels, array builders, decompression
Parquet metadata, object_store, tokio"] + ARROWR["Comet JVM Arrow
CometArrowAllocator, unbounded"] + NONHEAP["JVM non-heap
metaspace, code cache, thread stacks
GC structures, Netty direct buffers"] + PAGEC["page cache charged to the cgroup
file I/O, including spill files"] + FRAG["allocator overhead
fragmentation, padding
jemalloc retained and dirty pages"] + end + end +``` + 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 @@ -320,7 +442,7 @@ much they matter: with the `jemalloc` feature and compare `jemalloc_allocated` against summed reservations after the fact. There is no runtime value that an operator, a metric, or a policy could read. - **`spark.comet.exec.memoryPool.fraction` is a manual proxy for the gap.** It asks operators to - guess a per-workload haircut rather than measuring anything. + guess a per-workload margin rather than measuring anything. - **`CometArrowAllocator` is unbounded** and participates in no budget. - **Buffer and reservation lifetimes are independent across the FFI boundary.** A batch can be resident on either side with no reservation covering it, because reservations are made and diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 25b805d1ccb..205721408dc 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -83,12 +83,27 @@ than before. See the [Determining How Much Memory to Allocate] section for more ### Configuring Comet Memory Comet shares an off-heap memory pool with Spark. The size of the pool is -specified by `spark.memory.offHeap.size`. - -Comet's memory accounting isn't 100% accurate and this can result in Comet using more memory than it reserves, -leading to out-of-memory exceptions. To work around this issue, it is possible to -set `spark.comet.exec.memoryPool.fraction` to a value less than `1.0` to restrict the amount of memory that can be -reserved by Comet. +specified by `spark.memory.offHeap.size`. The pool is a shared _budget_ rather than a shared +allocator: Comet's native operators allocate from the Rust heap rather than from JVM off-heap +memory, but every reservation they make is charged against this same pool, so Comet and Spark's own +off-heap consumers draw down one number. + +Comet's memory pool only tracks memory that an operator explicitly reserves, which in practice means the batches +an operator deliberately accumulates: the sort buffer, the build side of a hash join, hash aggregation state, and the +shuffle writer's buffered partitions. Memory that is not reserved is invisible to the pool no matter how much of it +there is. That includes: + +- per-batch working memory in expression kernels and Arrow array builders, +- decompression buffers and Parquet reader structures, +- object store request buffers and the async runtime's own machinery, +- Arrow buffers allocated on the JVM side, which no budget covers at all, +- allocator overhead: buffer padding, size-class rounding, fragmentation, and pages the allocator retains after a + free rather than returning to the operating system. + +Reserved memory is therefore a lower bound on what Comet really uses, and how far below it sits depends on the +workload. This is why Comet can stay within the pool's limit and still push the executor past its container limit. +To leave room for the part that is not counted, set `spark.comet.exec.memoryPool.fraction` to a value less than +`1.0`, which restricts the amount of memory Comet is allowed to reserve. For more details about Spark off-heap memory mode, please refer to [Spark documentation]. From 22d3078e57d521a698a2c27e4286e0d21b2c5a22 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 18 Sep 2026 11:02:04 -0600 Subject: [PATCH 2/2] docs: show declared native reservations as charged to Spark's budget Review feedback: the container diagram placed native reservations in a group of their own, outside Spark-visible accounting. In off-heap mode, which is the only mode this page covers, both valid pool types forward a successful try_grow to CometTaskMemoryManager.acquireMemory, which calls Spark's acquireExecutionMemory. The reservation therefore does consume Spark's off-heap execution budget even though Rust allocates the bytes. Move the reservation node into the Spark-visible group, labelled so the distinction between allocation location and budget accounting survives, and rewrite the caption to say that Spark's accounting covers the group in two different senses. Make the three other statements of the same claim consistent: the allocator table's visibility cell, the native-heap bullet, and the overview paragraph. --- .../contributor-guide/memory_management.md | 36 +++++++++++-------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 91b7403d10f..91199dec35e 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -53,10 +53,10 @@ calls `System.exit(SparkExitCode.OOM)` (exit code 52). Real heap exhaustion ther the whole executor, not one task. An executor loss on its own does not tell you which budget was exceeded; the exit code does. -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. +Comet's difficulty is that its allocations are made by Rust code, so no JVM allocator produces them +and no JVM metric measures them, yet they land squarely in container RSS. Comet therefore maintains +its own budget that is meant to shadow the physical one, and declares it to Spark so that the two +compete for a single number. The accuracy of that shadow is the central problem this page is about. ## Who allocates what @@ -68,7 +68,7 @@ page: | --------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- | | 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 native (Rust global allocator) | Native heap | `memory_limit` (see below), enforced only via the memory pool | Reservations only | | Comet JVM Arrow (`CometArrowAllocator`) | Off-heap | **Nothing**: a `RootAllocator(Long.MaxValue)` | No | | Comet JVM shuffle pages | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | @@ -83,8 +83,10 @@ allocates the bytes and in who is able to see them. JVM-side consumer is in a position to report it, and `spark.memory.offHeap.size` together with `TaskMemoryManager` arbitrates it. - **Native heap** is allocated by Rust through its global allocator, in the same process. No JVM - code is involved and no JVM metric counts it, so the only layer that sees any of it is Comet's own - memory pool, and then only the portion that operators explicitly reserve. + code makes the call, so no JVM-side allocator or metric ever measures these bytes. The only layer + that sees any of them is Comet's own memory pool, and then only the portion that operators + explicitly reserve. That portion is still charged to Spark, as the next paragraph describes; what + is never reserved is measured by nothing and budgeted by nobody. **The budget is shared even though the memory is not.** `memory_limit` is derived from `spark.memory.offHeap.size` (see [Where Comet's budget comes from](#where-comets-budget-comes-from)), @@ -397,9 +399,7 @@ flowchart TB HEAP["JVM heap
execution and storage
spark.executor.memory"] TUNG["Spark Tungsten off-heap
TaskMemoryManager"] SHUFP["Comet JVM shuffle pages
CometUnifiedShuffleMemoryAllocator"] - end - subgraph DECL["declared to Comet's native pool only"] - NATRES["Comet native heap
operators that call try_grow"] + NATRES["Comet native heap, reserved
operators that call try_grow
declared to Spark over JNI, never measured"] end subgraph NONE["accounted by nobody"] NATUND["Comet native heap, undeclared
kernels, array builders, decompression
Parquet metadata, object_store, tokio"] @@ -411,11 +411,17 @@ flowchart TB end ``` -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. +Spark's accounting covers the first group, though not in the same sense throughout it. The JVM +heap, Tungsten pages and Comet's shuffle pages are allocated by JVM code that reports what it +allocated. A native reservation is a number an operator declared before allocating: `try_grow` +succeeds only once `CometTaskMemoryManager` has charged Spark's off-heap execution pool over JNI, so +the budget really is spent, but nothing measured the bytes and the reservation is only a lower bound +on them. The second group is outside every accounting layer. + +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: