You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@auron.apache.org by ri...@apache.org on 2026/10/03 03:01:04 UTC
(auron) branch master updated: [AURON #2554] Support shared native memory manager and spill coordinator (#2556)
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git
The following commit(s) were added to refs/heads/master by this push:
new e60cb7e57 [AURON #2554] Support shared native memory manager and spill coordinator (#2556)
e60cb7e57 is described below
commit e60cb7e57f2aa140bdc9b35370d5d4e5680fa910
Author: Yizhong Zhang <zy...@gmail.com>
AuthorDate: Sat Oct 3 11:00:58 2026 +0800
[AURON #2554] Support shared native memory manager and spill coordinator (#2556)
# Which issue does this PR close?
Closes #2554
# Rationale for this change
Native memory accounting and spill scheduling are coupled to the host
memory manager. A shared framework allows memory domains to provide
their own pressure measurements and spill preparation while reusing
consumer accounting and concurrency control.
# What changes are included in this PR?
- Introduce a memory-domain interface, a generic memory manager, and a
spill coordinator.
- Add watermark-based spill selection and bounded spill concurrency.
- Handle spill completion, failure, and cancellation safely, prevent
recursive spills, and wake waiting consumers when memory is released.
- Add regression tests for accounting, spill scheduling, and lifecycle
handling.
# Are there any user-facing changes?
No.
# How was this patch tested?
UTs.
Co-authored-by: zhangyizhong <zh...@kuaishou.com>
---
Cargo.lock | 1 +
native-engine/auron-memmgr/Cargo.toml | 1 +
native-engine/auron-memmgr/src/lib.rs | 2 +
.../auron-memmgr/src/runtime_mem_manager.rs | 265 ++++++++++++
.../auron-memmgr/src/runtime_spill_coordinator.rs | 465 +++++++++++++++++++++
.../auron-memmgr/tests/runtime_spill_tests.rs | 441 +++++++++++++++++++
6 files changed, 1175 insertions(+)
diff --git a/Cargo.lock b/Cargo.lock
index 60f61d89f..8e54e46b4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -377,6 +377,7 @@ dependencies = [
"bytesize",
"datafusion",
"datafusion-ext-commons",
+ "futures",
"jni",
"log",
"once_cell",
diff --git a/native-engine/auron-memmgr/Cargo.toml b/native-engine/auron-memmgr/Cargo.toml
index 893daa1db..b84896029 100644
--- a/native-engine/auron-memmgr/Cargo.toml
+++ b/native-engine/auron-memmgr/Cargo.toml
@@ -31,6 +31,7 @@ datafusion-ext-commons = { workspace = true }
async-trait = { workspace = true }
bytesize = { workspace = true }
+futures = { workspace = true }
jni = { workspace = true }
log = { workspace = true }
once_cell = { workspace = true }
diff --git a/native-engine/auron-memmgr/src/lib.rs b/native-engine/auron-memmgr/src/lib.rs
index 32c89971d..60ff87314 100644
--- a/native-engine/auron-memmgr/src/lib.rs
+++ b/native-engine/auron-memmgr/src/lib.rs
@@ -16,6 +16,8 @@
#![feature(get_mut_unchecked)]
pub mod metrics;
+pub mod runtime_mem_manager;
+pub mod runtime_spill_coordinator;
pub mod spill;
use std::{
diff --git a/native-engine/auron-memmgr/src/runtime_mem_manager.rs b/native-engine/auron-memmgr/src/runtime_mem_manager.rs
new file mode 100644
index 000000000..4325f7bc8
--- /dev/null
+++ b/native-engine/auron-memmgr/src/runtime_mem_manager.rs
@@ -0,0 +1,265 @@
+// 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.
+
+use std::{
+ marker::PhantomData,
+ sync::{Arc, Once},
+};
+
+use async_trait::async_trait;
+use bytesize::ByteSize;
+use once_cell::sync::OnceCell;
+use parking_lot::{Condvar, Mutex};
+
+use crate::runtime_spill_coordinator::RuntimeSpillCoordinator;
+
+/// Domain-specific memory accounting and spill policy.
+#[async_trait]
+pub trait MemDomain: Send + Sync + Sized + 'static {
+ fn name() -> &'static str;
+
+ /// Minimum reservation for automatic spill checks.
+ fn min_trigger_size() -> usize {
+ 0
+ }
+
+ /// Ratios satisfying `0 < L1 <= L2 <= 1`.
+ fn watermark_config() -> (f64, f64);
+
+ /// Maximum active spills at or above L2.
+ fn num_max_spilling_consumers() -> usize;
+
+ /// Called before each spill; may await an async precondition.
+ async fn on_before_spill(_consumer_name: &str) {}
+
+ /// Called once after the manager singleton is created.
+ fn on_manager_initialized() {}
+
+ fn manager_cell() -> &'static OnceCell<Arc<RuntimeMemManager<Self>>>;
+
+ fn coordinator_cell() -> &'static OnceCell<RuntimeSpillCoordinator<Self>>;
+}
+
+pub struct ConsumerInfo<D: MemDomain> {
+ pub name: String,
+ pub status: Mutex<ConsumerStatus>,
+ _marker: PhantomData<D>,
+}
+
+impl<D: MemDomain> std::fmt::Debug for ConsumerInfo<D> {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("ConsumerInfo")
+ .field("name", &self.name)
+ .field("status", &self.status)
+ .finish()
+ }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct ConsumerStatus {
+ pub mem_used: usize,
+ pub spillable: bool,
+}
+
+pub struct ManagerStatus<D: MemDomain> {
+ pub num_consumers: usize,
+ pub total_used: usize,
+ pub num_spillables: usize,
+ pub mem_spillables: usize,
+ _marker: PhantomData<D>,
+}
+
+impl<D: MemDomain> Default for ManagerStatus<D> {
+ fn default() -> Self {
+ Self {
+ num_consumers: 0,
+ total_used: 0,
+ num_spillables: 0,
+ mem_spillables: 0,
+ _marker: PhantomData,
+ }
+ }
+}
+
+impl<D: MemDomain> Clone for ManagerStatus<D> {
+ fn clone(&self) -> Self {
+ *self
+ }
+}
+
+impl<D: MemDomain> Copy for ManagerStatus<D> {}
+
+impl<D: MemDomain> ManagerStatus<D> {
+ pub fn update_total_used_with_diff(&mut self, diff_used: isize) -> usize {
+ assert!(self.total_used as isize + diff_used >= 0);
+
+ let new_used = (self.total_used as isize + diff_used) as usize;
+ let old_used = std::mem::replace(&mut self.total_used, new_used);
+
+ if new_used < old_used {
+ RuntimeMemManager::<D>::get().cv.notify_all();
+ }
+ new_used
+ }
+}
+
+pub struct RuntimeMemManager<D: MemDomain> {
+ pub total: usize,
+ pub consumers: Mutex<Vec<Arc<ConsumerInfo<D>>>>,
+ pub status: Mutex<ManagerStatus<D>>,
+ pub cv: Condvar,
+ initialized_hook: Once,
+}
+
+impl<D: MemDomain> RuntimeMemManager<D> {
+ pub fn init_with_total(total: usize) {
+ D::manager_cell().get_or_init(|| {
+ log::info!(
+ "{} manager initialized with total memory: {}",
+ D::name(),
+ ByteSize(total as u64),
+ );
+ Arc::new(RuntimeMemManager {
+ total,
+ consumers: Mutex::default(),
+ status: Mutex::new(ManagerStatus::default()),
+ cv: Condvar::default(),
+ initialized_hook: Once::new(),
+ })
+ });
+ Self::get()
+ .initialized_hook
+ .call_once(D::on_manager_initialized);
+ }
+
+ pub fn initialized() -> bool {
+ D::manager_cell().get().is_some()
+ }
+
+ pub fn get() -> &'static RuntimeMemManager<D> {
+ D::manager_cell()
+ .get()
+ .expect("memory domain manager not initialized")
+ }
+
+ pub fn num_consumers(&self) -> usize {
+ self.consumers.lock().len()
+ }
+
+ pub fn total_used(&self) -> usize {
+ self.status.lock().total_used
+ }
+
+ pub fn mem_used_percent(&self) -> f64 {
+ let used = self.total_used();
+ if self.total == 0 {
+ if used == 0 { 0.0 } else { f64::INFINITY }
+ } else {
+ used as f64 / self.total as f64
+ }
+ }
+
+ pub fn register_consumer_info(name: String, spillable: bool) -> Arc<ConsumerInfo<D>> {
+ let consumer_info = Arc::new(ConsumerInfo {
+ name,
+ status: Mutex::new(ConsumerStatus {
+ mem_used: 0,
+ spillable,
+ }),
+ _marker: PhantomData,
+ });
+ log::info!(
+ "{} manager registering consumer: {}",
+ D::name(),
+ consumer_info.name
+ );
+
+ let mm = Self::get();
+ let mut consumers = mm.consumers.lock();
+ let mut status = mm.status.lock();
+ consumers.push(consumer_info.clone());
+ status.num_consumers += 1;
+ if spillable {
+ status.num_spillables += 1;
+ }
+ consumer_info
+ }
+
+ pub fn deregister_consumer_info(name: &str, consumer_info: &Arc<ConsumerInfo<D>>) {
+ let mm = Self::get();
+ let mut consumers = mm.consumers.lock();
+ let mut status = mm.status.lock();
+ let consumer_status = consumer_info.status.lock();
+
+ assert!(status.total_used >= consumer_status.mem_used);
+ status.num_consumers -= 1;
+ status.update_total_used_with_diff(-(consumer_status.mem_used as isize));
+
+ if consumer_status.spillable {
+ assert!(status.mem_spillables >= consumer_status.mem_used);
+ status.num_spillables -= 1;
+ status.mem_spillables -= consumer_status.mem_used;
+ }
+ drop(consumer_status);
+
+ for i in 0..consumers.len() {
+ if Arc::ptr_eq(&consumers[i], consumer_info) {
+ log::info!("{} manager deregistered consumer: {name}", D::name());
+ consumers.swap_remove(i);
+ drop(status);
+ drop(consumers);
+ if let Some(coordinator) = D::coordinator_cell().get() {
+ coordinator.wake_waiters();
+ }
+ return;
+ }
+ }
+ unreachable!("deregistering non-registered {} memory consumer", D::name());
+ }
+
+ pub fn set_consumer_spillable(consumer_info: &Arc<ConsumerInfo<D>>, spillable: bool) {
+ let mm = Self::get();
+ let mut status = mm.status.lock();
+ let mut consumer = consumer_info.status.lock();
+ if consumer.spillable != spillable {
+ if spillable {
+ status.num_spillables += 1;
+ status.mem_spillables += consumer.mem_used;
+ } else {
+ status.num_spillables -= 1;
+ status.mem_spillables -= consumer.mem_used;
+ }
+ consumer.spillable = spillable;
+ }
+ drop(consumer);
+ drop(status);
+ if let Some(coordinator) = D::coordinator_cell().get() {
+ coordinator.wake_waiters();
+ }
+ }
+
+ pub fn log_consumers(&self) {
+ for consumer in &*self.consumers.lock() {
+ let consumer_status = consumer.status.lock();
+ log::info!(
+ "* {} consumer: {}, spillable: {}, mem_used: {}",
+ D::name(),
+ consumer.name,
+ consumer_status.spillable,
+ ByteSize(consumer_status.mem_used as u64),
+ );
+ }
+ }
+}
diff --git a/native-engine/auron-memmgr/src/runtime_spill_coordinator.rs b/native-engine/auron-memmgr/src/runtime_spill_coordinator.rs
new file mode 100644
index 000000000..41a384434
--- /dev/null
+++ b/native-engine/auron-memmgr/src/runtime_spill_coordinator.rs
@@ -0,0 +1,465 @@
+// 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.
+
+//! Automatic spills: none below L1, largest consumer between L1 and L2,
+//! bounded concurrency at or above L2.
+
+use std::{
+ future::Future,
+ sync::{Arc, Weak},
+};
+
+use bytesize::ByteSize;
+use datafusion::common::Result;
+use futures::channel::oneshot;
+use parking_lot::Mutex;
+
+use crate::runtime_mem_manager::{ConsumerInfo, ConsumerStatus, MemDomain, RuntimeMemManager};
+
+pub enum SpillAction {
+ DoSpill,
+ DoNotSpill,
+ WaitAndRetry(oneshot::Receiver<()>),
+}
+
+pub async fn wait_for_spill(receiver: oneshot::Receiver<()>) -> bool {
+ receiver.await.is_ok()
+}
+
+struct RuntimeSpillCoordinatorState<D: MemDomain> {
+ spilling_consumers: Vec<Weak<ConsumerInfo<D>>>,
+ active_consumers: Vec<Weak<ConsumerInfo<D>>>,
+ pending_notifiers: Vec<oneshot::Sender<()>>,
+}
+
+impl<D: MemDomain> Default for RuntimeSpillCoordinatorState<D> {
+ fn default() -> Self {
+ Self {
+ spilling_consumers: Vec::new(),
+ active_consumers: Vec::new(),
+ pending_notifiers: Vec::new(),
+ }
+ }
+}
+
+pub struct RuntimeSpillCoordinator<D: MemDomain> {
+ state: Mutex<RuntimeSpillCoordinatorState<D>>,
+ watermark_l1: f64,
+ watermark_l2: f64,
+ num_max_spilling_consumers: usize,
+}
+
+impl<D: MemDomain + RuntimeWatermarkEvaluator<D>> RuntimeSpillCoordinator<D> {
+ pub fn get() -> &'static RuntimeSpillCoordinator<D> {
+ D::coordinator_cell().get_or_init(|| {
+ let (watermark_l1, watermark_l2) = D::watermark_config();
+ let num_max_spilling_consumers = D::num_max_spilling_consumers();
+
+ assert!(
+ watermark_l1.is_finite()
+ && watermark_l2.is_finite()
+ && 0.0 < watermark_l1
+ && watermark_l1 <= watermark_l2
+ && watermark_l2 <= 1.0,
+ "spill watermarks must satisfy 0 < L1 <= L2 <= 1"
+ );
+ assert!(
+ num_max_spilling_consumers > 0,
+ "spill concurrency must be positive"
+ );
+
+ log::info!(
+ "{} spill coordinator initialized: L1: {watermark_l1}, L2: {watermark_l2}, \
+ max spilling consumers: {num_max_spilling_consumers}",
+ D::name(),
+ );
+ RuntimeSpillCoordinator {
+ state: Mutex::new(RuntimeSpillCoordinatorState::default()),
+ watermark_l1,
+ watermark_l2,
+ num_max_spilling_consumers,
+ }
+ })
+ }
+
+ /// Lock order: `coordinator.state -> {manager.consumers, manager.status}
+ /// -> consumer.status`.
+ pub fn update_consumer_mem_used(&self, consumer_info: &Arc<ConsumerInfo<D>>) -> SpillAction {
+ let mut state = self.state.lock();
+
+ state.spilling_consumers.retain(|elected| {
+ elected.upgrade().is_some_and(|info| {
+ let status = info.status.lock();
+ status.spillable && status.mem_used > 0
+ })
+ });
+ if state
+ .active_consumers
+ .iter()
+ .any(|active| weak_ptr_eq(active, consumer_info))
+ {
+ return SpillAction::DoNotSpill;
+ }
+ let status = *consumer_info.status.lock();
+ if !status.spillable || status.mem_used == 0 {
+ return SpillAction::DoNotSpill;
+ }
+ let watermark = <D as RuntimeWatermarkEvaluator<D>>::watermark();
+ if watermark < self.watermark_l1 {
+ state.spilling_consumers.clear();
+ return SpillAction::DoNotSpill;
+ }
+ if watermark < self.watermark_l2 {
+ if !state.active_consumers.is_empty() {
+ return SpillAction::DoNotSpill;
+ }
+ // Re-elect because the previous consumer may have shrunk.
+ state.spilling_consumers.clear();
+ elect_largest_spillable_consumer::<D>(&mut state.spilling_consumers);
+ return if state
+ .spilling_consumers
+ .iter()
+ .any(|elected| weak_ptr_eq(elected, consumer_info))
+ {
+ SpillAction::DoSpill
+ } else {
+ SpillAction::DoNotSpill
+ };
+ }
+ if state.active_consumers.len() >= self.num_max_spilling_consumers {
+ return Self::park_reporter(&mut state);
+ }
+ // Spill the reporter: an idle consumer cannot be triggered here.
+ SpillAction::DoSpill
+ }
+
+ fn park_reporter(state: &mut RuntimeSpillCoordinatorState<D>) -> SpillAction {
+ state
+ .pending_notifiers
+ .retain(|sender| !sender.is_canceled());
+ let (sender, receiver) = oneshot::channel();
+ state.pending_notifiers.push(sender);
+ SpillAction::WaitAndRetry(receiver)
+ }
+
+ /// Forces a spill with a slot guard that wakes waiters on completion or
+ /// cancellation.
+ pub async fn spill_in_action<F, Fut>(
+ &self,
+ consumer_name: &str,
+ consumer_info: &Arc<ConsumerInfo<D>>,
+ do_spill: F,
+ ) -> Result<()>
+ where
+ F: FnOnce() -> Fut,
+ Fut: Future<Output = Result<()>>,
+ {
+ self.spill_with_mode(consumer_name, consumer_info, do_spill, true)
+ .await
+ }
+
+ async fn spill_with_mode<F, Fut>(
+ &self,
+ consumer_name: &str,
+ consumer_info: &Arc<ConsumerInfo<D>>,
+ do_spill: F,
+ forced: bool,
+ ) -> Result<()>
+ where
+ F: FnOnce() -> Fut,
+ Fut: Future<Output = Result<()>>,
+ {
+ // Reserve the slot under the same lock used for admission.
+ let _guard = loop {
+ let receiver = {
+ let mut state = self.state.lock();
+ let status = *consumer_info.status.lock();
+ if !status.spillable
+ || (!forced && status.mem_used == 0)
+ || state
+ .active_consumers
+ .iter()
+ .any(|active| weak_ptr_eq(active, consumer_info))
+ {
+ return Ok(());
+ }
+ let watermark = <D as RuntimeWatermarkEvaluator<D>>::watermark();
+ if !forced && watermark < self.watermark_l1 {
+ return Ok(());
+ }
+ let limit = if watermark < self.watermark_l2 {
+ 1
+ } else {
+ self.num_max_spilling_consumers
+ };
+ if state.active_consumers.len() < limit {
+ state.active_consumers.push(Arc::downgrade(consumer_info));
+ break SpillInActionGuard {
+ coordinator: self,
+ consumer: Arc::downgrade(consumer_info),
+ };
+ }
+ match Self::park_reporter(&mut state) {
+ SpillAction::WaitAndRetry(receiver) => receiver,
+ _ => unreachable!("parking must return a waiter"),
+ }
+ };
+ wait_for_spill(receiver).await;
+ };
+
+ if !consumer_info.status.lock().spillable {
+ log::info!(
+ "{} spill coordinator: {} was pinned before its spill started, skipping",
+ D::name(),
+ consumer_info.name,
+ );
+ return Ok(());
+ }
+
+ let mem_used = consumer_info.status.lock().mem_used;
+ let unspillable = {
+ let mm_status = *RuntimeMemManager::<D>::get().status.lock();
+ mm_status
+ .total_used
+ .saturating_sub(mm_status.mem_spillables)
+ };
+ D::on_before_spill(consumer_name).await;
+ let status = *consumer_info.status.lock();
+ if !status.spillable
+ || (!forced
+ && (status.mem_used == 0
+ || <D as RuntimeWatermarkEvaluator<D>>::watermark() < self.watermark_l1))
+ {
+ return Ok(());
+ }
+ log::info!(
+ "{} manager spilling {consumer_name} (consumer: {}), unspillable: {}, \
+ logical watermark: {:.3}, physical watermark: {:.3}",
+ D::name(),
+ ByteSize(mem_used as u64),
+ ByteSize(unspillable as u64),
+ <D as RuntimeWatermarkEvaluator<D>>::logical_watermark(),
+ <D as RuntimeWatermarkEvaluator<D>>::physical_watermark(),
+ );
+ do_spill().await
+ }
+
+ pub fn remove_consumer(&self, consumer_info: &Arc<ConsumerInfo<D>>) {
+ let mut state = self.state.lock();
+ state
+ .spilling_consumers
+ .retain(|elected| !weak_ptr_eq(elected, consumer_info));
+ }
+}
+
+impl<D: MemDomain> RuntimeSpillCoordinator<D> {
+ pub(crate) fn wake_waiters(&self) {
+ let notifiers = std::mem::take(&mut self.state.lock().pending_notifiers);
+ for notifier in notifiers {
+ let _ = notifier.send(());
+ }
+ }
+}
+
+struct SpillInActionGuard<'a, D: MemDomain> {
+ coordinator: &'a RuntimeSpillCoordinator<D>,
+ consumer: Weak<ConsumerInfo<D>>,
+}
+
+impl<D: MemDomain> Drop for SpillInActionGuard<'_, D> {
+ fn drop(&mut self) {
+ let notifiers = {
+ let mut state = self.coordinator.state.lock();
+ state
+ .active_consumers
+ .retain(|active| !Weak::ptr_eq(active, &self.consumer));
+ std::mem::take(&mut state.pending_notifiers)
+ };
+ for notifier in notifiers {
+ let _ = notifier.send(());
+ }
+ }
+}
+
+fn elect_largest_spillable_consumer<D: MemDomain + RuntimeWatermarkEvaluator<D>>(
+ elected: &mut Vec<Weak<ConsumerInfo<D>>>,
+) {
+ let mm = RuntimeMemManager::<D>::get();
+ let consumers = mm.consumers.lock();
+
+ let mut largest: Option<(usize, &Arc<ConsumerInfo<D>>)> = None;
+ let mut num_spillable = 0usize;
+ let mut mem_spillable = 0usize;
+ let mut mem_unreclaimable = 0usize;
+ for consumer_info in consumers.iter() {
+ let status = *consumer_info.status.lock();
+ if !status.spillable || status.mem_used == 0 {
+ mem_unreclaimable += status.mem_used;
+ continue;
+ }
+ num_spillable += 1;
+ mem_spillable += status.mem_used;
+ if elected
+ .iter()
+ .any(|already| weak_ptr_eq(already, consumer_info))
+ {
+ continue;
+ }
+ if largest.is_none_or(|(mem_used, _)| status.mem_used > mem_used) {
+ largest = Some((status.mem_used, consumer_info));
+ }
+ }
+
+ if let Some((mem_used, consumer_info)) = largest {
+ log::info!(
+ "{} spill coordinator elected {} (mem_used: {}) to spill; \
+ spillable: {}/{} consumers holding {}, unreclaimable holding {}",
+ D::name(),
+ consumer_info.name,
+ ByteSize(mem_used as u64),
+ num_spillable,
+ consumers.len(),
+ ByteSize(mem_spillable as u64),
+ ByteSize(mem_unreclaimable as u64),
+ );
+ elected.push(Arc::downgrade(consumer_info));
+ } else {
+ static LAST_WARN: Mutex<Option<std::time::Instant>> = Mutex::new(None);
+ let mut last_warn = LAST_WARN.lock();
+ let now = std::time::Instant::now();
+ let should_warn = last_warn
+ .is_none_or(|prev| now.duration_since(prev) >= std::time::Duration::from_secs(1));
+ if should_warn {
+ *last_warn = Some(now);
+ drop(last_warn);
+ log::warn!(
+ "{} spill coordinator found NO spillable consumer to elect: \
+ {} consumers registered, {} spillable holding {}, unreclaimable holding {}, \
+ already elected: {}, logical watermark: {:.3}",
+ D::name(),
+ consumers.len(),
+ num_spillable,
+ ByteSize(mem_spillable as u64),
+ ByteSize(mem_unreclaimable as u64),
+ elected.len(),
+ <D as RuntimeWatermarkEvaluator<D>>::logical_watermark(),
+ );
+ }
+ }
+}
+
+fn weak_ptr_eq<D: MemDomain>(weak: &Weak<ConsumerInfo<D>>, strong: &Arc<ConsumerInfo<D>>) -> bool {
+ std::ptr::eq(weak.as_ptr(), Arc::as_ptr(strong))
+}
+
+/// Combines logical usage with domain-specific physical pressure.
+pub trait RuntimeWatermarkEvaluator<D: MemDomain> {
+ fn physical_watermark() -> f64;
+ fn logical_watermark() -> f64;
+
+ fn watermark() -> f64 {
+ Self::physical_watermark().max(Self::logical_watermark())
+ }
+}
+
+pub async fn update_consumer_mem_used_with_custom_updater<D, F, Fut>(
+ consumer_name: &str,
+ consumer_info: &Arc<ConsumerInfo<D>>,
+ updater: impl Fn(&mut ConsumerStatus) -> (usize, usize),
+ forced: bool,
+ do_spill: F,
+) -> Result<()>
+where
+ D: MemDomain + RuntimeWatermarkEvaluator<D>,
+ F: FnOnce() -> Fut,
+ Fut: Future<Output = Result<()>>,
+{
+ let mm = RuntimeMemManager::<D>::get();
+
+ {
+ let mut status = mm.status.lock();
+ let mut consumer_status = consumer_info.status.lock();
+
+ let (old_used, new_used) = updater(&mut consumer_status);
+ let spillable = consumer_status.spillable;
+ let diff_used = new_used as isize - old_used as isize;
+ assert!(
+ !forced || spillable,
+ "forced spilling an unspillable {} consumer",
+ D::name()
+ );
+
+ status.update_total_used_with_diff(diff_used);
+
+ if spillable {
+ assert!(status.mem_spillables as isize + diff_used >= 0);
+ status.mem_spillables = (status.mem_spillables as isize + diff_used) as usize;
+ }
+ drop(consumer_status);
+ drop(status);
+
+ if new_used < old_used {
+ if let Some(coordinator) = D::coordinator_cell().get() {
+ coordinator.wake_waiters();
+ }
+ }
+ if !forced
+ && (new_used == 0
+ || !spillable
+ || new_used <= old_used
+ || new_used <= D::min_trigger_size())
+ {
+ return Ok(());
+ }
+ }
+
+ if forced {
+ return RuntimeSpillCoordinator::<D>::get()
+ .spill_in_action(consumer_name, consumer_info, do_spill)
+ .await;
+ }
+
+ spill_consumer_if_needed(consumer_name, consumer_info, do_spill).await
+}
+
+async fn spill_consumer_if_needed<D, F, Fut>(
+ consumer_name: &str,
+ consumer_info: &Arc<ConsumerInfo<D>>,
+ do_spill: F,
+) -> Result<()>
+where
+ D: MemDomain + RuntimeWatermarkEvaluator<D>,
+ F: FnOnce() -> Fut,
+ Fut: Future<Output = Result<()>>,
+{
+ let consumer_status = *consumer_info.status.lock();
+ if !consumer_status.spillable || consumer_status.mem_used == 0 {
+ return Ok(());
+ }
+ let coordinator = RuntimeSpillCoordinator::<D>::get();
+ loop {
+ match coordinator.update_consumer_mem_used(consumer_info) {
+ SpillAction::DoNotSpill => return Ok(()),
+ SpillAction::DoSpill => {
+ return coordinator
+ .spill_with_mode(consumer_name, consumer_info, do_spill, false)
+ .await;
+ }
+ SpillAction::WaitAndRetry(receiver) => {
+ wait_for_spill(receiver).await;
+ }
+ }
+ }
+}
diff --git a/native-engine/auron-memmgr/tests/runtime_spill_tests.rs b/native-engine/auron-memmgr/tests/runtime_spill_tests.rs
new file mode 100644
index 000000000..ffc18b3f0
--- /dev/null
+++ b/native-engine/auron-memmgr/tests/runtime_spill_tests.rs
@@ -0,0 +1,441 @@
+// 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.
+
+use std::{
+ future::Future,
+ sync::{
+ Arc,
+ atomic::{AtomicUsize, Ordering},
+ },
+ task::Context,
+};
+
+use async_trait::async_trait;
+use auron_memmgr::{
+ runtime_mem_manager::{ConsumerInfo, MemDomain, RuntimeMemManager},
+ runtime_spill_coordinator::{
+ RuntimeSpillCoordinator, RuntimeWatermarkEvaluator, SpillAction,
+ update_consumer_mem_used_with_custom_updater,
+ },
+};
+use futures::{FutureExt, channel::oneshot, executor::block_on, future::pending, task::noop_waker};
+use once_cell::sync::OnceCell;
+use parking_lot::Mutex;
+
+struct TestDomain;
+static TEST_LOCK: Mutex<()> = Mutex::new(());
+static PRESSURE: AtomicUsize = AtomicUsize::new(100);
+static INITIALIZATIONS: AtomicUsize = AtomicUsize::new(0);
+static PREPARATION: Mutex<Option<oneshot::Receiver<()>>> = Mutex::new(None);
+
+fn test_lock() -> parking_lot::MutexGuard<'static, ()> {
+ let lock = TEST_LOCK.lock();
+ PRESSURE.store(100, Ordering::SeqCst);
+ lock
+}
+#[async_trait]
+impl MemDomain for TestDomain {
+ async fn on_before_spill(_consumer_name: &str) {
+ let receiver = PREPARATION.lock().take();
+ if let Some(receiver) = receiver {
+ receiver.await.expect("finish spill preparation");
+ }
+ }
+
+ fn on_manager_initialized() {
+ INITIALIZATIONS.fetch_add(1, Ordering::SeqCst);
+ }
+ fn name() -> &'static str {
+ "test"
+ }
+ fn watermark_config() -> (f64, f64) {
+ (0.8, 0.9)
+ }
+ fn num_max_spilling_consumers() -> usize {
+ 1
+ }
+ fn manager_cell() -> &'static OnceCell<Arc<RuntimeMemManager<Self>>> {
+ static CELL: OnceCell<Arc<RuntimeMemManager<TestDomain>>> = OnceCell::new();
+ &CELL
+ }
+ fn coordinator_cell() -> &'static OnceCell<RuntimeSpillCoordinator<Self>> {
+ static CELL: OnceCell<RuntimeSpillCoordinator<TestDomain>> = OnceCell::new();
+ &CELL
+ }
+}
+impl RuntimeWatermarkEvaluator<TestDomain> for TestDomain {
+ fn physical_watermark() -> f64 {
+ PRESSURE.load(Ordering::SeqCst) as f64 / 100.0
+ }
+ fn logical_watermark() -> f64 {
+ 0.0
+ }
+}
+struct Consumer(Arc<ConsumerInfo<TestDomain>>);
+impl Consumer {
+ fn new(name: &str) -> Self {
+ RuntimeMemManager::<TestDomain>::init_with_total(1000);
+ Self(RuntimeMemManager::<TestDomain>::register_consumer_info(
+ name.to_owned(),
+ true,
+ ))
+ }
+}
+impl Drop for Consumer {
+ fn drop(&mut self) {
+ RuntimeSpillCoordinator::<TestDomain>::get().remove_consumer(&self.0);
+ RuntimeMemManager::<TestDomain>::deregister_consumer_info(&self.0.name, &self.0);
+ }
+}
+#[test]
+fn active_spill_counts_toward_concurrency_limit() {
+ let _lock = test_lock();
+ let small = Consumer::new("small");
+ let large = Consumer::new("large");
+ // Directly account reservations without triggering a spill.
+ small.0.status.lock().mem_used = 100;
+ large.0.status.lock().mem_used = 200;
+ let mm = RuntimeMemManager::<TestDomain>::get();
+ mm.status.lock().total_used = 300;
+ mm.status.lock().mem_spillables = 300;
+ let coordinator = RuntimeSpillCoordinator::<TestDomain>::get();
+ let mut active = Box::pin(coordinator.spill_in_action("small", &small.0, || pending()));
+ let waker = noop_waker();
+ assert!(
+ active
+ .as_mut()
+ .poll(&mut Context::from_waker(&waker))
+ .is_pending()
+ );
+ assert!(
+ matches!(
+ coordinator.update_consumer_mem_used(&large.0),
+ SpillAction::WaitAndRetry(_)
+ ),
+ "a selected consumer must not bypass an active spill"
+ );
+}
+#[test]
+fn growing_report_inside_spill_does_not_reenter_spill() {
+ let _lock = test_lock();
+ let consumer = Consumer::new("recursive");
+ let nested_spills = AtomicUsize::new(0);
+ block_on(update_consumer_mem_used_with_custom_updater(
+ "recursive",
+ &consumer.0,
+ |status| {
+ let old = std::mem::replace(&mut status.mem_used, 100);
+ (old, 100)
+ },
+ false,
+ || async {
+ update_consumer_mem_used_with_custom_updater(
+ "recursive",
+ &consumer.0,
+ |status| {
+ let old = std::mem::replace(&mut status.mem_used, 200);
+ (old, 200)
+ },
+ false,
+ || async {
+ nested_spills.fetch_add(1, Ordering::SeqCst);
+ Ok(())
+ },
+ )
+ .await
+ },
+ ))
+ .expect("usage report");
+ assert_eq!(
+ nested_spills.load(Ordering::SeqCst),
+ 0,
+ "spill callbacks must not recursively spill their own consumer"
+ );
+}
+
+fn account(consumer: &Consumer, used: usize) {
+ let manager = RuntimeMemManager::<TestDomain>::get();
+ let mut status = manager.status.lock();
+ let mut consumer_status = consumer.0.status.lock();
+ status.total_used = status.total_used - consumer_status.mem_used + used;
+ if consumer_status.spillable {
+ status.mem_spillables = status.mem_spillables - consumer_status.mem_used + used;
+ }
+ consumer_status.mem_used = used;
+}
+
+#[test]
+fn cancellation_releases_slot_and_wakes_waiter() {
+ let _lock = test_lock();
+ let first = Consumer::new("first");
+ let second = Consumer::new("second");
+ account(&first, 100);
+ account(&second, 200);
+ let coordinator = RuntimeSpillCoordinator::<TestDomain>::get();
+ let mut active = Box::pin(coordinator.spill_in_action("first", &first.0, pending));
+ assert!(active.as_mut().now_or_never().is_none());
+ let waiter = match coordinator.update_consumer_mem_used(&second.0) {
+ SpillAction::WaitAndRetry(receiver) => receiver,
+ _ => unreachable!("expected waiter"),
+ };
+ drop(active);
+ assert!(
+ waiter
+ .now_or_never()
+ .expect("cancellation must wake waiter")
+ .is_ok()
+ );
+ let completed = coordinator
+ .spill_in_action("second", &second.0, || async { Ok(()) })
+ .now_or_never();
+ assert!(
+ completed
+ .expect("cancelled spill must release its slot")
+ .is_ok()
+ );
+}
+
+#[test]
+fn concurrent_admission_cannot_exceed_limit() {
+ let _lock = test_lock();
+ let first = Consumer::new("first");
+ let second = Consumer::new("second");
+ account(&first, 100);
+ account(&second, 200);
+ let coordinator = RuntimeSpillCoordinator::<TestDomain>::get();
+ // Both callers can pass the advisory check before either callback starts.
+ assert!(matches!(
+ coordinator.update_consumer_mem_used(&first.0),
+ SpillAction::DoSpill
+ ));
+ assert!(matches!(
+ coordinator.update_consumer_mem_used(&second.0),
+ SpillAction::DoSpill
+ ));
+ let mut active = Box::pin(coordinator.spill_in_action("first", &first.0, pending));
+ assert!(active.as_mut().now_or_never().is_none());
+ let started = AtomicUsize::new(0);
+ let mut next = Box::pin(coordinator.spill_in_action("second", &second.0, || async {
+ started.fetch_add(1, Ordering::SeqCst);
+ Ok(())
+ }));
+ assert!(next.as_mut().now_or_never().is_none());
+ assert_eq!(started.load(Ordering::SeqCst), 0);
+ drop(active);
+ assert!(
+ next.as_mut()
+ .now_or_never()
+ .expect("next spill starts after cancellation")
+ .is_ok()
+ );
+ assert_eq!(started.load(Ordering::SeqCst), 1);
+}
+
+#[test]
+fn pinning_waiting_consumer_skips_its_callback() {
+ let _lock = test_lock();
+ let first = Consumer::new("first");
+ let second = Consumer::new("second");
+ account(&first, 100);
+ account(&second, 200);
+ let coordinator = RuntimeSpillCoordinator::<TestDomain>::get();
+ let mut active = Box::pin(coordinator.spill_in_action("first", &first.0, pending));
+ assert!(active.as_mut().now_or_never().is_none());
+ let mut next = Box::pin(coordinator.spill_in_action("second", &second.0, || async {
+ unreachable!("a pinned consumer must not spill")
+ }));
+ assert!(next.as_mut().now_or_never().is_none());
+ RuntimeMemManager::<TestDomain>::set_consumer_spillable(&second.0, false);
+ assert!(
+ next.as_mut()
+ .now_or_never()
+ .expect("pinning wakes the reporter")
+ .is_ok()
+ );
+ assert_eq!(
+ RuntimeMemManager::<TestDomain>::get()
+ .status
+ .lock()
+ .mem_spillables,
+ 100
+ );
+}
+
+#[test]
+fn releasing_memory_wakes_reporters_before_spill_finishes() {
+ let _lock = test_lock();
+ let first = Consumer::new("first");
+ let second = Consumer::new("second");
+ account(&first, 100);
+ let coordinator = RuntimeSpillCoordinator::<TestDomain>::get();
+ let mut active = Box::pin(coordinator.spill_in_action("first", &first.0, pending));
+ assert!(active.as_mut().now_or_never().is_none());
+ let mut next = Box::pin(update_consumer_mem_used_with_custom_updater(
+ "second",
+ &second.0,
+ |status| {
+ let old = std::mem::replace(&mut status.mem_used, 200);
+ (old, 200)
+ },
+ false,
+ || async { unreachable!("pressure has already fallen") },
+ ));
+ assert!(next.as_mut().now_or_never().is_none());
+ PRESSURE.store(50, Ordering::SeqCst);
+ block_on(update_consumer_mem_used_with_custom_updater(
+ "first",
+ &first.0,
+ |status| {
+ let old = std::mem::replace(&mut status.mem_used, 0);
+ (old, 0)
+ },
+ false,
+ || async { unreachable!("releasing memory must not spill") },
+ ))
+ .expect("release");
+ assert!(
+ next.as_mut()
+ .now_or_never()
+ .expect("memory release wakes waiter")
+ .is_ok()
+ );
+}
+
+#[test]
+fn low_watermark_elects_largest_and_respects_active_spill() {
+ let _lock = test_lock();
+ let small = Consumer::new("small");
+ let large = Consumer::new("large");
+ account(&small, 100);
+ account(&large, 200);
+ let coordinator = RuntimeSpillCoordinator::<TestDomain>::get();
+ PRESSURE.store(79, Ordering::SeqCst);
+ assert!(matches!(
+ coordinator.update_consumer_mem_used(&large.0),
+ SpillAction::DoNotSpill
+ ));
+ PRESSURE.store(85, Ordering::SeqCst);
+ assert!(matches!(
+ coordinator.update_consumer_mem_used(&small.0),
+ SpillAction::DoNotSpill
+ ));
+ assert!(matches!(
+ coordinator.update_consumer_mem_used(&large.0),
+ SpillAction::DoSpill
+ ));
+ let mut active = Box::pin(coordinator.spill_in_action("small", &small.0, pending));
+ assert!(active.as_mut().now_or_never().is_none());
+ assert!(matches!(
+ coordinator.update_consumer_mem_used(&large.0),
+ SpillAction::DoNotSpill
+ ));
+ drop(active);
+ account(&large, 0);
+ assert!(matches!(
+ coordinator.update_consumer_mem_used(&small.0),
+ SpillAction::DoSpill
+ ));
+}
+
+#[test]
+fn spill_error_releases_slot() {
+ let _lock = test_lock();
+ let first = Consumer::new("first");
+ let second = Consumer::new("second");
+ let coordinator = RuntimeSpillCoordinator::<TestDomain>::get();
+ assert!(
+ block_on(coordinator.spill_in_action("first", &first.0, || async {
+ Err(datafusion::common::DataFusionError::Execution(
+ "injected spill error".into(),
+ ))
+ }))
+ .is_err()
+ );
+ assert!(
+ coordinator
+ .spill_in_action("second", &second.0, || async { Ok(()) })
+ .now_or_never()
+ .expect("failed spill must release slot")
+ .is_ok()
+ );
+}
+
+#[test]
+fn initialization_hook_runs_once() {
+ let _lock = test_lock();
+ RuntimeMemManager::<TestDomain>::init_with_total(1000);
+ RuntimeMemManager::<TestDomain>::init_with_total(2000);
+ assert_eq!(INITIALIZATIONS.load(Ordering::SeqCst), 1);
+ assert_eq!(RuntimeMemManager::<TestDomain>::get().total, 1000);
+}
+
+fn spill_after_preparation(forced: bool, relieve_pressure: impl FnOnce(&Consumer)) -> usize {
+ let consumer = Consumer::new("preparing");
+ let spills = AtomicUsize::new(0);
+ let (sender, receiver) = oneshot::channel();
+ *PREPARATION.lock() = Some(receiver);
+ let mut report = Box::pin(update_consumer_mem_used_with_custom_updater(
+ "preparing",
+ &consumer.0,
+ |status| {
+ let old = std::mem::replace(&mut status.mem_used, 100);
+ (old, 100)
+ },
+ forced,
+ || async {
+ spills.fetch_add(1, Ordering::SeqCst);
+ Ok(())
+ },
+ ));
+ assert!(report.as_mut().now_or_never().is_none());
+ relieve_pressure(&consumer);
+ sender.send(()).expect("resume spill preparation");
+ report
+ .as_mut()
+ .now_or_never()
+ .expect("report finishes")
+ .expect("spill decision");
+ spills.load(Ordering::SeqCst)
+}
+
+#[test]
+fn automatic_spill_skips_empty_consumer_after_preparation() {
+ let _lock = test_lock();
+ assert_eq!(
+ spill_after_preparation(false, |consumer| account(consumer, 0)),
+ 0
+ );
+}
+
+#[test]
+fn automatic_spill_rechecks_pressure_after_preparation() {
+ let _lock = test_lock();
+ assert_eq!(
+ spill_after_preparation(false, |_| PRESSURE.store(50, Ordering::SeqCst)),
+ 0
+ );
+}
+
+#[test]
+fn forced_spill_bypasses_pressure_after_preparation() {
+ let _lock = test_lock();
+ assert_eq!(
+ spill_after_preparation(true, |consumer| {
+ account(consumer, 0);
+ PRESSURE.store(50, Ordering::SeqCst);
+ }),
+ 1
+ );
+}