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

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

@@ -273,6 +273,20 @@ these events, and one without Comet on its classpath logs once that it dropped t
 `spark.eventLog.excludedPatterns=org.apache.comet.CometExecutorMemoryUsage` leaves them out of the
 event log.
 
+## Returning Freed Native Memory Between Stages
+
+On Linux, glibc's allocator keeps much of the memory that Comet's native operators free instead of
+returning it to the operating system, and it keeps it per thread. 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 the peak of one stage while the next stage allocates on top of it. The executor
+can then outgrow its container while the memory usage log shows little allocated.
+
+Setting `spark.comet.memory.trimNativeHeapWhenIdle.enabled=true` makes an executor call glibc's
+`malloc_trim` when its last running native plan finishes, which is typically between stages, and log
+at INFO level how long that took and its resident memory before and after. Set it when the
+application is submitted. It does nothing on other platforms, or in builds with the `jemalloc` or
+`mimalloc` feature.
+
 ## Batch Size
 
 Comet processes data in columnar batches. The batch size is controlled by `spark.comet.batchSize` (default

native/core/src/execution/jni_api.rs

@@ -125,6 +125,7 @@ use tokio::task::JoinHandle;
 use crate::execution::memory_pools::{
     create_memory_pool, overcommit, parse_memory_pool_config, MemoryPoolConfig, PlanMemoryPool,
 };
+use crate::execution::native_heap_trim;
 use crate::execution::operators::{ScanExec, ShuffleScanExec};
 use crate::execution::shuffle::{
     decode_remote_shuffle_batch, read_ipc_compressed, CompressionCodec, ShuffleWriterExec,
@@ -141,8 +142,9 @@ use crate::execution::spark_config::{
     SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY,
     COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXEC_SPILL_COMPRESSION_CODEC,
     COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE,
-    COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED,
-    COMET_STRICT_TIMESTAMP_LABELS, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES,
+    COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE, COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED,
+    COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_STRICT_TIMESTAMP_LABELS,
+    COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES,
 };
 use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID};
 use crate::parquet::parquet_exec::CometParquetScanOptions;
@@ -284,6 +286,15 @@ fn sum_reserved(pools: &[Arc<dyn MemoryPool>]) -> usize {
     pools.iter().map(|pool| pool.reserved()).sum()
 }
 
+/// The number of native plans registered on this executor, that is created and not yet released.
+fn live_plan_count() -> usize {
+    get_thread_memory_pools()
+        .lock()
+        .values()
+        .map(HashMap::len)
+        .sum()
+}
+
 /// Bytes reserved across `pools`, less the part that Spark has not granted them; see
 /// [`MemoryUsage::pools_reserved`].
 fn sum_reserved_less_overcommit(pools: &[Arc<dyn MemoryPool>]) -> usize {
@@ -597,6 +608,9 @@ struct ExecutionContext {
     pub output_stream: Option<Arc<dyn OutputStreamSink>>,
     /// Removes this context's tracing memory-pool entry on every exit path.
     memory_pool_registration: Option<ThreadMemoryPoolRegistration>,
+    /// Whether releasing this plan trims the native heap when no other plan is left running on
+    /// the executor; see [`crate::execution::native_heap_trim`].
+    trim_native_heap_when_idle: bool,
 }
 
 /// Accept serialized query plan and return the address of the native query plan.
@@ -728,6 +742,7 @@ struct PlanSettings {
     tracing_enabled: bool,
     max_temp_directory_size: u64,
     logging_memory_pool: bool,
+    trim_native_heap_when_idle: bool,
 }
 
 /// Deserializes the Spark configs of a plan and initializes the Tokio runtime from them. The
@@ -750,6 +765,7 @@ fn prepare_plan_settings(serialized_spark_configs: &[u8]) -> CometResult<PlanSet
         max_temp_directory_size: spark_config
             .get_u64(COMET_MAX_TEMP_DIRECTORY_SIZE, 100 * 1024 * 1024 * 1024),
         logging_memory_pool: spark_config.get_bool(COMET_DEBUG_MEMORY),
+        trim_native_heap_when_idle: spark_config.get_bool(COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE),
         spark_config,
     })
 }
@@ -909,6 +925,7 @@ fn create_execution_context(
         shuffle_partition_pusher: None,
         output_stream: None,
         memory_pool_registration,
+        trim_native_heap_when_idle: settings.trim_native_heap_when_idle,
     }))
 }
 
@@ -1671,6 +1688,7 @@ fn release_plan(
     // to other tasks. Dropping the plan aborts the tasks its operators spawned, and they give
     // back what they hold the next time they yield, so wait for them.
     let id = execution_context.id;
+    let trim_native_heap_when_idle = execution_context.trim_native_heap_when_idle;
     let plan_memory = Arc::clone(&execution_context.plan_memory);
     drop(execution_context);
     let held = plan_memory.wait_until_released(Instant::now() + PLAN_MEMORY_RELEASE_TIMEOUT);
@@ -1683,6 +1701,12 @@ fn release_plan(
         );
     }
 
+    // The plan has freed its memory. If no other plan is left running, hand what the allocator
+    // kept of it back to the operating system before the next stage allocates on top of it.
+    if trim_native_heap_when_idle {
+        native_heap_trim::trim_if_idle(live_plan_count());
+    }
+
     producer_stopped.and(metrics_flushed)
 }
 

native/core/src/execution/mod.rs

@@ -24,6 +24,7 @@ pub mod expressions;
 pub mod jni_api;
 pub(crate) mod merge_as_partial;
 pub(crate) mod metrics;
+pub(crate) mod native_heap_trim;
 pub mod operators;
 pub(crate) mod planner;
 #[cfg(feature = "python-udf")]

native/core/src/execution/native_heap_trim.rs

@@ -0,0 +1,111 @@
+// 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.
+
+//! Hands the memory that the C allocator keeps cached back to the operating system once an
+//! executor has no Comet native plan left running.
+//!
+//! glibc keeps much of the memory that native operators free in its per-thread arenas instead of
+//! returning it, and Comet's operators free a lot: every spill drops the batches it buffered, which
+//! are Arrow buffers of tens of KiB to a few MiB. A stage's plans also run on different threads
+//! than the next stage's (a plan with no JVM input runs on Tokio workers, one that reads shuffle
+//! on the Spark task thread), so the next stage does not reuse what the previous one left in the
+//! other threads' arenas. The executor's resident memory then stays near the previous stage's peak
+//! and the next stage's allocations add to it, which can exceed the container even though Comet's
+//! memory pools and allocation counter stay within their limits. Spark's own off-heap pages are
+//! large enough that glibc maps and unmaps each of them directly, so they do not build up this way.
+//!
+//! `malloc_trim` returns the free pages in every arena. It runs only when the last plan on the
+//! executor is released, typically between stages, so it does not hold the arena locks of
+//! operators that are still running.
+
+/// Trims the native heap if `live_plans`, the number of native plans still registered on this
+/// executor, is zero and this build can trim.
+pub(crate) fn trim_if_idle(live_plans: usize) {
+    if should_trim(
+        platform::SUPPORTED && crate::ALLOCATOR_BACKEND == "system",
+        live_plans,
+    ) {
+        platform::trim_and_log();
+    }
+}
+
+/// Whether releasing a plan should trim: trimming is `supported` and no other plan is running.
+/// `malloc_trim` is glibc's, so only a Linux build that allocates through glibc, which means the
+/// system allocator rather than the `jemalloc` or `mimalloc` feature, supports it.
+fn should_trim(supported: bool, live_plans: usize) -> bool {
+    supported && live_plans == 0
+}
+
+#[cfg(all(target_os = "linux", target_env = "gnu"))]
+mod platform {
+    use log::info;
+    use std::time::Instant;
+
+    pub(super) const SUPPORTED: bool = true;
+
+    pub(super) fn trim_and_log() {
+        extern "C" {
+            fn malloc_trim(pad: usize) -> std::os::raw::c_int;
+        }
+        let before = anonymous_rss_mib();
+        let start = Instant::now();
+        // SAFETY: glibc's malloc_trim takes the lock of each arena it trims and can be called from
+        // any thread. A pad of 0 keeps nothing beyond what the allocator needs at the top of the
+        // main heap.
+        unsafe {
+            malloc_trim(0);
+        }
+        let elapsed_ms = start.elapsed().as_millis();
+        match (before, anonymous_rss_mib()) {
+            (Some(before), Some(after)) => info!(
+                "Trimmed the native heap after the last Comet native plan on this executor \
+                 finished, in {elapsed_ms} ms: anonymous resident memory {before:.1} MiB before, \
+                 {after:.1} MiB after"
+            ),
+            _ => info!(
+                "Trimmed the native heap after the last Comet native plan on this executor \
+                 finished, in {elapsed_ms} ms"
+            ),
+        }
+    }
+
+    /// The process's anonymous resident memory, which the container's memory limit counts, or
+    /// None when `/proc` cannot be read.
+    fn anonymous_rss_mib() -> Option<f64> {
+        let status = procfs::process::Process::myself().ok()?.status().ok()?;
+        status.rssanon.map(|kib| kib as f64 / 1024.0)
+    }
+}
+
+#[cfg(not(all(target_os = "linux", target_env = "gnu")))]
+mod platform {
+    pub(super) const SUPPORTED: bool = false;
+
+    pub(super) fn trim_and_log() {}
+}
+
+#[cfg(test)]
+mod tests {
+    use super::should_trim;
+
+    #[test]
+    fn trims_only_when_supported_and_no_plan_is_left() {
+        assert!(should_trim(true, 0));
+        assert!(!should_trim(true, 1), "another plan is still running");
+        assert!(!should_trim(false, 0), "the build cannot trim");
+    }
+}

native/core/src/execution/spark_config.rs

@@ -22,6 +22,8 @@ pub(crate) const COMET_DEBUG_ENABLED: &str = "spark.comet.debug.enabled";
 pub(crate) const COMET_EXPLAIN_NATIVE_ENABLED: &str = "spark.comet.explain.native.enabled";
 pub(crate) const COMET_MAX_TEMP_DIRECTORY_SIZE: &str = "spark.comet.maxTempDirectorySize";
 pub(crate) const COMET_DEBUG_MEMORY: &str = "spark.comet.debug.memory";
+pub(crate) const COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE: &str =
+    "spark.comet.memory.trimNativeHeapWhenIdle.enabled";
 pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str =
     "spark.comet.parquet.rowFilterPushdown.enabled";
 pub(crate) const COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED: &str =

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -583,6 +583,26 @@ object CometConf extends ShimCometConf {
     .checkValue(_ >= 0, "The memory usage log interval must not be negative")
     .createWithDefault(TimeUnit.SECONDS.toMillis(10))
 
+  val COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE: ConfigEntry[Boolean] =
+    conf("spark.comet.memory.trimNativeHeapWhenIdle.enabled")
+      .category(CATEGORY_TUNING)
+      .doc(
+        "Whether an executor hands the memory that the C allocator keeps cached back to the " +
+          "operating system when its last running Comet native plan finishes, which is " +
+          "typically between stages. glibc keeps much of the memory that native operators free " +
+          "instead of returning it, so the executor's resident memory can stay near the peak " +
+          "of one stage while the next stage allocates on top of it, and the container can run " +
+          "out of memory while Comet's memory pools stay within their limits. With this " +
+          "enabled, Comet calls glibc's malloc_trim when no native plan is left running on the " +
+          "executor, and logs at INFO level how long it took and the resident memory before " +
+          "and after. Only Linux builds that use glibc and the default system allocator trim. " +
+          "Elsewhere, and in builds with the jemalloc or mimalloc feature, this does nothing. " +
+          "This is an executor setting, read from the application's configuration, so it must " +
+          "be set when the application is submitted. An invalid value disables trimming with " +
+          s"a warning. $TUNING_GUIDE.")
+      .booleanConf
+      .createWithDefault(false)
+
   val COMET_SHUFFLE_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.shuffle.enabled")
       .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.enabled")

spark/src/main/scala/org/apache/comet/CometExecIterator.scala

@@ -553,6 +553,32 @@ object CometExecIterator extends Logging {
         }
     }
 
+  private val invalidTrimSettingWarned = new AtomicBoolean(false)
+
+  /**
+   * Whether the executor trims its native heap when its last Comet native plan finishes, for the
+   * executor's configured value, if any. It is read for every native plan, so a value that does
+   * not parse disables trimming with a warning, logged once, rather than failing every Comet
+   * task.
+   */
+  def trimNativeHeapWhenIdle(configured: Option[String]): Boolean =
+    configured match {
+      case None => COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE.defaultValue.get
+      case Some(value) =>
+        try {
+          COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE.valueConverter(value)
+        } catch {
+          case NonFatal(e) =>
+            if (invalidTrimSettingWarned.compareAndSet(false, true)) {
+              logWarning(
+                s"Not trimming the native heap: invalid value '$value' for " +
+                  s"${COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE.key}. Expected true or false. " +
+                  e.getMessage)
+            }
+            false
+        }
+    }
+
   /**
    * Whether the cluster manager for `master` runs each executor in a container whose memory it
    * sizes from the memory overhead settings: YARN and Kubernetes. Local mode, local-cluster
@@ -891,6 +917,13 @@ object CometExecIterator extends Logging {
     val executorCores = numDriverOrExecutorCores(SparkEnv.get.conf)
     builder.putEntries("spark.executor.cores", executorCores.toString)
 
+    // Trimming acts on the whole executor, so it is read from the executor's configuration rather
+    // than the session's, and overrides a value set in the session.
+    val trimKey = COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE.key
+    builder.putEntries(
+      trimKey,
+      trimNativeHeapWhenIdle(SparkEnv.get.conf.getOption(trimKey)).toString)
+
     // Any Comet config that the native side reads must be added here manually, resolved.
     // `cometSqlConfs` only carries values that were explicitly set, exactly as they were
     // written, so defaults from `createWithDefault(...)` would otherwise not cross JNI, and

spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala

@@ -116,6 +116,14 @@ class CometExecSuite extends CometTestBase {
     withSQLConf(CometConf.COMET_STRICT_TIMESTAMP_LABELS.key -> "TRUE") {
       assert(entries.get(CometConf.COMET_STRICT_TIMESTAMP_LABELS.key) == "true")
     }
+
+    // Trimming the native heap is an executor setting: it crosses resolved from the executor's
+    // configuration, and a value set in the session does not change it.
+    val trimKey = CometConf.COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE.key
+    assert(defaults.get(trimKey) == "false")
+    withSQLConf(trimKey -> "true") {
+      assert(entries.get(trimKey) == "false")
+    }
   }
 
   test("native sort spill files are compressed with the default spill codec") {

spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala

@@ -845,6 +845,17 @@ class CometExecIteratorLifecycleSuite extends CometTestBase {
     assert(memoryUsageLogInterval(Some("10 seconds please")) == 0L)
   }
 
+  test("trimming the native heap is off by default and on a value it cannot use") {
+    import CometExecIterator.trimNativeHeapWhenIdle
+    assert(!trimNativeHeapWhenIdle(None))
+    assert(trimNativeHeapWhenIdle(Some("true")))
+    assert(trimNativeHeapWhenIdle(Some(" TRUE ")))
+    assert(!trimNativeHeapWhenIdle(Some("false")))
+    // Values that would otherwise throw from the creation of every native plan, failing its task.
+    assert(!trimNativeHeapWhenIdle(Some("yes")))
+    assert(!trimNativeHeapWhenIdle(Some("")))
+  }
+
   test("the native runtime gets a worker thread for every core the executor runs tasks on") {
     import CometExecIterator.numDriverOrExecutorCores
     def cores(master: String, executorCores: Option[String] = None): Option[Int] =