feat: observe real native memory usage in the off-heap memory pools

.github/workflows/pr_build_linux.yml

@@ -518,6 +518,7 @@ jobs:
               org.apache.spark.sql.comet.CometTPCDSV1_4_PlanStabilitySuite
               org.apache.spark.sql.comet.CometTPCDSV2_7_PlanStabilitySuite
               org.apache.spark.sql.comet.CometTaskMetricsSuite
+              org.apache.comet.exec.CometMemoryPoolNativeUsageSuite
               org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite
               org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite
               org.apache.spark.sql.comet.PlanDataInjectorSuite

.github/workflows/pr_build_macos.yml

@@ -199,6 +199,7 @@ jobs:
               org.apache.spark.sql.comet.CometTPCDSV1_4_PlanStabilitySuite
               org.apache.spark.sql.comet.CometTPCDSV2_7_PlanStabilitySuite
               org.apache.spark.sql.comet.CometTaskMetricsSuite
+              org.apache.comet.exec.CometMemoryPoolNativeUsageSuite
               org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite
               org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite
               org.apache.spark.sql.comet.PlanDataInjectorSuite

docs/source/contributor-guide/memory_management.md

@@ -97,10 +97,12 @@ rather than asking for a separate allocation:
 memory_limit = spark.memory.offHeap.size * spark.comet.exec.memoryPool.fraction
 ```
 
-`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.
+`spark.comet.exec.memoryPool.fraction` defaults to `1.0`. Lowering it is one 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. It bounds only what Comet may reserve; the
+off-heap pools compare real allocator usage against the whole `spark.memory.offHeap.size`, so
+lowering the fraction to provoke spilling does not also lower the ceiling on real usage.
 
 A second value, `memory_limit_per_task`, is computed and passed alongside it, but only the on-heap
 pool types read it.
@@ -260,8 +262,11 @@ diverge for several structural reasons:
   JVM closes 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.
+gap is workload-dependent. `spark.comet.exec.memoryPool.fraction` lets operators hold back a slice
+of the pool that covers the gap for their workload, and the off-heap pools additionally compare
+the allocator's real usage against `spark.memory.offHeap.size`, logging a crossing by default and
+refusing the reservation when `spark.comet.exec.memoryPool.enforceNativeUsage` is set, so the gap
+is at least measured rather than only estimated.
 
 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
@@ -316,11 +321,17 @@ has bounds _declared reservations_, and the sections above describe several stru
 declared reservations are a lower bound on physical usage. The known gaps, roughly in order of how
 much they matter:
 
-- **No signal for real native usage.** The only way to observe the gap today is to enable tracing
-  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.
+- **Real native usage is process-wide, with no per-task attribution.** `alloc_accounting` reports
+  one balance for the whole executor, so the off-heap pools' check cannot tell which task caused
+  an overrun: once any task pushes real usage past the budget, every task's next reservation sees
+  it, and under enforcement every one of them is denied.
+- **The check gates reservations, not allocations.** An allocation that never goes through the
+  pool is counted after the fact and is never refused, so real usage can still exceed the budget
+  between reservations. Enforcement therefore falls on the operators that do reserve, which are
+  not necessarily the ones responsible for the overshoot, and spilling releases only reserved
+  bytes so it may not relieve an overshoot that lives in untracked allocations.
+- **`spark.comet.exec.memoryPool.fraction` is still set by hand.** It asks operators to guess how
+  much of the pool to hold back, even though the overrun it guards against is now measured.
 - **`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

docs/source/contributor-guide/tracing.md

@@ -37,12 +37,17 @@ The `alloc-accounting` feature adds a second, allocator-independent measure of n
 wraps whichever global allocator the build selected and reports the bytes it has handed out as
 `native_allocated`. Unlike `jemalloc_allocated` it does not require jemalloc, and it counts only
 what Rust code allocated, so it can be compared against the memory pool's reservations without the
-allocator's own caching in the way. The two features are independent and can be combined:
+allocator's own caching in the way. It also backs the off-heap memory pools' check of real usage
+against Comet's budget, described in the [tuning guide]. It is on by default, so the command above
+already includes it; the two features are independent, and a build that drops the defaults can
+re-add either one:
 
 ```shell
-make release COMET_FEATURES="jemalloc,alloc-accounting"
+cd native && cargo build --release --no-default-features --features hdfs-opendal,jemalloc
 ```
 
+[tuning guide]: ../user-guide/latest/tuning.md
+
 Example output:
 
 ```json
@@ -121,10 +126,10 @@ not being tracked by the pool.
 
 ## Definition of Labels
 
-| Label                            | Meaning                                                                                                                                                              |
-| -------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| jvm_heap_used                    | JVM heap memory usage of live objects for the executor process                                                                                                       |
-| jemalloc_allocated               | Native memory usage for the executor process (requires `jemalloc` feature)                                                                                           |
-| native_allocated                 | Bytes handed out by the Rust global allocator, process-wide (requires `alloc-accounting` feature). Approximate to within 64 KiB of un-flushed delta per live thread. |
-| thread_NNN_comet_memory_reserved | Memory reserved by Comet's DataFusion memory pool (summed across all contexts on the thread). NNN is the Rust thread ID.                                             |
-| thread_NNN_comet_jvm_shuffle     | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID.                                                                                  |
+| Label                            | Meaning                                                                                                                                                                    |
+| -------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| jvm_heap_used                    | JVM heap memory usage of live objects for the executor process                                                                                                             |
+| jemalloc_allocated               | Native memory usage for the executor process (requires `jemalloc` feature)                                                                                                 |
+| native_allocated                 | Bytes handed out by the Rust global allocator, process-wide (`alloc-accounting` feature, on by default). Approximate to within 64 KiB of un-flushed delta per live thread. |
+| thread_NNN_comet_memory_reserved | Memory reserved by Comet's DataFusion memory pool (summed across all contexts on the thread). NNN is the Rust thread ID.                                                   |
+| thread_NNN_comet_jvm_shuffle     | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID.                                                                                        |

docs/source/user-guide/latest/tuning.md

@@ -114,6 +114,28 @@ when there is sufficient memory in order to leave enough memory for other operat
 The `greedy_unified` pool type implements a greedy first-come first-serve limit. This pool works well for queries that do not
 need to spill or have a single spillable operator.
 
+Both pools reserve memory against Spark's ledger, which only counts what operators explicitly asked for. Native memory
+that never went through the pool, such as scratch buffers inside kernels or intermediate Arrow arrays, stays invisible
+until the executor exceeds its container limit and is killed. To make it visible, both pools also compare the memory
+the native allocator has actually handed out against `spark.memory.offHeap.size`.
+
+By default a crossing is only logged, once per task. Setting `spark.comet.exec.memoryPool.enforceNativeUsage` to
+`true` makes the pools refuse the reservation instead, so operators that can spill do so and those that cannot fail
+the task rather than the executor. Enforcement is off by default while the rate of false positives on real workloads
+is established, and because spilling releases only reserved bytes: a denial provoked by untracked allocations may not
+relieve the pressure it reports.
+
+Two things about the comparison are worth knowing. The budget is the whole off-heap size, not the reservable portion:
+`spark.comet.exec.memoryPool.fraction` bounds what Comet may reserve, and lowering it to provoke spilling
+deliberately does not lower this ceiling too. The two sides also do not measure the same population, because
+`spark.memory.offHeap.size` is shared with Spark's own off-heap allocations and with Comet's JVM-side shuffle pages
+while the measured usage counts only Comet's native allocations. It is a loose backstop against losing the executor,
+not a bound on total off-heap usage.
+
+The comparison is process-wide on both sides, so there is no per-task attribution: once any task pushes real usage
+past the budget, every task's next reservation sees it. It gates reservations only; allocations themselves are never
+refused.
+
 [shuffle]: #shuffle
 [Advanced Memory Tuning]: #advanced-memory-tuning
 

native/core/Cargo.toml

@@ -105,7 +105,7 @@ datafusion-functions-nested = { version = "55.1.0" }
 
 [features]
 backtrace = ["datafusion/backtrace"]
-default = ["hdfs-opendal"]
+default = ["hdfs-opendal", "alloc-accounting"]
 contrib-lance = ["dep:comet-contrib-lance"]
 hdfs-opendal = ["opendal", "object_store_opendal", "hdfs-sys"]
 jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"]
@@ -114,10 +114,11 @@ jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"]
 # Default builds carry zero Delta surface.
 contrib-delta = ["dep:comet-contrib-delta"]
 
-# Observability for real native memory usage. Wraps the global allocator to track the bytes it
-# hands out, and reports the total as the `native_allocated` tracing metric so it can be compared
-# against the memory pool's reservations. Never rejects an allocation. Off by default; a build
-# without it has no wrapper and no per-allocation work.
+# Real native memory usage. Wraps the global allocator to track the bytes it hands out, reports
+# the total as the `native_allocated` tracing metric, and backs the off-heap memory pools' check
+# of real usage against Comet's budget. Never rejects an allocation. On by default; build with
+# `--no-default-features` (re-adding the other defaults) to drop the wrapper and its
+# per-allocation work, which also leaves the pools' check reporting zero bytes in use.
 alloc-accounting = []
 
 # exclude optional packages from cargo machete verifications

native/core/src/alloc_accounting.rs

@@ -181,24 +181,32 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
 }
 
 #[cfg(test)]
-mod tests {
-    use super::*;
-    use std::alloc::System;
-    use std::sync::atomic::AtomicUsize;
+pub(crate) mod test_support {
     use std::sync::{Mutex, MutexGuard};
 
     /// `BALANCE` is process-wide and the crate's tests run in parallel, so a test that reads it
     /// sees every other test's allocations. The tests that move it by tens of megabytes take this
     /// lock so they cannot land inside each other's windows; the rest of the crate is kept out by
     /// making each window microseconds wide and each expected move far larger than anything else
     /// allocates in that time.
+    ///
+    /// Lives outside the tests module because the memory pool's own gate test moves the balance
+    /// the same way and has to share the lock.
     static SERIAL: Mutex<()> = Mutex::new(());
 
-    fn serial() -> MutexGuard<'static, ()> {
+    pub(crate) fn serial() -> MutexGuard<'static, ()> {
         SERIAL
             .lock()
             .unwrap_or_else(|poisoned| poisoned.into_inner())
     }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::test_support::serial;
+    use super::*;
+    use std::alloc::System;
+    use std::sync::atomic::AtomicUsize;
 
     const MIB: usize = 1024 * 1024;
 

native/core/src/execution/jni_api.rs

@@ -112,8 +112,9 @@ use crate::execution::tracing::{
 use crate::execution::memory_pools::logging_pool::LoggingMemoryPool;
 use crate::execution::spark_config::{
     SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXPLAIN_NATIVE_ENABLED,
-    COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED,
-    COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES,
+    COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_MEMORY_POOL_ENFORCE_NATIVE_USAGE,
+    COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES,
+    SPARK_MEMORY_OFFHEAP_SIZE,
 };
 use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID};
 use datafusion_comet_proto::spark_operator::operator::OpStruct;
@@ -495,6 +496,15 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
         let max_temp_directory_size =
             spark_config.get_u64(COMET_MAX_TEMP_DIRECTORY_SIZE, 100 * 1024 * 1024 * 1024);
         let logging_memory_pool = spark_config.get_bool(COMET_DEBUG_MEMORY);
+        // Defaults to true, so it has to be read with an explicit default: the config map only
+        // carries values Spark actually holds, and `CometExecIterator` injects this one for that
+        // reason.
+        // Defaults to false: the pools observe and log a crossing but do not act on it yet.
+        let enforce_native_usage =
+            spark_config.get_bool_with_default(COMET_MEMORY_POOL_ENFORCE_NATIVE_USAGE, false);
+        // Injected by `CometExecIterator` because it is a Spark config rather than a Comet one.
+        // Absent (0) in on-heap mode, which leaves the check off.
+        let off_heap_size = spark_config.get_usize(SPARK_MEMORY_OFFHEAP_SIZE, 0);
 
         with_trace("createPlan", tracing_enabled, || {
             // Init JVM classes
@@ -527,6 +537,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
                 memory_pool_type,
                 memory_limit,
                 memory_limit_per_task,
+                enforce_native_usage,
+                off_heap_size,
             )?;
             let memory_pool =
                 create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id);