feat: add native allocation accounting for memory observability

.github/actions/rust-test/action.yaml

@@ -70,3 +70,16 @@ runs:
         export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH}
         RUST_BACKTRACE=1 cargo nextest run
 
+    # The `alloc-accounting` feature is off by default, so nothing else in CI compiles the
+    # allocator wrapper, its backend selection, or the benchmark's liveness guards. Lint them with
+    # the accounting wrapper over jemalloc, run the accounting tests with the wrapper installed, and
+    # check the system-allocator arm.
+    - name: Check and test the alloc-accounting feature
+      shell: bash
+      run: |
+        cd native
+        export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH}
+        cargo clippy --color=never -p datafusion-comet --all-targets --features jemalloc,alloc-accounting -- -D warnings
+        RUST_BACKTRACE=1 cargo nextest run -p datafusion-comet --lib --features jemalloc,alloc-accounting alloc_accounting
+        cargo check -p datafusion-comet --features alloc-accounting
+

docs/source/contributor-guide/tracing.md

@@ -33,6 +33,16 @@ Additionally, enabling the `jemalloc` feature will enable tracing of native memo
 make release COMET_FEATURES="jemalloc"
 ```
 
+The `alloc-accounting` feature adds a second, allocator-independent measure of native memory. It
+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:
+
+```shell
+make release COMET_FEATURES="jemalloc,alloc-accounting"
+```
+
 Example output:
 
 ```json
@@ -55,9 +65,9 @@ Example trace visualization:
 
 ## Analyzing Memory Usage
 
-The `analyze_trace` tool parses a trace log and compares jemalloc usage against the sum of per-thread
-Comet memory pool reservations. This is useful for detecting untracked native memory growth where jemalloc
-allocations exceed what the memory pools account for.
+The `analyze_trace` tool parses a trace log and compares the process-wide native allocation counter against
+the sum of per-thread Comet memory pool reservations. This is useful for detecting untracked native memory
+growth where native allocations exceed what the memory pools account for.
 
 Build and run:
 
@@ -66,31 +76,34 @@ cd native
 cargo run --bin analyze_trace -- /path/to/comet-event-trace.json
 ```
 
-The tool reads counter events from the trace log. Because tracing logs metrics per thread, `jemalloc_allocated`
-is a process-wide value (the same global allocation reported from whichever thread logs it), while
-`thread_NNN_comet_memory_reserved` values are per-thread pool reservations that are summed to get the total
-tracked memory.
+The tool reads counter events from the trace log. Because tracing logs metrics per thread, `native_allocated`
+and `jemalloc_allocated` are process-wide values (the same global allocation reported from whichever thread
+logs it), while `thread_NNN_comet_memory_reserved` values are per-thread pool reservations that are summed to
+get the total tracked memory. The tool analyzes `native_allocated` when the trace contains it, since that
+counts only what Rust code holds from the allocator, and otherwise falls back to `jemalloc_allocated`. The
+output names the counter it used. A trace with neither counter is rejected.
 
 Sample output:
 
 ```
 === Comet Trace Memory Analysis ===
 
 Counter events parsed: 193104
+Allocation counter:    jemalloc_allocated
 Threads with memory pools: 8
-Peak jemalloc allocated:   3068.2 MB
+Peak jemalloc_allocated:   3068.2 MB
 Peak pool total:           2864.6 MB
-Peak excess (jemalloc - pool): 364.6 MB
+Peak excess (jemalloc_allocated - pool): 364.6 MB
 
-WARNING: jemalloc exceeded pool reservation at 138 sampled points:
+WARNING: jemalloc_allocated exceeded pool reservation at 138 sampled points:
 
-     Time (us)        jemalloc      pool_total          excess
---------------------------------------------------------------
-        179578        210.8 MB          0.1 MB        210.7 MB
-        429663        420.5 MB        145.1 MB        275.5 MB
-       1304969       2122.5 MB       1797.2 MB        325.2 MB
-      21974838        407.0 MB         42.3 MB        364.6 MB
-      33543599          5.5 MB          0.1 MB          5.3 MB
+     Time (us)  jemalloc_allocated      pool_total          excess
+------------------------------------------------------------------
+        179578            210.8 MB          0.1 MB        210.7 MB
+        429663            420.5 MB        145.1 MB        275.5 MB
+       1304969           2122.5 MB       1797.2 MB        325.2 MB
+      21974838            407.0 MB         42.3 MB        364.6 MB
+      33543599              5.5 MB          0.1 MB          5.3 MB
 
 --- Final per-thread pool reservations ---
 
@@ -102,14 +115,16 @@ WARNING: jemalloc exceeded pool reservation at 138 sampled points:
   Total: 0.0 MB
 ```
 
-Some excess is expected (jemalloc metadata, fragmentation, non-pool allocations like Arrow IPC buffers).
-Large or growing excess may indicate memory that is not being tracked by the pool.
+Some excess is expected (allocator metadata and fragmentation for `jemalloc_allocated`, and non-pool
+allocations like Arrow IPC buffers for either counter). Large or growing excess may indicate memory that is
+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)                                               |
-| 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 (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.                                                                                  |

native/common/src/bin/analyze_trace.rs

@@ -16,8 +16,9 @@
 // under the License.
 
 //! Analyzes a Comet chrome trace event log (`comet-event-trace.json`) and
-//! compares jemalloc usage against the sum of per-thread Comet memory pool
-//! reservations. Reports any points where jemalloc exceeds the total pool size.
+//! compares the process-wide native allocation counter against the sum of
+//! per-thread Comet memory pool reservations. Reports any points where the
+//! allocated bytes exceed the total pool size.
 //!
 //! Usage:
 //!   cargo run --bin analyze_trace -- <path-to-comet-event-trace.json>
@@ -27,6 +28,15 @@ use std::collections::HashMap;
 use std::io::{BufRead, BufReader};
 use std::{env, fs::File};
 
+/// The process-wide allocation counters the tool understands, most preferred first.
+///
+/// `native_allocated` (the `alloc-accounting` feature) counts only the bytes Rust code holds from
+/// the global allocator, so it is the tighter comparison against pool reservations.
+/// `jemalloc_allocated` (the `jemalloc` feature) also includes jemalloc's own metadata. A trace
+/// that carries both is analyzed against `native_allocated` alone; a trace with neither cannot be
+/// analyzed.
+const ALLOCATED_COUNTERS: [&str; 2] = ["native_allocated", "jemalloc_allocated"];
+
 /// A single Chrome trace event (only the fields we care about).
 #[derive(Deserialize)]
 struct TraceEvent {
@@ -42,7 +52,7 @@ struct TraceEvent {
 /// Snapshot of memory state at a given timestamp.
 struct MemorySnapshot {
     ts: u64,
-    jemalloc: u64,
+    allocated: u64,
     pool_total: u64,
 }
 
@@ -61,14 +71,16 @@ fn main() {
     let file = File::open(&args[1]).expect("Failed to open trace file");
     let reader = BufReader::new(file);
 
-    // Latest jemalloc value (global, not per-thread)
-    let mut latest_jemalloc: u64 = 0;
+    // Index into ALLOCATED_COUNTERS of the counter being analyzed, once one has been seen
+    let mut source: Option<usize> = None;
+    // Latest allocated value (global, not per-thread)
+    let mut latest_allocated: u64 = 0;
     // Per-thread pool reservations: thread_NNN -> bytes
     let mut pool_by_thread: HashMap<String, u64> = HashMap::new();
-    // Points where jemalloc exceeded pool total
+    // Points where allocated exceeded pool total
     let mut violations: Vec<MemorySnapshot> = Vec::new();
     // Track peak values
-    let mut peak_jemalloc: u64 = 0;
+    let mut peak_allocated: u64 = 0;
     let mut peak_pool_total: u64 = 0;
     let mut peak_excess: u64 = 0;
     let mut counter_events: u64 = 0;
@@ -110,11 +122,28 @@ fn main() {
 
         counter_events += 1;
 
-        if event.name == "jemalloc_allocated" {
-            if let Some(val) = event.args.get("jemalloc_allocated") {
-                latest_jemalloc = val.as_u64().unwrap_or(0);
-                if latest_jemalloc > peak_jemalloc {
-                    peak_jemalloc = latest_jemalloc;
+        if let Some(rank) = ALLOCATED_COUNTERS
+            .iter()
+            .position(|name| *name == event.name)
+        {
+            match source {
+                // A preferred counter is present in this trace; ignore the other one.
+                Some(current) if current < rank => continue,
+                Some(current) if current == rank => {}
+                // First sighting of a more preferred counter. Start over so the peaks and
+                // violations reported all come from a single source.
+                _ => {
+                    source = Some(rank);
+                    latest_allocated = 0;
+                    peak_allocated = 0;
+                    peak_excess = 0;
+                    violations.clear();
+                }
+            }
+            if let Some(val) = event.args.get(&event.name) {
+                latest_allocated = val.as_u64().unwrap_or(0);
+                if latest_allocated > peak_allocated {
+                    peak_allocated = latest_allocated;
                 }
             }
         } else if event.name.contains("comet_memory_reserved") {
@@ -129,14 +158,16 @@ fn main() {
             continue;
         }
 
-        // After each jemalloc or pool update, check the current state
+        // After each allocated or pool update, check the current state. A comparison needs one
+        // sample of each side: an observed zero reservation is a real value that allocation can
+        // exceed, so only the absence of any pool sample defers the check.
         let pool_total: u64 = pool_by_thread.values().sum();
         if pool_total > peak_pool_total {
             peak_pool_total = pool_total;
         }
 
-        if latest_jemalloc > 0 && pool_total > 0 && latest_jemalloc > pool_total {
-            let excess = latest_jemalloc - pool_total;
+        if source.is_some() && !pool_by_thread.is_empty() && latest_allocated > pool_total {
+            let excess = latest_allocated - pool_total;
             if excess > peak_excess {
                 peak_excess = excess;
             }
@@ -147,46 +178,60 @@ fn main() {
             {
                 violations.push(MemorySnapshot {
                     ts: event.ts,
-                    jemalloc: latest_jemalloc,
+                    allocated: latest_allocated,
                     pool_total,
                 });
             }
         }
     }
 
+    let Some(source) = source.map(|rank| ALLOCATED_COUNTERS[rank]) else {
+        eprintln!(
+            "No process-wide allocation counter found in the trace: expected one of {}. \
+             Build the native library with the `alloc-accounting` or `jemalloc` feature.",
+            ALLOCATED_COUNTERS.join(", ")
+        );
+        std::process::exit(1);
+    };
+
     // Print summary
     println!("=== Comet Trace Memory Analysis ===\n");
     println!("Counter events parsed: {counter_events}");
+    println!("Allocation counter:    {source}");
     println!("Threads with memory pools: {}", pool_by_thread.len());
-    println!("Peak jemalloc allocated:   {}", format_bytes(peak_jemalloc));
+    println!("Peak {source}:   {}", format_bytes(peak_allocated));
     println!(
         "Peak pool total:           {}",
         format_bytes(peak_pool_total)
     );
     println!(
-        "Peak excess (jemalloc - pool): {}",
+        "Peak excess ({source} - pool): {}",
         format_bytes(peak_excess)
     );
     println!();
 
-    if violations.is_empty() {
-        println!("OK: jemalloc never exceeded the total pool reservation.");
+    if pool_by_thread.is_empty() {
+        println!(
+            "No pool reservation samples in the trace, so there is nothing to compare against."
+        );
+    } else if violations.is_empty() {
+        println!("OK: {source} never exceeded the total pool reservation.");
     } else {
         println!(
-            "WARNING: jemalloc exceeded pool reservation at {} sampled points:\n",
+            "WARNING: {source} exceeded pool reservation at {} sampled points:\n",
             violations.len()
         );
         println!(
-            "{:>14}  {:>14}  {:>14}  {:>14}",
-            "Time (us)", "jemalloc", "pool_total", "excess"
+            "{:>14}  {:>18}  {:>14}  {:>14}",
+            "Time (us)", source, "pool_total", "excess"
         );
-        println!("{}", "-".repeat(62));
+        println!("{}", "-".repeat(66));
         for snap in &violations {
-            let excess = snap.jemalloc - snap.pool_total;
+            let excess = snap.allocated - snap.pool_total;
             println!(
-                "{:>14}  {:>14}  {:>14}  {:>14}",
+                "{:>14}  {:>18}  {:>14}  {:>14}",
                 snap.ts,
-                format_bytes(snap.jemalloc),
+                format_bytes(snap.allocated),
                 format_bytes(snap.pool_total),
                 format_bytes(excess),
             );

native/core/Cargo.toml

@@ -114,6 +114,12 @@ 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.
+alloc-accounting = []
+
 # exclude optional packages from cargo machete verifications
 [package.metadata.cargo-machete]
 ignored = ["hdfs-sys", "paste"]
@@ -123,6 +129,10 @@ name = "comet"
 # "rlib" is for benchmarking with criterion.
 crate-type = ["cdylib", "rlib"]
 
+[[bench]]
+name = "alloc_overhead"
+harness = false
+
 [[bench]]
 name = "array_element_append"
 harness = false