Skip to content

feat: trim the native heap when an executor's last native plan finishes - #7304

Draft
comphead wants to merge 1 commit into
apache:mainfrom
comphead:idle-native-heap-trim
Draft

comphead wants to merge 1 commit into
apache:mainfrom
comphead:idle-native-heap-trim

Conversation

@comphead

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

None. Related to #7108, which stays open: it asks whether release builds should default to a different allocator, and this PR adds an opt-in mitigation for the default glibc build instead.

Rationale for this change

An executor can be OOMKilled by its container while Comet's memory pools and allocation counter stay well within their limits. Much of the difference is native memory that has already been freed: glibc keeps a large part of what Comet's operators free in its per-thread arenas instead of returning it to the operating system. #7108 measured this on Linux. Live native memory dropped to about 2 MiB between queries while RSS stayed near 12 GiB, and one malloc_trim released 5 to 6 GiB.

It gets worse between stages:

  • Comet's operators free a lot of memory. Each sort spill and each shuffle writer spill drops the batches it buffered, which are Arrow buffers of tens of KiB to a few MiB. glibc serves those from its arenas rather than with mmap.
  • One stage's native plans run on different threads than the next stage's. A plan with no JVM input runs on Tokio workers (BatchProducer::spawn), and a plan that reads shuffle runs on the Spark task thread. The next stage does not reuse what the previous one left in the other threads' arenas, so its allocations add to the retained memory.

Spark's own off-heap memory does not build up this way, because its pages are 64 MiB and glibc maps and unmaps each of them directly.

I saw this pattern on a wide group-by (87 sort keys, sort-based aggregation) that Spark finishes and Comet does not: every result-stage wave was OOMKilled. A local reproduction at the same per-task size shows the mechanism. Comet's result stage alone needs about what Spark's does (9.1 vs 8.9 GiB), but right after the map stage it reaches 16.2 GiB, because about 7 GiB that the map stage freed is still resident when it starts.

What changes are included in this PR?

  • A new setting, spark.comet.memory.trimNativeHeapWhenIdle.enabled, false by default. When it is set, releasing the last native plan running on an executor calls glibc's malloc_trim(0) and logs at INFO level how long that took and the process's anonymous resident memory before and after. It only runs when no plan is left, which is typically between stages, so it does not take the arena locks of running operators.
  • It only acts on Linux builds that use glibc and the default system allocator. On other platforms, and in builds with the jemalloc or mimalloc feature, it does nothing.
  • It is an executor setting. CometExecIterator.serializeCometSQLConfs reads it from the executor's SparkConf rather than the session's, so it applies however the job is launched (see Map stages of RDD-level jobs run Comet with default settings instead of the session's #7301). An invalid value logs a warning once and leaves trimming off.
  • The code is a new native_heap_trim.rs module and a hook at the end of release_plan in jni_api.rs, with a small live_plan_count() helper. The memory tuning guide gets a section on the setting.

How are these changes tested?

  • Unit tests: native_heap_trim::tests checks that it trims only when the build supports it and no plan is left. A CometExecIteratorLifecycleSuite test covers parsing the setting. CometExecSuite "SQLConf serde resolves the configs that native code parses" checks that the value always crosses JNI resolved, from the executor's configuration. These have not run locally yet.
  • malloc_trim does not exist on macOS, where I work, so I type-checked and linted the Linux path for aarch64-unknown-linux-gnu (clippy with -D warnings). It has not run on Linux yet. A Linux run with the setting on is the next step, and the INFO line will show its effect.
  • To exercise the hook itself, I ran the reproduction with a jemalloc build in which the hook purged jemalloc's arenas instead of calling malloc_trim (a test-only change, not in this PR), and with jemalloc keeping freed memory as glibc does. With the setting off, 4 of 4 runs exceeded the memory limit in the result stage. With it on, 4 of 4 passed. The hook ran twice per job, once after the map stage's last task and once at the end, and each purge took 0.3 to 0.7 s.

glibc keeps much of the memory that Comet's native operators free in its
per-thread arenas instead of returning it to the operating system. Spills
free a lot of memory, and one stage's native plans run on different
threads than the next stage's, so an executor's resident memory can stay
near one stage's peak while the next stage allocates on top of it. The
container can then run out of memory while Comet's memory pools stay
within their limits.

Add spark.comet.memory.trimNativeHeapWhenIdle.enabled, off by default.
When it is set, releasing the last native plan running on an executor
calls glibc's malloc_trim(0) and logs how long that took and the
anonymous resident memory before and after. It only acts on Linux builds
that use glibc and the default system allocator, and it is read from the
executor's configuration.
@github-actions github-actions Bot added the enhancement New feature or request label Oct 28, 2026
@andygrove

Copy link
Copy Markdown
Member

Thanks for digging into this. The analysis of what glibc retains between stages matches what #7108 measured, and gating the trim on the plan registry so it only runs when no operator can be holding an arena lock is a nice idea. Before we add a config for this, though, I want to step back and ask whether this is the right layer to fix it at.

The JDK already does this. -XX:TrimNativeHeapInterval=<ms> (JDK-8293114, a product flag since 17.0.12 and 21.0.3) calls malloc_trim on a timer and suspends itself around GC and other phases where a trim would hurt. Comet requires JDK 17+, so every supported JDK has it. It needs no Comet code, no config, and no release. It trims the whole process C heap rather than only Comet's plans, and it keeps working on executors that never go idle, such as a thrift server or concurrent jobs, where this PR's trigger never fires. The trade-off is the one this PR is built around: a timer can fire mid-stage and lock arenas while operators are running. I don't think anyone has measured that on a Comet workload yet.

The bigger question is the one #7108 asks: should the release build use a different allocator? People in #1648 already run Comet with jemalloc preloaded because it returns memory more predictably, and the jemalloc feature already carries the disable_initial_exec_tls workaround for being loaded by the JVM. If the release build moves to jemalloc, the glibc arena problem goes away and this hook becomes a config that does nothing. I would rather make that decision first than add a glibc-specific mitigation that it would obsolete.

So a few asks:

  1. Could you run the Executor is OOM-killed while Comet's pool and allocation counter stay under the limit: glibc keeps several GiB of freed native memory #7108 reproduction three ways: this PR, -XX:TrimNativeHeapInterval at a sensible interval, and a jemalloc build? Peak RSS, whether it passes, and total time for each would tell us which of these is worth pursuing.
  2. Could you add the JDK flag to the tuning guide in a small separate PR regardless of the outcome? It is the thing users can do today without upgrading Comet.
  3. Let's hold this PR until we have that data. If idle-triggered trimming beats the timer by enough to justify owning the code and the config, I'm happy to review it properly. In that case the hook would also need to run on an actual glibc build before merge, since right now the jemalloc stand-in is the only end-to-end evidence and CI never enables the setting.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants